# 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
```