110 lines
5.7 KiB
Python
110 lines
5.7 KiB
Python
"""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)
|