Diffusion¶

diffusion ¶

Diffusion primitives: timestep embeddings, schedulers, samplers, caching.

PerHeadAttentionCache ¶

PerHeadAttentionCache(num_heads: int, num_steps: int, rel_l1_thresh: float)

Stateful per-head attention output cache.

Returns a (num_heads,) bool decision per step. Inputs are assumed to have the head axis at position 1 — i.e. shape (B, num_heads, ...).

previous_output `property` ¶

previous_output: array

Last cached attention output. Raises before the first cache_output call.

reset ¶

reset() -> None

Clear all state. Call at the start of each new generation.

should_compute ¶

should_compute(step_index: int, attn_input: array) -> array

Per-head decide whether to recompute attention at step_index.

Side-effects: advances the stored previous input and per-head summary. Must be called once per step in order.

should_compute_from_summary ¶

should_compute_from_summary(step_index: int, summary: array) -> array

Decide per head using a caller-supplied (num_heads,) summary.

summary[h] is the analogue of the per-head delta ratio. Do not interleave with :meth:should_compute within a single denoising run: the two methods write semantically different values into the internal previous-summary slot. Pick one mode per run.

cache_output ¶

cache_output(output: array) -> None

Store the full (B, num_heads, ...) attention output for per-head splicing on skip.

PerLayerAttentionCache ¶

PerLayerAttentionCache(num_steps: int, rel_l1_thresh: float)

Stateful per-layer attention output cache.

previous_output `property` ¶

previous_output: array

Last cached attention output. Raises before the first cache_output call.

reset ¶

reset() -> None

Clear all state. Call at the start of each new generation.

should_compute ¶

should_compute(step_index: int, attn_input: array) -> bool

Decide whether to recompute attention at step_index.

Side-effects: advances the stored previous input and summary. Must be called once per step in order.

should_compute_from_summary ¶

should_compute_from_summary(step_index: int, summary: float) -> bool

Decide using a caller-supplied scalar summary instead of a tensor.

summary is the analogue of mean(abs(input - prev_input)) / mean(abs(prev_input)) — the caller has already done the math.

Do not interleave with :meth:should_compute within a single denoising run: the two methods write semantically different values into the internal previous-summary slot. Pick one mode per run.

cache_output ¶

cache_output(output: array) -> None

Store the attention output from the just-computed step for reuse on skip.

CFGSimilarityProfiler ¶

CFGSimilarityProfiler(num_blocks: int, num_heads: int, *, metric: Metric = 'cosine')

Per-(block, head) running-mean similarity accumulator.

Records the per-head similarity of cond and uncond attention outputs during a warmup pass and produces a static (num_blocks, num_heads) skip schedule.

scores `property` ¶

scores: array

(num_blocks, num_heads) running mean. Zero-count blocks → 0.0.

call_counts `property` ¶

call_counts: array

(num_blocks,) int32 count of :meth:record calls per block.

reset ¶

reset() -> None

Clear accumulated scores and counts. Call before a new warmup pass.

record ¶

record(block_idx: int, cond: array, uncond: array) -> None

Accumulate per-head cond/uncond similarity for block block_idx.

build_skip_schedule ¶

build_skip_schedule(threshold: float) -> array

Threshold :attr:scores into a (num_blocks, num_heads) skip mask.

CFGSkipController ¶

CFGSkipController(schedule: array)

Wraps a static (num_blocks, num_heads) bool skip schedule.

True at (b, h) means: for block b, head h skips the unconditional branch and reuses the conditional output. False means the head computes uncond normally.

num_blocks `property` ¶

num_blocks: int

Number of transformer blocks in the schedule.

num_heads `property` ¶

num_heads: int

Number of attention heads in the schedule.

from_profiler `classmethod` ¶

from_profiler(profiler: CFGSimilarityProfiler, threshold: float) -> 'CFGSkipController'

Build a controller from a profiler by thresholding its scores.

should_skip_uncond ¶

should_skip_uncond(block_idx: int) -> array

(num_heads,) bool mask — True heads skip uncond for this block.

apply ¶

apply(block_idx: int, cond_output: array, uncond_output: array) -> array

Apply the cached schedule to splice cond/uncond outputs (wraps :func:splice_heads).

DDIMScheduler ¶

DDIMScheduler(num_train_timesteps: int = 1000, beta_start: float = 0.00085, beta_end: float = 0.012, beta_schedule: BetaSchedule = 'scaled_linear', prediction_type: PredictionType = 'v_prediction', rescale_betas_zero_snr: bool = True, timestep_spacing: TimestepSpacing = 'trailing', set_alpha_to_one: bool = True, clip_sample: bool = False, num_inference_steps: int = 50)

Deterministic DDIM scheduler (eta=0).

Parameters:

Name	Type	Description	Default
`num_train_timesteps`	`int`	Number of diffusion steps used during training. Defaults to `1000`.	`1000`
`beta_start`	`float`	First beta value. Defaults to `0.00085` (Stable Diffusion convention).	`0.00085`
`beta_end`	`float`	Last beta value. Defaults to `0.012`.	`0.012`
`beta_schedule`	`BetaSchedule`	`"scaled_linear"` (default) or `"linear"`. `"scaled_linear"` matches Stable Diffusion's `linspace(sqrt(b0), sqrt(bT), N)**2`.	`'scaled_linear'`
`prediction_type`	`PredictionType`	`"epsilon"` (vanilla) or `"v_prediction"` (Salimans & Ho 2022, used by CogVideoX, SD 2.x v-pred).	`'v_prediction'`
`rescale_betas_zero_snr`	`bool`	Whether to apply zero-terminal-SNR rescaling to the beta schedule.	`True`
`timestep_spacing`	`TimestepSpacing`	`"leading"` (Stable Diffusion convention) or `"trailing"` (CogVideoX convention).	`'trailing'`
`set_alpha_to_one`	`bool`	Whether the final `alpha_cumprod` (used as `prev` when stepping past index 0) is forced to `1.0`.	`True`
`clip_sample`	`bool`	Whether to clip the predicted `x0` to `[-1, 1]` before computing the step. Set `False` for latent-space diffusion (default).	`False`
`num_inference_steps`	`int`	Default schedule length. Call :meth:`set_timesteps` to change it later.	`50`

timesteps `property` ¶

timesteps: array

Schedule indices in iteration order (descending in t).

