[Model] Add Tencent AuK speech generation and editing (encoder + diffusion pipeline) (#7385)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Sy03 <1370724210@qq.com>

tests/diffusion/models/auk/test_auk_transformer_parity.py

@@ -0,0 +1,439 @@
+#!/usr/bin/env python
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""AuK transformer parity: the correctness oracle for the in-tree port.
+
+Checks that :class:`AuKTransformer` reproduces the upstream ``auk`` reference backbone
+on the released ``auk_base`` weights. Both models are built from the same
+state dict and fed the same inputs, under the reference's own inference
+configuration: fp32 weights with ``torch.autocast(bfloat16)``. That
+configuration matters. The reference keeps its rotary frequencies in a buffer,
+so casting its weights to bf16 would round both the frequencies and the
+integer positions and move the numbers for reasons unrelated to the port.
+
+Cases:
+
+* ``single`` - one conditional forward
+* ``cfg`` - the doubled ``cfg_infer`` forward, both branches
+* ``masked`` - batch of 2 with padded targets, references and text
+* ``rope`` - the local rotary implementation against ``x_transformers``, and
+  its survival of a cast to bf16
+* ``sample`` - an 8-step Euler trajectory against the reference backbone
+  driven by a transcription of upstream's ODE loop
+* ``generator`` - the sampler's two noise paths agreeing
+* ``cache`` - the text projection cache and ``clear_cache()``
+* ``mask flag`` - ``attn_mask_enabled=False`` reproducing the reference with
+  its own flag off
+
+This directory deliberately has no ``__init__.py``, unlike its siblings.
+pytest's prepend import mode puts the nearest non-package directory on
+``sys.path``, and an ``auk`` package here would shadow the upstream ``auk``
+distribution this test compares against.
+
+Needs the released checkpoint and a GPU, so it runs locally rather than in
+CI. Point ``AUK_CKPT_DIR`` at a directory holding ``auk_base.safetensors`` and
+``config.yaml``; the test skips when that is unset or when the upstream
+``auk`` package is not importable::
+
+    AUK_CKPT_DIR=/path/to/ckpts/AuK CUDA_VISIBLE_DEVICES=6 \
+        python -m pytest tests/diffusion/models/auk/test_auk_transformer_parity.py -v -s
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import math
+import os
+from pathlib import Path
+
+import pytest
+import torch
+import yaml
+
+from vllm_omni.diffusion.models.auk.auk_transformer import (
+    AuKTransformer,
+    Rotary,
+    _apply_rope,
+    dit_state_dict,
+    sample_latents,
+)
+
+CKPT_DIR = os.environ.get("AUK_CKPT_DIR")
+
+# Tolerances. The port evaluates rotary embeddings in fp32 and skips the
+# reference's redundant uncond text projection, so the two graphs are not
+# bit-identical under autocast; they agree well inside bf16 noise.
+FORWARD_ATOL = 3e-2
+FORWARD_COS = 0.999
+SAMPLE_ATOL = 5e-2
+SAMPLE_COS = 0.999
+ROPE_ATOL = 1e-5
+
+TARGET_FRAMES = 300
+REF_FRAMES = 150
+TEXT_TOKENS = 40
+SHORT_LENS = (220, 90, 25)  # target, reference, text for the padded batch row
+TIME = 0.37
+SAMPLE_STEPS = 8
+SWAY_COEF = -1.0
+SEED = 1234
+DEVICE = "cuda:0"
+AUTOCAST_DTYPE = torch.bfloat16
+
+
+def _skip_reason() -> str | None:
+    if CKPT_DIR is None:
+        return "AUK_CKPT_DIR is not set"
+    if not Path(CKPT_DIR, "auk_base.safetensors").is_file():
+        return f"no auk_base.safetensors under {CKPT_DIR}"
+    try:
+        # Resolving the submodule, not just "auk", also catches the shadow an
+        # __init__.py in this directory would create.
+        upstream = importlib.util.find_spec("auk.model")
+    except ModuleNotFoundError:
+        upstream = None
+    if upstream is None:
+        return "the upstream 'auk' package is not importable or is shadowed by this directory"
+    if not torch.cuda.is_available():
+        return "no CUDA device"
+    return None
+
+
+SKIP_REASON = _skip_reason()
+
+pytestmark = [
+    pytest.mark.local_model,
+    pytest.mark.diffusion,
+    pytest.mark.gpu,
+    pytest.mark.cuda,
+    pytest.mark.skipif(SKIP_REASON is not None, reason=str(SKIP_REASON)),
+]
+
+
+def _arch() -> dict:
+    """Read the backbone geometry out of the checkpoint's config."""
+    config = yaml.safe_load(Path(CKPT_DIR, "config.yaml").read_text())["model"]
+    return {
+        "dim": config["arch"]["dim"],
+        "heads": config["arch"]["heads"],
+        "dim_head": config["arch"].get("dim_head", 64),
+        "ff_mult": config["arch"]["ff_mult"],
+        "text_hidden_dim": config["arch"]["text_hidden_dim"],
+        "num_layers": config["arch"]["num_layers"],
+        "num_single_layers": config["arch"]["num_single_layers"],
+        "latent_dim": config["vae"]["latent_dim"],
+    }
+
+
+def _stats(port: torch.Tensor, ref: torch.Tensor) -> tuple[float, float, float]:
+    """Return max abs diff, cosine similarity and the reference's RMS."""
+    a = port.flatten().float()
+    b = ref.flatten().float()
+    return (
+        (a - b).abs().max().item(),
+        torch.nn.functional.cosine_similarity(a, b, dim=0).item(),
+        b.pow(2).mean().sqrt().item(),
+    )
+
+
+def _assert_close(name: str, port: torch.Tensor, ref: torch.Tensor, atol: float, min_cos: float) -> None:
+    max_abs, cos, rms = _stats(port, ref)
+    print(f"[{name}] max|Δ|={max_abs:.4e} cos={cos:.8f} ref_rms={rms:.4f} rel={max_abs / rms:.3e}")
+    assert max_abs < atol, f"{name}: max|Δ|={max_abs:.4e} >= {atol:.1e} (reference rms {rms:.4f})"
+    assert cos > min_cos, f"{name}: cos={cos:.8f} <= {min_cos}"
+
+
+def _autocast():
+    return torch.autocast("cuda", dtype=AUTOCAST_DTYPE)
+
+
+@pytest.fixture(scope="module")
+def models() -> tuple[torch.nn.Module, AuKTransformer]:
+    """The reference and the port, same weights, fp32 on the accelerator."""
+    from auk.model.flux2_edit import Flux2Edit
+    from safetensors.torch import load_file
+
+    arch = _arch()
+    device = torch.device(DEVICE)
+    state = dit_state_dict(load_file(str(Path(CKPT_DIR, "auk_base.safetensors")), device="cpu"))
+
+    reference = Flux2Edit(dropout=0.0, attn_backend="torch", attn_mask_enabled=True, **arch)
+    reference.load_state_dict(state, strict=True)
+    reference = reference.to(device).eval().requires_grad_(False)
+
+    port = AuKTransformer(**arch)
+    # strict=True is the point of the port: every released tensor lands.
+    port.load_state_dict(state, strict=True)
+    port = port.to(device).eval().requires_grad_(False)
+
+    del state
+    return reference, port
+
+
+@pytest.fixture(scope="module")
+def inputs() -> dict[str, torch.Tensor]:
+    """Fixed random inputs, as a batch of 2 whose second row is padded."""
+    device = torch.device(DEVICE)
+    generator = torch.Generator(device="cpu").manual_seed(SEED)
+    arch = _arch()
+
+    def randn(*shape: int) -> torch.Tensor:
+        return torch.randn(*shape, generator=generator, dtype=torch.float32).to(device)
+
+    def pad_mask(width: int, short: int) -> torch.Tensor:
+        mask = torch.zeros(2, width, dtype=torch.bool, device=device)
+        mask[0, :] = True
+        mask[1, :short] = True
+        return mask
+
+    mask2 = pad_mask(TARGET_FRAMES, SHORT_LENS[0])
+    ref_mask2 = pad_mask(REF_FRAMES, SHORT_LENS[1])
+    c_mask2 = pad_mask(TEXT_TOKENS, SHORT_LENS[2])
+
+    # Padding rows are all-zero in the text encoder's output, so keep the
+    # inputs consistent with the masks.
+    x2 = randn(2, TARGET_FRAMES, arch["latent_dim"]) * mask2.unsqueeze(-1)
+    ref2 = randn(2, REF_FRAMES, arch["latent_dim"]) * ref_mask2.unsqueeze(-1)
+    text2 = randn(2, TEXT_TOKENS, arch["text_hidden_dim"]) * c_mask2.unsqueeze(-1)
+
+    return {
+        "x": x2[:1].contiguous(),
+        "ref": ref2[:1].contiguous(),
+        "text": text2[:1].contiguous(),
+        "c_mask": c_mask2[:1].contiguous(),
+        "x2": x2,
+        "ref2": ref2,
+        "text2": text2,
+        "mask2": mask2,
+        "ref_mask2": ref_mask2,
+        "c_mask2": c_mask2,
+        "time": torch.tensor(TIME, dtype=torch.float32, device=device),
+    }
+
+
+def test_single_forward(models, inputs) -> None:
+    """One conditional forward: reference against port."""
+    reference, port = models
+    kwargs = dict(text=inputs["text"], time=inputs["time"], c_mask=inputs["c_mask"], ref=inputs["ref"])
+    with torch.no_grad(), _autocast():
+        want = reference(x=inputs["x"], **kwargs)
+        got = port(inputs["x"], **kwargs)
+
+    assert got.shape == want.shape == (1, TARGET_FRAMES, port.latent_dim)
+    _assert_close("single", got, want, FORWARD_ATOL, FORWARD_COS)
+
+
+def test_cfg_forward(models, inputs) -> None:
+    """The doubled cfg_infer forward, cond branch first."""
+    reference, port = models
+    kwargs = dict(
+        text=inputs["text"],
+        time=inputs["time"],
+        c_mask=inputs["c_mask"],
+        ref=inputs["ref"],
+        cfg_infer=True,
+    )
+    with torch.no_grad(), _autocast():
+        want = reference(x=inputs["x"], **kwargs)
+        got = port(inputs["x"], **kwargs)
+
+    assert got.shape == want.shape == (2, TARGET_FRAMES, port.latent_dim)
+    _assert_close("cfg-cond", got[:1], want[:1], FORWARD_ATOL, FORWARD_COS)
+    _assert_close("cfg-uncond", got[1:], want[1:], FORWARD_ATOL, FORWARD_COS)
+
+
+def test_masked_batch_forward(models, inputs) -> None:
+    """Batch of 2 with padded targets, references and text.
+
+    Only valid target frames are compared. Padded frames are unconstrained:
+    the reference leaves them to whatever masked attention emits there.
+    """
+    reference, port = models
+    kwargs = dict(
+        text=inputs["text2"],
+        time=inputs["time"],
+        mask=inputs["mask2"],
+        c_mask=inputs["c_mask2"],
+        ref=inputs["ref2"],
+        ref_mask=inputs["ref_mask2"],
+    )
+    with torch.no_grad(), _autocast():
+        want = reference(x=inputs["x2"], **kwargs)
+        got = port(inputs["x2"], **kwargs)
+
+    assert got.shape == want.shape == (2, TARGET_FRAMES, port.latent_dim)
+    valid = inputs["mask2"].unsqueeze(-1).expand_as(got)
+    _assert_close("masked", got[valid], want[valid], FORWARD_ATOL, FORWARD_COS)
+
+
+def _set_mask_flag(reference, port, enabled: bool) -> None:
+    """Flip attn_mask_enabled on both models in place, so no second copy is loaded."""
+    for module in port.modules():
+        if hasattr(module, "attn_mask_enabled"):
+            module.attn_mask_enabled = enabled
+    for blocks in (reference.transformer_blocks, reference.single_transformer_blocks):
+        for block in blocks:
+            block.attn.processor.attn_mask_enabled = enabled
+
+
+def test_mask_flag_disabled(models, inputs) -> None:
+    """``attn_mask_enabled=False`` must track the reference with its own flag off.
+
+    With the flag off both models run attention unmasked and still zero the
+    padded outputs, so the padded batch is compared in full.
+    """
+    reference, port = models
+    kwargs = dict(
+        text=inputs["text2"],
+        time=inputs["time"],
+        mask=inputs["mask2"],
+        c_mask=inputs["c_mask2"],
+        ref=inputs["ref2"],
+        ref_mask=inputs["ref_mask2"],
+    )
+    with torch.no_grad(), _autocast():
+        masked = port(inputs["x2"], **kwargs)
+        _set_mask_flag(reference, port, False)
+        try:
+            want = reference(x=inputs["x2"], **kwargs)
+            got = port(inputs["x2"], **kwargs)
+        finally:
+            _set_mask_flag(reference, port, True)
+
+    _assert_close("mask-flag-off", got, want, FORWARD_ATOL, FORWARD_COS)
+    # Guard against a vacuous pass: the flag has to actually change something.
+    assert (got.float() - masked.float()).abs().max().item() > 1e-3
+
+
+def test_sampler_generator_matches_seed(models, inputs) -> None:
+    """A generator seeded with 0 must draw exactly what ``seed=0`` draws."""
+    _, port = models
+    common = dict(
+        gen_frames=TARGET_FRAMES,
+        nfe=2,
+        cfg_strength=0.0,
+        sway_sampling_coef=SWAY_COEF,
+        text=inputs["text"],
+        c_mask=inputs["c_mask"],
+        ref=inputs["ref"],
+        ref_mask=None,
+    )
+    with torch.no_grad(), _autocast():
+        by_seed = sample_latents(port, seed=0, **common)
+        by_generator = sample_latents(port, generator=torch.Generator(device=DEVICE).manual_seed(0), **common)
+
+    print(f"[generator] max|Δ|={(by_generator.float() - by_seed.float()).abs().max().item():.4e}")
+    torch.testing.assert_close(by_generator, by_seed, rtol=0, atol=0)
+
+
+def test_cache_needs_clearing(models, inputs) -> None:
+    """The cached text projection is reused until ``clear_cache()``.
+
+    ``sample_latents`` clears on the way out, so a pipeline that reuses the
+    module across requests is safe without doing anything. A caller driving
+    ``forward(cache=True)`` itself has to clear when the text changes.
+    """
+    _, port = models
+    other_text = inputs["text"] * 0.5
+    kwargs = dict(time=inputs["time"], c_mask=inputs["c_mask"], ref=inputs["ref"], cfg_infer=True)
+
+    port.clear_cache()
+    with torch.no_grad(), _autocast():
+        want = port(inputs["x"], inputs["text"], **kwargs)
+        port(inputs["x"], other_text, cache=True, **kwargs)
+        stale = port(inputs["x"], inputs["text"], cache=True, **kwargs)
+        port.clear_cache()
+        fresh = port(inputs["x"], inputs["text"], cache=True, **kwargs)
+        port.clear_cache()
+
+    # The cache is real: the second call reused the first call's projection.
+    assert (stale.float() - want.float()).abs().max().item() > 1e-3
+    torch.testing.assert_close(fresh, want, rtol=0, atol=0)
+    assert port.text_cond is None and port.text_uncond is None
+
+
+def test_rope_matches_x_transformers() -> None:
+    """The local rotary implementation against x_transformers, in fp32."""
+    from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb
+
+    device = torch.device(DEVICE)
+    dim_head, seq_len = 64, TEXT_TOKENS + REF_FRAMES + TARGET_FRAMES
+    q = torch.randn(2, 24, seq_len, dim_head, dtype=torch.float32, device=device)
+
+    freqs = RotaryEmbedding(dim_head).to(device).forward_from_seq_len(seq_len)[0]
+    want = apply_rotary_pos_emb(q, freqs)
+    got = _apply_rope(q, Rotary(dim_head).to(device)(seq_len))
+
+    max_abs, cos, _ = _stats(got, want)
+    print(f"[rope] max|Δ|={max_abs:.4e} cos={cos:.10f}")
+    assert max_abs < ROPE_ATOL, f"rope: max|Δ|={max_abs:.4e} >= {ROPE_ATOL:.1e}"
+
+
+def test_rope_survives_half_precision_cast() -> None:
+    """Casting the module to bf16 must not round the rotary frequencies.
+
+    The reference stores them in a buffer, so ``.to(bfloat16)`` rounds both the
+    frequencies and the integer positions and silently changes every rotation.
+    The port recomputes them instead.
+    """
+    dim_head, seq_len = 64, TEXT_TOKENS + REF_FRAMES + TARGET_FRAMES
+    exact = Rotary(dim_head)(seq_len)
+    got = Rotary(dim_head).to(torch.bfloat16)(seq_len)
+
+    assert got.dtype == torch.float32
+    torch.testing.assert_close(got, exact, rtol=0, atol=0)
+
+
+def _reference_euler(reference, *, x, text, c_mask, ref, steps, cfg_strength, sway) -> torch.Tensor:
+    """Upstream's ODE loop, transcribed.
+
+    ``CFMEdit.sample`` hands the backbone to ``torchdiffeq.odeint(method=
+    "euler")``, which on a fixed grid is ``y += (t1 - t0) * f(t0, y)``.
+    Constructing a real ``CFMEdit`` would drag in the Qwen-Omni text encoder,
+    so only the backbone here comes from upstream.
+    """
+    t = torch.linspace(0, 1, steps + 1, device=x.device, dtype=torch.float32)
+    if sway is not None:
+        t = t + sway * (torch.cos(math.pi / 2 * t) - 1 + t)
+
+    shared = dict(text=text, c_mask=c_mask, ref=ref, mask=None, ref_mask=None, cache=True)
+    for i in range(steps):
+        if cfg_strength < 1e-5:
+            v = reference(x=x, time=t[i], **shared)
+        else:
+            pred = reference(x=x, time=t[i], cfg_infer=True, **shared)
+            v_cond, v_uncond = torch.chunk(pred, 2, dim=0)
+            v = v_cond + (v_cond - v_uncond) * cfg_strength
+        x = x + (t[i + 1] - t[i]) * v
+    reference.clear_cache()
+    return x
+
+
+@pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
+def test_euler_sample(models, inputs, cfg_strength: float) -> None:
+    """An 8-step Euler trajectory, with and without guidance."""
+    reference, port = models
+    shared = dict(text=inputs["text"], c_mask=inputs["c_mask"], ref=inputs["ref"])
+
+    with torch.no_grad(), _autocast():
+        got = sample_latents(
+            port,
+            gen_frames=TARGET_FRAMES,
+            ref_mask=None,
+            nfe=SAMPLE_STEPS,
+            cfg_strength=cfg_strength,
+            sway_sampling_coef=SWAY_COEF,
+            seed=SEED,
+            **shared,
+        )
+        # Re-seeding reproduces the noise sample_latents drew for itself.
+        torch.manual_seed(SEED)
+        noise = torch.randn(
+            TARGET_FRAMES, port.latent_dim, device=inputs["ref"].device, dtype=inputs["ref"].dtype
+        ).unsqueeze(0)
+        want = _reference_euler(
+            reference, x=noise, steps=SAMPLE_STEPS, cfg_strength=cfg_strength, sway=SWAY_COEF, **shared
+        )
+
+    assert got.shape == want.shape == (1, TARGET_FRAMES, port.latent_dim)
+    _assert_close(f"sample-cfg{cfg_strength:g}", got, want, SAMPLE_ATOL, SAMPLE_COS)

tests/diffusion/models/auk/test_auk_vae_parity.py

@@ -0,0 +1,254 @@
+#!/usr/bin/env python
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""AuK VAE parity: the numeric oracle for the vendored BigVGAN-flow codec.
+
+Checks ``AuKVAE`` against the AuK reference implementation on real speech, with
+both models loaded from the same released ``vae.safetensors``:
+
+  * encoder statistics (mean / log-std) agree,
+  * ``encode(sample=False)`` reproduces the reference mean path,
+  * ``encode(sample=True)`` reproduces ``encoding_and_normalization`` for a
+    matching RNG seed,
+  * ``decode`` reproduces ``denormalize`` + ``inference_from_latents``, in fp32
+    and (reported, not gated tightly) under bf16 autocast,
+  * the waveform survives a mean-path round trip.
+
+Runs on GPU in fp32 and needs both the released checkpoint and the reference
+``auk`` package importable, so it is a local_model test rather than CI. Point
+``AUK_CKPT_DIR`` at the directory holding ``vae.safetensors`` and ``AUK_REF_WAV``
+at a speech clip (default ``$AUK_CKPT_DIR/ref.wav``)::
+
+    AUK_CKPT_DIR=... AUK_REF_WAV=... python -m pytest \
+        tests/diffusion/models/auk/test_auk_vae_parity.py -v -s
+
+Unlike its sibling directories this one carries no ``__init__.py`` on purpose: it would put
+``tests/diffusion/models`` on ``sys.path`` ahead of site-packages, and this directory's own name
+would then shadow the reference ``auk`` package the parity check imports.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+from pathlib import Path
+
+import pytest
+import soundfile as sf
+import torch
+import torchaudio.functional as AF
+
+from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE
+
+pytestmark = [pytest.mark.local_model, pytest.mark.diffusion, pytest.mark.gpu, pytest.mark.cuda]
+
+# The released model.vae.model_init_kwargs, shared by both models under test.
+VAE_CONFIG = {
+    "upsample_rates": [5, 4, 3, 2, 2, 2],
+    "upsample_kernel_sizes": [10, 8, 6, 4, 4, 4],
+    "upsample_initial_channel": 1536,
+    "resblock_kernel_sizes": [3, 7, 11],
+    "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+    "downsample_rates": [2, 2, 2, 3, 4, 5],
+    "downsample_channels": [12, 24, 48, 96, 192, 384, 768],
+    "snake_logscale": True,
+    "latent_dim": 64,
+    "use_vae": True,
+    "causal": True,
+    "flow_hidden_channels": 256,
+    "act_causal": True,
+}
+
+SAMPLE_RATE = 24000
+HOP_SIZE = 480
+LATENT_DIM = 64
+SEED = 1234
+# Keep the clip short: the parity run shares its GPU with other jobs.
+MAX_SECONDS = 4.0
+TOLERANCE = 1e-5
+DEVICE = "cuda"
+
+_CKPT_DIR = os.environ.get("AUK_CKPT_DIR")
+_VAE_PATH = Path(_CKPT_DIR) / "vae.safetensors" if _CKPT_DIR else None
+_REF_WAV = Path(os.environ.get("AUK_REF_WAV") or (Path(_CKPT_DIR) / "ref.wav" if _CKPT_DIR else "unset"))
+
+
+def _reference_importable() -> bool:
+    """Resolve ``auk.model``, not ``auk``.
+
+    An ``__init__.py`` in this directory would make its own name resolve as ``auk`` and satisfy a
+    top-level probe, leaving the fixtures to fail at setup instead of skipping. The submodule
+    probe fails in that case, which is the honest answer.
+    """
+    try:
+        return importlib.util.find_spec("auk.model") is not None
+    except (ImportError, ValueError):
+        return False
+
+
+def _skip_reason() -> str | None:
+    if not _CKPT_DIR:
+        return "set AUK_CKPT_DIR to the directory holding vae.safetensors"
+    if _VAE_PATH is None or not _VAE_PATH.is_file():
+        return f"no vae.safetensors under {_CKPT_DIR}"
+    if not _REF_WAV.is_file():
+        return f"no reference clip at {_REF_WAV} (set AUK_REF_WAV)"
+    if not _reference_importable():
+        return "cannot import the reference auk.model package (an __init__.py here would shadow it)"
+    if not torch.cuda.is_available():
+        return "parity runs in fp32 on GPU"
+    return None
+
+
+pytestmark.append(pytest.mark.skipif(_skip_reason() is not None, reason=_skip_reason() or ""))
+
+
+def _max_abs(lhs: torch.Tensor, rhs: torch.Tensor) -> float:
+    return (lhs.float() - rhs.float()).abs().max().item()
+
+
+def _snr_db(signal: torch.Tensor, estimate: torch.Tensor) -> float:
+    noise = (signal.float() - estimate.float()).pow(2).sum()
+    return (10.0 * torch.log10(signal.float().pow(2).sum() / noise)).item()
+
+
+@pytest.fixture(scope="module")
+def clip() -> torch.Tensor:
+    """Mono 24 kHz speech, trimmed to a whole number of latent frames."""
+    audio, sample_rate = sf.read(str(_REF_WAV), dtype="float32", always_2d=True)  # (T, C)
+    wav = torch.from_numpy(audio).mean(dim=1, keepdim=True).transpose(0, 1)
+    if sample_rate != SAMPLE_RATE:
+        wav = AF.resample(wav, sample_rate, SAMPLE_RATE)
+    samples = min(int(MAX_SECONDS * SAMPLE_RATE), wav.size(-1)) // HOP_SIZE * HOP_SIZE
+    wav = wav[:, :samples].contiguous().to(DEVICE)
+    print(f"\n[clip] {_REF_WAV.name}: {samples} samples, {samples / SAMPLE_RATE:.2f} s, {samples // HOP_SIZE} frames")
+    return wav
+
+
+@pytest.fixture(scope="module")
+def reference(clip: torch.Tensor) -> dict[str, torch.Tensor]:
+    """Reference tensors from the upstream model, which is then released so one model fits at a time."""
+    from auk.model.vae import load_vae_model
+    from auk.model.vae.bigvgan_flow_vae import BigVGANFlowVAEConfig
+
+    model = load_vae_model(
+        vae_name="BigVGANFlowVAE",
+        vae_cfg=BigVGANFlowVAEConfig.from_dict(VAE_CONFIG),
+        vae_ckpt=str(_VAE_PATH),
+        map_location="cpu",
+    )
+    model = model.to(DEVICE).eval()
+    model.requires_grad_(False)
+
+    with torch.no_grad():
+        mean, log_std = model.audio_encoder(clip.unsqueeze(1)).chunk(2, 1)
+        scale = torch.sqrt(model.global_log_std.float())
+        mean_latents = (mean.transpose(1, 2).float() - model.global_mean.float()) / scale
+        torch.manual_seed(SEED)
+        sampled_latents, _ = model.encoding_and_normalization(clip.unsqueeze(1))
+        decoded = model.inference_from_latents(model.denormalize(mean_latents).permute(0, 2, 1))
+
+    tensors = {
+        "mean": mean.clone(),
+        "log_std": log_std.clone(),
+        "mean_latents": mean_latents.clone(),
+        "sampled_latents": sampled_latents.clone(),
+        "decoded": decoded.squeeze(1).clone(),
+    }
+    del model
+    torch.accelerator.empty_cache()
+    return tensors
+
+
+@pytest.fixture(scope="module")
+def port(reference: dict[str, torch.Tensor]) -> tuple[AuKVAE, list[str], list[str]]:
+    """The vendored module, loaded from the same checkpoint. Depends on ``reference`` for load ordering.
+
+    Placed on the accelerator before loading so weight norm is folded there; see ``load_weights``.
+    """
+    model = AuKVAE.from_config(VAE_CONFIG).to(DEVICE)
+    missing, unexpected = model.load_weights(_VAE_PATH)
+    model = model.eval()
+    model.requires_grad_(False)
+    return model, missing, unexpected
+
+
+def test_checkpoint_covers_every_inference_tensor(port: tuple[AuKVAE, list[str], list[str]]) -> None:
+    model, missing, unexpected = port
+    print(f"[keys] missing={len(missing)} unexpected={len(unexpected)}")
+    print(f"[keys] unexpected prefixes={sorted({key.split('.')[0] for key in unexpected})}")
+    assert missing == []
+    assert unexpected and all(key.startswith("flow.") for key in unexpected)
+    assert (model.hop_size, model.sample_rate, model.latent_dim) == (HOP_SIZE, SAMPLE_RATE, LATENT_DIM)
+
+
+def test_encoder_statistics_match(
+    clip: torch.Tensor, reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]]
+) -> None:
+    model = port[0]
+    with torch.no_grad():
+        mean, log_std = model.audio_encoder(clip.unsqueeze(1)).chunk(2, 1)
+    mean_diff = _max_abs(mean, reference["mean"])
+    log_std_diff = _max_abs(log_std, reference["log_std"])
+    print(f"[encoder] max|dmean|={mean_diff:.3e} max|dlog_std|={log_std_diff:.3e}")
+    assert mean_diff < TOLERANCE
+    assert log_std_diff < TOLERANCE
+    # Folded and hook-computed weights are bit-identical on one device, so the whole encoder is.
+    assert (mean_diff, log_std_diff) == (0.0, 0.0)
+
+
+def test_encode_mean_path_matches(
+    clip: torch.Tensor, reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]]
+) -> None:
+    model = port[0]
+    with torch.no_grad():
+        latents = model.encode(clip)
+    assert latents.shape == (1, clip.size(-1) // HOP_SIZE, LATENT_DIM)
+    diff = _max_abs(latents, reference["mean_latents"])
+    print(f"[encode mean] shape={tuple(latents.shape)} max|d|={diff:.3e}")
+    assert diff == 0.0
+
+
+def test_encode_sampled_matches_seeded_reference(
+    clip: torch.Tensor, reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]]
+) -> None:
+    model = port[0]
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(SEED)
+    with torch.no_grad():
+        latents = model.encode(clip, sample=True, generator=generator)
+    diff = _max_abs(latents, reference["sampled_latents"])
+    spread = _max_abs(latents, reference["mean_latents"])
+    print(f"[encode sample] max|d vs seeded reference|={diff:.3e} max|d vs mean path|={spread:.3e}")
+    assert diff < TOLERANCE
+    assert diff == 0.0  # a fresh device Generator replays the default one after torch.manual_seed
+    assert spread > TOLERANCE  # the draw actually moved the latents
+
+
+def test_decode_matches(reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]]) -> None:
+    model = port[0]
+    latents = reference["mean_latents"]
+    with torch.no_grad():
+        wav = model.decode(latents)
+    assert wav.shape == (1, latents.size(1) * HOP_SIZE)
+    diff = _max_abs(wav, reference["decoded"])
+    print(f"[decode fp32] shape={tuple(wav.shape)} max|d|={diff:.3e}")
+    assert diff < TOLERANCE
+    assert diff == 0.0
+
+    with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        wav_bf16 = model.decode(latents)
+    bf16_diff = _max_abs(wav_bf16, reference["decoded"])
+    print(f"[decode bf16 autocast] max|d| vs fp32 reference={bf16_diff:.3e} snr={_snr_db(wav, wav_bf16):.2f} dB")
+    assert torch.isfinite(wav_bf16).all()
+    assert bf16_diff < 0.5  # sanity only: bf16 is reported, not held to fp32 parity
+
+
+def test_round_trip_snr(clip: torch.Tensor, port: tuple[AuKVAE, list[str], list[str]]) -> None:
+    model = port[0]
+    with torch.no_grad():
+        wav = model.decode(model.encode(clip))
+    assert wav.shape == clip.shape
+    snr = _snr_db(clip, wav)
+    print(f"[round trip] snr={snr:.2f} dB peak={wav.abs().max().item():.4f}")
+    assert snr > 10.0

