ai-toolkit/extensions/custom_shift/README.md

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:

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):

python -m unittest testing.test_custom_shift -v