set_timesteps ¶

set_timesteps(num_inference_steps: int) -> None

Recompute the timestep schedule for num_inference_steps steps.

step ¶

step(model_output: array, timestep: int | array, sample: array) -> array

Run one deterministic DDIM step.

Parameters:

Name	Type	Description	Default
`model_output`	`array`	The model's prediction at `timestep` (`epsilon` or `v` depending on `prediction_type`).	required
`timestep`	`int \| array`	Current step's index (int or 0-d `mx.array`).	required
`sample`	`array`	Current noisy sample.	required

Returns:

Type	Description
`array`	The denoised sample for the previous timestep.

add_noise ¶

add_noise(original: array, noise: array, timestep: int | array) -> array

Forward-diffuse original to noise level timestep.

FlowMatchEulerDiscreteScheduler ¶

FlowMatchEulerDiscreteScheduler(num_train_timesteps: int = 1000, shift: float = 1.0)

Stateful flow-matching Euler scheduler.

Mirrors the diffusers FlowMatchEulerDiscreteScheduler API: set_timesteps then step through timesteps calling step with the model's velocity output. add_noise forward-interpolates a clean sample toward pure noise.

Convention: sigmas ascend from ~0 (clean) to 1 (noise), with a terminal 1.0 appended so the final step has a valid sigma_next.

Parameters:

Name	Type	Description	Default
`num_train_timesteps`	`int`	Number of training timesteps.	`1000`
`shift`	`float`	Shift factor applied to the sigma schedule.	`1.0`

set_timesteps ¶

set_timesteps(num_inference_steps: int, sigmas: ndarray | list[float] | None = None) -> None

Configure the scheduler for num_inference_steps denoising steps.

Parameters:

Name	Type	Description	Default
`num_inference_steps`	`int`	Number of denoising steps.	required
`sigmas`	`ndarray \| list[float] \| None`	Optional custom sigmas (ascending, `[0,1]`). If omitted, a linspace schedule is generated.	`None`

step ¶

step(model_output: array, timestep: array, sample: array) -> array

Advance the sample by one Euler step using a velocity prediction.

Parameters:

Name	Type	Description	Default
`model_output`	`array`	Predicted velocity (same shape as `sample`).	required
`timestep`	`array`	Current timestep (only used to lazy-init the step index).	required
`sample`	`array`	Current noisy sample.	required

Returns:

Type	Description
`array`	Updated sample after one Euler step.

add_noise ¶

add_noise(original: array, noise: array, sigma: array) -> array

Flow-matching interpolation: sigma * noise + (1 - sigma) * original.

TeaCacheController ¶

TeaCacheController(num_steps: int, rel_l1_thresh: float, coefficients: Sequence[float])

Stateful controller deciding when to skip a transformer forward.

Usage per denoising step::

if controller.should_compute(step_index, modulated_input):
    x_in = x
    x = transformer_blocks(x)
    controller.cache_residual(x - x_in)
else:
    x = x + controller.previous_residual

Boundary steps (step_index == 0 and step_index == num_steps - 1) always compute and reset the accumulator.

Parameters:

Name	Type	Description	Default
`num_steps`	`int`	Total number of denoising steps in the generation.	required
`rel_l1_thresh`	`float`	Skip threshold on the accumulated rescaled L1 distance.	required
`coefficients`	`Sequence[float]`	Polynomial coefficients in `numpy.poly1d` order (highest degree first), calibrated to map raw L1 distances to a quality budget.	required

previous_residual `property` ¶

previous_residual: Any

Last cached payload. Raises before the first cache_residual call.

reset ¶

reset() -> None

Clear all state. Call at the start of each new generation.

should_compute ¶

should_compute(step_index: int, modulated_input: array) -> bool

Decide whether to run the transformer at step_index.

Side-effects: advances the stored previous_modulated_input and the internal accumulator. Must be called once per step in order.

cache_residual ¶

cache_residual(residual: Any) -> None

Store the residual from the just-computed step for reuse on skip.

residual is whatever the caller wants to retrieve later via previous_residual. Single-tensor models pass an mx.array; multi-stream models (e.g. LTX-2) pass a tuple or dict. The controller does not inspect or copy the value.

TimestepEmbedding ¶

TimestepEmbedding(in_channels: int, time_embed_dim: int)

Bases: Module

MLP that projects sinusoidal timestep embeddings into a feature space.

Weight keys: linear1.{weight,bias}, linear2.{weight,bias}.

VerifiedFeatureCache ¶

VerifiedFeatureCache(num_steps: int, tau_0: float, beta: float, *, order: int = 2, epsilon: float = 1e-08)

Lagrange-extrapolate a tracked feature and verify before accepting.

Usage per step::

cache = VerifiedFeatureCache(num_steps=50, tau_0=0.1, beta=0.5)

for step in range(num_steps):
    if cache.can_predict(step):
        predicted = cache.extrapolate(step)
        actual = run_verifier_layer(state)  # caller-supplied
        if cache.accept(step, predicted, actual):
            # use predicted for downstream layers / skip full forward
            cache.record(step, predicted)
            continue
    feature = run_full_forward(state)  # fallback
    cache.record(step, feature)

Boundary policy: can_predict returns False at step == 0 and step == num_steps - 1, forcing a full compute at both ends. This matches the convention used by other step-aware controllers in mlx_arsenal.diffusion.

Parameters:

Name	Type	Description	Default
`num_steps`	`int`	Total number of iterative steps in the generation.	required
`tau_0`	`float`	Base relative-L2² threshold for accepting a draft.	required
`beta`	`float`	Geometric schedule base. `beta < 1` → strict early / tolerant late (SpeCa default regime).	required
`order`	`int`	Lagrange polynomial order. Requires `order + 1` anchors before predictions become available. Typical values 1-3.	`2`
`epsilon`	`float`	Numerical floor in the relative-L2 denominator.	`1e-08`

reset ¶

reset() -> None

Drop all anchors. Call at the start of each new generation.

threshold ¶

threshold(step_index: int) -> float

Adaptive threshold at step_index (relative-L2² units).

can_predict ¶

can_predict(step_index: int) -> bool

True if a draft is available for step_index.

Returns False at boundary steps (0 and num_steps - 1) and whenever fewer than order + 1 anchors have been recorded.

extrapolate ¶

extrapolate(step_index: int) -> array

Lagrange-extrapolate the tracked feature at step_index.

