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