From 8bbd0516676c96ec1717bf099fbe0ecee8af3365 Mon Sep 17 00:00:00 2001
From: fatalis <fatalis@users.noreply.github.com>
Date: Wed, 15 Jul 2026 12:15:50 -0500
Subject: [PATCH] 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    | 1 +
 .../modal/modal_train_lora_flux_schnell_24gb.yaml        | 1 +
 config/examples/train_flex_redux.yaml                    | 1 +
 config/examples/train_full_fine_tune_flex.yaml           | 1 +
 config/examples/train_full_fine_tune_lumina.yaml         | 1 +
 config/examples/train_lora_chroma_24gb.yaml              | 1 +
 config/examples/train_lora_flex2_24gb.yaml               | 1 +
 config/examples/train_lora_flex_24gb.yaml                | 1 +
 config/examples/train_lora_flux_24gb.yaml                | 1 +
 config/examples/train_lora_flux_kontext_24gb.yaml        | 1 +
 config/examples/train_lora_flux_schnell_24gb.yaml        | 1 +
 config/examples/train_lora_hidream_48.yaml               | 1 +
 config/examples/train_lora_lumina.yaml                   | 1 +
 config/examples/train_lora_omnigen2_24gb.yaml            | 1 +
 config/examples/train_lora_qwen_image_24gb.yaml          | 1 +
 .../examples/train_lora_qwen_image_edit_2509_32gb.yaml   | 1 +
 config/examples/train_lora_qwen_image_edit_32gb.yaml     | 1 +
 config/examples/train_lora_sd35_large_24gb.yaml          | 1 +
 config/examples/train_lora_wan21_14b_24gb.yaml           | 1 +
 config/examples/train_lora_wan21_1b_24gb.yaml            | 1 +
 config/examples/train_lora_wan22_14b_24gb.yaml           | 1 +
 .../advanced_generator/config/train.example.yaml         | 1 +
 .../concept_replacer/config/train.example.yaml           | 1 +
 .../config/train.example.yaml                            | 1 +
 extensions_built_in/sd_trainer/config/train.example.yaml | 1 +
 .../ultimate_slider_trainer/config/train.example.yaml    | 1 +
 jobs/process/BaseSDTrainProcess.py                       | 6 +++++-
 toolkit/config_modules.py                                | 1 +
 ui/src/app/jobs/new/SimpleJob.tsx                        | 9 +++++++++
 ui/src/helpers/defaultSamples.ts                         | 5 ++++-
 ui/src/types.ts                                          | 1 +
 31 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/config/examples/modal/modal_train_lora_flux_24gb.yaml b/config/examples/modal/modal_train_lora_flux_24gb.yaml