Raises RuntimeError if called when can_predict would return False — the caller is expected to gate on it first.

accept ¶

accept(step_index: int, predicted: array, actual: array) -> bool

Compare draft to ground truth on the verification layer.

Uses squared relative-L2 to match the SpeCa formulation: e = ‖predicted − actual‖²₂ / (‖actual‖²₂ + ε). Returns True if e <= threshold(step_index).

record ¶

record(step_index: int, feature: array) -> None

Append feature as a new anchor at step_index.

Anchors are stored in a fixed-capacity FIFO of size order + 1. step_index must be strictly greater than the last recorded step (anchors are monotonic by construction).

WindowResidualController ¶

WindowResidualController(num_steps: int)

Step-aware controller for the WA-RS residual cache.

Construct via :meth:fixed, :meth:scheduled, or :meth:adaptive — the bare __init__ is intentionally not part of the public API.

previous_residual `property` ¶

previous_residual: array

Last cached residual. Raises before the first cache_residual call.

fixed `classmethod` ¶

fixed(num_steps: int, *, refresh_every: int) -> WindowResidualController

Refresh on step 0, num_steps - 1, and every refresh_every step.

scheduled `classmethod` ¶

scheduled(num_steps: int, *, refresh_steps: Sequence[int]) -> WindowResidualController

Refresh on step 0, num_steps - 1, and every step in refresh_steps.

adaptive `classmethod` ¶

adaptive(num_steps: int, *, rel_l1_thresh: float) -> WindowResidualController

Refresh when relative-L1 input delta crosses rel_l1_thresh.

Mirrors :class:TeaCacheController semantics: at non-boundary step i with previous input p, refresh iff mean(|input - p|) / mean(|p|) >= rel_l1_thresh. A zero-norm previous input also forces a refresh.

reset ¶

reset() -> None

Clear all state. Call at the start of each new generation.

should_refresh ¶

should_refresh(step_index: int, attn_input: array | None = None) -> bool

Decide whether to recompute full attention at step_index.

attn_input is required in adaptive mode and ignored otherwise. Boundary steps (0 and num_steps - 1) always return True.

cache_residual ¶

cache_residual(residual: array) -> None

Store the full - window residual from the just-refreshed step for reuse.

splice_heads ¶

splice_heads(new_output: array, cached_output: array, recompute_mask: array) -> array

Per-head select-between-tensors.

Where recompute_mask[h] is True, take new_output[:, h]; where False, take cached_output[:, h]. The head axis is assumed to be axis 1.

Parameters:

Name	Type	Description	Default
`new_output`	`array`	`(B, num_heads, ...)` newly computed attention output.	required
`cached_output`	`array`	`(B, num_heads, ...)` previously cached output, same shape as `new_output`.	required
`recompute_mask`	`array`	1D bool array of length `num_heads`.	required

Returns:

Type	Description
`array`	Spliced tensor with the same shape as `new_output`.

cfg_head_similarity ¶

cfg_head_similarity(cond: array, uncond: array, *, metric: Metric = 'cosine') -> array

Per-head similarity between conditional and unconditional outputs.

Parameters:

Name	Type	Description	Default
`cond`	`array`	`(B, num_heads, S, D)` conditional attention output.	required
`uncond`	`array`	`(B, num_heads, S, D)` unconditional output, same shape.	required
`metric`	`Metric`	`"cosine"` (default, `[-1, 1]`) or `"relative_l1"` (`≥ 0`).	`'cosine'`

Returns:

Type	Description
`array`	`(num_heads,)` float array. Reduced over batch, sequence, and
`array`	feature axes.

cfg_skip_mask ¶

cfg_skip_mask(scores: array, threshold: float, *, metric: Metric = 'cosine') -> array

Convert similarity scores to a skip-uncond mask.

Parameters:

Name	Type	Description	Default
`scores`	`array`	`(num_heads,)` or `(num_blocks, num_heads)` scores from :func:`cfg_head_similarity` or :class:`CFGSimilarityProfiler.scores`.	required
`threshold`	`float`	Cut-off. With `cosine`, skip when `scores >= threshold`; with `relative_l1`, skip when `scores <= threshold`.	required
`metric`	`Metric`	Which inversion to apply.	`'cosine'`

Returns:

Type	Description
`array`	Bool array with the same shape as `scores`. `True` means "this
`array`	head can skip the uncond branch".

classifier_free_guidance ¶

classifier_free_guidance(cond: array, uncond: array, scale: float) -> array

Apply classifier-free guidance.

Returns uncond + scale * (cond - uncond). A scale of 1.0 yields the conditioned prediction, 0.0 the unconditioned one, and values greater than 1.0 amplify the conditioning signal.

Parameters:

Name	Type	Description	Default
`cond`	`array`	Conditioned prediction.	required
`uncond`	`array`	Unconditioned prediction (same shape as `cond`).	required
`scale`	`float`	Guidance scale.	required

Returns:

Type	Description
`array`	Guided prediction.

euler_step ¶

euler_step(x: array, x0: array, sigma: float, sigma_next: float) -> array

Single stateless Euler step on an x0-prediction model.

Implements x_{t-1} = x + (sigma_next - sigma) * (x - x0) / sigma. When sigma == 0 the current sample is already clean, so x0 is returned directly.

Parameters:

Name	Type	Description	Default
`x`	`array`	Current noisy sample.	required
`x0`	`array`	Predicted clean sample.	required
`sigma`	`float`	Current noise level.	required
`sigma_next`	`float`	Next noise level.	required

Returns:

Type	Description
`array`	Updated sample at `sigma_next`.

dynamic_shift_schedule ¶

dynamic_shift_schedule(num_steps: int, num_tokens: int, base_shift: float = 0.95, max_shift: float = 2.05, base_tokens: int = 1024, max_tokens: int = 4096, stretch: bool = True, terminal: float = 0.1) -> list[float]

Sigma schedule with token-count-dependent shift (LTX-style).

Interpolates a shift factor linearly between base_shift at base_tokens and max_shift at max_tokens, then applies it to a descending linspace. Optional terminal stretching matches the last non-zero sigma to 1 - terminal.

Parameters:

