Add sample_start_step setting to configure when sampling starts (#949)

Co-authored-by: Jaret Burkett <jaretburkett@gmail.com>

config/examples/modal/modal_train_lora_flux_24gb.yaml

@@ -70,6 +70,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/modal/modal_train_lora_flux_schnell_24gb.yaml

@@ -72,6 +72,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_flex_redux.yaml

@@ -85,6 +85,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         # I leave half blank to test prompt and unprompted

config/examples/train_full_fine_tune_flex.yaml

@@ -81,6 +81,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_full_fine_tune_lumina.yaml

@@ -73,6 +73,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_chroma_24gb.yaml

@@ -78,6 +78,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_flex2_24gb.yaml

@@ -129,6 +129,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_flex_24gb.yaml

@@ -75,6 +75,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_flux_24gb.yaml

@@ -70,6 +70,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_flux_kontext_24gb.yaml

@@ -79,6 +79,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_flux_schnell_24gb.yaml

@@ -72,6 +72,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_hidream_48.yaml

@@ -86,6 +86,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_lumina.yaml

@@ -70,6 +70,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_omnigen2_24gb.yaml

@@ -68,6 +68,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_qwen_image_24gb.yaml

@@ -69,6 +69,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_qwen_image_edit_2509_32gb.yaml

@@ -70,6 +70,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         # you can provide up to 3 control images here

config/examples/train_lora_qwen_image_edit_32gb.yaml

@@ -68,6 +68,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         samples:

config/examples/train_lora_sd35_large_24gb.yaml

@@ -71,6 +71,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_wan21_14b_24gb.yaml

@@ -80,6 +80,7 @@ config:
       sample:
         sampler: "flowmatch"
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 832
         height: 480
         num_frames: 40

config/examples/train_lora_wan21_1b_24gb.yaml

@@ -69,6 +69,7 @@ config:
       sample:
         sampler: "flowmatch"
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 832
         height: 480
         num_frames: 40

config/examples/train_lora_wan22_14b_24gb.yaml

@@ -80,6 +80,7 @@ config:
       sample:
         sampler: "flowmatch"
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         # set to 1 for images

extensions_built_in/advanced_generator/config/train.example.yaml

@@ -52,6 +52,7 @@ config:
       sample:
         sampler: "ddpm" # must match train.noise_scheduler
         sample_every: 100 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 512
         height: 512
         prompts:

extensions_built_in/concept_replacer/config/train.example.yaml

@@ -52,6 +52,7 @@ config:
       sample:
         sampler: "ddpm" # must match train.noise_scheduler
         sample_every: 100 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 512
         height: 512
         prompts:

extensions_built_in/image_reference_slider_trainer/config/train.example.yaml

@@ -41,6 +41,7 @@ config:
       sample:
         sampler: "ddpm" # must match train.noise_scheduler
         sample_every: 100 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 512
         height: 512
         prompts:

extensions_built_in/sd_trainer/config/train.example.yaml

@@ -52,6 +52,7 @@ config:
       sample:
         sampler: "ddpm" # must match train.noise_scheduler
         sample_every: 100 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 512
         height: 512
         prompts:

extensions_built_in/ultimate_slider_trainer/config/train.example.yaml

@@ -41,6 +41,7 @@ config:
       sample:
         sampler: "ddpm" # must match train.noise_scheduler
         sample_every: 100 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 512
         height: 512
         prompts:

jobs/process/BaseSDTrainProcess.py

@@ -2371,7 +2371,11 @@ class BaseSDTrainProcess(BaseTrainProcess):
                 # todo improve this logic to send one of each through if we can buckets and batch size might be an issue
                 is_reg_step = False
                 is_save_step = self.save_config.save_every and self.step_num % self.save_config.save_every == 0
-                is_sample_step = self.sample_config.sample_every and self.step_num % self.sample_config.sample_every == 0
+                is_sample_step = (
+                    self.sample_config.sample_every
+                    and self.step_num >= self.sample_config.sample_start_step
+                    and self.step_num % self.sample_config.sample_every == 0
+                )
                 if self.train_config.disable_sampling:
                     is_sample_step = False
 

toolkit/config_modules.py

@@ -80,6 +80,7 @@ class SampleConfig:
     def __init__(self, **kwargs):
         self.sampler: str = kwargs.get('sampler', 'ddpm')
         self.sample_every: int = kwargs.get('sample_every', 100)
+        self.sample_start_step: int = kwargs.get('sample_start_step', 0)
         self.width: int = kwargs.get('width', 512)
         self.height: int = kwargs.get('height', 512)
         self.neg = kwargs.get('neg', False)

ui/src/app/jobs/new/SimpleJob.tsx

@@ -1182,6 +1182,15 @@ export default function SimpleJob({
                   min={1}
                   required
                 />
+                <NumberInput
+                  label="Sample Start Step"
+                  value={jobConfig.config.process[0].sample.sample_start_step ?? 0}
+                  onChange={value => setJobConfig(value, 'config.process[0].sample.sample_start_step')}
+                  placeholder="eg. 0"
+                  className="pt-2"
+                  min={0}
+                  required
+                />
                 <SelectInput
                   label="Sampler"
                   className="pt-2"

ui/src/helpers/defaultSamples.ts

@@ -3,6 +3,7 @@ import { SampleConfig } from "@/types";
 export const defaultSampleConfig: SampleConfig = {
   sampler: 'flowmatch',
   sample_every: 250,
+  sample_start_step: 0,
   width: 1024,
   height: 1024,
   samples: [
@@ -52,6 +53,7 @@ export const defaultSampleConfig: SampleConfig = {
 export const defaultAudioSampleConfig: SampleConfig = {
   sampler: 'flowmatch',
   sample_every: 250,
+  sample_start_step: 0,
   width: 1024,
   height: 1024,
   samples: [
@@ -406,6 +408,7 @@ It's Converging!
 export const defaultIdeogramSamplesConfig: SampleConfig = {
   sampler: 'flowmatch',
   sample_every: 250,
+  sample_start_step: 0,
   width: 1024,
   height: 1024,
   samples: [

ui/src/types.ts

@@ -201,6 +201,7 @@ export interface SampleItem {
 export interface SampleConfig {
   sampler: string;
   sample_every: number;
+  sample_start_step: number;
   width: number;
   height: number;
   prompts?: string[];