From c111f3a2d6048848cd239b9cbce37cd55001d57c Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Fri, 11 Sep 2026 13:09:58 -0700 Subject: [PATCH] [Model] Add Tencent AuK speech generation and editing (encoder + diffusion pipeline) (#7385) Signed-off-by: Yueqian Lin Co-authored-by: Sy03 <1370724210@qq.com> --- docs/models/supported_models.md | 1 + .../text_to_speech/README.md | 1 + recipes/README.md | 1 + recipes/Tencent/AuK-H100.md | 167 +++++ .../models/auk/test_auk_transformer_parity.py | 439 ++++++++++++ .../models/auk/test_auk_vae_parity.py | 254 +++++++ .../diffusion/models/auk/test_pipeline_auk.py | 514 +++++++++++++ tests/e2e/offline_inference/test_auk.py | 121 ++++ .../stage_input_processors/test_auk.py | 188 +++++ tests/model_extras/test_auk_requests.py | 157 ++++ tests/model_tests/diffusion/test_alignment.py | 4 + tools/auk_parity_reference.py | 227 ++++++ tools/prepare_auk_checkpoint.py | 369 ++++++++++ vllm_omni/config/pipeline_registry.py | 2 + vllm_omni/deploy/auk.yaml | 51 ++ vllm_omni/diffusion/models/auk/__init__.py | 16 + .../diffusion/models/auk/auk_transformer.py | 677 ++++++++++++++++++ vllm_omni/diffusion/models/auk/auk_vae.py | 505 +++++++++++++ .../diffusion/models/auk/pipeline_auk.py | 486 +++++++++++++ vllm_omni/diffusion/registry.py | 6 + .../model_executor/models/auk/__init__.py | 5 + vllm_omni/model_executor/models/auk/auk.py | 304 ++++++++ .../model_executor/models/auk/pipeline.py | 60 ++ vllm_omni/model_executor/models/registry.py | 5 + .../stage_input_processors/auk.py | 148 ++++ vllm_omni/model_extras/auk.py | 145 ++++ .../transformers_utils/configs/__init__.py | 3 + vllm_omni/transformers_utils/configs/auk.py | 130 ++++ 28 files changed, 4986 insertions(+) create mode 100644 recipes/Tencent/AuK-H100.md create mode 100644 tests/diffusion/models/auk/test_auk_transformer_parity.py create mode 100644 tests/diffusion/models/auk/test_auk_vae_parity.py create mode 100644 tests/diffusion/models/auk/test_pipeline_auk.py create mode 100644 tests/e2e/offline_inference/test_auk.py create mode 100644 tests/model_executor/stage_input_processors/test_auk.py create mode 100644 tests/model_extras/test_auk_requests.py create mode 100644 tools/auk_parity_reference.py create mode 100755 tools/prepare_auk_checkpoint.py create mode 100644 vllm_omni/deploy/auk.yaml create mode 100644 vllm_omni/diffusion/models/auk/__init__.py create mode 100644 vllm_omni/diffusion/models/auk/auk_transformer.py create mode 100644 vllm_omni/diffusion/models/auk/auk_vae.py create mode 100644 vllm_omni/diffusion/models/auk/pipeline_auk.py create mode 100644 vllm_omni/model_executor/models/auk/__init__.py create mode 100644 vllm_omni/model_executor/models/auk/auk.py create mode 100644 vllm_omni/model_executor/models/auk/pipeline.py create mode 100644 vllm_omni/model_executor/stage_input_processors/auk.py create mode 100644 vllm_omni/model_extras/auk.py create mode 100644 vllm_omni/transformers_utils/configs/auk.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 6b0952eac..1f4093828 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -68,6 +68,7 @@ th { | `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — | | `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — | +| `AuKForConditionalGeneration` | AuK / AuK-Flash (instruction-driven speech generation and editing) | `tencent/AuK`, `tencent/AuK-Flash` (with `Qwen/Qwen2.5-Omni-3B`; assemble with `tools/prepare_auk_checkpoint.py`) | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/AuK-H100.md) | | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — | | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — | | `DotsTTSForConditionalGeneration` | dots.tts | `dots-studio/dots.tts-soar` | ✅︎ | | | | — | diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index a8a273026..89ff2dbce 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -14,6 +14,7 @@ list of supported architectures across all modalities, see | Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate | |---|---|---|---|---|---|---| +| AuK | `tencent/AuK`, `tencent/AuK-Flash` (assembled with `tools/prepare_auk_checkpoint.py`) | 2 (thinker encoder + DiT) | ✓ | — | instruction-driven editing, enhancement, separation; no example script (use `vllm_omni.model_extras.auk`) | 24 kHz | | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | dual-AR | ✓ | ✓ | 11 languages | 44.1 kHz | | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz | diff --git a/recipes/README.md b/recipes/README.md index ea4576201..d9f4b6822 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -84,6 +84,7 @@ recipes/ | [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X | | [`SenseNova/SenseNova-U1.5.md`](./SenseNova/SenseNova-U1.5.md) | Unified image generation and understanding, with an 8-step distilled LoRA | 1x A800 80GB | | [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X | +| [`Tencent/AuK-H100.md`](./Tencent/AuK-H100.md) | Offline instruction-driven TTS, speech editing, enhancement and separation (AuK base and Flash) | 1x H100 80GB | | [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB | | [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 | | [`Wan-AI/Wan2.2-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 1x datacenter Blackwell (B200/B300) | diff --git a/recipes/Tencent/AuK-H100.md b/recipes/Tencent/AuK-H100.md new file mode 100644 index 000000000..5acd01cb0 --- /dev/null +++ b/recipes/Tencent/AuK-H100.md @@ -0,0 +1,167 @@ +# AuK speech generation and editing — H100 + +> Instruction-driven zero-shot TTS, instruct TTS, speech editing, enhancement +> and separation with Tencent's AuK, served as a two-stage vLLM-Omni pipeline +> (Qwen2.5-Omni thinker encoder, then a rectified-flow DiT with a BigVGAN-flow +> VAE) on one H100. + +## Summary + +- Vendor: Tencent Hunyuan +- Model: `tencent/AuK` (32-step base) and `tencent/AuK-Flash` (distilled 4-step), + both with `Qwen/Qwen2.5-Omni-3B` as the frozen encoder +- Task: zero-shot TTS, instruct TTS, content and acoustic editing, + paralinguistic editing, speech enhancement, speaker and music separation +- Mode: offline inference (online chat serving is not yet qualified) +- Hardware: 1x H100 80GB +- Maintainer: Community + +## When to use this recipe + +Use this recipe to run AuK through vLLM-Omni for any of its instruction-driven +tasks with a single shared request shape: an instruction, an optional source or +reference clip, and a target duration. The task is chosen by the instruction +text, exactly as in the upstream repository's cookbook. + +## Supported model contract + +| Task family | Instruction | Audio input | Duration | +| --- | --- | --- | --- | +| Zero-shot TTS | "Say the following with the same voice: '...'" | reference clip | `gen_seconds` required | +| Instruct TTS | `Generate speech based on the following description: "{voice description}". The content to speak is: "{text}".` (upstream template; other orderings can make the model speak the description) | none | `gen_seconds` required | +| Content / lyric editing | "Replace '...' with '...'" and similar | source clip | defaults to the source length | +| Acoustic / paralinguistic editing | pitch, speed, volume, emotion, timbre, de-accent, nonverbal, whisper | source clip | defaults to the source length (speed edits scale it) | +| Enhancement / separation | "Keep pure speech voice, remove noise and reverberation." and similar | source clip | defaults to the source length | + +| Property | Value | +| --- | --- | +| Input audio | mono or stereo (averaged), any sample rate (resampled to 24 kHz for the VAE and to 16 kHz for the encoder), one clip per request | +| Output | mono float32 waveform at 24 kHz, `ceil(gen_seconds * 50)` latent frames | +| Sequence limit | reference plus target latents up to 65536 frames (about 21 minutes) | +| Sampling knobs | `num_inference_steps` (default 32), `guidance_scale` (default 2.0), `seed`; Flash pins 4 steps, CFG 0 and its own time grid | +| Request knobs | `gen_seconds`, `sway` (default -1.0), `t_grid`, `vae_sample` via `additional_information["auk"]` | +| Concurrency | one request per DiT forward; the encoder stage runs with `max_num_seqs: 1`, prefix caching and chunked prefill off | +| Streaming | none: the ODE runs over the whole target | + +## References + +- Upstream repository and cookbook: +- Technical report: +- Weights: , , + +- Pipeline topology: `vllm_omni/model_executor/models/auk/pipeline.py`; + deploy defaults: `vllm_omni/deploy/auk.yaml` +- Request helpers: `vllm_omni/model_extras/auk.py` +- Test: `tests/e2e/offline_inference/test_auk.py` + +## Hardware + +- Accelerator model and per-device memory: NVIDIA H100 80GB +- Number of devices: 1 (both stages share the device) +- Device interconnect: not applicable +- Host memory: 64 GB is ample; the assembled directory is symlinks only +- Qualification scope: offline `Omni` inference for the four cookbook tasks + below, base and Flash checkpoints, seed-reproducible, four concurrent + requests bit-identical to sequential runs + +## Software environment + +- OS: Linux +- Python: 3.12 +- Driver / runtime: CUDA 13 driver with the vLLM 0.29.0 wheel (torch 2.13) +- vLLM version: 0.29.0 +- vLLM-Omni version or commit: branch `auk-support` on top of `aff7d6494` + +## Command + +Assemble the checkpoint directory once (symlinks by default): + +```bash +hf download tencent/AuK --local-dir ckpts/AuK +hf download Qwen/Qwen2.5-Omni-3B --local-dir ckpts/Qwen2.5-Omni-3B +python tools/prepare_auk_checkpoint.py --auk-dir ckpts/AuK --qwen-dir ckpts/Qwen2.5-Omni-3B --out ckpts/auk-omni +# For AuK-Flash: --auk-dir ckpts/AuK-Flash --out ckpts/auk-omni-flash +``` + +Offline inference (any task; the instruction selects it): + +```python +import soundfile as sf +from vllm.multimodal.media.audio import load_audio + +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.model_extras.auk import auk_prompt, auk_sampling_params + +omni = Omni(model="ckpts/auk-omni") # loads vllm_omni/deploy/auk.yaml +wav, sr = load_audio("ref.wav", sr=None) +prompt = auk_prompt( + "Say the following with the same voice: 'Ladies and gentlemen, welcome.'", + (wav, sr), + gen_seconds=4.0, +) +out = omni.generate(prompt, auk_sampling_params(nfe=32, cfg=2.0, seed=0))[0] +sf.write("out.wav", out.multimodal_output["audio"], out.multimodal_output["audio_sample_rate"]) +``` + +Text-only instruct TTS passes `None` for the audio and must set `gen_seconds`. +Editing and enhancement requests pass the source clip and may omit +`gen_seconds` to keep the source length. + +## Verification + +```bash +VLLM_OMNI_AUK_MODEL_DIR=ckpts/auk-omni python -m pytest tests/e2e/offline_inference/test_auk.py -q +``` + +The parity suites compare the port against the upstream implementation and +need reference artifacts produced by the upstream package (install it with +`pip install -e /path/to/AuK` in its own environment; it pins its own torch): + +```bash +python tools/auk_parity_reference.py --auk-repo /path/to/AuK --ckpt-dir ckpts/AuK \ + --qwen-dir ckpts/Qwen2.5-Omni-3B --out auk-parity +AUK_CKPT_DIR=ckpts/AuK AUK_REF_WAV=/path/to/AuK/assets/demo-input-audio/zero-shot-tts/ref.wav \ + python -m pytest tests/diffusion/models/auk/test_auk_transformer_parity.py tests/diffusion/models/auk/test_auk_vae_parity.py -q +AUK_OMNI_CKPT_DIR=ckpts/auk-omni AUK_PARITY_REF=auk-parity/parity_ref/base \ + python -m pytest tests/diffusion/models/auk/test_pipeline_auk.py -q +``` + +Measured on one H100 against the upstream implementation with the same seeds +and mean VAE latents (four cookbook cases, base checkpoint): identical +transcripts, speaker similarity to the upstream output 0.993 to 0.999, log-mel +L1 0.11 to 0.47 (the upstream VAE-resampling noise floor is 0.39). Flash: +0.989 to 0.999 and 0.24 to 0.51. Wall per request on the second call: base +1.0 to 2.0 s for 3.5 to 10.9 s of audio, Flash 0.17 to 0.30 s; engine start +about 47 s with a warm page cache. + +## Notes + +- Memory usage: about 25 GB peak on the device for a 12 s generation with + both stages resident (deploy defaults: encoder 0.45, diffusion stage 0.35 + of device memory). +- Key flags: `enforce_eager` on both stages (the encoder walks the decoder + layers itself for the layer fusion). `enable_prefix_caching` must stay off + for the encoder: a cache hit skips prompt positions that the fused + condition needs. `enable_chunked_prefill` is off by default: forcing it + (128-token chunks, so two to three chunks per prompt) reproduces the + unchunked condition to bf16 rounding (per-token cosine 0.99999) and leaves + request walls unchanged, because the encoder is about 3% of a base request, + while concurrent runs stop being bit-identical to sequential ones. +- Known limitations: the encoder's audio tower runs without `flash_attn` in a + plain vLLM install, which shifts the encoder output on audio token positions + (per-token cosine 0.96 vs the upstream fp32 fusion; text positions 0.9999) + with waveform-level effects at the VAE noise floor; installing `flash-attn` + removes the gap. No streaming. Digits are read the way upstream reads them + (spell numbers out in the text). Instruct TTS without a reference can leak + the description into the speech, as upstream does. + +## Supported features + +| Feature | Status | Guide | +| --- | --- | --- | +| Offline `Omni.generate` | supported (base and Flash) | `docs/getting_started/quickstart.md` | +| Online `/v1/chat/completions` with audio | not yet qualified | `docs/serving/` | +| `/v1/audio/speech` | not supported (needs a TTS adapter) | `docs/contributing/model/adding_tts_model.md` | +| Streaming / async chunk | not supported | `docs/design/feature/async_chunk.md` | +| Batching across requests | one request per DiT forward | `docs/user_guide/diffusion/` | +| Tensor / sequence parallelism | not supported | `docs/configuration/composable_parallel.md` | diff --git a/tests/diffusion/models/auk/test_auk_transformer_parity.py b/tests/diffusion/models/auk/test_auk_transformer_parity.py new file mode 100644 index 000000000..937369d05 --- /dev/null +++ b/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) diff --git a/tests/diffusion/models/auk/test_auk_vae_parity.py b/tests/diffusion/models/auk/test_auk_vae_parity.py new file mode 100644 index 000000000..cafc99614 --- /dev/null +++ b/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 diff --git a/tests/diffusion/models/auk/test_pipeline_auk.py b/tests/diffusion/models/auk/test_pipeline_auk.py new file mode 100644 index 000000000..261c5240d --- /dev/null +++ b/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/`` 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 diff --git a/tests/e2e/offline_inference/test_auk.py b/tests/e2e/offline_inference/test_auk.py new file mode 100644 index 000000000..746af7bd2 --- /dev/null +++ b/tests/e2e/offline_inference/test_auk.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Offline E2E smoke test for AuK (encoder stage -> diffusion stage). + +AuK ships as ``config.yaml`` plus loose safetensors and needs the Qwen2.5-Omni-3B +snapshot next to it, so the test expects an assembled directory produced by +``tools/prepare_auk_checkpoint.py`` and skips when it is not available: + + VLLM_OMNI_AUK_MODEL_DIR=/path/to/auk-omni python -m pytest tests/e2e/offline_inference/test_auk.py + +Two requests are exercised: zero-shot TTS from a reference clip and text-only +instruct TTS. The checks are structural (duration, sample rate, level, +finiteness); fidelity against the upstream implementation is covered by the +parity suites under ``tests/diffusion/models/auk``. +""" + +from __future__ import annotations + +import functools +import math +import os +from pathlib import Path + +import numpy as np +import pytest +import soundfile as sf + +from tests.helpers.mark import hardware_test +from tests.helpers.media import get_asset_path +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.model_extras.auk import HOP, SAMPLE_RATE, auk_prompt, auk_sampling_params + +MODEL_DIR_ENV = "VLLM_OMNI_AUK_MODEL_DIR" + +# Any clean speech clip works as the cloning reference; reuse the vendored +# CosyVoice3 prompt so no new binary asset is needed. +REFERENCE_WAV_PATH = get_asset_path("cosyvoice3/zero_shot_prompt.wav") +ZERO_SHOT_TEXT = "The quick brown fox jumps over the lazy dog." +ZERO_SHOT_SECONDS = 3.0 +# The upstream cookbook template: description first, then the content. +INSTRUCT_TEXT = ( + 'Generate speech based on the following description: "A calm young woman speaking warmly and slowly.". ' + 'The content to speak is: "Welcome back, how was your day?".' +) +INSTRUCT_SECONDS = 2.5 +# 8 Euler steps keep the test short; the base recipe is 32. +TEST_NFE = 8 + +_model_dir = os.environ.get(MODEL_DIR_ENV) +pytestmark = [ + pytest.mark.slow, + pytest.mark.tts, + pytest.mark.skipif(not _model_dir, reason=f"set {MODEL_DIR_ENV} to an assembled AuK directory"), + pytest.mark.parametrize( + "omni_runner", + [pytest.param((str(Path(_model_dir or ".").resolve()), get_deploy_config_path("auk.yaml"), {}), id="auk")], + indirect=True, + ), +] + + +@functools.lru_cache(maxsize=1) +def _load_reference_wav() -> tuple[np.ndarray, int]: + audio, sr = sf.read(str(REFERENCE_WAV_PATH), dtype="float32", always_2d=False) + if isinstance(audio, np.ndarray) and audio.ndim > 1: + audio = np.mean(audio, axis=-1) + return np.asarray(audio, dtype=np.float32), int(sr) + + +def _audio_from_output(output) -> tuple[np.ndarray, int]: + mm = getattr(output, "multimodal_output", None) or {} + assert "audio" in mm, f"no audio in the final output; keys={sorted(mm)}" + audio = mm["audio"] + if isinstance(audio, list): + audio = np.concatenate([np.asarray(a, dtype=np.float32).reshape(-1) for a in audio]) + if hasattr(audio, "detach"): + audio = audio.detach().cpu().float().numpy() + audio = np.asarray(audio, dtype=np.float32).reshape(-1) + sr = mm.get("audio_sample_rate", SAMPLE_RATE) + if isinstance(sr, list) and sr: + sr = sr[-1] + if hasattr(sr, "item"): + sr = sr.item() + return audio, int(sr) + + +def _check_waveform(audio: np.ndarray, sr: int, *, gen_seconds: float) -> None: + assert sr == SAMPLE_RATE + expected = math.ceil(gen_seconds * SAMPLE_RATE / HOP) * HOP + assert audio.size == expected, f"expected {expected} samples for {gen_seconds}s, got {audio.size}" + assert np.isfinite(audio).all(), "non-finite samples" + assert np.abs(audio).max() <= 1.0 + 1e-6, "audio exceeds full scale" + assert float(np.sqrt(np.mean(audio**2))) > 1e-3, "audio is silent" + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_auk_zero_shot_tts(omni_runner: OmniRunner) -> None: + """A reference clip plus an instruction yields speech of the requested length.""" + wav, sr = _load_reference_wav() + prompt = auk_prompt( + f"Say the following with the same voice: '{ZERO_SHOT_TEXT}'", + (wav, sr), + gen_seconds=ZERO_SHOT_SECONDS, + ) + outputs = omni_runner.omni.generate(prompt, auk_sampling_params(nfe=TEST_NFE, cfg=2.0, seed=0)) + assert outputs, "no outputs returned" + audio, out_sr = _audio_from_output(outputs[0]) + _check_waveform(audio, out_sr, gen_seconds=ZERO_SHOT_SECONDS) + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_auk_instruct_tts_text_only(omni_runner: OmniRunner) -> None: + """A text-only request (no reference audio) takes the no-prompt-audio path.""" + prompt = auk_prompt(INSTRUCT_TEXT, None, gen_seconds=INSTRUCT_SECONDS) + assert prompt["prompt"].rstrip().endswith("<|im_start|>assistant\n".rstrip()) + assert "||" in prompt["prompt"] + outputs = omni_runner.omni.generate(prompt, auk_sampling_params(nfe=TEST_NFE, cfg=2.0, seed=0)) + assert outputs, "no outputs returned" + audio, out_sr = _audio_from_output(outputs[0]) + _check_waveform(audio, out_sr, gen_seconds=INSTRUCT_SECONDS) diff --git a/tests/model_executor/stage_input_processors/test_auk.py b/tests/model_executor/stage_input_processors/test_auk.py new file mode 100644 index 000000000..b2904e1f7 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_auk.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for the AuK encoder -> diffusion stage input processor.""" + +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.auk import encoder2dit + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +HIDDEN = 2048 + + +def _request_output( + mm_output: Any, + completion_mm_output: Any = None, + request_id: str | None = None, +) -> SimpleNamespace: + """Build a minimal encoder-stage output mock.""" + return SimpleNamespace( + outputs=[SimpleNamespace(token_ids=[0], multimodal_output=completion_mm_output)], + multimodal_output=mm_output, + request_id=request_id, + ) + + +def _payload(text_cond: Any) -> dict[str, Any]: + """The slot the encoder actually emits: hidden_states.output.""" + return {"hidden_states": {"output": text_cond}} + + +def _prompt( + *, + audio: bool = False, + gen_seconds: float | None = None, + has_audio: bool | None = None, +) -> dict[str, Any]: + knobs: dict[str, Any] = { + "gen_seconds": gen_seconds, + "sway": -1.0, + "t_grid": None, + "vae_sample": False, + } + if has_audio is not None: + knobs["has_audio"] = has_audio + prompt: dict[str, Any] = { + "prompt": "<|im_start|>system\n", + "additional_information": {"auk": knobs}, + } + if audio: + prompt["multi_modal_data"] = {"audio": (np.zeros(24000, dtype=np.float32), 24000)} + return prompt + + +class TestTextCondExtraction: + def test_tensor_output(self): + cond = torch.zeros(7, HIDDEN, dtype=torch.bfloat16) + out = encoder2dit([_request_output(_payload(cond))], _prompt()) + assert out is not None + assert out["prompt"] == "" + assert out["prompt_embeds"].shape == (7, HIDDEN) + assert out["prompt_embeds"].dtype == torch.bfloat16 + assert not out["prompt_embeds"].is_cuda + + def test_list_of_tensors_is_concatenated(self): + parts = [torch.zeros(3, HIDDEN), torch.ones(4, HIDDEN)] + out = encoder2dit([_request_output(_payload(parts))], _prompt()) + assert out["prompt_embeds"].shape == (7, HIDDEN) + assert torch.equal(out["prompt_embeds"][3:], torch.ones(4, HIDDEN)) + + def test_single_element_list(self): + out = encoder2dit([_request_output(_payload([torch.zeros(5, HIDDEN)]))], _prompt()) + assert out["prompt_embeds"].shape == (5, HIDDEN) + + def test_numpy_output(self): + cond = np.zeros((4, HIDDEN), dtype=np.float32) + out = encoder2dit([_request_output(_payload(cond))], _prompt()) + assert isinstance(out["prompt_embeds"], torch.Tensor) + assert out["prompt_embeds"].shape == (4, HIDDEN) + + def test_completion_output_fallback(self): + cond = torch.zeros(6, HIDDEN) + out = encoder2dit( + [_request_output(None, completion_mm_output=_payload(cond))], + _prompt(), + ) + assert out["prompt_embeds"].shape == (6, HIDDEN) + + def test_hidden_states_as_a_struct(self): + """A deserialized payload arrives as a struct, not a mapping.""" + cond = torch.zeros(9, HIDDEN) + mm_output = {"hidden_states": SimpleNamespace(output=cond, layers=None)} + out = encoder2dit([_request_output(mm_output)], _prompt()) + assert out["prompt_embeds"].shape == (9, HIDDEN) + + def test_flattened_payload_key(self): + cond = torch.zeros(8, HIDDEN) + out = encoder2dit([_request_output({"hidden_states.output": cond})], _prompt()) + assert out["prompt_embeds"].shape == (8, HIDDEN) + + def test_direct_text_cond_key(self): + cond = torch.zeros(2, HIDDEN) + out = encoder2dit([_request_output({"text_cond": cond})], _prompt()) + assert out["prompt_embeds"].shape == (2, HIDDEN) + + def test_no_source_outputs_returns_none(self): + assert encoder2dit([], _prompt()) is None + + def test_missing_payload_raises(self): + with pytest.raises(ValueError, match="hidden_states.output"): + encoder2dit([_request_output({})], _prompt()) + + def test_empty_hidden_states_slot_raises(self): + with pytest.raises(ValueError, match="hidden_states.output"): + encoder2dit([_request_output({"hidden_states": {}})], _prompt()) + + def test_empty_multimodal_output_raises(self): + with pytest.raises(ValueError, match="hidden_states.output"): + encoder2dit([_request_output(None)], _prompt()) + + def test_empty_list_raises(self): + with pytest.raises(ValueError, match="empty text condition"): + encoder2dit([_request_output(_payload([]))], _prompt()) + + def test_wrong_rank_raises(self): + with pytest.raises(ValueError, match=r"\[tokens, hidden\]"): + encoder2dit([_request_output(_payload(torch.zeros(2, 7, HIDDEN)))], _prompt()) + + +class TestKnobPropagation: + def test_knobs_are_forwarded(self): + prompt = _prompt(audio=True, gen_seconds=6.0) + prompt["additional_information"]["auk"]["t_grid"] = [0.0, 0.5, 1.0] + prompt["additional_information"]["auk"]["vae_sample"] = True + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], prompt) + knobs = out["additional_information"]["auk"] + assert knobs["gen_seconds"] == 6.0 + assert knobs["sway"] == -1.0 + assert knobs["t_grid"] == [0.0, 0.5, 1.0] + assert knobs["vae_sample"] is True + + def test_sampling_knobs_are_not_copied(self): + """nfe/cfg/seed ride the stage-1 sampling params, not the prompt.""" + prompt = _prompt() + prompt["additional_information"]["auk"].update({"nfe": 8, "cfg": 1.0, "seed": 3}) + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], prompt) + knobs = out["additional_information"]["auk"] + assert set(knobs) == {"gen_seconds", "sway", "t_grid", "vae_sample", "has_audio"} + + def test_multi_modal_data_is_forwarded(self): + prompt = _prompt(audio=True) + out = encoder2dit( + [_request_output(_payload(torch.zeros(3, HIDDEN)))], + prompt, + requires_multimodal_data=True, + ) + assert out["multi_modal_data"] is prompt["multi_modal_data"] + + def test_text_only_request_has_no_multi_modal_data(self): + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], _prompt()) + assert out["multi_modal_data"] is None + assert out["additional_information"]["auk"]["has_audio"] is False + + def test_has_audio_declared_by_the_prompt_wins(self): + prompt = _prompt(audio=True, has_audio=True) + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], prompt) + assert out["additional_information"]["auk"]["has_audio"] is True + + def test_has_audio_inferred_from_multi_modal_data(self): + prompt = _prompt(audio=True) + prompt["additional_information"]["auk"].pop("has_audio", None) + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], prompt) + assert out["additional_information"]["auk"]["has_audio"] is True + + def test_prompt_as_list(self): + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], [_prompt(gen_seconds=2.0)]) + assert out["additional_information"]["auk"]["gen_seconds"] == 2.0 + + def test_missing_prompt_yields_empty_knobs(self): + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], None) + knobs = out["additional_information"]["auk"] + assert knobs["gen_seconds"] is None + assert knobs["has_audio"] is False diff --git a/tests/model_extras/test_auk_requests.py b/tests/model_extras/test_auk_requests.py new file mode 100644 index 000000000..1f319d092 --- /dev/null +++ b/tests/model_extras/test_auk_requests.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for the AuK request helpers.""" + +from __future__ import annotations + +import numpy as np +import pytest +from vllm import SamplingParams + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_extras.auk import ( + NO_PROMPT_AUDIO, + auk_prompt, + auk_sampling_params, + resolve_gen_frames, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +INSTRUCTION = "Say the following with the same voice: 'Ladies and gentlemen'" + +EXPECTED_WITH_AUDIO = ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + f"{INSTRUCTION}<|audio_bos|><|AUDIO|><|audio_eos|><|im_end|>\n" + "<|im_start|>assistant\n" +) + +EXPECTED_TEXT_ONLY = ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + f"{INSTRUCTION}{NO_PROMPT_AUDIO}<|im_end|>\n" + "<|im_start|>assistant\n" +) + + +def _audio() -> tuple[np.ndarray, int]: + return np.zeros(24000, dtype=np.float32), 24000 + + +class TestPromptRendering: + def test_with_audio(self): + audio = _audio() + prompt = auk_prompt(INSTRUCTION, audio, gen_seconds=6.0) + assert prompt["prompt"] == EXPECTED_WITH_AUDIO + assert prompt["multi_modal_data"] == {"audio": audio} + assert prompt["additional_information"]["auk"]["has_audio"] is True + + def test_text_only_appends_the_marker(self): + prompt = auk_prompt(INSTRUCTION, gen_seconds=6.0) + assert prompt["prompt"] == EXPECTED_TEXT_ONLY + assert "multi_modal_data" not in prompt + assert prompt["additional_information"]["auk"]["has_audio"] is False + + def test_marker_is_not_appended_twice(self): + once = auk_prompt(INSTRUCTION, gen_seconds=6.0) + twice = auk_prompt(INSTRUCTION + NO_PROMPT_AUDIO, gen_seconds=6.0) + assert once["prompt"] == twice["prompt"] + assert twice["prompt"].count(NO_PROMPT_AUDIO) == 1 + + def test_marker_is_not_added_when_audio_is_present(self): + prompt = auk_prompt(INSTRUCTION, _audio()) + assert NO_PROMPT_AUDIO not in prompt["prompt"] + + def test_knob_defaults(self): + knobs = auk_prompt(INSTRUCTION, _audio())["additional_information"]["auk"] + assert knobs == { + "gen_seconds": None, + "sway": -1.0, + "t_grid": None, + "vae_sample": False, + "has_audio": True, + } + + def test_knobs_are_carried(self): + grid = [0.0, 0.25, 1.0] + knobs = auk_prompt( + INSTRUCTION, + _audio(), + gen_seconds=3.5, + sway=0.0, + t_grid=grid, + vae_sample=True, + )["additional_information"]["auk"] + assert knobs["gen_seconds"] == 3.5 + assert knobs["sway"] == 0.0 + assert knobs["vae_sample"] is True + assert knobs["t_grid"] == grid + # A copy, so a caller mutating its list cannot reach the request. + assert knobs["t_grid"] is not grid + + def test_sampling_knobs_are_absent_from_the_prompt(self): + knobs = auk_prompt(INSTRUCTION, _audio())["additional_information"]["auk"] + assert not {"nfe", "cfg", "seed"} & set(knobs) + + +class TestSamplingParams: + def test_two_stages_in_order(self): + params = auk_sampling_params() + assert len(params) == 2 + assert isinstance(params[0], SamplingParams) + assert isinstance(params[1], OmniDiffusionSamplingParams) + + def test_encoder_stage_is_prefill_only(self): + encoder = auk_sampling_params()[0] + assert encoder.max_tokens == 1 + assert encoder.temperature == 0.0 + assert encoder.detokenize is False + + def test_defaults_are_the_base_recipe(self): + diffusion = auk_sampling_params()[1] + assert diffusion.num_inference_steps == 32 + assert diffusion.guidance_scale == 2.0 + assert diffusion.seed == 0 + + def test_overrides(self): + encoder, diffusion = auk_sampling_params(nfe=4, cfg=0.0, seed=None) + assert diffusion.num_inference_steps == 4 + assert diffusion.guidance_scale == 0.0 + assert diffusion.seed is None + assert encoder.seed is None + + +class TestDurationPolicy: + @pytest.mark.parametrize( + ("gen_seconds", "expected"), + [ + (6.0, 300), + (0.02, 1), + (0.001, 1), + (5.59, 280), + (1.0, 50), + ], + ) + def test_explicit_duration_rounds_up(self, gen_seconds: float, expected: int): + assert resolve_gen_frames(gen_seconds, 0) == expected + + def test_explicit_duration_wins_over_the_source(self): + assert resolve_gen_frames(2.0, 500) == 100 + + def test_source_length_is_the_default(self): + assert resolve_gen_frames(None, 137) == 137 + + def test_text_only_without_a_duration_raises(self): + with pytest.raises(ValueError, match="gen_seconds is required"): + resolve_gen_frames(None, 0) + + @pytest.mark.parametrize("gen_seconds", [0.0, -1.0]) + def test_non_positive_duration_raises(self, gen_seconds: float): + with pytest.raises(ValueError, match="must be positive"): + resolve_gen_frames(gen_seconds, 100) + + def test_custom_rate(self): + assert resolve_gen_frames(1.0, 0, sample_rate=16000, hop=320) == 50 diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 183ef8659..f0a1285bc 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -83,6 +83,10 @@ EXCLUDED_MODELS = [ "Gr00tN1d7Pipeline", "Pi0Pipeline", "SanaWmPipeline", + # Audio pipeline conditioned by a separate Qwen2.5-Omni encoder stage; + # covered by tests/diffusion/models/auk (CPU pipeline tests, parity) and + # the real-checkpoint E2E test instead of a tiny model. + "AuKPipeline", ] diff --git a/tools/auk_parity_reference.py b/tools/auk_parity_reference.py new file mode 100644 index 000000000..33b664232 --- /dev/null +++ b/tools/auk_parity_reference.py @@ -0,0 +1,227 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Generate the upstream AuK reference artifacts consumed by the parity tests. + +This runs the *upstream* implementation, so it must be executed in an +environment where the upstream package is installed +(``pip install -e /path/to/AuK``; it pins its own torch/transformers line and +does not need vLLM-Omni). It writes: + + /parity_ref//.pt, .wav, meta.json + Upstream generations for four cookbook cases (zero-shot TTS, instruct + TTS, content edit, denoise) with the reference latents, the generated + latents and the waveform. Point ``AUK_PARITY_REF`` at + ``/parity_ref/`` for tests/diffusion/models/auk/test_pipeline_auk.py. + /fusion_ref/zs_tts_audio.pt, instruct_text_only.pt + The fp32 layer-fused text condition from the HF thinker for the + zero-shot and text-only prompts (``AUK_FUSION_REF``). + +By default the VAE encode uses the posterior mean, matching the port's default +and making the artifacts deterministic; ``--vae-sample`` reproduces the +upstream stochastic draw (seeded through ``--global-seed``). The ODE noise is +drawn from ``--seed`` on both sides. + +Usage: + python tools/auk_parity_reference.py \\ + --auk-repo /path/to/AuK --ckpt-dir /path/to/AuK \\ + --qwen-dir /path/to/Qwen2.5-Omni-3B --out /path/to/auk-parity \\ + [--variant base|flash] [--vae-sample] [--fusion-device cpu|cuda] [--skip-fusion] +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as F + +WEIGHTS = {"base": "auk_base.safetensors", "flash": "auk_flash.safetensors"} + +# Cookbook cases; audio paths are relative to /assets/demo-input-audio. +CASES: dict[str, dict[str, Any]] = { + "zs_tts_en": { + "text": ( + "Say the following with the same voice: 'Ladies and gentlemen, it's an honor to have the " + "opportunity to address such a distinguished audience'" + ), + "audio": "zero-shot-tts/ref.wav", + "gen_seconds": 6.0, + }, + "instruct_en": { + "text": ( + 'Generate speech based on the following description: "A calm young woman speaking warmly and slowly.". ' + 'The content to speak is: "Welcome back, how was your day at work?".' + ), + "audio": None, + "gen_seconds": 3.5, + }, + "content_edit": { + "text": "Replace 'but accepting what we cannot have' with 'and living well with dreams unmet'.", + "audio": "content-edit/content.wav", + "gen_seconds": 7.0, + }, + "denoise": { + "text": "Keep pure speech voice, remove noise and reverberation.", + "audio": "vocal-extraction/vocal-1-input.wav", + "gen_seconds": None, + }, +} +# The two prompts the fusion reference covers (they must match the cases above). +FUSION_CASES = {"zs_tts_audio": "zs_tts_en", "instruct_text_only": "instruct_en"} +NO_PROMPT_AUDIO = "||" + + +def _messages(case: dict[str, Any], assets: Path) -> list[dict[str, Any]]: + content: list[dict[str, Any]] = [{"type": "text", "text": case["text"]}] + if case["audio"]: + content.append({"type": "audio", "audio": str(assets / case["audio"])}) + return [{"role": "user", "content": content}] + + +def _mean_encode(vae): + """Posterior-mean version of the upstream ``encoding_and_normalization``.""" + + def encode(sample, sample_lengths=None): + latent_stats = vae.audio_encoder(sample) + if sample_lengths is None: + sample_lengths = torch.LongTensor([sample.size(-1)] * sample.size(0)).to(sample.device) + latent_lens = sample_lengths // vae.hop_size + mean, _log_std = latent_stats.chunk(2, 1) + latents = mean.transpose(1, 2).float() + latents = (latents - vae.global_mean.float()) / torch.sqrt(vae.global_log_std.float()) + return latents, torch.clamp(latent_lens, max=latents.size(1)) + + return encode + + +def generate_references(args: argparse.Namespace) -> None: + from auk.infer.infer_auk import AukInfer, save_audio + + ckpt = Path(args.ckpt_dir) + out = Path(args.out) / "parity_ref" / args.variant + out.mkdir(parents=True, exist_ok=True) + assets = Path(args.auk_repo) / "assets" / "demo-input-audio" + + engine = AukInfer(str(ckpt / "config.yaml"), str(ckpt / WEIGHTS[args.variant]), qwen_path=args.qwen_dir) + encode = engine.vae_model.encoding_and_normalization if args.vae_sample else _mean_encode(engine.vae_model) + recorded: dict[str, torch.Tensor] = {} + + def encode_hook(sample, sample_lengths=None): + latents, lens = encode(sample, sample_lengths) + recorded["ref_latents"] = latents.detach().float().cpu() + recorded["ref_latent_lens"] = lens.detach().cpu() + return latents, lens + + engine.vae_model.encoding_and_normalization = encode_hook + upstream_sample = engine.model.sample + + def sample_hook(*a, **k): + latents, trajectory = upstream_sample(*a, **k) + recorded["generated"] = latents.detach().float().cpu() + return latents, trajectory + + engine.model.sample = sample_hook + + meta: dict[str, Any] = {} + for name, case in CASES.items(): + messages = _messages(case, assets) + recorded.clear() + torch.manual_seed(args.global_seed) + started = time.perf_counter() + audio, sr = engine.generate(messages, gen_seconds=case["gen_seconds"], seed=args.seed) + wall = time.perf_counter() - started + save_audio(audio, sr, str(out / f"{name}.wav")) + torch.save( + { + "messages": messages, + "gen_seconds": case["gen_seconds"], + "seed": args.seed, + "global_seed": args.global_seed, + "vae_sample": bool(args.vae_sample), + "audio": audio.float().cpu(), + "sr": sr, + **recorded, + }, + out / f"{name}.pt", + ) + ref_frames = int(recorded["ref_latent_lens"][0]) if "ref_latent_lens" in recorded else 0 + meta[name] = { + "wall_s": wall, + "out_s": audio.shape[-1] / sr, + "sr": sr, + "ref_frames": ref_frames, + "gen_frames": int(recorded["generated"].shape[1]) - ref_frames, + } + print(name, meta[name], flush=True) + (out / "meta.json").write_text(json.dumps(meta, indent=1) + "\n") + + +def generate_fusion(args: argparse.Namespace) -> None: + from auk.model.cfm_edit import CFMEdit + from safetensors import safe_open + from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniThinkerForConditionalGeneration + + out = Path(args.out) / "fusion_ref" + out.mkdir(parents=True, exist_ok=True) + assets = Path(args.auk_repo) / "assets" / "demo-input-audio" + with safe_open(str(Path(args.ckpt_dir) / WEIGHTS[args.variant]), "pt") as f: + layer_weights = f.get_tensor("layer_weights").float() + layer_scale = f.get_tensor("layer_scale").float() + + processor = Qwen2_5OmniProcessor.from_pretrained(args.qwen_dir) + thinker = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(args.qwen_dir, torch_dtype=torch.float32) + thinker.visual = None + thinker = thinker.to(args.fusion_device).eval() + + for out_name, case_name in FUSION_CASES.items(): + case = dict(CASES[case_name]) + if not case["audio"] and not case["text"].endswith(NO_PROMPT_AUDIO): + case["text"] = case["text"] + NO_PROMPT_AUDIO + cond = CFMEdit.build_cond_inputs([_messages(case, assets)], processor).to(args.fusion_device) + with torch.no_grad(): + hidden_states = thinker(**cond, output_hidden_states=True).hidden_states + width = hidden_states[0].shape[-1] + stacked = torch.stack([F.layer_norm(h.float(), [width]) for h in hidden_states[1:]], 0) + weights = F.softmax(layer_weights.to(stacked.device), 0) + fused = (stacked * weights[:, None, None, None]).sum(0) * layer_scale.to(stacked.device) + torch.save( + { + "input_ids": cond["input_ids"].cpu(), + "attention_mask": cond["attention_mask"].cpu(), + "fused": fused[0].float().cpu(), + "final_normed": hidden_states[-1][0].float().cpu(), + "text": case["text"], + }, + out / f"{out_name}.pt", + ) + print(out_name, {"tokens": int(cond["input_ids"].shape[1]), "fused_norm": float(fused.norm())}, flush=True) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--auk-repo", required=True, help="checkout of the upstream AuK repository (for demo clips)") + parser.add_argument("--ckpt-dir", required=True, help="tencent/AuK or tencent/AuK-Flash snapshot") + parser.add_argument("--qwen-dir", required=True, help="Qwen/Qwen2.5-Omni-3B snapshot") + parser.add_argument("--out", required=True, help="output root; writes parity_ref/ and fusion_ref") + parser.add_argument("--variant", choices=sorted(WEIGHTS), default="base") + parser.add_argument("--seed", type=int, default=0, help="ODE noise seed passed to the upstream sampler") + parser.add_argument("--global-seed", type=int, default=1234, help="global RNG seed set before each generate()") + parser.add_argument("--vae-sample", action="store_true", help="draw the VAE posterior sample as upstream does") + parser.add_argument("--fusion-device", default="cpu", help="device for the fp32 HF thinker forward") + parser.add_argument("--skip-fusion", action="store_true", help="only write the generation references") + args = parser.parse_args() + + generate_references(args) + if not args.skip_fusion: + generate_fusion(args) + print(f"reference artifacts written under {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/prepare_auk_checkpoint.py b/tools/prepare_auk_checkpoint.py new file mode 100755 index 000000000..7f55dab8a --- /dev/null +++ b/tools/prepare_auk_checkpoint.py @@ -0,0 +1,369 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Assemble an AuK checkpoint directory that ``--model `` can load. + +The released AuK repos ship ``config.yaml`` plus loose safetensors and depend +on a separate Qwen2.5-Omni-3B snapshot for the frozen text encoder. This tool +writes the ``config.json`` the vLLM-Omni pipeline expects and links every +artifact both stages need into one flat directory: + + / + config.json the Qwen2.5-Omni-3B config with model_type "auk", + AuK's architecture, and the dit / vae / variant / + flash_t_grid / defaults sections added at the top + level (see transformers_utils/configs/auk.py) + auk.safetensors -> auk_base.safetensors or auk_flash.safetensors + vae.safetensors -> the AuK VAE weights + thinker/ -> the Qwen2.5-Omni-3B snapshot, for reference only + model-0000*.safetensors, model.safetensors.index.json, + tokenizer*, vocab.json, merges.txt, added_tokens.json, + special_tokens_map.json, chat_template.json, + preprocessor_config.json, generation_config.json, spk_dict.pt + -> the same snapshot's entries + +The Qwen entries are linked at the top level rather than left under +``thinker/`` because both stages load from this one directory: the encoder +reads the thinker shards, the tokenizer and the audio processor from here, +and takes the layer-fusion tensors out of ``auk.safetensors``. The weight +index keeps ``auk.safetensors`` and ``vae.safetensors`` out of the thinker's +own weight load, and the DiT and VAE weights are read by model code directly. + +Usage: + python tools/prepare_auk_checkpoint.py \ + --auk-dir /path/to/AuK \ + --qwen-dir /path/to/Qwen2.5-Omni-3B \ + --out /path/to/auk-omni-base + +Re-running against an existing directory is safe: config.json is rewritten +and the links are repointed. Pass ``--copy`` on filesystems without symlinks. +""" + +from __future__ import annotations + +import argparse +import ast +import json +import re +import shutil +import sys +from pathlib import Path +from typing import Any + +# The distilled student's 4-step grid. Duplicated from +# vllm_omni/transformers_utils/configs/auk.py so this tool stays runnable +# without importing the package (and its torch dependency). +FLASH_T_GRID = [ + 0.0, + 0.07612049579620361, + 0.2928932309150696, + 0.6173166036605835, + 1.0, +] + +# Names the assembled directory owns; everything else in it is a Qwen entry. +WEIGHTS_LINK = "auk.safetensors" +VAE_LINK = "vae.safetensors" +THINKER_LINK = "thinker" +CONFIG_NAME = "config.json" + +# Snapshot entries that are not reused: the config is rewritten, and the repo +# paperwork would misdescribe the assembled directory. +_SKIP_QWEN_ENTRIES = frozenset({CONFIG_NAME, "README.md", "LICENSE"}) + +DIT_KEYS = ("dim", "heads", "ff_mult", "text_hidden_dim", "num_layers", "num_single_layers") +VAE_KEYS = ("latent_dim", "downsample_rate", "target_sample_rate", "model_init_kwargs") + +_INTERPOLATION = re.compile(r"^\$\{([A-Za-z0-9_.]+)\}$") + + +def _strip_comment(line: str) -> str: + """Drop a trailing ``#`` comment that is not inside a quoted scalar.""" + out: list[str] = [] + quote: str | None = None + for ch in line: + if quote is not None: + if ch == quote: + quote = None + out.append(ch) + continue + if ch in "\"'": + quote = ch + out.append(ch) + continue + if ch == "#": + break + out.append(ch) + return "".join(out).rstrip() + + +def _parse_scalar(text: str) -> Any: + """Parse one YAML scalar or inline sequence.""" + text = text.strip() + if not text: + return None + lowered = text.lower() + if lowered in ("true", "yes"): + return True + if lowered in ("false", "no"): + return False + if lowered in ("null", "~"): + return None + if text[0] in "[{": + normalized = re.sub(r"\b(True|true|yes)\b", "True", text) + normalized = re.sub(r"\b(False|false|no)\b", "False", normalized) + normalized = re.sub(r"\b(null|~)\b", "None", normalized) + try: + return ast.literal_eval(normalized) + except (SyntaxError, ValueError): + return text + try: + return ast.literal_eval(text) + except (SyntaxError, ValueError): + return text.strip("\"'") + + +def _parse_yaml_minimal(text: str) -> dict[str, Any]: + """Parse the small subset of YAML the AuK configs use. + + Nested mappings by indentation, scalars, and inline sequences. Used only + when PyYAML is unavailable. + """ + root: dict[str, Any] = {} + stack: list[tuple[int, dict[str, Any]]] = [(-1, root)] + for raw in text.splitlines(): + line = _strip_comment(raw) + if not line.strip(): + continue + indent = len(line) - len(line.lstrip(" ")) + body = line.strip() + if ":" not in body: + raise ValueError(f"unsupported YAML line (no key): {raw!r}") + key, _, rest = body.partition(":") + key = key.strip() + while stack and indent <= stack[-1][0]: + stack.pop() + if not stack: + raise ValueError(f"unsupported YAML indentation at: {raw!r}") + parent = stack[-1][1] + if rest.strip(): + parent[key] = _parse_scalar(rest) + else: + child: dict[str, Any] = {} + parent[key] = child + stack.append((indent, child)) + return root + + +def _resolve_interpolations(node: Any, root: dict[str, Any]) -> Any: + """Replace ``${a.b.c}`` scalars with the value at that dotted path.""" + if isinstance(node, dict): + return {k: _resolve_interpolations(v, root) for k, v in node.items()} + if isinstance(node, list): + return [_resolve_interpolations(v, root) for v in node] + if isinstance(node, str): + match = _INTERPOLATION.match(node.strip()) + if match is None: + return node + target: Any = root + for part in match.group(1).split("."): + if not isinstance(target, dict) or part not in target: + raise ValueError(f"cannot resolve interpolation ${{{match.group(1)}}}") + target = target[part] + return target + return node + + +def load_auk_yaml(path: Path) -> dict[str, Any]: + """Load the AuK ``config.yaml`` with interpolations resolved.""" + text = path.read_text() + try: + import yaml + + parsed = yaml.safe_load(text) + except ImportError: + parsed = _parse_yaml_minimal(text) + if not isinstance(parsed, dict) or "model" not in parsed: + raise ValueError(f"{path} has no top-level 'model' section") + return _resolve_interpolations(parsed, parsed) + + +def detect_variant(auk_dir: Path, model_section: dict[str, Any]) -> str: + """Infer base vs flash from the weight filename, then the model name.""" + if (auk_dir / "auk_flash.safetensors").exists(): + return "flash" + if (auk_dir / "auk_base.safetensors").exists(): + return "base" + name = str(model_section.get("name", "")).lower() + return "flash" if "flash" in name else "base" + + +def weights_name(variant: str) -> str: + return "auk_flash.safetensors" if variant == "flash" else "auk_base.safetensors" + + +def build_config( + auk_yaml: dict[str, Any], + qwen_config: dict[str, Any], + variant: str, +) -> dict[str, Any]: + """Assemble the ``config.json`` body for the output directory. + + The Qwen config is kept whole, including ``talker_config`` and + ``token2wav_config``: ``AuKConfig`` subclasses ``Qwen2_5OmniConfig`` and + the encoder stage reuses the in-tree thinker processor, which asserts that + type. AuK's own sections are added at the top level. + """ + model = auk_yaml["model"] + arch = model.get("arch", {}) + vae = model.get("vae", {}) + + missing = [key for key in DIT_KEYS if key not in arch] + if missing: + raise ValueError(f"AuK config.yaml arch section is missing {missing}") + missing = [key for key in VAE_KEYS if key not in vae] + if missing: + raise ValueError(f"AuK config.yaml vae section is missing {missing}") + if not isinstance(qwen_config.get("thinker_config"), dict): + raise ValueError("Qwen snapshot config.json has no 'thinker_config' mapping") + + dit = {key: arch[key] for key in DIT_KEYS} + # dim_head is implicit upstream (dim // heads) but the port takes it + # explicitly, so it is recorded rather than recomputed at load time. + dit["dim_head"] = int(arch["dim"]) // int(arch["heads"]) + # Carried from the released config: the DiT builds an attention mask over + # the padded [audio | text] sequence and parity depends on it. + dit["attn_mask_enabled"] = bool(arch.get("attn_mask_enabled", True)) + + config = dict(qwen_config) + config["model_type"] = "auk" + config["architectures"] = ["AuKForConditionalGeneration"] + config["dit"] = dit + config["vae"] = {key: vae[key] for key in VAE_KEYS} + config["variant"] = variant + config["flash_t_grid"] = FLASH_T_GRID + config["defaults"] = {"nfe": 32, "cfg": 2.0, "sway": -1.0} + return config + + +def _replace_path(target: Path) -> None: + """Remove an existing file, directory, or link at ``target``.""" + if target.is_symlink() or target.is_file(): + target.unlink() + elif target.is_dir(): + shutil.rmtree(target) + + +def link_or_copy(source: Path, target: Path, *, copy: bool) -> None: + """Point ``target`` at ``source``, replacing whatever is already there.""" + if not source.exists(): + raise FileNotFoundError(f"missing source artifact: {source}") + if target.is_symlink() and not copy and target.readlink() == source.resolve(): + return + _replace_path(target) + if copy: + if source.is_dir(): + shutil.copytree(source, target) + else: + shutil.copy2(source, target) + return + target.symlink_to(source.resolve(), target_is_directory=source.is_dir()) + + +MANIFEST_NAME = ".auk-assembled.json" + + +def _reject_overlap(out: Path, *inputs: Path) -> None: + """Refuse an output directory that is, contains, or lives inside an input.""" + out_r = out.resolve() + for src in inputs: + src_r = src.resolve() + if out_r == src_r or out_r in src_r.parents or src_r in out_r.parents: + raise ValueError(f"--out {out} overlaps input directory {src}; choose a separate directory") + + +def assemble(auk_dir: Path, qwen_dir: Path, out: Path, variant: str, *, copy: bool) -> None: + """Link the AuK weights and the Qwen snapshot entries into ``out``.""" + _reject_overlap(out, auk_dir, qwen_dir) + out.mkdir(parents=True, exist_ok=True) + manifest_path = out / MANIFEST_NAME + previous: set[str] = set() + if manifest_path.is_file(): + previous = set(json.loads(manifest_path.read_text()).get("entries", [])) + weights = auk_dir / weights_name(variant) + + entries: dict[str, Path] = { + WEIGHTS_LINK: weights, + VAE_LINK: auk_dir / VAE_LINK, + THINKER_LINK: qwen_dir, + } + # The assembled config replaces the Qwen one; every other entry is reused. + for entry in qwen_dir.iterdir(): + if entry.name.startswith(".") or entry.name in _SKIP_QWEN_ENTRIES: + continue + entries[entry.name] = entry + + for name, source in entries.items(): + link_or_copy(source, out / name, copy=copy) + + # Drop only entries this tool created in an earlier run and no longer needs; + # anything else in the directory belongs to the user. + for name in previous - set(entries) - {CONFIG_NAME}: + stale = out / name + if stale.exists() or stale.is_symlink(): + _replace_path(stale) + manifest_path.write_text(json.dumps({"entries": sorted(entries)}, indent=1) + "\n") + + +def print_tree(out: Path) -> None: + """Print the assembled directory, one entry per line.""" + print(f"\n{out}/") + for entry in sorted(out.iterdir()): + if entry.is_symlink(): + print(f" {entry.name} -> {entry.readlink()}") + elif entry.is_dir(): + print(f" {entry.name}/") + else: + print(f" {entry.name} ({entry.stat().st_size / 1e9:.2f} GB)") + + +def main() -> int: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--auk-dir", required=True, help="HF snapshot of the AuK or AuK-Flash repo") + parser.add_argument("--qwen-dir", required=True, help="HF snapshot of Qwen2.5-Omni-3B") + parser.add_argument("--out", required=True, help="Directory to assemble (created if absent)") + parser.add_argument( + "--variant", + choices=("base", "flash"), + help="Override the variant; detected from the AuK snapshot by default", + ) + parser.add_argument( + "--copy", + action="store_true", + help="Copy the weights and the snapshot entries instead of symlinking", + ) + args = parser.parse_args() + + auk_dir = Path(args.auk_dir).expanduser().resolve() + qwen_dir = Path(args.qwen_dir).expanduser().resolve() + out = Path(args.out).expanduser().resolve() + _reject_overlap(out, auk_dir, qwen_dir) + + auk_yaml = load_auk_yaml(auk_dir / "config.yaml") + variant = args.variant or detect_variant(auk_dir, auk_yaml["model"]) + qwen_config = json.loads((qwen_dir / CONFIG_NAME).read_text()) + + config = build_config(auk_yaml, qwen_config, variant) + assemble(auk_dir, qwen_dir, out, variant, copy=args.copy) + (out / CONFIG_NAME).write_text(json.dumps(config, indent=2) + "\n") + + print(f"variant: {variant} (weights {weights_name(variant)})") + print_tree(out) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index a69fb08c6..2378ce5b9 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -40,6 +40,7 @@ from vllm_omni.model_executor.models.audex.pipeline import ( AUDEX_TTS_PIPELINE, ) from vllm_omni.model_executor.models.audio8_tts.pipeline import AUDIO8_TTS_PIPELINE +from vllm_omni.model_executor.models.auk.pipeline import AUK_PIPELINE from vllm_omni.model_executor.models.aura_omni.pipeline import AURA_OMNI_PIPELINE from vllm_omni.model_executor.models.bagel.pipeline import ( BAGEL_PIPELINE, @@ -164,6 +165,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "wan2_2_ti2v": WAN2_2_TI2V_PIPELINE, "voxcpm2": VOXCPM2_PIPELINE, "dots_tts": DOTS_TTS_PIPELINE, + "auk": AUK_PIPELINE, "cosyvoice3": COSYVOICE3_PIPELINE, "audex_tts": AUDEX_TTS_PIPELINE, "audex_tta": AUDEX_TTA_PIPELINE, diff --git a/vllm_omni/deploy/auk.yaml b/vllm_omni/deploy/auk.yaml new file mode 100644 index 000000000..bbe544919 --- /dev/null +++ b/vllm_omni/deploy/auk.yaml @@ -0,0 +1,51 @@ +# AuK deploy: thinker encoder (stage 0) + rectified-flow DiT (stage 1). +# Topology declared in vllm_omni/model_executor/models/auk/pipeline.py. +# +# Fields omitted from a stage fall back to StageDeployConfig dataclass +# defaults (see vllm_omni/config/stage_config.py). trust_remote_code, +# enable_prefix_caching and enable_chunked_prefill are pipeline-wide +# DeployConfig fields and apply to both stages. +# +# The ODE runs over the whole target segment, so there is nothing to stream +# between the stages: async_chunk stays off. +async_chunk: false +trust_remote_code: true +# Prefix caching must stay off: a cache hit skips the prefix positions, and +# the fused [tokens, 2048] condition needs every position of the prompt. +# Chunked prefill works (the runner concatenates the per-chunk hidden states) +# but measured no wall-time gain, since the encoder is a few percent of a +# request, and it costs bit-exact reproducibility across scheduling. +enable_prefix_caching: false +enable_chunked_prefill: false + +stages: + # Stage 0: Qwen2.5-Omni thinker as a frozen text encoder. + # Prefill only — one sampled token, discarded downstream. eager because the + # fused-hidden forward walks the decoder layers itself and so bypasses the + # compiled graph. + - stage_id: 0 + devices: "0" + max_num_seqs: 1 + gpu_memory_utilization: 0.45 + enforce_eager: true + max_num_batched_tokens: 32768 + # AuK takes text plus at most one clip; no image or video tower is built. + limit_mm_per_prompt: {audio: 1} + default_sampling_params: + temperature: 0.0 + max_tokens: 1 + detokenize: false + seed: 0 + + # Stage 1: AuKPipeline (DiT + BigVGAN-flow VAE encoder and decoder). + # Shares GPU 0 with the encoder; 32 Euler steps at CFG 2.0 is the Base + # recipe, and AuK-Flash overrides both with its distilled 4-step grid. + - stage_id: 1 + devices: "0" + max_num_seqs: 1 + gpu_memory_utilization: 0.35 + enforce_eager: true + default_sampling_params: + seed: 0 + num_inference_steps: 32 + guidance_scale: 2.0 diff --git a/vllm_omni/diffusion/models/auk/__init__.py b/vllm_omni/diffusion/models/auk/__init__.py new file mode 100644 index 000000000..fd987d42b --- /dev/null +++ b/vllm_omni/diffusion/models/auk/__init__.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK diffusion-stage components: the flow transformer, the codec, the pipeline.""" + +from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, dit_state_dict, sample_latents +from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE +from vllm_omni.diffusion.models.auk.pipeline_auk import AuKPipeline, get_auk_post_process_func + +__all__ = [ + "AuKPipeline", + "AuKTransformer", + "AuKVAE", + "dit_state_dict", + "get_auk_post_process_func", + "sample_latents", +] diff --git a/vllm_omni/diffusion/models/auk/auk_transformer.py b/vllm_omni/diffusion/models/auk/auk_transformer.py new file mode 100644 index 000000000..5aca8f39b --- /dev/null +++ b/vllm_omni/diffusion/models/auk/auk_transformer.py @@ -0,0 +1,677 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The AuK audio DiT as a self-contained inference module. + +Two phases. ``num_layers`` double-stream blocks attend jointly over text and +audio while keeping each stream's residual separate, then +``num_single_layers`` single-stream blocks run over the concatenated +``[text | ref audio | target audio]`` sequence. Only the target-audio slice is +projected back to latent space, so the module predicts a flow-matching +velocity for the target frames and :func:`sample_latents` integrates it. + +Parameter names and shapes match the AuK reference backbone, so a released +checkpoint loads with ``load_state_dict(strict=True)`` once the +``transformer.`` prefix is stripped (see :func:`dit_state_dict`). + +Deliberate differences from the reference, all inference-only: + +* ``x_transformers`` and ``torchdiffeq`` are gone. Rotary embeddings are + computed here (interleaved GPT-J pairs, no xpos scaling) and the ODE is an + explicit Euler loop. +* Dropout, activation checkpointing, the ``flash_attn`` backend switch and the + zero-init of the modulation layers are training concerns and are dropped. +* The timestep sinusoid is cast to the time MLP's weight dtype rather than to + the timestep's own dtype, so an fp32 timestep works against half-precision + weights outside autocast. +* ``attn_mask_enabled`` defaults to ``True``, the value the released config + sets, rather than the reference's ``False``. +""" + +import math +from collections.abc import Iterable, Mapping + +import torch +import torch.nn.functional as F +from torch import nn + +__all__ = ["AuKTransformer", "dit_state_dict", "sample_latents"] + + +def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor: + """The single attention call site, so the backend can be swapped in one place. + + ``q``/``k``/``v`` are ``[B, H, N, D]``. ``mask`` is a boolean key-padding + mask ``[B, K]`` in which ``True`` marks a valid key; it is broadcast to + ``[B, H, Q, K]`` only when given. + """ + attn_mask = None + if mask is not None: + attn_mask = mask[:, None, None, :].expand(q.shape[0], q.shape[1], q.shape[-2], k.shape[-2]) + return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + """Rotate interleaved GPT-J pairs: ``(x1, x2) -> (-x2, x1)``.""" + pairs = x.reshape(*x.shape[:-1], -1, 2) + x1, x2 = pairs.unbind(dim=-1) + return torch.stack((-x2, x1), dim=-1).flatten(-2) + + +def _apply_rope(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Apply rotary frequencies ``[1, 1, N, D]`` to ``x`` ``[B, H, N, D]`` in fp32.""" + out_dtype = x.dtype + x32 = x.float() + out = x32 * freqs.cos() + _rotate_half(x32) * freqs.sin() + return out.to(out_dtype) + + +class Rotary(nn.Module): + """Rotary position frequencies for one stream, evaluated in fp32. + + The frequencies live in a buffer so that a released checkpoint's own copy + of them loads with ``strict=True``. They are recomputed whenever the + module has been cast below fp32: rounding either the frequencies or the + integer positions to bf16 wrecks the phase at long positions. + """ + + def __init__(self, dim: int, base: float = 10000.0) -> None: + super().__init__() + if dim % 2 != 0: + raise ValueError(f"rotary dim must be even, got {dim}") + self.dim = dim + self.base = base + self.register_buffer("inv_freq", self._frequencies(torch.device("cpu"))) + + def _frequencies(self, device: torch.device) -> torch.Tensor: + exponents = torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim + return 1.0 / (self.base**exponents) + + def forward(self, seq_len: int) -> torch.Tensor: + """Return frequencies ``[1, 1, seq_len, dim]`` for positions ``0..seq_len-1``.""" + inv_freq = self.inv_freq + if inv_freq.dtype != torch.float32: + inv_freq = self._frequencies(inv_freq.device) + pos = torch.arange(seq_len, device=inv_freq.device, dtype=torch.float32) + freqs = torch.outer(pos, inv_freq) + return torch.stack((freqs, freqs), dim=-1).flatten(-2)[None, None] + + +class TimeEmbedding(nn.Module): + """Sinusoidal timestep features followed by a two-layer MLP.""" + + def __init__(self, dim: int, freq_embed_dim: int = 256) -> None: + super().__init__() + self.freq_embed_dim = freq_embed_dim + self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + + def forward(self, timestep: torch.Tensor, scale: float = 1000.0) -> torch.Tensor: + half = self.freq_embed_dim // 2 + decay = math.log(10000) / (half - 1) + freqs = torch.exp(torch.arange(half, device=timestep.device, dtype=torch.float32) * -decay) + angles = scale * timestep.unsqueeze(1) * freqs.unsqueeze(0) + hidden = torch.cat((angles.sin(), angles.cos()), dim=-1) + return self.time_mlp(hidden.to(self.time_mlp[0].weight.dtype)) + + +class ConvPosEmbedding(nn.Module): + """Depthwise-grouped conv stack that adds local position information.""" + + def __init__(self, dim: int, kernel_size: int = 31, groups: int = 16) -> None: + super().__init__() + if kernel_size % 2 == 0: + raise ValueError(f"kernel_size must be odd, got {kernel_size}") + padding = kernel_size // 2 + self.conv1d = nn.Sequential( + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=padding), + nn.Mish(), + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=padding), + nn.Mish(), + ) + + def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: + keep = None if mask is None else mask.unsqueeze(1) + x = x.transpose(1, 2) + if keep is not None: + x = x.masked_fill(~keep, 0.0) + for layer in self.conv1d: + x = layer(x) + # Re-zero padding after each convolution so it cannot bleed into + # valid frames through the kernel window. + if keep is not None and isinstance(layer, nn.Conv1d): + x = x.masked_fill(~keep, 0.0) + return x.transpose(1, 2) + + +class AudioEmbedding(nn.Module): + """Project audio latents to model width and add conv position information.""" + + def __init__(self, latent_dim: int, dim: int) -> None: + super().__init__() + self.linear = nn.Linear(latent_dim, dim) + self.conv_pos_embed = ConvPosEmbedding(dim) + + def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: + x = self.linear(x) + return self.conv_pos_embed(x, mask=mask) + x + + +class AdaLayerNorm(nn.Module): + """Timestep-conditioned modulation for a block: six chunks from one projection.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 6) + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward( + self, x: torch.Tensor, emb: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.linear(self.silu(emb)).chunk(6, dim=1) + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +class AdaLayerNormFinal(nn.Module): + """Timestep-conditioned modulation before the output projection.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 2) + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: + scale, shift = self.linear(self.silu(emb)).chunk(2, dim=1) + return self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + + +class FeedForward(nn.Module): + """SwiGLU feed-forward with the gate projection fused into ``linear_in``.""" + + def __init__(self, dim: int, mult: float = 4.0) -> None: + super().__init__() + inner_dim = int(dim * mult) + self.linear_in = nn.Linear(dim, inner_dim * 2, bias=False) + self.linear_out = nn.Linear(inner_dim, dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate, up = self.linear_in(x).chunk(2, dim=-1) + return self.linear_out(F.silu(gate) * up) + + +class Attention(nn.Module): + """Self-attention with per-head QK RMSNorm and rotary positions. + + ``to_out`` stays a ``ModuleList`` because the released checkpoints name the + output projection ``to_out.0``; the reference's trailing dropout carries no + parameters and is dropped. + """ + + def __init__(self, dim: int, heads: int, dim_head: int, attn_mask_enabled: bool = True) -> None: + super().__init__() + self.heads = heads + self.dim_head = dim_head + self.attn_mask_enabled = attn_mask_enabled + inner_dim = heads * dim_head + self.to_qkv = nn.Linear(dim, 3 * inner_dim) + self.q_norm = nn.RMSNorm(dim_head, elementwise_affine=True) + self.k_norm = nn.RMSNorm(dim_head, elementwise_affine=True) + self.to_out = nn.ModuleList([nn.Linear(inner_dim, dim)]) + + def _merge_heads(self, x: torch.Tensor) -> torch.Tensor: + return x.transpose(1, 2).reshape(x.shape[0], -1, self.heads * self.dim_head) + + def _qkv( + self, + packed: torch.Tensor, + q_norm: nn.Module, + k_norm: nn.Module, + rope: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Split a packed QKV projection into heads, RMSNorm Q/K, then rotate.""" + q, k, v = (t.view(t.shape[0], -1, self.heads, self.dim_head).transpose(1, 2) for t in packed.chunk(3, dim=-1)) + q = q_norm(q) + k = k_norm(k) + if rope is not None: + q = _apply_rope(q, rope) + k = _apply_rope(k, rope) + return q, k, v + + def _project(self, x: torch.Tensor, rope: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return self._qkv(self.to_qkv(x), self.q_norm, self.k_norm, rope) + + def forward( + self, x: torch.Tensor, mask: torch.Tensor | None = None, rope: torch.Tensor | None = None + ) -> torch.Tensor: + q, k, v = self._project(x, rope) + attn_mask = mask if self.attn_mask_enabled else None + out = self._merge_heads(_sdpa(q, k, v, attn_mask).to(q.dtype)) + out = self.to_out[0](out) + if mask is not None: + out = out.masked_fill(~mask.unsqueeze(-1), 0.0) + return out + + +class JointAttention(Attention): + """Attention over the audio and text streams jointly, with separate projections.""" + + def __init__(self, dim: int, heads: int, dim_head: int, context_dim: int, attn_mask_enabled: bool = True) -> None: + super().__init__(dim, heads, dim_head, attn_mask_enabled) + inner_dim = heads * dim_head + self.to_qkv_c = nn.Linear(context_dim, 3 * inner_dim) + self.c_q_norm = nn.RMSNorm(dim_head, elementwise_affine=True) + self.c_k_norm = nn.RMSNorm(dim_head, elementwise_affine=True) + self.to_out_c = nn.Linear(inner_dim, context_dim) + + def _project_context(self, c: torch.Tensor, c_rope: torch.Tensor | None) -> tuple[torch.Tensor, ...]: + return self._qkv(self.to_qkv_c(c), self.c_q_norm, self.c_k_norm, c_rope) + + def forward( # type: ignore[override] # joint attention takes the context stream too + self, + x: torch.Tensor, + c: torch.Tensor, + mask: torch.Tensor | None = None, + rope: torch.Tensor | None = None, + c_rope: torch.Tensor | None = None, + c_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + audio_len = x.shape[1] + q, k, v = self._project(x, rope) + c_q, c_k, c_v = self._project_context(c, c_rope) + + q = torch.cat([q, c_q], dim=2) + k = torch.cat([k, c_k], dim=2) + v = torch.cat([v, c_v], dim=2) + + joint_mask = None + if self.attn_mask_enabled and mask is not None: + if c_mask is not None: + joint_mask = torch.cat([mask, c_mask], dim=1) + else: + joint_mask = F.pad(mask, (0, c.shape[1]), value=True) + + out = self._merge_heads(_sdpa(q, k, v, joint_mask).to(q.dtype)) + x_out = self.to_out[0](out[:, :audio_len]) + c_out = self.to_out_c(out[:, audio_len:]) + + if mask is not None: + x_out = x_out.masked_fill(~mask.unsqueeze(-1), 0.0) + if c_mask is not None: + c_out = c_out.masked_fill(~c_mask.unsqueeze(-1), 0.0) + return x_out, c_out + + +class DoubleBlock(nn.Module): + """MM-DiT block: joint attention, separate modulation and feed-forward per stream.""" + + def __init__(self, dim: int, heads: int, dim_head: int, ff_mult: float, attn_mask_enabled: bool) -> None: + super().__init__() + self.attn_norm_c = AdaLayerNorm(dim) + self.attn_norm_x = AdaLayerNorm(dim) + self.attn = JointAttention( + dim=dim, heads=heads, dim_head=dim_head, context_dim=dim, attn_mask_enabled=attn_mask_enabled + ) + self.ff_norm_c = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff_c = FeedForward(dim=dim, mult=ff_mult) + self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff_x = FeedForward(dim=dim, mult=ff_mult) + + def forward( + self, + x: torch.Tensor, + c: torch.Tensor, + t: torch.Tensor, + mask: torch.Tensor | None, + rope: torch.Tensor, + c_rope: torch.Tensor, + c_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, t) + norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, t) + + x_attn, c_attn = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope, c_mask=c_mask) + + c = c + c_gate_msa[:, None] * c_attn + norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] + c = c + c_gate_mlp[:, None] * self.ff_c(norm_c) + + x = x + x_gate_msa[:, None] * x_attn + norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None] + x = x + x_gate_mlp[:, None] * self.ff_x(norm_x) + return c, x + + +class SingleBlock(nn.Module): + """DiT block over the concatenated text and audio sequence.""" + + def __init__(self, dim: int, heads: int, dim_head: int, ff_mult: float, attn_mask_enabled: bool) -> None: + super().__init__() + self.attn_norm = AdaLayerNorm(dim) + self.attn = Attention(dim=dim, heads=heads, dim_head=dim_head, attn_mask_enabled=attn_mask_enabled) + self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(dim=dim, mult=ff_mult) + + def forward(self, x: torch.Tensor, t: torch.Tensor, mask: torch.Tensor | None, rope: torch.Tensor) -> torch.Tensor: + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, t) + x = x + gate_msa[:, None] * self.attn(x=norm, mask=mask, rope=rope) + norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + return x + gate_mlp[:, None] * self.ff(norm) + + +class AuKTransformer(nn.Module): + """Flow-matching velocity predictor for AuK audio latents. + + Args: + dim: Model width. + heads: Attention heads. + dim_head: Per-head width. + ff_mult: Feed-forward expansion factor. + latent_dim: Audio VAE latent channels, the input and output width. + text_hidden_dim: Width of the pre-encoded text hidden states. + num_layers: Double-stream (joint attention) block count. + num_single_layers: Single-stream block count. + attn_mask_enabled: Apply padding masks inside attention. When ``False`` + attention runs unmasked, though padded outputs are still zeroed, + which is what the reference does with the flag off. + """ + + 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.dim = dim + self.latent_dim = latent_dim + self.text_cond: torch.Tensor | None = None + self.text_uncond: torch.Tensor | None = None + + self.time_embed = TimeEmbedding(dim) + self.txt_norm = nn.RMSNorm(dim, elementwise_affine=True) + self.txt_proj = nn.Linear(text_hidden_dim, dim) + self.audio_embed = AudioEmbedding(latent_dim, dim) + self.rotary_embed = Rotary(dim_head) + + self.transformer_blocks = nn.ModuleList( + [ + DoubleBlock( + dim=dim, heads=heads, dim_head=dim_head, ff_mult=ff_mult, attn_mask_enabled=attn_mask_enabled + ) + for _ in range(num_layers) + ] + ) + self.single_transformer_blocks = nn.ModuleList( + [ + SingleBlock( + dim=dim, heads=heads, dim_head=dim_head, ff_mult=ff_mult, attn_mask_enabled=attn_mask_enabled + ) + for _ in range(num_single_layers) + ] + ) + + self.norm_out = AdaLayerNormFinal(dim) + self.proj_out = nn.Linear(dim, latent_dim) + + def project_text(self, text: torch.Tensor) -> torch.Tensor: + """Project LLM hidden states ``[B, nt, text_hidden_dim]`` to model width.""" + return self.txt_norm(self.txt_proj(text)) + + def clear_cache(self) -> None: + """Drop the cached text projection. Call this when the conditioning changes.""" + self.text_cond = None + self.text_uncond = None + + def _embed_audio( + self, + x: torch.Tensor, + ref: torch.Tensor | None, + drop_audio_cond: bool, + mask: torch.Tensor | None, + ref_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None, int]: + """Embed the target latents, prepending the reference prompt when present. + + Returns the ``[ref | target]`` sequence, its padding mask and the + reference length, which the caller uses to slice the target back out. + """ + if ref is not None and ref.shape[1] == 0: + ref = None + + target = self.audio_embed(x, mask=mask) + if ref is None: + return target, mask, 0 + + if drop_audio_cond: + ref = torch.zeros_like(ref) + prompt = self.audio_embed(ref, mask=ref_mask) + audio = torch.cat([prompt, target], dim=1) + prompt_len = prompt.shape[1] + + audio_mask = None + if mask is not None or ref_mask is not None: + if mask is None: + mask = x.new_ones(x.shape[:2], dtype=torch.bool) + if ref_mask is None: + ref_mask = ref.new_ones(ref.shape[:2], dtype=torch.bool) + audio_mask = torch.cat([ref_mask, mask], dim=1) + return audio, audio_mask, prompt_len + + def forward( + self, + x: torch.Tensor, + text: torch.Tensor, + time: torch.Tensor, + *, + mask: torch.Tensor | None = None, + c_mask: torch.Tensor | None = None, + ref: torch.Tensor | None = None, + ref_mask: torch.Tensor | None = None, + drop_audio_cond: bool = False, + drop_text: bool = False, + cfg_infer: bool = False, + cache: bool = False, + ) -> torch.Tensor: + """Predict the flow-matching velocity for the target frames of ``x``. + + Args: + x: Noised target latents ``[B, n, latent_dim]``. + text: Pre-encoded text hidden states ``[B, nt, text_hidden_dim]``. + time: Flow-matching timestep, a scalar or ``[B]``. + mask: Target padding mask ``[B, n]``, ``True`` where valid. + c_mask: Text padding mask ``[B, nt]``. Inferred from all-zero text + rows when omitted. + ref: Reference prompt latents ``[B, np, latent_dim]``, prepended to + the audio sequence. A zero-length ``ref`` counts as absent. + ref_mask: Reference padding mask ``[B, np]``. + drop_audio_cond: Zero the reference prompt, for the CFG uncond branch. + drop_text: Zero the projected text, for the CFG uncond branch. + cfg_infer: Run the cond and uncond branches as one batch of ``2B``, + cond first. ``drop_audio_cond`` and ``drop_text`` are then + driven per branch and ignored. + cache: Reuse the projected text across calls, for an ODE loop over + fixed conditioning. Only the ``cfg_infer`` path caches, as in + the reference. Call :meth:`clear_cache` when the text changes. + + Returns: + Velocity ``[B, n, latent_dim]``, or ``[2B, n, latent_dim]`` under + ``cfg_infer``. + """ + batch = x.shape[0] + if time.ndim == 0: + time = time.repeat(batch) + t = self.time_embed(time) + + if c_mask is None: + # Padding rows are all-zero in the text encoder's output. + c_mask = text.abs().sum(-1) > 0 + + if cfg_infer: + if cache and self.text_cond is not None: + c_cond, c_uncond = self.text_cond, self.text_uncond + else: + c_cond = self.project_text(text) + # The uncond branch drops the text, which is exactly a zeroed projection. + c_uncond = torch.zeros_like(c_cond) + if cache: + self.text_cond, self.text_uncond = c_cond, c_uncond + x_cond, mask_cond, prompt_len = self._embed_audio(x, ref, False, mask, ref_mask) + x_uncond, mask_uncond, _ = self._embed_audio(x, ref, True, mask, ref_mask) + + x = torch.cat((x_cond, x_uncond), dim=0) + c = torch.cat((c_cond, c_uncond), dim=0) + t = torch.cat((t, t), dim=0) + c_mask = torch.cat((c_mask, c_mask), dim=0) + audio_mask = None + if mask_cond is not None and mask_uncond is not None: + audio_mask = torch.cat((mask_cond, mask_uncond), dim=0) + else: + c = self.project_text(text) + if drop_text: + c = torch.zeros_like(c) + x, audio_mask, prompt_len = self._embed_audio(x, ref, drop_audio_cond, mask, ref_mask) + + seq_len = x.shape[1] + text_len = c.shape[1] + rope_audio = self.rotary_embed(seq_len) + rope_text = self.rotary_embed(text_len) + + for block in self.transformer_blocks: + c, x = block(x, c, t, mask=audio_mask, rope=rope_audio, c_rope=rope_text, c_mask=c_mask) + + x = torch.cat([c, x], dim=1) + rope = self.rotary_embed(text_len + seq_len) + single_mask = None if audio_mask is None else torch.cat([c_mask, audio_mask], dim=1) + + for block in self.single_transformer_blocks: + x = block(x, t, mask=single_mask, rope=rope) + + x = x[:, text_len + prompt_len :] + return self.proj_out(self.norm_out(x, t)) + + +def dit_state_dict( + checkpoint: Mapping[str, torch.Tensor] | Iterable[tuple[str, torch.Tensor]], + prefix: str = "transformer.", +) -> dict[str, torch.Tensor]: + """Select the DiT tensors out of a full AuK checkpoint and strip ``prefix``. + + Released checkpoints wrap the backbone under ``transformer.`` alongside the + text-encoder layer-fusion parameters, which this module does not own. + """ + items = checkpoint.items() if isinstance(checkpoint, Mapping) else checkpoint + return {name[len(prefix) :]: tensor for name, tensor in items if name.startswith(prefix)} + + +@torch.no_grad() +def sample_latents( + dit: AuKTransformer, + *, + text: torch.Tensor, + c_mask: torch.Tensor | None, + ref: torch.Tensor, + ref_mask: torch.Tensor | None, + gen_frames: int, + nfe: int = 32, + cfg_strength: float = 1.0, + sway_sampling_coef: float | None = None, + t_grid: list[float] | None = None, + seed: int | None = None, + generator: torch.Generator | None = None, + latent_dim: int | None = None, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + """Integrate the flow from noise to audio latents with explicit Euler steps. + + Single-request only: the batch dimension carries the CFG branches, not + separate prompts, so no target padding mask is needed. + + Args: + dit: The velocity model. + text: Pre-encoded text hidden states ``[1, nt, text_hidden_dim]``. + c_mask: Text padding mask ``[1, nt]``. + ref: Reference prompt latents ``[1, np, latent_dim]``. + ref_mask: Reference padding mask ``[1, np]``. + gen_frames: Target latent frames to generate. + nfe: Euler steps, ignored when ``t_grid`` is given. + cfg_strength: Classifier-free guidance weight. Below ``1e-5`` the + uncond branch is skipped entirely. + sway_sampling_coef: Reshapes the uniform time grid towards ``t=0``. + t_grid: Explicit timesteps, overriding ``nfe`` and the sway reshape. + seed: Seeds the global RNG before drawing the noise. Ignored when + ``generator`` is given. + generator: Draws the initial noise from its own RNG, leaving the global + one untouched. + latent_dim: Latent channels; defaults to the model's. + device: Device for the noise; defaults to the reference's. + dtype: Dtype for the noise and the Euler accumulator; defaults to the + reference's, which is the reference implementation's rule. Pass + ``torch.float32`` explicitly when serving half-precision weights. + Drawing the noise at bf16 instead of fp32 moves a 32-step CFG + trajectory by 0.025 relative MSE, roughly 17 times what bf16 + weights themselves cost, because the draw consumes the generator + differently and starts the ODE somewhere else. Keeping it fp32 + costs a few tens of KB and no measurable time. + + Returns: + The target latents at ``t=1``, ``[1, gen_frames, latent_dim]``. + """ + device = torch.device(device) if device is not None else ref.device + dtype = dtype if dtype is not None else ref.dtype + latent_dim = latent_dim if latent_dim is not None else dit.latent_dim + + if generator is None: + if seed is not None: + torch.manual_seed(seed) + x = torch.randn(gen_frames, latent_dim, device=device, dtype=dtype).unsqueeze(0) + else: + x = ( + torch.randn(gen_frames, latent_dim, generator=generator, device=generator.device, dtype=dtype) + .to(device) + .unsqueeze(0) + ) + + if t_grid is not None: + t = torch.tensor(t_grid, device=device, dtype=torch.float32) + else: + t = torch.linspace(0, 1, nfe + 1, device=device, dtype=torch.float32) + if sway_sampling_coef is not None: + t = t + sway_sampling_coef * (torch.cos(math.pi / 2 * t) - 1 + t) + + guided = cfg_strength >= 1e-5 + if t.numel() < 2 or not bool(torch.all(t[1:] > t[:-1])): + raise ValueError( + f"AuK sampling needs a strictly increasing time grid with at least two points; got {t.tolist()}" + ) + try: + for i in range(t.shape[0] - 1): + if guided: + pred = dit( + x, + text, + t[i], + c_mask=c_mask, + ref=ref, + ref_mask=ref_mask, + cfg_infer=True, + cache=True, + ) + v_cond, v_uncond = pred.chunk(2, dim=0) + v = v_cond + (v_cond - v_uncond) * cfg_strength + else: + v = dit(x, text, t[i], c_mask=c_mask, ref=ref, ref_mask=ref_mask) + x = x + (t[i + 1] - t[i]) * v + finally: + # The cached text projections belong to this request only; a failed + # step must not leak them into the next one. + dit.clear_cache() + return x diff --git a/vllm_omni/diffusion/models/auk/auk_vae.py b/vllm_omni/diffusion/models/auk/auk_vae.py new file mode 100644 index 000000000..1fabeb327 --- /dev/null +++ b/vllm_omni/diffusion/models/auk/auk_vae.py @@ -0,0 +1,505 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# +# Adapted from the AuK reference release (MIT), which in turn adapts NVIDIA +# BigVGAN (MIT), jik876/hifi-gan (MIT), junjun3518/alias-free-torch +# (Apache-2.0), EdwardDixon/snake (MIT) and adefossez/julius (MIT). +"""BigVGAN-flow VAE for AuK: 24 kHz mono waveform <-> 50 Hz, 64-channel latents. + +Inference only. The reference module also carries a residual coupling flow, a KL +term and a discriminator; all three are training-only, so the released +``flow.*`` tensors are skipped (see :meth:`AuKVAE.load_weights`). + +Parameter and buffer names follow the reference release, so the published +``vae.safetensors`` loads without a key map. That is also why the convolutions +are built with the legacy ``weight_norm`` parametrization (``weight_g`` / +``weight_v``): :meth:`remove_weight_norm` folds them into plain weights once +the checkpoint is in. +""" + +from __future__ import annotations + +import inspect +import math +import warnings +from pathlib import Path + +import torch +import torch.nn.functional as F +from safetensors.torch import load_file +from torch import nn +from torch.nn.utils import remove_weight_norm as _fold_weight_norm +from torch.nn.utils import weight_norm as _apply_weight_norm +from vllm.logger import init_logger + +logger = init_logger(__name__) + +__all__ = ["AuKVAE"] + +# Checkpoint prefix owned by the training-only coupling flow. +TRAINING_ONLY_PREFIX = "flow." + + +def _weight_norm(module: nn.Module) -> nn.Module: + """Legacy ``weight_g``/``weight_v`` split, which is how the release stores these tensors.""" + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message=".*weight_norm.*deprecated.*", category=FutureWarning) + return _apply_weight_norm(module) + + +def _same_padding(kernel_size: int, dilation: int = 1) -> int: + return int((kernel_size * dilation - dilation) / 2) + + +def _kaiser_sinc_filter(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: + """Kaiser-windowed sinc low-pass, returned as a ``[1, 1, kernel_size]`` conv kernel.""" + half_size = kernel_size // 2 + attenuation = 2.285 * (half_size - 1) * math.pi * (4 * half_width) + 7.95 + if attenuation > 50.0: + beta = 0.1102 * (attenuation - 8.7) + elif attenuation >= 21.0: + beta = 0.5842 * (attenuation - 21.0) ** 0.4 + 0.07886 * (attenuation - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + if kernel_size % 2 == 0: + time = torch.arange(-half_size, half_size) + 0.5 + else: + time = torch.arange(kernel_size) - half_size + if cutoff == 0: + return torch.zeros_like(time).view(1, 1, kernel_size) + taps = 2 * cutoff * window * torch.sinc(2 * cutoff * time) + # Normalize to unit sum so a constant input passes through unscaled. + return (taps / taps.sum()).view(1, 1, kernel_size) + + +class Conv(nn.Conv1d): + """1-D convolution over ``[B, C, T]``; ``causal`` swaps centred padding for a left pad.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int = 1, + stride: int = 1, + dilation: int = 1, + bias: bool = True, + causal: bool = False, + ) -> None: + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=0 if causal else _same_padding(kernel_size, dilation), + dilation=dilation, + bias=bias, + ) + self.causal = causal + self.left_padding = dilation * (kernel_size - 1) if causal else 0 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.causal: + x = F.pad(x, (self.left_padding, 0)) + return super().forward(x) + + +class ConvTranspose(nn.ConvTranspose1d): + """Transposed 1-D convolution; ``causal`` drops the trailing ``stride`` samples instead of padding.""" + + def __init__( + self, in_channels: int, out_channels: int, kernel_size: int, stride: int, causal: bool = False + ) -> None: + if causal and kernel_size != 2 * stride: + raise ValueError(f"causal ConvTranspose needs kernel_size == 2 * stride, got {kernel_size} and {stride}") + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=0 if causal else (kernel_size - stride) // 2, + ) + self.causal = causal + self.trim = stride + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = super().forward(x) + return x[:, :, : -self.trim] if self.causal else x + + +class NormConv(nn.Module): + """Weight-normalized conv held under ``layer``; the attribute name is load-bearing for the checkpoint.""" + + def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 1, stride: int = 1) -> None: + super().__init__() + self.layer = _weight_norm( + nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, padding=(kernel_size - 1) // 2) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.layer(x) + + +class ResStack(nn.Module): + """Stack of dilated residual conv pairs. The leaky slope is torch's 0.01 default, as upstream.""" + + def __init__(self, channels: int, kernel_size: int = 3, base: int = 3, nums: int = 4) -> None: + super().__init__() + self.layers = nn.ModuleList( + [ + nn.Sequential( + nn.LeakyReLU(), + _weight_norm(nn.Conv1d(channels, channels, kernel_size, dilation=base**i, padding=base**i)), + nn.LeakyReLU(), + _weight_norm(nn.Conv1d(channels, channels, kernel_size, dilation=1, padding=1)), + ) + for i in range(nums) + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for layer in self.layers: + x = x + layer(x) + return x + + +class Encoder(nn.Module): + """Strided conv encoder over ``[B, 1, T]``, emitting ``[B, 2 * latent_dim, T / hop]`` mean/log-std pairs.""" + + def __init__( + self, + latent_dim: int, + channels: tuple[int, ...], + down_sample_factors: tuple[int, ...], + in_channels: int = 1, + proj_kernel_size: int = 3, + stack_kernel_size: int = 3, + stack_dilation_base: int = 2, + stacks: int = 6, + ) -> None: + super().__init__() + layers: list[nn.Module] = [ + NormConv(in_channels, channels[0], kernel_size=proj_kernel_size), + nn.LeakyReLU(0.2, True), + ] + for in_c, out_c, factor in zip(channels[:-1], channels[1:], down_sample_factors, strict=True): + layers += [ + NormConv(in_c, out_c, kernel_size=factor * 2, stride=factor), + ResStack(out_c, stack_kernel_size, stack_dilation_base, stacks), + nn.LeakyReLU(0.2, True), + ] + layers.append(NormConv(channels[-1], latent_dim * 2, kernel_size=proj_kernel_size)) + self.generator = nn.Sequential(*layers) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.generator(x) + + +class SnakeBeta(nn.Module): + """``x + sin^2(alpha * x) / beta`` with per-channel learned alpha and beta (log-scale in AuK).""" + + def __init__(self, channels: int, alpha_logscale: bool = False) -> None: + super().__init__() + self.alpha_logscale = alpha_logscale + init = torch.zeros(channels) if alpha_logscale else torch.ones(channels) + self.alpha = nn.Parameter(init.clone()) + self.beta = nn.Parameter(init.clone()) + self.eps = 1e-9 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + return x + (1.0 / (beta + self.eps)) * torch.sin(x * alpha).pow(2) + + +class LowPass(nn.Module): + """Fixed FIR low-pass with replicate padding, applied per channel.""" + + def __init__( + self, + cutoff: float = 0.5, + half_width: float = 0.6, + stride: int = 1, + kernel_size: int = 12, + causal: bool = False, + ) -> None: + super().__init__() + if not 0.0 <= cutoff <= 0.5: + raise ValueError(f"cutoff must lie in [0, 0.5], got {cutoff}") + if causal: + self.pad_left, self.pad_right = kernel_size - 1, 0 + else: + even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.register_buffer("filter", _kaiser_sinc_filter(cutoff, half_width, kernel_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + channels = x.size(1) + x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate") + return F.conv1d(x, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels) + + +class Upsample(nn.Module): + """Band-limited interpolation by ``ratio``. Always non-causal, matching how AuK builds it.""" + + def __init__(self, ratio: int = 2, kernel_size: int = 12) -> None: + super().__init__() + self.ratio = ratio + self.pad = kernel_size // ratio - 1 + self.pad_left = self.pad * ratio + (kernel_size - ratio) // 2 + self.pad_right = self.pad * ratio + (kernel_size - ratio + 1) // 2 + self.register_buffer("filter", _kaiser_sinc_filter(0.5 / ratio, 0.6 / ratio, kernel_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + channels = x.size(1) + x = F.pad(x, (self.pad, self.pad), mode="replicate") + x = self.ratio * F.conv_transpose1d(x, self.filter.expand(channels, -1, -1), stride=self.ratio, groups=channels) + return x[..., self.pad_left : -self.pad_right] + + +class Downsample(nn.Module): + """Decimation by ``ratio`` behind a low-pass; ``lowpass`` holds the filter buffer.""" + + def __init__(self, ratio: int = 2, kernel_size: int = 12, causal: bool = False) -> None: + super().__init__() + self.lowpass = LowPass( + cutoff=0.5 / ratio, half_width=0.6 / ratio, stride=ratio, kernel_size=kernel_size, causal=causal + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.lowpass(x) + + +class AliasFreeActivation(nn.Module): + """Oversample, apply the pointwise activation, decimate, so snake harmonics do not fold back.""" + + def __init__( + self, + activation: nn.Module, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + causal: bool = False, + ) -> None: + super().__init__() + self.act = activation + self.upsample = Upsample(up_ratio, up_kernel_size) + self.downsample = Downsample(down_ratio, down_kernel_size, causal=causal) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.downsample(self.act(self.upsample(x))) + + +class AmpBlock(nn.Module): + """Anti-aliased multi-periodicity residual block: alternating alias-free snake and dilated convs.""" + + def __init__( + self, + channels: int, + kernel_size: int, + dilations: tuple[int, ...], + snake_logscale: bool = True, + causal: bool = True, + act_causal: bool = False, + ) -> None: + super().__init__() + self.convs1 = nn.ModuleList( + [_weight_norm(Conv(channels, channels, kernel_size, 1, dilation=d, causal=causal)) for d in dilations] + ) + self.convs2 = nn.ModuleList( + [_weight_norm(Conv(channels, channels, kernel_size, 1, dilation=1, causal=causal)) for _ in dilations] + ) + self.activations = nn.ModuleList( + [ + AliasFreeActivation(SnakeBeta(channels, alpha_logscale=snake_logscale), causal=act_causal) + for _ in range(len(self.convs1) + len(self.convs2)) + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + acts1, acts2 = self.activations[::2], self.activations[1::2] + for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2, strict=True): + x = conv2(act2(conv1(act1(x)))) + x + return x + + +class AuKVAE(nn.Module): + """AuK's audio codec: :meth:`encode` waveform to normalized latents, :meth:`decode` back. + + Build from the released ``model.vae.model_init_kwargs`` with :meth:`from_config`, then + :meth:`load_weights`. Latents are normalized by the checkpoint's global statistics, which is + the space the rectified-flow transformer works in. + """ + + def __init__( + self, + *, + upsample_rates: tuple[int, ...] = (5, 4, 3, 2, 2, 2), + upsample_kernel_sizes: tuple[int, ...] = (10, 8, 6, 4, 4, 4), + upsample_initial_channel: int = 1536, + resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11), + resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ((1, 3, 5), (1, 3, 5), (1, 3, 5)), + downsample_rates: tuple[int, ...] = (2, 2, 2, 3, 4, 5), + downsample_channels: tuple[int, ...] = (12, 24, 48, 96, 192, 384, 768), + snake_logscale: bool = True, + latent_dim: int = 64, + use_vae: bool = True, + causal: bool = True, + act_causal: bool = True, + sample_rate: int = 24000, + ) -> None: + super().__init__() + if not use_vae: + raise ValueError("AuKVAE requires use_vae=True: the encoder head emits mean/log-std pairs") + self.latent_dim = latent_dim + self.sample_rate = sample_rate + self.hop_size = math.prod(downsample_rates) + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + self._norm_folded = False + + self.register_buffer("global_mean", torch.zeros(latent_dim, dtype=torch.float32)) + self.register_buffer("global_log_std", torch.ones(latent_dim, dtype=torch.float32)) + + self.audio_encoder = Encoder( + latent_dim=latent_dim, + channels=tuple(downsample_channels), + down_sample_factors=tuple(downsample_rates), + ) + + self.conv_pre = _weight_norm(Conv(latent_dim, upsample_initial_channel, 7, 1, causal=False)) + self.ups = nn.ModuleList() + for i, (rate, kernel) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True)): + # One-element inner list: the checkpoint keys are ``ups..0.*``. + self.ups.append( + nn.ModuleList( + [ + _weight_norm( + ConvTranspose( + upsample_initial_channel // (2**i), + upsample_initial_channel // (2 ** (i + 1)), + kernel, + rate, + causal=causal, + ) + ) + ] + ) + ) + + self.resblocks = nn.ModuleList() + channels = upsample_initial_channel + for i in range(self.num_upsamples): + channels = upsample_initial_channel // (2 ** (i + 1)) + for kernel, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True): + self.resblocks.append( + AmpBlock( + channels, + kernel, + tuple(dilations), + snake_logscale=snake_logscale, + causal=causal, + act_causal=act_causal, + ) + ) + + self.activation_post = AliasFreeActivation( + SnakeBeta(channels, alpha_logscale=snake_logscale), causal=act_causal + ) + self.conv_post = _weight_norm(Conv(channels, 1, 7, 1, causal=causal, bias=False)) + + @classmethod + def from_config(cls, cfg: dict) -> AuKVAE: + """Build from the released ``model_init_kwargs`` mapping, ignoring training-only entries.""" + accepted = set(inspect.signature(cls.__init__).parameters) - {"self"} + return cls(**{key: value for key, value in cfg.items() if key in accepted}) + + def load_weights( + self, path: str | Path, *, device: str | torch.device | None = None, fold_weight_norm: bool = True + ) -> tuple[list[str], list[str]]: + """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. + """ + target = self.global_mean.device if device is None else torch.device(device) + missing, unexpected = self.load_state_dict(load_file(str(path), device=str(target)), strict=False) + missing, unexpected = list(missing), list(unexpected) + stray = [key for key in missing + unexpected if not key.startswith(TRAINING_ONLY_PREFIX)] + if stray: + raise ValueError(f"AuK VAE weights at {path} do not match this module: {sorted(stray)[:8]}") + logger.info( + "Loaded AuK VAE from %s (skipped %d training-only %s* tensors)", path, len(unexpected), TRAINING_ONLY_PREFIX + ) + if fold_weight_norm: + self.remove_weight_norm() + return missing, unexpected + + def remove_weight_norm(self) -> None: + """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. + """ + if self._norm_folded: + return + folded = 0 + for module in self.modules(): + if hasattr(module, "weight_g"): + _fold_weight_norm(module) + folded += 1 + self._norm_folded = True + logger.debug("Folded weight norm on %d AuK VAE convolutions", folded) + + def encode( + self, wav: torch.Tensor, *, sample: bool = False, generator: torch.Generator | None = None + ) -> torch.Tensor: + """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. + """ + if wav.dim() != 2 or wav.size(0) != 1: + raise ValueError(f"encode expects a mono [1, T] waveform, got {tuple(wav.shape)}") + # Statistics stay in fp32 even under an ambient autocast, as in the reference. + with torch.autocast(device_type=wav.device.type, enabled=False): + mean, log_std = self.audio_encoder(wav.unsqueeze(1)).chunk(2, dim=1) + if sample: + latents = mean + self._noise_like(mean, generator) * torch.exp(log_std) + else: + latents = mean + latents = latents.transpose(1, 2).float() + return (latents - self.global_mean.float()) / torch.sqrt(self.global_log_std.float()) + + def decode(self, latents: torch.Tensor) -> torch.Tensor: + """Decode normalized ``[B, Np, latent_dim]`` latents into a ``[B, Np * hop_size]`` waveform in [-1, 1].""" + if latents.dim() != 3 or latents.size(-1) != self.latent_dim: + raise ValueError(f"decode expects [B, Np, {self.latent_dim}] latents, got {tuple(latents.shape)}") + x = latents.float() * torch.sqrt(self.global_log_std.float()) + self.global_mean.float() + x = self.conv_pre(x.permute(0, 2, 1)) + for i in range(self.num_upsamples): + for up in self.ups[i]: + x = up(x) + first = i * self.num_kernels + residual = self.resblocks[first](x) + for j in range(1, self.num_kernels): + residual = residual + self.resblocks[first + j](x) + x = residual / self.num_kernels + x = self.conv_post(self.activation_post(x)) + return torch.clamp(x, min=-1.0, max=1.0).squeeze(1) + + @staticmethod + def _noise_like(reference: torch.Tensor, generator: torch.Generator | None) -> torch.Tensor: + """Draw standard normal noise shaped like ``reference``; ``generator`` decides which RNG stream.""" + device = reference.device if generator is None else generator.device + noise = torch.randn(reference.shape, generator=generator, dtype=reference.dtype, device=device) + return noise.to(reference.device) diff --git a/vllm_omni/diffusion/models/auk/pipeline_auk.py b/vllm_omni/diffusion/models/auk/pipeline_auk.py new file mode 100644 index 000000000..6ad12b02d --- /dev/null +++ b/vllm_omni/diffusion/models/auk/pipeline_auk.py @@ -0,0 +1,486 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""AuK audio-editing pipeline for vLLM-Omni. + +Turns a layer-fused text condition (produced by the encoder stage) plus an +optional source clip into a 24 kHz waveform: the source clip is encoded to VAE +latents, a rectified-flow DiT integrates the target latents conditioned on +both, and the BigVGAN-flow decoder renders the waveform. +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Iterable +from typing import Any, ClassVar + +import numpy as np +import torch +import torchaudio +from safetensors import safe_open +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, dit_state_dict, sample_latents +from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE +from vllm_omni.diffusion.models.interface import ( + SupportAudioInput, + SupportAudioOutput, + SupportsComponentDiscovery, +) +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.model_extras.auk import resolve_gen_frames + +logger = init_logger(__name__) + +# Longest ref + target latent sequence the released checkpoints accept. +MAX_LATENT_FRAMES = 65536 + +# Weight-file layout: the DiT lives under this prefix; the layer-fusion +# parameters sit next to it but belong to the encoder stage. +_DIT_PREFIX = "transformer." +_FUSION_KEYS = frozenset({"layer_weights", "layer_scale"}) + +# The distilled variant only reproduces its training recipe at these settings. +_FLASH_NFE = 4 +_FLASH_CFG = 0.0 +# Decorrelates the VAE posterior draw from the initial latent for the same request seed. +_VAE_SEED_OFFSET = 0x5EED_0A0C + + +def get_auk_post_process_func(od_config: OmniDiffusionConfig): + """Create the post-processing function for AuK audio output. + + Tensor output types pass through; anything else becomes a numpy waveform. + The sample rate is not attached here: the output formatter reads + ``AuKPipeline.audio_sample_rate`` for audio-output pipelines. + """ + + del od_config # The conversion does not depend on the config. + + def post_process_func( + audio: torch.Tensor, + output_type: str = "np", + ): + if output_type in ("pt", "latent"): + return audio + return audio.cpu().float().numpy() + + return post_process_func + + +def _prompt_mapping(prompt: Any) -> dict[str, Any]: + """Return the prompt dict, or an empty mapping for a bare text prompt.""" + + if isinstance(prompt, dict): + return prompt + return {} + + +def _unwrap_single(value: Any) -> Any: + """Unwrap the one-element lists the serving layer sometimes produces.""" + + if isinstance(value, list) and len(value) == 1: + return value[0] + return value + + +def _split_audio(audio: Any, default_sample_rate: int) -> tuple[Any, int]: + """Split a multimodal audio item into raw samples and a sample rate.""" + + if isinstance(audio, (tuple, list)) and len(audio) == 2 and isinstance(audio[1], (int, float)): + return audio[0], int(audio[1]) + return audio, default_sample_rate + + +class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComponentDiscovery): + """Instruction-driven audio generation and editing with AuK. + + One request per forward: the rectified-flow ODE runs over the whole target + span and the batch dimension carries the CFG branches, so there is nothing + to share between requests yet. + + Args: + od_config: OmniDiffusion configuration. ``od_config.model`` must be an + assembled AuK directory (``config.json``, ``auk.safetensors``, + ``vae.safetensors``), as produced by + ``tools/prepare_auk_checkpoint.py``. + prefix: Unused; kept for the pipeline construction contract. + """ + + supports_request_batch = False + + # Picked up by ``supports_audio_output`` in the diffusion engine so the + # default stage metadata reports ``final_output_type="audio"`` and the + # ``multimodal_output`` payload includes the sample rate. + support_audio_output: ClassVar[bool] = True + support_audio_input: ClassVar[bool] = True + audio_sample_rate: ClassVar[int] = 24000 + + # The DiT cannot run without an encoder-stage text condition, which the + # engine's synthetic warmup request has no way to produce. + dummy_run_num_frames: ClassVar[int] = 0 + + _dit_modules: ClassVar[list[str]] = ["dit"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ) -> None: + super().__init__() + del prefix # Weights are not namespaced: one checkpoint, one pipeline. + self.od_config = od_config + self.device = get_local_device() + self.dtype = getattr(od_config, "dtype", None) or torch.bfloat16 + + model_dir = od_config.model + if not model_dir or not os.path.isdir(model_dir): + raise ValueError( + f"AuK needs an assembled local checkpoint directory, got {model_dir!r}. " + "Build one with tools/prepare_auk_checkpoint.py." + ) + config = _read_config(model_dir) + + vae_config = dict(config["vae"]) + self.latent_dim = int(vae_config["latent_dim"]) + self.hop_size = int(vae_config["downsample_rate"]) + self.sample_rate = int(vae_config["target_sample_rate"]) + + self.variant = str(config.get("variant", "base")) + self.is_flash = self.variant == "flash" + self.flash_t_grid = [float(t) for t in config.get("flash_t_grid") or ()] + if self.is_flash and len(self.flash_t_grid) != _FLASH_NFE + 1: + raise ValueError(f"The flash variant needs a {_FLASH_NFE + 1}-point flash_t_grid, got {self.flash_t_grid}.") + defaults = dict(config.get("defaults") or {}) + self.default_nfe = int(defaults.get("nfe", 32)) + self.default_cfg = float(defaults.get("cfg", 2.0)) + self.default_sway = defaults.get("sway", -1.0) + self._flash_lock_logged = False + + # The codec is small and numerically sensitive: it stays in fp32. It + # must sit on the inference device BEFORE load_weights folds weight + # norm: a CPU fold lands one ULP off the accelerator's and the encoder + # amplifies that into a 1e-2 latent drift. + self.vae = AuKVAE.from_config(vae_config["model_init_kwargs"]).to(device=self.device, dtype=torch.float32) + self.vae.load_weights(os.path.join(model_dir, "vae.safetensors")) + self.vae = self.vae.eval() + self.vae.requires_grad_(False) + self._check_vae_geometry() + + # The transformer owns every key the checkpoint tool writes into `dit`, + # `attn_mask_enabled` included, so the section passes straight through. + self.dit = AuKTransformer(latent_dim=self.latent_dim, **config["dit"]) + self.dit = self.dit.to(dtype=self.dtype) + self.dit.load_state_dict(_read_dit_weights(model_dir, self.dtype), strict=True) + self.dit = self.dit.to(device=self.device).eval() + self.dit.requires_grad_(False) + + logger.info( + "AuK pipeline ready: variant=%s dtype=%s latent_dim=%d hop=%d sample_rate=%d", + self.variant, + self.dtype, + self.latent_dim, + self.hop_size, + self.sample_rate, + ) + + # The assembled checkpoint is not a diffusers layout: __init__ reads + # auk.safetensors and vae.safetensors directly, so the loader has no + # component sources to stream and load_weights only reports what is loaded. + weights_sources: ClassVar[tuple] = () + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stray = [name for name, _ in weights] + if stray: + logger.warning("AuKPipeline ignores %d loader-provided weights (e.g. %s)", len(stray), stray[:3]) + return {name for name, _ in self.named_parameters()} + + def _check_vae_geometry(self) -> None: + """Fail fast when config.json and the VAE disagree on latent geometry.""" + + declared = { + "latent_dim": (self.latent_dim, int(self.vae.latent_dim)), + "downsample_rate": (self.hop_size, int(self.vae.hop_size)), + "target_sample_rate": (self.sample_rate, int(self.vae.sample_rate)), + } + bad = {name: pair for name, pair in declared.items() if pair[0] != pair[1]} + if bad: + raise ValueError(f"config.json and the AuK VAE disagree (config, vae): {bad}") + if self.sample_rate != self.audio_sample_rate: + raise ValueError( + f"AuK advertises {self.audio_sample_rate} Hz output but the checkpoint is {self.sample_rate} Hz." + ) + + def _resolve_generator(self, sampling_params: Any) -> torch.Generator | None: + """Per-request generator; the process-global RNG is never seeded.""" + + generator = sampling_params.generator + if isinstance(generator, list): + if len(generator) > 1: + logger.warning( + "AuKPipeline runs one request per forward; using the first of %d generators", len(generator) + ) + generator = generator[0] if generator else None + if generator is None and sampling_params.seed is not None: + generator = torch.Generator(device=self.device).manual_seed(int(sampling_params.seed)) + return generator + + def _prepare_waveform(self, audio: Any) -> torch.Tensor: + """Return mono float32 samples ``[1, T]`` at the VAE sample rate. + + ``T`` is truncated to a whole number of latent frames. The encoder is + causal, so dropping a trailing partial frame leaves every emitted frame + unchanged. + """ + + data, sample_rate = _split_audio(audio, self.sample_rate) + if isinstance(data, np.ndarray): + wav = torch.from_numpy(np.ascontiguousarray(data)) + else: + wav = torch.as_tensor(data) + wav = wav.detach().to(device="cpu", dtype=torch.float32) + + if wav.ndim == 1: + wav = wav.unsqueeze(0) + elif wav.ndim == 2: + # Accept both [channels, samples] and soundfile's [samples, channels]; + # the channel axis is the short one (at most 8 wide). + if wav.shape[0] > wav.shape[1] and wav.shape[1] <= 8: + wav = wav.transpose(0, 1) + if wav.shape[0] > 1: + wav = wav.mean(dim=0, keepdim=True) + else: + raise ValueError(f"AuK audio input must be [samples] or [channels, samples], got {tuple(wav.shape)}.") + + if wav.shape[-1] == 0: + raise ValueError("AuK audio input is empty.") + if not torch.isfinite(wav).all(): + raise ValueError("AuK audio input contains NaN or Inf.") + + if sample_rate != self.sample_rate: + wav = torchaudio.functional.resample(wav, sample_rate, self.sample_rate) + + frames = wav.shape[-1] // self.hop_size + if frames < 1: + raise ValueError( + f"AuK audio input is shorter than one latent frame ({self.hop_size} samples at {self.sample_rate} Hz)." + ) + return wav[..., : frames * self.hop_size] + + def _encode_source( + self, + audio: Any, + *, + sample: bool, + generator: torch.Generator | None, + ) -> torch.Tensor: + """Encode the source clip to normalized latents ``[1, np, latent_dim]``.""" + + if audio is None: + return torch.zeros(1, 0, self.latent_dim, device=self.device, dtype=torch.float32) + wav = self._prepare_waveform(_unwrap_single(audio)).to(self.device) + return self.vae.encode(wav, sample=sample, generator=generator) + + def _target_frames(self, gen_seconds: Any, ref_frames: int) -> int: + """Target latent length, capped by the remaining context.""" + + frames = resolve_gen_frames( + None if gen_seconds is None else float(gen_seconds), + ref_frames, + sample_rate=self.sample_rate, + hop=self.hop_size, + ) + budget = MAX_LATENT_FRAMES - ref_frames + if budget < 1: + raise ValueError( + f"The source clip already fills the {MAX_LATENT_FRAMES}-frame context ({ref_frames} frames)." + ) + if frames > budget: + logger.warning("AuK target of %d frames exceeds the remaining context; clamping to %d.", frames, budget) + frames = budget + return frames + + def _resolve_schedule( + self, + sampling_params: Any, + knobs: dict[str, Any], + ) -> tuple[int, float, float | None, list[float] | None]: + """Resolve (nfe, cfg, sway, t_grid), locking the distilled recipe on Flash. + + A knob present but ``None`` counts as unspecified: the stage input + processor fills every key it knows about, whether the caller set it or + not. + """ + + steps = sampling_params.num_inference_steps + nfe = self.default_nfe if steps is None else int(steps) + if nfe < 1: + raise ValueError(f"num_inference_steps must be >= 1 for AuK; got {steps}") + cfg = float(sampling_params.guidance_scale) if sampling_params.guidance_scale_provided else self.default_cfg + sway = knobs.get("sway") + sway = self.default_sway if sway is None else float(sway) + raw_grid = knobs.get("t_grid") + t_grid = [float(t) for t in raw_grid] if raw_grid else None + + if not self.is_flash: + return nfe, cfg, sway, t_grid + + locked = (_FLASH_NFE, _FLASH_CFG, None, self.flash_t_grid) + if (nfe, cfg, sway, t_grid) != locked and not self._flash_lock_logged: + self._flash_lock_logged = True + logger.warning( + "AuK-Flash ignores per-request sampling knobs (asked nfe=%s cfg=%s sway=%s t_grid=%s); " + "using the distilled recipe nfe=%d cfg=%s with the checkpoint time grid.", + nfe, + cfg, + sway, + t_grid, + _FLASH_NFE, + _FLASH_CFG, + ) + return locked + + def _dit_autocast(self): + """Match the reference, which runs the DiT under bf16 autocast.""" + + return torch.autocast( + device_type=self.device.type, + dtype=self.dtype, + enabled=self.dtype in (torch.bfloat16, torch.float16), + ) + + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: + """Generate one waveform. + + Args: + req: Request batch holding exactly one request. The prompt carries + ``prompt_embeds`` (the fused text condition ``[nt, 2048]``), an + optional ``multi_modal_data["audio"]`` source clip, and + ``additional_information["auk"]`` with ``gen_seconds``, + ``sway``, ``t_grid`` and ``vae_sample``. ``num_inference_steps``, + ``guidance_scale`` and ``seed``/``generator`` come from the + sampling params. + + Returns: + One ``DiffusionOutput`` whose ``output`` is a float32 mono waveform + ``[T]`` at 24 kHz, or the normalized target latents + ``[1, gen_frames, latent_dim]`` when ``output_type`` is ``latent``. + """ + + assert req.num_reqs == 1, f"AuKPipeline runs one request per forward, got {req.num_reqs}." + prompt = _prompt_mapping(req.prompts[0]) + sampling_params = req.sampling_params + knobs = dict((prompt.get("additional_information") or {}).get("auk") or {}) + + text = _unwrap_single(prompt.get("prompt_embeds")) + if text is None: + raise ValueError("AuK stage 1 needs `prompt_embeds`, the fused text condition from the encoder stage.") + text = torch.as_tensor(text).to(device=self.device, dtype=self.dtype) + if text.ndim == 3 and text.shape[0] == 1: + text = text[0] + if text.ndim != 2: + raise ValueError(f"AuK `prompt_embeds` must be [nt, text_hidden_dim], got {tuple(text.shape)}.") + text = text.unsqueeze(0) + c_mask = torch.ones(text.shape[:2], dtype=torch.bool, device=self.device) + + generator = self._resolve_generator(sampling_params) + audio = (prompt.get("multi_modal_data") or {}).get("audio") + output_type = sampling_params.output_type or "np" + + with torch.inference_mode(): + # The VAE posterior draw gets its own generator, seeded with a + # fixed offset from the request seed: the initial latent below then + # still equals the reference's fresh manual_seed draw, and the two + # noise streams are not copies of each other. + vae_generator = None + if knobs.get("vae_sample") and generator is not None: + vae_seed = (int(generator.initial_seed()) + _VAE_SEED_OFFSET) % (2**63) + vae_generator = torch.Generator(device=self.device).manual_seed(vae_seed) + ref = self._encode_source(audio, sample=bool(knobs.get("vae_sample")), generator=vae_generator) + ref_frames = ref.shape[1] + ref_mask = torch.ones(ref.shape[:2], dtype=torch.bool, device=self.device) + gen_frames = self._target_frames(knobs.get("gen_seconds"), ref_frames) + nfe, cfg, sway, t_grid = self._resolve_schedule(sampling_params, knobs) + + with self._dit_autocast(): + latents = sample_latents( + self.dit, + text=text, + c_mask=c_mask, + ref=ref, + ref_mask=ref_mask, + gen_frames=gen_frames, + nfe=nfe, + cfg_strength=cfg, + sway_sampling_coef=sway, + t_grid=t_grid, + seed=None, + latent_dim=self.latent_dim, + device=self.device, + dtype=torch.float32, + generator=generator, + ) + + latents = latents.float() + if not torch.isfinite(latents).all(): + raise RuntimeError("AuK generated latents contain NaN or Inf.") + if output_type == "latent": + return [DiffusionOutput(output=latents.detach().cpu())] + + wav = self.vae.decode(latents) + + # One mono waveform per request; the formatter expects [T]. + wav = wav.detach().to(device="cpu", dtype=torch.float32).reshape(-1) + if not torch.isfinite(wav).all(): + raise RuntimeError("AuK generated audio contains NaN or Inf.") + return [DiffusionOutput(output=wav)] + + +def _read_config(model_dir: str) -> dict[str, Any]: + """Read the assembled AuK ``config.json``.""" + + path = os.path.join(model_dir, "config.json") + try: + with open(path) as handle: + config = json.load(handle) + except FileNotFoundError as exc: + raise ValueError(f"AuK checkpoint is missing config.json: {path}") from exc + except json.JSONDecodeError as exc: + logger.error("AuK config.json is not valid JSON: %s", path) + raise ValueError(f"AuK config.json is not valid JSON: {path}") from exc + + missing = [key for key in ("dit", "vae") if key not in config] + if missing: + raise ValueError(f"AuK config.json is missing required sections {missing}: {path}") + return config + + +def _read_dit_weights(model_dir: str, dtype: torch.dtype) -> dict[str, torch.Tensor]: + """Read the DiT state dict from ``auk.safetensors``, cast to ``dtype``. + + Tensors are cast one at a time so the fp32 checkpoint is never fully + resident alongside the model. + """ + + path = os.path.join(model_dir, "auk.safetensors") + if not os.path.isfile(path): + raise ValueError(f"AuK checkpoint is missing auk.safetensors: {path}") + + with safe_open(path, framework="pt", device="cpu") as checkpoint: + keys = list(checkpoint.keys()) + state = dit_state_dict(((key, checkpoint.get_tensor(key).to(dtype)) for key in keys), _DIT_PREFIX) + + if not state: + raise ValueError(f"AuK checkpoint has no '{_DIT_PREFIX}*' weights: {path}") + unexpected = sorted(key for key in keys if key not in _FUSION_KEYS and not key.startswith(_DIT_PREFIX)) + if unexpected: + logger.warning("Ignoring %d unrecognized keys in %s, e.g. %s", len(unexpected), path, unexpected[:5]) + return state diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 39fff7df2..26b052a3a 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -116,6 +116,11 @@ _DIFFUSION_MODELS = { "pipeline_minimax_h3", "MiniMaxH3Pipeline", ), + "AuKPipeline": ( + "auk", + "pipeline_auk", + "AuKPipeline", + ), "StableAudioPipeline": ( "stable_audio", "pipeline_stable_audio", @@ -567,6 +572,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func", "MiniMaxH3Pipeline": "get_minimax_h3_post_process_func", "MiniMaxH3ModularPipeline": "get_minimax_h3_post_process_func", + "AuKPipeline": "get_auk_post_process_func", "StableAudioPipeline": "get_stable_audio_post_process_func", "WanImageToVideoPipeline": "get_wan22_i2v_post_process_func", "WanS2VPipeline": "get_wan22_s2v_post_process_func", diff --git a/vllm_omni/model_executor/models/auk/__init__.py b/vllm_omni/model_executor/models/auk/__init__.py new file mode 100644 index 000000000..d77d6a64d --- /dev/null +++ b/vllm_omni/model_executor/models/auk/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from vllm_omni.model_executor.models.auk.auk import AuKForConditionalGeneration + +__all__ = ["AuKForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/auk/auk.py b/vllm_omni/model_executor/models/auk/auk.py new file mode 100644 index 000000000..ebb20711d --- /dev/null +++ b/vllm_omni/model_executor/models/auk/auk.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK (Tencent): instruction-driven speech generation and editing. + +Stage 0 (``model_stage="encoder"``) runs the frozen Qwen2.5-Omni thinker as a +prefill-only encoder. The outputs of all decoder layers are fused with learned +softmax weights into one ``[tokens, hidden]`` text condition, emitted through +``multimodal_outputs["hidden_states"]["output"]`` for the diffusion stage +(``vllm_omni/diffusion/models/auk``). The stage samples exactly one token +(``max_tokens=1``, forced to EOS); that token is discarded downstream. +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable, Mapping +from functools import cached_property +from itertools import islice + +import torch +import torch.nn as nn +import torch.nn.functional as F +from safetensors import safe_open +from transformers import PretrainedConfig +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP +from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFeatureSpec +from vllm.sequence import IntermediateTensors +from vllm.v1.sample.sampler import Sampler + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni import Qwen2_5OmniForConditionalGeneration +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( + Qwen2_5OmniThinkerDummyInputsBuilder, + Qwen2_5OmniThinkerForConditionalGeneration, + Qwen2_5OmniThinkerMultiModalProcessor, + Qwen2_5OmniThinkerProcessingInfo, +) +from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights + +logger = init_logger(__name__) + +# File next to the thinker shards in the assembled checkpoint directory that +# carries the DiT (``transformer.*``) and the two layer-fusion tensors. +AUK_WEIGHTS_FILE = "auk.safetensors" +LAYER_WEIGHTS_KEY = "layer_weights" +LAYER_SCALE_KEY = "layer_scale" +# The fused text condition travels as ``multimodal_outputs["hidden_states"]["output"]`` +# (``HiddenStatesStruct.output``), the same stage-wire slot MiniMax H3 uses. +TEXT_COND_KEY = "output" +# ``F.layer_norm`` epsilon used by the upstream fusion (torch default). +_FUSION_LN_EPS = 1e-5 + + +class AuKProcessingInfo(Qwen2_5OmniThinkerProcessingInfo): + """Qwen2.5-Omni thinker processing restricted to AuK's single audio input. + + AuK conditions on text plus at most one reference or source clip; images + and video are not part of its contract, so they are not advertised (and the + thinker then builds no vision tower). + """ + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"audio": 1} + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen2_5OmniThinkerMultiModalProcessor, + info=AuKProcessingInfo, + dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder, +) +class AuKForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP, SupportsMRoPE): + """Encoder stage of AuK: Qwen2.5-Omni thinker with learned layer fusion. + + The forward pass re-implements the thinker language model's layer loop so + the per-layer outputs can be fused on the fly: vLLM's decoder layers carry + the residual stream separately, and ``hidden + residual`` after layer ``i`` + is the HF ``output_hidden_states[i + 1]`` entry, while the last entry is + the final-normed state. Only one accumulator lives at a time. + """ + + have_multimodal_outputs = True + # The fused condition is the payload; the runner's default last-hidden + # copy and prefix-cached hidden states are not needed (MiniMax H3 pattern). + omni_pooler_payload_include_hidden = False + requires_full_prefix_cached_hidden_states = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.has_preprocess = False + self.vllm_config = vllm_config + config = vllm_config.model_config.hf_config + + self.model_stage = getattr(vllm_config.model_config, "model_stage", None) or "encoder" + if self.model_stage != "encoder": + raise ValueError( + f"AuKForConditionalGeneration only hosts the 'encoder' stage, got model_stage={self.model_stage!r}; " + "the DiT stage is the AuKPipeline diffusion model." + ) + + thinker_config = config.thinker_config + self.thinker_config = thinker_config + self.thinker = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "thinker"), + hf_config=thinker_config, + architectures=["Qwen2_5OmniThinkerModel"], + ) + # The audio-only mm limits keep the thinker from building its vision + # tower; drop it explicitly as well, as the upstream implementation does, + # so no vision weights are loaded or resident. + if getattr(self.thinker, "visual", None) is not None: + del self.thinker.visual + self.thinker.visual = None + self.model = self.thinker + self.make_empty_intermediate_tensors = self.thinker.make_empty_intermediate_tensors + + num_layers = thinker_config.text_config.num_hidden_layers + # Filled by load_weights() from AUK_WEIGHTS_FILE; buffers so the vLLM + # loader does not expect them among the thinker shards. + self.register_buffer("layer_weights", torch.zeros(num_layers, dtype=torch.float32)) + self.register_buffer("layer_scale", torch.ones(1, dtype=torch.float32)) + self._fusion_loaded = False + + # -------------------- delegations to the thinker -------------------- + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + return Qwen2_5OmniThinkerForConditionalGeneration.get_placeholder_str(modality, i) + + def get_language_model(self) -> nn.Module: + return self.thinker.get_language_model() + + @cached_property + def sampler(self): + if hasattr(self.thinker, "sampler"): + return self.thinker.sampler + return Sampler() + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + is_multimodal=None, + ) -> torch.Tensor: + return self.thinker.embed_input_ids( + input_ids=input_ids, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def embed_multimodal(self, **kwargs): + return self.thinker.embed_multimodal(**kwargs) + + def get_mrope_input_positions( + self, + input_tokens: list[int], + mm_features: list[MultiModalFeatureSpec] | None = None, + *, + hf_config: PretrainedConfig, + image_grid_thw: list[list[int]] | torch.Tensor, + video_grid_thw: list[list[int]] | torch.Tensor, + second_per_grid_ts: list[float] | None = None, + context_len: int = 0, + seq_len: int | None = None, + audio_feature_lengths: torch.Tensor | None = None, + use_audio_in_video: bool = False, + ) -> tuple[torch.Tensor, int]: + # The omni runner filters kwargs by this signature (the extended + # Qwen2.5-Omni form); the unified Qwen2.5-Omni class implements the + # audio-aware position math without touching instance state. + return Qwen2_5OmniForConditionalGeneration.get_mrope_input_positions( + self, + input_tokens, + mm_features, + hf_config=hf_config, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + context_len=context_len, + seq_len=seq_len, + audio_feature_lengths=audio_feature_lengths, + use_audio_in_video=use_audio_in_video, + ) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, **kwargs: object) -> torch.Tensor | None: + """Emit EOS immediately; the useful result is the prefill text condition.""" + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + text_config = self.thinker_config.text_config + logits = hidden_states.new_full( + (hidden_states.shape[0], int(text_config.vocab_size)), + torch.finfo(hidden_states.dtype).min, + ) + logits[:, self._eos_token_id()] = 0 + return logits + + def _eos_token_id(self) -> int: + for cfg in (self.thinker_config.text_config, self.thinker_config, self.vllm_config.model_config.hf_config): + eos = getattr(cfg, "eos_token_id", None) + if isinstance(eos, (list, tuple)) and eos: + eos = eos[0] + if isinstance(eos, int): + return eos + # Qwen2.5-Omni tokenizers end turns with <|im_end|>. + return 151645 + + # -------------------- forward -------------------- + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> OmniOutput: + if intermediate_tensors is not None: + raise NotImplementedError("AuK encoder does not support pipeline parallelism") + language_model = self.thinker.language_model.model + num_layers = len(language_model.layers) + if language_model.start_layer != 0 or language_model.end_layer != num_layers: + raise NotImplementedError("AuK encoder needs all decoder layers on one rank") + + if inputs_embeds is None: + hidden_states = language_model.embed_input_ids(input_ids) + else: + hidden_states = inputs_embeds + residual: torch.Tensor | None = None + + weights = torch.softmax(self.layer_weights, dim=0) + fused: torch.Tensor | None = None + normed_shape = (hidden_states.shape[-1],) + layers = islice(language_model.layers, language_model.start_layer, language_model.end_layer) + for idx, layer in enumerate(layers): + hidden_states, residual = layer(positions, hidden_states, residual) + if idx + 1 < num_layers: + layer_output = hidden_states + residual + else: + # HF's last hidden_states entry is post final norm. + hidden_states, _ = language_model.norm(hidden_states, residual) + layer_output = hidden_states + normed = F.layer_norm(layer_output.float(), normed_shape, eps=_FUSION_LN_EPS) + contribution = normed * weights[idx] + fused = contribution if fused is None else fused + contribution + + assert fused is not None + text_cond = (fused * self.layer_scale).to(hidden_states.dtype) + # Stage-wire layout: [tokens, features] under hidden_states.output so + # the runner slices it per request (and concatenates across chunked + # prefill steps); the stage-1 processor reads the same path. + return OmniOutput( + text_hidden_states=hidden_states, + multimodal_outputs={"hidden_states": {"output": text_cond}}, + ) + + # -------------------- weights -------------------- + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + thinker_weights: list[tuple[str, torch.Tensor]] = [] + skipped = 0 + for name, tensor in weights: + if name.startswith("thinker."): + thinker_weights.append((name, tensor)) + elif name.startswith(("talker.", "token2wav.", "transformer.")): + skipped += 1 + else: + logger.warning("AuK encoder: ignoring unexpected weight %s", name) + if skipped: + logger.debug("AuK encoder: skipped %d non-thinker weights", skipped) + + loaded = self.thinker.load_weights(thinker_weights) + loaded = add_prefix_to_loaded_weights(loaded, "thinker") + self._load_fusion_params() + return loaded + + def _load_fusion_params(self) -> None: + model_dir = self.vllm_config.model_config.model + path = os.path.join(model_dir, AUK_WEIGHTS_FILE) + if not os.path.isfile(path): + raise FileNotFoundError( + f"AuK encoder: {AUK_WEIGHTS_FILE} not found in {model_dir!r}; assemble the checkpoint with " + "tools/prepare_auk_checkpoint.py" + ) + with safe_open(path, framework="pt", device="cpu") as f: + layer_weights = f.get_tensor(LAYER_WEIGHTS_KEY).float() + layer_scale = f.get_tensor(LAYER_SCALE_KEY).float().reshape(1) + if layer_weights.shape != self.layer_weights.shape: + raise ValueError( + f"AuK encoder: {LAYER_WEIGHTS_KEY} has {tuple(layer_weights.shape)} entries, " + f"thinker has {tuple(self.layer_weights.shape)} layers" + ) + self.layer_weights.copy_(layer_weights.to(self.layer_weights.device)) + self.layer_scale.copy_(layer_scale.to(self.layer_scale.device)) + self._fusion_loaded = True + logger.info( + "AuK encoder: loaded layer fusion (scale %.3f, top layer weight %.3f at layer %d)", + float(layer_scale[0]), + float(torch.softmax(layer_weights, 0).max()), + int(torch.softmax(layer_weights, 0).argmax()) + 1, + ) diff --git a/vllm_omni/model_executor/models/auk/pipeline.py b/vllm_omni/model_executor/models/auk/pipeline.py new file mode 100644 index 000000000..8b9d6855b --- /dev/null +++ b/vllm_omni/model_executor/models/auk/pipeline.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK pipeline topology (frozen). + +Stage 0: Encoder — frozen Qwen2.5-Omni thinker, prefill only, emits the + learned layer fusion as ``multimodal_outputs["hidden_states"]["output"]`` +Stage 1: DiT — rectified-flow transformer plus the BigVGAN-flow VAE, + emits a 24 kHz waveform +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.auk" + +AUK_PIPELINE = PipelineConfig( + model_type="auk", + default_deploy_config_name="auk.yaml", + model_arch="AuKForConditionalGeneration", + hf_architectures=("AuKForConditionalGeneration",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="encoder", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=False, + owns_tokenizer=True, + requires_multimodal_data=True, + hf_config_name="thinker_config", + engine_output_type="latent", + model_arch="AuKForConditionalGeneration", + # No model_subdir: tools/prepare_auk_checkpoint.py assembles one + # flat directory, so the encoder finds the thinker shards, the + # tokenizer, the audio processor and the layer-fusion tensors in + # auk.safetensors all under the pipeline root. + # One prefill step produces the text condition; the sampled token + # is discarded, and detokenizing it would only cost latency. + sampling_constraints={"max_tokens": 1, "temperature": 0.0, "detokenize": False}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(0,), + requires_multimodal_data=True, + final_output=True, + final_output_type="audio", + model_arch="AuKPipeline", + custom_process_input_func=f"{_PROC}.encoder2dit", + omni_kv_config={"need_recv_cache": False}, + # Single replica, and the whole ODE runs inside one forward, so + # the stage stays in the orchestrator process. + inline_diffusion=True, + ), + ), +) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index b9373e5f2..febda68af 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -87,6 +87,11 @@ _OMNI_MODELS = { "cosyvoice3", "CosyVoice3Model", ), + "AuKForConditionalGeneration": ( + "auk", + "auk", + "AuKForConditionalGeneration", + ), "NemotronDenseForCausalLM": ( "audex", "audex_thinker", diff --git a/vllm_omni/model_executor/stage_input_processors/auk.py b/vllm_omni/model_executor/stage_input_processors/auk.py new file mode 100644 index 000000000..7c417aee1 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/auk.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Stage input processor for AuK: thinker encoder -> diffusion transition.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch +from vllm.inputs import TextPrompt +from vllm.logger import init_logger + +from vllm_omni.inputs.data import OmniTokensPrompt + +logger = init_logger(__name__) + +# The encoder emits the fused condition as +# ``multimodal_outputs["hidden_states"]["output"]`` (``HiddenStatesStruct.output``, +# the same stage-wire slot MiniMax H3 uses). The payload flattener may deliver +# it under the dotted spelling instead, and a plain ``text_cond`` key is +# accepted so an encoder that emits the condition directly still works. +HIDDEN_STATES_KEY = "hidden_states" +TEXT_COND_KEY = "output" +_FALLBACK_KEYS = ("hidden_states.output", "text_cond") + +# Per-request knobs the prompt carries for stage 1. nfe/cfg/seed are not here: +# they ride the stage-1 OmniDiffusionSamplingParams instead. +_AUK_KNOBS = ("gen_seconds", "sway", "t_grid", "vae_sample") + + +def _as_dict(prompt: Any) -> dict[str, Any]: + """Normalize the many prompt spellings the orchestrator may pass.""" + if isinstance(prompt, list): + prompt = prompt[0] if prompt else {} + if prompt is None: + return {} + if isinstance(prompt, dict): + return prompt + if hasattr(prompt, "_asdict"): + return prompt._asdict() + if hasattr(prompt, "__dict__"): + return vars(prompt) + return {} + + +def _has_audio(mm_data: Any) -> bool: + return isinstance(mm_data, Mapping) and mm_data.get("audio") is not None + + +def _from_hidden_states(hidden: Any) -> Any: + """Read the fused condition out of the payload's hidden-states slot.""" + if isinstance(hidden, Mapping): + return hidden.get(TEXT_COND_KEY) + # A deserialized payload arrives as a msgspec struct rather than a dict. + return getattr(hidden, TEXT_COND_KEY, None) + + +def _extract_text_cond(ar_output: Any) -> Any: + """Read the fused text condition from the encoder output. + + ``multimodal_output`` is attached to the RequestOutput; the CompletionOutput + is checked as a fallback for vLLM versions that attach it there. + """ + for holder in (ar_output, *(getattr(ar_output, "outputs", None) or ())): + mm_output = getattr(holder, "multimodal_output", None) + if not isinstance(mm_output, Mapping): + continue + text_cond = _from_hidden_states(mm_output.get(HIDDEN_STATES_KEY)) + if text_cond is not None: + return text_cond + for key in _FALLBACK_KEYS: + if mm_output.get(key) is not None: + return mm_output[key] + return None + + +def _to_cpu_tensor(text_cond: Any) -> torch.Tensor: + """Coerce the emitted condition to a single 2-D CPU tensor.""" + if isinstance(text_cond, list): + if not text_cond: + raise ValueError("AuK encoder emitted an empty text condition list") + parts = [_to_cpu_tensor(part) for part in text_cond] + return torch.cat(parts, dim=0) if len(parts) > 1 else parts[0] + if isinstance(text_cond, torch.Tensor): + tensor = text_cond + else: + # numpy array or any array-like the connector round-tripped. + tensor = torch.as_tensor(text_cond) + tensor = tensor.detach().cpu() + if tensor.ndim != 2: + raise ValueError(f"AuK text condition must be [tokens, hidden]; got shape {tuple(tensor.shape)}") + return tensor + + +def encoder2dit( + source_outputs: list[Any], + prompt: OmniTokensPrompt | TextPrompt | list | None = None, + requires_multimodal_data: bool = False, + streaming_context: Any | None = None, +) -> dict[str, Any] | None: + """Turn the AuK encoder stage's fused hidden states into a diffusion prompt. + + The source clip is not decoded here: the original ``multi_modal_data`` is + forwarded so stage 1 resamples and VAE-encodes it itself. It is forwarded + unconditionally because the editing and enhancement tasks cannot run + without it; ``requires_multimodal_data`` only records the stage's + declaration. + """ + del streaming_context, requires_multimodal_data + if not source_outputs: + return None + + ar_output = source_outputs[0] + text_cond = _extract_text_cond(ar_output) + if text_cond is None: + raise ValueError( + "AuK encoder stage produced no hidden_states.output multimodal payload; " + "stage 1 cannot run without the fused thinker hidden states" + ) + prompt_embeds = _to_cpu_tensor(text_cond) + + original = _as_dict(prompt) + original_knobs = original.get("additional_information") or {} + if isinstance(original_knobs, Mapping): + original_knobs = original_knobs.get("auk") or {} + if not isinstance(original_knobs, Mapping): + original_knobs = {} + + mm_data = original.get("multi_modal_data") + knobs: dict[str, Any] = {key: original_knobs.get(key) for key in _AUK_KNOBS} + declared_audio = original_knobs.get("has_audio") + knobs["has_audio"] = bool(declared_audio) if declared_audio is not None else _has_audio(mm_data) + + logger.debug( + "[encoder2dit] text_cond=%s dtype=%s has_audio=%s gen_seconds=%s", + tuple(prompt_embeds.shape), + prompt_embeds.dtype, + knobs["has_audio"], + knobs["gen_seconds"], + ) + + return { + "prompt": "", + "prompt_embeds": prompt_embeds, + "multi_modal_data": mm_data, + "additional_information": {"auk": knobs}, + } diff --git a/vllm_omni/model_extras/auk.py b/vllm_omni/model_extras/auk.py new file mode 100644 index 000000000..16d6af52a --- /dev/null +++ b/vllm_omni/model_extras/auk.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Request helpers for AuK: prompt rendering, sampling params, duration policy. + +AuK's contract is one instruction plus an optional reference clip in, audio +out; the task (zero-shot TTS, instruct TTS, content or acoustic editing, +enhancement, separation) is carried by the instruction text alone, so there is +no per-task request surface. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, Any + +from vllm import SamplingParams + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +if TYPE_CHECKING: + import numpy as np + import torch + +# The upstream ChatML render of a single user turn through +# ``Qwen2_5OmniProcessor.apply_chat_template(..., add_generation_prompt=True)`` +# with the template's default system turn. +_CHATML = ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + "{instruction}{audio_part}<|im_end|>\n" + "<|im_start|>assistant\n" +) +_AUDIO_PART = "<|audio_bos|><|AUDIO|><|audio_eos|>" + +# Literal marker the reference appends to the instruction when no reference +# clip is supplied. The model was trained with it, so it is not decorative. +NO_PROMPT_AUDIO = "||" + +# Latent frame rate: 24 kHz waveform, VAE downsample 480 -> 50 frames/s. +SAMPLE_RATE = 24000 +HOP = 480 + +DEFAULT_NFE = 32 +DEFAULT_CFG = 2.0 +DEFAULT_SWAY = -1.0 + + +def auk_prompt( + instruction: str, + audio: tuple[np.ndarray | torch.Tensor, int] | None = None, + *, + gen_seconds: float | None = None, + sway: float = DEFAULT_SWAY, + t_grid: list[float] | None = None, + vae_sample: bool = False, +) -> dict[str, Any]: + """Build the stage-0 prompt for one AuK request. + + Args: + instruction: The task instruction; for text-only requests the + no-reference marker is appended if it is not already there. + audio: Reference or source clip as ``(waveform, sample_rate)``. + ``None`` selects the text-only (instruct TTS) path. + gen_seconds: Target duration. ``None`` reuses the source clip's + length and is an error for text-only requests. + sway: Sway coefficient for the Euler time grid. + t_grid: Explicit time grid, overriding ``sway``. + vae_sample: Draw the reference latent from the VAE posterior instead + of taking its mean. + + Returns: + A prompt dict for ``Omni.generate``. + """ + has_audio = audio is not None and not (isinstance(audio, (tuple, list)) and (len(audio) == 0 or audio[0] is None)) + if not has_audio and not instruction.endswith(NO_PROMPT_AUDIO): + instruction = instruction + NO_PROMPT_AUDIO + + prompt: dict[str, Any] = { + "prompt": _CHATML.format( + instruction=instruction, + audio_part=_AUDIO_PART if has_audio else "", + ), + "additional_information": { + "auk": { + "gen_seconds": gen_seconds, + "sway": sway, + "t_grid": list(t_grid) if t_grid is not None else None, + "vae_sample": vae_sample, + "has_audio": has_audio, + } + }, + } + if has_audio: + prompt["multi_modal_data"] = {"audio": audio} + return prompt + + +def auk_sampling_params( + *, + nfe: int = DEFAULT_NFE, + cfg: float = DEFAULT_CFG, + seed: int | None = 0, +) -> list[SamplingParams | OmniDiffusionSamplingParams]: + """Build the per-stage sampling params for one AuK request. + + Stage 0 is a prefill-only encoder, so its vLLM sampling is a no-op that + only has to satisfy the scheduler. Stage 1 carries the real knobs: ``nfe`` + Euler steps and CFG strength. AuK-Flash ignores both and pins its + distilled recipe. + """ + return [ + SamplingParams(max_tokens=1, temperature=0.0, detokenize=False, seed=seed), + OmniDiffusionSamplingParams(num_inference_steps=nfe, guidance_scale=cfg, seed=seed), + ] + + +def resolve_gen_frames( + gen_seconds: float | None, + ref_frames: int, + *, + sample_rate: int = SAMPLE_RATE, + hop: int = HOP, +) -> int: + """Resolve the target latent length in frames. + + An explicit duration wins; otherwise the target is as long as the source + clip. A text-only request with no duration has nothing to generate, which + is an error rather than a silent one-frame output. + """ + if gen_seconds is not None: + if gen_seconds <= 0: + raise ValueError(f"gen_seconds must be positive; got {gen_seconds}") + return max(1, math.ceil(gen_seconds * sample_rate / hop)) + if ref_frames > 0: + return ref_frames + raise ValueError("gen_seconds is required when the request carries no source audio") + + +__all__ = [ + "NO_PROMPT_AUDIO", + "auk_prompt", + "auk_sampling_params", + "resolve_gen_frames", +] diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 6beb91018..fdbfb4b60 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -34,11 +34,13 @@ _CLASS_TO_MODULE: dict[str, str] = { "MingFlashOmniConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", "Qwen3VLMoeVisionConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", "WhisperEncoderConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", + "AuKConfig": "vllm_omni.transformers_utils.configs.auk", "SenseNovaU1Config": "vllm_omni.transformers_utils.configs.sensenova_u1", } __all__ = [ "Audio8TTSConfig", + "AuKConfig", "Audio8TTSFastARConfig", "Audio8TTSSlowARConfig", "HiggsAudioV3Config", @@ -82,6 +84,7 @@ def __dir__(): # Eagerly import all config modules so their AutoConfig.register() side-effects # run as soon as `vllm_omni.transformers_utils.configs` is imported. from vllm_omni.transformers_utils.configs import audio8_tts as _audio8_tts # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import auk as _auk # noqa: F401, E402 from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import dots_tts as _dots_tts # noqa: F401, E402 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/auk.py b/vllm_omni/transformers_utils/configs/auk.py new file mode 100644 index 000000000..d1e8fa805 --- /dev/null +++ b/vllm_omni/transformers_utils/configs/auk.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Config for the AuK instruction-driven audio generation model. + +The upstream checkpoint ships ``config.yaml`` plus loose safetensors, so an +assembled directory is produced by ``tools/prepare_auk_checkpoint.py``: the +Qwen2.5-Omni-3B ``config.json`` with this ``model_type``, AuK's architecture, +and the five sections below added at the top level. + +Subclassing ``Qwen2_5OmniConfig`` is load-bearing, not stylistic. The encoder +stage reuses the in-tree Qwen2.5-Omni thinker processor, whose +``get_hf_config`` resolves ``self.ctx.get_hf_config(Qwen2_5OmniConfig)`` and +asserts the type, so a bare ``PretrainedConfig`` subclass would fail before +the first request. Inheriting also leaves the parent's nested +``thinker_config`` / ``talker_config`` / ``token2wav_config`` handling and its +``get_text_config`` in place, which is what the encoder stage reads. +""" + +from typing import Any + +from transformers import AutoConfig +from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniConfig + +# The distilled AuK-Flash checkpoint pins a 4-step time grid recorded in its +# metadata. The values are exact and must never be rounded: the student was +# trained on them, and a rounded grid changes the output amplitude. +AUK_FLASH_T_GRID: tuple[float, ...] = ( + 0.0, + 0.07612049579620361, + 0.2928932309150696, + 0.6173166036605835, + 1.0, +) + +DEFAULT_DIT: dict[str, Any] = { + "dim": 1536, + "heads": 24, + "dim_head": 64, + "ff_mult": 2, + "text_hidden_dim": 2048, + "num_layers": 10, + "num_single_layers": 20, + "attn_mask_enabled": True, +} + +DEFAULT_VAE: dict[str, Any] = { + "latent_dim": 64, + "downsample_rate": 480, + "target_sample_rate": 24000, + "model_init_kwargs": { + "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, + }, +} + +# Base AuK sampling recipe; AuK-Flash replaces nfe and cfg with its own. +DEFAULT_SAMPLING: dict[str, Any] = {"nfe": 32, "cfg": 2.0, "sway": -1.0} + +VARIANTS = ("base", "flash") + + +class AuKConfig(Qwen2_5OmniConfig): + """Assembled AuK checkpoint config: a Qwen2.5-Omni config plus DiT and VAE.""" + + model_type = "auk" + + def __init__( + self, + dit: dict[str, Any] | None = None, + vae: dict[str, Any] | None = None, + variant: str = "base", + flash_t_grid: list[float] | None = None, + defaults: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + if variant not in VARIANTS: + raise ValueError(f"AuK variant must be one of {VARIANTS}; got {variant!r}") + + # Assigned before the parent's init so nothing it runs during + # validation sees a half-built config. + self.dit = {**DEFAULT_DIT, **(dit or {})} + self.vae = {**DEFAULT_VAE, **(vae or {})} + self.variant = variant + self.flash_t_grid = list(flash_t_grid) if flash_t_grid is not None else list(AUK_FLASH_T_GRID) + self.defaults = {**DEFAULT_SAMPLING, **(defaults or {})} + + kwargs.setdefault("architectures", ["AuKForConditionalGeneration"]) + # The parent owns thinker_config / talker_config / token2wav_config and + # get_text_config; nothing here overrides them. + super().__init__(**kwargs) + + @property + def is_flash(self) -> bool: + """Whether this checkpoint is the distilled 4-step student.""" + return getattr(self, "variant", "base") == "flash" + + @property + def latent_dim(self) -> int: + return int(self.vae["latent_dim"]) + + @property + def downsample_rate(self) -> int: + return int(self.vae["downsample_rate"]) + + @property + def target_sample_rate(self) -> int: + return int(self.vae["target_sample_rate"]) + + +AutoConfig.register("auk", AuKConfig) + +__all__ = [ + "AUK_FLASH_T_GRID", + "DEFAULT_DIT", + "DEFAULT_SAMPLING", + "DEFAULT_VAE", + "VARIANTS", + "AuKConfig", +]