vllm_omni.diffusion.models.auk ¶

Generate one waveform.

Parameters:

Returns:

Compile the DiT as configured, then capture the codec decode buckets.

Defining this hook replaces the runner's generic transformer compile, so the DiT's diffusion_compile_granularity is honoured here: full compiles the whole denoise step; regional compiles the double- and single-stream blocks. Compilation is lazy and happens in the DiT CUDA graph's warm-up, before capture, so each graph replays the compiled kernels. The startup cost worth paying is the VAE decode, whose buckets are compiled and captured before the first request.

Drop the cached text projection. Call this when the conditioning changes.

Predict the flow-matching velocity for the target frames of x.

Parameters:

Returns:

adaLN modulations for every timestep, [len(timesteps), sum of widths].

The modulations depend only on the timestep, and a request's time grid is known before its first step. Computing them for the whole grid at once reads each modulation weight once per request instead of once per step, which at batch 1 is what those matrix-vector products cost.

Everything a denoise step needs that does not depend on x or the time.

An Euler loop over fixed conditioning calls this once and then :meth:step per timestep, so the text projection, the reference prompt embedding, the padding masks and biases and the rotary tables are not recomputed on every step. Arguments are those of :meth:forward; target_len is the target frame count n. timesteps are the times the steps will run at (the grid without its end point); when given, their adaLN modulations are precomputed and :meth:step selects them by step_index.

Project LLM hidden states [B, nt, text_hidden_dim] to model width.

Predict the velocity of x at time under a :meth:prepare context.

The target embedding does not depend on the CFG branch, so it is computed once and shared by both branches. When the context carries a modulation table, step_index selects this step's row of it and time is unused; the row broadcasts over the CFG branches.

Decode normalized [B, Np, latent_dim] latents into a [B, Np * hop_size] waveform in [-1, 1].

Latent frames of (left, right) context a decoded frame can depend on.

Summed layer by layer from the modules' own padding and kernel geometry, so it is an upper bound that holds for any weights: a window decoded on its own reproduces the full decode exactly except within these many frames of a window edge that is not a true clip edge. The decoder is causal apart from conv_pre and the band-limited upsamplers, so the right context is a few frames while the left one is dominated by the dilated convolutions of the first (lowest-rate) stage.

Encode a mono [1, T] waveform at :attr:sample_rate into [1, T / hop_size, latent_dim].

sample=False takes the posterior mean, which is what a reproducible pipeline wants; sample=True draws mean + randn * exp(log_std) from generator. Latents come back normalized by the checkpoint's global statistics.

Build from the released model_init_kwargs mapping, ignoring training-only entries.

Load a released vae.safetensors, then fold weight norm away. Returns (missing, unexpected).

device defaults to where the module already lives, which is the placement that matters: the fold reduces weight_v to a norm, and a CPU reduction lands up to one ULP away from the accelerator's. The encoder amplifies its activations roughly fiftyfold, so that one ULP grows into a ~4e-4 relative drift by the latent head. Move the module first, then load.

Fold weight_g/weight_v into plain weights, on whichever device the module is on.

Idempotent, and meant to run once the checkpoint is in. See :meth:load_weights for why the device this runs on is a numeric decision rather than a convenience.

Toggle the FIR modules' per-channel filter caching.

On by default. The switch exists so the cache's cost can be measured against the plain expand-per-call path. Flip it before any CUDA graph is captured: it reallocates the cached taps, and a captured graph keeps reading the old buffers.