custom shift

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)