custom shift

.gitignore

@@ -179,6 +179,7 @@ cython_debug/
 !/output/.gitkeep
 /extensions/*
 !/extensions/example
+!/extensions/custom_shift
 /temp
 /wandb
 .vscode/settings.json

extensions/custom_shift/README.md

@@ -0,0 +1,60 @@
+# Custom Shift
+
+Independent shifted logit-normal training timestep sampling. Musubi is an
+algorithm reference only: this extension does not import it or require its
+repository, packages, or installation.
+
+Select **Custom Shift** under **Timestep Type** in the training UI. Its five
+fields are saved with the job and retained when switching to another type.
+Alternatively:
+
+```yaml
+train:
+  noise_scheduler: flowmatch
+  timestep_type: custom_shift
+  timestep_sampler:
+    discrete_flow_shift: 1.0
+    sigmoid_scale: 1.0
+    min_timestep: 0
+    max_timestep: 1000
+    preserve_distribution_shape: false
+```
+
+All five settings are optional and default to the values above. Shift and scale
+must be finite and positive. Bounds must be integers satisfying
+`0 <= min_timestep < max_timestep <= 1000`.
+
+For `z ~ Normal(0, 1)`, draw `u = sigmoid(sigmoid_scale * z)`, then
+`t = shift * u / (1 + (shift - 1) * u)`. The implementation uses the numerically
+stable equivalent `sigmoid(sigmoid_scale * z + log(shift))`. Model timesteps are
+`1000 * t`, where 1000 is pure noise. The original **Shift** type is unchanged.
+
+- With shape preservation off, scale the distribution into the configured range.
+- With shape preservation on, reject values outside the range, yielding the
+  original distribution conditioned on that range. No boundary clamping is used.
+- For multistage models, further condition the final distribution on the current
+  expert's range. Empty intersections fail explicitly. Rejection sampling fails
+  after 1000 attempts rather than returning values outside the requested range.
+- Timestep Bias (`content_or_style`, including regularization data) is ignored.
+  Sampling applies to both ordinary and regularization batches.
+
+Supported path: the standard ai-toolkit flowmatch scheduler (including compatible
+subclasses), with 1000 training timesteps and non-inverted sigmas. Inference
+sampling settings are independent. Keep `min_denoising_steps/max_denoising_steps`
+at their defaults (0/999); use the new bounds instead.
+
+Discrete table weighting (`linear_timesteps`, `linear_timesteps2`),
+`first_timestep_chance`, SNR weighting, turbo/refiner training, source/unaugmented
+loss targets, diffusion feature extractor losses, and stepped/mean-flow losses are rejected because they require
+different timestep handling. These checks apply only to `custom_shift`.
+
+The integration is deliberately small: `TrainConfig` activates the extension,
+`BaseSDTrainProcess` requests final per-batch timesteps, and guidance initializes
+its scheduler with the resolved base type. The scheduler and model implementations
+are not patched or replaced. A `.gitignore` exception keeps this extension tracked.
+
+CPU validation (requires torch, diffusers and numpy):
+
+```shell
+python -m unittest testing.test_custom_shift -v
+```

extensions/custom_shift/__init__.py

@@ -0,0 +1,3 @@
+"""Independent custom-shift timestep sampling for ai-toolkit."""
+
+from .sampler import CustomShiftSampler

extensions/custom_shift/sampler.py

@@ -0,0 +1,109 @@
+"""Shifted logit-normal sampling. No dependency on musubi-tuner."""
+
+import math
+from collections.abc import Mapping
+
+import torch
+
+
+DEFAULTS = {
+    "discrete_flow_shift": 1.0,
+    "sigmoid_scale": 1.0,
+    "min_timestep": 0,
+    "max_timestep": 1000,
+    "preserve_distribution_shape": False,
+}
+
+
+class CustomShiftSampler:
+    def __init__(self, train_config):
+        raw = train_config.timestep_sampler
+        if raw is None:
+            raw = {}
+        if not isinstance(raw, Mapping):
+            raise ValueError("train.timestep_sampler must be a mapping for custom_shift")
+        unknown = set(raw) - set(DEFAULTS)
+        if unknown:
+            raise ValueError(f"Unknown custom_shift options: {sorted(unknown)}")
+        self.config = {**DEFAULTS, **raw}
+        for key in ("discrete_flow_shift", "sigmoid_scale", "min_timestep", "max_timestep"):
+            value = self.config[key]
+            if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
+                raise ValueError(f"custom_shift {key} must be a finite number")
+        for key in ("discrete_flow_shift", "sigmoid_scale"):
+            if self.config[key] <= 0:
+                raise ValueError(f"custom_shift {key} must be greater than zero")
+        lo, hi = self.config["min_timestep"], self.config["max_timestep"]
+        if not (0 <= lo < hi <= 1000) or int(lo) != lo or int(hi) != hi:
+            raise ValueError("custom_shift requires integer timesteps: 0 <= min_timestep < max_timestep <= 1000")
+        if not isinstance(self.config["preserve_distribution_shape"], bool):
+            raise ValueError("custom_shift preserve_distribution_shape must be a boolean")
+        if train_config.noise_scheduler != "flowmatch" or train_config.num_train_timesteps != 1000:
+            raise ValueError("custom_shift requires noise_scheduler: flowmatch and num_train_timesteps: 1000")
+        for key in ("linear_timesteps", "linear_timesteps2", "first_timestep_chance", "train_turbo",
+                    "learnable_snr_gos", "snr_gamma", "min_snr_gamma"):
+            if getattr(train_config, key, False):
+                raise ValueError(f"custom_shift is incompatible with {key}; disable it")
+        if train_config.min_denoising_steps != 0 or train_config.max_denoising_steps != 999:
+            raise ValueError("custom_shift uses min_timestep/max_timestep; leave min/max_denoising_steps at 0/999")
+        if train_config.loss_target in ("source", "unaugmented") or train_config.loss_type in ("stepped", "mean_flow"):
+            raise ValueError("custom_shift does not support losses requiring discrete scheduler lookup")
+        if getattr(train_config, "diffusion_feature_extractor_path", None) is not None:
+            raise ValueError("custom_shift does not support diffusion feature extractor losses")
+
+    def sample(self, count, device, *, stage_range=None):
+        """Sample final float32 timesteps; stage ranges condition, never rescale."""
+        lo = self.config["min_timestep"] / 1000.0
+        hi = self.config["max_timestep"] / 1000.0
+        lower, upper = lo, hi
+        if stage_range is not None:
+            stage_min, stage_max = stage_range
+            if not 0 <= stage_min < stage_max <= 1:
+                raise ValueError("custom_shift stage range must satisfy 0 <= min < max <= 1")
+            lower, upper = max(lo, stage_min), min(hi, stage_max)
+        if lower >= upper:
+            raise ValueError("custom_shift timestep range has no overlap with the current model stage")
+        if count < 1:
+            raise ValueError("custom_shift batch size must be positive")
+        scale = self.config["sigmoid_scale"]
+        log_shift = math.log(self.config["discrete_flow_shift"])
+        preserve = self.config["preserve_distribution_shape"]
+        draw_dtype = torch.float32 if torch.device(device).type == "mps" else torch.float64
+
+        def draw(n):
+            # Equivalent to s*u/(1+(s-1)*u), u=sigmoid(scale*z).
+            # Float64 keeps large valid shifts stable; output stays float32.
+            t = torch.sigmoid(scale * torch.randn(n, device=device, dtype=draw_dtype) + log_shift)
+            if not preserve:
+                t = lo + (hi - lo) * t
+            return t
+
+        accepted = []
+        remaining = count
+        for _ in range(1000):
+            candidates = draw(max(count, remaining))
+            valid = candidates[(candidates >= lower) & (candidates <= upper)]
+            taken = valid[:remaining]
+            accepted.append(taken)
+            remaining -= taken.numel()
+            if remaining == 0:
+                return (torch.cat(accepted) * 1000.0).to(torch.float32)
+        raise RuntimeError(
+            f"custom_shift could not sample {count} timesteps in [{lower * 1000}, {upper * 1000}] "
+            "after 1000 attempts; widen the range or adjust shift/scale"
+        )
+
+    def sample_for_model(self, model, count, device, boundary_index=0):
+        from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
+
+        if not model.is_flow_matching or not isinstance(model.noise_scheduler, CustomFlowMatchEulerDiscreteScheduler):
+            raise ValueError("custom_shift requires the standard ai-toolkit flowmatch scheduler")
+        if model.noise_scheduler.config.num_train_timesteps != 1000:
+            raise ValueError("custom_shift requires a scheduler with 1000 training timesteps")
+        if getattr(model.noise_scheduler.config, "invert_sigmas", False):
+            raise ValueError("custom_shift does not support inverted-sigma schedulers")
+        stage_range = None
+        if model.is_multistage:
+            boundaries = [1] + model.multistage_boundaries
+            stage_range = (boundaries[boundary_index + 1], boundaries[boundary_index])
+        return self.sample(count, device, stage_range=stage_range)

jobs/process/BaseSDTrainProcess.py

@@ -116,6 +116,8 @@ class BaseSDTrainProcess(BaseTrainProcess):
         # update modelconfig dtype to match train
         model_config['dtype'] = self.train_config.dtype
         self.model_config = ModelConfig(**model_config)
+        if self.train_config.custom_shift_sampler is not None and self.model_config.refiner_name_or_path is not None:
+            raise ValueError("custom_shift does not support refiner training")
 
         self.save_config = SaveConfig(**self.get_conf('save', {}))
         self.sample_config = SampleConfig(**self.get_conf('sample', {}))
@@ -1183,7 +1185,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
                     
                     timestep_type = 'linear' if linear_timesteps else None
                     if timestep_type is None:
-                        timestep_type = self.train_config.timestep_type
+                        timestep_type = self.train_config.scheduler_timestep_type
                     
                     if self.train_config.timestep_type == 'next_sample':
                         # simulate a sample
@@ -1208,7 +1210,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
                     self.sd.noise_scheduler.set_timesteps(
                         num_train_timesteps, device=self.device_torch
                     )
-            if self.sd.is_multistage:
+            if self.sd.is_multistage and self.train_config.custom_shift_sampler is None:
                 with self.timer('adjust_multistage_timesteps'):
                     # get our current sample range
                     boundaries = [1] + self.sd.multistage_boundaries
@@ -1232,7 +1234,10 @@ class BaseSDTrainProcess(BaseTrainProcess):
                 if is_reg:
                     content_or_style = self.train_config.content_or_style_reg
 
-                if self.train_config.timestep_type in ['two_step', 'four_step', 'eight_step']:
+                if self.train_config.custom_shift_sampler is not None:
+                    # Final continuous timesteps are drawn below, without index bias.
+                    timestep_indices = torch.zeros((batch_size,), device=self.device_torch, dtype=torch.long)
+                elif self.train_config.timestep_type in ['two_step', 'four_step', 'eight_step']:
                     if self.train_config.timestep_type == 'two_step':
                         indice_choices = [0, 499]
                     elif self.train_config.timestep_type == 'four_step':
@@ -1306,7 +1311,13 @@ class BaseSDTrainProcess(BaseTrainProcess):
                     timestep_indices = torch.where(force_first, torch.zeros_like(timestep_indices), timestep_indices)
             with self.timer('convert_timestep_indices_to_timesteps'):
                 # convert the timestep_indices to a timestep
-                timesteps = self.sd.noise_scheduler.timesteps[timestep_indices.long()]
+                if self.train_config.custom_shift_sampler is not None:
+                    timesteps = self.train_config.custom_shift_sampler.sample_for_model(
+                        self.sd, batch_size, self.device_torch,
+                        boundary_index=getattr(self, 'current_boundary_index', 0),
+                    )
+                else:
+                    timesteps = self.sd.noise_scheduler.timesteps[timestep_indices.long()]
                 
             with self.timer('prepare_noise'):
                 # get noise

testing/test_custom_shift.py

@@ -0,0 +1,222 @@
+"""CPU checks: python -m unittest testing.test_custom_shift -v"""
+
+import ast
+from contextlib import nullcontext
+import math
+from pathlib import Path
+from types import SimpleNamespace
+import typing
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from extensions.custom_shift import CustomShiftSampler
+from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
+
+
+def load_train_config():
+    # Execute the real config classes without loading unrelated audio/model stacks.
+    path = Path(__file__).resolve().parents[1] / "toolkit" / "config_modules.py"
+    tree = ast.parse(path.read_text(encoding="utf-8"))
+    classes = [node for node in tree.body if isinstance(node, ast.ClassDef)
+               and node.name in {"TrainConfig", "EMAConfig", "ValidationConfig", "ValidationItem"}]
+    namespace = dict(vars(typing))
+    exec(compile(ast.Module(body=classes, type_ignores=[]), str(path), "exec"), namespace)
+    return namespace["TrainConfig"]
+
+
+TrainConfig = load_train_config()
+
+
+def load_prepare_timesteps():
+    # Exercise the actual training preparation block with a real scheduler and
+    # synthetic model context, without loading models or the training dependency stack.
+    path = Path(__file__).resolve().parents[1] / "jobs" / "process" / "BaseSDTrainProcess.py"
+    tree = ast.parse(path.read_text(encoding="utf-8"))
+
+    def timer_name(node):
+        if isinstance(node, ast.With):
+            call = node.items[0].context_expr
+            if isinstance(call, ast.Call) and isinstance(call.func, ast.Attribute) and call.func.attr == "timer":
+                return call.args[0].value
+        return None
+
+    for parent in ast.walk(tree):
+        for _, nodes in ast.iter_fields(parent):
+            if not isinstance(nodes, list):
+                continue
+            start = next((i for i, node in enumerate(nodes) if timer_name(node) == "prepare_scheduler"), None)
+            if start is None:
+                continue
+            stop = next(i for i in range(start, len(nodes)) if timer_name(nodes[i]) == "prepare_noise")
+            function = ast.parse("def prepare(self, batch, latents, is_reg): pass").body[0]
+            function.body = nodes[start:stop] + [ast.Return(value=ast.Name(id="timesteps", ctx=ast.Load()))]
+            module = ast.fix_missing_locations(ast.Module(body=[function], type_ignores=[]))
+            namespace = {"torch": torch}
+            exec(compile(module, str(path), "exec"), namespace)
+            return namespace["prepare"]
+    raise AssertionError("Training timestep preparation block not found")
+
+
+prepare_timesteps = load_prepare_timesteps()
+
+
+def sampler(**options):
+    return TrainConfig(noise_scheduler="flowmatch", timestep_type="custom_shift",
+                       timestep_sampler=options).custom_shift_sampler
+
+
+def cdf(t, shift=1.0, scale=1.0):
+    if t <= 0:
+        return 0.0
+    if t >= 1:
+        return 1.0
+    z = (math.log(t / (1 - t)) - math.log(shift)) / scale
+    return (1 + math.erf(z / math.sqrt(2))) / 2
+
+
+class CustomShiftTests(unittest.TestCase):
+    def setUp(self):
+        torch.manual_seed(42)
+
+    def test_original_types_stay_inactive_even_with_saved_options(self):
+        for mode in ("sigmoid", "linear", "shift", "weighted", "next_sample"):
+            config = TrainConfig(timestep_type=mode, timestep_sampler={"sigmoid_scale": -1})
+            self.assertIsNone(config.custom_shift_sampler)
+            self.assertEqual(config.scheduler_timestep_type, mode)
+
+    def test_custom_type_uses_linear_initialization(self):
+        config = TrainConfig(noise_scheduler="flowmatch", timestep_type="custom_shift")
+        self.assertEqual(config.timestep_type, "custom_shift")
+        self.assertEqual(config.scheduler_timestep_type, "linear")
+        self.assertIsInstance(config.custom_shift_sampler, CustomShiftSampler)
+
+    def test_matches_rational_shift_formula(self):
+        z = torch.tensor([-3., -1., 0., 1., 3.], dtype=torch.float64)
+        with patch("extensions.custom_shift.sampler.torch.randn", return_value=z):
+            actual = sampler(discrete_flow_shift=3.0, sigmoid_scale=1.7).sample(5, "cpu")
+        u = torch.sigmoid(z * 1.7)
+        expected = (3 * u / (1 + 2 * u) * 1000).float()
+        torch.testing.assert_close(actual, expected)
+        self.assertEqual(actual.dtype, torch.float32)
+
+    def test_scaled_distribution_matches_analytic_cdf(self):
+        values = sampler(discrete_flow_shift=3, sigmoid_scale=1.4,
+                         min_timestep=200, max_timestep=800).sample(50000, "cpu")
+        self.assertTrue(bool(((values >= 200) & (values <= 800)).all()))
+        for x in (300, 400, 600, 700):
+            self.assertAlmostEqual((values <= x).double().mean().item(),
+                                   cdf((x - 200) / 600, 3, 1.4), delta=0.01)
+
+    def test_preserved_distribution_matches_conditional_cdf(self):
+        values = sampler(discrete_flow_shift=3, sigmoid_scale=1.4, min_timestep=200,
+                         max_timestep=800, preserve_distribution_shape=True).sample(50000, "cpu")
+        self.assertTrue(bool(((values >= 200) & (values <= 800)).all()))
+        for x in (300, 400, 600, 700):
+            expected = (cdf(x / 1000, 3, 1.4) - cdf(.2, 3, 1.4)) / (cdf(.8, 3, 1.4) - cdf(.2, 3, 1.4))
+            self.assertAlmostEqual((values <= x).double().mean().item(), expected, delta=0.01)
+
+    def test_stage_conditions_scaled_distribution(self):
+        values = sampler(discrete_flow_shift=2, min_timestep=100, max_timestep=900).sample(
+            30000, "cpu", stage_range=(.6, 1))
+        self.assertTrue(bool(((values >= 600) & (values <= 900)).all()))
+        base_min = cdf((.6 - .1) / .8, 2)
+        expected = (cdf((.75 - .1) / .8, 2) - base_min) / (1 - base_min)
+        self.assertAlmostEqual((values <= 750).double().mean().item(), expected, delta=.015)
+
+    def test_no_stage_overlap_fails_before_drawing(self):
+        with patch("extensions.custom_shift.sampler.torch.randn") as draw:
+            with self.assertRaisesRegex(ValueError, "no overlap"):
+                sampler(min_timestep=0, max_timestep=500).sample(1, "cpu", stage_range=(.8, 1))
+            draw.assert_not_called()
+
+    def test_rejection_exhaustion_never_returns_out_of_range(self):
+        with patch("extensions.custom_shift.sampler.torch.randn", return_value=torch.zeros(1)):
+            with self.assertRaisesRegex(RuntimeError, "1000 attempts"):
+                sampler(min_timestep=900, max_timestep=1000,
+                        preserve_distribution_shape=True).sample(1, "cpu")
+
+    def test_invalid_options(self):
+        for options in ({"discrete_flow_shift": 0}, {"discrete_flow_shift": float("inf")},
+                        {"sigmoid_scale": float("nan")}, {"sigmoid_scale": -1},
+                        {"sigmoid_scale": True}, {"min_timestep": 1.5},
+                        {"min_timestep": 1000}, {"max_timestep": 1001},
+                        {"preserve_distribution_shape": "false"}, {"typo": 1}):
+            with self.subTest(options=options), self.assertRaises(ValueError):
+                sampler(**options)
+        with self.assertRaises(ValueError):
+            TrainConfig(noise_scheduler="flowmatch", timestep_type="custom_shift", timestep_sampler=[])
+
+    def test_incompatible_training_options(self):
+        for options in ({"noise_scheduler": "ddpm"}, {"num_train_timesteps": 500},
+                        {"linear_timesteps": True}, {"linear_timesteps2": True},
+                        {"first_timestep_chance": .1}, {"min_denoising_steps": 10},
+                        {"max_denoising_steps": 800}, {"loss_type": "stepped"},
+                        {"loss_type": "mean_flow"}, {"loss_target": "source"},
+                        {"diffusion_feature_extractor_path": "test"},
+                        {"train_turbo": True}, {"snr_gamma": 5}):
+            with self.subTest(options=options), self.assertRaises(ValueError):
+                TrainConfig(**{"noise_scheduler": "flowmatch", "timestep_type": "custom_shift", **options})
+
+    def test_model_stage_and_actual_scheduler_add_noise(self):
+        scheduler = CustomFlowMatchEulerDiscreteScheduler()
+        scheduler.set_train_timesteps(1000, "cpu", timestep_type="linear")
+        model = SimpleNamespace(is_flow_matching=True, is_multistage=True,
+                                multistage_boundaries=[.875, 0.0], noise_scheduler=scheduler)
+        for boundary_index, lower, upper in ((0, 875, 1000), (1, 0, 875)):
+            values = sampler().sample_for_model(model, 32, "cpu", boundary_index)
+            self.assertTrue(bool(((values >= lower) & (values <= upper)).all()))
+            clean = torch.ones(32, 4, 2, 2)
+            noise = torch.zeros_like(clean)
+            noisy = scheduler.add_noise(clean, noise, values.view(-1, 1, 1, 1))
+            expected = (1 - values / 1000).view(-1, 1, 1, 1).expand_as(clean)
+            torch.testing.assert_close(noisy, expected)
+
+    def test_custom_draw_does_not_change_inference_scheduler(self):
+        scheduler = CustomFlowMatchEulerDiscreteScheduler()
+        scheduler.set_timesteps(20)
+        times, sigmas = scheduler.timesteps.clone(), scheduler.sigmas.clone()
+        model = SimpleNamespace(is_flow_matching=True, is_multistage=False, noise_scheduler=scheduler)
+        sampler().sample_for_model(model, 4, "cpu")
+        torch.testing.assert_close(scheduler.timesteps, times)
+        torch.testing.assert_close(scheduler.sigmas, sigmas)
+
+    def test_scheduler_capability_checks(self):
+        for scheduler in (object(), CustomFlowMatchEulerDiscreteScheduler(invert_sigmas=True),
+                          CustomFlowMatchEulerDiscreteScheduler(num_train_timesteps=500)):
+            model = SimpleNamespace(is_flow_matching=True, is_multistage=False, noise_scheduler=scheduler)
+            with self.assertRaises(ValueError):
+                sampler().sample_for_model(model, 1, "cpu")
+
+    def test_training_entry_point_ignores_bias_and_respects_expert(self):
+        batch = SimpleNamespace(file_items=[None] * 16)
+        latents = torch.zeros(16, 4, 4, 4)
+        for is_reg in (False, True):
+            for stage, lo, hi in ((0, 875, 1000), (1, 200, 875)):
+                config = TrainConfig(noise_scheduler="flowmatch", timestep_type="custom_shift",
+                                     content_or_style="style", timestep_sampler={"min_timestep": 200})
+                model = SimpleNamespace(is_flow_matching=True, is_multistage=True, is_flux=True, arch="flux",
+                                        multistage_boundaries=[.875, 0],
+                                        noise_scheduler=CustomFlowMatchEulerDiscreteScheduler())
+                process = SimpleNamespace(train_config=config, sd=model, current_boundary_index=stage,
+                                          device_torch=torch.device("cpu"), timer=lambda _: nullcontext(),
+                                          model_config=SimpleNamespace(refiner_name_or_path=None))
+                values = prepare_timesteps(process, batch, latents, is_reg)
+                self.assertTrue(bool(((values >= lo) & (values <= hi)).all()))
+                self.assertEqual(values.shape, (16,))
+                self.assertEqual(model.noise_scheduler.timestep_type, "linear")
+
+    def test_original_training_entry_point_keeps_table_sampling(self):
+        config = TrainConfig(noise_scheduler="flowmatch", timestep_type="shift")
+        model = SimpleNamespace(is_flow_matching=True, is_multistage=False, is_flux=True, arch="flux",
+                                noise_scheduler=CustomFlowMatchEulerDiscreteScheduler(shift=3))
+        process = SimpleNamespace(train_config=config, sd=model, device_torch=torch.device("cpu"),
+                                  timer=lambda _: nullcontext(), model_config=SimpleNamespace(refiner_name_or_path=None))
+        values = prepare_timesteps(process, SimpleNamespace(file_items=[None] * 8), torch.zeros(8, 4, 4, 4), False)
+        self.assertEqual(model.noise_scheduler.timestep_type, "shift")
+        self.assertTrue(bool(torch.isin(values, model.noise_scheduler.timesteps).all()))
+
+
+if __name__ == "__main__":
+    unittest.main()

toolkit/config_modules.py

@@ -626,6 +626,14 @@ class TrainConfig:
         if validation is not None:
             self.validation_config: ValidationConfig = ValidationConfig(**validation)
 
+        self.timestep_sampler = kwargs.get('timestep_sampler', None)
+        self.custom_shift_sampler = None
+        self.scheduler_timestep_type = self.timestep_type
+        if self.timestep_type == 'custom_shift':
+            from extensions.custom_shift import CustomShiftSampler
+            self.custom_shift_sampler = CustomShiftSampler(self)
+            self.scheduler_timestep_type = 'linear'
+
 
 ModelArch = Literal['sd1', 'sd2', 'sd3', 'sdxl', 'pixart', 'pixart_sigma', 'auraflow', 'flux', 'flex1', 'flex2', 'lumina2', 'vega', 'ssd', 'wan21', 'anima']
 

toolkit/guidance.py

@@ -437,7 +437,7 @@ def get_guided_loss_polarity(
             
             timestep_type = 'linear' if linear_timesteps else None
             if timestep_type is None:
-                timestep_type = train_config.timestep_type
+                timestep_type = train_config.scheduler_timestep_type
             
             sd.noise_scheduler.set_train_timesteps(
                 1000,

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

@@ -22,6 +22,7 @@ import {
   CreatableSelectInput,
 } from '@/components/formInputs';
 import Card from '@/components/Card';
+import CustomShiftFields, { customShiftDefaults } from '@/components/CustomShiftFields';
 import { X, Copy, Wand2, SquareDashed, Info } from 'lucide-react';
 import { openDoc } from '@/components/DocModal';
 import { openUpsamplePromptsModal, toAspectRatio } from '@/components/UpsamplePromptsModal';
@@ -776,17 +777,36 @@ export default function SimpleJob({
                       label="Timestep Type"
                       value={jobConfig.config.process[0].train.timestep_type}
                       disabled={disableSections.includes('train.timestep_type') || false}
-                      onChange={value => setJobConfig(value, 'config.process[0].train.timestep_type')}
+                      onChange={value => {
+                        if (value === 'custom_shift') {
+                          setJobConfig(
+                            { ...customShiftDefaults, ...jobConfig.config.process[0].train.timestep_sampler },
+                            'config.process[0].train.timestep_sampler',
+                          );
+                        }
+                        setJobConfig(value, 'config.process[0].train.timestep_type');
+                      }}
                       options={[
                         { value: 'sigmoid', label: 'Sigmoid' },
                         { value: 'linear', label: 'Linear' },
                         { value: 'shift', label: 'Shift' },
+                        ...(jobConfig.config.process[0].train.noise_scheduler === 'flowmatch' ||
+                        jobConfig.config.process[0].train.timestep_type === 'custom_shift'
+                          ? [{ value: 'custom_shift', label: 'Custom Shift' }] : []),
                         { value: 'weighted', label: 'Weighted' },
                       ]}
                     />
                   )}
+                  {!disableSections.includes('train.timestep_type') &&
+                    jobConfig.config.process[0].train.timestep_type === 'custom_shift' && (
+                      <CustomShiftFields
+                        value={jobConfig.config.process[0].train.timestep_sampler}
+                        onChange={value => setJobConfig(value, 'config.process[0].train.timestep_sampler')}
+                      />
+                    )}
                   <SelectInput
                     label="Timestep Bias"
+                    disabled={jobConfig.config.process[0].train.timestep_type === 'custom_shift'}
                     className="pt-2"
                     value={jobConfig.config.process[0].train.content_or_style}
                     onChange={value => setJobConfig(value, 'config.process[0].train.content_or_style')}

ui/src/components/CustomShiftFields.tsx

@@ -0,0 +1,84 @@
+'use client';
+
+import { useEffect, useRef, useState } from 'react';
+import { Checkbox } from '@/components/formInputs';
+import { CustomShiftConfig } from '@/types';
+
+export const customShiftDefaults: Required<CustomShiftConfig> = {
+  discrete_flow_shift: 1.0,
+  sigmoid_scale: 1.0,
+  min_timestep: 0,
+  max_timestep: 1000,
+  preserve_distribution_shape: false,
+};
+
+function NumericField({
+  label, value, onChange, min, max, integer = false, error = '',
+}: {
+  label: string;
+  value: number;
+  onChange: (value: number) => void;
+  min: number;
+  max?: number;
+  integer?: boolean;
+  error?: string;
+}) {
+  const ref = useRef<HTMLInputElement>(null);
+  const [text, setText] = useState(String(value));
+  useEffect(() => setText(String(value)), [value]);
+  useEffect(() => { ref.current?.setCustomValidity(error); }, [error]);
+
+  return (
+    <label className="block pt-2 text-xs text-gray-300">
+      {label}
+      <input
+        ref={ref}
+        type="number"
+        className="mt-1 w-full text-sm px-3 py-1 bg-gray-950 dark:bg-gray-800 border border-gray-700 rounded-sm text-gray-100"
+        value={text}
+        onChange={event => {
+          const next = event.target.value;
+          setText(next);
+          if (next !== '' && Number.isFinite(Number(next))) onChange(Number(next));
+        }}
+        min={min}
+        max={max}
+        step={integer ? 1 : 'any'}
+        required
+      />
+    </label>
+  );
+}
+
+export default function CustomShiftFields({
+  value, onChange,
+}: {
+  value?: CustomShiftConfig;
+  onChange: (value: CustomShiftConfig) => void;
+}) {
+  const config = { ...customShiftDefaults, ...value };
+  const update = (key: keyof CustomShiftConfig, next: number | boolean) => onChange({ ...config, [key]: next });
+  const rangeError = config.min_timestep >= config.max_timestep ? 'Min Timestep must be less than Max Timestep.' : '';
+
+  return (
+    <div className="pt-2">
+      <NumericField label="Discrete Flow Shift" value={config.discrete_flow_shift} min={0}
+        error={config.discrete_flow_shift <= 0 ? 'Discrete Flow Shift must be greater than zero.' : ''}
+        onChange={next => update('discrete_flow_shift', next)} />
+      <NumericField label="Sigmoid Scale" value={config.sigmoid_scale} min={0}
+        error={config.sigmoid_scale <= 0 ? 'Sigmoid Scale must be greater than zero.' : ''}
+        onChange={next => update('sigmoid_scale', next)} />
+      <NumericField label="Min Timestep" value={config.min_timestep} min={0} max={999} integer
+        error={rangeError} onChange={next => update('min_timestep', next)} />
+      <NumericField label="Max Timestep" value={config.max_timestep} min={1} max={1000} integer
+        error={rangeError} onChange={next => update('max_timestep', next)} />
+      <Checkbox label="Preserve Distribution Shape" className="pt-3"
+        checked={config.preserve_distribution_shape}
+        onChange={next => update('preserve_distribution_shape', next)} />
+      <p className="pt-2 text-xs text-gray-400">
+        When enabled, samples outside the range are rejected. Otherwise, the distribution is scaled into the range.
+        Model stage boundaries still apply. Timestep Bias is ignored.
+      </p>
+    </div>
+  );
+}

ui/src/types.ts

@@ -155,6 +155,14 @@ export interface ValidationConfig {
   validation_sigmas?: number[];
 }
 
+export interface CustomShiftConfig {
+  discrete_flow_shift?: number;
+  sigmoid_scale?: number;
+  min_timestep?: number;
+  max_timestep?: number;
+  preserve_distribution_shape?: boolean;
+}
+
 export interface TrainConfig {
   batch_size: number;
   bypass_guidance_embedding?: boolean;
@@ -165,6 +173,7 @@ export interface TrainConfig {
   gradient_checkpointing: boolean;
   noise_scheduler: string;
   timestep_type: string;
+  timestep_sampler?: CustomShiftConfig;
   content_or_style: string;
   optimizer: string;
   lr: number;