From 5f95115e703ffe51e03e40ae23c78ff841c3335e Mon Sep 17 00:00:00 2001 From: Nikita Osterov <90456300+osipovmekete@users.noreply.github.com> Date: Tue, 6 Oct 2026 10:55:15 +0300 Subject: [PATCH] Kandinsky6 (#8537) Signed-off-by: leffff Signed-off-by: Artem Borisov Co-authored-by: leffff Co-authored-by: Artem Borisov --- .../common/ci_source_file_dependencies.yml | 9 + .buildkite/cuda/test-merge.yml | 9 + docs/models/supported_models.md | 1 + docs/user_guide/diffusion_features.md | 1 + .../image_to_video/README.md | 3 +- .../text_to_video/text_to_video.md | 19 + .../text_to_video/text_to_video.py | 2 +- pyproject.toml | 12 + recipes/Kandinsky/Kandinsky6-TI2VA.md | 247 ++ recipes/README.md | 1 + tests/diffusion/cache/test_cache_dit.py | 1 + tests/diffusion/models/kandinsky6/__init__.py | 0 .../models/kandinsky6/test_acceleration.py | 89 + .../kandinsky6/test_pipeline_kandinsky6.py | 246 ++ .../kandinsky6/test_transformer_kandinsky6.py | 325 +++ .../kandinsky6/test_vae_patch_parallel.py | 118 + .../e2e/offline_inference/test_kandinsky6.py | 79 + tests/e2e/online_serving/test_kandinsky6.py | 89 + .../openai_api/test_video_server.py | 32 + tests/model_tests/diffusion/test_alignment.py | 3 + .../cache/cachedit/model_specific.py | 30 + vllm_omni/diffusion/cache/magcache/backend.py | 15 +- vllm_omni/diffusion/cache/magcache/hook.py | 5 +- vllm_omni/diffusion/cache/teacache/backend.py | 21 +- vllm_omni/diffusion/cache/teacache/config.py | 11 +- vllm_omni/diffusion/model_metadata.py | 9 + .../diffusion/models/kandinsky6/__init__.py | 22 + .../models/kandinsky6/cache_accel.py | 199 ++ .../kandinsky6/kandinsky6_transformer.py | 2006 +++++++++++++++++ .../kandinsky6/modeling_kandinsky6_audio.py | 1753 ++++++++++++++ .../kandinsky6/modeling_kandinsky6_vae.py | 1189 ++++++++++ .../models/kandinsky6/pipeline_kandinsky6.py | 1851 +++++++++++++++ .../kandinsky6/scheduling_kandinsky6.py | 103 + vllm_omni/diffusion/registry.py | 7 + vllm_omni/diffusion/utils/hf_utils.py | 8 +- vllm_omni/entrypoints/openai/serving_video.py | 7 +- vllm_omni/model_extras/kandinsky6.py | 48 + vllm_omni/model_extras/registry.py | 12 + 38 files changed, 8574 insertions(+), 8 deletions(-) create mode 100644 recipes/Kandinsky/Kandinsky6-TI2VA.md create mode 100644 tests/diffusion/models/kandinsky6/__init__.py create mode 100644 tests/diffusion/models/kandinsky6/test_acceleration.py create mode 100644 tests/diffusion/models/kandinsky6/test_pipeline_kandinsky6.py create mode 100644 tests/diffusion/models/kandinsky6/test_transformer_kandinsky6.py create mode 100644 tests/diffusion/models/kandinsky6/test_vae_patch_parallel.py create mode 100644 tests/e2e/offline_inference/test_kandinsky6.py create mode 100644 tests/e2e/online_serving/test_kandinsky6.py create mode 100644 vllm_omni/diffusion/models/kandinsky6/__init__.py create mode 100644 vllm_omni/diffusion/models/kandinsky6/cache_accel.py create mode 100644 vllm_omni/diffusion/models/kandinsky6/kandinsky6_transformer.py create mode 100644 vllm_omni/diffusion/models/kandinsky6/modeling_kandinsky6_audio.py create mode 100644 vllm_omni/diffusion/models/kandinsky6/modeling_kandinsky6_vae.py create mode 100644 vllm_omni/diffusion/models/kandinsky6/pipeline_kandinsky6.py create mode 100644 vllm_omni/diffusion/models/kandinsky6/scheduling_kandinsky6.py create mode 100644 vllm_omni/model_extras/kandinsky6.py diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index b34e9d271..d7068cfca 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -149,6 +149,10 @@ x-lingbot-video: &lingbot_video x-lingbot-world: &lingbot_world - vllm_omni/diffusion/models/lingbot_world/ +x-kandinsky6: &kandinsky6 + - vllm_omni/diffusion/models/kandinsky6/ + - vllm_omni/model_extras/kandinsky6.py + source_file_dependencies: # Hardware presets and shared test helpers. Not a job key — do not attach this @@ -491,6 +495,11 @@ source_file_dependencies: - tests/dfx/perf/scripts/run_benchmark.py - tests/dfx/perf/tests/test_lingbot_video_vllm_omni.json + diffusion_kandinsky6_function: + - *kandinsky6 + - tests/e2e/offline_inference/test_kandinsky6.py + - tests/e2e/online_serving/test_kandinsky6.py + # Invalid-param jobs: model business code only, plus the reliability scripts. reliability_invalid_param_h100: diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 021faaef5..3b002d9ff 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -189,6 +189,15 @@ steps: ' mirror_hardwares: h100_1 + - label: "Diffusion · Kandinsky6 Test" + source_file_dependencies: diffusion_kandinsky6_function + commands: + - | + timeout 40m bash -c " + pytest -s -v tests/e2e/offline_inference/test_kandinsky6.py tests/e2e/online_serving/test_kandinsky6.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' + " + mirror_hardwares: h100_1 + - label: "Diffusion · Wan22 Test" source_file_dependencies: diffusion_wan22_function commands: diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 5d2178cc2..f5d3034a5 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -106,6 +106,7 @@ th { | `SenseNovaU1Pipeline` | SenseNova-U1-A3B (MoE) | `sensenova/SenseNova-U1-A3B-MoT` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.md#a3b-moe-variant) | | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — | | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanVideo-1.5.md) | +| `Kandinsky6TI2VAPipeline` | Kandinsky 6 TI2VA | `kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Kandinsky/Kandinsky6-TI2VA.md) | | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanVideo-1.5.md) | | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — | | `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 7bbb08951..2a2f54aa9 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -196,6 +196,7 @@ The following tables show which models support each feature: | **SANA-Video-2B T2V I2V** | ❌ | ❌ | ✅ (`--usp`, frame-sharded) | ✅ | ✅ (TP=2 only) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | +| **Kandinsky 6 TI2VA** | ✅* (uncalibrated) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ (FP8) | ✅ (1 request) | | **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **LongCat-Video-Avatar-1.5** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **MiniMax-H3** | ✅ (FL2VA) | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ | diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md index 9386264b5..19dd31d5e 100644 --- a/examples/offline_inference/image_to_video/README.md +++ b/examples/offline_inference/image_to_video/README.md @@ -1,7 +1,7 @@ # Image-To-Video This shared example generates videos from images with VACE, Wan2.2, LTX-2, -HunyuanVideo-1.5, SANA-Video, Cosmos3, MAGI-2, and other compatible pipelines. +HunyuanVideo-1.5, SANA-Video, Cosmos3, MAGI-2, Kandinsky 6, and other compatible pipelines. - `image_to_video.py`: command-line script for single video generation with advanced options. @@ -30,6 +30,7 @@ This folder provides a unified CLI script for image-to-video generation using vL | `sand-ai/MAGI-2-preview` | 512 x 896 | 125 | 100 | Model-fixed | Native four-GPU TP/SP; resident SP4 default; DLO available | | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480 x 832 | 81 | 50 | 6.0 | Native `SanaImageToVideoPipeline`; Wan VAE | | `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704 x 1280 | 81 | 50 | 6.0 | Native `SanaImageToVideoPipeline`; LTX-2 Video VAE | +| `kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers` | 480 x 864 | 125 | 50 | 5.0 | ~75 GiB reserved (H100, CPU offload); gated Hub repo | MAGI-2 native Preview setup, four-GPU topology and DLO choices, I2VA commands, and eight-GPU validation status are documented in the diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index f051d8f88..cdbbac6c1 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -14,6 +14,7 @@ A unified script for text-to-video generation. Supports multiple models with mod | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | | `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) | | `BestWishYsh/Helios-Base` / `Helios-Mid` / `Helios-Distilled` | 384x640 | 99 | 50 | 5.0 / 5.0 / 1.0 | — | +| `kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers` | 480x864 | 125 | 50 | 5.0 | ~75 GiB reserved (H100, CPU offload, SDPA) | | `sand-ai/MAGI-2-preview` | 512x896 | 125 | 100 | Model-fixed | Native four-GPU TP/SP; resident SP4 default; DLO available | | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480x832 | 81 | 50 | 6.0 | BF16 DiT + FP32 Wan VAE wrapper | | `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704x1280 | 81 | 50 | 6.0 | BF16 DiT + LTX-2 Video VAE wrapper | @@ -249,6 +250,24 @@ python text_to_video.py \ > conditioning tensors that cannot be passed through the JSON `--extra-body` > flag; they are out of scope for this text-to-video example. +### Kandinsky 6 (TI2VA) + +The Hub checkpoint is gated; export `HF_TOKEN` if download returns 401. +`--enable-cpu-offload` is required on an 80 GB H100. Audio is on by default. + +```bash +python text_to_video.py \ + --model kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers \ + --prompt "A golden retriever runs along a sunny beach, waves crashing, cinematic footage" \ + --enable-cpu-offload \ + --output kandinsky6_t2va.mp4 +``` + +Video only: add `--extra-body '{"sample_audio": false}'`. Image-to-video uses +the shared [`image_to_video.py`](../image_to_video/README.md) example with +`--image`. Serving numbers and flags are in the +[`Kandinsky 6 TI2VA recipe`](../../../recipes/Kandinsky/Kandinsky6-TI2VA.md). + ## Key Arguments ### Common diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 7b7bcb890..40eac3bfa 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -235,7 +235,7 @@ def parse_extra_body(value: str) -> dict[str, Any]: def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Generate a video from a text prompt. " - "Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, MAGI-2, " + "Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, MAGI-2, Kandinsky 6, " "and other text-to-video models." ) parser.add_argument( diff --git a/pyproject.toml b/pyproject.toml index c65255f99..77ef14ea0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -232,6 +232,11 @@ ignore = [ # F-correctness rules still apply. "vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/**" = ["E501", "E721", "E741", "F401", "F841", "I001", "N803", "N814", "UP006", "UP007", "UP008", "UP011", "UP032", "UP035", "W291", "W293"] "vllm_omni/experimental/fullduplex/joyvl/decision/prompts.py" = ["E501"] # System/summarizer prompts are long prose lines +# Kandinsky 6 keeps native/BigVGAN names (T/H/W, `l`) and NVIDIA vocoder +# docstring length so the port stays diffable against the source modules. +"vllm_omni/diffusion/models/kandinsky6/kandinsky6_transformer.py" = ["N803", "E741"] +"vllm_omni/diffusion/models/kandinsky6/pipeline_kandinsky6.py" = ["N803"] +"vllm_omni/diffusion/models/kandinsky6/modeling_kandinsky6_audio.py" = ["E501", "E741", "N803", "UP008"] [tool.ruff.lint.flake8-tidy-imports.banned-api] "librosa".msg = "The librosa module is banned, use vllm.multimodal helpers instead" @@ -364,6 +369,8 @@ extend-ignore-identifiers-re = [ # (output feature dim), not "doubt". ".*dout.*", ".*DOUT.*", + # Kandinsky nabla threshold argument (`thr`), not the word "the". + "^thr$", ] extend-ignore-re = [ "\\.\\.\\.n,d->\\.\\.\\.nd", @@ -404,3 +411,8 @@ IY = "IY" # Speech DAC codec identifiers (convnet_type, transconvnet_type) and is not a # typo of "convent". convnet = "convnet" +# Upstream Kandinsky system prompt. "promt" and "scren" are part of the +# conditioning text and must not be corrected. +thr = "thr" +promt = "promt" +scren = "scren" diff --git a/recipes/Kandinsky/Kandinsky6-TI2VA.md b/recipes/Kandinsky/Kandinsky6-TI2VA.md new file mode 100644 index 000000000..b4e0e3a09 --- /dev/null +++ b/recipes/Kandinsky/Kandinsky6-TI2VA.md @@ -0,0 +1,247 @@ +# Kandinsky 6 TI2VA — H100 (v1, single GPU, CPU offload) + +> Text/image-to-video-and-audio serving with Kandinsky 6 Pro + +## Summary + +- Vendor: Kandinsky Lab ([`kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers)) +- Model: Kandinsky 6 Pro (TI2VA — text/image-to-video-and-audio) +- Task: Joint text/image-to-video-and-audio generation +- Mode: Offline (`Omni(...)`) and online serving with the OpenAI-compatible + `/v1/videos` API +- Hardware: 1x NVIDIA H100 80GB HBM3 with component-level CPU offload +- Maintainer: Community + +## When to use this recipe + +Use this recipe to serve Kandinsky 6's TI2VA pipeline — a single DiT that +jointly denoises video and audio latents, conditioned on a text prompt and, +optionally, a reference image (image-to-video). Tensor parallel, CFG +parallel, Ulysses, ring, pipeline parallel, HSDP, layerwise offload, +MagCache, TeaCache, Cache-DiT, FP8, and tiled VAE patch parallel are wired. +NaviCache and a registered NABLA backend are not. TeaCache coefficients are +uncalibrated. See "Known limitations". + +## Supported model contract + +| Task | Entrypoint | Input | Output | +|---|---|---|---| +| Text-to-video-and-audio | `vllm serve kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers --omni` / `Omni(model="kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers")` | text prompt | synchronized MP4 (H.264 video + AAC 44.1 kHz mono audio) | +| Image-to-video-and-audio | same | text prompt + reference image | synchronized MP4, first frame conditioned on the reference image | + +- Audio generation is on by default (`sample_audio=True`); pass + `extra_args.sample_audio=false` to generate video only. +- Batch size is 1 request per generation call (upstream limitation of the + ported denoise loop; concurrent requests are still served by vLLM-Omni's + own request scheduler, just not batched together within one DiT forward + pass yet). +- Serving defaults follow the Pro production geometry registered in + `vllm_omni/model_extras/kandinsky6.py`: **864x480, 125 frames @ 24 fps + (5.2 s), 50 steps, CFG 5.0**. Any `4k+1` frame count and 16-divisible + resolution can be requested explicitly via `height`/`width`/`num_frames`/ + `num_inference_steps` — the model is *not* fixed-duration. +- Audio is always 44.1 kHz; its length is derived from the video duration. +- Negative prompt defaults to the pipeline's built-in K6 negative prompt + when omitted. + +## References + +- Hub checkpoint: [`kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers) +- Joint video+audio native-model reference: [`recipes/MiniMaxAI/MiniMax-H3.md`](../MiniMaxAI/MiniMax-H3.md) +- Offline example: [`examples/offline_inference/text_to_video/text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py) +- Image-to-video example: [`examples/offline_inference/image_to_video/image_to_video.py`](../../examples/offline_inference/image_to_video/image_to_video.py) + +## Hardware + +- Accelerator model and per-device memory: NVIDIA H100 80GB HBM3 + (driver 570.133.20) +- Number of devices: 1 (multi-GPU tensor parallelism is implemented in the + transformer's linear layers but not yet validated end-to-end) +- Device interconnect: N/A (single-device path) +- Host memory: 1.4 TB installed; the bf16 DiT (~56 GB) plus the text + encoders are staged on the host under `--enable-cpu-offload`, so budget + at least ~90 GB of free host RAM. +- Qualification scope: single-GPU T2VA at the Pro default geometry and at a + reduced smoke geometry, offline and via `vllm serve`. Numbers below were + measured once on the shared box described above; treat them as + indicative, not as a benchmark. + +## Software environment + +- OS: Ubuntu 22.04.5 LTS +- Python: 3.12.13 +- Driver / runtime: CUDA 12.9 (torch 2.13.0+cu129) +- vLLM version: 0.29.0 +- vLLM-Omni version or commit: this repository, `vllm_omni/diffusion/models/kandinsky6/` +- diffusers 0.40.0, transformers 5.14.1 +- Attention: FlashAttention-3 on this box (`flash_attn_interface`). The Pro + timings below use it. Without FA3, `attention_engine: "auto"` falls back + to PyTorch SDPA; an earlier SDPA Pro run on the same GPU was 26.9 s/it + and 1377 s. + +## Weights + +`Kandinsky6TI2VAPipeline._load_components` loads the public Diffusers layout +from [`kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers) +(`model_index.json`, `transformer/`, `vae/`, `text_encoder/`, `tokenizer/`, +`text_encoder_2/`, `tokenizer_2/`, `audio_vae/`, `scheduler/`). Pass that +repo id (or a local snapshot of it) to `vllm serve` / `Omni(model=...)`. +The Hub repo is gated; export `HF_TOKEN` if download returns 401. + +A raw Kandinsky 6 weight tree is not required. If you already have a local +Diffusers-layout snapshot, point `--model` at that directory instead of the +Hub id. + +Each weight folder is loaded on its own (`transformer/`, `vae/`, `text_encoder/`, `text_encoder_2/`, `audio_vae/`). The DiT is built with initialization skipped, then the framework loader fills it from `transformer/`. + +## Command + +```bash +vllm serve kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers --omni \ + --host 127.0.0.1 --port 8091 \ + --num-gpus 1 --enable-cpu-offload +``` + +`--enable-cpu-offload` enables model-level offload with mutual exclusion +between the transformer and the two text encoders (they never co-reside on +the GPU). Without it the bf16 DiT (~56 GB) plus Qwen2.5-VL-7B (~16 GB) plus +the VAEs and activations do not fit an 80 GB device at the Pro geometry. + +Offline (writes an MP4 with the audio track muxed in): + +```bash +python examples/offline_inference/text_to_video/text_to_video.py \ + --model kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers \ + --prompt "A golden retriever runs along a sunny beach, waves crashing, cinematic footage" \ + --seed 42 --enable-cpu-offload --output kandinsky6_t2va.mp4 +# Video only: add --extra-body '{"sample_audio": false}' +``` + +Image-to-video-and-audio uses the shared image-to-video example: + +```bash +python examples/offline_inference/image_to_video/image_to_video.py \ + --model kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers \ + --image first_frame.png \ + --prompt "A golden retriever runs along a sunny beach, waves crashing, cinematic footage" \ + --seed 42 --enable-cpu-offload --output kandinsky6_i2va.mp4 +``` + +The generic scripts pick up the Pro defaults from +`vllm_omni/model_extras/kandinsky6.py`. + +## Verification + +`/v1/videos` takes multipart form fields (not a JSON body). Submit, poll, +then download: + +```bash +VID=$(curl -s -X POST http://127.0.0.1:8091/v1/videos \ + -F prompt="A golden retriever runs along a sunny beach, waves crashing" \ + -F size=864x480 -F num_frames=125 -F num_inference_steps=50 -F seed=42 \ + | python3 -c "import sys,json;print(json.load(sys.stdin)['id'])") + +# Poll until "status": "completed" +curl -s http://127.0.0.1:8091/v1/videos/$VID + +curl -s -o kandinsky6.mp4 http://127.0.0.1:8091/v1/videos/$VID/content +# Expect an MP4 with an H.264 video stream (125 frames, 24 fps, 864x480) +# and an AAC audio stream (44100 Hz, mono) of matching duration. +``` + +Smoke-size request for quick validation (~30 s end-to-end on one H100): +`-F size=512x320 -F num_frames=25 -F num_inference_steps=10`. + +## Measurements + +Single H100 80GB, seed 42, prompt +"A golden retriever runs along a sunny beach, waves crashing, cinematic footage", +negative prompt left at the pipeline default, MagCache off, prompt expansion +off. Pro geometry: 864x480, 125 frames, 50 steps, CFG 5.0, audio on. + +vLLM-Omni: `text_to_video.py --model kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers --enable-cpu-offload --model-class-name Kandinsky6TI2VAPipeline`, +FlashAttention-3. + +| Wall / generation | Step time | Peak GPU | After load | +|---|---|---|---| +| 751.7 s generation (12.5 min); engine init 41 s | 14.4 s/it average | 78.39 GiB reserved (78.0 GiB nvidia-smi) | 17.74 GiB | + +An earlier SDPA vLLM-Omni Pro run on the same GPU was 26.9 s/it and 1377 s, +peak reserved 74.97 GiB. A 512x320 / 25-frame / 10-step SDPA smoke was 1.4 s/it +and 29.5 s at the same reserved peak. + +The MP4 is H.264 864x480, 125 frames at 24 fps (5.22 s), plus AAC 44.1 kHz +mono (230400 samples): a golden retriever on a sunny beach with breaking waves. + +vLLM-Omni's VAE decode logged two failed CUDA allocations (7.15 GiB, then +4.27 GiB) and still wrote a complete file. + +Audio is peak-normalized to full scale by default; pass +`extra_args.audio_normalization="clip"` to keep the raw decoded amplitude +instead. + +## Notes + +- Key flags: `extra_args.sample_audio` (bool, default true), + `extra_args.audio_normalization` (`"normalize"` default | `"clip"`), + `extra_args.visual_cond_scheme` (defaults to `pretrain` for text-only, + `tail_cond_first_frame` when an image is supplied). +- The post-process payload is a flat dict (`video`, `audio` as float32 in + `[-1, 1]`, `audio_sample_rate=44100`, `fps`) so the serving muxer picks + up the correct sample rate. +- Acceleration that is wired: + - Tensor parallel: the DiT loader narrows full checkpoint tensors onto + each rank's `ColumnParallelLinear` / `RowParallelLinear` shard. + - CFG parallel (`--cfg-parallel-size 2`): rank 0 runs the conditional + DiT, rank 1 the unconditional one, then both all-gather video and + audio velocities. + - HSDP: `_hsdp_shard_conditions` covers the indexed transformer blocks. + Do not combine with tensor parallel. + - Sequence parallel: visual tokens (and their RoPE) are sharded; + Ulysses all-to-all runs inside visual self-attention, Ring uses the + PyTorch ring kernel. Text stays replicated. Audio queries gather the + full visual sequence. + - Pipeline parallel: `visual_transformer_blocks` are split with + `PPMissingLayer`. Embeddings stay on the first stage, output heads on + the last. Activations and the final velocity are sent between stages. + - VAE patch parallel: Hunyuan decode splits the latent width when + `vae_patch_parallel_size > 1` and tiling is on. + - MagCache / TeaCache / Cache-DiT are registered. MagCache reuses the + Pro `mag_ratios`. TeaCache coefficients are **uncalibrated** (copied + from Qwen-Image) so skips are not quality-neutral. Cache-DiT targets + `visual_transformer_blocks` only and falls back to the same step cache + if the fused `(video, audio)` block return is rejected. + - FP8: online quantization runs in the DiT loader via + `process_weights_after_loading` after the sharded assign. + - `attention_engine: auto` uses FlashAttention-3 when + `flash_attn_interface` imports. +- Known limitations: + - NABLA sparse attention is not registered in the framework attention + backend. NaviCache is not wired. + - Prompt expansion, NF4 Qwen, and the super-resolution cascade are not + part of this port. + - Distributed layerwise offload still uses one process unless another + parallel axis sets `world_size > 1` (for example tensor parallel). + - `_load_components` loads the Diffusers checkpoint directly rather than the + streamed-prefetch loader. + +## Supported features + +| Feature | Status | Guide | +|---|---|---| +| Text-to-video-and-audio | Supported (validated on H100) | — | +| Image-to-video-and-audio | Supported (not yet validated on hardware) | — | +| CPU offload (`--enable-cpu-offload`) | Supported (required on 80 GB) | — | +| Layerwise offload (`--enable-layerwise-offload`) | Supported (about 36 GiB peak on the smoke geometry) | — | +| Tensor parallelism | Wired (sharded DiT load) | — | +| CFG parallelism | Wired (`--cfg-parallel-size 2`) | — | +| Ulysses / Ring sequence parallel | Wired (visual tokens only) | — | +| Pipeline parallelism | Wired (visual block split) | — | +| HSDP | Wired (not with tensor parallel) | — | +| VAE patch parallel | Wired (width-split decode) | — | +| MagCache | Wired (Pro ratios, step reuse) | — | +| TeaCache | Wired, coefficients uncalibrated | — | +| Cache-DiT | Wired on `visual_transformer_blocks` | — | +| FP8 | Wired (online quant; smoke peak about 46 GiB vs 75 GiB bf16) | [`FP8 quantization`](../../docs/user_guide/quantization/fp8.md) | +| NaviCache | Not supported | — | +| LoRA | Not supported | — | diff --git a/recipes/README.md b/recipes/README.md index 9a3837367..e7fe44236 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -64,6 +64,7 @@ recipes/ | [`MiniMaxAI/MiniMax-H3-Spark-GB10.md`](./MiniMaxAI/MiniMax-H3-Spark-GB10.md) | MiniMax H3 T2VA and Ref2VA on DGX Spark | 1x or 2x DGX Spark (GB10) | | [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) | | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB / 1x Intel Arc Pro B70 32GB | +| [`Kandinsky/Kandinsky6-TI2VA.md`](./Kandinsky/Kandinsky6-TI2VA.md) | Text/image-to-video-and-audio, offline and `/v1/videos` | 1x H100 80GB with CPU offload; 2x H100 for parallel modes | | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU / 1x Intel Arc Pro B70 32GB | | [`LTX/LTX-2.5.md`](./LTX/LTX-2.5.md) | LTX-2.5-Diffusers: Full/SFT one-stage and distilled two-stage T2V/I2V with synchronized audio | NVIDIA B300; cuDNN-qualified / 1x Intel Arc Pro B70 32GB | | [`m-a-p/SheetSage2-H200.md`](./m-a-p/SheetSage2-H200.md) | Official audio-to-score preprocessing and YuE2 request export | 1x H200 141GB | diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py index 7de2aaa38..5377356b0 100644 --- a/tests/diffusion/cache/test_cache_dit.py +++ b/tests/diffusion/cache/test_cache_dit.py @@ -59,6 +59,7 @@ def test_custom_cache_dit_enablers_are_registered_explicitly(): "Cosmos3OmniDiffusersPipeline": cd_model_specific.enable_cache_for_cosmos3, "Cosmos3OmniPipeline": cd_model_specific.enable_cache_for_cosmos3, "Krea2Pipeline": cd_model_specific.enable_cache_for_krea2, + "Kandinsky6TI2VAPipeline": cd_model_specific.enable_cache_for_kandinsky6, "Magi2Pipeline": cd_model_specific.enable_cache_for_magi2, "MammothModa2DiTPipeline": cd_model_specific.enable_cache_for_mammothmoda2, } diff --git a/tests/diffusion/models/kandinsky6/__init__.py b/tests/diffusion/models/kandinsky6/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/diffusion/models/kandinsky6/test_acceleration.py b/tests/diffusion/models/kandinsky6/test_acceleration.py new file mode 100644 index 000000000..3ed58f86b --- /dev/null +++ b/tests/diffusion/models/kandinsky6/test_acceleration.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CPU checks that Kandinsky 6 acceleration hooks are actually wired.""" + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.models.kandinsky6.cache_accel import K6_PRO_MAG_RATIOS, K6StepCache +from vllm_omni.diffusion.models.kandinsky6.kandinsky6_transformer import Kandinsky6Transformer3DModel +from vllm_omni.diffusion.models.kandinsky6.pipeline_kandinsky6 import _shard_loaded_weight +from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_shard_loaded_weight_narrows_column_parallel_output(): + class _Owner: + tp_rank = 1 + + def weight_loader(self, param, loaded): + del param, loaded + + param = nn.Parameter(torch.empty(2, 4)) + param.output_dim = 0 + param.weight_loader = _Owner().weight_loader + loaded = torch.arange(16, dtype=torch.float32).reshape(4, 4) + + shard = _shard_loaded_weight(param, loaded) + + assert tuple(shard.shape) == (2, 4) + torch.testing.assert_close(shard, loaded[2:4]) + + +def test_row_bias_is_not_sharded_when_shapes_match(): + class _Owner: + tp_rank = 1 + + def weight_loader(self, param, loaded): + del param, loaded + + param = nn.Parameter(torch.empty(4)) + param.output_dim = 0 + param.weight_loader = _Owner().weight_loader + loaded = torch.arange(4, dtype=torch.float32) + + shard = _shard_loaded_weight(param, loaded) + + torch.testing.assert_close(shard, loaded) + + +def test_hsdp_predicate_matches_visual_blocks_only_when_indexed(): + assert Kandinsky6Transformer3DModel._is_transformer_block("visual_transformer_blocks.3", nn.Identity()) + assert Kandinsky6Transformer3DModel._is_transformer_block("video_text_transformer_blocks.0", nn.Identity()) + assert not Kandinsky6Transformer3DModel._is_transformer_block("visual_embeddings", nn.Identity()) + assert Kandinsky6Transformer3DModel._hsdp_shard_conditions + + +def test_sp_plan_shards_visual_tokens_and_gathers_before_the_head(): + plan = Kandinsky6Transformer3DModel._sp_plan + assert "_sp_visual_shard" in plan + assert "_sp_visual_rope" in plan + assert plan["_sp_visual_gather"].gather_dim == 1 + assert Kandinsky6Transformer3DModel._magcache_block_attrs == ("visual_transformer_blocks",) + + +def test_k6_is_not_on_the_no_cache_list_and_enablers_are_registered(): + import vllm_omni.diffusion.cache.cachedit # noqa: F401 registers enablers + from vllm_omni.diffusion.cache.cachedit.backend import CUSTOM_DIT_ENABLERS + from vllm_omni.diffusion.cache.teacache.backend import CUSTOM_TEACACHE_ENABLERS + + assert "Kandinsky6TI2VAPipeline" not in _NO_CACHE_ACCELERATION + assert "Kandinsky6TI2VAPipeline" in CUSTOM_TEACACHE_ENABLERS + assert "Kandinsky6TI2VAPipeline" in CUSTOM_DIT_ENABLERS + assert len(K6_PRO_MAG_RATIOS) == 100 + + +def test_step_cache_skips_only_after_a_real_step(): + cache = K6StepCache( + kind="mag", + threshold=0.24, + max_skip_steps=3, + retention_ratio=0.0, + mag_ratios=(1.0, 1.0), + ) + assert cache.should_skip(0) is False + cache.store(torch.zeros(1), torch.zeros(1)) + assert cache.should_skip(1) is True diff --git a/tests/diffusion/models/kandinsky6/test_pipeline_kandinsky6.py b/tests/diffusion/models/kandinsky6/test_pipeline_kandinsky6.py new file mode 100644 index 000000000..883beca8f --- /dev/null +++ b/tests/diffusion/models/kandinsky6/test_pipeline_kandinsky6.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import numpy as np +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_kandinsky6_pipeline_import_and_registry(): + from vllm_omni.diffusion.models.kandinsky6 import ( + AutoencoderKLHunyuanVideo, + Kandinsky6AudioVAE, + Kandinsky6TI2VAPipeline, + Kandinsky6Transformer3DModel, + KandinskyFlowMatchScheduler, + get_kandinsky6_post_process_func, + get_kandinsky6_pre_process_func, + ) + from vllm_omni.diffusion.registry import ( + _DIFFUSION_MODELS, + _DIFFUSION_POST_PROCESS_FUNCS, + _DIFFUSION_PRE_PROCESS_FUNCS, + _NO_CACHE_ACCELERATION, + ) + + assert Kandinsky6TI2VAPipeline is not None + assert Kandinsky6Transformer3DModel is not None + assert AutoencoderKLHunyuanVideo is not None + assert Kandinsky6AudioVAE is not None + assert KandinskyFlowMatchScheduler is not None + assert get_kandinsky6_post_process_func is not None + assert get_kandinsky6_pre_process_func is not None + + assert _DIFFUSION_MODELS["Kandinsky6TI2VAPipeline"] == ( + "kandinsky6", + "pipeline_kandinsky6", + "Kandinsky6TI2VAPipeline", + ) + assert _DIFFUSION_POST_PROCESS_FUNCS["Kandinsky6TI2VAPipeline"] == "get_kandinsky6_post_process_func" + assert _DIFFUSION_PRE_PROCESS_FUNCS["Kandinsky6TI2VAPipeline"] == "get_kandinsky6_pre_process_func" + assert "Kandinsky6TI2VAPipeline" not in _NO_CACHE_ACCELERATION + + +def test_kandinsky6_component_discovery_declarations(): + from vllm_omni.diffusion.models.kandinsky6 import Kandinsky6TI2VAPipeline + + assert Kandinsky6TI2VAPipeline._dit_modules == ["transformer"] + assert Kandinsky6TI2VAPipeline._encoder_modules == ["text_encoder", "text_encoder_2"] + assert Kandinsky6TI2VAPipeline._vae_modules == ["vae", "audio_vae"] + assert Kandinsky6TI2VAPipeline.supports_step_execution is True + assert Kandinsky6TI2VAPipeline.support_audio_output is True + assert Kandinsky6TI2VAPipeline.support_image_input is True + + +def test_kandinsky6_pipeline_satisfies_capability_protocols(): + """isinstance checks against the @runtime_checkable protocols the + framework actually uses to detect these capabilities (io_support.py's + supports_audio_output, etc.) — stronger than just checking the class + attribute exists.""" + from vllm_omni.diffusion.models.interface import ( + SupportAudioOutput, + SupportImageInput, + SupportsComponentDiscovery, + ) + from vllm_omni.diffusion.models.kandinsky6 import Kandinsky6TI2VAPipeline + + assert isinstance(Kandinsky6TI2VAPipeline, SupportAudioOutput) + assert isinstance(Kandinsky6TI2VAPipeline, SupportImageInput) + assert isinstance(Kandinsky6TI2VAPipeline, SupportsComponentDiscovery) + + +class _FakeAudioVAE: + downsample_factor = 1024 + + +def test_kandinsky6_post_process_func_packages_video_and_audio(): + """Pure-function test: no model/weights needed. Confirms the post-process + payload shape matches what io_support.py -> output_formatter.py -> + media_utils.py's PyAV muxing expects (the same flat {"video", "audio", + "audio_sample_rate", "fps"} shape MiniMax H3's own post-process function + produces), and that int16 PCM from ``postprocess_audio`` is rescaled to + the float32 [-1, 1] waveform the muxer consumes.""" + from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func + + post_process = get_kandinsky6_post_process_func(od_config=None) + video = np.zeros((1, 4, 8, 8, 3), dtype=np.uint8) + audio = np.array([0, 32767, -32767], dtype=np.int16) + + result = post_process({"video": video, "audio": audio, "audio_sample_rate": 44100}, output_type="np") + + assert result["video"] is video + assert result["audio"].dtype == np.float32 + np.testing.assert_allclose(result["audio"], [0.0, 1.0, -1.0]) + assert result["audio_sample_rate"] == 44100 + assert result["fps"] == 24.0 + + +def test_kandinsky6_post_process_sample_rate_survives_output_formatter(): + """The formatter only lifts ``audio_sample_rate`` into metadata from the + flat payload form; an envelope would silently lose it and the MP4 would + be muxed at the 24 kHz default.""" + from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func + from vllm_omni.diffusion.output_formatter import normalize_diffusion_postprocess_output + + post_process = get_kandinsky6_post_process_func(od_config=None) + video = np.zeros((1, 2, 4, 4, 3), dtype=np.uint8) + audio = np.zeros((10,), dtype=np.int16) + + normalized = normalize_diffusion_postprocess_output( + post_process({"video": video, "audio": audio, "audio_sample_rate": 44100}, output_type="np") + ) + + assert normalized.primary_key == "video" + assert normalized.metadata["audio"]["sample_rate"] == 44100 + assert normalized.metadata["video"]["fps"] == 24.0 + assert "audio_sample_rate" not in normalized.outputs + + +def test_kandinsky6_post_process_func_omits_audio_when_none(): + from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func + + post_process = get_kandinsky6_post_process_func(od_config=None) + video = np.zeros((1, 2, 4, 4, 3), dtype=np.uint8) + + result = post_process({"video": video, "audio": None, "audio_sample_rate": None}, output_type="np") + + assert result["video"] is video + assert "audio" not in result + assert "audio_sample_rate" not in result + + +def test_kandinsky6_post_process_func_unwraps_batched_audio_list(): + from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func + + post_process = get_kandinsky6_post_process_func(od_config=None) + video = np.zeros((1, 2, 4, 4, 3), dtype=np.uint8) + audio_item = np.full((50,), 16384, dtype=np.int16) + + result = post_process({"video": video, "audio": [audio_item], "audio_sample_rate": 44100}, output_type="np") + + assert result["audio"].shape == (50,) + np.testing.assert_allclose(result["audio"], 16384 / 32767, rtol=1e-6) + + +class _StubAudioVAE: + """Audio VAE stand-in: identity "decode" that returns the latent as a waveform.""" + + scaling_factor = 0.5 + mean_value = 0.0 + device = "cpu" + + def wrapped_decode(self, latents): + # (1, audio_dim, A) -> (A,) waveform: take the first channel. + return latents[0, 0] + + +def test_postprocess_audio_normalizes_by_default_and_clips_on_request(): + """The default mode must match the k6_video production pipeline, which + peak-normalizes each decoded waveform to full scale; ``clip`` keeps the + raw amplitude (saturated to [-1, 1]) for callers that want it.""" + import torch + + from vllm_omni.diffusion.models.kandinsky6.pipeline_kandinsky6 import LatentBundle, postprocess_audio + + # Latent (A=4, audio_dim=1); the stub scales by 1/scaling_factor -> [0.2, -0.4, 0.6, 3.0]. + audio = torch.tensor([[0.1], [-0.2], [0.3], [1.5]], dtype=torch.float32) + bundle = LatentBundle( + video=None, + audio=audio, + video_cu_seqlens=None, + audio_cu_seqlens=torch.tensor([0, 4], dtype=torch.int32), + ) + vae = _StubAudioVAE() + + normalized = postprocess_audio(bundle, vae) + assert normalized is not None and len(normalized) == 1 + assert normalized[0].dtype == np.int16 + np.testing.assert_allclose( + normalized[0].astype(np.float32) / 32767, + np.array([0.2, -0.4, 0.6, 3.0], dtype=np.float32) / 3.0, + atol=1e-4, + ) + + clipped = postprocess_audio(bundle, vae, normalization_mode="clip") + np.testing.assert_allclose( + clipped[0].astype(np.float32) / 32767, + np.array([0.2, -0.4, 0.6, 1.0], dtype=np.float32), + atol=1e-4, + ) + + assert postprocess_audio(LatentBundle(None, None, None, None), vae) is None + with pytest.raises(ValueError, match="normalization_mode"): + postprocess_audio(bundle, vae, normalization_mode="loud") + + +def test_kandinsky6_pre_process_func_is_identity_for_now(): + """v1 scope: the pre-process hook is registered (matching the framework's + convention for image-conditioned models like SanaImageToVideoPipeline) + but does not yet validate/transform the request — documented follow-up, + not a silent gap.""" + from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_pre_process_func + + pre_process = get_kandinsky6_pre_process_func(od_config=None) + sentinel = object() + + assert pre_process(sentinel) is sentinel + + +def test_hub_component_keys_map_onto_pipeline_modules(): + """Hub folder keys from Kandinsky-6.0-Pro-5s-Diffusers land on the modules.""" + from vllm_omni.diffusion.models.kandinsky6.pipeline_kandinsky6 import ( + _WEIGHT_SUBFOLDERS, + _adapt_k6_weight_name, + ) + + assert _WEIGHT_SUBFOLDERS == ( + ("transformer", "transformer."), + ("vae", "vae."), + ("text_encoder", "text_encoder."), + ("text_encoder_2", "text_encoder_2."), + ("audio_vae", "audio_vae."), + ) + assert _adapt_k6_weight_name("transformer.visual_transformer_blocks.0.videoT.self_attention.to_query.weight") == ( + "transformer.visual_transformer_blocks.0.video_dec_block.self_attention.to_query.weight" + ) + assert _adapt_k6_weight_name("transformer.visual_transformer_blocks.0.audioT.feed_forward.in_layer.weight") == ( + "transformer.visual_transformer_blocks.0.audio_dec_block.feed_forward.in_layer.weight" + ) + assert _adapt_k6_weight_name("vae.decoder.conv_in.conv.weight") == "vae.decoder.conv_in.conv.weight" + assert _adapt_k6_weight_name("text_encoder_2.encoder.layers.0.mlp.fc1.weight") == ( + "text_encoder_2.encoder.layers.0.mlp.fc1.weight" + ) + assert ( + _adapt_k6_weight_name("text_encoder.model.layers.0.input_layernorm.weight") + == "text_encoder.model.language_model.layers.0.input_layernorm.weight" + ) + assert _adapt_k6_weight_name("text_encoder.visual.blocks.0.norm1.weight") == ( + "text_encoder.model.visual.blocks.0.norm1.weight" + ) + assert ( + _adapt_k6_weight_name("audio_vae.vae.decoder.conv_in.weight") + == "audio_vae.native.tod.vae.decoder.conv_in.weight" + ) + assert _adapt_k6_weight_name("audio_vae.vocoder.conv_pre.weight") == "audio_vae.native.tod.vocoder.conv_pre.weight" + assert _adapt_k6_weight_name("audio_vae.mel_converter.hann_window") == "audio_vae.native.mel_converter.hann_window" diff --git a/tests/diffusion/models/kandinsky6/test_transformer_kandinsky6.py b/tests/diffusion/models/kandinsky6/test_transformer_kandinsky6.py new file mode 100644 index 000000000..130f1dca0 --- /dev/null +++ b/tests/diffusion/models/kandinsky6/test_transformer_kandinsky6.py @@ -0,0 +1,325 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import os +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.fixture(autouse=True) +def _init_distributed(monkeypatch): + """The native transformer uses vLLM parallel linear layers, which require a + tensor-parallel group; initialize a single-process group for CPU tests + (mirrors tests/diffusion/models/sana_video/test_transformer_sana_video.py).""" + from vllm.distributed.parallel_state import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, + ) + from vllm.model_executor.layers.utils import default_unquantized_gemm + + monkeypatch.setattr( + "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", + lambda *_args, **_kwargs: default_unquantized_gemm, + ) + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29502") + init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://") + initialize_model_parallel() + yield + cleanup_dist_env_and_memory() + + +# model_dim must be an exact multiple of head_dim = sum(axes_dims); this is +# an architectural constraint of the native model (DiffusionTransformer3D), +# not specific to this port. +_TINY_T2V_CONFIG = { + "in_visual_dim": 4, + "out_visual_dim": 4, + "in_text_dim": 8, + "in_text_dim2": 6, + "time_dim": 16, + "patch_size": (1, 2, 2), + "model_dim": 24, + "ff_dim": 32, + "num_text_blocks": 1, + "num_visual_blocks": 2, + "axes_dims": (4, 4, 4), + "visual_cond": False, + "is_multimodal": False, + "attention_engine": "sdpa", +} + +# time_dim_a is intentionally omitted (defaults to time_dim): the fused +# block's cross-modal modulation (va_mod/av_mod) consumes the *other* +# modality's time embedding by default (fix_modulation=False), so time_dim +# and time_dim_a must match unless fix_modulation=True. +_TINY_T2VA_CONFIG = { + **_TINY_T2V_CONFIG, + "is_multimodal": True, + "in_audio_dim": 6, + "model_dim_a": 12, + "ff_dim_a": 16, + "axes_dims_a": (2, 2, 2), +} + + +def _sdpa_config(): + from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec + + return SimpleNamespace( + diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA")), + parallel_config=SimpleNamespace(ring_degree=1, allgather_degree=1), + ) + + +def _build_transformer(config): + from vllm_omni.diffusion.config import set_current_diffusion_config + from vllm_omni.diffusion.models.kandinsky6 import Kandinsky6Transformer3DModel + + with set_current_diffusion_config(_sdpa_config()): + model = Kandinsky6Transformer3DModel(**config) + model.eval() + return model + + +def test_t2v_only_forward_shape(): + model = _build_transformer(_TINY_T2V_CONFIG) + time_frames, height, width = 2, 4, 4 + x_video = torch.randn(time_frames, height, width, 4) + text_embed = torch.randn(5, 8) + pooled = torch.randn(1, 6) + time = torch.tensor([500.0]) + visual_rope = model.visual_rope_embeddings( + shape=(time_frames, height // 2, width // 2), + pos=[torch.arange(time_frames), torch.arange(height // 2), torch.arange(width // 2)], + scale_factor=(1.0, 1.0, 1.0), + ) + text_rope = model.text_rope_embeddings(torch.arange(5)) + + with torch.no_grad(): + out = model( + x_video=x_video, + x_audio=None, + text_embed=text_embed, + pooled_text_embed=pooled, + time=time, + visual_rope=visual_rope, + audio_rope=None, + text_rope=text_rope, + ) + + assert out.shape == (time_frames, height, width, 4) + + +def test_visual_cond_channel_concat_path(): + """visual_cond=True doubles+1 the visual input channel count (cond + + mask channels appended by _build_video_input, stripped again by + OutLayer) — the I2V conditioning path.""" + config = {**_TINY_T2V_CONFIG, "visual_cond": True} + model = _build_transformer(config) + time_frames, height, width = 2, 4, 4 + x_video = torch.randn(time_frames, height, width, 2 * 4 + 1) + text_embed = torch.randn(5, 8) + pooled = torch.randn(1, 6) + time = torch.tensor([500.0]) + visual_rope = model.visual_rope_embeddings( + shape=(time_frames, height // 2, width // 2), + pos=[torch.arange(time_frames), torch.arange(height // 2), torch.arange(width // 2)], + scale_factor=(1.0, 1.0, 1.0), + ) + text_rope = model.text_rope_embeddings(torch.arange(5)) + + with torch.no_grad(): + out = model( + x_video=x_video, + x_audio=None, + text_embed=text_embed, + pooled_text_embed=pooled, + time=time, + visual_rope=visual_rope, + audio_rope=None, + text_rope=text_rope, + ) + + assert out.shape == (time_frames, height, width, 4) + + +def test_t2va_joint_forward_shapes(): + """Exercises the fused cross-modal decoder blocks (Kandinsky6FusedTransformerDecoderBlock).""" + model = _build_transformer(_TINY_T2VA_CONFIG) + time_frames, height, width = 2, 4, 4 + audio_len = 7 + x_video = torch.randn(time_frames, height, width, 4) + x_audio = torch.randn(audio_len, 6) + text_embed = torch.randn(5, 8) + pooled = torch.randn(1, 6) + time = torch.tensor([500.0]) + visual_rope = model.visual_rope_embeddings( + shape=(time_frames, height // 2, width // 2), + pos=[torch.arange(time_frames), torch.arange(height // 2), torch.arange(width // 2)], + scale_factor=(1.0, 1.0, 1.0), + ) + audio_rope = model.audio_rope_embeddings(torch.arange(audio_len)) + video_text_rope = model.video_text_rope_embeddings(torch.arange(5)) + audio_text_rope = model.audio_text_rope_embeddings(torch.arange(5)) + + with torch.no_grad(): + video_out, audio_out = model( + x_video=x_video, + x_audio=x_audio, + text_embed=[text_embed, text_embed], + pooled_text_embed=[pooled, pooled], + time=[time, time], + visual_rope=visual_rope, + audio_rope=audio_rope, + text_rope=[video_text_rope, audio_text_rope], + ) + + assert video_out.shape == (time_frames, height, width, 4) + assert audio_out.shape == (audio_len, 6) + + +@pytest.mark.parametrize("modality", ["video", "audio"]) +def test_single_modality_through_multimodal_model(modality): + """A multimodal (is_multimodal=True) model must still support a pure + video-only or audio-only forward call (denoise_loop's partial-sampling + path uses this).""" + model = _build_transformer(_TINY_T2VA_CONFIG) + time_frames, height, width = 2, 4, 4 + audio_len = 7 + text_embed = torch.randn(5, 8) + pooled = torch.randn(1, 6) + time = torch.tensor([500.0]) + + if modality == "video": + x_video = torch.randn(time_frames, height, width, 4) + visual_rope = model.visual_rope_embeddings( + shape=(time_frames, height // 2, width // 2), + pos=[torch.arange(time_frames), torch.arange(height // 2), torch.arange(width // 2)], + scale_factor=(1.0, 1.0, 1.0), + ) + text_rope = model.video_text_rope_embeddings(torch.arange(5)) + with torch.no_grad(): + out = model( + x_video=x_video, + x_audio=None, + text_embed=text_embed, + pooled_text_embed=pooled, + time=time, + visual_rope=visual_rope, + audio_rope=None, + text_rope=text_rope, + ) + assert out.shape == (time_frames, height, width, 4) + else: + x_audio = torch.randn(audio_len, 6) + audio_rope = model.audio_rope_embeddings(torch.arange(audio_len)) + text_rope = model.audio_text_rope_embeddings(torch.arange(5)) + with torch.no_grad(): + out = model( + x_video=None, + x_audio=x_audio, + text_embed=text_embed, + pooled_text_embed=pooled, + time=time, + visual_rope=None, + audio_rope=audio_rope, + text_rope=text_rope, + ) + assert out.shape == (audio_len, 6) + + +@pytest.mark.parametrize("batched", [False, True]) +def test_self_attention_attends_over_tokens_for_batched_and_unbatched_inputs(batched): + """Self-attention must contract over the *token* axis whether its input is + unbatched text ``(S, C)`` or batched=1 visual/audio ``(1, N, C)``. + + Regression: an unconditional ``unsqueeze(0)`` on the already-batched + visual path handed SDPA a 5-D ``(1, 1, N, H, D)`` tensor, which it + accepts silently but then treats the head axis as the sequence — every + generated video decoded to pure noise.""" + from vllm_omni.diffusion.config import set_current_diffusion_config + from vllm_omni.diffusion.models.kandinsky6.kandinsky6_transformer import Kandinsky6Attention, apply_rotary + + torch.manual_seed(0) + channels, head_dim, tokens = 24, 12, 7 + with set_current_diffusion_config(_sdpa_config()): + attn = Kandinsky6Attention(channels, head_dim, engine="sdpa").eval() + with torch.no_grad(): + # vLLM parallel linears allocate with torch.empty (no init). + for param in attn.parameters(): + param.normal_(std=0.2) + x = torch.randn(tokens, channels) + rope = torch.randn(tokens, 1, head_dim // 2, 2, 2) + hidden = x.unsqueeze(0) if batched else x + + with torch.no_grad(): + out = attn(hidden, rotary_emb=rope) + + heads = channels // head_dim + q = attn.query_norm(attn.to_query(x).reshape(tokens, heads, head_dim)) + k = attn.key_norm(attn.to_key(x).reshape(tokens, heads, head_dim)) + v = attn.to_value(x).reshape(tokens, heads, head_dim) + q = apply_rotary(q, rope).type_as(q) + k = apply_rotary(k, rope).type_as(k) + ref = torch.nn.functional.scaled_dot_product_attention( + q.transpose(0, 1), k.transpose(0, 1), v.transpose(0, 1) + ).transpose(0, 1) + expected = attn.out_layer(ref.reshape(tokens, channels)) + + assert out.shape == (hidden.shape[:-1] + (channels,)) + torch.testing.assert_close(out.reshape(tokens, channels), expected, rtol=1e-4, atol=1e-5) + + +def test_modulation_projections_are_zero_initialized(): + """Modulation layers are zero-initialized at construction (standard + AdaLN-zero init) — confirms the zero-init in Kandinsky6Modulation.__init__ + reaches the tensor-parallel out_layer correctly (it's a + ColumnParallelLinear, not plain nn.Linear).""" + model = _build_transformer(_TINY_T2V_CONFIG) + for block in model.text_transformer_blocks: + assert torch.all(block.text_modulation.out_layer.weight == 0) + for block in model.visual_transformer_blocks: + assert torch.all(block.visual_modulation.out_layer.weight == 0) + + +def test_dit_constructs_from_a_diffusers_config_dict(): + from vllm_omni.diffusion.config import set_current_diffusion_config + from vllm_omni.diffusion.models.kandinsky6 import Kandinsky6Transformer3DModel + + with set_current_diffusion_config(_sdpa_config()): + model = Kandinsky6Transformer3DModel.from_diffusers_config( + {**_TINY_T2VA_CONFIG, "patch_size": [1, 2, 2], "axes_dims": [4, 4, 4]} + ) + assert tuple(model.patch_size) == (1, 2, 2) + assert model._repeated_blocks == [ + "Kandinsky6FusedTransformerDecoderBlock", + "Kandinsky6TransformerEncoderBlock", + "Kandinsky6TransformerDecoderBlock", + ] + visual = model.visual_transformer_blocks[0].video_dec_block.self_attention.attn + audio = model.visual_transformer_blocks[0].audio_dec_block.self_attention.attn + text = model.video_text_transformer_blocks[0].self_attention.attn + cross = model.visual_transformer_blocks[0].va_cross_attention.attn + assert visual.role == "kandinsky6.visual_self" + assert visual.skip_sequence_parallel is False + assert audio.skip_sequence_parallel is True + assert text.role == "kandinsky6.text_self" + assert text.skip_sequence_parallel is True + assert cross.role == "kandinsky6.video_audio_cross" + assert cross.skip_sequence_parallel is True + + +def test_cfg_combine_matches_apply_cfg_without_normalization(): + from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin + from vllm_omni.diffusion.models.kandinsky6.pipeline_kandinsky6 import apply_cfg + + positive = torch.tensor([1.0, 3.0]) + negative = torch.tensor([0.25, -1.0]) + combined = CFGParallelMixin.combine_cfg_noise(object(), positive, negative, 5.0, False) + torch.testing.assert_close(combined, apply_cfg(positive, negative, 5.0)) diff --git a/tests/diffusion/models/kandinsky6/test_vae_patch_parallel.py b/tests/diffusion/models/kandinsky6/test_vae_patch_parallel.py new file mode 100644 index 000000000..4863e1e89 --- /dev/null +++ b/tests/diffusion/models/kandinsky6/test_vae_patch_parallel.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Kandinsky 6 VAE patch parallel keeps the single-GPU tile grid.""" + +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from vllm_omni.diffusion.models.kandinsky6.modeling_kandinsky6_vae import AutoencoderKLHunyuanVideo + + +def _stub_vae() -> AutoencoderKLHunyuanVideo: + vae = AutoencoderKLHunyuanVideo.__new__(AutoencoderKLHunyuanVideo) + vae.spatial_compression_ratio = 8 + vae.temporal_compression_ratio = 4 + vae.use_tiling = True + vae.use_framewise_decoding = True + vae.tile_sample_min_height = 32 + vae.tile_sample_min_width = 32 + vae.tile_sample_min_num_frames = 16 + vae.tile_sample_stride_height = 24 + vae.tile_sample_stride_width = 24 + vae.tile_sample_stride_num_frames = 12 + vae.tile_size = None + vae.distributed_executor = None + vae.post_quant_conv = lambda tensor: tensor + vae._seen_widths = [] + + def decoder(tensor: torch.Tensor) -> torch.Tensor: + vae._seen_widths.append(int(tensor.shape[-1])) + return tensor + 1 + + vae.decoder = decoder + return vae + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_distributed_decode_tiles_match_local_tiled_decode(): + vae = _stub_vae() + latent = torch.arange(1 * 4 * 2 * 8 * 12, dtype=torch.float32).reshape(1, 4, 2, 8, 12) + + local = vae.tiled_decode(latent).sample + tasks, spec = vae._decode_tile_split(latent) + assert len(tasks) > 1 + assert max(int(task.tensor.shape[-1]) for task in tasks) < latent.shape[-1] + + decoded = {task.grid_coord: vae._decode_tile_exec(task) for task in tasks} + merged = vae._decode_tile_merge(decoded, spec) + + torch.testing.assert_close(merged, local) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_decode_still_applies_optimal_tiling_when_patch_parallel_is_on(): + vae = _stub_vae() + seen = {} + + def tiling(shape, device=None): + del device + seen["shape"] = tuple(shape) + return (1, 17, 32, 32), (8, 24, 24) + + from diffusers.models.autoencoders.vae import DecoderOutput + + vae.get_dec_optimal_tiling = tiling + vae._decode = lambda z, return_dict=True: DecoderOutput(sample=z) + vae.distributed_executor = SimpleNamespace(parallel_size=2) + vae.is_distributed_enabled = lambda: True + + out = vae.decode(torch.zeros(1, 16, 8, 8, 12)) + + assert seen["shape"] == (1, 16, 8, 8, 12) + assert tuple(out.sample.shape) == (1, 16, 8, 8, 12) + + +def _patch_parallel_worker(rank: int, port: str, shapes: dict) -> None: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = "2" + dist.init_process_group("gloo", rank=rank, world_size=2) + try: + import vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor as executor_mod + + world = SimpleNamespace(device_group=dist.group.WORLD) + executor_mod.get_world_group = lambda: world + + vae = _stub_vae() + latent = torch.arange(1 * 4 * 2 * 8 * 12, dtype=torch.float32).reshape(1, 4, 2, 8, 12) + reference = vae.tiled_decode(latent).sample.clone() + vae._seen_widths.clear() + vae.set_parallel_size(2, mode="tile") + decoded = vae.tiled_decode(latent).sample + torch.testing.assert_close(decoded, reference) + if rank == 0: + shapes["widths"] = list(vae._seen_widths) + shapes["full_width"] = int(latent.shape[-1]) + finally: + dist.destroy_process_group() + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_patch_parallel_decode_splits_tiles_across_ranks(): + manager = mp.Manager() + shapes = manager.dict() + port = str(29500 + (os.getpid() % 1000)) + mp.spawn(_patch_parallel_worker, args=(port, shapes), nprocs=2, join=True) + + assert shapes["widths"] + assert max(shapes["widths"]) < shapes["full_width"] diff --git a/tests/e2e/offline_inference/test_kandinsky6.py b/tests/e2e/offline_inference/test_kandinsky6.py new file mode 100644 index 000000000..c4296fffb --- /dev/null +++ b/tests/e2e/offline_inference/test_kandinsky6.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Offline smoke for Kandinsky 6 TI2VA. + +Short request (256x256, 5 frames, 2 steps) so it finishes inside the default +online poll window. Pro geometry (480x864, 125 frames, 50 steps) stays the +serving default, not this smoke. +""" + +import os +from pathlib import Path + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OfflineOmniClient, get_model_prefix +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +_HUB_MODEL = "kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers" +_LOCAL_BUNDLE = Path(__file__).resolve().parents[4] / "kandinsky6_bundle" + + +def _resolve_model() -> str: + env = os.environ.get("KANDINSKY6_MODEL") + if env: + return env + if _LOCAL_BUNDLE.is_dir(): + return str(_LOCAL_BUNDLE) + return _HUB_MODEL + + +def _checkpoint_on_disk(model: str) -> bool: + if Path(model).is_dir(): + return True + prefix = get_model_prefix() + return bool(prefix) and Path(prefix + model).is_dir() + + +MODEL = _resolve_model() +PROMPT = "A golden retriever runs along a sunny beach, waves crashing." + +pytestmark = [ + pytest.mark.skipif( + not _checkpoint_on_disk(MODEL), + reason=( + "Kandinsky 6 Diffusers weights are gated. Set KANDINSKY6_MODEL or " + "MODEL_PREFIX to a local copy of " + "kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers." + ), + ), +] + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1}) +@pytest.mark.parametrize( + "omni_runner", + [(MODEL, None, {"enable_cpu_offload": True, "enforce_eager": True})], + indirect=True, +) +def test_text_to_video_and_audio_001(offline_client: OfflineOmniClient): + sampling = OmniDiffusionSamplingParams( + height=256, + width=256, + num_frames=5, + fps=24, + num_inference_steps=2, + guidance_scale=5.0, + seed=42, + extra_args={"sample_audio": True}, + ) + request_config = { + "model": MODEL, + "prompt": PROMPT, + "sampling_params": sampling, + } + offline_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_kandinsky6.py b/tests/e2e/online_serving/test_kandinsky6.py new file mode 100644 index 000000000..13e3f903d --- /dev/null +++ b/tests/e2e/online_serving/test_kandinsky6.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Online smoke for Kandinsky 6 TI2VA via ``/v1/videos``. + +Short geometry so the job finishes inside the helper's default 300 s poll. +""" + +import os +from pathlib import Path + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, get_model_prefix + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +_HUB_MODEL = "kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers" +_LOCAL_BUNDLE = Path(__file__).resolve().parents[4] / "kandinsky6_bundle" + + +def _resolve_model() -> str: + env = os.environ.get("KANDINSKY6_MODEL") + if env: + return env + if _LOCAL_BUNDLE.is_dir(): + return str(_LOCAL_BUNDLE) + return _HUB_MODEL + + +def _checkpoint_on_disk(model: str) -> bool: + if Path(model).is_dir(): + return True + prefix = get_model_prefix() + return bool(prefix) and Path(prefix + model).is_dir() + + +MODEL = _resolve_model() +PROMPT = "A golden retriever runs along a sunny beach, waves crashing." + +pytestmark = [ + pytest.mark.skipif( + not _checkpoint_on_disk(MODEL), + reason=( + "Kandinsky 6 Diffusers weights are gated. Set KANDINSKY6_MODEL or " + "MODEL_PREFIX to a local copy of " + "kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers." + ), + ), +] + +SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) + + +def _get_diffusion_feature_cases(model: str): + return [ + pytest.param( + OmniServerParams( + model=model, + server_args=["--enable-cpu-offload", "--enforce-eager"], + init_timeout=3600, + stage_init_timeout=3600, + ), + id="cpu_offload", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + ] + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) +def test_text_to_video_and_audio_001(omni_server: OmniServer, online_client: OnlineOmniClient) -> None: + """Short TI2VA job completes and returns an MP4.""" + request_config = { + "model": omni_server.model, + "form_data": { + "prompt": PROMPT, + "height": 256, + "width": 256, + "num_frames": 5, + "fps": 24, + "num_inference_steps": 2, + "guidance_scale": 5.0, + "seed": 42, + }, + } + online_client.send_video_diffusion_request(request_config) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 740961f81..ba3569a8c 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -999,6 +999,38 @@ def test_magi2_serving_applies_native_defaults_and_rejects_explicit_frame_mismat ) +def test_kandinsky6_serving_applies_defaults_but_allows_explicit_frame_count(): + """Kandinsky 6 publishes Pro-geometry defaults (480x864, 125 frames, 50 + steps) but is not a fixed-duration model: an explicit ``num_frames`` + must be honoured rather than rejected.""" + engine = FakeAsyncOmni() + engine.model_class_name = "Kandinsky6TI2VAPipeline" + handler = OmniOpenAIServingVideo.for_diffusion( + diffusion_engine=engine, + model_name="kandinsky6-bundle", + ) + + asyncio.run(handler._run_and_extract(VideoGenerationRequest(prompt="A dog runs on a beach"), "defaults")) + sampling = engine.captured_sampling_params_list[0] + assert (sampling.width, sampling.height) == (864, 480) + assert sampling.num_frames == 125 + assert sampling.num_inference_steps == 50 + assert sampling.fps == 24.0 + + asyncio.run( + handler._run_and_extract( + VideoGenerationRequest( + prompt="A dog runs on a beach", width=512, height=320, num_frames=25, num_inference_steps=10 + ), + "explicit", + ) + ) + sampling = engine.captured_sampling_params_list[-1] + assert (sampling.width, sampling.height) == (512, 320) + assert sampling.num_frames == 25 + assert sampling.num_inference_steps == 10 + + def test_i2v_video_generation_with_image_reference_form(test_client, mocker: MockerFixture): mocker.patch( "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index d3ef4fd12..3b2b96adc 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -76,6 +76,9 @@ EXCLUDED_MODELS = [ # The 270+ GiB multimodal checkpoint needs 4/8-way native sequence/head # parallel execution; covered by focused native contract and GPU tests. "Magi2Pipeline", + # Pro DiT is ~56 GiB bf16 and the Hub checkpoint is gated. Covered by + # tests/diffusion/models/kandinsky6 and the real-checkpoint E2E smoke. + "Kandinsky6TI2VAPipeline", "OmniVoicePipeline", "OmniVoice", "Cosmos3OmniDiffusersPipeline", diff --git a/vllm_omni/diffusion/cache/cachedit/model_specific.py b/vllm_omni/diffusion/cache/cachedit/model_specific.py index 8bc36a8b0..3bcf94e1a 100644 --- a/vllm_omni/diffusion/cache/cachedit/model_specific.py +++ b/vllm_omni/diffusion/cache/cachedit/model_specific.py @@ -908,6 +908,35 @@ def enable_cache_for_mammothmoda2(pipeline: Any, cache_config: Any) -> CacheDiTE return CacheDiTEnableResult(refresh=refresh, targets=(block_adapter,)) +def enable_cache_for_kandinsky6(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc: + """Cache-DiT over K6 visual blocks only. + + Text and audio block lists are not one residual chain with the visual + stack. If the block adapter rejects the fused ``(video, audio)`` return, + fall back to the same step-level velocity cache MagCache uses. + """ + transformer = pipeline.transformer + block_adapter = BlockAdapter( + transformer=transformer, + blocks=[transformer.visual_transformer_blocks], + forward_pattern=[ForwardPattern.Pattern_0], + has_separate_cfg=True, + check_forward_pattern=False, + ) + try: + return enable_cache_for_dit(pipeline, cache_config, block_adapter) + except Exception as exc: + logger.warning("Cache-DiT block adapter failed for Kandinsky 6 (%s); using step cache", exc) + from vllm_omni.diffusion.models.kandinsky6.cache_accel import attach_mag_cache + + attach_mag_cache(transformer, threshold=0.24, max_skip_steps=3, retention_ratio=0.1) + + def _refresh(*_args, **_kwargs): + return None + + return _refresh + + def register_custom_dit_enablers() -> None: """Register model-specific Cache-DiT enablers. @@ -925,6 +954,7 @@ def register_custom_dit_enablers() -> None: "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3, "Cosmos3OmniPipeline": enable_cache_for_cosmos3, "Krea2Pipeline": enable_cache_for_krea2, + "Kandinsky6TI2VAPipeline": enable_cache_for_kandinsky6, "Magi2Pipeline": enable_cache_for_magi2, "MammothModa2DiTPipeline": enable_cache_for_mammothmoda2, } diff --git a/vllm_omni/diffusion/cache/magcache/backend.py b/vllm_omni/diffusion/cache/magcache/backend.py index 4ad427d3a..965e5c2e1 100644 --- a/vllm_omni/diffusion/cache/magcache/backend.py +++ b/vllm_omni/diffusion/cache/magcache/backend.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ MagCache backend implementation. @@ -73,6 +73,19 @@ class MagCacheBackend(CacheBackend): transformer = pipeline.transformer transformer_type = transformer.__class__.__name__ + if transformer_type == "Kandinsky6Transformer3DModel": + from vllm_omni.diffusion.models.kandinsky6.cache_accel import attach_mag_cache + + attach_mag_cache( + transformer, + threshold=self.config.mag_threshold, + max_skip_steps=self.config.mag_max_skip_steps, + retention_ratio=self.config.mag_retention_ratio, + ) + self._registered = True + logger.info("MagCache: Kandinsky 6 step-cache attached (visual/audio velocity reuse)") + return + num_inference_steps = self.config.num_inference_steps or 28 mag_ratios = self.config.mag_ratios diff --git a/vllm_omni/diffusion/cache/magcache/hook.py b/vllm_omni/diffusion/cache/magcache/hook.py index 4b2b4ea06..2c741a6d3 100644 --- a/vllm_omni/diffusion/cache/magcache/hook.py +++ b/vllm_omni/diffusion/cache/magcache/hook.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Hook-based MagCache implementation for vLLM-Omni. @@ -543,9 +543,12 @@ def apply_mag_cache_hook( state_manager = StateManager(MagCacheState, (), {}) remaining_blocks = [] + allowed = getattr(module, "_magcache_block_attrs", None) for name, submodule in module.named_children(): if not isinstance(submodule, torch.nn.ModuleList): continue + if allowed is not None and name not in allowed: + continue for index, block in enumerate(submodule): remaining_blocks.append((f"{name}.{index}", block)) diff --git a/vllm_omni/diffusion/cache/teacache/backend.py b/vllm_omni/diffusion/cache/teacache/backend.py index 6da9d953b..c745b666f 100644 --- a/vllm_omni/diffusion/cache/teacache/backend.py +++ b/vllm_omni/diffusion/cache/teacache/backend.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ TeaCache backend implementation. @@ -143,11 +143,30 @@ def enable_flux2_klein_teacache(pipeline: Any, config: DiffusionCacheConfig) -> ) +def enable_kandinsky6_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None: + """TeaCache for Kandinsky 6 via the denoise-loop step cache. + + Coefficients are uncalibrated (see ``K6_TEACACHE_COEFFICIENTS``). + """ + from vllm_omni.diffusion.models.kandinsky6.cache_accel import attach_tea_cache + + attach_tea_cache( + pipeline.transformer, + rel_l1_thresh=config.rel_l1_thresh if config.rel_l1_thresh is not None else 0.2, + coefficients=tuple(config.coefficients) if config.coefficients is not None else None, + ) + logger.info( + "TeaCache: Kandinsky 6 step-cache attached (uncalibrated coefficients, thresh=%s)", + config.rel_l1_thresh, + ) + + CUSTOM_TEACACHE_ENABLERS = { "BagelPipeline": enable_bagel_teacache, "Flux2KleinPipeline": enable_flux2_klein_teacache, "HunyuanImage3Pipeline": enable_hunyuan_image3_teacache, "MammothModa2DiTPipeline": enable_mammoth_moda2_teacache, + "Kandinsky6TI2VAPipeline": enable_kandinsky6_teacache, "MiniMaxH3Pipeline": enable_minimax_h3_teacache, "SenseNovaU1Pipeline": enable_sensenova_u1_teacache, } diff --git a/vllm_omni/diffusion/cache/teacache/config.py b/vllm_omni/diffusion/cache/teacache/config.py index 832e13ca0..3c23d84e4 100644 --- a/vllm_omni/diffusion/cache/teacache/config.py +++ b/vllm_omni/diffusion/cache/teacache/config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from dataclasses import dataclass @@ -8,6 +8,15 @@ from dataclasses import dataclass # Source: TeaCache paper and ComfyUI-TeaCache empirical tuning _MODEL_COEFFICIENTS = { # FLUX transformer coefficients from TeaCache paper + # Uncalibrated stand-in so Kandinsky 6 TeaCache can be enabled. + # Copied from Qwen-Image; refit before treating skips as quality-neutral. + "Kandinsky6Transformer3DModel": [ + -4.50000000e02, + 2.80000000e02, + -4.50000000e01, + 3.20000000e00, + -2.00000000e-02, + ], "FluxTransformer2DModel": [ 4.98651651e02, -2.83781631e02, diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index aaefb2cee..9e924370d 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -81,6 +81,15 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = { max_multimodal_image_inputs=1, final_output_type="video", ), + # Joint text/image-to-video-and-audio, same shape as MiniMaxH3Pipeline + # above (an MP4 with both tracks) — declared "video" for the same reason: + # the final container is a video file, so /v1/videos* is the right API + # surface even though the model also produces audio. + "Kandinsky6TI2VAPipeline": DiffusionModelMetadata( + supports_multimodal_inputs=True, + max_multimodal_image_inputs=1, + final_output_type="video", + ), "WanPipeline": DiffusionModelMetadata( attention_mask_free=True, final_output_type="video", diff --git a/vllm_omni/diffusion/models/kandinsky6/__init__.py b/vllm_omni/diffusion/models/kandinsky6/__init__.py new file mode 100644 index 000000000..29c8811c1 --- /dev/null +++ b/vllm_omni/diffusion/models/kandinsky6/__init__.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from .kandinsky6_transformer import Kandinsky6Transformer3DModel +from .modeling_kandinsky6_audio import Kandinsky6AudioVAE +from .modeling_kandinsky6_vae import AutoencoderKLHunyuanVideo +from .pipeline_kandinsky6 import ( + Kandinsky6TI2VAPipeline, + get_kandinsky6_post_process_func, + get_kandinsky6_pre_process_func, +) +from .scheduling_kandinsky6 import KandinskyFlowMatchScheduler + +__all__ = [ + "AutoencoderKLHunyuanVideo", + "Kandinsky6AudioVAE", + "Kandinsky6TI2VAPipeline", + "Kandinsky6Transformer3DModel", + "KandinskyFlowMatchScheduler", + "get_kandinsky6_post_process_func", + "get_kandinsky6_pre_process_func", +] diff --git a/vllm_omni/diffusion/models/kandinsky6/cache_accel.py b/vllm_omni/diffusion/models/kandinsky6/cache_accel.py new file mode 100644 index 000000000..322d51431 --- /dev/null +++ b/vllm_omni/diffusion/models/kandinsky6/cache_accel.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Kandinsky 6 cache adapters. + +Generic MagCache / TeaCache hooks assume a single residual stream named +``hidden_states``. K6's fused visual blocks return ``(video, audio)``, so +the denoise loop consults :class:`K6StepCache` instead and reuses the +previous velocity when the schedule says the step can be skipped. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +# Pro geometry ratios from k6_video +# ``k6_pro_125_480_864_mCache_mOffload.yaml`` (cache.magcache.mag_ratios). +K6_PRO_MAG_RATIOS: tuple[float, ...] = ( + 1.0, + 1.0, + 0.94526286, + 0.94622554, + 1.13130749, + 1.11243134, + 1.02926022, + 1.02675114, + 1.011996, + 1.01235353, + 1.03039566, + 1.02848639, + 1.00101278, + 1.00307236, + 1.02238594, + 1.02108294, + 1.0248162, + 1.02491556, + 0.99771836, + 0.99882495, + 1.01694329, + 1.01367628, + 1.01096435, + 1.00959183, + 1.01537554, + 1.01504827, + 0.99237553, + 0.99345238, + 1.02368983, + 1.02088727, + 1.01396087, + 1.01278108, + 0.99773024, + 0.99927632, + 1.01573298, + 1.01489232, + 1.00129728, + 1.00120319, + 1.01759963, + 1.01691798, + 0.99335877, + 0.99548372, + 1.02925427, + 1.02299559, + 0.99749375, + 0.99682709, + 0.99694342, + 1.00064403, + 1.01101129, + 1.0117984, + 1.04420642, + 1.03624643, + 0.9599572, + 0.97176547, + 1.00452036, + 1.0046107, + 1.01447715, + 1.01351055, + 0.9999695, + 1.00270806, + 1.00770242, + 1.00271775, + 0.98795461, + 0.98880837, + 1.01113352, + 1.00860993, + 1.00116034, + 1.00251894, + 1.05147006, + 1.04040677, + 0.99651072, + 0.99212138, + 0.99040673, + 0.98881456, + 1.01265486, + 1.01241698, + 0.97343865, + 0.97302804, + 1.02539269, + 1.02113221, + 1.02117306, + 1.02987275, + 0.98607016, + 0.98262938, + 0.98404577, + 0.98453141, + 1.01065119, + 1.00807871, + 0.99986077, + 0.99668861, + 0.98597074, + 0.98133325, + 0.98239303, + 0.98370113, + 0.93155548, + 0.93534231, + 0.94231988, + 0.93388087, + 0.84200656, + 0.83657581, +) + +# Uncalibrated. Borrowed from the Qwen-Image TeaCache polynomial so the flag +# skips steps instead of raising. Quality will not match a fitted K6 curve. +K6_TEACACHE_COEFFICIENTS: tuple[float, ...] = ( + -4.50000000e02, + 2.80000000e02, + -4.50000000e01, + 3.20000000e00, + -2.00000000e-02, +) + + +@dataclass +class K6StepCache: + """Step-level skip state shared by MagCache and TeaCache.""" + + kind: str + threshold: float + max_skip_steps: int + retention_ratio: float + mag_ratios: tuple[float, ...] | None = None + coefficients: tuple[float, ...] | None = None + rel_l1_thresh: float = 0.2 + last_cond: object = None + last_uncond: object = None + accumulated_err: float = 0.0 + accumulated_ratio: float = 1.0 + accumulated_steps: int = 0 + + def should_skip(self, step_index: int) -> bool: + if self.last_cond is None or step_index <= 0: + return False + if self.kind == "mag": + ratios = self.mag_ratios or () + scale = float(ratios[step_index]) if step_index < len(ratios) else 1.0 + self.accumulated_ratio *= scale + self.accumulated_steps += 1 + self.accumulated_err += abs(1.0 - self.accumulated_ratio) + if ( + self.accumulated_err <= self.threshold + and self.accumulated_steps <= self.max_skip_steps + and step_index >= max(1, int(self.retention_ratio * max(step_index, 1))) + ): + return True + self.accumulated_ratio = 1.0 + self.accumulated_steps = 0 + self.accumulated_err = 0.0 + return False + # TeaCache stand-in: uncalibrated polynomial of the step index. + # The real TeaCache distance is the modulated input; this only makes + # the flag skip steps instead of raising. + coeff = self.coefficients or K6_TEACACHE_COEFFICIENTS + x = float(step_index) + predicted = sum(weight * (x ** (len(coeff) - 1 - i)) for i, weight in enumerate(coeff)) + return abs(predicted) < self.rel_l1_thresh * 1.0e4 + + def store(self, cond: object, uncond: object) -> None: + self.last_cond = cond + self.last_uncond = uncond + + +def attach_mag_cache(transformer, *, threshold: float, max_skip_steps: int, retention_ratio: float) -> None: + transformer._k6_step_cache = K6StepCache( + kind="mag", + threshold=threshold, + max_skip_steps=max_skip_steps, + retention_ratio=retention_ratio, + mag_ratios=K6_PRO_MAG_RATIOS, + ) + + +def attach_tea_cache(transformer, *, rel_l1_thresh: float, coefficients: tuple[float, ...] | None) -> None: + transformer._k6_step_cache = K6StepCache( + kind="tea", + threshold=rel_l1_thresh if rel_l1_thresh is not None else 0.2, + max_skip_steps=3, + retention_ratio=0.0, + coefficients=coefficients or K6_TEACACHE_COEFFICIENTS, + rel_l1_thresh=rel_l1_thresh if rel_l1_thresh is not None else 0.2, + ) diff --git a/vllm_omni/diffusion/models/kandinsky6/kandinsky6_transformer.py b/vllm_omni/diffusion/models/kandinsky6/kandinsky6_transformer.py new file mode 100644 index 000000000..703e3427b --- /dev/null +++ b/vllm_omni/diffusion/models/kandinsky6/kandinsky6_transformer.py @@ -0,0 +1,2006 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Generated standalone Kandinsky 6 vLLM-Omni transformer. + +The template records source modules only. ``kandinsky.ports.assemble`` reads +those modules and pastes their current definitions into the generated file; +the generated artifact therefore has no native-package import dependency +(only ``torch`` and ``vllm``/``vllm_omni``, matching vLLM-Omni's other +native diffusion models). +""" + +from __future__ import annotations + +import contextvars +import math +from typing import Any + +import torch +from torch import Tensor, nn +from torch.nn.attention.flex_attention import BlockMask +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput + +# Set while visual blocks run so self-attention can apply Ulysses/Ring only +# to the sequence-sharded visual stream (text self-attention stays local). +_VISUAL_SP: contextvars.ContextVar[bool] = contextvars.ContextVar("k6_visual_sp", default=False) + + +def _parallel_size(kind: str) -> int: + """World size / rank for a parallel axis. 1 (or rank 0) when unset.""" + try: + from vllm_omni.diffusion.distributed import parallel_state as ps + except Exception: + return 0 if kind == "pp_rank" else 1 + try: + if kind == "pp": + return int(ps.get_pipeline_parallel_world_size()) + if kind == "pp_rank": + return int(ps.get_pipeline_parallel_rank()) + if kind == "sp": + return int(ps.get_sequence_parallel_world_size()) + if kind == "ulysses": + return int(ps.get_ulysses_parallel_world_size()) + if kind == "ring": + return int(ps.get_ring_parallel_world_size()) + except Exception: + return 0 if kind == "pp_rank" else 1 + return 1 + + +def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: + return torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) + + +def apply_scale_shift_norm(norm, x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: + """AdaLN-style affine in fp32, cast back to ``x.dtype``.""" + return (norm(x.float()) * (scale.float() + 1.0) + shift.float()).to(dtype=x.dtype) + + +def apply_gate_sum(x: Tensor, out: Tensor, gate: Tensor) -> Tensor: + """Residual gate in fp32, cast back to ``x.dtype``.""" + return (x.float() + gate.float() * out.float()).to(dtype=x.dtype) + + +def apply_rotary(x: Tensor, rope: Tensor) -> Tensor: + """RoPE apply in fp32 (rope tables are fp32), cast back to ``x.dtype``.""" + x_ = x.reshape(*x.shape[:-1], -1, 1, 2).float() + return (rope.float() * x_).sum(dim=-1).reshape(*x.shape).to(dtype=x.dtype) + + +def _local_patch(x: Tensor, shape: tuple, group_size: tuple, dim: int = 0) -> Tensor: + T, H, W = shape + g1, g2, g3 = group_size + x = x.reshape(*x.shape[:dim], T // g1, g1, H // g2, g2, W // g3, g3, *x.shape[dim + 3 :]) + d = len(x.shape[:dim]) + x = x.permute(*range(d), d, d + 2, d + 4, d + 1, d + 3, d + 5, *range(d + 6, len(x.shape))) + return x.flatten(dim, dim + 2).flatten(dim + 1, dim + 3) + + +def _local_merge(x: Tensor, shape: tuple, group_size: tuple, dim: int = 0) -> Tensor: + T, H, W = shape + g1, g2, g3 = group_size + x = x.reshape(*x.shape[:dim], T // g1, H // g2, W // g3, g1, g2, g3, *x.shape[dim + 2 :]) + d = len(x.shape[:dim]) + x = x.permute(*range(d), d, d + 3, d + 1, d + 4, d + 2, d + 5, *range(d + 6, len(x.shape))) + return x.flatten(dim, dim + 1).flatten(dim + 1, dim + 2).flatten(dim + 2, dim + 3) + + +def fractal_flatten(x: Tensor, rope: Tensor, shape: tuple, block_mask: bool = False): + if block_mask: + ps = 8 + x = _local_patch(x, shape, (1, ps, ps), dim=0) + rope = _local_patch(rope, shape, (1, ps, ps), dim=0) + return x.flatten(0, 1), rope.flatten(0, 1) + return x.flatten(0, 2), rope.flatten(0, 2) + + +def fractal_unflatten(x: Tensor, shape: tuple, block_mask: bool = False) -> Tensor: + if block_mask: + ps = 8 + x = x.reshape(-1, ps * ps, x.shape[-1]) + return _local_merge(x, shape, (1, ps, ps), dim=0) + return x.reshape(*shape, x.shape[-1]) + + +def fast_sta_nabla( + T: int, + H: int, + W: int, + wT: int = 3, + wH: int = 3, + wW: int = 3, + device: str | torch.device = "cuda", +) -> Tensor: + """Precomputes the Sliding Tile Attention (STA) boolean mask for nabla attention.""" + l = max(T, H, W) + r = torch.arange(l, dtype=torch.int16, device=device) + mat = (r.unsqueeze(1) - r.unsqueeze(0)).abs() + + sta_t = mat[:T, :T].flatten() <= wT // 2 + sta_h = mat[:H, :H].flatten() <= wH // 2 + sta_w = mat[:W, :W].flatten() <= wW // 2 + + sta_hw = (sta_h.unsqueeze(1) * sta_w.unsqueeze(0)).reshape(H, H, W, W).transpose(1, 2).flatten() + sta = (sta_t.unsqueeze(1) * sta_hw.unsqueeze(0)).reshape(T, T, H * W, H * W).transpose(1, 2) + return sta.reshape(T * H * W, T * H * W) + + +def nabla_block_mask( + q: Tensor, + k: Tensor, + sta: Tensor, + thr: float = 0.9, + block_size: int = 64, +) -> BlockMask: + """Builds a dynamic nabla BlockMask from query/key statistics + STA prior.""" + B, h, S, D = q.shape + s1 = S // block_size + qa = q.reshape(B, h, s1, block_size, D).mean(-2) + ka = k.reshape(B, h, s1, block_size, D).mean(-2).transpose(-2, -1) + attn_map = torch.softmax((qa @ ka) / math.sqrt(D), dim=-1) + + vals, inds = attn_map.sort(-1) + mask = (vals.cumsum_(-1) >= 1 - thr).int().gather(-1, inds.argsort(-1)) + mask = torch.logical_or(mask, sta) + + kv_nb = mask.sum(-1).to(torch.int32) + kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) + return BlockMask.from_kv_blocks( + torch.zeros_like(kv_nb), + kv_inds, + kv_nb, + kv_inds, + BLOCK_SIZE=block_size, + mask_mod=None, + ) + + +class RoPE1D(nn.Module): + """1-D Rotary Position Embedding — used for text and audio sequences.""" + + def __init__( + self, + dim: int, + max_pos: int = 2048, + max_period: float = 10000.0, + freqs_scaling: float = 1.0, + ): + super().__init__() + self.dim = dim + self.max_pos = max_pos + self.max_period = max_period + self.freqs_scaling = freqs_scaling + freq = get_freqs(dim // 2, max_period) * freqs_scaling + self.register_buffer("args", torch.outer(torch.arange(max_pos, dtype=freq.dtype), freq), persistent=False) + + def forward(self, pos: Tensor) -> Tensor: + # RoPE tables are fp32; keep trig in fp32. + args = self.args[pos] # (seq_len, dim//2) + rope = torch.stack([torch.cos(args), -torch.sin(args), torch.sin(args), torch.cos(args)], dim=-1) + return rope.view(*rope.shape[:-1], 2, 2).unsqueeze(-4) + + def reset_parameters(self) -> None: + freq = get_freqs(self.dim // 2, self.max_period).to(self.args.device) * self.freqs_scaling + self.args = torch.outer(torch.arange(self.max_pos, dtype=freq.dtype, device=freq.device), freq) + + +class RoPE3D(nn.Module): + """3-D Rotary Position Embedding — used for video spatial-temporal tokens (T, H, W).""" + + def __init__( + self, + axes_dims: tuple[int, int, int], + max_pos: tuple[int, int, int] = (128, 128, 128), + max_period: float = 10000.0, + ): + super().__init__() + self.axes_dims = axes_dims + self.max_pos = max_pos + self.max_period = max_period + for i, (d, mp) in enumerate(zip(axes_dims, max_pos)): + freq = get_freqs(d // 2, max_period) + self.register_buffer(f"args_{i}", torch.outer(torch.arange(mp, dtype=freq.dtype), freq), persistent=False) + + def forward( + self, + shape: tuple, + pos: list[Tensor], + scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), + ) -> Tensor: + T, H, W = shape + args_t = getattr(self, "args_0")[pos[0]] / scale_factor[0] # (T, d//2) + args_h = getattr(self, "args_1")[pos[1]] / scale_factor[1] # (H, d//2) + args_w = getattr(self, "args_2")[pos[2]] / scale_factor[2] # (W, d//2) + + args = torch.cat( + [ + args_t.view(T, 1, 1, -1).expand(T, H, W, -1), + args_h.view(1, H, 1, -1).expand(T, H, W, -1), + args_w.view(1, 1, W, -1).expand(T, H, W, -1), + ], + dim=-1, + ) + cos, sin = torch.cos(args), torch.sin(args) + rope = torch.stack([cos, -sin, sin, cos], dim=-1) # (T, H, W, total_dim, 4) + rope = rope.view(*rope.shape[:-1], 2, 2) # (T, H, W, total_dim, 2, 2) + return rope.unsqueeze(-4) # (T, H, W, 1, total_dim, 2, 2) + + def reset_parameters(self) -> None: + for i, (d, mp) in enumerate(zip(self.axes_dims, self.max_pos)): + freq = get_freqs(d // 2, self.max_period).to(getattr(self, f"args_{i}").device) + setattr(self, f"args_{i}", torch.outer(torch.arange(mp, dtype=freq.dtype, device=freq.device), freq)) + + +"""vLLM-Omni transformer for Kandinsky 6 — tensor-parallel port of ``DiffusionTransformer3D``. + +The port assembler extracts every top-level class from this module into the +generated vLLM-Omni pipeline (``inline_module``). The block structure and +AdaLN-modulation math are unchanged from ``core/components/dit.py`` — but +unlike an earlier revision of this file, submodule attribute names and +nesting here match the **Diffusers-exported checkpoint layout** exactly +(``in_layer``/``out_layer``, ``to_query``/``to_key``/``to_value``, +``self_attention``/``cross_attention``, ``text_transformer_blocks``/ +``visual_transformer_blocks``, etc. — see +``ports/overrides/diffusers/kandinsky6_transformer.py`` for the reference +this mirrors), not core's terse native names. This is deliberate: real +Kandinsky 6 checkpoints are converted and published in the "patched +Diffusers" bundle layout (``convert_checkpoint.py`` with +``--use-patched-diffusers``, i.e. ``uses_patched_diffusers: true`` in +``model_index.json``), which saves weights under the Diffusers vocabulary. +Matching that vocabulary here — including using SEPARATE +``to_query``/``to_key``/``to_value`` projections rather than a fused QKV +matrix — means ``ModelMixin.from_pretrained()`` loads such a checkpoint +directly via ordinary 1:1 state-dict key matching, with no custom +weight-fusion loader required. (An earlier revision used a fused +``QKVParallelLinear`` for self-attention, which is structurally +incompatible with a checkpoint that stores separate Q/K/V tensors — fixed +here in favor of three separate ``ColumnParallelLinear``s, matching +cross-attention's layout.) + +Sharding scheme (Megatron-style, matching vLLM-Omni's other native +transformers — Wan2.2, MiniMax H3): +- A "pair" gets Column -> Row: the first Linear scatters to an + intermediate that's sharded per rank (``gather_output=False``), the + second gathers/reduces it back to a full, replicated ``model_dim`` + output that's added straight onto the residual stream + (``input_is_parallel=True``). Used for ``feed_forward`` and + ``*_time_embeddings``. +- A standalone Linear whose result must be a full (replicated) tensor + immediately — because it either seeds the residual stream or its output + is consumed elementwise (AdaLN shift/scale/gate, applied directly to the + full residual tensor) rather than matrix-multiplied further — uses + ``ColumnParallelLinear(..., gather_output=True)``. Used for + ``*_text_embeddings``/``visual_embeddings`` (seed the residual stream) + and ``*_modulation``/``out_layer``/``audio_out_layer`` (elementwise- + consumed or final output). +- Attention: ``to_query``/``to_key``/``to_value`` are three separate + ``ColumnParallelLinear``s (``gather_output=False``), the output + projection is a ``RowParallelLinear``. +- ``query_norm``/``key_norm`` stay plain (non-distributed) RMSNorm over + ``head_dim``: K6 reshapes into per-head vectors *before* normalizing, so + the norm never crosses a head boundary and needs no cross-rank + reduction — each rank's local heads are already independent under + standard head-parallel tensor parallelism. + +Attention *computation* itself (dense SDPA/flash dispatch and NABLA +block-sparse attention) is reused unmodified from +``core/components/attention/dispatch.py`` and +``core/components/attention/nabla.py`` (inlined by the template) rather +than routed through vLLM-Omni's pluggable ``Attention``/backend-registry +layer or Diffusers' ``dispatch_attention_fn`` — that deeper integration +(paged KV-cache metadata, a registered NABLA backend analogous to +FASTVIDEO_VSA) is follow-up work once this TP-sharded weight path is +validated, not a prerequisite for correct tensor-parallel serving. +""" + +# The assembler inlines core/utils/tensors.py (apply_scale_shift_norm, +# apply_gate_sum, apply_rotary, fractal_flatten, fractal_unflatten, +# get_freqs), core/components/rope.py (RoPE1D, RoPE3D), +# core/components/attention/dispatch.py (SelfAttentionEngine), and +# core/components/attention/nabla.py (nabla_block_mask) into the same +# generated module. They are intentionally resolved there. +# ruff: noqa: F821 + + +# --------------------------------------------------------------------------- +# Small embedding / projection modules +# --------------------------------------------------------------------------- + + +class Kandinsky6TimeEmbeddings(nn.Module): + """Sinusoidal time embedding -> tensor-parallel MLP (Column -> Row pair).""" + + def __init__( + self, + model_dim: int, + time_dim: int, + max_period: float = 10000.0, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + if model_dim % 2: + raise ValueError("model_dim must be even") + self.register_buffer("freqs", get_freqs(model_dim // 2, max_period), persistent=False) + self.in_layer = ColumnParallelLinear( + model_dim, + time_dim, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.in_layer" if prefix else "in_layer", + ) + self.activation = nn.SiLU() + self.out_layer = RowParallelLinear( + time_dim, + time_dim, + bias=True, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.out_layer" if prefix else "out_layer", + ) + + def forward(self, time: Tensor) -> Tensor: + # Sinusoidal embedding is built in fp32 (its args have large + # magnitude). The MLP itself runs through the layers' own forward + # (their native, e.g. bf16, dtype) rather than core's manual fp32 + # upcast trick, so tensor-parallel all-reduce/bias-add for + # `out_layer` (RowParallelLinear) stay correct without hand-rolling + # its reduction internals. + freqs = self.freqs.to(device=time.device, dtype=torch.float32) + args = torch.outer(time.float(), freqs) + embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + h = self.activation(self.in_layer(embed.to(dtype=self.in_layer.weight.dtype))) + return self.out_layer(h) + + +class Kandinsky6TextEmbeddings(nn.Module): + """Projects raw text/pooled features into the (full, replicated) residual stream.""" + + def __init__( + self, + text_dim: int, + model_dim: int, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.in_layer = ColumnParallelLinear( + text_dim, + model_dim, + bias=True, + gather_output=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.in_layer" if prefix else "in_layer", + ) + self.norm = nn.LayerNorm(model_dim, elementwise_affine=True) + + def forward(self, x: Tensor) -> Tensor: + wdtype = self.in_layer.weight.dtype + x = self.in_layer(x.to(dtype=wdtype)) + return self.norm(x).to(dtype=x.dtype) + + +class Kandinsky6VisualEmbeddings(nn.Module): + """Patchifies video latents and projects into the (full) residual stream.""" + + def __init__( + self, + visual_dim: int, + model_dim: int, + patch_size: tuple, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.patch_size = patch_size + self.in_layer = ColumnParallelLinear( + math.prod(patch_size) * visual_dim, + model_dim, + bias=True, + gather_output=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.in_layer" if prefix else "in_layer", + ) + + def forward(self, x: Tensor) -> Tensor: + duration, height, width, channels = x.shape + p_t, p_h, p_w = self.patch_size + x = ( + x.view(duration // p_t, p_t, height // p_h, p_h, width // p_w, p_w, channels) + .permute(0, 2, 4, 1, 3, 5, 6) + .flatten(3, 6) + ) + return self.in_layer(x.to(dtype=self.in_layer.weight.dtype)) + + +class Kandinsky6Modulation(nn.Module): + """AdaLN modulation. Always gathers a full output: its ``shift``/``scale``/ + ``gate`` chunks are applied elementwise to the full residual-stream + tensor, not matrix-multiplied further, so a sharded output would chunk + incorrectly across ranks.""" + + def __init__( + self, + time_dim: int, + model_dim: int, + num_params: int, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.activation = nn.SiLU() + self.out_layer = ColumnParallelLinear( + time_dim, + num_params * model_dim, + bias=True, + gather_output=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.out_layer" if prefix else "out_layer", + ) + nn.init.zeros_(self.out_layer.weight) + if self.out_layer.bias is not None: + nn.init.zeros_(self.out_layer.bias) + + def forward(self, x: Tensor) -> Tensor: + out = self.out_layer(self.activation(x.float()).to(dtype=self.out_layer.weight.dtype)) + return out.to(dtype=x.dtype) + + +class Kandinsky6FeedForward(nn.Module): + """Tensor-parallel FFN (Column -> Row pair).""" + + def __init__( + self, + dim: int, + ff_dim: int, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.in_layer = ColumnParallelLinear( + dim, + ff_dim, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.in_layer" if prefix else "in_layer", + ) + self.activation = nn.GELU() + self.out_layer = RowParallelLinear( + ff_dim, + dim, + bias=False, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.out_layer" if prefix else "out_layer", + ) + + def forward(self, x: Tensor) -> Tensor: + return self.out_layer(self.activation(self.in_layer(x))) + + +# --------------------------------------------------------------------------- +# Attention (unified self/cross, matching the Diffusers-exported layout) +# --------------------------------------------------------------------------- + + +def _partition_visual_blocks(num_blocks: int, factory) -> tuple[int, int, nn.ModuleList]: + """Keep every rank's block index stable; non-local slots are ``PPMissingLayer``.""" + from vllm.model_executor.models.utils import PPMissingLayer + + world = _parallel_size("pp") + if world <= 1: + start, end = 0, num_blocks + else: + from vllm.distributed.utils import get_pp_indices + + start, end = get_pp_indices(num_blocks, _parallel_size("pp_rank"), world) + blocks = [factory(i) if start <= i < end else PPMissingLayer() for i in range(num_blocks)] + return start, end, nn.ModuleList(blocks) + + +def _maybe_sequence_parallel_qkv( + query: Tensor, + key: Tensor, + value: Tensor, + *, + is_self_attention: bool, +) -> tuple[Tensor, Tensor, Tensor, Any, Any]: + """Ulysses all-to-all (and optional Ring) for sharded visual self-attention. + + Text self-attention is not sequence-sharded, so the all-to-all runs only + while ``_VISUAL_SP`` is set. Returns ``(q, k, v, ring_group, ulysses_group)``. + """ + if not is_self_attention or not _VISUAL_SP.get() or query.dim() != 4: + return query, key, value, None, None + ulysses = _parallel_size("ulysses") + ring = _parallel_size("ring") + if ulysses <= 1 and ring <= 1: + return query, key, value, None, None + from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D + from vllm_omni.diffusion.distributed.parallel_state import get_sp_group + + sp = get_sp_group() + ulysses_group = None + if ulysses > 1: + ulysses_group = sp.ulysses_group + query = SeqAllToAll4D.apply(ulysses_group, query, 2, 1, False) + key = SeqAllToAll4D.apply(ulysses_group, key, 2, 1, False) + value = SeqAllToAll4D.apply(ulysses_group, value, 2, 1, False) + ring_group = sp.ring_group if ring > 1 else None + return query, key, value, ring_group, ulysses_group + + +class Kandinsky6Attention(nn.Module): + """K6 attention — self-attention when ``encoder_hidden_states`` is + omitted, cross-attention otherwise. Separate ``to_query``/``to_key``/ + ``to_value`` projections (not a fused QKV matrix) to match the + Diffusers-exported checkpoint's per-tensor layout.""" + + def __init__( + self, + num_channels: int, + head_dim: int, + kv_dim: int | None = None, + engine: str = "auto", + text_token_padding: bool = False, + visual: bool = False, + *, + sequence_parallel: bool = False, + role: str = "kandinsky6.text_self", + role_category: str = "self", + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + if num_channels % head_dim: + raise ValueError("num_channels must be divisible by head_dim") + kv_dim = kv_dim or num_channels + tp_size = get_tensor_model_parallel_world_size() + self.head_dim = head_dim + self.num_heads = (num_channels // head_dim) // tp_size + self.to_query = ColumnParallelLinear( + num_channels, + num_channels, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_query" if prefix else "to_query", + ) + self.to_key = ColumnParallelLinear( + kv_dim, + num_channels, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_key" if prefix else "to_key", + ) + self.to_value = ColumnParallelLinear( + kv_dim, + num_channels, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_value" if prefix else "to_value", + ) + self.query_norm = nn.RMSNorm(head_dim) + self.key_norm = nn.RMSNorm(head_dim) + self.out_layer = RowParallelLinear( + num_channels, + num_channels, + bias=True, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.out_layer" if prefix else "out_layer", + ) + self.visual = visual + self.attention_engine = engine + self.text_token_padding = text_token_padding + self.sequence_parallel = sequence_parallel + self.attn = Attention( + num_heads=self.num_heads, + head_size=self.head_dim, + causal=False, + softmax_scale=1.0 / (self.head_dim**0.5), + num_kv_heads=self.num_heads, + prefix=prefix, + role=role, + role_category=role_category, + scatter_idx=2, + gather_idx=1, + skip_sequence_parallel=not sequence_parallel, + ) + + def forward( + self, + hidden_states: Tensor, + encoder_hidden_states: Tensor | None = None, + attn_mask: Tensor | None = None, + rotary_emb: Tensor | None = None, + sparse_params: dict[str, Any] | None = None, + rope_q: Tensor | None = None, + rope_kv: Tensor | None = None, + **kwargs: Any, + ) -> Tensor: + if attn_mask is None: + attn_mask = kwargs.get("attention_mask") + rotary_emb_q = rope_q if rope_q is not None else rotary_emb + rotary_emb_kv = rope_kv if rope_kv is not None else rotary_emb + + is_self_attention = encoder_hidden_states is None + kv_input = hidden_states if is_self_attention else encoder_hidden_states + + query = self.to_query(hidden_states).reshape(*hidden_states.shape[:-1], self.num_heads, self.head_dim) + key = self.to_key(kv_input).reshape(*kv_input.shape[:-1], self.num_heads, self.head_dim) + value = self.to_value(kv_input).reshape(*kv_input.shape[:-1], self.num_heads, self.head_dim) + query = self.query_norm(query) + key = self.key_norm(key) + if rotary_emb_q is not None: + query = apply_rotary(query, rotary_emb_q).type_as(query) + if rotary_emb_kv is not None: + key = apply_rotary(key, rotary_emb_kv).type_as(key) + + if is_self_attention: + # FA3/FA2/SDPA dispatch requires exactly (B, S, H, D). Text + # self-attention's `hidden_states` is unbatched ((S, H, D) after + # reshape) and needs a batch dim added; visual/audio self-attention + # is already batched=1 ((1, N, H, D)) because `_embed_visual` / + # `_embed_audio` unsqueeze upstream. Adding a batch dim + # unconditionally would hand SDPA a 5-D (1, 1, N, H, D) tensor, + # which it happily accepts — but then it treats the *head* axis + # as the sequence and attends across heads instead of tokens + # (matches core/components/dit.py: MultiheadSelfAttentionEnc + # unsqueezes, MultiheadSelfAttentionDec does not). + strip_output_batch = query.dim() == 3 + if strip_output_batch: + query, key, value = query.unsqueeze(0), key.unsqueeze(0), value.unsqueeze(0) + else: + # Cross-attention: the query side (always visual, or the + # opposite modality for cross-modal attention) already carries + # a batch=1 dim; only the kv side may need one added to match + # (core/components/dit.py's MultiheadCrossAttention: "k/v may + # lack the batch dim (text has no batch, visual has batch=1)"). + # No dim is added beyond what's needed, so no stripping after. + if key.dim() < query.dim(): + key, value = key.unsqueeze(0), value.unsqueeze(0) + strip_output_batch = False + + if sparse_params is not None: + query, key, value, _ring_group, ulysses_group = _maybe_sequence_parallel_qkv( + query, key, value, is_self_attention=is_self_attention and self.sequence_parallel + ) + from torch.nn.attention.flex_attention import flex_attention + + q_ = query.transpose(1, 2).contiguous() + k_ = key.transpose(1, 2).contiguous() + v_ = value.transpose(1, 2).contiguous() + block_mask = nabla_block_mask(q_, k_, sparse_params["sta_mask"], thr=sparse_params["P"]) + out = flex_attention(q_, k_, v_, block_mask=block_mask).transpose(1, 2).contiguous() + if ulysses_group is not None: + from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D + + out = SeqAllToAll4D.apply(ulysses_group, out, 1, 2, False) + else: + metadata = AttentionMetadata(attn_mask=attn_mask) if attn_mask is not None else None + out = self.attn(query, key, value, metadata) + + if strip_output_batch: + out = out[0] + return self.out_layer(out.flatten(-2, -1)) + + +# --------------------------------------------------------------------------- +# Output layers +# --------------------------------------------------------------------------- + + +class Kandinsky6OutLayer(nn.Module): + """Projects model_dim -> pixel patches for the video output (full/gathered).""" + + def __init__( + self, + model_dim: int, + time_dim: int, + visual_dim: int, + patch_size: tuple, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.patch_size = patch_size + self.modulation = Kandinsky6Modulation( + time_dim, model_dim, 2, quant_config=quant_config, prefix=f"{prefix}.modulation" if prefix else "modulation" + ) + self.norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.out_layer = ColumnParallelLinear( + model_dim, + math.prod(patch_size) * visual_dim, + bias=True, + gather_output=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.out_layer" if prefix else "out_layer", + ) + + def forward(self, visual_embed: Tensor, time_embed: Tensor) -> Tensor: + shift, scale = torch.chunk(self.modulation(time_embed), 2, dim=-1) + x = apply_scale_shift_norm(self.norm, visual_embed, scale[:, None, None], shift[:, None, None]).type_as( + visual_embed + ) + x = self.out_layer(x.to(dtype=self.out_layer.weight.dtype)) + + duration, height, width, _ = x.shape + p_t, p_h, p_w = self.patch_size + return ( + x.view(duration, height, width, -1, p_t, p_h, p_w) + .permute(0, 4, 1, 5, 2, 6, 3) + .flatten(0, 1) + .flatten(1, 2) + .flatten(2, 3) + ) + + +class Kandinsky6OutLayerAudio(nn.Module): + """Projects model_dim_a -> audio latent channels (full/gathered).""" + + def __init__( + self, + model_dim: int, + time_dim: int, + audio_dim: int, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.modulation = Kandinsky6Modulation( + time_dim, model_dim, 2, quant_config=quant_config, prefix=f"{prefix}.modulation" if prefix else "modulation" + ) + self.norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.out_layer = ColumnParallelLinear( + model_dim, + audio_dim, + bias=True, + gather_output=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.out_layer" if prefix else "out_layer", + ) + + def forward(self, audio_embed: Tensor, time_embed: Tensor) -> Tensor: + shift, scale = torch.chunk(self.modulation(time_embed), 2, dim=-1) + x = apply_scale_shift_norm(self.norm, audio_embed, scale, shift).type_as(audio_embed) + x = self.norm(x) # matches reference training (double norm — do not remove for parity) + return self.out_layer(x.to(dtype=self.out_layer.weight.dtype)) + + +# --------------------------------------------------------------------------- +# Transformer blocks +# --------------------------------------------------------------------------- + + +class Kandinsky6TransformerEncoderBlock(nn.Module): + """Text-only self-attention + FFN block.""" + + def __init__( + self, + model_dim: int, + time_dim: int, + ff_dim: int, + head_dim: int, + engine: str = "auto", + text_token_padding: bool = False, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.text_modulation = Kandinsky6Modulation( + time_dim, + model_dim, + 6, + quant_config=quant_config, + prefix=f"{prefix}.text_modulation" if prefix else "text_modulation", + ) + self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.self_attention = Kandinsky6Attention( + model_dim, + head_dim, + engine=engine, + text_token_padding=text_token_padding, + role="kandinsky6.text_self", + role_category="self", + quant_config=quant_config, + prefix=f"{prefix}.self_attention" if prefix else "self_attention", + ) + self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.feed_forward = Kandinsky6FeedForward( + model_dim, ff_dim, quant_config=quant_config, prefix=f"{prefix}.feed_forward" if prefix else "feed_forward" + ) + + def forward(self, x: Tensor, time_embed: Tensor, rope: Tensor, attn_mask: Tensor | None = None) -> Tensor: + sa_params, ff_params = torch.chunk(self.text_modulation(time_embed), 2, dim=-1) + shift, scale, gate = torch.chunk(sa_params, 3, dim=-1) + x = apply_gate_sum( + x, + self.self_attention( + apply_scale_shift_norm(self.self_attention_norm, x, scale, shift), + rotary_emb=rope, + attn_mask=attn_mask, + ), + gate, + ) + shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) + return apply_gate_sum( + x, self.feed_forward(apply_scale_shift_norm(self.feed_forward_norm, x, scale, shift)), gate + ) + + +class Kandinsky6TransformerDecoderBlock(nn.Module): + """Visual self-attention + cross-attention to text + FFN block.""" + + def __init__( + self, + model_dim: int, + time_dim: int, + ff_dim: int, + head_dim: int, + engine: str = "auto", + text_token_padding: bool = False, + *, + self_sequence_parallel: bool = False, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.visual_modulation = Kandinsky6Modulation( + time_dim, + model_dim, + 9, + quant_config=quant_config, + prefix=f"{prefix}.visual_modulation" if prefix else "visual_modulation", + ) + self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.self_attention = Kandinsky6Attention( + model_dim, + head_dim, + engine=engine, + visual=True, + sequence_parallel=self_sequence_parallel, + role="kandinsky6.visual_self" if self_sequence_parallel else "kandinsky6.audio_self", + role_category="self", + quant_config=quant_config, + prefix=f"{prefix}.self_attention" if prefix else "self_attention", + ) + self.cross_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.cross_attention = Kandinsky6Attention( + model_dim, + head_dim, + kv_dim=model_dim, + engine=engine, + text_token_padding=text_token_padding, + role="kandinsky6.text_cross", + role_category="cross", + quant_config=quant_config, + prefix=f"{prefix}.cross_attention" if prefix else "cross_attention", + ) + self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.feed_forward = Kandinsky6FeedForward( + model_dim, ff_dim, quant_config=quant_config, prefix=f"{prefix}.feed_forward" if prefix else "feed_forward" + ) + + def forward( + self, + vis: Tensor, + text: Tensor, + time_embed: Tensor, + rope: Tensor, + sparse_params: dict | None, + attn_mask=None, + ) -> Tensor: + sa_params, ca_params, ff_params = torch.chunk(self.visual_modulation(time_embed), 3, dim=-1) + shift, scale, gate = torch.chunk(sa_params, 3, dim=-1) + vis = apply_gate_sum( + vis, + self.self_attention( + apply_scale_shift_norm(self.self_attention_norm, vis, scale, shift), + rotary_emb=rope, + sparse_params=sparse_params, + ), + gate, + ) + shift, scale, gate = torch.chunk(ca_params, 3, dim=-1) + vis = apply_gate_sum( + vis, + self.cross_attention( + apply_scale_shift_norm(self.cross_attention_norm, vis, scale, shift), + encoder_hidden_states=text, + attn_mask=attn_mask, + ), + gate, + ) + shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) + return apply_gate_sum( + vis, self.feed_forward(apply_scale_shift_norm(self.feed_forward_norm, vis, scale, shift)), gate + ) + + +class Kandinsky6FusedTransformerDecoderBlock(nn.Module): + """Fused video + audio block with cross-modal attention (T2VA backbone).""" + + def __init__( + self, + model_dim: int, + time_dim: int, + ff_dim: int, + head_dim: int, + model_dim_a: int, + time_dim_a: int, + ff_dim_a: int, + head_dim_a: int, + engine: str = "auto", + text_token_padding: bool = False, + ca_rope: bool = False, + cross_gates: bool = False, + fix_modulation: bool = False, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.video_dec_block = Kandinsky6TransformerDecoderBlock( + model_dim, + time_dim, + ff_dim, + head_dim, + engine, + text_token_padding, + self_sequence_parallel=True, + quant_config=quant_config, + prefix=f"{prefix}.video_dec_block" if prefix else "video_dec_block", + ) + self.audio_dec_block = Kandinsky6TransformerDecoderBlock( + model_dim_a, + time_dim_a, + ff_dim_a, + head_dim_a, + engine, + text_token_padding, + self_sequence_parallel=False, + quant_config=quant_config, + prefix=f"{prefix}.audio_dec_block" if prefix else "audio_dec_block", + ) + + self.va_cross_attention = Kandinsky6Attention( + model_dim, + head_dim, + model_dim_a, + engine, + role="kandinsky6.video_audio_cross", + role_category="cross", + quant_config=quant_config, + prefix=f"{prefix}.va_cross_attention" if prefix else "va_cross_attention", + ) + self.av_cross_attention = Kandinsky6Attention( + model_dim_a, + head_dim_a, + model_dim, + engine, + role="kandinsky6.audio_video_cross", + role_category="cross", + quant_config=quant_config, + prefix=f"{prefix}.av_cross_attention" if prefix else "av_cross_attention", + ) + + self.va_modulation = Kandinsky6Modulation( + time_dim, + model_dim if not cross_gates else model_dim * 2 + model_dim_a, + 1 if cross_gates else 3, + quant_config=quant_config, + prefix=f"{prefix}.va_modulation" if prefix else "va_modulation", + ) + self.av_modulation = Kandinsky6Modulation( + time_dim_a, + model_dim_a if not cross_gates else model_dim_a * 2 + model_dim, + 1 if cross_gates else 3, + quant_config=quant_config, + prefix=f"{prefix}.av_modulation" if prefix else "av_modulation", + ) + self.va_normalization = nn.LayerNorm(model_dim, elementwise_affine=False) + self.av_normalization = nn.LayerNorm(model_dim_a, elementwise_affine=False) + + self.ca_rope = ca_rope + self.cross_gates = cross_gates + self.fix_modulation = fix_modulation + self.model_dim = model_dim + self.model_dim_a = model_dim_a + + def forward( + self, + vis: Tensor, + aud: Tensor, + text_v: Tensor, + text_a: Tensor, + time_embed, # (video_time, audio_time) tuple + vis_rope: Tensor, + aud_rope: Tensor, + sparse_params: dict | None, + attn_mask=None, + modality_mask=None, + av_gate_scale: float = 1.0, + va_gate_scale: float = 1.0, + ): + fake_audio = modality_mask[0] if modality_mask is not None else 0 + fake_video = modality_mask[1] if modality_mask is not None else 0 + t_v, t_a = time_embed + + # ---- video backbone ---- + if vis is not None: + sa_p, ca_p, ff_p = torch.chunk(self.video_dec_block.visual_modulation(t_v), 3, dim=-1) + + shift, scale, gate = torch.chunk(sa_p, 3, dim=-1) + vis = apply_gate_sum( + vis, + self.video_dec_block.self_attention( + apply_scale_shift_norm(self.video_dec_block.self_attention_norm, vis, scale, shift), + rotary_emb=vis_rope, + sparse_params=sparse_params, + ), + gate, + ).type_as(vis) + + shift, scale, gate_v = torch.chunk(ca_p, 3, dim=-1) + vis_pre_ca = apply_scale_shift_norm(self.video_dec_block.cross_attention_norm, vis, scale, shift).type_as( + vis + ) + vis_out_t = self.video_dec_block.cross_attention( + vis_pre_ca, encoder_hidden_states=text_v, attn_mask=attn_mask + ) + + # ---- audio backbone ---- + if aud is not None: + sa_p, ca_p, ff_p_a = torch.chunk(self.audio_dec_block.visual_modulation(t_a), 3, dim=-1) + + shift, scale, gate = torch.chunk(sa_p, 3, dim=-1) + aud = apply_gate_sum( + aud, + self.audio_dec_block.self_attention( + apply_scale_shift_norm(self.audio_dec_block.self_attention_norm, aud, scale, shift), + rotary_emb=aud_rope, + ), + gate, + ).type_as(aud) + + shift, scale, gate_a = torch.chunk(ca_p, 3, dim=-1) + aud_pre_ca = apply_scale_shift_norm(self.audio_dec_block.cross_attention_norm, aud, scale, shift).type_as( + aud + ) + aud_out_t = self.audio_dec_block.cross_attention( + aud_pre_ca, encoder_hidden_states=text_a, attn_mask=attn_mask + ) + aud = apply_gate_sum(aud, aud_out_t, gate_a).type_as(aud) + + # ---- cross-modal attention ---- + if vis is not None: + t_va_mod = t_a if not self.fix_modulation else t_v + t_av_mod = t_v if not self.fix_modulation else t_a + va_params = self.va_modulation(t_va_mod) + av_params = self.av_modulation(t_av_mod) + + if self.cross_gates: + va_shift, va_scale, va_gate = torch.split( + va_params, [self.model_dim, self.model_dim, self.model_dim_a], dim=-1 + ) + av_shift, av_scale, av_gate = torch.split( + av_params, [self.model_dim_a, self.model_dim_a, self.model_dim], dim=-1 + ) + else: + va_shift, va_scale, va_gate = torch.chunk(va_params, 3, dim=-1) + av_shift, av_scale, av_gate = torch.chunk(av_params, 3, dim=-1) + + vis = apply_gate_sum(vis, vis_out_t, gate_v).type_as(vis) + vis_for_va = apply_scale_shift_norm(self.va_normalization, vis, va_scale, va_shift).type_as(vis) + aud_for_av = apply_scale_shift_norm(self.av_normalization, aud, av_scale, av_shift).type_as(aud) + + # V->A and A->V attention. Visual tokens may be sequence-sharded; + # audio queries need the full visual key/value sequence. + rq_v = vis_rope if self.ca_rope else None + rk_a = aud_rope if self.ca_rope else None + vis_kv = vis_pre_ca + vis_rope_kv = rq_v + if _parallel_size("sp") > 1 and vis_pre_ca is not None: + from vllm_omni.diffusion.distributed.sp_sharding import sp_gather + + vis_kv = sp_gather(vis_pre_ca, dim=1) + if vis_rope_kv is not None: + vis_rope_kv = sp_gather(vis_rope, dim=0) + vis_from_aud = ( + self.va_cross_attention(vis_for_va, encoder_hidden_states=aud_pre_ca, rope_q=rq_v, rope_kv=rk_a) + * (1 - fake_audio) + * (1 - fake_video) + ) + aud_from_vis = ( + self.av_cross_attention(aud_for_av, encoder_hidden_states=vis_kv, rope_q=rk_a, rope_kv=vis_rope_kv) + * (1 - fake_audio) + * (1 - fake_video) + ) + + va_g = (va_gate if not self.cross_gates else av_gate) * va_gate_scale + av_g = (av_gate if not self.cross_gates else va_gate) * av_gate_scale + vis = apply_gate_sum(vis, vis_from_aud, va_g).type_as(vis) + aud = apply_gate_sum(aud, aud_from_vis, av_g).type_as(aud) + else: + vis = apply_gate_sum(vis, vis_out_t, gate_v).type_as(vis) + + # ---- FFN ---- + if vis is not None: + shift, scale, gate = torch.chunk(ff_p, 3, dim=-1) + vis = apply_gate_sum( + vis, + self.video_dec_block.feed_forward( + apply_scale_shift_norm(self.video_dec_block.feed_forward_norm, vis, scale, shift) + ), + gate, + ).type_as(vis) + + if aud is not None: + shift, scale, gate = torch.chunk(ff_p_a, 3, dim=-1) + aud = apply_gate_sum( + aud, + self.audio_dec_block.feed_forward( + apply_scale_shift_norm(self.audio_dec_block.feed_forward_norm, aud, scale, shift) + ), + gate, + ).type_as(aud) + + return vis, aud + + +# --------------------------------------------------------------------------- +# Unified DiffusionTransformer3D (tensor-parallel) +# --------------------------------------------------------------------------- + + +class Kandinsky6Transformer3DModel(nn.Module): + """Kandinsky 6 DiT — handles T2V (is_multimodal=False) and T2VA (is_multimodal=True). + + Tensor-parallel port of ``core/components/dit.py::DiffusionTransformer3D``, + with submodule names/nesting matching the Diffusers-exported checkpoint + layout (see module docstring) so ``from_pretrained`` loads a + ``convert_checkpoint.py --use-patched-diffusers`` bundle directly. + + The class is a plain ``nn.Module``. ``from_diffusers_config`` reads the + bundle ``config.json``. Safetensors from ``transformer/`` are applied + later by the pipeline ``load_weights`` path, which narrows full + checkpoint tensors onto each rank's ``ColumnParallelLinear`` / + ``RowParallelLinear`` shard. + """ + + _repeated_blocks = [ + "Kandinsky6FusedTransformerDecoderBlock", + "Kandinsky6TransformerEncoderBlock", + "Kandinsky6TransformerDecoderBlock", + ] + + # MagCache must not treat text and visual ModuleLists as one residual chain. + _magcache_block_attrs = ("visual_transformer_blocks",) + + @classmethod + def from_diffusers_config( + cls, + config: dict, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> Kandinsky6Transformer3DModel: + """Build the DiT from a Diffusers ``config.json`` dict.""" + fields = ( + "in_visual_dim", + "out_visual_dim", + "in_text_dim", + "in_text_dim2", + "time_dim", + "patch_size", + "model_dim", + "ff_dim", + "num_text_blocks", + "num_visual_blocks", + "axes_dims", + "visual_cond", + "is_multimodal", + "in_audio_dim", + "model_dim_a", + "time_dim_a", + "ff_dim_a", + "axes_dims_a", + "audio_freqs_scaling", + "attention_engine", + "text_token_padding", + "ca_rope", + "cross_gates", + "fix_modulation", + "visual_token_type_num_embeddings", + ) + kwargs = {key: config[key] for key in fields if key in config} + for key in ("patch_size", "axes_dims", "axes_dims_a"): + if isinstance(kwargs.get(key), list): + kwargs[key] = tuple(kwargs[key]) + return cls(**kwargs, quant_config=quant_config, prefix=prefix) + + @staticmethod + def _is_transformer_block(name: str, module: nn.Module) -> bool: + del module + leaf = name.rsplit(".", 1)[-1] + if not leaf.isdigit(): + return False + return any( + part in name + for part in ( + "visual_transformer_blocks", + "text_transformer_blocks", + "video_text_transformer_blocks", + "audio_text_transformer_blocks", + ) + ) + + _hsdp_shard_conditions = [_is_transformer_block] + + # Visual tokens are (1, N, D). RoPE after fractal flatten is (N, 1, C, 2, 2). + # Text and audio stay replicated; cross-modal KV is gathered in the fused block. + _sp_plan = { + "_sp_visual_shard": { + 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True), + }, + "_sp_visual_rope": { + 0: SequenceParallelInput(split_dim=0, expected_dims=5, split_output=True, auto_pad=True), + }, + "_sp_visual_gather": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + + def __init__( + self, + in_visual_dim: int = 16, + out_visual_dim: int = 16, + in_text_dim: int = 3584, + in_text_dim2: int = 768, + time_dim: int = 1024, + patch_size: tuple = (1, 2, 2), + model_dim: int = 4096, + ff_dim: int = 16384, + num_text_blocks: int = 4, + num_visual_blocks: int = 60, + axes_dims: tuple = (32, 48, 48), + visual_cond: bool = True, + is_multimodal: bool = False, + # Audio (T2VA only) + in_audio_dim: int = 20, + model_dim_a: int | None = None, + time_dim_a: int | None = None, + ff_dim_a: int | None = None, + axes_dims_a: tuple | None = None, + audio_freqs_scaling: float = 1.0, + # Misc + attention_engine: str = "auto", + text_token_padding: bool = False, + ca_rope: bool = False, + cross_gates: bool = False, + fix_modulation: bool = False, + # I2VA: 0 = off; 2 = generated vs reference frame (tail_cond_first_frame) + visual_token_type_num_embeddings: int = 0, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.patch_size = patch_size + self.visual_cond = visual_cond + self.is_multimodal = is_multimodal + self.in_visual_dim = in_visual_dim + self.in_audio_dim = in_audio_dim + self.text_token_padding = text_token_padding + self.visual_token_type_num_embeddings = int(visual_token_type_num_embeddings or 0) + # Time-independent text/pooled projections (cleared each generation). + self._text_proj_cache: dict[tuple, object] = {} + + head_dim = sum(axes_dims) + + # Effective audio dims (default to video dims) + model_dim_a = model_dim_a or model_dim + time_dim_a = time_dim_a or time_dim + ff_dim_a = ff_dim_a or ff_dim + axes_dims_a = axes_dims_a or axes_dims + head_dim_a = sum(axes_dims_a) + + # ---- visual backbone (shared) ---- + vis_in_dim = (2 * in_visual_dim + 1) if visual_cond else in_visual_dim + self.visual_embeddings = Kandinsky6VisualEmbeddings( + vis_in_dim, model_dim, patch_size, quant_config=quant_config, prefix=f"{prefix}.visual_embeddings" + ) + if self.visual_token_type_num_embeddings > 0: + self.visual_token_type_embeddings = nn.Embedding(self.visual_token_type_num_embeddings, model_dim) + self.visual_rope_embeddings = RoPE3D(axes_dims) + self.out_layer = Kandinsky6OutLayer( + model_dim, time_dim, out_visual_dim, patch_size, quant_config=quant_config, prefix=f"{prefix}.out_layer" + ) + + if not is_multimodal: + # T2V: single text/time embedding branch + self.time_embeddings = Kandinsky6TimeEmbeddings( + model_dim, time_dim, quant_config=quant_config, prefix=f"{prefix}.time_embeddings" + ) + self.text_embeddings = Kandinsky6TextEmbeddings( + in_text_dim, model_dim, quant_config=quant_config, prefix=f"{prefix}.text_embeddings" + ) + self.pooled_text_embeddings = Kandinsky6TextEmbeddings( + in_text_dim2, time_dim, quant_config=quant_config, prefix=f"{prefix}.pooled_text_embeddings" + ) + self.text_rope_embeddings = RoPE1D(head_dim) + self.text_transformer_blocks = nn.ModuleList( + [ + Kandinsky6TransformerEncoderBlock( + model_dim, + time_dim, + ff_dim, + head_dim, + attention_engine, + text_token_padding, + quant_config=quant_config, + prefix=f"{prefix}.text_transformer_blocks.{i}", + ) + for i in range(num_text_blocks) + ] + ) + + def _decoder_block(i: int) -> Kandinsky6TransformerDecoderBlock: + return Kandinsky6TransformerDecoderBlock( + model_dim, + time_dim, + ff_dim, + head_dim, + attention_engine, + text_token_padding, + self_sequence_parallel=True, + quant_config=quant_config, + prefix=f"{prefix}.visual_transformer_blocks.{i}", + ) + + self._pp_block_start, self._pp_block_end, self.visual_transformer_blocks = _partition_visual_blocks( + num_visual_blocks, _decoder_block + ) + else: + # T2VA: dual (video / audio) text+time branches + fused blocks + self.audio_embeddings = Kandinsky6TextEmbeddings( + in_audio_dim, model_dim_a, quant_config=quant_config, prefix=f"{prefix}.audio_embeddings" + ) + self.audio_rope_embeddings = RoPE1D(head_dim_a, freqs_scaling=audio_freqs_scaling) + self.audio_out_layer = Kandinsky6OutLayerAudio( + model_dim_a, time_dim_a, in_audio_dim, quant_config=quant_config, prefix=f"{prefix}.audio_out_layer" + ) + + for mod_prefix, md, td, fd, hd in [ + ("video", model_dim, time_dim, ff_dim, head_dim), + ("audio", model_dim_a, time_dim_a, ff_dim_a, head_dim_a), + ]: + setattr( + self, + f"{mod_prefix}_time_embeddings", + Kandinsky6TimeEmbeddings( + md, td, quant_config=quant_config, prefix=f"{prefix}.{mod_prefix}_time_embeddings" + ), + ) + setattr( + self, + f"{mod_prefix}_text_embeddings", + Kandinsky6TextEmbeddings( + in_text_dim, md, quant_config=quant_config, prefix=f"{prefix}.{mod_prefix}_text_embeddings" + ), + ) + setattr( + self, + f"{mod_prefix}_pooled_text_embeddings", + Kandinsky6TextEmbeddings( + in_text_dim2, + td, + quant_config=quant_config, + prefix=f"{prefix}.{mod_prefix}_pooled_text_embeddings", + ), + ) + setattr(self, f"{mod_prefix}_text_rope_embeddings", RoPE1D(hd)) + setattr( + self, + f"{mod_prefix}_text_transformer_blocks", + nn.ModuleList( + [ + Kandinsky6TransformerEncoderBlock( + md, + td, + fd, + hd, + attention_engine, + text_token_padding, + quant_config=quant_config, + prefix=f"{prefix}.{mod_prefix}_text_transformer_blocks.{i}", + ) + for i in range(num_text_blocks) + ] + ), + ) + + def _fused_block(i: int) -> Kandinsky6FusedTransformerDecoderBlock: + return Kandinsky6FusedTransformerDecoderBlock( + model_dim, + time_dim, + ff_dim, + head_dim, + model_dim_a, + time_dim_a, + ff_dim_a, + head_dim_a, + attention_engine, + text_token_padding, + ca_rope=ca_rope, + cross_gates=cross_gates, + fix_modulation=fix_modulation, + quant_config=quant_config, + prefix=f"{prefix}.visual_transformer_blocks.{i}", + ) + + self._pp_block_start, self._pp_block_end, self.visual_transformer_blocks = _partition_visual_blocks( + num_visual_blocks, _fused_block + ) + + from vllm.model_executor.models.utils import PPMissingLayer + + # Embeddings live on the first PP stage, heads on the last. Text/time + # modules stay on every rank (they are small next to the visual stack). + if _parallel_size("pp") > 1 and _parallel_size("pp_rank") != 0: + self.visual_embeddings = PPMissingLayer() + if hasattr(self, "audio_embeddings"): + self.audio_embeddings = PPMissingLayer() + if _parallel_size("pp") > 1 and _parallel_size("pp_rank") != _parallel_size("pp") - 1: + self.out_layer = PPMissingLayer() + if hasattr(self, "audio_out_layer"): + self.audio_out_layer = PPMissingLayer() + self._sp_visual_shard = nn.Identity() + self._sp_visual_rope = nn.Identity() + self._sp_visual_gather = nn.Identity() + self._pp_final = None + + # ------------------------------------------------------------------ + # Stage helpers (shared by forward / MagCache) — unchanged from core. + # ------------------------------------------------------------------ + + def clear_text_proj_cache(self) -> None: + """Drop cached text/pooled projections (call once per generation).""" + self._text_proj_cache.clear() + + def _project_pooled(self, prefix: str | None, pooled: Tensor) -> Tensor: + key = ("pe", prefix, pooled.data_ptr(), tuple(pooled.shape)) + hit = self._text_proj_cache.get(key) + if hit is not None: + return hit # type: ignore[return-value] + if prefix is None: + pe = self.pooled_text_embeddings(pooled) + else: + pe = getattr(self, f"{prefix}_pooled_text_embeddings")(pooled) + self._text_proj_cache[key] = pe + return pe + + def _project_text_tokens( + self, + prefix: str | None, + text_embed: Tensor, + pooled: Tensor, + ) -> tuple[Tensor, Tensor]: + """Time-independent token + pooled Linears (cached within a generation).""" + key = ( + "te", + prefix, + text_embed.data_ptr(), + pooled.data_ptr(), + tuple(text_embed.shape), + tuple(pooled.shape), + ) + hit = self._text_proj_cache.get(key) + if hit is not None: + return hit # type: ignore[return-value] + if prefix is None: + te = self.text_embeddings(text_embed) + else: + te = getattr(self, f"{prefix}_text_embeddings")(text_embed) + pe = self._project_pooled(prefix, pooled) + self._text_proj_cache[key] = (te, pe) + return te, pe + + def _time_embed(self, prefix: str | None, time: Tensor, pooled_proj: Tensor) -> Tensor: + if prefix is None: + return self.time_embeddings(time) + pooled_proj + return getattr(self, f"{prefix}_time_embeddings")(time) + pooled_proj + + @staticmethod + def _normalize_attn_mask(attn_mask: Tensor | None) -> Tensor | None: + """HF/K5 key-padding (True=valid): ``(S,)`` -> ``(1, S)`` for SDPA.""" + if attn_mask is None: + return None + if attn_mask.dim() == 1: + return attn_mask.unsqueeze(0) + return attn_mask + + def _run_text_blocks( + self, + prefix: str | None, + te: Tensor, + tm: Tensor, + text_rope: Tensor, + attn_mask: Tensor | None = None, + ) -> Tensor: + blocks = self.text_transformer_blocks if prefix is None else getattr(self, f"{prefix}_text_transformer_blocks") + for blk in blocks: + te = blk(te, tm, text_rope, attn_mask) + return te + + def _encode_text( + self, + prefix: str, + text_embed: Tensor, + pooled: Tensor, + time: Tensor, + text_rope: Tensor, + attn_mask: Tensor | None = None, + ) -> tuple[Tensor, Tensor]: + te, pe = self._project_text_tokens(prefix, text_embed, pooled) + tm = self._time_embed(prefix, time, pe) + te = self._run_text_blocks(prefix, te, tm, text_rope, attn_mask) + return te, tm + + def _encode_t2v( + self, + text_embed: Tensor, + pooled: Tensor, + time: Tensor, + text_rope: Tensor, + attn_mask: Tensor | None = None, + ) -> tuple[Tensor, Tensor]: + te, pe = self._project_text_tokens(None, text_embed, pooled) + tm = self._time_embed(None, time, pe) + te = self._run_text_blocks(None, te, tm, text_rope, attn_mask) + return te, tm + + def _time_only(self, prefix: str | None, pooled: Tensor, time: Tensor) -> Tensor: + """Pooled (cached) + time — OutLayer path when MagCache skips text+visual blocks.""" + pe = self._project_pooled(prefix, pooled) + return self._time_embed(prefix, time, pe) + + def _embed_visual( + self, + x_video: Tensor, + visual_rope: Tensor, + sparse_params: dict | None, + *, + apply_fractal: bool = True, + visual_token_type_ids: Tensor | None = None, + ) -> tuple[Tensor, tuple, Tensor]: + vis_embed = self.visual_embeddings(x_video) + if hasattr(self, "visual_token_type_embeddings") and visual_token_type_ids is not None: + if visual_token_type_ids.shape[0] != vis_embed.shape[0]: + raise ValueError( + "visual_token_type_ids length must match visual latent frames: " + f"type_ids={tuple(visual_token_type_ids.shape)}, " + f"visual={tuple(vis_embed.shape)}" + ) + vis_embed = ( + vis_embed + + self.visual_token_type_embeddings(visual_token_type_ids.to(device=vis_embed.device))[:, None, None, :] + ) + vis_shape = vis_embed.shape[:-1] + vis_rope = visual_rope + if apply_fractal: + to_fractal = sparse_params["to_fractal"] if sparse_params else False + vis_embed, vis_rope = fractal_flatten(vis_embed, vis_rope, vis_shape, block_mask=to_fractal) + else: + vis_embed, vis_rope = fractal_flatten(vis_embed, vis_rope, vis_shape) + vis_embed = vis_embed.unsqueeze(0) + return vis_embed, vis_shape, vis_rope + + def _embed_audio(self, x_audio: Tensor, audio_rope: Tensor) -> tuple[Tensor, Tensor]: + aud_embed = self.audio_embeddings(x_audio).unsqueeze(0) + return aud_embed, audio_rope + + def _sp_enter(self, vis: Tensor | None, rope: Tensor | None) -> tuple[Tensor | None, Tensor | None]: + if vis is None or (_parallel_size("pp") > 1 and _parallel_size("pp_rank") != 0): + return vis, rope + vis = self._sp_visual_shard(vis) + if rope is not None: + rope = self._sp_visual_rope(rope) + return vis, rope + + def _sp_exit(self, vis: Tensor | None) -> Tensor | None: + if vis is None: + return None + if _parallel_size("pp") > 1 and _parallel_size("pp_rank") != _parallel_size("pp") - 1: + return vis + return self._sp_visual_gather(vis) + + def _pp_recv_hidden( + self, + vis: Tensor | None, + aud: Tensor | None, + vis_rope: Tensor | None = None, + aud_rope: Tensor | None = None, + ) -> tuple[Tensor | None, Tensor | None, Tensor | None, Tensor | None]: + pp = _parallel_size("pp") + rank = _parallel_size("pp_rank") + if pp <= 1 or rank == 0: + return vis, aud, vis_rope, aud_rope + from vllm_omni.diffusion.distributed.parallel_state import get_pp_group + + payload = get_pp_group().recv_tensor_dict(src=rank - 1) + assert payload is not None + if payload.get("shape") is not None: + self._pp_vis_shape = payload["shape"] + return ( + payload.get("vis"), + payload.get("aud"), + payload.get("vis_rope", vis_rope), + payload.get("aud_rope", aud_rope), + ) + + def _pp_send_hidden_or_wait( + self, + vis: Tensor | None, + aud: Tensor | None, + vis_rope: Tensor | None = None, + aud_rope: Tensor | None = None, + ): + """Non-last PP ranks forward hidden states and wait for the velocity.""" + pp = _parallel_size("pp") + if pp <= 1: + return None + rank = _parallel_size("pp_rank") + if rank == pp - 1: + return None + from vllm_omni.diffusion.distributed.parallel_state import get_pp_group + + group = get_pp_group() + group.send_tensor_dict( + { + "vis": vis, + "aud": aud, + "vis_rope": vis_rope, + "aud_rope": aud_rope, + "shape": getattr(self, "_pp_vis_shape", None), + }, + dst=rank + 1, + ) + final = group.recv_tensor_dict(src=pp - 1) + assert final is not None + if "pair" in final: + return final["video"], final["audio"] + return final["video"] + + def _pp_publish(self, result: Tensor | tuple[Tensor, Tensor]): + pp = _parallel_size("pp") + if pp <= 1 or _parallel_size("pp_rank") != pp - 1: + return result + from vllm_omni.diffusion.distributed.parallel_state import get_pp_group + + group = get_pp_group() + if isinstance(result, tuple): + payload = {"pair": True, "video": result[0], "audio": result[1]} + else: + payload = {"video": result} + for dst in range(pp - 1): + group.send_tensor_dict(payload, dst=dst) + return result + + def _iter_local_visual_blocks(self): + start = getattr(self, "_pp_block_start", 0) + end = getattr(self, "_pp_block_end", len(self.visual_transformer_blocks)) + for index, block in enumerate(self.visual_transformer_blocks): + if index < start or index >= end: + continue + if type(block).__name__ == "PPMissingLayer": + continue + yield block + + def _run_visual_blocks_single( + self, + vis_embed: Tensor | None, + aud_embed: Tensor | None, + te: Tensor, + tm: Tensor, + vis_rope: Tensor | None, + aud_rope: Tensor | None, + sparse_params: dict | None, + attn_mask: Tensor | None = None, + ) -> tuple[Tensor | None, Tensor | None]: + self._pp_final = None + vis_embed, aud_embed, vis_rope, aud_rope = self._pp_recv_hidden(vis_embed, aud_embed, vis_rope, aud_rope) + vis_embed, vis_rope = self._sp_enter(vis_embed, vis_rope) + token = _VISUAL_SP.set(_parallel_size("sp") > 1 and vis_embed is not None) + try: + for blk in self._iter_local_visual_blocks(): + if self.is_multimodal: + if vis_embed is not None and aud_embed is None: + vis_embed, _ = blk( + vis_embed, + None, + te, + te, + (tm, tm), + vis_rope, + None, + sparse_params, + attn_mask, + ) + elif aud_embed is not None and vis_embed is None: + _, aud_embed = blk( + None, + aud_embed, + te, + te, + (tm, tm), + None, + aud_rope, + None, + attn_mask, + ) + else: + raise RuntimeError("single-modality fused path expects exactly one of video/audio") + elif vis_embed is not None: + vis_embed = blk(vis_embed, te, tm, vis_rope, sparse_params, attn_mask) + else: + aud_embed = blk(aud_embed, te, tm, aud_rope, None, attn_mask) + finally: + _VISUAL_SP.reset(token) + early = self._pp_send_hidden_or_wait(vis_embed, aud_embed, vis_rope, aud_rope) + if early is not None: + self._pp_final = early + return None, None + return vis_embed, aud_embed + + def _run_visual_blocks_fused( + self, + vis_embed: Tensor, + aud_embed: Tensor, + video_te: Tensor, + audio_te: Tensor, + video_tm: Tensor, + audio_tm: Tensor, + vis_rope: Tensor, + aud_rope: Tensor, + sparse_params: dict | None, + attn_mask: Tensor | None = None, + ) -> tuple[Tensor, Tensor]: + self._pp_final = None + vis_embed, aud_embed, vis_rope, aud_rope = self._pp_recv_hidden(vis_embed, aud_embed, vis_rope, aud_rope) + vis_embed, vis_rope = self._sp_enter(vis_embed, vis_rope) + token = _VISUAL_SP.set(_parallel_size("sp") > 1) + try: + for blk in self._iter_local_visual_blocks(): + vis_embed, aud_embed = blk( + vis_embed, + aud_embed, + video_te, + audio_te, + (video_tm, audio_tm), + vis_rope, + aud_rope, + sparse_params, + attn_mask, + ) + finally: + _VISUAL_SP.reset(token) + early = self._pp_send_hidden_or_wait(vis_embed, aud_embed, vis_rope, aud_rope) + if early is not None: + self._pp_final = early + return None, None # type: ignore[return-value] + return vis_embed, aud_embed + + def _project_video( + self, + vis_embed: Tensor, + vis_shape: tuple, + tm: Tensor, + sparse_params: dict | None, + ) -> Tensor: + to_fractal = sparse_params["to_fractal"] if sparse_params else False + if vis_shape is None: + vis_shape = getattr(self, "_pp_vis_shape", None) + vis_embed = self._sp_exit(vis_embed) + vis_embed = fractal_unflatten(vis_embed, vis_shape, block_mask=to_fractal) + return self._pp_publish(self.out_layer(vis_embed, tm)) + + def _project_audio(self, aud_embed: Tensor, tm: Tensor) -> Tensor: + return self._pp_publish(self.audio_out_layer(aud_embed.squeeze(0), tm)) + + def _project_fused( + self, + vis_embed: Tensor, + aud_embed: Tensor, + vis_shape: tuple, + video_tm: Tensor, + audio_tm: Tensor, + ) -> tuple[Tensor, Tensor]: + if vis_shape is None: + vis_shape = self._pp_vis_shape + vis_embed = self._sp_exit(vis_embed) + vis_embed = fractal_unflatten(vis_embed, vis_shape) + video_vel = self.out_layer(vis_embed, video_tm) + audio_vel = self.audio_out_layer(aud_embed.squeeze(0), audio_tm) + return self._pp_publish((video_vel, audio_vel)) + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + + def forward( + self, + x_video: Tensor | None, + x_audio: Tensor | None, + text_embed: Tensor | list[Tensor], + pooled_text_embed: Tensor | list[Tensor], + time: Tensor | list[Tensor], + visual_rope: Tensor | None, + audio_rope: Tensor | None, + text_rope: Tensor | list[Tensor], + sparse_params: dict | None = None, + attention_mask: Tensor | None = None, + visual_token_type_ids: Tensor | None = None, + ) -> Tensor | tuple[Tensor, Tensor]: + """RoPE tensors are precomputed outside (pipeline / export); not from positions.""" + both = x_video is not None and x_audio is not None and self.is_multimodal + attn_mask = self._normalize_attn_mask(attention_mask) + + if not both: + te_in = text_embed[0] if isinstance(text_embed, list) else text_embed + pe_in = pooled_text_embed[0] if isinstance(pooled_text_embed, list) else pooled_text_embed + rope_in = text_rope[0] if isinstance(text_rope, list) else text_rope + t_in = time[0] if isinstance(time, list) else time + + if self.is_multimodal: + prefix = "audio" if x_audio is not None else "video" + # Multimodal single-modality path: list text_rope is [video, audio]. + if isinstance(text_rope, list): + rope_in = text_rope[1] if prefix == "audio" else text_rope[0] + te, tm = self._encode_text(prefix, te_in, pe_in, t_in, rope_in, attn_mask) + else: + te, tm = self._encode_t2v(te_in, pe_in, t_in, rope_in, attn_mask) + + if x_video is not None: + if _parallel_size("pp") > 1 and _parallel_size("pp_rank") != 0: + vis_embed, vis_shape, vis_rope = None, getattr(self, "_pp_vis_shape", None), visual_rope + else: + vis_embed, vis_shape, vis_rope = self._embed_visual( + x_video, + visual_rope, + sparse_params, + apply_fractal=True, + visual_token_type_ids=visual_token_type_ids, + ) + self._pp_vis_shape = vis_shape + vis_embed, _ = self._run_visual_blocks_single( + vis_embed, + None, + te, + tm, + vis_rope, + None, + sparse_params, + attn_mask, + ) + if self._pp_final is not None: + final = self._pp_final + self._pp_final = None + return final + return self._project_video(vis_embed, self._pp_vis_shape, tm, sparse_params) + + aud_embed, aud_rope = self._embed_audio(x_audio, audio_rope) + _, aud_embed = self._run_visual_blocks_single( + None, + aud_embed, + te, + tm, + None, + aud_rope, + sparse_params, + attn_mask, + ) + if self._pp_final is not None: + final = self._pp_final + self._pp_final = None + return final + return self._project_audio(aud_embed, tm) + + te_v, pe_v = ( + (text_embed[0], pooled_text_embed[0]) if isinstance(text_embed, list) else (text_embed, pooled_text_embed) + ) + te_a, pe_a = ( + (text_embed[1], pooled_text_embed[1]) if isinstance(text_embed, list) else (text_embed, pooled_text_embed) + ) + if isinstance(text_rope, list): + rope_v, rope_a = text_rope[0], text_rope[1] + else: + rope_v = rope_a = text_rope + t_v, t_a = (time[0], time[1]) if isinstance(time, list) else (time, time) + + video_te, video_tm = self._encode_text("video", te_v, pe_v, t_v, rope_v, attn_mask) + audio_te, audio_tm = self._encode_text("audio", te_a, pe_a, t_a, rope_a, attn_mask) + + if _parallel_size("pp") > 1 and _parallel_size("pp_rank") != 0: + vis_embed, vis_shape, vis_rope = None, getattr(self, "_pp_vis_shape", None), visual_rope + aud_embed, aud_rope = None, audio_rope + else: + vis_embed, vis_shape, vis_rope = self._embed_visual( + x_video, + visual_rope, + sparse_params, + apply_fractal=False, + visual_token_type_ids=visual_token_type_ids, + ) + self._pp_vis_shape = vis_shape + aud_embed, aud_rope = self._embed_audio(x_audio, audio_rope) + + vis_embed, aud_embed = self._run_visual_blocks_fused( + vis_embed, + aud_embed, + video_te, + audio_te, + video_tm, + audio_tm, + vis_rope, + aud_rope, + sparse_params, + attn_mask, + ) + if self._pp_final is not None: + final = self._pp_final + self._pp_final = None + return final + return self._project_fused( + vis_embed, + aud_embed, + self._pp_vis_shape, + video_tm, + audio_tm, + ) + + def reset_parameters(self) -> None: + for m in self.modules(): + if m is not self and hasattr(m, "reset_parameters"): + m.reset_parameters() + if hasattr(self, "visual_token_type_embeddings"): + nn.init.zeros_(self.visual_token_type_embeddings.weight) + + +__all__ = [ + "Kandinsky6Transformer3DModel", + "Kandinsky6TransformerEncoderBlock", + "Kandinsky6TransformerDecoderBlock", + "Kandinsky6FusedTransformerDecoderBlock", +] diff --git a/vllm_omni/diffusion/models/kandinsky6/modeling_kandinsky6_audio.py b/vllm_omni/diffusion/models/kandinsky6/modeling_kandinsky6_audio.py new file mode 100644 index 000000000..742647b82 --- /dev/null +++ b/vllm_omni/diffusion/models/kandinsky6/modeling_kandinsky6_audio.py @@ -0,0 +1,1753 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Generated standalone Kandinsky 6 audio VAE component for vLLM-Omni. + +The template contains source references for the frontend, time-domain VAE, +and BigVGAN implementation. The assembler expands those references into this +single module, keeping the export synchronized with native source changes. + +Identical to the Diffusers port's audio VAE module: ``Kandinsky6AudioVAE`` +only uses ``diffusers`` as a modeling-utility library (ModelMixin/ +ConfigMixin checkpoint serialization), not for pipeline coupling, so it +needs no vLLM-specific override. +""" + +from __future__ import annotations + +import json +import math +import os +import shutil +from os import PathLike +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from einops import rearrange +from huggingface_hub import PyTorchModelHubMixin +from torch import pow, sin +from torch.nn import Conv1d, ConvTranspose1d, Parameter +from torch.nn.utils import weight_norm as legacy_weight_norm +from torch.nn.utils.parametrizations import weight_norm +from torch.nn.utils.parametrize import remove_parametrizations + +from vllm_omni.transformers_utils.repo_utils import hf_api +from vllm_omni.utils.audio import mel_filter_bank + +# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# This work is licensed under a Creative Commons +# Attribution-NonCommercial-ShareAlike 4.0 International License. +# You should have received a copy of the license along with this +# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/ +"""Improved diffusion model architecture proposed in the paper +"Analyzing and Improving the Training Dynamics of Diffusion Models".""" + + +# ---------------------------------------------------------------------------- +# Variant of constant() that inherits dtype and device from the given +# reference tensor by default. + +_constant_cache = dict() + + +def constant(value, shape=None, dtype=None, device=None, memory_format=None): + value = np.asarray(value) + if shape is not None: + shape = tuple(shape) + if dtype is None: + dtype = torch.get_default_dtype() + if device is None: + device = torch.device("cpu") + if memory_format is None: + memory_format = torch.contiguous_format + + key = (value.shape, value.dtype, value.tobytes(), shape, dtype, device, memory_format) + tensor = _constant_cache.get(key, None) + if tensor is None: + tensor = torch.as_tensor(value.copy(), dtype=dtype, device=device) + if shape is not None: + tensor, _ = torch.broadcast_tensors(tensor, torch.empty(shape)) + tensor = tensor.contiguous(memory_format=memory_format) + _constant_cache[key] = tensor + return tensor + + +def const_like(ref, value, shape=None, dtype=None, device=None, memory_format=None): + if dtype is None: + dtype = ref.dtype + if device is None: + device = ref.device + return constant(value, shape=shape, dtype=dtype, device=device, memory_format=memory_format) + + +# ---------------------------------------------------------------------------- +# Normalize given tensor to unit magnitude with respect to the given +# dimensions. Default = all dimensions except the first. + + +def normalize(x, dim=None, eps=1e-4): + if dim is None: + dim = list(range(1, x.ndim)) + norm = torch.linalg.vector_norm(x, dim=dim, keepdim=True, dtype=torch.float32) + norm = torch.add(eps, norm, alpha=np.sqrt(norm.numel() / x.numel())) + return x / norm.to(x.dtype) + + +class Normalize(torch.nn.Module): + def __init__(self, dim=None, eps=1e-4): + super().__init__() + self.dim = dim + self.eps = eps + + def forward(self, x): + return normalize(x, dim=self.dim, eps=self.eps) + + +# ---------------------------------------------------------------------------- +# Upsample or downsample the given tensor with the given filter, +# or keep it as is. + + +def resample(x, f=[1, 1], mode="keep"): + if mode == "keep": + return x + f = np.float32(f) + assert f.ndim == 1 and len(f) % 2 == 0 + pad = (len(f) - 1) // 2 + f = f / f.sum() + f = np.outer(f, f)[np.newaxis, np.newaxis, :, :] + f = const_like(x, f) + c = x.shape[1] + if mode == "down": + return torch.nn.functional.conv2d(x, f.tile([c, 1, 1, 1]), groups=c, stride=2, padding=(pad,)) + assert mode == "up" + return torch.nn.functional.conv_transpose2d(x, (f * 4).tile([c, 1, 1, 1]), groups=c, stride=2, padding=(pad,)) + + +# ---------------------------------------------------------------------------- +# Magnitude-preserving SiLU (Equation 81). + + +def mp_silu(x): + return torch.nn.functional.silu(x) / 0.596 + + +class MPSiLU(torch.nn.Module): + def forward(self, x): + return mp_silu(x) + + +# ---------------------------------------------------------------------------- +# Magnitude-preserving sum (Equation 88). + + +def mp_sum(a, b, t=0.5): + return a.lerp(b, t) / np.sqrt((1 - t) ** 2 + t**2) + + +# ---------------------------------------------------------------------------- +# Magnitude-preserving concatenation (Equation 103). + + +def mp_cat(a, b, dim=1, t=0.5): + Na = a.shape[dim] + Nb = b.shape[dim] + C = np.sqrt((Na + Nb) / ((1 - t) ** 2 + t**2)) + wa = C / np.sqrt(Na) * (1 - t) + wb = C / np.sqrt(Nb) * t + return torch.cat([wa * a, wb * b], dim=dim) + + +# ---------------------------------------------------------------------------- +# Magnitude-preserving convolution or fully-connected layer (Equation 47) +# with force weight normalization (Equation 66). + + +class MPConv1D(torch.nn.Module): + def __init__(self, in_channels, out_channels, kernel_size): + super().__init__() + self.out_channels = out_channels + self.weight = torch.nn.Parameter(torch.randn(out_channels, in_channels, kernel_size)) + + self.weight_norm_removed = False + + def forward(self, x, gain=1): + assert self.weight_norm_removed, "call remove_weight_norm() before inference" + + w = self.weight * gain + if w.ndim == 2: + return x @ w.t() + assert w.ndim == 3 + return torch.nn.functional.conv1d(x, w, padding=(w.shape[-1] // 2,)) + + def remove_weight_norm(self): + w = self.weight.to(torch.float32) + w = normalize(w) # traditional weight normalization + w = w / np.sqrt(w[0].numel()) + w = w.to(self.weight.dtype) + self.weight.data.copy_(w) + + self.weight_norm_removed = True + return self + + +class DiagonalGaussianDistribution: + def __init__(self, parameters, deterministic=False): + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device) + + def sample(self, rng: torch.Generator | None = None): + # x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device) + + r = torch.empty_like(self.mean).normal_(generator=rng) + x = self.mean + self.std * r + + return x + + def kl(self, other=None): + if self.deterministic: + return torch.Tensor([0.0]) + else: + if other is None: + return 0.5 * torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar + else: + return 0.5 * ( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var + - 1.0 + - self.logvar + + other.logvar + ) + + def nll(self, sample, dims=[1, 2, 3]): + if self.deterministic: + return torch.Tensor([0.0]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, dim=dims) + + def mode(self): + return self.mean + + +def nonlinearity(x): + # swish + return mp_silu(x) + + +class ResnetBlock1D(nn.Module): + def __init__(self, *, in_dim, out_dim=None, conv_shortcut=False, kernel_size=3, use_norm=True): + super().__init__() + self.in_dim = in_dim + out_dim = in_dim if out_dim is None else out_dim + self.out_dim = out_dim + self.use_conv_shortcut = conv_shortcut + self.use_norm = use_norm + + self.conv1 = MPConv1D(in_dim, out_dim, kernel_size=kernel_size) + self.conv2 = MPConv1D(out_dim, out_dim, kernel_size=kernel_size) + if self.in_dim != self.out_dim: + if self.use_conv_shortcut: + self.conv_shortcut = MPConv1D(in_dim, out_dim, kernel_size=kernel_size) + else: + self.nin_shortcut = MPConv1D(in_dim, out_dim, kernel_size=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # pixel norm + if self.use_norm: + x = normalize(x, dim=1) + + h = x + h = nonlinearity(h) + h = self.conv1(h) + + h = nonlinearity(h) + h = self.conv2(h) + + if self.in_dim != self.out_dim: + if self.use_conv_shortcut: + x = self.conv_shortcut(x) + else: + x = self.nin_shortcut(x) + + return mp_sum(x, h, t=0.3) + + +class AttnBlock1D(nn.Module): + def __init__(self, in_channels, num_heads=1): + super().__init__() + self.in_channels = in_channels + + self.num_heads = num_heads + self.qkv = MPConv1D(in_channels, in_channels * 3, kernel_size=1) + self.proj_out = MPConv1D(in_channels, in_channels, kernel_size=1) + + def forward(self, x): + h = x + y = self.qkv(h) + y = y.reshape(y.shape[0], self.num_heads, -1, 3, y.shape[-1]) + q, k, v = normalize(y, dim=2).unbind(3) + + q = rearrange(q, "b h c l -> b h l c") + k = rearrange(k, "b h c l -> b h l c") + v = rearrange(v, "b h c l -> b h l c") + + h = F.scaled_dot_product_attention(q, k, v) + h = rearrange(h, "b h l c -> b (h c) l") + + h = self.proj_out(h) + + return mp_sum(x, h, t=0.3) + + +class Upsample1D(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + self.conv = MPConv1D(in_channels, in_channels, kernel_size=3) + + def forward(self, x): + x = F.interpolate(x, scale_factor=2.0, mode="nearest-exact") # support 3D tensor(B,C,T) + if self.with_conv: + x = self.conv(x) + return x + + +class Downsample1D(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + # no asymmetric padding in torch conv, must do it ourselves + self.conv1 = MPConv1D(in_channels, in_channels, kernel_size=1) + self.conv2 = MPConv1D(in_channels, in_channels, kernel_size=1) + + def forward(self, x): + if self.with_conv: + x = self.conv1(x) + + x = F.avg_pool1d(x, kernel_size=2, stride=2) + + if self.with_conv: + x = self.conv2(x) + + return x + + +# fmt: off +DATA_MEAN_128D = [ + -3.3462, -2.6723, -2.4893, -2.3143, -2.2664, -2.3317, -2.1802, -2.4006, -2.2357, -2.4597, + -2.3717, -2.4690, -2.5142, -2.4919, -2.6610, -2.5047, -2.7483, -2.5926, -2.7462, -2.7033, + -2.7386, -2.8112, -2.7502, -2.9594, -2.7473, -3.0035, -2.8891, -2.9922, -2.9856, -3.0157, + -3.1191, -2.9893, -3.1718, -3.0745, -3.1879, -3.2310, -3.1424, -3.2296, -3.2791, -3.2782, + -3.2756, -3.3134, -3.3509, -3.3750, -3.3951, -3.3698, -3.4505, -3.4509, -3.5089, -3.4647, + -3.5536, -3.5788, -3.5867, -3.6036, -3.6400, -3.6747, -3.7072, -3.7279, -3.7283, -3.7795, + -3.8259, -3.8447, -3.8663, -3.9182, -3.9605, -3.9861, -4.0105, -4.0373, -4.0762, -4.1121, + -4.1488, -4.1874, -4.2461, -4.3170, -4.3639, -4.4452, -4.5282, -4.6297, -4.7019, -4.7960, + -4.8700, -4.9507, -5.0303, -5.0866, -5.1634, -5.2342, -5.3242, -5.4053, -5.4927, -5.5712, + -5.6464, -5.7052, -5.7619, -5.8410, -5.9188, -6.0103, -6.0955, -6.1673, -6.2362, -6.3120, + -6.3926, -6.4797, -6.5565, -6.6511, -6.8130, -6.9961, -7.1275, -7.2457, -7.3576, -7.4663, + -7.6136, -7.7469, -7.8815, -8.0132, -8.1515, -8.3071, -8.4722, -8.7418, -9.3975, -9.6628, + -9.7671, -9.8863, -9.9992, -10.0860, -10.1709, -10.5418, -11.2795, -11.3861 +] + +DATA_STD_128D = [ + 2.3804, 2.4368, 2.3772, 2.3145, 2.2803, 2.2510, 2.2316, 2.2083, 2.1996, 2.1835, 2.1769, 2.1659, + 2.1631, 2.1618, 2.1540, 2.1606, 2.1571, 2.1567, 2.1612, 2.1579, 2.1679, 2.1683, 2.1634, 2.1557, + 2.1668, 2.1518, 2.1415, 2.1449, 2.1406, 2.1350, 2.1313, 2.1415, 2.1281, 2.1352, 2.1219, 2.1182, + 2.1327, 2.1195, 2.1137, 2.1080, 2.1179, 2.1036, 2.1087, 2.1036, 2.1015, 2.1068, 2.0975, 2.0991, + 2.0902, 2.1015, 2.0857, 2.0920, 2.0893, 2.0897, 2.0910, 2.0881, 2.0925, 2.0873, 2.0960, 2.0900, + 2.0957, 2.0958, 2.0978, 2.0936, 2.0886, 2.0905, 2.0845, 2.0855, 2.0796, 2.0840, 2.0813, 2.0817, + 2.0838, 2.0840, 2.0917, 2.1061, 2.1431, 2.1976, 2.2482, 2.3055, 2.3700, 2.4088, 2.4372, 2.4609, + 2.4731, 2.4847, 2.5072, 2.5451, 2.5772, 2.6147, 2.6529, 2.6596, 2.6645, 2.6726, 2.6803, 2.6812, + 2.6899, 2.6916, 2.6931, 2.6998, 2.7062, 2.7262, 2.7222, 2.7158, 2.7041, 2.7485, 2.7491, 2.7451, + 2.7485, 2.7233, 2.7297, 2.7233, 2.7145, 2.6958, 2.6788, 2.6439, 2.6007, 2.4786, 2.2469, 2.1877, + 2.1392, 2.0717, 2.0107, 1.9676, 1.9140, 1.7102, 0.9101, 0.7164 +] +# fmt: on + + +class VAE(nn.Module): + def __init__( + self, + *, + data_dim: int, + embed_dim: int, + hidden_dim: int, + ): + super().__init__() + + if data_dim != 128: + raise ValueError(f"Only 44k audio VAE data is supported, got data_dim={data_dim}") + self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_128D, dtype=torch.float32)) + self.data_std = nn.Buffer(torch.tensor(DATA_STD_128D, dtype=torch.float32)) + + self.data_mean = self.data_mean.view(1, -1, 1) + self.data_std = self.data_std.view(1, -1, 1) + + self.encoder = Encoder1D( + dim=hidden_dim, + ch_mult=(1, 2, 4), + num_res_blocks=2, + attn_layers=[3], + down_layers=[0], + in_dim=data_dim, + embed_dim=embed_dim, + ) + self.decoder = Decoder1D( + dim=hidden_dim, + ch_mult=(1, 2, 4), + num_res_blocks=2, + attn_layers=[3], + down_layers=[0], + in_dim=data_dim, + out_dim=data_dim, + embed_dim=embed_dim, + ) + + self.embed_dim = embed_dim + # self.quant_conv = nn.Conv1d(2 * embed_dim, 2 * embed_dim, 1) + # self.post_quant_conv = nn.Conv1d(embed_dim, embed_dim, 1) + + self.initialize_weights() + + def initialize_weights(self): + pass + + def encode(self, x: torch.Tensor, normalize: bool = True) -> DiagonalGaussianDistribution: + if normalize: + x = self.normalize(x) + moments = self.encoder(x) + posterior = DiagonalGaussianDistribution(moments) + return posterior + + def decode(self, z: torch.Tensor, unnormalize: bool = True) -> torch.Tensor: + dec = self.decoder(z) + if unnormalize: + dec = self.unnormalize(dec) + return dec + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x - self.data_mean) / self.data_std + + def unnormalize(self, x: torch.Tensor) -> torch.Tensor: + return x * self.data_std + self.data_mean + + def forward( + self, + x: torch.Tensor, + sample_posterior: bool = True, + rng: torch.Generator | None = None, + normalize: bool = True, + unnormalize: bool = True, + ) -> tuple[torch.Tensor, DiagonalGaussianDistribution]: + posterior = self.encode(x, normalize=normalize) + if sample_posterior: + z = posterior.sample(rng) + else: + z = posterior.mode() + dec = self.decode(z, unnormalize=unnormalize) + return dec, posterior + + def load_weights(self, src_dict) -> set[str]: + """Load a state dict, or the ``(name, tensor)`` stream from AutoWeightsLoader.""" + if not isinstance(src_dict, dict): + src_dict = dict(src_dict) + self.load_state_dict(src_dict, strict=True) + return set(src_dict) + + @property + def device(self) -> torch.device: + return next(self.parameters()).device + + def get_last_layer(self): + return self.decoder.conv_out.weight + + def remove_weight_norm(self): + for name, m in self.named_modules(): + if isinstance(m, MPConv1D): + m.remove_weight_norm() + + return self + + +class Encoder1D(nn.Module): + def __init__( + self, + *, + dim: int, + ch_mult: tuple[int] = (1, 2, 4, 8), + num_res_blocks: int, + attn_layers: list[int] = [], + down_layers: list[int] = [], + resamp_with_conv: bool = True, + in_dim: int, + embed_dim: int, + double_z: bool = True, + kernel_size: int = 3, + clip_act: float = 256.0, + ): + super().__init__() + self.dim = dim + self.num_layers = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.in_channels = in_dim + self.clip_act = clip_act + self.down_layers = down_layers + self.attn_layers = attn_layers + self.conv_in = MPConv1D(in_dim, self.dim, kernel_size=kernel_size) + + in_ch_mult = (1,) + tuple(ch_mult) + self.in_ch_mult = in_ch_mult + # downsampling + self.down = nn.ModuleList() + for i_level in range(self.num_layers): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = dim * in_ch_mult[i_level] + block_out = dim * ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append(ResnetBlock1D(in_dim=block_in, out_dim=block_out, kernel_size=kernel_size, use_norm=True)) + block_in = block_out + if i_level in attn_layers: + attn.append(AttnBlock1D(block_in)) + down = nn.Module() + down.block = block + down.attn = attn + if i_level in down_layers: + down.downsample = Downsample1D(block_in, resamp_with_conv) + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, kernel_size=kernel_size, use_norm=True) + self.mid.attn_1 = AttnBlock1D(block_in) + self.mid.block_2 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, kernel_size=kernel_size, use_norm=True) + + # end + self.conv_out = MPConv1D(block_in, 2 * embed_dim if double_z else embed_dim, kernel_size=kernel_size) + + self.learnable_gain = nn.Parameter(torch.zeros([])) + + def forward(self, x): + # downsampling + hs = [self.conv_in(x)] + for i_level in range(self.num_layers): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1]) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + h = h.clamp(-self.clip_act, self.clip_act) + hs.append(h) + if i_level in self.down_layers: + hs.append(self.down[i_level].downsample(hs[-1])) + + # middle + h = hs[-1] + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + h = h.clamp(-self.clip_act, self.clip_act) + + # end + h = nonlinearity(h) + h = self.conv_out(h, gain=(self.learnable_gain + 1)) + return h + + +class Decoder1D(nn.Module): + def __init__( + self, + *, + dim: int, + out_dim: int, + ch_mult: tuple[int] = (1, 2, 4, 8), + num_res_blocks: int, + attn_layers: list[int] = [], + down_layers: list[int] = [], + kernel_size: int = 3, + resamp_with_conv: bool = True, + in_dim: int, + embed_dim: int, + clip_act: float = 256.0, + ): + super().__init__() + self.ch = dim + self.num_layers = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.in_channels = in_dim + self.clip_act = clip_act + self.down_layers = [i + 1 for i in down_layers] # each downlayer add one + + # compute in_ch_mult, block_in and curr_res at lowest res + block_in = dim * ch_mult[self.num_layers - 1] + + # z to block_in + self.conv_in = MPConv1D(embed_dim, block_in, kernel_size=kernel_size) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, use_norm=True) + self.mid.attn_1 = AttnBlock1D(block_in) + self.mid.block_2 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, use_norm=True) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_layers)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = dim * ch_mult[i_level] + for i_block in range(self.num_res_blocks + 1): + block.append(ResnetBlock1D(in_dim=block_in, out_dim=block_out, use_norm=True)) + block_in = block_out + if i_level in attn_layers: + attn.append(AttnBlock1D(block_in)) + up = nn.Module() + up.block = block + up.attn = attn + if i_level in self.down_layers: + up.upsample = Upsample1D(block_in, resamp_with_conv) + self.up.insert(0, up) # prepend to get consistent order + + # end + self.conv_out = MPConv1D(block_in, out_dim, kernel_size=kernel_size) + self.learnable_gain = nn.Parameter(torch.zeros([])) + + def forward(self, z): + # z to block_in + h = self.conv_in(z) + + # middle + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + h = h.clamp(-self.clip_act, self.clip_act) + + # upsampling + for i_level in reversed(range(self.num_layers)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + h = h.clamp(-self.clip_act, self.clip_act) + if i_level in self.down_layers: + h = self.up[i_level].upsample(h) + + h = nonlinearity(h) + h = self.conv_out(h, gain=(self.learnable_gain + 1)) + return h + + +def VAE_44k(**kwargs) -> VAE: + return VAE(data_dim=128, embed_dim=40, hidden_dim=512, **kwargs) + + +def get_my_vae(name: str, **kwargs) -> VAE: + if name == "44k": + return VAE_44k(**kwargs) + raise ValueError(f"Unknown model: {name}") + + +if __name__ == "__main__": + network = get_my_vae("44k") + + # print the number of parameters in terms of millions + num_params = sum(p.numel() for p in network.parameters()) / 1e6 + print(f"Number of parameters: {num_params:.2f}M") + +# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license. +# LICENSE is in incl_licenses directory. + + +class Snake(nn.Module): + """ + Implementation of a sine-based periodic activation function + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter + References: + - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://arxiv.org/abs/2006.08195 + Examples: + >>> a1 = snake(256) + >>> x = torch.randn(256) + >>> x = a1(x) + """ + + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + """ + Initialization. + INPUT: + - in_features: shape of the input + - alpha: trainable parameter + alpha is initialized to 1 by default, higher values = higher-frequency. + alpha will be trained along with the rest of your model. + """ + super(Snake, self).__init__() + self.in_features = in_features + + # Initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # Log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + else: # Linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + """ + Forward pass of the function. + Applies the function to the input elementwise. + Snake ∶= x + 1/a * sin^2 (xa) + """ + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # Line up with x to [B, C, T] + if self.alpha_logscale: + alpha = torch.exp(alpha) + x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + return x + + +class SnakeBeta(nn.Module): + """ + A modified Snake function which uses separate parameters for the magnitude of the periodic components + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + References: + - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://arxiv.org/abs/2006.08195 + Examples: + >>> a1 = snakebeta(256) + >>> x = torch.randn(256) + >>> x = a1(x) + """ + + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + """ + Initialization. + INPUT: + - in_features: shape of the input + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + alpha is initialized to 1 by default, higher values = higher-frequency. + beta is initialized to 1 by default, higher values = higher-magnitude. + alpha will be trained along with the rest of your model. + """ + super(SnakeBeta, self).__init__() + self.in_features = in_features + + # Initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # Log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + self.beta = Parameter(torch.zeros(in_features) * alpha) + else: # Linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + self.beta = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + """ + Forward pass of the function. + Applies the function to the input elementwise. + SnakeBeta ∶= x + 1/b * sin^2 (xa) + """ + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # Line up with x to [B, C, T] + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + return x + + +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 +# LICENSE is in incl_licenses directory. + + +if "sinc" in dir(torch): + sinc = torch.sinc +else: + # This code is adopted from adefossez's julius.core.sinc under the MIT License + # https://adefossez.github.io/julius/julius/core.html + # LICENSE is in incl_licenses directory. + def sinc(x: torch.Tensor): + """ + Implementation of sinc, i.e. sin(pi * x) / (pi * x) + __Warning__: Different to julius.sinc, the input is multiplied by `pi`! + """ + return torch.where( + x == 0, + torch.tensor(1.0, device=x.device, dtype=x.dtype), + torch.sin(math.pi * x) / math.pi / x, + ) + + +# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License +# https://adefossez.github.io/julius/julius/lowpass.html +# LICENSE is in incl_licenses directory. +def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): # return filter [1,1,kernel_size] + even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + + # For kaiser window + delta_f = 4 * half_width + A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 + if A > 50.0: + beta = 0.1102 * (A - 8.7) + elif A >= 21.0: + beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + + # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio + if even: + time = torch.arange(-half_size, half_size) + 0.5 + else: + time = torch.arange(kernel_size) - half_size + if cutoff == 0: + filter_ = torch.zeros_like(time) + else: + filter_ = 2 * cutoff * window * sinc(2 * cutoff * time) + """ + Normalize filter to have sum = 1, otherwise we will have a small leakage of the constant component in the input signal. + """ + filter_ /= filter_.sum() + filter = filter_.view(1, 1, kernel_size) + + return filter + + +class LowPassFilter1d(nn.Module): + def __init__( + self, + cutoff=0.5, + half_width=0.6, + stride: int = 1, + padding: bool = True, + padding_mode: str = "replicate", + kernel_size: int = 12, + ): + """ + kernel_size should be even number for stylegan3 setup, in this implementation, odd number is also possible. + """ + super().__init__() + if cutoff < -0.0: + raise ValueError("Minimum cutoff must be larger than zero.") + if cutoff > 0.5: + raise ValueError("A cutoff above 0.5 does not make sense.") + self.kernel_size = kernel_size + self.even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(self.even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.padding = padding + self.padding_mode = padding_mode + filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) + self.register_buffer("filter", filter) + + # Input [B, C, T] + def forward(self, x): + _, C, _ = x.shape + + if self.padding: + x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) + out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + + return out + + +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 +# LICENSE is in incl_licenses directory. + + +class UpSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None): + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.stride = ratio + self.pad = self.kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 + self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 + filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size) + self.register_buffer("filter", filter) + + # x: [B, C, T] + def forward(self, x): + _, C, _ = x.shape + + x = F.pad(x, (self.pad, self.pad), mode="replicate") + x = self.ratio * F.conv_transpose1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + x = x[..., self.pad_left : -self.pad_right] + + return x + + +class DownSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None): + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.lowpass = LowPassFilter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + stride=ratio, + kernel_size=self.kernel_size, + ) + + def forward(self, x): + xx = self.lowpass(x) + + return xx + + +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 +# LICENSE is in incl_licenses directory. + + +class Activation1d(nn.Module): + def __init__( + self, + activation, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + ): + super().__init__() + self.up_ratio = up_ratio + self.down_ratio = down_ratio + self.act = activation + self.upsample = UpSample1d(up_ratio, up_kernel_size) + self.downsample = DownSample1d(down_ratio, down_kernel_size) + + # x: [B,C,T] + def forward(self, x): + x = self.upsample(x) + x = self.act(x) + x = self.downsample(x) + + return x + + +# Adapted from https://github.com/jik876/hifi-gan under the MIT license. +# LICENSE is in incl_licenses directory. + + +class AttrDict(dict): + def __init__(self, *args, **kwargs): + super(AttrDict, self).__init__(*args, **kwargs) + self.__dict__ = self + + +def build_env(config, config_name, path): + t_path = os.path.join(path, config_name) + if config != t_path: + os.makedirs(path, exist_ok=True) + shutil.copyfile(config, os.path.join(path, config_name)) + + +# Adapted from https://github.com/jik876/hifi-gan under the MIT license. +# LICENSE is in incl_licenses directory. + + +def init_weights(m, mean=0.0, std=0.01): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + m.weight.data.normal_(mean, std) + + +def apply_weight_norm(m): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + legacy_weight_norm(m) + + +def get_padding(kernel_size, dilation=1): + return int((kernel_size * dilation - dilation) / 2) + + +def load_checkpoint(filepath, device): + assert os.path.isfile(filepath) + print(f"Loading '{filepath}'") + checkpoint_dict = torch.load(filepath, map_location=device) + print("Complete.") + return checkpoint_dict + + +# Copyright (c) 2024 NVIDIA CORPORATION. +# Licensed under the MIT license. + +# Adapted from https://github.com/jik876/hifi-gan under the MIT license. +# LICENSE is in incl_licenses directory. + + +def load_hparams_from_json(path) -> AttrDict: + with open(path) as f: + data = f.read() + return AttrDict(json.loads(data)) + + +class AMPBlock1(torch.nn.Module): + """ + AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer. + AMPBlock1 has additional self.convs2 that contains additional Conv1d layers with a fixed dilation=1 followed by each layer in self.convs1 + + Args: + h (AttrDict): Hyperparameters. + channels (int): Number of convolution channels. + kernel_size (int): Size of the convolution kernel. Default is 3. + dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5). + activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Default is None. + """ + + def __init__( + self, + h: AttrDict, + channels: int, + kernel_size: int = 3, + dilation: tuple = (1, 3, 5), + activation: str = None, + ): + super().__init__() + + self.h = h + + self.convs1 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + stride=1, + dilation=d, + padding=get_padding(kernel_size, d), + ) + ) + for d in dilation + ] + ) + self.convs1.apply(init_weights) + + self.convs2 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + stride=1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ) + for _ in range(len(dilation)) + ] + ) + self.convs2.apply(init_weights) + + self.num_layers = len(self.convs1) + len(self.convs2) # Total number of conv layers + + # Select which Activation1d, lazy-load cuda version to ensure backward compatibility + if self.h.get("use_cuda_kernel", False): + raise RuntimeError("use_cuda_kernel=True is unavailable in the standalone Diffusers audio VAE") + Activation1d = TorchActivation1d + + # Activation functions + if activation == "snake": + self.activations = nn.ModuleList( + [ + Activation1d(activation=Snake(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ] + ) + elif activation == "snakebeta": + self.activations = nn.ModuleList( + [ + Activation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ] + ) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + def forward(self, x): + acts1, acts2 = self.activations[::2], self.activations[1::2] + for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2): + xt = a1(x) + xt = c1(xt) + xt = a2(xt) + xt = c2(xt) + x = xt + x + + return x + + def remove_weight_norm(self): + for l in self.convs1: + remove_parametrizations(l, "weight") + for l in self.convs2: + remove_parametrizations(l, "weight") + + +class AMPBlock2(torch.nn.Module): + """ + AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer. + Unlike AMPBlock1, AMPBlock2 does not contain extra Conv1d layers with fixed dilation=1 + + Args: + h (AttrDict): Hyperparameters. + channels (int): Number of convolution channels. + kernel_size (int): Size of the convolution kernel. Default is 3. + dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5). + activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Default is None. + """ + + def __init__( + self, + h: AttrDict, + channels: int, + kernel_size: int = 3, + dilation: tuple = (1, 3, 5), + activation: str = None, + ): + super().__init__() + + self.h = h + + self.convs = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + stride=1, + dilation=d, + padding=get_padding(kernel_size, d), + ) + ) + for d in dilation + ] + ) + self.convs.apply(init_weights) + + self.num_layers = len(self.convs) # Total number of conv layers + + # Select which Activation1d, lazy-load cuda version to ensure backward compatibility + if self.h.get("use_cuda_kernel", False): + raise RuntimeError("use_cuda_kernel=True is unavailable in the standalone Diffusers audio VAE") + Activation1d = TorchActivation1d + + # Activation functions + if activation == "snake": + self.activations = nn.ModuleList( + [ + Activation1d(activation=Snake(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ] + ) + elif activation == "snakebeta": + self.activations = nn.ModuleList( + [ + Activation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ] + ) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + def forward(self, x): + for c, a in zip(self.convs, self.activations): + xt = a(x) + xt = c(xt) + x = xt + x + return x + + def remove_weight_norm(self): + for l in self.convs: + remove_parametrizations(l, "weight") + + +TorchActivation1d = Activation1d + + +class BigVGANv2( + torch.nn.Module, + PyTorchModelHubMixin, + library_name="bigvgan", + repo_url="https://github.com/NVIDIA/BigVGAN", + docs_url="https://github.com/NVIDIA/BigVGAN/blob/main/README.md", + pipeline_tag="audio-to-audio", + license="mit", + tags=["neural-vocoder", "audio-generation", "arxiv:2206.04658"], +): + """ + BigVGAN is a neural vocoder model that applies anti-aliased periodic activation for residual blocks (resblocks). + New in BigVGAN-v2: it can optionally use optimized CUDA kernels for AMP (anti-aliased multi-periodicity) blocks. + + Args: + h (AttrDict): Hyperparameters. + use_cuda_kernel (bool): If set to True, loads optimized CUDA kernels for AMP. This should be used for inference only, as training is not supported with CUDA kernels. + + Note: + - The `use_cuda_kernel` parameter should be used for inference only, as training with CUDA kernels is not supported. + - Ensure that the activation function is correctly specified in the hyperparameters (h.activation). + """ + + def __init__(self, h: AttrDict, use_cuda_kernel: bool = False): + super().__init__() + self.h = h + self.h["use_cuda_kernel"] = use_cuda_kernel + + # Select which Activation1d, lazy-load cuda version to ensure backward compatibility + if self.h.get("use_cuda_kernel", False): + raise RuntimeError("use_cuda_kernel=True is unavailable in the standalone Diffusers audio VAE") + Activation1d = TorchActivation1d + + self.num_kernels = len(h.resblock_kernel_sizes) + self.num_upsamples = len(h.upsample_rates) + + # Pre-conv + self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)) + + # Define which AMPBlock to use. BigVGAN uses AMPBlock1 as default + if h.resblock == "1": + resblock_class = AMPBlock1 + elif h.resblock == "2": + resblock_class = AMPBlock2 + else: + raise ValueError(f"Incorrect resblock class specified in hyperparameters. Got {h.resblock}") + + # Transposed conv-based upsamplers. does not apply anti-aliasing + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)): + self.ups.append( + nn.ModuleList( + [ + weight_norm( + ConvTranspose1d( + h.upsample_initial_channel // (2**i), + h.upsample_initial_channel // (2 ** (i + 1)), + k, + u, + padding=(k - u) // 2, + ) + ) + ] + ) + ) + + # Residual blocks using anti-aliased multi-periodicity composition modules (AMP) + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = h.upsample_initial_channel // (2 ** (i + 1)) + for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)): + self.resblocks.append(resblock_class(h, ch, k, d, activation=h.activation)) + + # Post-conv + activation_post = ( + Snake(ch, alpha_logscale=h.snake_logscale) + if h.activation == "snake" + else (SnakeBeta(ch, alpha_logscale=h.snake_logscale) if h.activation == "snakebeta" else None) + ) + if activation_post is None: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + self.activation_post = Activation1d(activation=activation_post) + + # Whether to use bias for the final conv_post. Default to True for backward compatibility + self.use_bias_at_final = h.get("use_bias_at_final", True) + self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final)) + + # Weight initialization + for i in range(len(self.ups)): + self.ups[i].apply(init_weights) + self.conv_post.apply(init_weights) + + # Final tanh activation. Defaults to True for backward compatibility + self.use_tanh_at_final = h.get("use_tanh_at_final", True) + + def forward(self, x): + # Pre-conv + x = self.conv_pre(x) + + for i in range(self.num_upsamples): + # Upsampling + for i_up in range(len(self.ups[i])): + x = self.ups[i][i_up](x) + # AMP blocks + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs / self.num_kernels + + # Post-conv + x = self.activation_post(x) + x = self.conv_post(x) + # Final tanh activation + if self.use_tanh_at_final: + x = torch.tanh(x) + else: + x = torch.clamp(x, min=-1.0, max=1.0) # Bound the output to [-1, 1] + + return x + + def remove_weight_norm(self): + try: + # print("Removing weight norm...") + for l in self.ups: + for l_i in l: + remove_parametrizations(l_i, "weight") + for l in self.resblocks: + l.remove_weight_norm() + remove_parametrizations(self.conv_pre, "weight") + remove_parametrizations(self.conv_post, "weight") + except ValueError: + print("[INFO] Model already removed weight norm. Skipping!") + pass + + # Additional methods for huggingface_hub support + def _save_pretrained(self, save_directory: Path) -> None: + """Save weights and config.json from a Pytorch model to a local directory.""" + + model_path = save_directory / "bigvgan_generator.pt" + torch.save({"generator": self.state_dict()}, model_path) + + config_path = save_directory / "config.json" + with open(config_path, "w") as config_file: + json.dump(self.h, config_file, indent=4) + + @classmethod + def _from_pretrained( + cls, + *, + model_id: str, + revision: str, + cache_dir: str, + force_download: bool, + proxies: dict | None, + resume_download: bool, + local_files_only: bool, + token: str | bool | None, + map_location: str = "cpu", # Additional argument + strict: bool = False, # Additional argument + use_cuda_kernel: bool = False, + **model_kwargs, + ): + """Load Pytorch pretrained weights and return the loaded model.""" + + # Download and load hyperparameters (h) used by BigVGAN + if os.path.isdir(model_id): + # print("Loading config.json from local directory") + config_file = os.path.join(model_id, "config.json") + else: + config_file = hf_api().hf_hub_download( + repo_id=model_id, + filename="config.json", + revision=revision, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + token=token, + local_files_only=local_files_only, + ) + h = load_hparams_from_json(config_file) + + # instantiate BigVGAN using h + if use_cuda_kernel: + print( + "[WARNING] You have specified use_cuda_kernel=True during BigVGAN.from_pretrained(). Only inference is supported (training is not implemented)!" + ) + print( + "[WARNING] You need nvcc and ninja installed in your system that matches your PyTorch build is using to build the kernel. If not, the model will fail to initialize or generate incorrect waveform!" + ) + print( + "[WARNING] For detail, see the official GitHub repository: https://github.com/NVIDIA/BigVGAN?tab=readme-ov-file#using-custom-cuda-kernel-for-synthesis" + ) + model = cls(h, use_cuda_kernel=use_cuda_kernel) + + # Download and load pretrained generator weight + if os.path.isdir(model_id): + # print("Loading weights from local directory") + model_file = os.path.join(model_id, "bigvgan_generator.pt") + else: + print(f"Loading weights from {model_id}") + model_file = hf_api().hf_hub_download( + repo_id=model_id, + filename="bigvgan_generator.pt", + revision=revision, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + token=token, + local_files_only=local_files_only, + ) + + checkpoint_dict = torch.load(model_file, map_location=map_location, weights_only=True) + + try: + model.load_state_dict(checkpoint_dict["generator"]) + except RuntimeError: + print( + "[INFO] the pretrained checkpoint does not contain weight norm. Loading the checkpoint after removing weight norm!" + ) + model.remove_weight_norm() + model.load_state_dict(checkpoint_dict["generator"]) + + return model + + +def _load_bigvgan_v2( + vocoder_ckpt_path: str | None = None, + *, + vocoder_config: dict | None = None, + use_cuda_kernel: bool = False, +) -> BigVGANv2: + """Load BigVGAN-v2 from a local directory (avoids HubMixin/huggingface_hub API drift).""" + model_file = None if vocoder_ckpt_path is None else os.path.join(vocoder_ckpt_path, "bigvgan_generator.pt") + if vocoder_config is None: + if vocoder_ckpt_path is None: + raise ValueError("Either vocoder_ckpt_path or vocoder_config is required") + config_file = os.path.join(vocoder_ckpt_path, "config.json") + if not os.path.isfile(config_file) or not os.path.isfile(model_file): + raise FileNotFoundError(f"Expected config.json and bigvgan_generator.pt under {vocoder_ckpt_path}") + h = load_hparams_from_json(config_file) + else: + if model_file is not None and not os.path.isfile(model_file): + raise FileNotFoundError(f"BigVGAN checkpoint does not exist: {model_file}") + h = AttrDict(vocoder_config) + + model = BigVGANv2(h, use_cuda_kernel=use_cuda_kernel) + if vocoder_ckpt_path is not None: + assert model_file is not None + checkpoint_dict = torch.load(model_file, map_location="cpu", weights_only=True) + try: + model.load_state_dict(checkpoint_dict["generator"]) + except RuntimeError: + model.remove_weight_norm() + model.load_state_dict(checkpoint_dict["generator"]) + model.remove_weight_norm() + else: + model.remove_weight_norm() + return model + + +class AutoEncoderModule(nn.Module): + def __init__( + self, + *, + vae_ckpt_path: str | None = None, + vocoder_ckpt_path: str | None = None, + vocoder_config: dict | None = None, + mode: Literal["44k"], + need_vae_encoder: bool = True, + need_vae_decoder: bool = True, + ): + super().__init__() + self.vae: VAE = get_my_vae(mode).eval() + if vae_ckpt_path is not None: + vae_state_dict = torch.load(vae_ckpt_path, weights_only=True, map_location="cpu") + self.vae.load_state_dict(vae_state_dict) + self.vae.remove_weight_norm() + + if need_vae_decoder: + self.vocoder = _load_bigvgan_v2( + vocoder_ckpt_path, + vocoder_config=vocoder_config, + use_cuda_kernel=False, + ).eval() + else: + del self.vae.decoder + + if not need_vae_encoder: + del self.vae.encoder + + for param in self.parameters(): + param.requires_grad = False + + @torch.inference_mode() + def encode(self, x: torch.Tensor) -> DiagonalGaussianDistribution: + return self.vae.encode(x) + + @torch.inference_mode() + def decode(self, z: torch.Tensor) -> torch.Tensor: + return self.vae.decode(z) + + @torch.inference_mode() + def vocode(self, spec: torch.Tensor) -> torch.Tensor: + return self.vocoder(spec) + + +# Reference: # https://github.com/bytedance/Make-An-Audio-2 + + +def dynamic_range_compression_torch(x, C=1, clip_val=1e-5, *, norm_fn): + return norm_fn(torch.clamp(x, min=clip_val) * C) + + +def spectral_normalize_torch(magnitudes, norm_fn): + output = dynamic_range_compression_torch(magnitudes, norm_fn=norm_fn) + return output + + +class MelConverter(nn.Module): + def __init__( + self, + *, + sampling_rate: float, + n_fft: int, + num_mels: int, + hop_size: int, + win_size: int, + fmin: float, + fmax: float, + norm_fn, + ): + super().__init__() + self.sampling_rate = sampling_rate + self.n_fft = n_fft + self.num_mels = num_mels + self.hop_size = hop_size + self.win_size = win_size + self.fmin = fmin + self.fmax = fmax + self.norm_fn = norm_fn + + # torchaudio Slaney filterbank — numerically matches librosa.filters.mel + # (the repo bans librosa; see vllm_omni/utils/audio.py). + mel_basis = mel_filter_bank( + sr=self.sampling_rate, n_fft=self.n_fft, n_mels=self.num_mels, fmin=self.fmin, fmax=self.fmax + ).float() + hann_window = torch.hann_window(self.win_size) + + self.register_buffer("mel_basis", mel_basis) + self.register_buffer("hann_window", hann_window) + + @property + def device(self): + return self.mel_basis.device + + def forward(self, waveform: torch.Tensor, center: bool = False) -> torch.Tensor: + waveform = waveform.clamp(min=-1.0, max=1.0).to(self.device) + + waveform = torch.nn.functional.pad( + waveform.unsqueeze(1), + [int((self.n_fft - self.hop_size) / 2), int((self.n_fft - self.hop_size) / 2)], + mode="reflect", + ) + waveform = waveform.squeeze(1) + + spec = torch.stft( + waveform, + self.n_fft, + hop_length=self.hop_size, + win_length=self.win_size, + window=self.hann_window, + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + + spec = torch.view_as_real(spec) + spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)).float() + spec = torch.matmul(self.mel_basis, spec) + spec = spectral_normalize_torch(spec, self.norm_fn) + + return spec + + +def get_mel_converter(mode: Literal["44k"]) -> MelConverter: + return MelConverter( + sampling_rate=44_100, + n_fft=2048, + num_mels=128, + hop_size=512, + win_size=2048, + fmin=0, + fmax=44100 / 2, + norm_fn=torch.log, + ) + + +"""Audio feature / VAE facade (ported from kandinsky-5-inference MMAudio).""" + + +class FeaturesUtils(nn.Module): + def __init__( + self, + *, + tod_vae_ckpt: str | None = None, + bigvgan_vocoder_ckpt: str | None = None, + vocoder_config: dict | None = None, + mode: Literal["44k"] = "44k", + need_vae_encoder: bool = True, + need_vae_decoder: bool = True, + scaling_factor: float = 1.0, + ): + super().__init__() + self.mel_converter = get_mel_converter(mode) + self.tod = AutoEncoderModule( + vae_ckpt_path=tod_vae_ckpt, + vocoder_ckpt_path=bigvgan_vocoder_ckpt, + vocoder_config=vocoder_config, + mode=mode, + need_vae_encoder=need_vae_encoder, + need_vae_decoder=need_vae_decoder, + ) + self.scaling_factor = scaling_factor + self.downsample_factor = 1024 + + def compile(self): + self.decode = torch.compile(self.decode) + self.vocode = torch.compile(self.vocode) + + def train(self, mode: bool = True) -> FeaturesUtils: + return super().train(False) + + @torch.inference_mode() + def encode_audio(self, x) -> DiagonalGaussianDistribution: + assert self.tod is not None, "VAE is not loaded" + mel = self.mel_converter(x) + return self.tod.encode(mel) + + @torch.inference_mode() + def vocode(self, mel: torch.Tensor) -> torch.Tensor: + assert self.tod is not None, "VAE is not loaded" + return self.tod.vocode(mel) + + @torch.inference_mode() + def decode(self, z: torch.Tensor) -> torch.Tensor: + assert self.tod is not None, "VAE is not loaded" + return self.tod.decode(z) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + @torch.no_grad() + def wrapped_decode(self, z): + mel_decoded = self.decode(z.to(dtype=self.dtype)) + return self.vocode(mel_decoded) + + @torch.no_grad() + def wrapped_encode(self, audio): + dist = self.encode_audio(audio.to(dtype=self.dtype)) + return dist.mean + + +class Kandinsky6AudioVAE(ModelMixin, ConfigMixin): + """vLLM-Omni-serializable K6 audio VAE, mel frontend, and vocoder. + + The conversion-only checkpoint paths are accepted for materializing a + bundle, while ``vocoder_config`` makes a saved bundle self-contained. + """ + + ignore_for_config = ["tod_vae_ckpt", "bigvgan_vocoder_ckpt", "device"] + + @register_to_config + def __init__( + self, + *, + tod_vae_ckpt: str | PathLike[str] | None = None, + bigvgan_vocoder_ckpt: str | PathLike[str] | None = None, + vocoder_config: dict[str, Any] | None = None, + mode: str = "44k", + need_vae_encoder: bool = True, + need_vae_decoder: bool = True, + scaling_factor: float = 1.0, + device: str | torch.device | None = None, + ) -> None: + super().__init__() + if tod_vae_ckpt is None and vocoder_config is None: + raise ValueError("Kandinsky6AudioVAE requires a VAE checkpoint or embedded vocoder config") + self.native = FeaturesUtils( + tod_vae_ckpt=None if tod_vae_ckpt is None else str(tod_vae_ckpt), + bigvgan_vocoder_ckpt=None if bigvgan_vocoder_ckpt is None else str(bigvgan_vocoder_ckpt), + vocoder_config=vocoder_config, + mode=mode, + need_vae_encoder=need_vae_encoder, + need_vae_decoder=need_vae_decoder, + scaling_factor=scaling_factor, + ).eval() + if device is not None: + self.to(device) + self.downsample_factor = int(self.native.downsample_factor) + self.scaling_factor = float(scaling_factor) + + @property + def device(self) -> torch.device: + return self.native.device + + @property + def dtype(self) -> torch.dtype: + return self.native.dtype + + @torch.no_grad() + def encode_audio(self, audio: torch.Tensor): + return self.native.encode_audio(audio) + + @torch.no_grad() + def decode(self, latents: torch.Tensor) -> torch.Tensor: + return self.native.decode(latents) + + @torch.no_grad() + def vocode(self, mel: torch.Tensor) -> torch.Tensor: + return self.native.vocode(mel) + + @torch.no_grad() + def wrapped_decode(self, latents: torch.Tensor) -> torch.Tensor: + return self.native.wrapped_decode(latents) + + @torch.no_grad() + def wrapped_encode(self, audio: torch.Tensor) -> torch.Tensor: + return self.native.wrapped_encode(audio) + + +__all__ = ["Kandinsky6AudioVAE"] diff --git a/vllm_omni/diffusion/models/kandinsky6/modeling_kandinsky6_vae.py b/vllm_omni/diffusion/models/kandinsky6/modeling_kandinsky6_vae.py new file mode 100644 index 000000000..603adb2a3 --- /dev/null +++ b/vllm_omni/diffusion/models/kandinsky6/modeling_kandinsky6_vae.py @@ -0,0 +1,1189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Generated standalone Kandinsky 6 video VAE component for vLLM-Omni. + +Identical to the Diffusers port's VAE module: ``AutoencoderKLHunyuanVideo`` +only uses ``diffusers`` as a modeling-utility library (ModelMixin/ +ConfigMixin checkpoint serialization), not for pipeline coupling, so it +needs no vLLM-specific override — vLLM-Omni's own native models (e.g. +Wan2.2) already import from ``diffusers`` the same way. + +HunyuanVideo VAE with free-VRAM-aware spatial/temporal tiling, ported from +kandinsky-5-inference ``models/vae_orig.py`` (not the Diffusers +``AutoencoderKLHunyuanVideo``, which uses fixed tile sizes and peaks much higher). +""" + +from __future__ import annotations + +from math import ceil, sqrt + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.activations import get_activation +from diffusers.models.attention_processor import Attention +from diffusers.models.autoencoders.vae import DecoderOutput, DiagonalGaussianDistribution +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.utils.accelerate_utils import apply_forward_hook + +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedOperator, + DistributedVaeMixin, + GridSpec, + TileTask, +) +from vllm_omni.platforms import current_omni_platform + + +def prepare_causal_attention_mask(f: int, s: int, dtype: torch.dtype, device: torch.device, b: int) -> torch.Tensor: + return ( + torch.ones((f, f), dtype=dtype, device=device) + .tril_() + .log_() + .repeat_interleave(s, dim=0) + .repeat_interleave(s, dim=1) + .unsqueeze(0) + .expand(b, -1, -1) + .contiguous() + ) + + +class HunyuanVideoCausalConv3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int, int] = 3, + stride: int | tuple[int, int, int] = 1, + padding: int | tuple[int, int, int] = 0, + dilation: int | tuple[int, int, int] = 1, + bias: bool = True, + pad_mode: str = "replicate", + ) -> None: + super().__init__() + + kernel_size = (kernel_size, kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size + + self.pad_mode = pad_mode + self.time_causal_padding = ( + kernel_size[0] // 2, + kernel_size[0] // 2, + kernel_size[1] // 2, + kernel_size[1] // 2, + kernel_size[2] - 1, + 0, + ) + + self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = F.pad(hidden_states, self.time_causal_padding, mode=self.pad_mode) + return self.conv(hidden_states) + + +class HunyuanVideoUpsampleCausal3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int | None = None, + kernel_size: int = 3, + stride: int = 1, + bias: bool = True, + upsample_factor: tuple[float, float, float] = (2, 2, 2), + ) -> None: + super().__init__() + + out_channels = out_channels or in_channels + self.upsample_factor = upsample_factor + + self.conv = HunyuanVideoCausalConv3d(in_channels, out_channels, kernel_size, stride, bias=bias) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_frames = hidden_states.size(2) + dtp = hidden_states.dtype + first_frame, other_frames = hidden_states.split((1, num_frames - 1), dim=2) + first_frame = ( + F.interpolate( + first_frame.squeeze(2), + scale_factor=self.upsample_factor[1:], + mode="nearest", + ) + .unsqueeze(2) + .to(dtp) + ) # force cast + + if num_frames > 1: + other_frames = other_frames.contiguous() + other_frames = F.interpolate(other_frames, scale_factor=self.upsample_factor, mode="nearest").to( + dtp + ) # force cast + hidden_states = torch.cat((first_frame, other_frames), dim=2) + del first_frame + del other_frames + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + else: + hidden_states = first_frame + + hidden_states = self.conv(hidden_states) + return hidden_states + + +class HunyuanVideoDownsampleCausal3D(nn.Module): + def __init__( + self, + channels: int, + out_channels: int | None = None, + padding: int = 1, + kernel_size: int = 3, + bias: bool = True, + stride=2, + ) -> None: + super().__init__() + out_channels = out_channels or channels + + self.conv = HunyuanVideoCausalConv3d(channels, out_channels, kernel_size, stride, padding, bias=bias) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv(hidden_states) + return hidden_states + + +class HunyuanVideoResnetBlockCausal3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int | None = None, + dropout: float = 0.0, + groups: int = 32, + eps: float = 1e-6, + non_linearity: str = "swish", + ) -> None: + super().__init__() + out_channels = out_channels or in_channels + + self.nonlinearity = get_activation(non_linearity) + + self.norm1 = nn.GroupNorm(groups, in_channels, eps=eps, affine=True) + self.conv1 = HunyuanVideoCausalConv3d(in_channels, out_channels, 3, 1, 0) + + self.norm2 = nn.GroupNorm(groups, out_channels, eps=eps, affine=True) + self.dropout = nn.Dropout(dropout) + self.conv2 = HunyuanVideoCausalConv3d(out_channels, out_channels, 3, 1, 0) + + self.conv_shortcut = None + if in_channels != out_channels: + self.conv_shortcut = HunyuanVideoCausalConv3d(in_channels, out_channels, 1, 1, 0) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + dtp = hidden_states.dtype + hidden_states = hidden_states.contiguous() + residual = hidden_states + + hidden_states = self.norm1(hidden_states).to(dtp) # force cast + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv1(hidden_states) + + hidden_states = self.norm2(hidden_states).to(dtp) # force cast + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + residual = self.conv_shortcut(residual) + + hidden_states = hidden_states + residual + return hidden_states + + +class HunyuanVideoMidBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + add_attention: bool = True, + attention_head_dim: int = 1, + ) -> None: + super().__init__() + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + self.add_attention = add_attention + + # There is always at least one resnet + resnets = [ + HunyuanVideoResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + ) + ] + attentions = [] + + for _ in range(num_layers): + if self.add_attention: + attentions.append( + Attention( + in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + eps=resnet_eps, + norm_num_groups=resnet_groups, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + else: + attentions.append(None) + + resnets.append( + HunyuanVideoResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.resnets[0](hidden_states) + + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + batch_size, _, num_frames, height, width = hidden_states.shape + hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3) + mask = prepare_causal_attention_mask( + num_frames, + height * width, + hidden_states.dtype, + hidden_states.device, + batch_size, + ) + hidden_states = attn(hidden_states, attention_mask=mask) + hidden_states = hidden_states.unflatten(1, (num_frames, height, width)).permute(0, 4, 1, 2, 3) + + hidden_states = resnet(hidden_states) + + return hidden_states + + +class HunyuanVideoDownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + add_downsample: bool = True, + downsample_stride: int = 2, + downsample_padding: int = 1, + ) -> None: + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + HunyuanVideoResnetBlockCausal3D( + in_channels=in_channels, + out_channels=out_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + HunyuanVideoDownsampleCausal3D( + out_channels, + out_channels=out_channels, + padding=downsample_padding, + stride=downsample_stride, + ) + ] + ) + else: + self.downsamplers = None + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + return hidden_states + + +class HunyuanVideoUpBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + add_upsample: bool = True, + upsample_scale_factor: tuple[int, int, int] = (2, 2, 2), + ) -> None: + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + HunyuanVideoResnetBlockCausal3D( + in_channels=input_channels, + out_channels=out_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList( + [ + HunyuanVideoUpsampleCausal3D( + out_channels, + out_channels=out_channels, + upsample_factor=upsample_scale_factor, + ) + ] + ) + else: + self.upsamplers = None + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +class HunyuanVideoEncoder3D(nn.Module): + r""" + Causal encoder for 3D video-like data introduced + in [Hunyuan Video](https://huggingface.co/papers/2412.03603). + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: tuple[str, ...] = ( + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + ), + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + double_z: bool = True, + mid_block_add_attention=True, + temporal_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + ) -> None: + super().__init__() + + self.conv_in = HunyuanVideoCausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1) + self.mid_block = None + self.down_blocks = nn.ModuleList([]) + + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + if down_block_type != "HunyuanVideoDownBlock3D": + raise ValueError(f"Unsupported down_block_type: {down_block_type}") + + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio)) + num_time_downsample_layers = int(np.log2(temporal_compression_ratio)) + + if temporal_compression_ratio == 4: + add_spatial_downsample = bool(i < num_spatial_downsample_layers) + add_time_downsample = bool( + i >= (len(block_out_channels) - 1 - num_time_downsample_layers) and not is_final_block + ) + elif temporal_compression_ratio == 8: + add_spatial_downsample = bool(i < num_spatial_downsample_layers) + add_time_downsample = bool(i < num_time_downsample_layers) + else: + raise ValueError(f"Unsupported time_compression_ratio: {temporal_compression_ratio}") + + downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1) + downsample_stride_T = (2,) if add_time_downsample else (1,) + downsample_stride = tuple(downsample_stride_T + downsample_stride_HW) + + down_block = HunyuanVideoDownBlock3D( + num_layers=layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + add_downsample=bool(add_spatial_downsample or add_time_downsample), + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + downsample_stride=downsample_stride, + downsample_padding=0, + ) + + self.down_blocks.append(down_block) + + self.mid_block = HunyuanVideoMidBlock3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + add_attention=mid_block_add_attention, + ) + + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + + conv_out_channels = 2 * out_channels if double_z else out_channels + self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv_in(hidden_states) + + for down_block in self.down_blocks: + hidden_states = down_block(hidden_states) + + hidden_states = self.mid_block(hidden_states) + + hidden_states = self.conv_norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states) + + return hidden_states + + +class HunyuanVideoDecoder3D(nn.Module): + r""" + Causal decoder for 3D video-like data introduced + in [Hunyuan Video](https://huggingface.co/papers/2412.03603). + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + up_block_types: tuple[str, ...] = ( + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + ), + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + mid_block_add_attention=True, + time_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.conv_in = HunyuanVideoCausalConv3d(in_channels, block_out_channels[-1], kernel_size=3, stride=1) + self.up_blocks = nn.ModuleList([]) + + # mid + self.mid_block = HunyuanVideoMidBlock3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + add_attention=mid_block_add_attention, + ) + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + if up_block_type != "HunyuanVideoUpBlock3D": + raise ValueError(f"Unsupported up_block_type: {up_block_type}") + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio)) + num_time_upsample_layers = int(np.log2(time_compression_ratio)) + + if time_compression_ratio == 4: + add_spatial_upsample = bool(i < num_spatial_upsample_layers) + add_time_upsample = bool( + i >= len(block_out_channels) - 1 - num_time_upsample_layers and not is_final_block + ) + else: + raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}") + + upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1) + upsample_scale_factor_T = (2,) if add_time_upsample else (1,) + upsample_scale_factor = tuple(upsample_scale_factor_T + upsample_scale_factor_HW) + + up_block = HunyuanVideoUpBlock3D( + num_layers=self.layers_per_block + 1, + in_channels=prev_output_channel, + out_channels=output_channel, + add_upsample=bool(add_spatial_upsample or add_time_upsample), + upsample_scale_factor=upsample_scale_factor, + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + ) + + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[0], out_channels, kernel_size=3) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + dtp = hidden_states.dtype + hidden_states = self.conv_in(hidden_states) + + hidden_states = self.mid_block(hidden_states) + + for up_block in self.up_blocks: + hidden_states = up_block(hidden_states) + + hidden_states = self.conv_norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states).to(dtp) # force cast + hidden_states = self.conv_out(hidden_states) + + return hidden_states + + +class AutoencoderKLHunyuanVideo(ModelMixin, ConfigMixin, DistributedVaeMixin): + r""" + A VAE model with KL loss for encoding videos into latents + and decoding latent representations into videos. + Introduced in [HunyuanVideo](https://huggingface.co/papers/2412.03603). + + This model inherits from [`ModelMixin`]. Check the superclass + documentation for it's generic methods implemented + for all models (such as downloading or saving). + """ + + @register_to_config + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + latent_channels: int = 16, + down_block_types: tuple[str, ...] = ( + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + ), + up_block_types: tuple[str, ...] = ( + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + ), + block_out_channels: tuple[int] = (128, 256, 512, 512), + layers_per_block: int = 2, + act_fn: str = "silu", + norm_num_groups: int = 32, + scaling_factor: float = 0.476986, + spatial_compression_ratio: int = 8, + temporal_compression_ratio: int = 4, + mid_block_add_attention: bool = True, + ) -> None: + super().__init__() + + self.time_compression_ratio = temporal_compression_ratio + + self.encoder = HunyuanVideoEncoder3D( + in_channels=in_channels, + out_channels=latent_channels, + down_block_types=down_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + double_z=True, + mid_block_add_attention=mid_block_add_attention, + temporal_compression_ratio=temporal_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + ) + + self.decoder = HunyuanVideoDecoder3D( + in_channels=latent_channels, + out_channels=out_channels, + up_block_types=up_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + time_compression_ratio=temporal_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + mid_block_add_attention=mid_block_add_attention, + ) + + self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1) + self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1) + + self.spatial_compression_ratio = spatial_compression_ratio + self.temporal_compression_ratio = temporal_compression_ratio + + self.use_slicing = False + + self.use_tiling = True + + self.use_framewise_encoding = True + self.use_framewise_decoding = True + + self.tile_sample_min_height = 256 + self.tile_sample_min_width = 256 + self.tile_sample_min_num_frames = 16 + + self.tile_sample_stride_height = 192 + self.tile_sample_stride_width = 192 + self.tile_sample_stride_num_frames = 12 + + self.tile_size = None + self.distributed_executor = None + + def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None: + if self.distributed_executor is None: + self.init_distributed() + super().set_parallel_size(parallel_size, mode=mode) + + def _decode_tile_params(self) -> tuple[int, int, int, int]: + ratio = int(self.spatial_compression_ratio) + return ( + self.tile_sample_min_height // ratio, + self.tile_sample_min_width // ratio, + self.tile_sample_stride_height // ratio, + self.tile_sample_stride_width // ratio, + ) + + def _decode_tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + """Same spatial tiles as ``tiled_decode``, one task per tile.""" + _, _, _, height, width = z.shape + tile_h, tile_w, stride_h, stride_w = self._decode_tile_params() + h_starts = list(range(0, height - tile_h + 1, stride_h)) + w_starts = list(range(0, width - tile_w + 1, stride_w)) + tasks: list[TileTask] = [] + for row, i in enumerate(h_starts): + for col, j in enumerate(w_starts): + tile = z[:, :, :, i : i + tile_h, j : j + tile_w] + tasks.append( + TileTask( + tile_id=len(tasks), + grid_coord=(row, col), + tensor=tile, + workload=int(tile.shape[-2] * tile.shape[-1]), + ) + ) + spec = GridSpec( + split_dims=(3, 4), + grid_shape=(len(h_starts), len(w_starts)), + tile_spec={ + "sample_height": height * int(self.spatial_compression_ratio), + "sample_width": width * int(self.spatial_compression_ratio), + }, + output_dtype=z.dtype, + ) + return tasks, spec + + def _decode_tile_exec(self, task: TileTask) -> torch.Tensor: + return self.decoder(self.post_quant_conv(task.tensor)) + + def _decode_tile_merge( + self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec + ) -> torch.Tensor: + n_rows, n_cols = grid_spec.grid_shape + blend_height = self.tile_sample_min_height - self.tile_sample_stride_height + blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + rows = [[coord_tensor_map[(i, j)] for j in range(n_cols)] for i in range(n_rows)] + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + height_lim = self.tile_sample_min_height if i == n_rows - 1 else self.tile_sample_stride_height + width_lim = self.tile_sample_min_width if j == n_cols - 1 else self.tile_sample_stride_width + result_row.append(tile[:, :, :, :height_lim, :width_lim]) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=3) + return dec[:, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]] + + def _encode(self, x: torch.Tensor) -> torch.Tensor: + _, _, num_frames, height, width = x.shape + + if self.use_framewise_decoding and num_frames > (self.tile_sample_min_num_frames + 1): + return self._temporal_tiled_encode(x) + + if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): + return self.tiled_encode(x) + + x = self.encoder(x) + enc = self.quant_conv(x) + return enc + + @apply_forward_hook + def encode( + self, x: torch.Tensor, opt_tiling: bool = True, return_dict: bool = True + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: + r""" + Encode a batch of images into latents. + + Args: + x (`torch.Tensor`): Input batch of images. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] + instead of a plain tuple. + + Returns: + The latent representations of the encoded videos. If `return_dict` is True, a + [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, + otherwise a plain `tuple` is returned. + """ + if opt_tiling: + tile_size, tile_stride = self.get_enc_optimal_tiling(x.shape, device=x.device) + else: + b, _, f, h, w = x.shape + tile_size, tile_stride = (b, f, h, w), (f, h, w) + if tile_size != self.tile_size: + self.tile_size = tile_size + self.apply_tiling(tile_size, tile_stride) + + h = self._encode(x) + + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) + + def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: + _, _, num_frames, height, width = z.shape + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + + if self.use_framewise_decoding and num_frames > (tile_latent_min_num_frames + 1): + return self._temporal_tiled_decode(z, return_dict=return_dict) + + if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height): + return self.tiled_decode(z, return_dict=return_dict) + + z = self.post_quant_conv(z) + dec = self.decoder(z) + + if not return_dict: + return (dec,) + + return DecoderOutput(sample=dec) + + @apply_forward_hook + def decode(self, z: torch.Tensor, return_dict: bool = True, **ignore_kwargs) -> DecoderOutput | torch.Tensor: + r""" + Decode a batch of images. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, + otherwise a plain `tuple` is returned. + """ + tile_size, tile_stride = self.get_dec_optimal_tiling(z.shape, device=z.device) + if tile_size != self.tile_size: + self.tile_size = tile_size + self.apply_tiling(tile_size, tile_stride) + + decoded = self._decode(z).sample + + if not return_dict: + return (decoded,) + + return DecoderOutput(sample=decoded) + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( + y / blend_extent + ) + return b + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( + x / blend_extent + ) + return b + + def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * ( + x / blend_extent + ) + return b + + def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput: + r"""Encode a batch of images using a tiled encoder. + + Args: + x (`torch.Tensor`): Input batch of videos. + + Returns: + `torch.Tensor`: + The latent representation of the encoded videos. + """ + _, _, _, height, width = x.shape + latent_height = height // self.spatial_compression_ratio + latent_width = width // self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = tile_latent_min_height - tile_latent_stride_height + blend_width = tile_latent_min_width - tile_latent_stride_width + + rows = [] + for i in range(0, height - self.tile_sample_min_height + 1, self.tile_sample_stride_height): + row = [] + for j in range(0, width - self.tile_sample_min_width + 1, self.tile_sample_stride_width): + tile = x[ + :, + :, + :, + i : i + self.tile_sample_min_height, + j : j + self.tile_sample_min_width, + ] + tile = self.encoder(tile).clone() + tile = self.quant_conv(tile) + row.append(tile) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + height_lim = tile_latent_min_height if i == len(rows) - 1 else tile_latent_stride_height + width_lim = tile_latent_min_width if j == len(row) - 1 else tile_latent_stride_width + result_row.append(tile[:, :, :, :height_lim, :width_lim]) + result_rows.append(torch.cat(result_row, dim=4)) + + enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width] + return enc + + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: + r""" + Decode a batch of images using a tiled decoder. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, + otherwise a plain `tuple` is returned. + """ + + _, _, _, height, width = z.shape + sample_height = height * self.spatial_compression_ratio + sample_width = width * self.spatial_compression_ratio + + if self.distributed_executor is not None and self.is_distributed_enabled(): + decoded = self.distributed_executor.execute( + z, + DistributedOperator( + split=self._decode_tile_split, + exec=self._decode_tile_exec, + merge=self._decode_tile_merge, + ), + broadcast_result=True, + ) + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = self.tile_sample_min_height - self.tile_sample_stride_height + blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + + rows = [] + for i in range(0, height - tile_latent_min_height + 1, tile_latent_stride_height): + row = [] + for j in range(0, width - tile_latent_min_width + 1, tile_latent_stride_width): + tile = z[ + :, + :, + :, + i : i + tile_latent_min_height, + j : j + tile_latent_min_width, + ] + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile).clone() + row.append(decoded) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + height_lim = self.tile_sample_min_height if i == len(rows) - 1 else self.tile_sample_stride_height + width_lim = self.tile_sample_min_width if j == len(row) - 1 else self.tile_sample_stride_width + result_row.append(tile[:, :, :, :height_lim, :width_lim]) + result_rows.append(torch.cat(result_row, dim=-1)) + + dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width] + + if not return_dict: + return (dec,) + return DecoderOutput(sample=dec) + + def _temporal_tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput: + _, _, num_frames, height, width = x.shape + latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1 + + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio + blend_num_frames = tile_latent_min_num_frames - tile_latent_stride_num_frames + + row = [] + # for i in range(0, num_frames, self.tile_sample_stride_num_frames): + for i in range( + 0, + num_frames - self.tile_sample_min_num_frames + 1, + self.tile_sample_stride_num_frames, + ): + tile = x[:, :, i : i + self.tile_sample_min_num_frames + 1, :, :] + if self.use_tiling and (height > self.tile_sample_min_height or width > self.tile_sample_min_width): + tile = self.tiled_encode(tile) + else: + tile = self.encoder(tile).clone() + tile = self.quant_conv(tile) + if i > 0: + tile = tile[:, :, 1:, :, :] + row.append(tile) + + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_num_frames) + t_lim = tile_latent_min_num_frames if i == len(row) - 1 else tile_latent_stride_num_frames + result_row.append(tile[:, :, :t_lim, :, :]) + else: + result_row.append(tile[:, :, : tile_latent_stride_num_frames + 1, :, :]) + + enc = torch.cat(result_row, dim=2)[:, :, :latent_num_frames] + return enc + + def _temporal_tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: + _, _, num_frames, _, _ = z.shape + num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1 + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio + blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames + + row = [] + for i in range( + 0, + num_frames - tile_latent_min_num_frames + 1, + tile_latent_stride_num_frames, + ): + tile = z[:, :, i : i + tile_latent_min_num_frames + 1, :, :] + if self.use_tiling and (tile.shape[-1] > tile_latent_min_width or tile.shape[-2] > tile_latent_min_height): + decoded = self.tiled_decode(tile, return_dict=True).sample + else: + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile).clone() + if i > 0: + decoded = decoded[:, :, 1:, :, :] + row.append(decoded) + + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_num_frames) + t_lim = self.tile_sample_min_num_frames if i == len(row) - 1 else self.tile_sample_stride_num_frames + result_row.append(tile[:, :, :t_lim, :, :]) + else: + result_row.append(tile[:, :, : self.tile_sample_stride_num_frames + 1, :, :]) + + dec = torch.cat(result_row, dim=2)[:, :, :num_sample_frames] + + if not return_dict: + return (dec,) + return DecoderOutput(sample=dec) + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: + r""" + Args: + sample (`torch.Tensor`): Input sample. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the posterior. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + """ + x = sample + posterior = self.encode(x).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z, return_dict=return_dict) + return dec + + def apply_tiling(self, tile: tuple[int, int, int, int], stride: tuple[int, int, int]): + """Applies tiling.""" + _, ft, ht, wt = tile + fs, hs, ws = stride + + self.use_tiling = True + self.tile_sample_min_num_frames = ft - 1 + self.tile_sample_stride_num_frames = fs + self.tile_sample_min_height = ht + self.tile_sample_min_width = wt + self.tile_sample_stride_height = hs + self.tile_sample_stride_width = ws + + def get_enc_optimal_tiling( + self, shape: list[int], device: torch.device | str | None = None + ) -> tuple[tuple[int, int, int, int], tuple[int, int, int]]: + """Returns optimal tiling for given shape.""" + h, w = shape[3:] + + if device is None: + device = next(self.parameters()).device + device = torch.device(device) + free_mem = float(current_omni_platform.get_free_memory(device)) if device.type != "cpu" else float("inf") + executor = self.distributed_executor + if ( + device.type == "cuda" + and executor is not None + and executor.parallel_size > 1 + and dist.is_initialized() + and self.use_tiling + ): + free = torch.tensor([int(free_mem)], device=device, dtype=torch.int64) + dist.all_reduce(free, op=dist.ReduceOp.MIN, group=executor.group) + free_mem = float(free.item()) + max_area = free_mem / 256 / 17 / 8 + num_vals = 256 * 17 * (h + 32) * (w + 32) + + if h * w < max_area and num_vals < 2**31: + return (1, 17, h, w), (8, h, w) + + def factorize(n, k): + a = sqrt(n / k) + b = sqrt(n * k) + return ceil(a), ceil(b) + + k = max(h / w, w / h) + N = max(ceil(h * w / max_area), ceil(num_vals / 2**31)) + a, b = factorize(N, k) + if h >= w: + wn, hn = a, b + else: + wn, hn = b, a + + if wn > 1: + wt = ceil(w / wn / 8) * 8 + 16 + ws = wt - 32 + else: + wt = w + ws = w + if hn > 1: + ht = ceil(h / hn / 8) * 8 + 16 + hs = ht - 32 + else: + ht = h + hs = h + return (1, 17, ht, wt), (8, hs, ws) + + def get_dec_optimal_tiling( + self, shape: list[int], device: torch.device | str | None = None + ) -> tuple[tuple[int, int, int, int], tuple[int, int, int]]: + """Returns optimal tiling for given shape.""" + b, _, f, h, w = shape + enc_inp_shape = [b, 3, 4 * (f - 1) + 1, 8 * h, 8 * w] + return self.get_enc_optimal_tiling(enc_inp_shape, device=device) + + +__all__ = ["AutoencoderKLHunyuanVideo"] diff --git a/vllm_omni/diffusion/models/kandinsky6/pipeline_kandinsky6.py b/vllm_omni/diffusion/models/kandinsky6/pipeline_kandinsky6.py new file mode 100644 index 000000000..54d9ea267 --- /dev/null +++ b/vllm_omni/diffusion/models/kandinsky6/pipeline_kandinsky6.py @@ -0,0 +1,1851 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Generated standalone Kandinsky 6 vLLM-Omni pipeline. + +Unlike the Diffusers port (one fully self-contained file, shipped into an +external package), this file is one member of a normal multi-file +vLLM-Omni model package — it imports its sibling generated modules +(``modeling_kandinsky6``, ``modeling_kandinsky6_vae``, +``modeling_kandinsky6_audio``, ``scheduling_kandinsky6``) via relative +imports, exactly like vLLM-Omni's other native pipelines (e.g. +``sana_video/pipeline_sana_video.py`` importing from +``.transformer_sana_video``). +""" + +from __future__ import annotations + +import json +import math +import os +from collections.abc import Iterable +from contextlib import nullcontext +from dataclasses import dataclass +from typing import Any, TypedDict + +import numpy as np +import torch +from torch import Tensor, nn +from transformers import ( + AutoConfig, + CLIPTextModel, + CLIPTokenizer, + Qwen2_5_VLForConditionalGeneration, +) +from transformers import AutoProcessor as QwenAutoProcessor + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.interface import SupportAudioOutput, SupportImageInput, SupportsComponentDiscovery +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.models.utils import _load_json +from vllm_omni.diffusion.offloader.config import offload_enabled +from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + +from .kandinsky6_transformer import Kandinsky6Transformer3DModel +from .modeling_kandinsky6_audio import Kandinsky6AudioVAE +from .modeling_kandinsky6_vae import AutoencoderKLHunyuanVideo +from .scheduling_kandinsky6 import KandinskyFlowMatchScheduler + + +class TextEmbeds(TypedDict): + """Output of TextEmbedder.encode() — packed across batch.""" + + text_embeds: Tensor # (total_seq_len, text_dim) Qwen hidden states + pooled_embed: Tensor # (total_seq_len, clip_dim) CLIP pooled; packed per-sample + + +@dataclass +class LatentBundle: + """Packed latent state flowing through the denoising loop. + + Both video and audio tensors use varlen layout: segments from different + batch items are concatenated along dim-0, with cu_seqlens marking + boundaries. audio is None for T2V modes. + """ + + video: Tensor | None # (sum_T, H, W, C) e.g. sum_T = bs*duration + audio: Tensor | None # (sum_A, audio_dim); None for T2V + video_cu_seqlens: Tensor | None # (bs+1,) int32 + audio_cu_seqlens: Tensor | None # (bs+1,) int32; None for T2V + + +@dataclass +class Kandinsky6PipelineOutput: + frames: Tensor # (bs, 3, T, H, W) uint8 + audio: list[np.ndarray] | None = None # list[bs] of (samples,) int16; None for T2V + path: str | None = None # set when save_path was passed to the pipeline + + +def _build_video_input( + video: Tensor, + has_visual_cond: bool, + first_frames: Tensor | None, + video_cu_seqlens: Tensor | None, + visual_cond_scheme: str, +) -> Tensor: + """Appends visual-conditioning channels [cond, mask] to the video latent. + + Schemes (K5 parity): + - ``pretrain``: inject first_frames into cond channels at sequence starts + - ``i2v``: overwrite latents at sequence starts; cond unused + - ``tail_cond_first_frame``: overwrite / mask the last (reference) frame + """ + if not has_visual_cond: + return video + + cond = torch.zeros_like(video) + mask = torch.zeros([*video.shape[:-1], 1], dtype=video.dtype, device=video.device) + + if first_frames is not None: + ff = first_frames.to(device=video.device, dtype=video.dtype) + if video_cu_seqlens is None: + raise ValueError(f"{visual_cond_scheme} requires video_cu_seqlens") + starts = video_cu_seqlens[:-1] + tail_cond = visual_cond_scheme == "tail_cond_first_frame" + inject = video_cu_seqlens[1:] - 1 if tail_cond else starts + + if visual_cond_scheme == "i2v": + video[starts] = ff + elif tail_cond: + video[inject] = ff + elif visual_cond_scheme == "pretrain": + cond[starts] = ff + else: + raise ValueError(f"unknown visual_cond_scheme={visual_cond_scheme!r}") + mask[inject] = 1 + + return torch.cat([video, cond, mask], dim=-1) + + +def _resolve_null_embeds( + null_text_embeds: TextEmbeds | list[TextEmbeds | None], + null_text_rope: Tensor | list[Tensor | None], +) -> tuple[Tensor, Tensor, Tensor | list[Tensor]]: + """Returns (text_embed, pooled_embed, text_rope) for the uncond DiT pass. + + When null_text_embeds is a list of two items (T2VA per-modality nulls), + returns list arguments so the DiT can apply different null text per modality. + """ + if isinstance(null_text_embeds, list): + # null_text_embeds[0] = video null, null_text_embeds[1] = audio null (may be None) + null_v = null_text_embeds[0] + null_a = null_text_embeds[1] + null_rope_v = null_text_rope[0] if isinstance(null_text_rope, list) else null_text_rope + null_rope_a = null_text_rope[1] if isinstance(null_text_rope, list) else None + + if null_a is not None and null_rope_a is not None: + # Different null text per modality — pass as list to DiT + text_embed = [null_v["text_embeds"], null_a["text_embeds"]] + pooled_embed = [null_v["pooled_embed"], null_a["pooled_embed"]] + rope = [null_rope_v, null_rope_a] + else: + text_embed = null_v["text_embeds"] + pooled_embed = null_v["pooled_embed"] + rope = null_rope_v + else: + text_embed = null_text_embeds["text_embeds"] + pooled_embed = null_text_embeds["pooled_embed"] + rope = null_text_rope + + return text_embed, pooled_embed, rope + + +def _raw_dit(dit: nn.Module) -> nn.Module: + inner = getattr(dit, "module", None) + return inner if inner is not None and hasattr(dit, "set_cache") else dit + + +def _rebuild_text_rope(dit: nn.Module, template: Tensor | list[Tensor]) -> Tensor | list[Tensor]: + """Rebuild text RoPE using the standalone Diffusers transformer names.""" + raw = getattr(dit, "module", None) + raw = raw if raw is not None and hasattr(raw, "set_cache") else dit + if isinstance(template, list): + if getattr(raw, "is_multimodal", False): + modules = [raw.video_text_rope_embeddings, raw.audio_text_rope_embeddings] + else: + modules = [raw.text_rope_embeddings] * len(template) + return [compute_rope1d(module, int(item.shape[0])) for module, item in zip(modules, template, strict=True)] + module = raw.video_text_rope_embeddings if getattr(raw, "is_multimodal", False) else raw.text_rope_embeddings + return compute_rope1d(module, int(template.shape[0])) + + +def apply_cfg(cond: Tensor, uncond: Tensor, guidance_weight: float) -> Tensor: + """Classifier-free guidance: uncond + w * (cond - uncond).""" + return uncond + guidance_weight * (cond - uncond) + + +def prepare_video_latents( + bs: int, + duration: int, + H_lat: int, + W_lat: int, + C: int, + seed: int, + device: torch.device | str, + dtype: torch.dtype = torch.bfloat16, +) -> LatentBundle: + """Sample random video noise latent and compute cu_seqlens. + + Returns a LatentBundle with video=(bs*duration, H_lat, W_lat, C) and + uniform video_cu_seqlens=(bs+1,) int32. + """ + g = torch.Generator(device=device) + g.manual_seed(seed) + video = torch.randn(bs * duration, H_lat, W_lat, C, device=device, dtype=dtype, generator=g) + cu = duration * torch.arange(bs + 1, dtype=torch.int32, device=device) + return LatentBundle(video=video, audio=None, video_cu_seqlens=cu, audio_cu_seqlens=None) + + +def audio_latent_duration( + video_latent_frames: int, + *, + fps: float = 24.0, + audio_fps: int = 44100, + downsample_factor: int = 1024, +) -> int: + """Audio latent length matching K5 T2VA: ceil(sample_frames/fps * audio_fps / downsample).""" + sample_frames = (video_latent_frames - 1) * 4 + 1 + return int(math.ceil(sample_frames / fps * audio_fps / downsample_factor)) + + +def prepare_audio_latents( + bundle: LatentBundle, + audio_duration: int, + audio_dim: int, + seed: int, + device: torch.device | str, + dtype: torch.dtype = torch.bfloat16, +) -> LatentBundle: + """Attach random audio noise latent to an existing LatentBundle. + + audio_duration: number of latent audio frames per batch item. + """ + bs = bundle.video_cu_seqlens.shape[0] - 1 if bundle.video_cu_seqlens is not None else 1 + g = torch.Generator(device=device) + g.manual_seed(seed + 1) # offset from video seed for independence + audio = torch.randn(bs * audio_duration, audio_dim, device=device, dtype=dtype, generator=g) + cu = audio_duration * torch.arange(bs + 1, dtype=torch.int32, device=device) + return LatentBundle( + video=bundle.video, + audio=audio, + video_cu_seqlens=bundle.video_cu_seqlens, + audio_cu_seqlens=cu, + ) + + +def compute_rope1d(rope: nn.Module, length: int, device: torch.device | None = None) -> Tensor: + """Build 1-D RoPE for positions ``0..length-1``.""" + if length < 1: + raise ValueError(f"rope length must be positive, got {length}") + # Index the table where it lives (the DiT may still be CPU-resident under + # sequential offload) and move the small result to the requested device. + table_device = next(rope.buffers()).device + out = rope(torch.arange(length, device=table_device)) + return out if device is None or out.device == torch.device(device) else out.to(device) + + +def compute_visual_rope( + rope: nn.Module, + shape: tuple[int, int, int], + scale_factor: tuple[float, float, float], + device: torch.device | None = None, +) -> Tensor: + """Build 3-D RoPE grid for ``shape=(T,H,W)`` with prefix arange positions.""" + T, H, W = (int(shape[0]), int(shape[1]), int(shape[2])) + if T < 1 or H < 1 or W < 1: + raise ValueError(f"visual rope shape must be positive, got {shape}") + table_device = next(rope.buffers()).device + pos = [ + torch.arange(T, device=table_device), + torch.arange(H, device=table_device), + torch.arange(W, device=table_device), + ] + scale = (float(scale_factor[0]), float(scale_factor[1]), float(scale_factor[2])) + out = rope((T, H, W), pos, scale) + return out if device is None or out.device == torch.device(device) else out.to(device) + + +@torch.no_grad() +def postprocess_audio( + bundle: LatentBundle, + audio_vae, + normalization_mode: str = "normalize", +) -> list[np.ndarray] | None: + """Decode audio latents → list of (samples,) int16 numpy arrays. + + Returns None when bundle.audio is None (T2V mode). + + ``normalize`` (default, matches the k6_video production pipeline) + peak-normalizes each waveform before converting it to int16. + ``clip`` preserves the decoded amplitude and saturates it to [-1, 1]. + """ + audio = bundle.audio + if audio is None: + return None + + cu = bundle.audio_cu_seqlens + assert cu is not None + bs = cu.shape[0] - 1 + + # Reverse audio VAE normalization + audio_scaled = audio / getattr(audio_vae, "scaling_factor", 1.0) + audio_scaled = audio_scaled + getattr(audio_vae, "mean_value", 0.0) + + result: list[np.ndarray] = [] + for i in range(bs): + segment = audio_scaled[cu[i].item() : cu[i + 1].item()] # (A, audio_dim) + waveform = ( + audio_vae.wrapped_decode(segment.transpose(1, 0).unsqueeze(0).to(audio_vae.device)) + .squeeze() + .cpu() + .float() + .numpy() + ) + if normalization_mode == "normalize": + peak = np.max(np.abs(waveform)) + if peak > 0: + waveform = waveform / peak * 32767 + elif normalization_mode == "clip": + waveform = np.clip(waveform, -1.0, 1.0) * 32767 + else: + raise ValueError(f"unknown audio normalization_mode={normalization_mode}") + + result.append(waveform.astype(np.int16)) + + return result + + +@torch.no_grad() +def postprocess_video( + bundle: LatentBundle, + vae, + bs: int, +) -> Tensor: + """Decode video latents → (bs, 3, T, H, W) uint8 in [0, 255]. + + Input layout: (sum_T, H_lat, W_lat, C) — packed over batch*duration. + """ + video = bundle.video + assert video is not None + + # (sum_T, H, W, C) → (bs, T, H, W, C) + frames = video.reshape(bs, -1, video.shape[-3], video.shape[-2], video.shape[-1]) + # (bs, T, H, W, C) → (bs, C, T, H, W) for VAE input + frames = (frames / vae.config.scaling_factor).permute(0, 4, 1, 2, 3) + # Hunyuan VAE loads as fp16; DiT latents are bf16 — match weight dtype. + vae_dtype = next(vae.parameters()).dtype + frames = vae.decode(frames.to(dtype=vae_dtype)).sample + + return ((frames.clamp(-1.0, 1.0) + 1.0) * 127.5).to(torch.uint8) + + +@torch.no_grad() +def resize_image( + image: Tensor, + max_area: int, + divisibility: int = 16, + world_size: int = 1, +) -> tuple[Tensor, float]: + """Aspect-preserving resize to fit ``max_area`` with sides divisible by ``divisibility``. + + K5 ``i2v_pipeline.resize_image`` parity. ``image`` is ``(B, C, H, W)``. + Returns ``(resized, scale_k)``. + """ + from math import sqrt + + h, w = image.shape[2:] + area = h * w + div = divisibility + if div == 16: + if world_size in (2, 4): + div *= 2 + elif world_size == 8: + div *= 4 + k = sqrt(max_area / area) / div + new_h = int(round(h * k) * div) + new_w = int(round(w * k) * div) + import torchvision.transforms.functional as TF + + return TF.resize(image, (new_h, new_w)), k + + +def _load_pil_rgb(image: str | object): + from PIL import Image + + if isinstance(image, str): + try: + pil_image = Image.open(image).convert("RGB") + pil_image.load() + except Exception as exc: + raise ValueError(f"Cannot decode i2va input image {image!r}: {exc}") from exc + elif isinstance(image, Image.Image): + pil_image = image.convert("RGB") + else: + raise TypeError(f"i2va image must be a path or PIL image, got {type(image).__name__}") + return pil_image + + +@torch.no_grad() +def encode_i2va_first_frame( + image: str | object, + vae, + device: torch.device | str, + height: int | None = None, + width: int | None = None, + *, + max_area: int | None = None, + divisibility: int = 16, + world_size: int = 1, +) -> tuple[Tensor, int, int]: + """VAE-encode a reference image to one latent frame ``(1, H_lat, W_lat, C)``. + + Two modes (K5 parity): + - ``max_area`` set (and height/width omitted): aspect-preserving resize like + K5 ``resize_image`` / ``get_first_frame_from_image`` — output HxW from input. + - ``height`` + ``width`` set: cover-resize + center-crop to that canvas + (K5 ``encode_i2va_first_frame``). + + Returns ``(latent, pixel_height, pixel_width)``. + """ + import torchvision.transforms.functional as TF + + pil_image = _load_pil_rgb(image) + tensor = TF.pil_to_tensor(pil_image).unsqueeze(0) + + if max_area is not None and (height is None or width is None): + tensor, _ = resize_image(tensor, max_area=max_area, divisibility=divisibility, world_size=world_size) + height, width = int(tensor.shape[-2]), int(tensor.shape[-1]) + else: + if height is None or width is None: + raise ValueError("encode_i2va_first_frame needs height/width or max_area") + src_h, src_w = tensor.shape[-2:] + scale = min(src_h / height, src_w / width) + tensor = TF.resize(tensor, (int(src_h / scale), int(src_w / scale))) + cur_h, cur_w = tensor.shape[-2:] + tensor = TF.crop(tensor, (cur_h - height) // 2, (cur_w - width) // 2, height, width) + + tensor = tensor / 127.5 - 1.0 + vae_dtype = next(vae.parameters()).dtype + # (1, 3, H, W) → (1, 3, 1, H, W) video layout for Hunyuan VAE + tensor = tensor.to(device=device, dtype=vae_dtype).transpose(0, 1).unsqueeze(0) + latent = vae.encode(tensor, opt_tiling=False).latent_dist.sample() + latent = latent.squeeze(0).permute(1, 2, 3, 0) * vae.config.scaling_factor + return latent, int(height), int(width) + + +def append_i2va_tail_condition( + latent_visual: Tensor, + first_frames: Tensor, + batch_size: int, + video_duration: int, +) -> tuple[Tensor, Tensor, Tensor]: + """Append one clean first-frame latent and build generated/reference token ids. + + Returns ``(latent, token_types, generated_mask)`` where ``token_types`` is + 0=generated / 1=reference and ``generated_mask`` selects frames to decode. + """ + _, height, width, dim = latent_visual.shape + first_frames = first_frames.to(device=latent_visual.device, dtype=latent_visual.dtype) + expected = (batch_size, height, width, dim) + if tuple(first_frames.shape) != expected: + raise ValueError(f"first-frame latent shape mismatch: expected {expected}, got {tuple(first_frames.shape)}") + if latent_visual.shape[0] != batch_size * video_duration: + raise ValueError( + "generated visual latent length mismatch: expected " + f"{batch_size * video_duration}, got {latent_visual.shape[0]}" + ) + + latent_visual = torch.cat( + [ + latent_visual.reshape(batch_size, video_duration, height, width, dim), + first_frames[:, None], + ], + dim=1, + ).reshape(batch_size * (video_duration + 1), height, width, dim) + + token_types = torch.cat( + [ + torch.zeros(video_duration, dtype=torch.long, device=latent_visual.device), + torch.ones(1, dtype=torch.long, device=latent_visual.device), + ] + ).repeat(batch_size) + return latent_visual, token_types, token_types == 0 + + +@torch.no_grad() +def denoise_loop( # noqa: PLR0912, PLR0913, PLR0915 + bundle: LatentBundle, + dit: nn.Module, + cfg_model: CFGParallelMixin, + text_embeds: TextEmbeds, + null_text_embeds: TextEmbeds | list[TextEmbeds | None], + visual_rope: Tensor | None, + audio_rope: Tensor | None, + text_rope: Tensor | list[Tensor], + null_text_rope: Tensor | list[Tensor | None], + num_steps: int, + guidance_weight: float, + scheduler: Any, + first_frames: Tensor | None = None, + visual_cond_scheme: str = "pretrain", + sparse_params: dict | None = None, + sample_video: bool = True, + sample_audio: bool = True, + *, + attention_mask: Tensor | None = None, + null_attention_mask: Tensor | None = None, + visual_token_type_ids: Tensor | None = None, + recompute_ropes_each_step: bool = False, + scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), + progress_callback=None, +) -> LatentBundle: + """Single-device denoising loop for T2V(A) driven by a vLLM-Omni scheduler. + + RoPE tensors are normally precomputed by the caller (pipeline). Set + ``recompute_ropes_each_step=True`` to rebuild them every DiT forward (A/B + vs pipeline cache; not for production). + """ + video = bundle.video # (sum_T, H, W, C) or (sum_T, H, W, C+extra) for instruct + audio = bundle.audio # (sum_A, audio_dim) or None + is_multimodal = video is not None and audio is not None + + device = video.device if video is not None else audio.device + guidance_epsilon = 1e-6 + scheduler.set_timesteps(num_steps, device=device) + step_timesteps = scheduler.timesteps + if step_timesteps.shape[0] != num_steps: + raise ValueError( + "vLLM-Omni scheduler returned an unexpected number of timesteps: " + f"expected {num_steps}, got {step_timesteps.shape[0]}" + ) + + bs = ( + bundle.video_cu_seqlens.shape[0] - 1 + if bundle.video_cu_seqlens is not None + else bundle.audio_cu_seqlens.shape[0] - 1 + ) + + null_te, null_pe, null_rope = _resolve_null_embeds(null_text_embeds, null_text_rope) + + out_c: int | None = None # determined after first step, used to strip instruct channels + raw = _raw_dit(dit) + vis_shape = ( + (int(visual_rope.shape[0]), int(visual_rope.shape[1]), int(visual_rope.shape[2])) + if visual_rope is not None + else None + ) + scale = (float(scale_factor[0]), float(scale_factor[1]), float(scale_factor[2])) + tail_cond = visual_cond_scheme == "tail_cond_first_frame" + ref_positions = bundle.video_cu_seqlens[1:] - 1 if tail_cond and bundle.video_cu_seqlens is not None else None + + def _cache_scope(name: str): + cache_context = getattr(dit, "cache_context", None) + return cache_context(name) if callable(cache_context) else nullcontext() + + for step_index, t in enumerate(step_timesteps): + # The scheduler's `timesteps` are already model-scale (sigma * 1000). + t_step: Tensor | list[Tensor] = t.unsqueeze(0).expand(bs) # (bs,) + + model_input_v = ( + _build_video_input( + video, + dit.visual_cond, + first_frames, + bundle.video_cu_seqlens, + visual_cond_scheme, + ) + if video is not None + else None + ) + + # Freeze one modality at t=0 for partial sampling (T2VA only) + if is_multimodal: + t_frozen = scheduler.sigmas[-1].unsqueeze(0).expand(bs) * 1000 + if not sample_audio: + t_step = [t_step, t_frozen] + elif not sample_video: + t_step = [t_frozen, t_step] + + def _forward( + te, + pe, + rope, + attn_mask, + *, + _model_input_v=model_input_v, + _audio=audio, + _t_step=t_step, + ): + vr, ar, tr = visual_rope, audio_rope, rope + if recompute_ropes_each_step: + if vis_shape is not None: + vr = compute_visual_rope(raw.visual_rope_embeddings, vis_shape, scale) + if audio_rope is not None: + ar = compute_rope1d(raw.audio_rope_embeddings, int(audio_rope.shape[0])) + tr = _rebuild_text_rope(raw, rope) + return dit( + x_video=_model_input_v, + x_audio=_audio, + text_embed=te, + pooled_text_embed=pe, + time=_t_step, + visual_rope=vr, + audio_rope=ar, + text_rope=tr, + sparse_params=sparse_params, + attention_mask=attn_mask, + visual_token_type_ids=visual_token_type_ids, + ) + + cache = getattr(raw, "_k6_step_cache", None) + if cache is not None and cache.should_skip(step_index): + vel_cond, vel_uncond = cache.last_cond, cache.last_uncond + else: + use_cfg = abs(guidance_weight - 1.0) > guidance_epsilon + guided = cfg_model.predict_noise_maybe_with_cfg( + do_true_cfg=use_cfg, + true_cfg_scale=guidance_weight, + positive_kwargs={ + "forward": _forward, + "text_embeds": text_embeds["text_embeds"], + "pooled": text_embeds["pooled_embed"], + "rope": text_rope, + "attn_mask": attention_mask, + "cache_scope": _cache_scope("cond"), + }, + negative_kwargs={ + "forward": _forward, + "text_embeds": null_te, + "pooled": null_pe, + "rope": null_rope, + "attn_mask": null_attention_mask, + "cache_scope": _cache_scope("uncond"), + }, + cfg_normalize=False, + ) + # The mixin already applied CFG. Storing the guided velocity on + # both slots keeps a later apply_cfg call as the identity. + vel_cond = guided + vel_uncond = guided + if cache is not None: + cache.store(vel_cond, vel_uncond) + + # Step per modality. A single scheduler step is consumed for each + # denoising iteration; when both modalities are sampled, the same + # scalar step size (derived from `scheduler.sigmas`, not a second + # stateful `.step()` call) is applied to the second modality so the + # scheduler's internal step-index bookkeeping stays in sync with the + # loop counter. + if isinstance(vel_cond, tuple): + vel_v, vel_a = vel_cond + uvel_v, uvel_a = vel_uncond if vel_uncond is not vel_cond else (vel_v, vel_a) + if out_c is None: + out_c = vel_v.shape[-1] + if video is not None and sample_video: + guided_v = apply_cfg(vel_v, uvel_v, guidance_weight) + video = scheduler.step(guided_v, t, video, return_dict=False)[0] + if tail_cond and first_frames is not None and ref_positions is not None: + video[ref_positions] = first_frames.to(device=device, dtype=video.dtype) + if audio is not None and sample_audio: + guided_a = apply_cfg(vel_a, uvel_a, guidance_weight) + if not (video is not None and sample_video): + audio = scheduler.step(guided_a, t, audio, return_dict=False)[0] + else: + step_size = scheduler.sigmas[step_index + 1] - scheduler.sigmas[step_index] + audio += step_size.to(device=audio.device, dtype=audio.dtype) * guided_a + else: + if out_c is None: + out_c = vel_cond.shape[-1] + vel = apply_cfg(vel_cond, vel_uncond, guidance_weight) + if video is not None and sample_video: + video = scheduler.step(vel, t, video, return_dict=False)[0] + if tail_cond and first_frames is not None and ref_positions is not None: + video[ref_positions] = first_frames.to(device=device, dtype=video.dtype) + elif audio is not None and sample_audio: + audio = scheduler.step(vel, t, audio, return_dict=False)[0] + + if progress_callback is not None: + progress_callback() + + # I2V / I2VA: keep injected reference frames unchanged + if first_frames is not None and video is not None and bundle.video_cu_seqlens is not None: + ff = first_frames.to(device=device, dtype=video.dtype) + if visual_cond_scheme == "i2v": + video[bundle.video_cu_seqlens[:-1]] = ff + elif tail_cond and ref_positions is not None: + video[ref_positions] = ff + + # Strip instruct extra-channel padding from the video latent + if out_c is not None and video is not None: + video = video[..., :out_c] + + return LatentBundle( + video=video, + audio=audio, + video_cu_seqlens=bundle.video_cu_seqlens, + audio_cu_seqlens=bundle.audio_cu_seqlens, + ) + + +"""vLLM-Omni native pipeline for Kandinsky 6 TI2VA. + +The port assembler extracts this module's classes/functions into the +generated vLLM-Omni pipeline (``inline_module``). Unlike the Diffusers port +(whose pipeline class only needs to satisfy ``DiffusionPipeline`` and can +rely on ``DiffusionPipeline.from_pretrained()`` for generic component +loading), a vLLM-Omni native pipeline is constructed directly by +``DiffusionModelRegistry`` from an ``OmniDiffusionConfig`` and must own its +own component loading (``_load_components``), request-batch parsing +(``forward``), and ``load_weights``. This mirrors the shape of +vLLM-Omni's other native pipelines (e.g. SanaVideoPipeline, +MiniMaxH3Pipeline): ``encode_prompt`` / ``prepare_latents`` / ``diffuse`` / +``forward`` methods, ``SupportsComponentDiscovery`` for offload/sharding, +and ``SupportAudioOutput`` for the joint video+audio output path. + +Scope note: weights are loaded per Hub component after construction. +``weights_sources`` points at ``transformer/``, ``vae/``, ``text_encoder/``, +``text_encoder_2/``, and ``audio_vae/`` in +``kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers``. ``load_weights`` maps those +prefixes onto the pipeline modules. Tokenizers and the scheduler are not +weight tensors; they are still read from their subfolders here. +""" + +# The assembler inlines these functions/classes from core/algo, core/types, +# and the vLLM overrides package into the same generated module. They are +# intentionally resolved there. +# ruff: noqa: F821 + + +_PROMPT_TEMPLATE = "\n".join( + [ + "<|im_start|>system\nYou are a promt engineer. Describe the video in detail.", + "Describe how the camera moves or shakes, describe the zoom and view angle, whether it follows the objects.", + "Describe the location of the video, main characters or objects and their action.", + "Describe the dynamism of the video and presented actions.", + "Name the visual style of the video: whether it is a professional footage, user generated content, " + "some kind of animation, video game or scren content.", + "Describe the visual effects, postprocessing and transitions if they are presented in the video.", + "Pay attention to the order of key actions shown in the scene.<|im_end|>", + "<|im_start|>user\n{}<|im_end|>", + ] +) +_QWEN_CROP_START = 129 +_CLIP_MAX_LENGTH = 77 +_DEFAULT_NEGATIVE_PROMPT = ( + "Static, 2D cartoon, cartoon, 2d animation, paintings, images, " + "worst quality, low quality, ugly, deformed, walking backwards" +) + + +_TRANSFORMER_WEIGHTS_NAME = "diffusion_pytorch_model.safetensors" + + +_WEIGHT_SUBFOLDERS = ( + ("transformer", "transformer."), + ("vae", "vae."), + ("text_encoder", "text_encoder."), + ("text_encoder_2", "text_encoder_2."), + ("audio_vae", "audio_vae."), +) + + +def _resolve_bundle_root(model: str) -> str: + """Return a local Diffusers bundle directory for a path or Hub repo id.""" + if os.path.isdir(model): + return model + from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific + + return download_weights_from_hf_specific( + model_name_or_path=model, + cache_dir=None, + allow_patterns=["*"], + require_all=True, + ) + + +def _component_sources(model_root: str, *, include_audio_vae: bool) -> list: + """One loader source per weight folder in the Diffusers repo.""" + sources = [] + for subfolder, prefix in _WEIGHT_SUBFOLDERS: + if subfolder == "audio_vae" and not include_audio_vae: + continue + if not os.path.isdir(os.path.join(model_root, subfolder)): + continue + sources.append( + DiffusersPipelineLoader.ComponentSource( + model_or_path=model_root, + subfolder=subfolder, + revision=None, + prefix=prefix, + fall_back_to_pt=True, + ) + ) + return sources + + +def _adapt_k6_weight_name(name: str) -> str: + """Map Hub component keys onto this pipeline's module tree. + + ``text_encoder/`` is Qwen2.5-VL saved with the language stack at + ``model.layers`` and the vision tower at ``visual.*``. This transformers + build nests those under ``model.language_model`` and ``model.visual``. + ``audio_vae/`` stores MMAudio keys at the root (``vae.*``, ``vocoder.*``, + ``mel_converter.*``). ``vae`` and ``vocoder`` live under ``native.tod``; + ``mel_converter`` lives on ``native``. + """ + if name.startswith("transformer."): + # Released Pro-5s checkpoints use the pre-rename block names. + rest = name[len("transformer.") :] + rest = rest.replace(".videoT.", ".video_dec_block.").replace(".audioT.", ".audio_dec_block.") + return "transformer." + rest + if name.startswith("text_encoder."): + rest = name[len("text_encoder.") :] + if rest.startswith("visual."): + return "text_encoder.model." + rest + if rest.startswith("model.") and not rest.startswith(("model.language_model.", "model.visual.")): + return "text_encoder.model.language_model." + rest[len("model.") :] + return name + if not name.startswith("audio_vae."): + return name + rest = name[len("audio_vae.") :] + if rest.startswith("native."): + return name + if rest.startswith(("vae.", "vocoder.")): + rest = f"tod.{rest}" + return f"audio_vae.native.{rest}" + + +def _build_audio_vae(audio_vae_dir: str) -> nn.Module: + with open(os.path.join(audio_vae_dir, "config.json"), encoding="utf-8") as handle: + config = json.load(handle) + return Kandinsky6AudioVAE( + vocoder_config=config.get("vocoder_config"), + mode=str(config.get("mode", "44k")), + need_vae_encoder=bool(config.get("need_vae_encoder", True)), + need_vae_decoder=bool(config.get("need_vae_decoder", True)), + scaling_factor=float(config.get("scaling_factor", 1.0)), + ) + + +def _shard_loaded_weight(param: nn.Parameter, loaded_weight: Tensor) -> Tensor: + """Slice a full checkpoint tensor the way ``ColumnParallelLinear`` / + ``RowParallelLinear.weight_loader`` would, without requiring the + parameter storage to already be allocated. + + Column weights/biases carry ``output_dim``; row weights carry + ``input_dim``. Row biases are replicated (shapes already match). + """ + if getattr(param, "is_sharded_weight", False): + return loaded_weight + loader = getattr(param, "weight_loader", None) + owner = getattr(loader, "__self__", None) + tp_rank = int(getattr(owner, "tp_rank", 0) or 0) + input_dim = getattr(param, "input_dim", None) + output_dim = getattr(param, "output_dim", None) + if ( + input_dim is not None + and loaded_weight.ndim > input_dim + and loaded_weight.shape[input_dim] != param.shape[input_dim] + ): + shard = param.shape[input_dim] + loaded_weight = loaded_weight.narrow(input_dim, tp_rank * shard, shard) + elif ( + output_dim is not None + and loaded_weight.ndim > output_dim + and loaded_weight.shape[output_dim] != param.shape[output_dim] + ): + shard = param.shape[output_dim] + loaded_weight = loaded_weight.narrow(output_dim, tp_rank * shard, shard) + if not loaded_weight.is_contiguous(): + loaded_weight = loaded_weight.contiguous() + return loaded_weight + + +def get_kandinsky6_post_process_func(od_config: OmniDiffusionConfig): + """Factory returning the post-process function registered in + ``vllm_omni/diffusion/registry.py``'s ``_DIFFUSION_POST_PROCESS_FUNCS``. + + Unpacks the ``(video, audio)`` pair set on ``DiffusionOutput.output`` by + ``Kandinsky6TI2VAPipeline.forward`` into the flat ``{"video": ..., + "audio": ..., "audio_sample_rate": ..., "fps": ...}`` payload the + framework's shared audio-output plumbing (``io_support.py`` -> + ``output_formatter.py`` -> ``media_utils.py`` PyAV muxing) already + expects — the same shape MiniMax H3's post-process function produces. + ``output_formatter.normalize_diffusion_postprocess_output`` lifts + ``audio_sample_rate``/``fps`` from this flat form into metadata; wrapping + it in a ``{"payload", "metadata"}`` envelope instead would silently drop + the sample rate (the envelope path only reads + ``metadata["audio"]["sample_rate"]``) and the MP4 would be muxed at the + 24 kHz default, i.e. ~1.8x too slow. + + ``postprocess_audio`` emits int16 PCM; the muxer (and the OpenAI video + route) consume float32 in ``[-1, 1]``, so the waveform is rescaled here. + """ + model_config = dict(getattr(od_config, "model_config", None) or {}) + fps = float(model_config.get("sample_fps", 24.0)) + + def post_process_func(output, output_type: str = "np", sampling_params=None): + video, audio = output["video"], output["audio"] + payload: dict[str, object] = {"video": video, "fps": fps} + if audio is not None: + waveform = audio[0] if isinstance(audio, list) else audio + if isinstance(waveform, torch.Tensor): + waveform = waveform.detach().cpu().numpy() + waveform = np.asarray(waveform) + if np.issubdtype(waveform.dtype, np.integer): + waveform = waveform.astype(np.float32) / float(np.iinfo(waveform.dtype).max) + payload["audio"] = waveform.astype(np.float32, copy=False) + audio_sample_rate = output.get("audio_sample_rate") + if audio_sample_rate is not None: + payload["audio_sample_rate"] = int(audio_sample_rate) + return payload + + return post_process_func + + +def get_kandinsky6_pre_process_func(od_config: OmniDiffusionConfig): + """Factory for the request pre-process hook (image-conditioning input).""" + + def pre_process_func(request): + return request + + return pre_process_func + + +class Kandinsky6TI2VAPipeline( + nn.Module, + CFGParallelMixin, + ProgressBarMixin, + DiffusionPipelineProfilerMixin, + SupportsComponentDiscovery, + SupportAudioOutput, + SupportImageInput, +): + """vLLM-Omni pipeline for text/image-to-video-and-audio generation with Kandinsky 6. + + Args: + transformer: Multimodal K6 transformer used to denoise video and + audio latents (``Kandinsky6Transformer3DModel``). + vae: Video VAE used to decode generated video latents + (``AutoencoderKLHunyuanVideo``). + text_encoder: Qwen2.5-VL text encoder for token-level embeddings. + audio_vae: Audio VAE used to decode generated audio latents + (``Kandinsky6AudioVAE``). May be ``None`` only when + ``sample_audio=False`` is the caller's own default. + scheduler: Flow-matching scheduler used by the denoising loop + (``KandinskyFlowMatchScheduler``). + tokenizer: Qwen2.5-VL processor. + text_encoder_2: CLIP text encoder for pooled embeddings. + tokenizer_2: CLIP tokenizer. + """ + + _dit_modules = ["transformer"] + _encoder_modules = ["text_encoder", "text_encoder_2"] + _vae_modules = ["vae", "audio_vae"] + supports_step_execution = True + support_audio_output = True + support_image_input = True + # Startup warmup is a 512px image-to-video probe. Skip it for the 28B DiT. + dummy_run_num_frames = 0 + default_num_inference_steps = 50 + + def __init__( + self, + transformer: nn.Module | None = None, + vae: nn.Module | None = None, + text_encoder: nn.Module | None = None, + audio_vae: nn.Module | None = None, + scheduler: object | None = None, + tokenizer: object | None = None, + text_encoder_2: nn.Module | None = None, + tokenizer_2: object | None = None, + *, + scale_factor: tuple[float, float, float] = (1.0, 2.0, 2.0), + sample_fps: float = 24.0, + audio_sample_rate: int = 44100, + audio_downsample_factor: int = 1024, + max_sequence_length: int = 1024, + text_token_padding: bool = False, + od_config: OmniDiffusionConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.od_config = od_config + self.device = get_local_device() + + if od_config is not None: + model_config = dict(od_config.model_config or {}) + scale_factor = tuple(float(v) for v in model_config.get("scale_factor", scale_factor)) + sample_fps = float(model_config.get("sample_fps", sample_fps)) + audio_sample_rate = int(model_config.get("audio_sample_rate", audio_sample_rate)) + audio_downsample_factor = int(model_config.get("audio_downsample_factor", audio_downsample_factor)) + max_sequence_length = int(model_config.get("max_sequence_length", max_sequence_length)) + text_token_padding = bool(model_config.get("text_token_padding", text_token_padding)) + ( + transformer, + vae, + text_encoder, + audio_vae, + scheduler, + tokenizer, + text_encoder_2, + tokenizer_2, + ) = self._load_components(od_config, prefix) + + if transformer is None or vae is None or text_encoder is None or scheduler is None: + raise ValueError("Kandinsky6TI2VAPipeline requires transformer, vae, text_encoder, and scheduler.") + + self.transformer = transformer + self.vae = vae + self.text_encoder = text_encoder + self.audio_vae = audio_vae + self.scheduler = scheduler + self.tokenizer = tokenizer + self.text_encoder_2 = text_encoder_2 + self.tokenizer_2 = tokenizer_2 + + self.scale_factor = scale_factor + self.sample_fps = sample_fps + self.audio_sample_rate = audio_sample_rate + self.audio_downsample_factor = audio_downsample_factor + self.max_sequence_length = max_sequence_length + self.text_token_padding = text_token_padding + + # Keep ``self.device`` on the accelerator even when components were + # loaded to CPU for the offload backend: latents/RoPE live on the GPU + # and the sequential-offload hooks move the DiT there on first forward. + first_parameter = next(self.transformer.parameters(), None) + if first_parameter is not None and first_parameter.device.type != "cpu": + self.device = first_parameter.device + + # Layerwise offload streams the DiT block-by-block (~20 GiB footprint + # instead of the 56 GiB resident bf16 DiT). The block-list names depend + # on the checkpoint flavour (multimodal TI2VA vs. video-only), so the + # plan is declared per instance from what the transformer actually has. + block_attrs = tuple( + name + for name in ( + "text_transformer_blocks", + "video_text_transformer_blocks", + "audio_text_transformer_blocks", + "visual_transformer_blocks", + ) + if isinstance(getattr(_raw_dit(self.transformer), name, None), nn.ModuleList) + ) + if block_attrs: + self._offload_plan = OffloadPlan( + block_attrs={"transformer": block_attrs}, + resident_dit_paths=frozenset({"transformer"}), + ) + + if od_config is not None: + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler, + profiler_targets=["forward"], + ) + + # ------------------------------------------------------------------ + # Modules are constructed empty. ``weights_sources`` then loads each + # Hub folder (transformer, vae, text_encoder, text_encoder_2, audio_vae) + # on its own. + # ------------------------------------------------------------------ + + def _load_components( + self, + od_config: OmniDiffusionConfig, + prefix: str, + ) -> tuple[nn.Module, nn.Module, nn.Module, nn.Module | None, object, object, nn.Module | None, object | None]: + del prefix + from diffusers.models.modeling_utils import no_init_weights + + model_root = _resolve_bundle_root(od_config.model) + dtype = getattr(od_config, "dtype", torch.bfloat16) + model_config = dict(od_config.model_config or {}) + + # Keep the empty modules on CPU when offload will own GPU placement. + # Allocating the bf16 DiT (~56 GiB) plus Qwen (~17 GiB) on the GPU + # before weights arrive does not fit an 80 GB device. + load_device = torch.device("cpu") if offload_enabled(od_config) else self.device + + transformer_dir = os.path.join(model_root, "transformer") + with open(os.path.join(transformer_dir, "config.json"), encoding="utf-8") as handle: + transformer_config = json.load(handle) + with torch.device(load_device), no_init_weights(): + transformer = Kandinsky6Transformer3DModel.from_diffusers_config( + transformer_config, + quant_config=getattr(od_config, "quantization_config", None), + ) + + vae_dir = os.path.join(model_root, "vae") + with open(os.path.join(vae_dir, "config.json"), encoding="utf-8") as handle: + vae_config = json.load(handle) + with torch.device(load_device), no_init_weights(): + vae = AutoencoderKLHunyuanVideo.from_config(vae_config) + vae.to(dtype=torch.float16) + + text_encoder_dir = os.path.join(model_root, "text_encoder") + text_encoder_config = AutoConfig.from_pretrained(text_encoder_dir) + with torch.device(load_device), no_init_weights(): + text_encoder = Qwen2_5_VLForConditionalGeneration(text_encoder_config) + text_encoder.to(dtype=dtype) + + # Tokenizer/processor files live in the sibling `tokenizer/` directory + # (standard Diffusers multi-component layout), not inside + # `text_encoder/` itself — that dir only has the model weights/config. + tokenizer = QwenAutoProcessor.from_pretrained(os.path.join(model_root, "tokenizer")) + # encode_prompt() crops embeddings at a fixed offset (_QWEN_CROP_START) + # assuming the fixed system-prompt prefix leads and any padding trails + # — don't trust the checkpoint's own tokenizer_config.json for this; + # pin it explicitly like every sibling pipeline that does fixed-offset + # crops does (sana_video, sana_wm, boogu_image). + tokenizer.tokenizer.padding_side = "right" + + clip_dir = os.path.join(model_root, "text_encoder_2") + clip_config = AutoConfig.from_pretrained(clip_dir) + with torch.device(load_device), no_init_weights(): + text_encoder_2 = CLIPTextModel(clip_config) + text_encoder_2.to(dtype=dtype) + tokenizer_2 = CLIPTokenizer.from_pretrained(os.path.join(model_root, "tokenizer_2")) + + audio_vae = None + audio_vae_dir = os.path.join(model_root, "audio_vae") + include_audio_vae = bool(model_config.get("sample_audio", True)) and os.path.isdir(audio_vae_dir) + if include_audio_vae: + with torch.device(load_device), no_init_weights(): + audio_vae = _build_audio_vae(audio_vae_dir) + audio_vae.to(dtype=dtype) + + self.weights_sources = _component_sources(model_root, include_audio_vae=include_audio_vae) + + # model_config never actually carries a "scheduler_scale" key in + # practice (nothing populates it from the checkpoint), so this always + # fell back to a hardcoded 3.0 — wrong for every known K6 checkpoint, + # which persist their real shift as scheduler/scheduler_config.json's + # "shift" key (same convention lingbot_world's pipeline already reads + # at vllm_omni/diffusion/models/lingbot_world/pipeline.py:504). + checkpoint_scheduler_scale = 5.0 # matches every known K6 generation config and kandinsky-5's own CLI default + scheduler_config_path = os.path.join(model_root, "scheduler", "scheduler_config.json") + if os.path.isfile(scheduler_config_path): + checkpoint_scheduler_scale = float( + _load_json(model_root, "scheduler/scheduler_config.json").get("shift", checkpoint_scheduler_scale) + ) + scheduler_scale = float(model_config.get("scheduler_scale", checkpoint_scheduler_scale)) + scheduler = KandinskyFlowMatchScheduler(scheduler_scale=scheduler_scale, device=self.device) + + return transformer, vae, text_encoder, audio_vae, scheduler, tokenizer, text_encoder_2, tokenizer_2 + + def load_weights(self, weights: Iterable[tuple[str, Tensor]]) -> set[str]: + """Load each Hub folder into its module. + + The framework yields ``transformer.*``, ``vae.*``, ``text_encoder.*``, + ``text_encoder_2.*``, and ``audio_vae.*`` from ``weights_sources``. + """ + from vllm.model_executor.models.utils import AutoWeightsLoader, is_pp_missing_parameter + + quant_config = getattr(self.od_config, "quantization_config", None) + + def adapted(): + for name, tensor in weights: + if name.startswith("transformer."): + key = name[len("transformer.") :] + if is_pp_missing_parameter(key, self.transformer): + continue + if quant_config is not None and tensor.is_floating_point() and tensor.dtype != torch.bfloat16: + tensor = tensor.to(torch.bfloat16) + name = _adapt_k6_weight_name(name) + yield name, tensor + + return AutoWeightsLoader(self).load_weights(adapted()) + + # ------------------------------------------------------------------ + # Prompt encoding — same Qwen2.5-VL + CLIP dual-encoder logic as the + # native/Diffusers ports (framework-agnostic; no vLLM-specific change). + # ------------------------------------------------------------------ + + @staticmethod + def _as_text_embeds(value) -> TextEmbeds: + if not isinstance(value, dict): + raise TypeError("prompt embeddings must be a mapping with text_embeds and pooled_embed") + missing = {"text_embeds", "pooled_embed"} - set(value) + if missing: + raise ValueError(f"prompt embeddings are missing: {sorted(missing)}") + return value + + def encode_prompt(self, text: str) -> tuple[TextEmbeds, Tensor, Tensor | None]: + full_text = _PROMPT_TEMPLATE.format(text) + # Inputs go to the execution device, not the encoder's current one: + # under sequential offload the encoders rest on CPU between requests + # and are swapped onto the GPU by the hook when their forward runs. + qwen_device = self.device + inputs = self.tokenizer( + text=[full_text], + images=None, + videos=None, + max_length=self.max_sequence_length + _QWEN_CROP_START, + truncation=True, + return_tensors="pt", + padding="max_length", + ).to(qwen_device) + qwen_output = self.text_encoder( + input_ids=inputs["input_ids"], + return_dict=True, + output_hidden_states=True, + ) + embeds = qwen_output["hidden_states"][-1][:, _QWEN_CROP_START:] + attention = inputs["attention_mask"][:, _QWEN_CROP_START:].to(dtype=torch.bool) + if self.text_token_padding: + qwen_embeds = embeds[0] + qwen_attention = attention[0] + cu_seqlens = torch.tensor([0, qwen_embeds.shape[0]], dtype=torch.int32, device=qwen_embeds.device) + else: + qwen_embeds = embeds[attention] + qwen_attention = None + cu_seqlens = torch.tensor([0, int(attention.sum().item())], dtype=torch.int32, device=embeds.device) + + clip_device = self.device + clip_inputs = self.tokenizer_2( + [text], + max_length=_CLIP_MAX_LENGTH, + truncation=True, + add_special_tokens=True, + padding="max_length", + return_tensors="pt", + ).to(clip_device) + pooled_embed = self.text_encoder_2(**clip_inputs)["pooler_output"] + return {"text_embeds": qwen_embeds, "pooled_embed": pooled_embed}, cu_seqlens, qwen_attention + + def _encode_prompt_pair( + self, prompt: str, negative_prompt: str + ) -> tuple[TextEmbeds, Tensor, Tensor | None, TextEmbeds, Tensor, Tensor | None]: + prompt_embeds, prompt_cu, prompt_mask = self.encode_prompt(prompt) + negative_embeds, negative_cu, negative_mask = self.encode_prompt(negative_prompt) + return prompt_embeds, prompt_cu, prompt_mask, negative_embeds, negative_cu, negative_mask + + @staticmethod + def _move_text_embeds( + embeds: TextEmbeds, + cu_seqlens: Tensor, + attention_mask: Tensor | None, + *, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[TextEmbeds, Tensor, Tensor | None]: + moved = {key: value.to(device=device, dtype=dtype) for key, value in embeds.items()} + mask = attention_mask.to(device=device) if attention_mask is not None else None + return moved, cu_seqlens.to(device=device), mask + + def _text_ropes( + self, + raw_dit: nn.Module, + *, + text_length: int, + negative_text_length: int, + audio_length: int | None, + device: torch.device, + ) -> tuple[Tensor | list[Tensor], Tensor | list[Tensor], Tensor | None]: + if getattr(raw_dit, "is_multimodal", False): + text_rope = [ + compute_rope1d(raw_dit.video_text_rope_embeddings, text_length, device=device), + compute_rope1d(raw_dit.audio_text_rope_embeddings, text_length, device=device), + ] + negative_text_rope = [ + compute_rope1d(raw_dit.video_text_rope_embeddings, negative_text_length, device=device), + compute_rope1d(raw_dit.audio_text_rope_embeddings, negative_text_length, device=device), + ] + else: + text_rope = compute_rope1d(raw_dit.text_rope_embeddings, text_length, device=device) + negative_text_rope = compute_rope1d(raw_dit.text_rope_embeddings, negative_text_length, device=device) + audio_rope = ( + compute_rope1d(raw_dit.audio_rope_embeddings, audio_length, device=device) + if audio_length is not None + else None + ) + return text_rope, negative_text_rope, audio_rope + + # ------------------------------------------------------------------ + # Latent preparation + # ------------------------------------------------------------------ + + def prepare_latents( + self, + *, + latent_frames: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + seed: int, + sample_audio: bool, + ) -> LatentBundle: + raw_dit = self.transformer + channels = int(getattr(raw_dit, "in_visual_dim", 16)) + bundle = prepare_video_latents( + bs=1, + duration=latent_frames, + H_lat=height // 8, + W_lat=width // 8, + C=channels, + seed=seed, + device=device, + dtype=dtype, + ) + if not sample_audio: + return bundle + if self.audio_vae is None: + raise ValueError("sample_audio=True requires an audio_vae") + + audio_dim = int(getattr(raw_dit, "in_audio_dim", 20)) + downsample_factor = int(getattr(self.audio_vae, "downsample_factor", self.audio_downsample_factor)) + audio_duration = audio_latent_duration( + latent_frames, + fps=self.sample_fps, + audio_fps=self.audio_sample_rate, + downsample_factor=downsample_factor, + ) + return prepare_audio_latents( + bundle, + audio_duration=audio_duration, + audio_dim=audio_dim, + seed=seed, + device=device, + dtype=dtype, + ) + + # ------------------------------------------------------------------ + # Denoise + # ------------------------------------------------------------------ + + def predict_noise( + self, + *, + forward, + text_embeds, + pooled, + rope, + attn_mask, + cache_scope, + ): + """One CFG branch. ``forward`` is the denoise loop's DiT call.""" + with cache_scope: + return forward(text_embeds, pooled, rope, attn_mask) + + def diffuse( + self, + *, + bundle: LatentBundle, + positive: TextEmbeds, + negative: TextEmbeds, + positive_cu: Tensor, + negative_cu: Tensor, + positive_mask: Tensor | None, + negative_mask: Tensor | None, + visual_rope: Tensor, + audio_rope: Tensor | None, + num_inference_steps: int, + guidance_scale: float, + first_frames: Tensor | None, + visual_cond_scheme: str, + sample_audio: bool, + visual_token_type_ids: Tensor | None, + device: torch.device, + ) -> LatentBundle: + text_rope, negative_text_rope, resolved_audio_rope = self._text_ropes( + self.transformer, + text_length=int(positive_cu[-1].item()), + negative_text_length=int(negative_cu[-1].item()), + audio_length=int(bundle.audio_cu_seqlens[-1].item()) + if bundle.audio is not None and bundle.audio_cu_seqlens is not None + else None, + device=device, + ) + with self.progress_bar(total=num_inference_steps) as progress_bar: + return denoise_loop( + bundle=bundle, + dit=self.transformer, + cfg_model=self, + text_embeds=positive, + null_text_embeds=negative, + visual_rope=visual_rope, + audio_rope=resolved_audio_rope, + text_rope=text_rope, + null_text_rope=negative_text_rope, + num_steps=num_inference_steps, + guidance_weight=guidance_scale, + scheduler=self.scheduler, + first_frames=first_frames, + visual_cond_scheme=visual_cond_scheme, + sample_video=True, + sample_audio=sample_audio, + attention_mask=positive_mask, + null_attention_mask=negative_mask, + visual_token_type_ids=visual_token_type_ids, + scale_factor=self.scale_factor, + progress_callback=progress_bar.update, + ) + + # ------------------------------------------------------------------ + # Request entrypoint + # ------------------------------------------------------------------ + + def __call__(self, req: DiffusionRequestBatch) -> DiffusionOutput: + return self.forward(req) + + def _guided_velocity(self, state: Any) -> Tensor | tuple: + """One CFG-combined DiT velocity. Audio, when present, is the second element.""" + extra = state.extra + step_index = int(state.step_index) + scheduler = state.scheduler + t = state.current_timestep + video = state.latents + audio = extra["audio"] + bundle = extra["bundle"] + dit = self.transformer + raw = _raw_dit(dit) + bs = int(extra["bs"]) + guidance_weight = float(extra["guidance_scale"]) + sample_audio = bool(extra["sample_audio"]) + is_multimodal = video is not None and audio is not None + t_step: Tensor | list[Tensor] = t.unsqueeze(0).expand(bs) + model_input_v = ( + _build_video_input( + video, + dit.visual_cond, + extra["first_frames"], + bundle.video_cu_seqlens, + extra["visual_cond_scheme"], + ) + if video is not None + else None + ) + if is_multimodal: + t_frozen = scheduler.sigmas[-1].unsqueeze(0).expand(bs) * 1000 + if not sample_audio: + t_step = [t_step, t_frozen] + elif not extra["sample_video"]: + t_step = [t_frozen, t_step] + + def _cache_scope(name: str): + cache_context = getattr(dit, "cache_context", None) + return cache_context(name) if callable(cache_context) else nullcontext() + + def _forward(te, pe, rope, attn_mask, *, _model_input_v=model_input_v, _audio=audio, _t_step=t_step): + return dit( + x_video=_model_input_v, + x_audio=_audio, + text_embed=te, + pooled_text_embed=pe, + time=_t_step, + visual_rope=extra["visual_rope"], + audio_rope=extra["audio_rope"], + text_rope=rope, + sparse_params=extra["sparse_params"], + attention_mask=attn_mask, + visual_token_type_ids=extra["visual_token_type_ids"], + ) + + cache = getattr(raw, "_k6_step_cache", None) + if cache is not None and cache.should_skip(step_index): + return cache.last_cond + positive = extra["positive"] + guided = self.predict_noise_maybe_with_cfg( + do_true_cfg=bool(state.do_true_cfg), + true_cfg_scale=guidance_weight, + positive_kwargs={ + "forward": _forward, + "text_embeds": positive["text_embeds"], + "pooled": positive["pooled_embed"], + "rope": extra["text_rope"], + "attn_mask": extra["positive_mask"], + "cache_scope": _cache_scope("cond"), + }, + negative_kwargs={ + "forward": _forward, + "text_embeds": extra["null_te"], + "pooled": extra["null_pe"], + "rope": extra["null_rope"], + "attn_mask": extra["negative_mask"], + "cache_scope": _cache_scope("uncond"), + }, + cfg_normalize=False, + ) + if cache is not None: + cache.store(guided, guided) + return guided + + def prepare_encode(self, state: Any, **kwargs: Any) -> Any: + """Encode one request and store everything the denoise step needs.""" + del kwargs + if state.prompt is None: + raise ValueError("Prompt is required for Kandinsky 6 generation.") + prompt_obj = state.prompt + prompt = prompt_obj if isinstance(prompt_obj, str) else (prompt_obj.get("prompt") or "") + negative_prompt = "" if isinstance(prompt_obj, str) else (prompt_obj.get("negative_prompt") or "") + image = None + if not isinstance(prompt_obj, str): + image = (prompt_obj.get("multi_modal_data") or {}).get("image") + if image is None: + image = prompt_obj.get("image") + if isinstance(image, (list, tuple)): + image = image[0] if image else None + sampling = state.sampling + _raw_dit(self.transformer).clear_text_proj_cache() + height = sampling.height or 512 + width = sampling.width or 768 + num_frames = sampling.num_frames or 121 + num_inference_steps = sampling.num_inference_steps or self.default_num_inference_steps + guidance_scale = sampling.guidance_scale if sampling.guidance_scale_provided else 5.0 + extra_args = sampling.extra_args or {} + sample_audio = bool(extra_args.get("sample_audio", True)) + # The Diffusers TI2VA pipeline always appends a reference image as a masked tail frame. + visual_cond_scheme = "tail_cond_first_frame" if image else "pretrain" + generator = sampling.generator + if generator is None and sampling.seed is not None: + generator = torch.Generator(device=self.device).manual_seed(sampling.seed) + seed = ( + int(torch.randint(0, 2**31, (1,), generator=generator, device=generator.device).item()) + if generator is not None + else int(torch.randint(0, 2**31, (1,)).item()) + ) + device = self.device + dtype = getattr(self.transformer, "dtype", torch.bfloat16) + if not isinstance(dtype, torch.dtype) or dtype.itemsize < 2: + dtype = torch.bfloat16 + raw_dit = _raw_dit(self.transformer) + patch_size = tuple(int(v) for v in getattr(raw_dit, "patch_size", (1, 2, 2))) + first_frames = None + if image is not None: + first_frames, height, width = encode_i2va_first_frame(image, self.vae, device, height=height, width=width) + negative_prompt = negative_prompt or _DEFAULT_NEGATIVE_PROMPT + positive, positive_cu, positive_mask, negative, negative_cu, negative_mask = self._encode_prompt_pair( + prompt, negative_prompt + ) + positive, positive_cu, positive_mask = self._move_text_embeds( + positive, positive_cu, positive_mask, device=device, dtype=dtype + ) + negative, negative_cu, negative_mask = self._move_text_embeds( + negative, negative_cu, negative_mask, device=device, dtype=dtype + ) + latent_frames = (num_frames - 1) // 4 + 1 + bundle = self.prepare_latents( + latent_frames=latent_frames, + height=height, + width=width, + dtype=dtype, + device=device, + seed=seed, + sample_audio=sample_audio, + ) + visual_token_type_ids = None + generated_visual_mask = None + if image is not None and visual_cond_scheme == "tail_cond_first_frame": + video, visual_token_type_ids, generated_visual_mask = append_i2va_tail_condition( + bundle.video, first_frames, batch_size=1, video_duration=latent_frames + ) + bundle = LatentBundle( + video=video, + audio=bundle.audio, + video_cu_seqlens=torch.tensor([0, latent_frames + 1], dtype=torch.int32, device=device), + audio_cu_seqlens=bundle.audio_cu_seqlens, + ) + visual_shape = ( + latent_frames // patch_size[0], + (height // 8) // patch_size[1], + (width // 8) // patch_size[2], + ) + visual_rope = compute_visual_rope( + raw_dit.visual_rope_embeddings, visual_shape, self.scale_factor, device=device + ) + if generated_visual_mask is not None: + visual_rope = torch.cat([visual_rope, visual_rope[:1]], dim=0) + scheduler = KandinskyFlowMatchScheduler( + scheduler_scale=self.scheduler.scheduler_scale, + device=device, + ) + scheduler.set_timesteps(num_inference_steps, device=device) + text_rope, null_rope, audio_rope = self._text_ropes( + raw_dit, + text_length=int(positive_cu[-1].item()), + negative_text_length=int(negative_cu[-1].item()), + audio_length=int(bundle.audio_cu_seqlens[-1].item()) + if bundle.audio is not None and bundle.audio_cu_seqlens is not None + else None, + device=device, + ) + null_te, null_pe, null_rope = _resolve_null_embeds(negative, null_rope) + state.latents = bundle.video if bundle.video is not None else bundle.audio + state.timesteps = scheduler.timesteps + state.step_index = 0 + state.scheduler = scheduler + state.do_true_cfg = abs(guidance_scale - 1.0) > 1e-6 + state.extra = { + "bundle": bundle, + "audio": bundle.audio, + "positive": positive, + "positive_mask": positive_mask, + "negative_mask": negative_mask, + "null_te": null_te, + "null_pe": null_pe, + "null_rope": null_rope, + "text_rope": text_rope, + "visual_rope": visual_rope, + "audio_rope": audio_rope, + "sparse_params": None, + "first_frames": first_frames, + "visual_cond_scheme": visual_cond_scheme, + "visual_token_type_ids": visual_token_type_ids, + "generated_visual_mask": generated_visual_mask, + "sample_audio": sample_audio, + "sample_video": True, + "guidance_scale": guidance_scale, + "bs": 1, + "latent_frames": latent_frames, + "audio_normalization": str(extra_args.get("audio_normalization", "normalize")), + "audio_velocity": None, + "out_c": None, + } + return state + + def denoise_step(self, input_batch: Any, **kwargs: Any) -> Tensor | None: + """One CFG-combined video velocity. Audio velocity is stashed on the request.""" + del input_batch + states = kwargs.get("states") or () + if len(states) != 1: + raise ValueError("Kandinsky 6 step execution supports one request at a time.") + state = states[0] + guided = self._guided_velocity(state) + if isinstance(guided, tuple): + video_velocity, audio_velocity = guided + state.extra["audio_velocity"] = audio_velocity + state.extra["out_c"] = int(video_velocity.shape[-1]) + return video_velocity + state.extra["out_c"] = int(guided.shape[-1]) + return guided + + def step_scheduler(self, state: Any, noise_pred: Tensor | None, **kwargs: Any) -> None: + """Apply one scheduler update to video and, when present, audio.""" + del kwargs + if noise_pred is None: + state.step_index += 1 + return + extra = state.extra + scheduler = state.scheduler + t = state.current_timestep + step_index = int(state.step_index) + video = state.latents + scheme = extra["visual_cond_scheme"] + tail_cond = scheme == "tail_cond_first_frame" + bundle = extra["bundle"] + ref_positions = bundle.video_cu_seqlens[1:] - 1 if tail_cond and bundle.video_cu_seqlens is not None else None + first_frames = extra["first_frames"] + if extra["sample_video"] and video is not None and video.ndim >= 4: + video = scheduler.step(noise_pred, t, video, return_dict=False)[0] + if tail_cond and first_frames is not None and ref_positions is not None: + video[ref_positions] = first_frames.to(device=video.device, dtype=video.dtype) + state.latents = video + audio = extra["audio"] + audio_velocity = extra.get("audio_velocity") + if audio is not None and audio_velocity is not None and extra["sample_audio"]: + if not (video is not None and extra["sample_video"] and video.ndim >= 4): + extra["audio"] = scheduler.step(audio_velocity, t, audio, return_dict=False)[0] + else: + step_size = scheduler.sigmas[step_index + 1] - scheduler.sigmas[step_index] + extra["audio"] = audio + step_size.to(device=audio.device, dtype=audio.dtype) * audio_velocity + state.step_index += 1 + + def post_decode(self, state: Any, **kwargs: Any) -> DiffusionOutput: + """Decode the finished video and audio latents.""" + del kwargs + extra = state.extra + video = state.latents + out_c = extra.get("out_c") + if out_c is not None and video is not None and video.ndim >= 4: + video = video[..., :out_c] + mask = extra.get("generated_visual_mask") + if mask is not None and video is not None: + video = video[mask] + device = video.device if video is not None else self.device + bundle = LatentBundle( + video=video, + audio=extra["audio"], + video_cu_seqlens=torch.tensor([0, extra["latent_frames"]], dtype=torch.int32, device=device), + audio_cu_seqlens=extra["bundle"].audio_cu_seqlens, + ) + if state.sampling.output_type == "latent": + video_out: Tensor | np.ndarray = video.unsqueeze(0) + else: + decoded = postprocess_video(bundle, self.vae, bs=1) + video_out = decoded.permute(0, 2, 3, 4, 1).cpu().numpy() + audio_out = ( + postprocess_audio(bundle, self.audio_vae, normalization_mode=extra["audio_normalization"]) + if extra["sample_audio"] + else None + ) + audio_sample_rate = self.audio_sample_rate if audio_out is not None else None + return DiffusionOutput(output={"video": video_out, "audio": audio_out, "audio_sample_rate": audio_sample_rate}) + + @torch.no_grad() + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + if len(req.prompts) != 1: + raise ValueError("Kandinsky6TI2VAPipeline currently supports exactly one prompt per request.") + prompt_obj = req.prompts[0] + prompt = prompt_obj if isinstance(prompt_obj, str) else (prompt_obj.get("prompt") or "") + negative_prompt = "" if isinstance(prompt_obj, str) else (prompt_obj.get("negative_prompt") or "") + image = None + if not isinstance(prompt_obj, str): + # The framework's request envelope carries reference images under + # ``multi_modal_data.image`` (see sana_video / wan2_2 I2V); keep the + # bare ``image`` key for direct/legacy callers. + image = (prompt_obj.get("multi_modal_data") or {}).get("image") + if image is None: + image = prompt_obj.get("image") + if isinstance(image, (list, tuple)): + if len(image) > 1: + raise ValueError("Kandinsky6TI2VAPipeline accepts at most one reference image.") + image = image[0] if image else None + if not prompt: + raise ValueError("Prompt is required for Kandinsky 6 generation.") + + # The transformer's text-projection cache is keyed by tensor.data_ptr(), + # not a value/content key, and is documented as scoped "within a + # generation" — but nothing ever called clear_text_proj_cache() before + # this fix, so it persisted for the whole process lifetime. The pooled + # CLIP embedding's shape never varies, so a freed tensor's CUDA address + # getting reused by the allocator (e.g. after the startup dummy-run + # warmup, or the previous request) silently returns stale conditioning + # from an unrelated prompt. Clear it at the start of every generation. + _raw_dit(self.transformer).clear_text_proj_cache() + + sampling = req.sampling_params + height = sampling.height or 512 + width = sampling.width or 768 + num_frames = sampling.num_frames or 121 + num_inference_steps = sampling.num_inference_steps or self.default_num_inference_steps + guidance_scale = sampling.guidance_scale if sampling.guidance_scale_provided else 5.0 + extra_args = sampling.extra_args or {} + sample_audio = bool(extra_args.get("sample_audio", True)) + # ``normalize`` (peak-normalize each waveform) matches the k6_video + # production pipeline; ``clip`` keeps the raw decoded amplitude. + audio_normalization = str(extra_args.get("audio_normalization", "normalize")) + # The Diffusers TI2VA pipeline always appends a reference image as a masked tail frame. + visual_cond_scheme = "tail_cond_first_frame" if image else "pretrain" + + generator = sampling.generator + if generator is None and sampling.seed is not None: + generator = torch.Generator(device=self.device).manual_seed(sampling.seed) + seed = ( + int(torch.randint(0, 2**31, (1,), generator=generator, device=generator.device).item()) + if generator is not None + else int(torch.randint(0, 2**31, (1,)).item()) + ) + + device = self.device + dtype = getattr(self.transformer, "dtype", torch.bfloat16) + if not isinstance(dtype, torch.dtype) or dtype.itemsize < 2: + # FP8 weights must not become the latent/activation dtype. + dtype = torch.bfloat16 + + raw_dit = self.transformer + patch_size = tuple(int(v) for v in getattr(raw_dit, "patch_size", (1, 2, 2))) + + first_frames = None + if image is not None: + first_frames, height, width = encode_i2va_first_frame(image, self.vae, device, height=height, width=width) + + negative_prompt = negative_prompt or _DEFAULT_NEGATIVE_PROMPT + positive, positive_cu, positive_mask, negative, negative_cu, negative_mask = self._encode_prompt_pair( + prompt, negative_prompt + ) + positive, positive_cu, positive_mask = self._move_text_embeds( + positive, positive_cu, positive_mask, device=device, dtype=dtype + ) + negative, negative_cu, negative_mask = self._move_text_embeds( + negative, negative_cu, negative_mask, device=device, dtype=dtype + ) + + latent_frames = (num_frames - 1) // 4 + 1 + bundle = self.prepare_latents( + latent_frames=latent_frames, + height=height, + width=width, + dtype=dtype, + device=device, + seed=seed, + sample_audio=sample_audio, + ) + + visual_token_type_ids = None + generated_visual_mask = None + if image is not None and visual_cond_scheme == "tail_cond_first_frame": + video, visual_token_type_ids, generated_visual_mask = append_i2va_tail_condition( + bundle.video, first_frames, batch_size=1, video_duration=latent_frames + ) + bundle = LatentBundle( + video=video, + audio=bundle.audio, + video_cu_seqlens=torch.tensor([0, latent_frames + 1], dtype=torch.int32, device=device), + audio_cu_seqlens=bundle.audio_cu_seqlens, + ) + + visual_shape = ( + latent_frames // patch_size[0], + (height // 8) // patch_size[1], + (width // 8) // patch_size[2], + ) + if visual_shape[0] < 1: + raise ValueError(f"invalid visual latent shape {visual_shape}") + visual_rope = compute_visual_rope( + raw_dit.visual_rope_embeddings, visual_shape, self.scale_factor, device=device + ) + if generated_visual_mask is not None: + visual_rope = torch.cat([visual_rope, visual_rope[:1]], dim=0) + + result = self.diffuse( + bundle=bundle, + positive=positive, + negative=negative, + positive_cu=positive_cu, + negative_cu=negative_cu, + positive_mask=positive_mask, + negative_mask=negative_mask, + visual_rope=visual_rope, + audio_rope=None, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + first_frames=first_frames, + visual_cond_scheme=visual_cond_scheme, + sample_audio=sample_audio, + visual_token_type_ids=visual_token_type_ids, + device=device, + ) + + if generated_visual_mask is not None and result.video is not None: + result = LatentBundle( + video=result.video[generated_visual_mask], + audio=result.audio, + video_cu_seqlens=torch.tensor([0, latent_frames], dtype=torch.int32, device=device), + audio_cu_seqlens=result.audio_cu_seqlens, + ) + + if sampling.output_type == "latent": + video_out: Tensor | np.ndarray = result.video.unsqueeze(0) + else: + decoded = postprocess_video(result, self.vae, bs=1) + video_out = decoded.permute(0, 2, 3, 4, 1).cpu().numpy() + + audio_out = ( + postprocess_audio(result, self.audio_vae, normalization_mode=audio_normalization) if sample_audio else None + ) + audio_sample_rate = self.audio_sample_rate if audio_out is not None else None + return DiffusionOutput(output={"video": video_out, "audio": audio_out, "audio_sample_rate": audio_sample_rate}) + + +__all__ = [ + "Kandinsky6TI2VAPipeline", + "get_kandinsky6_post_process_func", + "get_kandinsky6_pre_process_func", +] diff --git a/vllm_omni/diffusion/models/kandinsky6/scheduling_kandinsky6.py b/vllm_omni/diffusion/models/kandinsky6/scheduling_kandinsky6.py new file mode 100644 index 000000000..88321d17c --- /dev/null +++ b/vllm_omni/diffusion/models/kandinsky6/scheduling_kandinsky6.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Native K6 shifted flow-matching Euler scheduler. + +Self-contained (no ``diffusers`` dependency), matching the duck-typed +scheduler interface (``set_timesteps`` / ``timesteps`` / ``sigmas`` / +``step``) that other native vLLM-Omni diffusion pipelines use for their own +per-model schedulers — see ``WanEulerScheduler`` in +``vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler`` for the sibling +pattern this mirrors. + +The math reproduces ``kandinsky.core.algo.flow_matching.flow_match_timesteps`` +exactly: for ``t`` linearly spaced from 1.0 down to 0.0, +``sigma = scale * t / (1 + (scale - 1) * t)``. This is the same shifted +rectified-flow schedule family Wan's own scheduler uses, parameterized by +``scheduler_scale`` (K6's name for what Wan calls ``shift``). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace + +import torch + + +@dataclass +class KandinskyFlowMatchSchedulerOutput: + prev_sample: torch.Tensor + + +class KandinskyFlowMatchScheduler: + """Stateful Euler stepper over K6's shifted flow-matching sigma schedule. + + ``timesteps`` holds the model-scale (``[0, 1000]``-ish) values fed to the + DiT's ``time`` input; ``sigmas`` holds the same schedule unscaled + (``num_inference_steps + 1`` entries, running 1.0 -> 0.0), needed to + derive a step size without consuming a stateful ``step()`` call when one + modality is stepped implicitly alongside another (see the ``vllm`` port's + ``denoise_loop`` override). + """ + + order = 1 + + def __init__(self, scheduler_scale: float = 1.0, device: torch.device | str = "cpu") -> None: + self.scheduler_scale = float(scheduler_scale) + self.device = device + self.config = SimpleNamespace(scheduler_scale=self.scheduler_scale) + + self._step_index: int | None = None + self.timesteps = torch.empty(0, dtype=torch.float32) + self.sigmas = torch.empty(0, dtype=torch.float32) + + self.set_timesteps(num_inference_steps=1, device=self.device) + + @property + def step_index(self) -> int | None: + return self._step_index + + def set_timesteps( + self, + num_inference_steps: int, + device: torch.device | str | None = None, + **kwargs, # noqa: ARG002 - kept for scheduler API compatibility + ) -> None: + device = device or self.device + t = torch.linspace(1.0, 0.0, int(num_inference_steps) + 1, device=device, dtype=torch.float32) + scale = self.scheduler_scale + self.sigmas = scale * t / (1 + (scale - 1) * t) + # DiT `time` input is model-scale, i.e. sigma * 1000 (matches the core + # algo loop's `t.unsqueeze(0).expand(bs) * 1000`). + self.timesteps = self.sigmas[:-1] * 1000.0 + self._step_index = None + + def index_for_timestep(self, timestep: torch.Tensor) -> int: + timestep_t = torch.as_tensor(timestep, device=self.timesteps.device, dtype=self.timesteps.dtype) + return int(torch.argmin(torch.abs(self.timesteps - timestep_t)).item()) + + def step( + self, + model_output: torch.Tensor, + timestep: torch.Tensor, + sample: torch.Tensor, + return_dict: bool = True, + **kwargs, # noqa: ARG002 - kept for scheduler API compatibility + ) -> KandinskyFlowMatchSchedulerOutput | tuple[torch.Tensor]: + if self._step_index is None: + self._step_index = self.index_for_timestep(timestep) + + sigma = self.sigmas[self._step_index] + sigma_next = self.sigmas[self._step_index + 1] + dt = (sigma_next - sigma).to(device=sample.device, dtype=sample.dtype) + prev_sample = sample + dt * model_output + + self._step_index += 1 + + if not return_dict: + return (prev_sample,) + return KandinskyFlowMatchSchedulerOutput(prev_sample=prev_sample) + + def __len__(self) -> int: + return int(self.timesteps.shape[0]) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index a12ec1633..9d05a8251 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -396,6 +396,11 @@ _DIFFUSION_MODELS = { "pipeline_krea2", "Krea2Pipeline", ), + "Kandinsky6TI2VAPipeline": ( + "kandinsky6", + "pipeline_kandinsky6", + "Kandinsky6TI2VAPipeline", + ), } _DIFFUSION_MODELS["MingImageLayeredDiffusionPipeline"] = _DIFFUSION_MODELS["MingImageDiffusionPipeline"] @@ -710,6 +715,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "StableDiffusionXLPipeline": "get_sdxl_image_post_process_func", "Krea2Pipeline": "get_krea2_post_process_func", "HunyuanImage3ForCausalMM": "get_hunyuan_image3_post_process_func", + "Kandinsky6TI2VAPipeline": "get_kandinsky6_post_process_func", } _DIFFUSION_POST_PROCESS_FUNCS["MingImageLayeredDiffusionPipeline"] = _DIFFUSION_POST_PROCESS_FUNCS[ "MingImageDiffusionPipeline" @@ -756,6 +762,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func", "Cosmos3OmniPipeline": "get_cosmos3_pre_process_func", "MammothModa2DiTPipeline": "get_mammoth_moda2_pre_process_func", + "Kandinsky6TI2VAPipeline": "get_kandinsky6_pre_process_func", } diff --git a/vllm_omni/diffusion/utils/hf_utils.py b/vllm_omni/diffusion/utils/hf_utils.py index 9f5e6f06f..c713d417a 100644 --- a/vllm_omni/diffusion/utils/hf_utils.py +++ b/vllm_omni/diffusion/utils/hf_utils.py @@ -158,7 +158,13 @@ def is_diffusion_model(model_name: str) -> bool: logger.debug("Failed to read local %s: %s", filename, e) # Strategy 2: Check using vllm's utility (works for both local and remote models) - config_dict = get_diffusion_model_index(model_name) + # get_hf_file_to_dict can raise for local paths that don't resolve to a Hub + # repo id (see resolve_model_class_name's identical guard for the same call). + try: + config_dict = get_diffusion_model_index(model_name) + except Exception as e: + logger.debug("Failed to fetch diffusion model index via vllm utility: %s", e) + config_dict = None if config_dict is not None and config_dict.get("_class_name") and config_dict.get("_diffusers_version"): logger.debug("Detected diffusion model via a standard Diffusers index") return True diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index d4ecf20a4..67d44760b 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -334,7 +334,12 @@ class OmniOpenAIServingVideo: ) vp.width = vp.width or video_defaults.width vp.height = vp.height or video_defaults.height - if num_frames_provided and vp.num_frames != video_defaults.num_frames: + # Only fixed-duration pipelines (those declaring + # ``duration_seconds``, e.g. MAGI-2) pin the frame count; models + # that merely publish serving defaults (e.g. Kandinsky 6, any + # 4k+1 frame count) still accept an explicit ``num_frames``. + fixed_frame_count = video_defaults.duration_seconds is not None + if fixed_frame_count and num_frames_provided and vp.num_frames != video_defaults.num_frames: raise HTTPException( status_code=HTTPStatus.BAD_REQUEST.value, detail=(f"This diffusion model requires {video_defaults.num_frames} frames; got {vp.num_frames}."), diff --git a/vllm_omni/model_extras/kandinsky6.py b/vllm_omni/model_extras/kandinsky6.py new file mode 100644 index 000000000..e1f73ee32 --- /dev/null +++ b/vllm_omni/model_extras/kandinsky6.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Request fields and example defaults for the Kandinsky 6 TI2VA pipeline.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from vllm_omni.model_extras.video_generation import VideoGenerationDefaults + +# Forwarded to ``OmniDiffusionSamplingParams.extra_args`` by the shared examples +# and the OpenAI video API's ``extra_body``. +KANDINSKY6_EXTRA_BODY_PARAMS = frozenset( + { + # bool, default True: also sample the audio track (TI2VA). False -> video only. + "sample_audio", + # "pretrain" (text-only) | "tail_cond_first_frame" (reference image). + "visual_cond_scheme", + } +) +KANDINSKY6_EXTRA_OUTPUT_PARAMS: frozenset[str] = frozenset() + + +def get_kandinsky6_video_generation_defaults( + extra_body: Mapping[str, Any] | None = None, +) -> VideoGenerationDefaults: + """Kandinsky 6 Pro production geometry: 480x864, 125 frames @ 24 fps, 50 steps, CFG 5. + + These are defaults, not a fixed contract: any ``4k+1`` frame count and + 16-divisible resolution can be requested explicitly, so ``duration_seconds`` + is left unset (the serving layer only pins ``num_frames`` for + fixed-duration pipelines such as MAGI-2). + """ + del extra_body + return VideoGenerationDefaults( + width=864, + height=480, + num_frames=125, + num_inference_steps=50, + fps=24.0, + guidance_scale=5.0, + output="kandinsky6_output.mp4", + # ``None`` lets the pipeline apply its own tuned negative prompt. + default_negative_prompt=None, + duration_seconds=None, + ) diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 2ba1ee01e..233516db8 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -40,6 +40,11 @@ from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as buil from vllm_omni.model_extras.hunyuan_image3 import ( validate_ar_tokenizer as validate_hunyuan_image3_ar_tokenizer, ) +from vllm_omni.model_extras.kandinsky6 import ( + KANDINSKY6_EXTRA_BODY_PARAMS, + KANDINSKY6_EXTRA_OUTPUT_PARAMS, + get_kandinsky6_video_generation_defaults, +) from vllm_omni.model_extras.lingbot_video import LINGBOT_VIDEO_EXTRA_BODY_PARAMS from vllm_omni.model_extras.ltx2 import ( LTX_EXTRA_BODY_PARAMS, @@ -241,6 +246,13 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { "extra_body_params": HELIOS_EXTRA_BODY_PARAMS, "extra_output_params": HELIOS_EXTRA_OUTPUT_PARAMS, }, + "Kandinsky6TI2VAPipeline": { + "extra_body_params": KANDINSKY6_EXTRA_BODY_PARAMS, + "extra_output_params": KANDINSKY6_EXTRA_OUTPUT_PARAMS, + "video_generation_defaults_builder": get_kandinsky6_video_generation_defaults, + # forward() returns uint8 frames already; declared for completeness. + "output_tensor_range": "zero_to_one", + }, "HeliosPyramidPipeline": { "extra_body_params": HELIOS_EXTRA_BODY_PARAMS, "extra_output_params": HELIOS_EXTRA_OUTPUT_PARAMS,