Name	Type	Description	Default
`num_steps`	`int`	Number of denoising steps.	required
`num_tokens`	`int`	Number of latent tokens (drives the shift).	required
`base_shift`	`float`	Shift at `base_tokens`.	`0.95`
`max_shift`	`float`	Shift at `max_tokens`.	`2.05`
`base_tokens`	`int`	Anchor for `base_shift`.	`1024`
`max_tokens`	`int`	Anchor for `max_shift`.	`4096`
`stretch`	`bool`	Rescale so the last non-zero sigma equals `1 - terminal`.	`True`
`terminal`	`float`	Target terminal for stretching.	`0.1`

Returns:

Type	Description
`list[float]`	List of `num_steps + 1` sigma values ending at `0.0`.

get_sampling_sigmas ¶

get_sampling_sigmas(num_steps: int, shift: float = 1.0) -> list[float]

Flow-matching sigma schedule: linspace(1, 0, steps+1) with optional shift.

The schedule includes the terminal 0.0 so pairs are formed via zip(sigmas[:-1], sigmas[1:]).

Parameters:

Name	Type	Description	Default
`num_steps`	`int`	Number of denoising steps.	required
`shift`	`float`	Shift factor applied as `shift * s / (1 + (shift - 1) * s)`.	`1.0`

Returns:

Type	Description
`list[float]`	List of `num_steps + 1` sigma values descending from 1.0 to 0.0.

get_timestep_embedding ¶

get_timestep_embedding(timesteps: array, embedding_dim: int, flip_sin_to_cos: bool = True, downscale_freq_shift: float = 0.0, scale: float = 1.0, max_period: float = 10000.0) -> array

Create sinusoidal timestep embeddings.

Parameters:

Name	Type	Description	Default
`timesteps`	`array`	1D array of timestep values.	required
`embedding_dim`	`int`	Dimension of the output embeddings.	required
`flip_sin_to_cos`	`bool`	If True, output `[cos, sin]`; if False, `[sin, cos]`.	`True`
`downscale_freq_shift`	`float`	Controls delta between frequencies.	`0.0`
`scale`	`float`	Scaling factor applied before sin/cos.	`1.0`
`max_period`	`float`	Controls the minimum frequency.	`10000.0`

Returns:

Type	Description
`array`	Array of shape `(len(timesteps), embedding_dim)`.

geometric_threshold ¶

geometric_threshold(step_index: int, num_steps: int, tau_0: float, beta: float) -> float

SpeCa-style geometric threshold schedule.

τ_t = τ₀ · β^((T - 1 - t) / max(T - 1, 1))

With beta < 1: threshold grows from τ₀·β (strict early) to τ₀ (tolerant late). With beta > 1: the opposite. beta == 1 yields a constant threshold τ₀. The "right" regime is empirical and depends on the schedule and model — see the research note at docs/research/verified-feature-caching.md.

attention_cache ¶

Attention output cache (AST-style) for diffusion transformers.

Caches attention sub-layer output across denoising steps and reuses it on the next step when the input has barely changed. Two granularities:

:class:PerLayerAttentionCache — scalar similarity, one decision per layer per step. Simpler, mirrors :class:mlx_arsenal.diffusion.TeaCacheController but at the attention sub-layer instead of a whole transformer block.

Decision rule, at step i:

If i == 0 or i == num_steps - 1 → recompute (boundary).
If mean(abs(prev_input)) is zero → recompute (degenerate).
If mean(abs(input - prev_input)) / mean(abs(prev_input)) >= rel_l1_thresh → recompute.

References

DiTFastAttn — Attention Sharing across Timesteps (AST).

PerLayerAttentionCache ¶

PerLayerAttentionCache(num_steps: int, rel_l1_thresh: float)

Stateful per-layer attention output cache.

previous_output `property` ¶

previous_output: array

Last cached attention output. Raises before the first cache_output call.

reset ¶

reset() -> None

Clear all state. Call at the start of each new generation.

should_compute ¶

should_compute(step_index: int, attn_input: array) -> bool

Decide whether to recompute attention at step_index.

Side-effects: advances the stored previous input and summary. Must be called once per step in order.

should_compute_from_summary ¶

should_compute_from_summary(step_index: int, summary: float) -> bool

Decide using a caller-supplied scalar summary instead of a tensor.

summary is the analogue of mean(abs(input - prev_input)) / mean(abs(prev_input)) — the caller has already done the math.

Do not interleave with :meth:should_compute within a single denoising run: the two methods write semantically different values into the internal previous-summary slot. Pick one mode per run.

cache_output ¶

cache_output(output: array) -> None

Store the attention output from the just-computed step for reuse on skip.

PerHeadAttentionCache ¶

PerHeadAttentionCache(num_heads: int, num_steps: int, rel_l1_thresh: float)

Stateful per-head attention output cache.

Returns a (num_heads,) bool decision per step. Inputs are assumed to have the head axis at position 1 — i.e. shape (B, num_heads, ...).

previous_output `property` ¶

previous_output: array

Last cached attention output. Raises before the first cache_output call.

reset ¶

reset() -> None

Clear all state. Call at the start of each new generation.

should_compute ¶

should_compute(step_index: int, attn_input: array) -> array

Per-head decide whether to recompute attention at step_index.

Side-effects: advances the stored previous input and per-head summary. Must be called once per step in order.

should_compute_from_summary ¶

should_compute_from_summary(step_index: int, summary: array) -> array

Decide per head using a caller-supplied (num_heads,) summary.

summary[h] is the analogue of the per-head delta ratio. Do not interleave with :meth:should_compute within a single denoising run: the two methods write semantically different values into the internal previous-summary slot. Pick one mode per run.

cache_output ¶

cache_output(output: array) -> None

Store the full (B, num_heads, ...) attention output for per-head splicing on skip.

splice_heads ¶

splice_heads(new_output: array, cached_output: array, recompute_mask: array) -> array

Per-head select-between-tensors.

Where recompute_mask[h] is True, take new_output[:, h]; where False, take cached_output[:, h]. The head axis is assumed to be axis 1.

Parameters:

Name	Type	Description	Default
`new_output`	`array`	`(B, num_heads, ...)` newly computed attention output.	required
`cached_output`	`array`	`(B, num_heads, ...)` previously cached output, same shape as `new_output`.	required
`recompute_mask`	`array`	1D bool array of length `num_heads`.	required

Returns:

Type	Description
`array`	Spliced tensor with the same shape as `new_output`.

cfg_skip ¶

CFG-skip (Attention Sharing across CFG, DiTFastAttn ASC).

