custom shift

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)