From 73cb12e532efb25b64024dcce4913ed6b9275301 Mon Sep 17 00:00:00 2001 From: wenyifan Date: Tue, 3 Nov 2026 17:18:35 +0800 Subject: [PATCH] custom shift --- .gitignore | 3 +- extensions/custom_shift/README.md | 60 +++++++ extensions/custom_shift/__init__.py | 3 + extensions/custom_shift/sampler.py | 109 ++++++++++++ jobs/process/BaseSDTrainProcess.py | 19 +- testing/test_custom_shift.py | 222 ++++++++++++++++++++++++ toolkit/config_modules.py | 10 +- toolkit/guidance.py | 2 +- ui/src/app/jobs/new/SimpleJob.tsx | 22 ++- ui/src/components/CustomShiftFields.tsx | 84 +++++++++ ui/src/types.ts | 9 + 11 files changed, 535 insertions(+), 8 deletions(-) create mode 100644 extensions/custom_shift/README.md create mode 100644 extensions/custom_shift/__init__.py create mode 100644 extensions/custom_shift/sampler.py create mode 100644 testing/test_custom_shift.py create mode 100644 ui/src/components/CustomShiftFields.tsx diff --git a/.gitignore b/.gitignore index 578838e..06e8732 100644 --- a/.gitignore +++ b/.gitignore @@ -179,6 +179,7 @@ cython_debug/ !/output/.gitkeep /extensions/* !/extensions/example +!/extensions/custom_shift /temp /wandb .vscode/settings.json @@ -195,4 +196,4 @@ original_repo testing/.model_test_outputs testing/.engine_test_outputs testing/.* -/issue_review \ No newline at end of file +/issue_review diff --git a/extensions/custom_shift/README.md b/extensions/custom_shift/README.md new file mode 100644 index 0000000..13b6304 --- /dev/null +++ b/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 +``` diff --git a/extensions/custom_shift/__init__.py b/extensions/custom_shift/__init__.py new file mode 100644 index 0000000..607e288 --- /dev/null +++ b/extensions/custom_shift/__init__.py @@ -0,0 +1,3 @@ +"""Independent custom-shift timestep sampling for ai-toolkit.""" + +from .sampler import CustomShiftSampler diff --git a/extensions/custom_shift/sampler.py b/extensions/custom_shift/sampler.py new file mode 100644 index 0000000..78b0d63 --- /dev/null +++ b/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) diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index dd395af..93e64cd 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/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 diff --git a/testing/test_custom_shift.py b/testing/test_custom_shift.py new file mode 100644 index 0000000..0ee5ab2 --- /dev/null +++ b/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() diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index c45a3ab..60ad487 100644 --- a/toolkit/config_modules.py +++ b/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'] @@ -1566,4 +1574,4 @@ def validate_configs( raise ValueError(f"Cannot cache unload text encoder with {model_config.arch} model. Control images are encoded with text embeddings. You can cache the text embeddings though") if train_config.diff_output_preservation and train_config.blank_prompt_preservation: - raise ValueError("Cannot use both differential output preservation and blank prompt preservation at the same time. Please set one of them to False.") \ No newline at end of file + raise ValueError("Cannot use both differential output preservation and blank prompt preservation at the same time. Please set one of them to False.") diff --git a/toolkit/guidance.py b/toolkit/guidance.py index 340b22e..4474016 100644 --- a/toolkit/guidance.py +++ b/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, diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index 9cc5e34..669401b 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/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' && ( + setJobConfig(value, 'config.process[0].train.timestep_sampler')} + /> + )} setJobConfig(value, 'config.process[0].train.content_or_style')} diff --git a/ui/src/components/CustomShiftFields.tsx b/ui/src/components/CustomShiftFields.tsx new file mode 100644 index 0000000..b47d2e6 --- /dev/null +++ b/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 = { + 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(null); + const [text, setText] = useState(String(value)); + useEffect(() => setText(String(value)), [value]); + useEffect(() => { ref.current?.setCustomValidity(error); }, [error]); + + return ( + + ); +} + +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 ( +
+ update('discrete_flow_shift', next)} /> + update('sigmoid_scale', next)} /> + update('min_timestep', next)} /> + update('max_timestep', next)} /> + update('preserve_distribution_shape', next)} /> +

+ 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. +

+
+ ); +} diff --git a/ui/src/types.ts b/ui/src/types.ts index b11d937..7350281 100644 --- a/ui/src/types.ts +++ b/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;