Profiles per-head similarity between the conditional and unconditional CFG branches, builds a static schedule, and applies that schedule at runtime to skip the unconditional branch on heads where cond and uncond outputs are near-identical.

Two metrics:

cosine — dot/norm of flattened (B, S, D) per head, in [-1, 1]. Skip when score is at least the threshold. Default; matches the DiTFastAttn ASC literature.
relative_l1 — mean(|c - u|) / mean(|c|) per head, ≥ 0. Skip when score is at most the threshold. Same family as :class:~mlx_arsenal.diffusion.TeaCacheController and the attention output caches, so existing thresholds carry over.

The runtime apply step (:meth:CFGSkipController.apply) bundles a :func:mlx_arsenal.diffusion.splice_heads call so callers do not have to import the splice helper directly.

CFGSimilarityProfiler ¶

CFGSimilarityProfiler(num_blocks: int, num_heads: int, *, metric: Metric = 'cosine')

Per-(block, head) running-mean similarity accumulator.

Records the per-head similarity of cond and uncond attention outputs during a warmup pass and produces a static (num_blocks, num_heads) skip schedule.

scores `property` ¶

scores: array

(num_blocks, num_heads) running mean. Zero-count blocks → 0.0.

call_counts `property` ¶

call_counts: array

(num_blocks,) int32 count of :meth:record calls per block.

reset ¶

reset() -> None

Clear accumulated scores and counts. Call before a new warmup pass.

record ¶

record(block_idx: int, cond: array, uncond: array) -> None

Accumulate per-head cond/uncond similarity for block block_idx.

build_skip_schedule ¶

build_skip_schedule(threshold: float) -> array

Threshold :attr:scores into a (num_blocks, num_heads) skip mask.

CFGSkipController ¶

CFGSkipController(schedule: array)

Wraps a static (num_blocks, num_heads) bool skip schedule.

True at (b, h) means: for block b, head h skips the unconditional branch and reuses the conditional output. False means the head computes uncond normally.

num_blocks `property` ¶

num_blocks: int

Number of transformer blocks in the schedule.

num_heads `property` ¶

num_heads: int

Number of attention heads in the schedule.

from_profiler `classmethod` ¶

from_profiler(profiler: CFGSimilarityProfiler, threshold: float) -> 'CFGSkipController'

Build a controller from a profiler by thresholding its scores.

should_skip_uncond ¶

should_skip_uncond(block_idx: int) -> array

(num_heads,) bool mask — True heads skip uncond for this block.

apply ¶

apply(block_idx: int, cond_output: array, uncond_output: array) -> array

Apply the cached schedule to splice cond/uncond outputs (wraps :func:splice_heads).

cfg_head_similarity ¶

cfg_head_similarity(cond: array, uncond: array, *, metric: Metric = 'cosine') -> array

Per-head similarity between conditional and unconditional outputs.

Parameters:

Name	Type	Description	Default
`cond`	`array`	`(B, num_heads, S, D)` conditional attention output.	required
`uncond`	`array`	`(B, num_heads, S, D)` unconditional output, same shape.	required
`metric`	`Metric`	`"cosine"` (default, `[-1, 1]`) or `"relative_l1"` (`≥ 0`).	`'cosine'`

Returns:

Type	Description
`array`	`(num_heads,)` float array. Reduced over batch, sequence, and
`array`	feature axes.

cfg_skip_mask ¶

cfg_skip_mask(scores: array, threshold: float, *, metric: Metric = 'cosine') -> array

Convert similarity scores to a skip-uncond mask.

Parameters:

Name	Type	Description	Default
`scores`	`array`	`(num_heads,)` or `(num_blocks, num_heads)` scores from :func:`cfg_head_similarity` or :class:`CFGSimilarityProfiler.scores`.	required
`threshold`	`float`	Cut-off. With `cosine`, skip when `scores >= threshold`; with `relative_l1`, skip when `scores <= threshold`.	required
`metric`	`Metric`	Which inversion to apply.	`'cosine'`

Returns:

Type	Description
`array`	Bool array with the same shape as `scores`. `True` means "this
`array`	head can skip the uncond branch".

ddim ¶

DDIM (Denoising Diffusion Implicit Models) scheduler — MLX port.

Matches the diffusers DDIMScheduler / CogVideoXDDIMScheduler behaviour for the deterministic case (eta=0). Supports the two common prediction types (epsilon and v_prediction), the two common spacing strategies (leading, trailing), and optional zero-SNR rescaling of the beta schedule.

Use this for ports of DDPM-trained models (CogVideoX, Stable Diffusion 1.x / 2.x, ...). For flow-matching models (LTX, Hunyuan-DiT, ERNIE-Image), see :class:FlowMatchEulerDiscreteScheduler instead.

DDIMScheduler ¶

DDIMScheduler(num_train_timesteps: int = 1000, beta_start: float = 0.00085, beta_end: float = 0.012, beta_schedule: BetaSchedule = 'scaled_linear', prediction_type: PredictionType = 'v_prediction', rescale_betas_zero_snr: bool = True, timestep_spacing: TimestepSpacing = 'trailing', set_alpha_to_one: bool = True, clip_sample: bool = False, num_inference_steps: int = 50)

Deterministic DDIM scheduler (eta=0).

Parameters:

Name	Type	Description	Default
`num_train_timesteps`	`int`	Number of diffusion steps used during training. Defaults to `1000`.	`1000`
`beta_start`	`float`	First beta value. Defaults to `0.00085` (Stable Diffusion convention).	`0.00085`
`beta_end`	`float`	Last beta value. Defaults to `0.012`.	`0.012`
`beta_schedule`	`BetaSchedule`	`"scaled_linear"` (default) or `"linear"`. `"scaled_linear"` matches Stable Diffusion's `linspace(sqrt(b0), sqrt(bT), N)**2`.	`'scaled_linear'`
`prediction_type`	`PredictionType`	`"epsilon"` (vanilla) or `"v_prediction"` (Salimans & Ho 2022, used by CogVideoX, SD 2.x v-pred).	`'v_prediction'`
`rescale_betas_zero_snr`	`bool`	Whether to apply zero-terminal-SNR rescaling to the beta schedule.	`True`
`timestep_spacing`	`TimestepSpacing`	`"leading"` (Stable Diffusion convention) or `"trailing"` (CogVideoX convention).	`'trailing'`
`set_alpha_to_one`	`bool`	Whether the final `alpha_cumprod` (used as `prev` when stepping past index 0) is forced to `1.0`.	`True`
`clip_sample`	`bool`	Whether to clip the predicted `x0` to `[-1, 1]` before computing the step. Set `False` for latent-space diffusion (default).	`False`
`num_inference_steps`	`int`	Default schedule length. Call :meth:`set_timesteps` to change it later.	`50`