tests/diffusion/models/auk/test_pipeline_auk.py

@@ -0,0 +1,514 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""AuK stage-1 pipeline: request parsing on CPU, generation parity on a GPU.
+
+The CPU tests stub the transformer, the codec and the ODE so they exercise only
+what the pipeline itself owns: how a request turns into a target length, a
+sampling schedule and a noise generator.
+
+The parity test replays the saved upstream reference for the zero-shot TTS case
+and reports latent, per-frame and waveform distances rather than asserting bit
+equality, because two differences live in the conditioning rather than in the
+port: the saved text condition is an fp32 thinker forward while the reference
+generated under bf16 autocast, and the reference's stochastic VAE draw came from
+the process-global RNG seeded at 1234.
+
+It replays with ``vae_sample=False``. That is the deliberate choice: the
+reference re-seeds the global RNG immediately before drawing the ODE noise, so a
+per-request generator reproduces that noise exactly (verified bit-identical) as
+long as nothing else has drawn from it. Sampling the VAE posterior consumes the
+generator first, which moves the noise and produces a different, equally valid
+realization; comparing latents then measures the noise, not the port. The
+posterior mean costs a reference latent that is 0.045 MSE (0.969 cosine) away
+from the draw the reference used, which is the smaller of the two errors.
+Generate the reference artifacts once with the upstream package installed
+(``tools/auk_parity_reference.py --auk-repo ... --ckpt-dir ... --qwen-dir ...
+--out /path/to/auk-parity``; it writes ``parity_ref/<variant>`` and
+``fusion_ref``), then run::
+
+    AUK_OMNI_CKPT_DIR=/path/to/auk-omni-base \
+    AUK_PARITY_REF=/path/to/auk-parity/parity_ref/base \
+    python -m pytest -s tests/diffusion/models/auk/test_pipeline_auk.py
+
+``AUK_OMNI_CKPT_DIR`` is the assembled directory this pipeline loads, which is
+not the same thing as the released AuK snapshot that the transformer and codec
+parity tests in this directory read from ``AUK_CKPT_DIR``. ``AUK_CKPT_DIR`` is
+accepted as a fallback, and the test skips rather than fails when the directory
+it resolves to is not an assembled one.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+import soundfile as sf
+import torch
+import torch.nn.functional as F
+import torchaudio
+from torch import nn
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.models.auk import pipeline_auk
+from vllm_omni.diffusion.models.auk.pipeline_auk import AuKPipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+LATENT_DIM = 64
+HOP = 480
+SAMPLE_RATE = 24000
+TEXT_HIDDEN_DIM = 2048
+
+# The distilled student's grid, exact values from the released checkpoint.
+FLASH_T_GRID = [0.0, 0.07612049579620361, 0.2928932309150696, 0.6173166036605835, 1.0]
+
+CKPT_DIR = os.environ.get("AUK_OMNI_CKPT_DIR") or os.environ.get("AUK_CKPT_DIR")
+PARITY_REF = os.environ.get("AUK_PARITY_REF")
+# An assembled directory is the only thing this pipeline can load.
+IS_ASSEMBLED = bool(CKPT_DIR) and (Path(CKPT_DIR).expanduser() / "config.json").is_file()
+
+# Measured on the released base checkpoint: mean frame cosine 0.961 to 0.968,
+# relative latent MSE 0.051 to 0.063, log-mel distance 0.34. That spread is
+# bf16 against fp32 DiT weights plus where the codec folds its weight norm,
+# neither of which is worth gating on. A port that has lost the conditioning
+# sits at 0.110 cosine, 1.82 relative MSE and 2.88 log-mel, so these gates
+# still separate the two by an order of magnitude.
+PARITY_MIN_FRAME_COSINE = 0.93
+PARITY_MAX_RELATIVE_MSE = 0.15
+PARITY_MAX_MEL_DISTANCE = 0.6
+
+
+class _StubTransformer(nn.Module):
+    """Records its construction arguments; owns one weight so strict loading runs."""
+
+    def __init__(
+        self,
+        *,
+        dim: int,
+        heads: int,
+        dim_head: int,
+        ff_mult: float,
+        latent_dim: int,
+        text_hidden_dim: int,
+        num_layers: int,
+        num_single_layers: int,
+        attn_mask_enabled: bool = True,
+    ) -> None:
+        super().__init__()
+        self.init_kwargs = {
+            "dim": dim,
+            "heads": heads,
+            "dim_head": dim_head,
+            "ff_mult": ff_mult,
+            "latent_dim": latent_dim,
+            "text_hidden_dim": text_hidden_dim,
+            "num_layers": num_layers,
+            "num_single_layers": num_single_layers,
+            "attn_mask_enabled": attn_mask_enabled,
+        }
+        self.latent_dim = latent_dim
+        self.proj = nn.Linear(latent_dim, latent_dim, bias=False)
+
+
+class _StubVAE(nn.Module):
+    """Deterministic stand-in for the codec, one latent frame per hop."""
+
+    hop_size = HOP
+    sample_rate = SAMPLE_RATE
+    latent_dim = LATENT_DIM
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.encode_calls: list[dict[str, Any]] = []
+        self.decode_calls = 0
+        self.weights_path: str | None = None
+        self.config: dict[str, Any] = {}
+
+    @classmethod
+    def from_config(cls, cfg: dict[str, Any]) -> _StubVAE:
+        vae = cls()
+        vae.config = dict(cfg)
+        return vae
+
+    def load_weights(self, path: str, **_: Any) -> tuple[list[str], list[str]]:
+        self.weights_path = str(path)
+        return [], []
+
+    def encode(
+        self,
+        wav: torch.Tensor,
+        *,
+        sample: bool = False,
+        generator: torch.Generator | None = None,
+    ) -> torch.Tensor:
+        frames = wav.shape[-1] // HOP
+        self.encode_calls.append({"samples": int(wav.shape[-1]), "sample": sample, "generator": generator})
+        if sample:
+            return torch.randn(1, frames, LATENT_DIM, generator=generator, device=wav.device)
+        return torch.zeros(1, frames, LATENT_DIM, device=wav.device)
+
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        self.decode_calls += 1
+        return torch.zeros(1, latents.shape[1] * HOP, device=latents.device)
+
+
+def _stub_sampler(calls: list[dict[str, Any]]):
+    """Record the ODE arguments and return noise drawn from the request generator."""
+
+    def sample_latents(dit: nn.Module, **kwargs: Any) -> torch.Tensor:
+        calls.append(kwargs)
+        shape = (1, kwargs["gen_frames"], kwargs["latent_dim"])
+        generator = kwargs.get("generator")
+        if generator is None:
+            return torch.zeros(*shape, device=kwargs["device"], dtype=kwargs["dtype"])
+        return torch.randn(*shape, generator=generator, device=kwargs["device"], dtype=kwargs["dtype"])
+
+    return sample_latents
+
+
+def _write_checkpoint(root: Path, variant: str) -> Path:
+    """Write a tiny assembled checkpoint the stubs can load."""
+
+    from safetensors.torch import save_file
+
+    model_dir = root / f"auk-{variant}"
+    model_dir.mkdir(parents=True, exist_ok=True)
+    config = {
+        "model_type": "auk",
+        "architectures": ["AuKForConditionalGeneration"],
+        "dit": {
+            "dim": 16,
+            "heads": 2,
+            "dim_head": 8,
+            "ff_mult": 2,
+            "text_hidden_dim": TEXT_HIDDEN_DIM,
+            "num_layers": 1,
+            "num_single_layers": 1,
+            "attn_mask_enabled": True,
+        },
+        "vae": {
+            "latent_dim": LATENT_DIM,
+            "downsample_rate": HOP,
+            "target_sample_rate": SAMPLE_RATE,
+            "model_init_kwargs": {"latent_dim": LATENT_DIM},
+        },
+        "variant": variant,
+        "flash_t_grid": FLASH_T_GRID,
+        "defaults": {"nfe": 32, "cfg": 2.0, "sway": -1.0},
+    }
+    (model_dir / "config.json").write_text(json.dumps(config))
+    save_file(
+        {
+            "transformer.proj.weight": torch.zeros(LATENT_DIM, LATENT_DIM),
+            "layer_weights": torch.zeros(36),
+            "layer_scale": torch.ones(1),
+        },
+        str(model_dir / "auk.safetensors"),
+    )
+    save_file({"global_mean": torch.zeros(LATENT_DIM)}, str(model_dir / "vae.safetensors"))
+    return model_dir
+
+
+@pytest.fixture
+def build_pipeline(tmp_path, monkeypatch):
+    """Build an AuKPipeline whose transformer, codec and ODE are stubs."""
+
+    monkeypatch.setattr(pipeline_auk, "get_local_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(pipeline_auk, "AuKTransformer", _StubTransformer)
+    monkeypatch.setattr(pipeline_auk, "AuKVAE", _StubVAE)
+
+    def _build(variant: str = "base") -> tuple[AuKPipeline, list[dict[str, Any]]]:
+        calls: list[dict[str, Any]] = []
+        monkeypatch.setattr(pipeline_auk, "sample_latents", _stub_sampler(calls))
+        od_config = OmniDiffusionConfig(
+            model=str(_write_checkpoint(tmp_path, variant)),
+            dtype=torch.float32,
+            model_class_name="AuKPipeline",
+        )
+        return AuKPipeline(od_config=od_config), calls
+
+    return _build
+
+
+def _prompt(
+    *,
+    tokens: int = 8,
+    audio: Any = None,
+    knobs: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Build the diffusion prompt the stage input processor would emit."""
+
+    prompt: dict[str, Any] = {
+        "prompt": "",
+        "prompt_embeds": torch.zeros(tokens, TEXT_HIDDEN_DIM),
+        "additional_information": {"auk": dict(knobs or {})},
+    }
+    if audio is not None:
+        prompt["multi_modal_data"] = {"audio": audio}
+    return prompt
+
+
+def _batch(prompt: dict[str, Any], **sampling: Any) -> DiffusionRequestBatch:
+    request = OmniDiffusionRequest(
+        prompt=prompt,
+        sampling_params=OmniDiffusionSamplingParams(**sampling),
+        request_id="auk-test-0",
+    )
+    return DiffusionRequestBatch(requests=[request])
+
+
+def _silence(seconds: float, sample_rate: int = SAMPLE_RATE) -> tuple[np.ndarray, int]:
+    return np.zeros(int(seconds * sample_rate), dtype=np.float32), sample_rate
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+class TestRequestParsing:
+    def test_pipeline_declares_audio_output(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+
+        assert pipeline.support_audio_output is True
+        assert pipeline.audio_sample_rate == SAMPLE_RATE
+        assert pipeline.supports_request_batch is False
+        # The warmup request cannot carry an encoder-stage text condition.
+        assert pipeline.dummy_run_num_frames == 0
+
+    def test_config_reaches_the_transformer(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+
+        assert pipeline.dit.init_kwargs["latent_dim"] == LATENT_DIM
+        assert pipeline.dit.init_kwargs["text_hidden_dim"] == TEXT_HIDDEN_DIM
+        assert pipeline.dit.init_kwargs["attn_mask_enabled"] is True
+
+    def test_gen_seconds_sets_the_target_length(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+
+        outputs = pipeline.forward(_batch(_prompt(audio=_silence(2.0), knobs={"gen_seconds": 6.0}), seed=1))
+
+        assert calls[0]["gen_frames"] == math.ceil(6.0 * SAMPLE_RATE / HOP)
+        assert calls[0]["ref"].shape == (1, 100, LATENT_DIM)
+        assert calls[0]["ref_mask"].shape == (1, 100)
+        assert calls[0]["c_mask"].shape == (1, 8)
+        assert outputs[0].output.shape == (300 * HOP,)
+        assert outputs[0].output.dtype is torch.float32
+
+    def test_target_length_falls_back_to_the_source_clip(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+
+        pipeline.forward(_batch(_prompt(audio=_silence(2.0), knobs={"gen_seconds": None}), seed=1))
+
+        assert calls[0]["gen_frames"] == 100
+
+    def test_text_only_request_without_a_duration_is_rejected(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+
+        with pytest.raises(ValueError, match="gen_seconds"):
+            pipeline.forward(_batch(_prompt(knobs={"gen_seconds": None}), seed=1))
+        assert calls == []
+
+    def test_text_only_request_gets_an_empty_reference(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+
+        pipeline.forward(_batch(_prompt(knobs={"gen_seconds": 3.5}), seed=1))
+
+        assert calls[0]["ref"].shape == (1, 0, LATENT_DIM)
+        assert calls[0]["gen_frames"] == 175
+
+    def test_source_clip_is_resampled_to_the_codec_rate(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+
+        pipeline.forward(_batch(_prompt(audio=_silence(1.0, 16000), knobs={"gen_seconds": 1.0}), seed=1))
+
+        assert pipeline.vae.encode_calls[0]["samples"] == SAMPLE_RATE
+
+    def test_stereo_source_is_mixed_to_mono(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+        stereo = np.zeros((2, SAMPLE_RATE), dtype=np.float32)
+
+        pipeline.forward(_batch(_prompt(audio=(stereo, SAMPLE_RATE), knobs={"gen_seconds": 1.0}), seed=1))
+
+        assert pipeline.vae.encode_calls[0]["samples"] == SAMPLE_RATE
+
+    def test_missing_text_condition_is_rejected(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+        prompt = _prompt(knobs={"gen_seconds": 1.0})
+        del prompt["prompt_embeds"]
+
+        with pytest.raises(ValueError, match="prompt_embeds"):
+            pipeline.forward(_batch(prompt, seed=1))
+
+    def test_base_variant_uses_the_requested_schedule(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+
+        pipeline.forward(
+            _batch(
+                _prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "sway": -1.0}),
+                seed=1,
+                num_inference_steps=16,
+                guidance_scale=3.0,
+            )
+        )
+
+        assert calls[0]["nfe"] == 16
+        assert calls[0]["cfg_strength"] == 3.0
+        assert calls[0]["sway_sampling_coef"] == -1.0
+        assert calls[0]["t_grid"] is None
+
+    def test_schedule_defaults_come_from_the_checkpoint(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+
+        pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "sway": None}), seed=1))
+
+        assert calls[0]["nfe"] == 32
+        assert calls[0]["cfg_strength"] == 2.0
+        assert calls[0]["sway_sampling_coef"] == -1.0
+
+    def test_flash_variant_locks_the_distilled_recipe(self, build_pipeline):
+        pipeline, calls = build_pipeline("flash")
+
+        pipeline.forward(
+            _batch(
+                _prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "sway": -1.0}),
+                seed=1,
+                num_inference_steps=32,
+                guidance_scale=2.0,
+            )
+        )
+
+        assert calls[0]["nfe"] == 4
+        assert calls[0]["cfg_strength"] == 0.0
+        assert calls[0]["sway_sampling_coef"] is None
+        assert calls[0]["t_grid"] == FLASH_T_GRID
+
+    def test_noise_comes_from_a_seeded_request_generator(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+
+        outputs = pipeline.forward(
+            _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=0, output_type="latent"),
+        )
+
+        expected = torch.randn(1, 50, LATENT_DIM, generator=torch.Generator().manual_seed(0))
+        assert torch.equal(outputs[0].output, expected)
+        # The global RNG is never seeded on the pipeline's behalf.
+        assert calls[0]["seed"] is None
+        assert isinstance(calls[0]["generator"], torch.Generator)
+
+    def test_latent_output_type_skips_the_decoder(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+
+        outputs = pipeline.forward(_batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1, output_type="latent"))
+
+        assert outputs[0].output.shape == (1, 50, LATENT_DIM)
+        assert pipeline.vae.decode_calls == 0
+
+    def test_vae_sample_draws_from_the_request_generator(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+
+        pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "vae_sample": True}), seed=1))
+
+        call = pipeline.vae.encode_calls[0]
+        assert call["sample"] is True
+        assert isinstance(call["generator"], torch.Generator)
+
+    def test_one_request_per_forward(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+        batch = _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1)
+        batch.requests.append(batch.requests[0])
+
+        with pytest.raises(AssertionError, match="one request per forward"):
+            pipeline.forward(batch)
+
+
+def _reference_audio_path(messages: list[dict[str, Any]]) -> str:
+    for message in messages:
+        for item in message.get("content") or ():
+            if isinstance(item, dict) and item.get("type") == "audio":
+                path = item.get("audio") or item.get("audio_url")
+                if path:
+                    return str(path)
+    raise AssertionError("the saved reference case carries no source clip")
+
+
+def _log_mel_distance(left: torch.Tensor, right: torch.Tensor, sample_rate: int) -> float:
+    """Mean absolute log-mel difference between two mono waveforms."""
+
+    mel = torchaudio.transforms.MelSpectrogram(sample_rate=sample_rate, n_fft=1024, hop_length=256, n_mels=80)
+    length = min(left.shape[-1], right.shape[-1])
+    spectra = [torch.log(mel(wav[..., :length].float()) + 1e-5) for wav in (left, right)]
+    return float((spectra[0] - spectra[1]).abs().mean())
+
+
+@pytest.mark.local_model
+@pytest.mark.diffusion
+@pytest.mark.skipif(
+    not (IS_ASSEMBLED and PARITY_REF),
+    reason="needs AUK_PARITY_REF and AUK_OMNI_CKPT_DIR pointing at an assembled checkpoint directory",
+)
+def test_zero_shot_tts_tracks_the_upstream_reference():
+    """Replay the saved zero-shot TTS case and report the distances."""
+
+    parity_dir = Path(PARITY_REF)
+    fusion_dir = Path(os.environ.get("AUK_FUSION_REF") or parity_dir.parent.parent / "fusion_ref")
+    reference = torch.load(parity_dir / "zs_tts_en.pt", map_location="cpu", weights_only=False)
+    fusion = torch.load(fusion_dir / "zs_tts_audio.pt", map_location="cpu", weights_only=False)
+
+    # soundfile rather than torchaudio.load: the repo's test convention, and
+    # torchaudio's decoder needs an ffmpeg the port environment does not have.
+    samples, source_rate = sf.read(_reference_audio_path(reference["messages"]), dtype="float32", always_2d=True)
+    source = torch.from_numpy(samples).transpose(0, 1)
+    knobs = {"gen_seconds": reference["gen_seconds"], "sway": -1.0, "t_grid": None, "vae_sample": False}
+    prompt = {
+        "prompt": "",
+        # The encoder stage emits bf16; the saved fusion is an fp32 forward.
+        "prompt_embeds": fusion["fused"].to(torch.bfloat16),
+        "multi_modal_data": {"audio": (source, int(source_rate))},
+        "additional_information": {"auk": knobs},
+    }
+    od_config = OmniDiffusionConfig(model=CKPT_DIR, dtype=torch.bfloat16, model_class_name="AuKPipeline")
+    pipeline = AuKPipeline(od_config=od_config)
+
+    def run(output_type: str) -> torch.Tensor:
+        batch = _batch(
+            dict(prompt),
+            seed=reference["seed"],
+            num_inference_steps=32,
+            guidance_scale=2.0,
+            output_type=output_type,
+        )
+        return pipeline.forward(batch)[0].output
+
+    # Two deterministic runs: the ODE is seeded, so the latents belong to the
+    # waveform even though the decode happens in the second call.
+    latents = run("latent")
+    waveform = run("pt")
+
+    ref_frames = int(reference["ref_latent_lens"][0])
+    expected = reference["generated"][:, ref_frames:, :].float()
+    assert latents.shape == expected.shape
+    assert torch.isfinite(latents).all()
+    assert torch.isfinite(waveform).all()
+    assert waveform.shape[-1] == expected.shape[1] * HOP
+
+    cosine = F.cosine_similarity(latents[0], expected[0], dim=-1)
+    metrics = {
+        "gen_frames": int(expected.shape[1]),
+        "ref_frames": ref_frames,
+        "latent_mse": float(torch.mean((latents - expected) ** 2)),
+        "latent_relative_mse": float(torch.mean((latents - expected) ** 2) / expected.var()),
+        "frame_cosine_mean": float(cosine.mean()),
+        "frame_cosine_min": float(cosine.min()),
+        "log_mel_distance": _log_mel_distance(waveform, reference["audio"].reshape(-1), int(reference["sr"])),
+        "rms_port": float(waveform.pow(2).mean().sqrt()),
+        "rms_reference": float(reference["audio"].reshape(-1).pow(2).mean().sqrt()),
+    }
+    print("\nAuK zero-shot TTS parity:", json.dumps(metrics, indent=1))
+
+    assert metrics["frame_cosine_mean"] > PARITY_MIN_FRAME_COSINE
+    assert metrics["latent_relative_mse"] < PARITY_MAX_RELATIVE_MSE
+    assert metrics["log_mel_distance"] < PARITY_MAX_MEL_DISTANCE