index 51873de..fae41b2 100644
--- a/config/examples/modal/modal_train_lora_flux_24gb.yaml
+++ b/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:
diff --git a/config/examples/modal/modal_train_lora_flux_schnell_24gb.yaml b/config/examples/modal/modal_train_lora_flux_schnell_24gb.yaml
index 6d1e964..85a969a 100644
--- a/config/examples/modal/modal_train_lora_flux_schnell_24gb.yaml
+++ b/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:
diff --git a/config/examples/train_flex_redux.yaml b/config/examples/train_flex_redux.yaml
index 918de84..40bd1f8 100644
--- a/config/examples/train_flex_redux.yaml
+++ b/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
diff --git a/config/examples/train_full_fine_tune_flex.yaml b/config/examples/train_full_fine_tune_flex.yaml
index f386a3b..107ee42 100644
--- a/config/examples/train_full_fine_tune_flex.yaml
+++ b/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:
diff --git a/config/examples/train_full_fine_tune_lumina.yaml b/config/examples/train_full_fine_tune_lumina.yaml
index 51a6173..5d4877b 100644
--- a/config/examples/train_full_fine_tune_lumina.yaml
+++ b/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:
diff --git a/config/examples/train_lora_chroma_24gb.yaml b/config/examples/train_lora_chroma_24gb.yaml
index 807f2bb..0332209 100644
--- a/config/examples/train_lora_chroma_24gb.yaml
+++ b/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:
diff --git a/config/examples/train_lora_flex2_24gb.yaml b/config/examples/train_lora_flex2_24gb.yaml
index cc69801..98b0d68 100644
--- a/config/examples/train_lora_flex2_24gb.yaml
+++ b/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:
diff --git a/config/examples/train_lora_flex_24gb.yaml b/config/examples/train_lora_flex_24gb.yaml
index 92bf849..aa65605 100644
--- a/config/examples/train_lora_flex_24gb.yaml
+++ b/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:
diff --git a/config/examples/train_lora_flux_24gb.yaml b/config/examples/train_lora_flux_24gb.yaml
index 8e29402..cdb7a73 100644
--- a/config/examples/train_lora_flux_24gb.yaml
+++ b/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:
diff --git a/config/examples/train_lora_flux_kontext_24gb.yaml b/config/examples/train_lora_flux_kontext_24gb.yaml
index d570da8..750fc9c 100644
--- a/config/examples/train_lora_flux_kontext_24gb.yaml
+++ b/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:
diff --git a/config/examples/train_lora_flux_schnell_24gb.yaml b/config/examples/train_lora_flux_schnell_24gb.yaml
index a4aef07..4485386 100644
--- a/config/examples/train_lora_flux_schnell_24gb.yaml
+++ b/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:
diff --git a/config/examples/train_lora_hidream_48.yaml b/config/examples/train_lora_hidream_48.yaml
index f915f4c..55a2d53 100644
--- a/config/examples/train_lora_hidream_48.yaml
+++ b/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:
diff --git a/config/examples/train_lora_lumina.yaml b/config/examples/train_lora_lumina.yaml
index e5d2d75..655a630 100644
--- a/config/examples/train_lora_lumina.yaml
+++ b/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:
diff --git a/config/examples/train_lora_omnigen2_24gb.yaml b/config/examples/train_lora_omnigen2_24gb.yaml
index 6eb1530..8257cfe 100644
--- a/config/examples/train_lora_omnigen2_24gb.yaml
+++ b/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:
diff --git a/config/examples/train_lora_qwen_image_24gb.yaml b/config/examples/train_lora_qwen_image_24gb.yaml
index e464020..d6ad3bb 100644
--- a/config/examples/train_lora_qwen_image_24gb.yaml
+++ b/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:
diff --git a/config/examples/train_lora_qwen_image_edit_2509_32gb.yaml b/config/examples/train_lora_qwen_image_edit_2509_32gb.yaml
index 845ac92..06d48c6 100644
--- a/config/examples/train_lora_qwen_image_edit_2509_32gb.yaml
+++ b/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
diff --git a/config/examples/train_lora_qwen_image_edit_32gb.yaml b/config/examples/train_lora_qwen_image_edit_32gb.yaml
index 81a999d..64deb10 100644
--- a/config/examples/train_lora_qwen_image_edit_32gb.yaml
+++ b/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:
diff --git a/config/examples/train_lora_sd35_large_24gb.yaml b/config/examples/train_lora_sd35_large_24gb.yaml
index e1766c3..29cbc40 100644
--- a/config/examples/train_lora_sd35_large_24gb.yaml
+++ b/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:
diff --git a/config/examples/train_lora_wan21_14b_24gb.yaml b/config/examples/train_lora_wan21_14b_24gb.yaml
index 32babd1..919c5df 100644
--- a/config/examples/train_lora_wan21_14b_24gb.yaml
+++ b/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
diff --git a/config/examples/train_lora_wan21_1b_24gb.yaml b/config/examples/train_lora_wan21_1b_24gb.yaml
index 5e75f7a..1ff1bfb 100644
--- a/config/examples/train_lora_wan21_1b_24gb.yaml
+++ b/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
diff --git a/config/examples/train_lora_wan22_14b_24gb.yaml b/config/examples/train_lora_wan22_14b_24gb.yaml
index 966f184..a1c5952 100644
--- a/config/examples/train_lora_wan22_14b_24gb.yaml
+++ b/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
diff --git a/extensions_built_in/advanced_generator/config/train.example.yaml b/extensions_built_in/advanced_generator/config/train.example.yaml
index 793d5d5..95fb946 100644
--- a/extensions_built_in/advanced_generator/config/train.example.yaml
+++ b/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:
diff --git a/extensions_built_in/concept_replacer/config/train.example.yaml b/extensions_built_in/concept_replacer/config/train.example.yaml
index 793d5d5..95fb946 100644
--- a/extensions_built_in/concept_replacer/config/train.example.yaml
+++ b/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:
diff --git a/extensions_built_in/image_reference_slider_trainer/config/train.example.yaml b/extensions_built_in/image_reference_slider_trainer/config/train.example.yaml
index 8b0f473..1815927 100644
--- a/extensions_built_in/image_reference_slider_trainer/config/train.example.yaml
+++ b/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:
diff --git a/extensions_built_in/sd_trainer/config/train.example.yaml b/extensions_built_in/sd_trainer/config/train.example.yaml
index 793d5d5..95fb946 100644
--- a/extensions_built_in/sd_trainer/config/train.example.yaml
+++ b/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:
diff --git a/extensions_built_in/ultimate_slider_trainer/config/train.example.yaml b/extensions_built_in/ultimate_slider_trainer/config/train.example.yaml
index 8b0f473..1815927 100644
--- a/extensions_built_in/ultimate_slider_trainer/config/train.example.yaml
+++ b/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:
diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py
index 127a810..beb3686 100644
--- a/jobs/process/BaseSDTrainProcess.py
+++ b/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
 
diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py
index 9c9a04b..a84449b 100644
--- a/toolkit/config_modules.py
+++ b/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)
diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx
index 78d8580..e5d51b7 100644
--- a/ui/src/app/jobs/new/SimpleJob.tsx
+++ b/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"
diff --git a/ui/src/helpers/defaultSamples.ts b/ui/src/helpers/defaultSamples.ts
index 9b2e9f4..70881d9 100644
--- a/ui/src/helpers/defaultSamples.ts
+++ b/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: [
@@ -973,4 +976,4 @@ export const defaultIdeogramSamplesConfig: SampleConfig = {
   sample_steps: 30,
   num_frames: 1,
   fps: 1,
-}
\ No newline at end of file
+}
diff --git a/ui/src/types.ts b/ui/src/types.ts
index 0633cc5..07f070c 100644
--- a/ui/src/types.ts
+++ b/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[];