timesteps `property` ¶

timesteps: array

Schedule indices in iteration order (descending in t).

set_timesteps ¶

set_timesteps(num_inference_steps: int) -> None

Recompute the timestep schedule for num_inference_steps steps.

step ¶

step(model_output: array, timestep: int | array, sample: array) -> array

Run one deterministic DDIM step.

Parameters:

Name	Type	Description	Default
`model_output`	`array`	The model's prediction at `timestep` (`epsilon` or `v` depending on `prediction_type`).	required
`timestep`	`int \| array`	Current step's index (int or 0-d `mx.array`).	required
`sample`	`array`	Current noisy sample.	required

Returns:

Type	Description
`array`	The denoised sample for the previous timestep.

add_noise ¶

add_noise(original: array, noise: array, timestep: int | array) -> array

Forward-diffuse original to noise level timestep.

samplers ¶

Stateless samplers and guidance primitives for diffusion denoising.

euler_step ¶

euler_step(x: array, x0: array, sigma: float, sigma_next: float) -> array

Single stateless Euler step on an x0-prediction model.

Implements x_{t-1} = x + (sigma_next - sigma) * (x - x0) / sigma. When sigma == 0 the current sample is already clean, so x0 is returned directly.

Parameters:

Name	Type	Description	Default
`x`	`array`	Current noisy sample.	required
`x0`	`array`	Predicted clean sample.	required
`sigma`	`float`	Current noise level.	required
`sigma_next`	`float`	Next noise level.	required

Returns:

Type	Description
`array`	Updated sample at `sigma_next`.

classifier_free_guidance ¶

classifier_free_guidance(cond: array, uncond: array, scale: float) -> array

Apply classifier-free guidance.

Returns uncond + scale * (cond - uncond). A scale of 1.0 yields the conditioned prediction, 0.0 the unconditioned one, and values greater than 1.0 amplify the conditioning signal.

Parameters:

Name	Type	Description	Default
`cond`	`array`	Conditioned prediction.	required
`uncond`	`array`	Unconditioned prediction (same shape as `cond`).	required
`scale`	`float`	Guidance scale.	required

Returns:

Type	Description
`array`	Guided prediction.

schedulers ¶

Sigma schedules and noise schedulers for flow-matching diffusion.

FlowMatchEulerDiscreteScheduler ¶

FlowMatchEulerDiscreteScheduler(num_train_timesteps: int = 1000, shift: float = 1.0)

Stateful flow-matching Euler scheduler.

Mirrors the diffusers FlowMatchEulerDiscreteScheduler API: set_timesteps then step through timesteps calling step with the model's velocity output. add_noise forward-interpolates a clean sample toward pure noise.

Convention: sigmas ascend from ~0 (clean) to 1 (noise), with a terminal 1.0 appended so the final step has a valid sigma_next.

Parameters:

Name	Type	Description	Default
`num_train_timesteps`	`int`	Number of training timesteps.	`1000`
`shift`	`float`	Shift factor applied to the sigma schedule.	`1.0`

set_timesteps ¶

set_timesteps(num_inference_steps: int, sigmas: ndarray | list[float] | None = None) -> None

Configure the scheduler for num_inference_steps denoising steps.

Parameters:

Name	Type	Description	Default
`num_inference_steps`	`int`	Number of denoising steps.	required
`sigmas`	`ndarray \| list[float] \| None`	Optional custom sigmas (ascending, `[0,1]`). If omitted, a linspace schedule is generated.	`None`

step ¶

step(model_output: array, timestep: array, sample: array) -> array

Advance the sample by one Euler step using a velocity prediction.

Parameters:

Name	Type	Description	Default
`model_output`	`array`	Predicted velocity (same shape as `sample`).	required
`timestep`	`array`	Current timestep (only used to lazy-init the step index).	required
`sample`	`array`	Current noisy sample.	required

Returns:

Type	Description
`array`	Updated sample after one Euler step.

add_noise ¶

add_noise(original: array, noise: array, sigma: array) -> array

Flow-matching interpolation: sigma * noise + (1 - sigma) * original.

get_sampling_sigmas ¶

get_sampling_sigmas(num_steps: int, shift: float = 1.0) -> list[float]

Flow-matching sigma schedule: linspace(1, 0, steps+1) with optional shift.

The schedule includes the terminal 0.0 so pairs are formed via zip(sigmas[:-1], sigmas[1:]).

Parameters:

Name	Type	Description	Default
`num_steps`	`int`	Number of denoising steps.	required
`shift`	`float`	Shift factor applied as `shift * s / (1 + (shift - 1) * s)`.	`1.0`

Returns:

Type	Description
`list[float]`	List of `num_steps + 1` sigma values descending from 1.0 to 0.0.

dynamic_shift_schedule ¶

dynamic_shift_schedule(num_steps: int, num_tokens: int, base_shift: float = 0.95, max_shift: float = 2.05, base_tokens: int = 1024, max_tokens: int = 4096, stretch: bool = True, terminal: float = 0.1) -> list[float]

Sigma schedule with token-count-dependent shift (LTX-style).

Interpolates a shift factor linearly between base_shift at base_tokens and max_shift at max_tokens, then applies it to a descending linspace. Optional terminal stretching matches the last non-zero sigma to 1 - terminal.

Parameters:

Name	Type	Description	Default
`num_steps`	`int`	Number of denoising steps.	required
`num_tokens`	`int`	Number of latent tokens (drives the shift).	required
`base_shift`	`float`	Shift at `base_tokens`.	`0.95`
`max_shift`	`float`	Shift at `max_tokens`.	`2.05`
`base_tokens`	`int`	Anchor for `base_shift`.	`1024`
`max_tokens`	`int`	Anchor for `max_shift`.	`4096`
`stretch`	`bool`	Rescale so the last non-zero sigma equals `1 - terminal`.	`True`
`terminal`	`float`	Target terminal for stretching.	`0.1`

Returns:

Type	Description
`list[float]`	List of `num_steps + 1` sigma values ending at `0.0`.

teacache ¶

TeaCache: timestep-aware residual caching for diffusion transformers.

TeaCache (Liu et al., "Timestep Embedding Aware Cache") accelerates diffusion inference by reusing a cached residual across timesteps when the modulated input doesn't move much. The trade-off is governed by rel_l1_thresh: higher = more skipping = faster but lossier.

Architecture-agnostic mechanism — coefficients are model-specific and live with each model's port (e.g. inside the LTX-2 / Hunyuan / Flux MLX implementation). mlx-arsenal ships only the engine.

References

https://github.com/ali-vilab/TeaCache

TeaCacheController ¶

TeaCacheController(num_steps: int, rel_l1_thresh: float, coefficients: Sequence[float])

Stateful controller deciding when to skip a transformer forward.

Usage per denoising step::

if controller.should_compute(step_index, modulated_input):
    x_in = x
    x = transformer_blocks(x)
    controller.cache_residual(x - x_in)
else:
    x = x + controller.previous_residual

Boundary steps (step_index == 0 and step_index == num_steps - 1) always compute and reset the accumulator.

Parameters:

Name	Type	Description	Default
`num_steps`	`int`	Total number of denoising steps in the generation.	required
`rel_l1_thresh`	`float`	Skip threshold on the accumulated rescaled L1 distance.	required
`coefficients`	`Sequence[float]`	Polynomial coefficients in `numpy.poly1d` order (highest degree first), calibrated to map raw L1 distances to a quality budget.	required

previous_residual `property` ¶

previous_residual: Any

Last cached payload. Raises before the first cache_residual call.

reset ¶

reset() -> None

Clear all state. Call at the start of each new generation.

should_compute ¶

should_compute(step_index: int, modulated_input: array) -> bool

Decide whether to run the transformer at step_index.

Side-effects: advances the stored previous_modulated_input and the internal accumulator. Must be called once per step in order.

cache_residual ¶

cache_residual(residual: Any) -> None

Store the residual from the just-computed step for reuse on skip.

residual is whatever the caller wants to retrieve later via previous_residual. Single-tensor models pass an mx.array; multi-stream models (e.g. LTX-2) pass a tuple or dict. The controller does not inspect or copy the value.

timestep ¶

Timestep embeddings for diffusion models.

TimestepEmbedding ¶

TimestepEmbedding(in_channels: int, time_embed_dim: int)

Bases: Module

MLP that projects sinusoidal timestep embeddings into a feature space.

Weight keys: linear1.{weight,bias}, linear2.{weight,bias}.

get_timestep_embedding ¶

get_timestep_embedding(timesteps: array, embedding_dim: int, flip_sin_to_cos: bool = True, downscale_freq_shift: float = 0.0, scale: float = 1.0, max_period: float = 10000.0) -> array

Create sinusoidal timestep embeddings.

Parameters:

Name	Type	Description	Default
`timesteps`	`array`	1D array of timestep values.	required
`embedding_dim`	`int`	Dimension of the output embeddings.	required
`flip_sin_to_cos`	`bool`	If True, output `[cos, sin]`; if False, `[sin, cos]`.	`True`
`downscale_freq_shift`	`float`	Controls delta between frequencies.	`0.0`
`scale`	`float`	Scaling factor applied before sin/cos.	`1.0`
`max_period`	`float`	Controls the minimum frequency.	`10000.0`

Returns:

Type	Description
`array`	Array of shape `(len(timesteps), embedding_dim)`.

verified_cache ¶

Verified feature caching (SpeCa-style) for iterative generative models.

Forecast-then-verify cache: at each iterative step (denoising timestep, autoregressive step, etc.) the controller can extrapolate a tracked feature from previously observed anchors via Lagrange polynomial extrapolation, then accept or reject the prediction by comparing it to a freshly computed ground-truth feature (typically the output of a single inner layer — the "verification layer").

The pattern is lossy but bounded — see the SpeCa convergence analysis (Zou et al., 2025, arxiv:2509.11628). For bit-exact reproducibility, do not use.

Architecture-agnostic mechanism. The caller decides:

Which feature to track (the verification-layer output is a common pick).
How to map the controller's abstract step_index to its own iterative axis (denoising step, diffusion timestep, etc.).
The Lagrange order (number of anchors − 1) and the threshold schedule parameters tau_0 and beta.

References

SpeCa — https://arxiv.org/abs/2509.11628 TaylorSeer — https://arxiv.org/abs/2503.06923

VerifiedFeatureCache ¶

VerifiedFeatureCache(num_steps: int, tau_0: float, beta: float, *, order: int = 2, epsilon: float = 1e-08)

Lagrange-extrapolate a tracked feature and verify before accepting.

Usage per step::

cache = VerifiedFeatureCache(num_steps=50, tau_0=0.1, beta=0.5)

for step in range(num_steps):
    if cache.can_predict(step):
        predicted = cache.extrapolate(step)
        actual = run_verifier_layer(state)  # caller-supplied
        if cache.accept(step, predicted, actual):
            # use predicted for downstream layers / skip full forward
            cache.record(step, predicted)
            continue
    feature = run_full_forward(state)  # fallback
    cache.record(step, feature)

Boundary policy: can_predict returns False at step == 0 and step == num_steps - 1, forcing a full compute at both ends. This matches the convention used by other step-aware controllers in mlx_arsenal.diffusion.

Parameters:

Name	Type	Description	Default
`num_steps`	`int`	Total number of iterative steps in the generation.	required
`tau_0`	`float`	Base relative-L2² threshold for accepting a draft.	required
`beta`	`float`	Geometric schedule base. `beta < 1` → strict early / tolerant late (SpeCa default regime).	required
`order`	`int`	Lagrange polynomial order. Requires `order + 1` anchors before predictions become available. Typical values 1-3.	`2`
`epsilon`	`float`	Numerical floor in the relative-L2 denominator.	`1e-08`

reset ¶

reset() -> None

Drop all anchors. Call at the start of each new generation.

threshold ¶

threshold(step_index: int) -> float

Adaptive threshold at step_index (relative-L2² units).

can_predict ¶

can_predict(step_index: int) -> bool

True if a draft is available for step_index.

Returns False at boundary steps (0 and num_steps - 1) and whenever fewer than order + 1 anchors have been recorded.

extrapolate ¶

extrapolate(step_index: int) -> array

Lagrange-extrapolate the tracked feature at step_index.

Raises RuntimeError if called when can_predict would return False — the caller is expected to gate on it first.

accept ¶

accept(step_index: int, predicted: array, actual: array) -> bool

Compare draft to ground truth on the verification layer.

Uses squared relative-L2 to match the SpeCa formulation: e = ‖predicted − actual‖²₂ / (‖actual‖²₂ + ε). Returns True if e <= threshold(step_index).

record ¶

record(step_index: int, feature: array) -> None

Append feature as a new anchor at step_index.

Anchors are stored in a fixed-capacity FIFO of size order + 1. step_index must be strictly greater than the last recorded step (anchors are monotonic by construction).

geometric_threshold ¶

geometric_threshold(step_index: int, num_steps: int, tau_0: float, beta: float) -> float

SpeCa-style geometric threshold schedule.

τ_t = τ₀ · β^((T - 1 - t) / max(T - 1, 1))

With beta < 1: threshold grows from τ₀·β (strict early) to τ₀ (tolerant late). With beta > 1: the opposite. beta == 1 yields a constant threshold τ₀. The "right" regime is empirical and depends on the schedule and model — see the research note at docs/research/verified-feature-caching.md.

window_residual ¶

WA-RS controller (DiTFastAttn Window Attention + Residual Sharing).

Decides when to refresh the cached full - window attention residual:

:meth:WindowResidualController.fixed — refresh every K steps.
:meth:WindowResidualController.scheduled — refresh on an explicit list.
:meth:WindowResidualController.adaptive — refresh when the attention input has moved more than rel_l1_thresh since the previous step (same relative-L1 metric as :class:mlx_arsenal.diffusion.TeaCacheController).

Boundary steps (0 and num_steps - 1) always refresh regardless of mode. The cached residual itself is opaque — the controller does not look inside mx.array shapes.

References

DiTFastAttn — Window Attention with Residual Sharing (WA-RS).

WindowResidualController ¶

WindowResidualController(num_steps: int)

Step-aware controller for the WA-RS residual cache.

Construct via :meth:fixed, :meth:scheduled, or :meth:adaptive — the bare __init__ is intentionally not part of the public API.

previous_residual `property` ¶

previous_residual: array

Last cached residual. Raises before the first cache_residual call.

fixed `classmethod` ¶

fixed(num_steps: int, *, refresh_every: int) -> WindowResidualController

Refresh on step 0, num_steps - 1, and every refresh_every step.

scheduled `classmethod` ¶

scheduled(num_steps: int, *, refresh_steps: Sequence[int]) -> WindowResidualController

Refresh on step 0, num_steps - 1, and every step in refresh_steps.

adaptive `classmethod` ¶

adaptive(num_steps: int, *, rel_l1_thresh: float) -> WindowResidualController

Refresh when relative-L1 input delta crosses rel_l1_thresh.

Mirrors :class:TeaCacheController semantics: at non-boundary step i with previous input p, refresh iff mean(|input - p|) / mean(|p|) >= rel_l1_thresh. A zero-norm previous input also forces a refresh.

reset ¶

reset() -> None

Clear all state. Call at the start of each new generation.

should_refresh ¶

should_refresh(step_index: int, attn_input: array | None = None) -> bool

Decide whether to recompute full attention at step_index.

attn_input is required in adaptive mode and ignored otherwise. Boundary steps (0 and num_steps - 1) always return True.

cache_residual ¶

cache_residual(residual: array) -> None

Store the full - window residual from the just-refreshed step for reuse.

Diffusion¶

diffusion ¶

PerHeadAttentionCache ¶

previous_output property ¶

reset ¶

should_compute ¶

should_compute_from_summary ¶

cache_output ¶

PerLayerAttentionCache ¶

previous_output property ¶

reset ¶

should_compute ¶

should_compute_from_summary ¶

cache_output ¶

CFGSimilarityProfiler ¶

scores property ¶

call_counts property ¶

reset ¶

record ¶

build_skip_schedule ¶

CFGSkipController ¶

num_blocks property ¶

num_heads property ¶

from_profiler classmethod ¶

should_skip_uncond ¶

apply ¶

DDIMScheduler ¶

timesteps property ¶

set_timesteps ¶

step ¶

add_noise ¶

FlowMatchEulerDiscreteScheduler ¶

set_timesteps ¶

step ¶

add_noise ¶

TeaCacheController ¶

previous_residual property ¶

reset ¶

should_compute ¶

cache_residual ¶

TimestepEmbedding ¶

VerifiedFeatureCache ¶

reset ¶

threshold ¶

can_predict ¶

extrapolate ¶

accept ¶

record ¶

WindowResidualController ¶

previous_residual property ¶

fixed classmethod ¶

scheduled classmethod ¶

adaptive classmethod ¶

reset ¶

should_refresh ¶

cache_residual ¶

splice_heads ¶

cfg_head_similarity ¶

cfg_skip_mask ¶

classifier_free_guidance ¶

euler_step ¶

dynamic_shift_schedule ¶

get_sampling_sigmas ¶

get_timestep_embedding ¶

geometric_threshold ¶

attention_cache ¶

PerLayerAttentionCache ¶

previous_output property ¶

reset ¶

should_compute ¶

should_compute_from_summary ¶

cache_output ¶

PerHeadAttentionCache ¶

previous_output property ¶

reset ¶

should_compute ¶

should_compute_from_summary ¶

cache_output ¶

splice_heads ¶

cfg_skip ¶

previous_output `property` ¶

previous_output `property` ¶

scores `property` ¶

call_counts `property` ¶

num_blocks `property` ¶

num_heads `property` ¶

from_profiler `classmethod` ¶

timesteps `property` ¶

previous_residual `property` ¶

previous_residual `property` ¶

fixed `classmethod` ¶

scheduled `classmethod` ¶

adaptive `classmethod` ¶

previous_output `property` ¶

previous_output `property` ¶

scores `property` ¶

call_counts `property` ¶

num_blocks `property` ¶

num_heads `property` ¶

from_profiler `classmethod` ¶

timesteps `property` ¶

previous_residual `property` ¶

previous_residual `property` ¶

fixed `classmethod` ¶

scheduled `classmethod` ¶

adaptive `classmethod` ¶