diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 055ba0925..add8b86b3 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -51,6 +51,7 @@ th { | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — | | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | +| `LongCatVideoAvatarPipeline` | LongCat-Video-Avatar-1.5 A2V/AI2V (native single-speaker and multi-speaker AI2V/AVC) | `meituan-longcat/LongCat-Video-Avatar-1.5` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md) | | `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | | — | | `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index d5741a321..6d76cb690 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -159,6 +159,7 @@ The following tables show which models support each feature: | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | +| **LongCat-Video-Avatar-1.5** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **MiniMax-H3** | ✅ (FL2VA) | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ | **Frame Interpolation Support** diff --git a/examples/offline_inference/longcat_video/end2end.py b/examples/offline_inference/longcat_video/end2end.py new file mode 100644 index 000000000..8071dfc25 --- /dev/null +++ b/examples/offline_inference/longcat_video/end2end.py @@ -0,0 +1,319 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import argparse +import json +import re +import subprocess +import tempfile +import wave +from functools import lru_cache +from pathlib import Path +from typing import Any + +import imageio +import numpy as np +import torch + +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.diffusion.models.longcat_video.pipeline_longcat_video_avatar import ( + prepare_longcat_video_avatar_model_for_omni, +) +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.platforms import current_omni_platform + + +def parse_args(): + parser = argparse.ArgumentParser(description="Generate LongCat-Video-Avatar A2V/AI2V video.") + parser.add_argument("--model", default="meituan-longcat/LongCat-Video-Avatar-1.5") + parser.add_argument("--base-model-dir", default=None, help="Optional local LongCat-Video base model directory.") + parser.add_argument("--stage", choices=["at2v", "ai2v"], default="at2v") + parser.add_argument("--prompt", default="A person speaks calmly while facing the camera.") + parser.add_argument("--negative-prompt", default="low quality, blurry, watermark, text") + parser.add_argument("--audio", default=None, help="Input speech audio path.") + parser.add_argument("--image", default=None, help="Input reference image path for AI2V.") + parser.add_argument("--input-json", default=None, help="Official LongCat Avatar JSON case for AI2V.") + parser.add_argument("--output", default="longcat_avatar_output.mp4") + parser.add_argument("--resolution", choices=["480p", "720p"], default="480p") + parser.add_argument("--num-frames", type=int, default=93) + parser.add_argument("--num-segments", default="1", help="Number of AVC segments, or 'auto' to cover full audio.") + parser.add_argument("--num-cond-frames", type=int, default=13) + parser.add_argument("--ref-img-index", type=int, default=10) + parser.add_argument("--mask-frame-range", type=int, default=3) + parser.add_argument("--fps", type=int, default=25) + parser.add_argument("--num-inference-steps", type=int, default=8) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--use-distill", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--use-int8", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument( + "--build-components-on-gpu", + action="store_true", + default=False, + help="Build large Avatar components directly on GPU for faster startup. Requires more peak VRAM.", + ) + parser.add_argument("--use-kv-cache", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--offload-kv-cache", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("--enforce-eager", action="store_true") + return parser.parse_args() + + +def _extract_frames(output: Any): + if isinstance(output, list): + output = output[0] if output else None + if isinstance(output, OmniRequestOutput): + if not output.images: + raise ValueError("No video frames found in OmniRequestOutput.") + frames = output.images[0] if len(output.images) == 1 else output.images + if isinstance(frames, tuple) and len(frames) == 2: + frames = frames[0] + if isinstance(frames, dict): + frames = frames.get("frames") or frames.get("video") + else: + frames = output + if frames is None: + raise ValueError("No video frames found in output.") + return frames + + +def _frames_to_uint8(frames): + frames_np = np.asarray([np.asarray(frame) for frame in frames]) + if frames_np.ndim == 5: + frames_np = frames_np[0] + if np.issubdtype(frames_np.dtype, np.floating): + frames_np = (np.clip(frames_np, 0.0, 1.0) * 255).round() + return np.clip(frames_np, 0, 255).astype(np.uint8) + + +@lru_cache(maxsize=1) +def _ffmpeg_exe() -> str: + try: + import imageio_ffmpeg + + return imageio_ffmpeg.get_ffmpeg_exe() + except Exception as exc: + raise RuntimeError("LongCat Avatar example requires the Python imageio[ffmpeg] dependency.") from exc + + +def _wav_duration(path: Path) -> float: + with wave.open(str(path), "rb") as audio: + return audio.getnframes() / float(audio.getframerate()) + + +def _save_video_with_audio(frames, output_path: Path, audio_path: str, fps: int) -> None: + frames_np = _frames_to_uint8(frames) + output_path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix="longcat_avatar_export_") as tmp: + tmp_dir = Path(tmp) + temp_video = tmp_dir / "video.mp4" + crop_audio = tmp_dir / "audio.wav" + crop_video = tmp_dir / "video_cropped.mp4" + writer = imageio.get_writer(str(temp_video), fps=fps, quality=5) + try: + for frame in frames_np: + writer.append_data(frame) + finally: + writer.close() + + duration = len(frames_np) / fps + ffmpeg = _ffmpeg_exe() + subprocess.run([ffmpeg, "-y", "-i", audio_path, "-t", f"{duration}", str(crop_audio)], check=True) + audio_duration = _wav_duration(crop_audio) + subprocess.run( + [ + ffmpeg, + "-y", + "-i", + str(temp_video), + "-t", + f"{audio_duration}", + "-c:v", + "copy", + str(crop_video), + ], + check=True, + ) + subprocess.run( + [ + ffmpeg, + "-y", + "-i", + str(crop_video), + "-i", + str(crop_audio), + "-c:v", + "libx264", + "-c:a", + "aac", + "-shortest", + str(output_path), + ], + check=True, + ) + + +def _infer_asset_root_from_input_json(input_json: Path) -> Path: + resolved = input_json.resolve() + parts = resolved.parts + if "assets" not in parts: + return resolved.parent + asset_idx = parts.index("assets") + return Path(*parts[:asset_idx]) + + +def _resolve_case_path(path: str, asset_root: Path) -> str: + raw_path = Path(path) + if raw_path.is_absolute(): + return str(raw_path) + candidate = asset_root / raw_path + return str(candidate if candidate.exists() else raw_path) + + +def _speaker_sort_key(name: str) -> tuple[int, str]: + match = re.search(r"(\d+)$", str(name)) + return (int(match.group(1)) if match else 10_000, str(name)) + + +def _load_input_case(input_json: str) -> tuple[dict[str, Any], Path]: + input_json_path = Path(input_json) + with input_json_path.open("r", encoding="utf-8") as f: + case = json.load(f) + return case, _infer_asset_root_from_input_json(input_json_path) + + +def _create_merged_audio_for_case( + audio_paths: dict[str, str], + audio_type: str, + output_path: Path, +) -> str: + sorted_paths = [path for _, path in sorted(audio_paths.items(), key=lambda item: _speaker_sort_key(str(item[0])))] + if len(sorted_paths) == 1: + return sorted_paths[0] + inputs = [arg for path in sorted_paths for arg in ("-i", path)] + audio_type = audio_type.lower() + if audio_type == "para": + filter_complex = f"amix=inputs={len(sorted_paths)}:duration=longest:normalize=0" + elif audio_type == "add": + filter_complex = "".join(f"[{idx}:a]" for idx in range(len(sorted_paths))) + filter_complex += f"concat=n={len(sorted_paths)}:v=0:a=1" + else: + raise NotImplementedError(f"Unsupported LongCat Avatar audio_type {audio_type!r}.") + subprocess.run([_ffmpeg_exe(), "-y", *inputs, "-filter_complex", filter_complex, str(output_path)], check=True) + return str(output_path) + + +def main(): + args = parse_args() + input_case = None + input_asset_root = None + audio_for_mux = args.audio + if args.input_json: + input_case, input_asset_root = _load_input_case(args.input_json) + args.stage = "ai2v" + args.prompt = input_case.get("prompt") or args.prompt + args.image = _resolve_case_path(input_case["cond_image"], input_asset_root) + args.audio = None + if not args.input_json and not args.audio: + raise ValueError("--audio is required unless --input-json is provided.") + if args.stage == "ai2v" and not args.image: + raise ValueError("--image is required when --stage=ai2v.") + model_for_omni = prepare_longcat_video_avatar_model_for_omni(args.model, args.use_int8) + + additional_config = { + "model_type": "avatar-v1.5", + "resolution": args.resolution, + "use_distill": args.use_distill, + "use_int8": args.use_int8, + "build_components_on_gpu": args.build_components_on_gpu, + } + if args.base_model_dir: + additional_config["base_model_dir"] = args.base_model_dir + + omni = Omni( + model=model_for_omni, + model_class_name="LongCatVideoAvatarPipeline", + dtype="bfloat16", + enforce_eager=args.enforce_eager, + additional_config=additional_config, + parallel_config=DiffusionParallelConfig( + ulysses_degree=1, + ring_degree=1, + cfg_parallel_size=1, + tensor_parallel_size=1, + vae_patch_parallel_size=1, + pipeline_parallel_size=1, + ), + ) + + audio_input: str | dict[str, str] | None = args.audio + additional_information = {"stage": args.stage, "resolution": args.resolution} + if input_case is not None and input_asset_root is not None: + audio_input = { + speaker: _resolve_case_path(path, input_asset_root) + for speaker, path in sorted( + input_case["cond_audio"].items(), + key=lambda item: _speaker_sort_key(str(item[0])), + ) + } + additional_information.update( + { + "input_json": args.input_json, + "audio_type": input_case.get("audio_type", "para"), + "bbox": input_case.get("bbox", {}), + } + ) + + prompt = { + "prompt": args.prompt, + "negative_prompt": args.negative_prompt, + "multi_modal_data": {"audio": audio_input}, + "additional_information": additional_information, + } + if args.stage == "ai2v": + prompt["multi_modal_data"]["image"] = args.image + + try: + output = omni.generate( + prompt, + OmniDiffusionSamplingParams( + generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed), + seed=args.seed, + guidance_scale=1.0 if args.use_distill else 4.0, + guidance_scale_2=1.0 if args.use_distill else 4.0, + num_inference_steps=args.num_inference_steps, + num_frames=args.num_frames, + fps=args.fps, + extra_args={ + "stage": args.stage, + "resolution": args.resolution, + "save_fps": args.fps, + "use_distill": args.use_distill, + "use_int8": args.use_int8, + "input_json": args.input_json, + "num_segments": args.num_segments, + "num_cond_frames": args.num_cond_frames, + "ref_img_index": args.ref_img_index, + "mask_frame_range": args.mask_frame_range, + "use_kv_cache": args.use_kv_cache, + "offload_kv_cache": args.offload_kv_cache, + }, + ), + use_tqdm=True, + ) + finally: + omni.close() + + with tempfile.TemporaryDirectory(prefix="longcat_avatar_audio_mux_") as tmp: + if input_case is not None and isinstance(audio_input, dict): + audio_for_mux = _create_merged_audio_for_case( + audio_input, + input_case.get("audio_type", "para"), + Path(tmp) / "merged_audio.wav", + ) + _save_video_with_audio(_extract_frames(output), Path(args.output), audio_for_mux, args.fps) + print(f"Saved video to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index f8328e656..bbb8433f9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -89,6 +89,14 @@ local = [ "FlagEmbedding", ] +# LongCat-Video-Avatar A2V/AI2V runtime dependencies for vocal +# separation and loudness normalization. +longcat-video-avatar = [ + "audio-separator==0.44.2", + "pyloudnorm==0.1.1", + "soundfile>=0.13.1", +] + # Optional Blackwell-only fused-bias FP8 GEMM; package is `quack-kernels` (imported # as `quack`), not the unrelated `quack`. See docs/user_guide/quantization/fp8.md. quack = [ diff --git a/recipes/README.md b/recipes/README.md index 334f44c30..eb4016482 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -55,6 +55,7 @@ recipes/ | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB | | [`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 | | [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom | +| [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB | | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB | | [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) | | [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 | diff --git a/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md b/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md new file mode 100644 index 000000000..67898a0e8 --- /dev/null +++ b/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md @@ -0,0 +1,282 @@ +# LongCat-Video-Avatar-1.5 + +> Audio-driven avatar video generation (AT2V / AI2V) + +## Summary + +- Vendor: meituan-longcat +- Model: `meituan-longcat/LongCat-Video-Avatar-1.5` +- Task: Audio-to-video (AT2V) and audio-and-image-to-video (AI2V) avatar generation +- Mode: Offline inference +- Maintainer: Community + +## When to use this recipe + +Use this recipe to drive a talking or singing avatar from a speech clip, either +from text alone (AT2V) or from a reference image (AI2V), with the native +LongCat-Video-Avatar-1.5 pipeline in vLLM-Omni. It covers: + +1. **Single-speaker AT2V / AI2V** — one 93-frame segment from one audio track. +2. **AVC continuation** — chained segments that cover a longer audio track. +3. **Multi-speaker AI2V** — several audio tracks placed on a reference image + through speaker bounding boxes. + +The example saves an mp4 with the generated frames muxed with the input audio. + +## References + +- Upstream model card: +- Upstream repository and example assets: +- Related example under `examples/`: [`longcat_video/end2end.py`](../../examples/offline_inference/longcat_video/end2end.py) + +## Hardware Support + +## CUDA + +### 1× NVIDIA H100 (80 GB) + +#### Environment + +- OS: Linux +- Python: 3.10+ +- Driver: NVIDIA driver with CUDA 12.x +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Prerequisites + +Install the LongCat Avatar optional runtime dependencies: + +```bash +pip install "vllm-omni[longcat-video-avatar]" +``` + +The example uses the `imageio_ffmpeg` executable installed by vLLM-Omni's +`imageio[ffmpeg]` dependency. + +#### Command + +Audio-to-video (AT2V): + +```bash +python examples/offline_inference/longcat_video/end2end.py \ + --model meituan-longcat/LongCat-Video-Avatar-1.5 \ + --stage at2v \ + --audio /path/to/speech.wav \ + --prompt "A person speaks calmly while facing the camera." \ + --num-frames 93 \ + --num-inference-steps 8 \ + --output longcat_avatar_at2v.mp4 +``` + +Audio-and-image-to-video (AI2V): + +```bash +python examples/offline_inference/longcat_video/end2end.py \ + --model meituan-longcat/LongCat-Video-Avatar-1.5 \ + --stage ai2v \ + --audio /path/to/speech.wav \ + --image /path/to/reference.jpg \ + --prompt "A person speaks calmly while facing the camera." \ + --num-frames 93 \ + --num-inference-steps 8 \ + --output longcat_avatar_ai2v.mp4 +``` + +#### Official Example Assets + +Download the official Avatar example assets: + +```bash +mkdir -p longcat_avatar_assets/assets/avatar/single longcat_avatar_assets/assets/avatar/multi + +curl -L -o longcat_avatar_assets/assets/avatar/single/man.mp3 \ + https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/single/man.mp3 +curl -L -o longcat_avatar_assets/assets/avatar/single/man.png \ + https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/single/man.png +curl -L -o longcat_avatar_assets/assets/avatar/single_example_1.json \ + https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/single_example_1.json + +curl -L -o longcat_avatar_assets/assets/avatar/multi/sing.png \ + https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi/sing.png +curl -L -o longcat_avatar_assets/assets/avatar/multi/sing_man.WAV \ + https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi/sing_man.WAV +curl -L -o longcat_avatar_assets/assets/avatar/multi/sing_woman.WAV \ + https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi/sing_woman.WAV +curl -L -o longcat_avatar_assets/assets/avatar/multi/introduce.png \ + https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi/introduce.png +curl -L -o longcat_avatar_assets/assets/avatar/multi/introduce_man.mp3 \ + https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi/introduce_man.mp3 +curl -L -o longcat_avatar_assets/assets/avatar/multi/introduce_woman.mp3 \ + https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi/introduce_woman.mp3 +curl -L -o longcat_avatar_assets/assets/avatar/multi_example_1.json \ + https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi_example_1.json +curl -L -o longcat_avatar_assets/assets/avatar/multi_example_2.json \ + https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi_example_2.json + +export LONGCAT_VIDEO_ASSET_DIR="$PWD/longcat_avatar_assets/assets/avatar" +export AVATAR_MODEL=meituan-longcat/LongCat-Video-Avatar-1.5 +export PROMPT="$(python -c 'import json, os; print(json.load(open(os.environ["LONGCAT_VIDEO_ASSET_DIR"] + "/single_example_1.json"))["prompt"])')" +``` + +`--base-model-dir` points at a local LongCat-Video base checkpoint for the +tokenizer, text encoder and VAE; omit it to download those components from +Hugging Face. + +```bash +export BASE_MODEL_DIR=/cache/longcat_models/LongCat-Video +``` + +#### Official Asset Smoke Test + +Audio-to-video with the official audio: + +```bash +python examples/offline_inference/longcat_video/end2end.py \ + --model "$AVATAR_MODEL" \ + --base-model-dir "$BASE_MODEL_DIR" \ + --stage at2v \ + --audio "$LONGCAT_VIDEO_ASSET_DIR/single/man.mp3" \ + --prompt "$PROMPT" \ + --num-frames 93 \ + --num-inference-steps 8 \ + --output official_asset_longcat_avatar_at2v.mp4 +``` + +Audio-and-image-to-video with the official audio and image: + +```bash +python examples/offline_inference/longcat_video/end2end.py \ + --model "$AVATAR_MODEL" \ + --base-model-dir "$BASE_MODEL_DIR" \ + --stage ai2v \ + --audio "$LONGCAT_VIDEO_ASSET_DIR/single/man.mp3" \ + --image "$LONGCAT_VIDEO_ASSET_DIR/single/man.png" \ + --prompt "$PROMPT" \ + --num-frames 93 \ + --num-inference-steps 8 \ + --output official_asset_longcat_avatar_ai2v.mp4 +``` + +#### Official Single-Speaker AVC Continuation + +`--num-segments` chains AVC segments; `auto` covers the full input audio. + +```bash +python examples/offline_inference/longcat_video/end2end.py \ + --model "$AVATAR_MODEL" \ + --base-model-dir "$BASE_MODEL_DIR" \ + --input-json "$LONGCAT_VIDEO_ASSET_DIR/single_example_1.json" \ + --num-frames 93 \ + --num-cond-frames 13 \ + --num-segments 5 \ + --num-inference-steps 8 \ + --output official_asset_longcat_avatar_single_example_1_ai2v_5seg.mp4 +``` + +Replace `--num-segments 5` with `--num-segments auto` to run until the full +official audio is covered. That is a slow archival validation path rather than a +regular smoke test: about 24m28s on a Modal H100 for the full +`single_example_1.json` audio with 8 inference steps. + +#### Official Multi-Speaker AI2V + +The multi-speaker cases use several audio tracks with speaker masks and optional +bounding boxes. The command below generates the first 93-frame segment, matching +the first call to the official `generate_ai2v()`. + +```bash +python examples/offline_inference/longcat_video/end2end.py \ + --model "$AVATAR_MODEL" \ + --base-model-dir "$BASE_MODEL_DIR" \ + --input-json "$LONGCAT_VIDEO_ASSET_DIR/multi_example_1.json" \ + --num-frames 93 \ + --num-inference-steps 8 \ + --output official_asset_longcat_avatar_multi_example_1_ai2v.mp4 +``` + +Use `multi_example_2.json` for the sequential-speaking case, and add +`--num-cond-frames 13 --num-segments auto` to cover the full audio with AVC +continuation. + +#### Verification + +```bash +ffprobe -hide_banner longcat_avatar_ai2v.mp4 +``` + +The output must report one video stream and one audio stream: the example muxes +the generated frames with the input audio. + +#### Notes + +**Support matrix** + +| Mode | Speakers | Image required | AVC continuation | +| --- | --- | --- | --- | +| AT2V | Single | No | Yes | +| AI2V | Single | Yes | Yes | +| AT2V | Multi | Not supported | Not supported | +| AI2V | Multi | Yes | Yes | + +Multi-speaker Avatar requires AI2V because speaker bounding boxes and masks are +defined on a reference image; AT2V with several tracks fails with an explicit +error. AVC continuation uses full-audio embeddings only when `--num-segments` is +greater than `1` or set to `auto`; single-segment AT2V/AI2V keeps the shorter +first-segment audio path. + +**Sampling parameters** + +| Parameter | LongCat pipeline value | Description | +| --- | --- | --- | +| `--num-frames` | `num_frames` | Frames generated per segment. The default is 93 frames, or about 3.72 seconds at 25 FPS. LongCat normalizes other values to the temporal-VAE-compatible `4k+1` form. | +| `--num-inference-steps` | `steps` | Denoising steps. The distilled Avatar path uses 8 steps. | +| `--fps` | `save_fps` | Output frame rate, and the rate used to align audio embeddings with video frames. The default is 25 FPS. | +| `sampling_params.guidance_scale` | `text_guidance_scale` | Classifier-free guidance applied to the text prompt. | +| `sampling_params.guidance_scale_2` | `audio_guidance_scale` | Classifier-free guidance applied to the speech audio. | + +With the distilled LoRA enabled (the default), the pipeline uses 8 inference +steps and sets both guidance scales to `1.0`. With `--no-use-distill`, the +example sets both to `4.0`. For AVC continuation every segment holds +`num_frames` frames, but each segment after the first reuses `num_cond_frames` +conditioning frames, so the unique frame count is +`num_frames + (num_segments - 1) * (num_frames - num_cond_frames)`. + +**Key flags** + +- `--stage`: `at2v` for audio-to-video, `ai2v` for audio-and-image-to-video. +- `--resolution`: `480p` or `720p`. +- `--input-json`: official LongCat Avatar JSON case, used for the single-speaker + and multi-speaker AI2V examples. +- `--num-segments`: number of AVC segments, or `auto` to cover the full audio. +- `--num-cond-frames`: previous-frame conditioning window for AVC continuation. +- `--use-int8` / `--no-use-int8`: load the official INT8 Avatar DiT weights by + default, or full precision weights. +- `--use-distill` / `--no-use-distill`: enable the official distilled LoRA path + by default. +- `--build-components-on-gpu`: build large Avatar components directly on GPU for + faster startup, at a higher peak VRAM cost. Disabled by default. +- `--base-model-dir`: local LongCat-Video base model directory for the + tokenizer, text encoder and VAE components. + +**Memory usage** + +Measured on a Modal H100 running the official 93-frame AI2V example with the +INT8 DiT, distilled LoRA and 8 inference steps: + +| Component build | CLI option | Peak GPU memory | Model loading time | +| --- | --- | ---: | ---: | +| CPU (default) | None | About 41.0 GiB | About 130s | +| GPU | `--build-components-on-gpu` | About 56.8 GiB | About 33s | + +The reported peak covers the end-to-end run. Building components on CPU only +reduces initialization memory; denoising still runs on GPU. These numbers are +memory guidance rather than a performance benchmark. + +**Known limitations** + +- Single GPU only. Cache acceleration, sequence/CFG/tensor parallelism, HSDP, + CPU offload, VAE patch parallelism and quantization are not wired up for this + pipeline yet; see the diffusion feature matrix in + [`docs/user_guide/diffusion_features.md`](../../docs/user_guide/diffusion_features.md). +- Offline inference only. Online serving is not supported yet. diff --git a/tests/diffusion/models/longcat_video/test_longcat_video_avatar.py b/tests/diffusion/models/longcat_video/test_longcat_video_avatar.py new file mode 100644 index 000000000..900b51b8e --- /dev/null +++ b/tests/diffusion/models/longcat_video/test_longcat_video_avatar.py @@ -0,0 +1,266 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import json + +import numpy as np +import pytest +import torch +from PIL import Image + +from vllm_omni.diffusion.models.longcat_video.longcat_video_avatar_transformer import ( + LongCatVideoAvatarTransformer3DModel, + _read_config, + replace_linear_with_quantized, +) +from vllm_omni.diffusion.models.longcat_video.pipeline_longcat_video_avatar import ( + _avatar_model_allow_patterns, + _build_multi_speaker_ref_target_masks, + _default_at2v_shape, + _infer_asset_root_from_path, + _prepare_multi_speaker_audio_arrays, + _resolve_num_segments, + prepare_longcat_video_avatar_model_for_omni, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.mark.parametrize( + ("resolution", "expected_shape"), + [ + ("480p", (480, 832)), + ("720p", (736, 1248)), + ], +) +def test_longcat_video_avatar_at2v_default_shape_uses_resolution_bucket( + resolution: str, + expected_shape: tuple[int, int], +): + assert _default_at2v_shape(resolution) == expected_shape + + +@pytest.mark.parametrize( + ("use_int8", "expected_weight_dir", "unexpected_weight_dir"), + [ + (True, "base_model_int8/*", "base_model/*"), + (False, "base_model/*", "base_model_int8/*"), + ], +) +def test_longcat_video_avatar_allow_patterns_download_one_weight_set( + use_int8: bool, + expected_weight_dir: str, + unexpected_weight_dir: str, +): + allow_patterns = _avatar_model_allow_patterns(use_int8) + + assert expected_weight_dir in allow_patterns + assert unexpected_weight_dir not in allow_patterns + assert "whisper-large-v3/*" not in allow_patterns + assert "whisper-large-v3/model.safetensors" in allow_patterns + assert "whisper-large-v3/pytorch_model.bin" not in allow_patterns + assert "whisper-large-v3/flax_model.msgpack" not in allow_patterns + assert "vocal_separator/*" in allow_patterns + + +def test_longcat_video_avatar_prepare_model_adds_omni_metadata(tmp_path): + model_dir = tmp_path / "LongCat-Video-Avatar-1.5" + base_model_dir = model_dir / "base_model_int8" + base_model_dir.mkdir(parents=True) + (base_model_dir / "config.json").write_text('{"model_type": "longcat_avatar"}', encoding="utf-8") + (model_dir / "model_index.json").write_text("{}", encoding="utf-8") + + prepared = prepare_longcat_video_avatar_model_for_omni(str(model_dir), use_int8=True) + + assert prepared == str(model_dir) + model_index = json.loads((model_dir / "model_index.json").read_text(encoding="utf-8")) + assert model_index["_class_name"] == "LongCatVideoAvatarPipeline" + assert model_index["_diffusers_version"] == "0.38.0" + transformer_config = json.loads((model_dir / "transformer" / "config.json").read_text(encoding="utf-8")) + assert transformer_config == {"model_type": "longcat_avatar"} + + +def test_longcat_video_avatar_infers_repo_root_from_official_asset_path(tmp_path): + input_json = tmp_path / "assets" / "avatar" / "single_example_1.json" + input_json.parent.mkdir(parents=True) + input_json.write_text("{}", encoding="utf-8") + + assert _infer_asset_root_from_path(input_json) == tmp_path + + +def test_longcat_video_avatar_falls_back_to_json_parent_without_official_layout(tmp_path): + input_json = tmp_path / "single_example_1.json" + input_json.write_text("{}", encoding="utf-8") + + assert _infer_asset_root_from_path(input_json) == tmp_path + + +@pytest.mark.parametrize( + ("value", "audio_duration", "expected"), + [ + (None, 26.6, 1), + ("1", 26.6, 1), + (3, 26.6, 3), + ("auto", 3.0, 1), + ("auto", 23.8, 8), + ("auto", 26.6, 9), + ], +) +def test_longcat_video_avatar_resolve_num_segments( + value, + audio_duration: float, + expected: int, +): + assert ( + _resolve_num_segments( + value, + audio_duration=audio_duration, + num_frames=93, + num_cond_frames=13, + save_fps=25, + ) + == expected + ) + + +def test_longcat_video_avatar_read_config_filters_non_constructor_metadata(tmp_path): + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "_class_name": "LongCatVideoAvatarTransformer3DModel", + "architectures": ["LongCatVideoAvatarTransformer3DModel"], + "_diffusers_version": "0.35.1", + "model_max_length": 512, + "hidden_size": 4096, + "depth": 48, + "num_heads": 32, + } + ), + encoding="utf-8", + ) + + config = _read_config(config_path) + + assert config == { + "hidden_size": 4096, + "depth": 48, + "num_heads": 32, + } + + +def test_longcat_video_avatar_load_weights_supports_int8_buffers(): + model = LongCatVideoAvatarTransformer3DModel( + hidden_size=4, + depth=0, + num_heads=1, + caption_channels=4, + intermediate_dim=4, + output_dim=4, + audio_channel=4, + context_tokens=1, + ) + replace_linear_with_quantized(model) + + buffer_name = "t_embedder.mlp.0.weight_int8" + buffers = dict(model.named_buffers()) + loaded_weight = torch.ones_like(buffers[buffer_name]) + + loaded_params = model.load_weights([(buffer_name, loaded_weight)]) + + assert buffer_name in loaded_params + assert torch.equal(dict(model.named_buffers())[buffer_name], loaded_weight) + + +def test_longcat_video_avatar_multi_speaker_para_audio_arrays_are_aligned(): + left = np.ones(3, dtype=np.float32) + right = np.ones(5, dtype=np.float32) * 2 + + left_out, right_out = _prepare_multi_speaker_audio_arrays( + [left, right], + audio_type="para", + generate_duration=0.5, + sample_rate=10, + ) + + assert left_out.tolist() == [1, 1, 1, 0, 0] + assert right_out.tolist() == [2, 2, 2, 2, 2] + + +def test_longcat_video_avatar_multi_speaker_add_audio_arrays_are_sequential(): + left = np.ones(3, dtype=np.float32) + right = np.ones(2, dtype=np.float32) * 2 + + left_out, right_out = _prepare_multi_speaker_audio_arrays( + [left, right], + audio_type="add", + generate_duration=0.1, + sample_rate=10, + ) + + assert left_out.tolist() == [1, 1, 1, 0, 0] + assert right_out.tolist() == [0, 0, 0, 2, 2] + + +def test_longcat_video_avatar_multi_speaker_masks_use_explicit_bboxes(): + image = Image.new("RGB", (8, 6)) + masks, use_background = _build_multi_speaker_ref_target_masks( + image, + { + "person1": [1, 1, 4, 3], + "person2": [2, 4, 5, 7], + }, + ) + + assert not use_background + assert masks.shape == (3, 6, 8) + assert masks[0, 1:4, 1:3].sum() == 6 + assert masks[1, 2:5, 4:7].sum() == 9 + assert masks[2, 0, 0] == 1 + assert masks[2, 2, 2] == 0 + + +def test_longcat_video_avatar_multi_speaker_masks_default_to_left_right_split(): + image = Image.new("RGB", (10, 8)) + masks, use_background = _build_multi_speaker_ref_target_masks(image, None) + + assert not use_background + assert masks.shape == (3, 8, 10) + assert masks[0, :, :5].sum() > 0 + assert masks[0, :, 5:].sum() == 0 + assert masks[1, :, :5].sum() == 0 + assert masks[1, :, 5:].sum() > 0 + + +def test_longcat_video_avatar_transformer_accepts_multi_speaker_masks(): + model = LongCatVideoAvatarTransformer3DModel( + hidden_size=4, + depth=0, + num_heads=1, + caption_channels=4, + intermediate_dim=4, + output_dim=4, + audio_channel=4, + context_tokens=1, + ) + hidden_states = torch.randn(1, 16, 5, 4, 4) + timestep = torch.zeros(1, 5) + encoder_hidden_states = torch.randn(1, 1, 4, 4) + audio_embs = torch.randn(2, 5, 5, 12, 4) + ref_target_masks = torch.zeros(3, 4, 4) + ref_target_masks[0, :, :2] = 1 + ref_target_masks[1, :, 2:] = 1 + ref_target_masks[2] = torch.where(ref_target_masks[0] + ref_target_masks[1] > 0, 0, 1) + + output = model( + hidden_states=hidden_states, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + num_cond_latents=1, + audio_embs=audio_embs, + ref_target_masks=ref_target_masks, + ) + + assert output.shape == hidden_states.shape diff --git a/tests/e2e/offline_inference/test_longcat_video_avatar.py b/tests/e2e/offline_inference/test_longcat_video_avatar.py new file mode 100644 index 000000000..f0fa81eb3 --- /dev/null +++ b/tests/e2e/offline_inference/test_longcat_video_avatar.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Generator + +import numpy as np +import pytest +import soundfile as sf +from PIL import Image + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from vllm_omni.diffusion.models.longcat_video.pipeline_longcat_video_avatar import ( + prepare_longcat_video_avatar_model_for_omni, +) +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +MODEL = "meituan-longcat/LongCat-Video-Avatar-1.5" +PROMPT = "A person speaks calmly while facing the camera." +NEGATIVE_PROMPT = "low quality, blurry, watermark, text" + + +@pytest.fixture() +def synthetic_audio_path(tmp_path): + sample_rate = 16000 + duration_s = 2.0 + t = np.arange(int(sample_rate * duration_s), dtype=np.float32) / sample_rate + audio = (0.08 * np.sin(2 * np.pi * 220 * t)).astype(np.float32) + path = tmp_path / "avatar_reference.wav" + sf.write(path, audio, sample_rate) + return str(path) + + +@pytest.fixture() +def synthetic_multi_audio_paths(tmp_path): + sample_rate = 16000 + duration_s = 2.0 + t = np.arange(int(sample_rate * duration_s), dtype=np.float32) / sample_rate + left = (0.08 * np.sin(2 * np.pi * 220 * t)).astype(np.float32) + right = (0.08 * np.sin(2 * np.pi * 330 * t)).astype(np.float32) + left_path = tmp_path / "avatar_left.wav" + right_path = tmp_path / "avatar_right.wav" + sf.write(left_path, left, sample_rate) + sf.write(right_path, right, sample_rate) + return {"person1": str(left_path), "person2": str(right_path)} + + +@pytest.fixture() +def synthetic_image_path(tmp_path): + path = tmp_path / "avatar_reference.png" + Image.new("RGB", (768, 512), (130, 104, 88)).save(path) + return str(path) + + +def _sampling( + stage: str, + *, + num_frames: int = 5, + extra_args: dict | None = None, +) -> OmniDiffusionSamplingParams: + sampling_extra_args = { + "stage": stage, + "resolution": "480p", + "save_fps": 25, + "use_distill": True, + "use_int8": True, + } + if extra_args: + sampling_extra_args.update(extra_args) + return OmniDiffusionSamplingParams( + num_frames=num_frames, + fps=25, + num_inference_steps=8, + guidance_scale=1.0, + guidance_scale_2=1.0, + seed=42, + extra_args=sampling_extra_args, + ) + + +@pytest.fixture(scope="module") +def longcat_omni_runner() -> Generator[OmniRunner, None, None]: + model = prepare_longcat_video_avatar_model_for_omni(MODEL, use_int8=True) + with OmniRunner( + model, + seed=42, + stage_configs_path=None, + model_class_name="LongCatVideoAvatarPipeline", + additional_config={ + "model_type": "avatar-v1.5", + "resolution": "480p", + "use_distill": True, + "use_int8": True, + }, + ) as runner: + yield runner + + +def _generate_avatar_frames( + omni_runner: OmniRunner, + *, + stage: str, + audio_path: str | dict[str, str], + image_path: str | None = None, + additional_information: dict | None = None, + num_frames: int = 5, + sampling_extra_args: dict | None = None, +): + prompt = { + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "multi_modal_data": {"audio": audio_path}, + "additional_information": additional_information or {}, + } + prompt["additional_information"].update({"stage": stage, "resolution": "480p"}) + if image_path is not None: + prompt["multi_modal_data"]["image"] = image_path + + outputs = omni_runner.generate( + [prompt], + [_sampling(stage, num_frames=num_frames, extra_args=sampling_extra_args)], + ) + assert outputs + output = outputs[0] + assert output.images + frames = output.images[0] + if isinstance(frames, tuple) and len(frames) == 2: + frames = frames[0] + if isinstance(frames, dict): + frames = frames.get("frames") or frames.get("video") + assert isinstance(frames, list) + return frames + + +def _assert_frames(frames, *, expected_size: tuple[int, int], expected_count: int = 5): + assert len(frames) == expected_count + for frame in frames: + assert isinstance(frame, Image.Image) + assert frame.mode == "RGB" + assert frame.size == expected_size + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1}) +def test_longcat_video_avatar_at2v(longcat_omni_runner: OmniRunner, synthetic_audio_path): + frames = _generate_avatar_frames( + longcat_omni_runner, + stage="at2v", + audio_path=synthetic_audio_path, + ) + _assert_frames(frames, expected_size=(832, 480)) + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1}) +def test_longcat_video_avatar_ai2v( + longcat_omni_runner: OmniRunner, + synthetic_audio_path, + synthetic_image_path, +): + frames = _generate_avatar_frames( + longcat_omni_runner, + stage="ai2v", + audio_path=synthetic_audio_path, + image_path=synthetic_image_path, + ) + _assert_frames(frames, expected_size=(768, 512)) + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1}) +def test_longcat_video_avatar_multi_speaker_ai2v( + longcat_omni_runner: OmniRunner, + synthetic_multi_audio_paths, + synthetic_image_path, +): + frames = _generate_avatar_frames( + longcat_omni_runner, + stage="ai2v", + audio_path=synthetic_multi_audio_paths, + image_path=synthetic_image_path, + additional_information={ + "audio_type": "para", + "bbox": { + "person1": [32, 48, 384, 336], + "person2": [32, 432, 384, 720], + }, + }, + ) + _assert_frames(frames, expected_size=(768, 512)) + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1}) +def test_longcat_video_avatar_single_speaker_ai2v_avc_continuation( + longcat_omni_runner: OmniRunner, + synthetic_audio_path, + synthetic_image_path, +): + frames = _generate_avatar_frames( + longcat_omni_runner, + stage="ai2v", + audio_path=synthetic_audio_path, + image_path=synthetic_image_path, + num_frames=17, + sampling_extra_args={ + "num_segments": 2, + "num_cond_frames": 5, + "use_kv_cache": True, + }, + ) + _assert_frames(frames, expected_size=(768, 512), expected_count=29) + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1}) +def test_longcat_video_avatar_multi_speaker_ai2v_avc_continuation( + longcat_omni_runner: OmniRunner, + synthetic_multi_audio_paths, + synthetic_image_path, +): + frames = _generate_avatar_frames( + longcat_omni_runner, + stage="ai2v", + audio_path=synthetic_multi_audio_paths, + image_path=synthetic_image_path, + additional_information={ + "audio_type": "para", + "bbox": { + "person1": [32, 48, 384, 336], + "person2": [32, 432, 384, 720], + }, + }, + num_frames=17, + sampling_extra_args={ + "num_segments": 2, + "num_cond_frames": 5, + "use_kv_cache": True, + }, + ) + _assert_frames(frames, expected_size=(768, 512), expected_count=29) diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index f5d2768c5..539f25c72 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -40,6 +40,7 @@ EXCLUDED_MODELS = [ "WanT2VDMD2Pipeline", "WanI2VDMD2Pipeline", "LongCatImagePipeline", + "LongCatVideoAvatarPipeline", "BagelPipeline", "BooguImagePipeline", "LancePipeline", diff --git a/vllm_omni/diffusion/models/longcat_video/__init__.py b/vllm_omni/diffusion/models/longcat_video/__init__.py new file mode 100644 index 000000000..a1fa7e3b0 --- /dev/null +++ b/vllm_omni/diffusion/models/longcat_video/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.diffusion.models.longcat_video.longcat_video_avatar_transformer import ( + LongCatVideoAvatarTransformer3DModel, +) +from vllm_omni.diffusion.models.longcat_video.pipeline_longcat_video_avatar import ( + LongCatVideoAvatarPipeline, + get_longcat_video_avatar_post_process_func, + get_longcat_video_avatar_pre_process_func, +) + +__all__ = [ + "LongCatVideoAvatarPipeline", + "LongCatVideoAvatarTransformer3DModel", + "get_longcat_video_avatar_post_process_func", + "get_longcat_video_avatar_pre_process_func", +] diff --git a/vllm_omni/diffusion/models/longcat_video/longcat_video_avatar_transformer.py b/vllm_omni/diffusion/models/longcat_video/longcat_video_avatar_transformer.py new file mode 100644 index 000000000..c9928e1b0 --- /dev/null +++ b/vllm_omni/diffusion/models/longcat_video/longcat_video_avatar_transformer.py @@ -0,0 +1,1675 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import inspect +import json +import math +import os +from collections import OrderedDict +from collections.abc import Iterable +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange, repeat +from safetensors.torch import load_file +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention as VllmAttention + + +def linear_fp32(linear: nn.Module, x: torch.Tensor) -> torch.Tensor: + if not hasattr(linear, "weight"): + return linear(x.float()) + bias = None if linear.bias is None else linear.bias.float() + return F.linear(x.float(), linear.weight.float(), bias) + + +def silu_linear_fp32(module: nn.Sequential, x: torch.Tensor) -> torch.Tensor: + return linear_fp32(module[1], F.silu(x.float())) + + +class FeedForwardSwiGLU(nn.Module): + def __init__( + self, + dim: int, + hidden_dim: int, + multiple_of: int = 256, + ffn_dim_multiplier: float | None = None, + ): + super().__init__() + hidden_dim = int(2 * hidden_dim / 3) + if ffn_dim_multiplier is not None: + hidden_dim = int(ffn_dim_multiplier * hidden_dim) + hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + + self.dim = dim + self.hidden_dim = hidden_dim + self.w1 = nn.Linear(dim, hidden_dim, bias=False) + self.w2 = nn.Linear(hidden_dim, dim, bias=False) + self.w3 = nn.Linear(dim, hidden_dim, bias=False) + + def forward(self, x): + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +class RMSNorm_FP32(torch.nn.Module): + def __init__(self, dim: int, eps: float): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + output = self._norm(x.float()).type_as(x) + return output * self.weight + + +class LayerNorm_FP32(nn.LayerNorm): + def __init__(self, dim, eps, elementwise_affine): + super().__init__(dim, eps=eps, elementwise_affine=elementwise_affine) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + origin_dtype = inputs.dtype + out = F.layer_norm( + inputs.float(), + self.normalized_shape, + None if self.weight is None else self.weight.float(), + None if self.bias is None else self.bias.float(), + self.eps, + ).to(origin_dtype) + return out + + +class PatchEmbed3D(nn.Module): + def __init__( + self, + patch_size=(2, 4, 4), + in_chans=3, + embed_dim=96, + norm_layer=None, + flatten=True, + ): + super().__init__() + self.patch_size = patch_size + self.flatten = flatten + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + _, _, depth, height, width = x.size() + if width % self.patch_size[2] != 0: + x = F.pad(x, (0, self.patch_size[2] - width % self.patch_size[2])) + if height % self.patch_size[1] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[1] - height % self.patch_size[1])) + if depth % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - depth % self.patch_size[0])) + + x = self.proj(x) + if self.norm is not None: + depth, height, width = x.size(2), x.size(3), x.size(4) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + x = x.transpose(1, 2).view(-1, self.embed_dim, depth, height, width) + if self.flatten: + x = x.flatten(2).transpose(1, 2) + return x + + +def modulate_fp32(norm_func, x, shift, scale): + assert shift.dtype == torch.float32 and scale.dtype == torch.float32 + dtype = x.dtype + x = norm_func(x.to(torch.float32)) + x = x * (scale + 1) + shift + x = x.to(dtype) + return x + + +class FinalLayer_FP32(nn.Module): + def __init__(self, hidden_size, num_patch, out_channels, adaln_tembed_dim): + super().__init__() + self.hidden_size = hidden_size + self.num_patch = num_patch + self.out_channels = out_channels + self.adaln_tembed_dim = adaln_tembed_dim + + self.norm_final = LayerNorm_FP32(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, num_patch * out_channels, bias=True) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(adaln_tembed_dim, 2 * hidden_size, bias=True)) + + def forward(self, x, t, latent_shape): + assert t.dtype == torch.float32 + batch, seq_len, channels = x.shape + num_frames, _, _ = latent_shape + + shift, scale = silu_linear_fp32(self.adaLN_modulation, t).unsqueeze(2).chunk(2, dim=-1) + x = modulate_fp32(self.norm_final, x.view(batch, num_frames, -1, channels), shift, scale).view( + batch, seq_len, channels + ) + x = linear_fp32(self.linear, x) + return x + + +class TimestepEmbedder(nn.Module): + def __init__(self, t_embed_dim, frequency_embedding_size=256): + super().__init__() + self.t_embed_dim = t_embed_dim + self.frequency_embedding_size = frequency_embedding_size + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, t_embed_dim, bias=True), + nn.SiLU(), + nn.Linear(t_embed_dim, t_embed_dim, bias=True), + ) + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half) + freqs = freqs.to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t, dtype): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + if t_freq.dtype != dtype: + t_freq = t_freq.to(dtype) + t_emb = linear_fp32(self.mlp[0], t_freq) + t_emb = F.silu(t_emb) + t_emb = linear_fp32(self.mlp[2], t_emb) + return t_emb + + +class CaptionEmbedder(nn.Module): + def __init__(self, in_channels, hidden_size): + super().__init__() + self.in_channels = in_channels + self.hidden_size = hidden_size + self.y_proj = nn.Sequential( + nn.Linear(in_channels, hidden_size, bias=True), + nn.GELU(approximate="tanh"), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + + def forward(self, caption): + caption = self.y_proj(caption) + return caption + + +def broadcast_concat(tensors, dim=-1): + num_tensors = len(tensors) + shape_lens = set(map(lambda t: len(t.shape), tensors)) + assert len(shape_lens) == 1, "tensors must all have the same number of dimensions" + shape_len = list(shape_lens)[0] + dim = (dim + shape_len) if dim < 0 else dim + dims = list(zip(*map(lambda t: list(t.shape), tensors))) + expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim] + assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]), "invalid broadcast concat dimensions" + max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims)) + expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims)) + expanded_dims.insert(dim, (dim, dims[dim])) + expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims))) + tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes))) + return torch.cat(tensors, dim=dim) + + +def rotate_half(x): + x = rearrange(x, "... (d r) -> ... d r", r=2) + x1, x2 = x.unbind(dim=-1) + x = torch.stack((-x2, x1), dim=-1) + return rearrange(x, "... d r -> ... (d r)") + + +def normalize_and_scale(values, source_range, target_range): + source_min, source_max = source_range + target_min, target_max = target_range + denom = source_max - source_min + if torch.is_tensor(denom): + denom = torch.clamp(denom, min=torch.finfo(values.dtype).eps) + else: + denom = max(denom, torch.finfo(values.dtype).eps) + normalized = (values - source_min) / denom + return normalized * (target_max - target_min) + target_min + + +class RotaryPositionalEmbedding1D(nn.Module): + def __init__(self, head_dim): + super().__init__() + self.head_dim = head_dim + self.base = 10000 + + def precompute_freqs_cis_1d(self, pos_indices): + freqs = 1.0 / (self.base ** (torch.arange(0, self.head_dim, 2)[: (self.head_dim // 2)].float() / self.head_dim)) + freqs = freqs.to(pos_indices.device) + freqs = torch.einsum("..., f -> ... f", pos_indices.float(), freqs) + return repeat(freqs, "... n -> ... (n r)", r=2) + + def forward(self, x, pos_indices): + freqs_cis = self.precompute_freqs_cis_1d(pos_indices) + x_ = x.float() + freqs_cis = freqs_cis.float().to(x.device) + cos, sin = freqs_cis.cos(), freqs_cis.sin() + cos, sin = rearrange(cos, "n d -> 1 1 n d"), rearrange(sin, "n d -> 1 1 n d") + x_ = (x_ * cos) + (rotate_half(x_) * sin) + return x_.type_as(x) + + +def _calculate_x_ref_attn_map( + noise_q: torch.Tensor, + ref_k: torch.Tensor, + ref_target_masks: torch.Tensor, +) -> torch.Tensor: + ref_k = ref_k.to(noise_q.dtype).to(noise_q.device) + noise_q = noise_q * (1.0 / noise_q.shape[-1] ** 0.5) + noise_q = noise_q.transpose(1, 2) + ref_k = ref_k.transpose(1, 2) + attn = noise_q @ ref_k.transpose(-2, -1) + attn_source = attn.softmax(-1).to(noise_q.dtype) + + x_ref_attn_maps = [] + ref_target_masks = ref_target_masks.to(noise_q.dtype) + for ref_target_mask in ref_target_masks: + ref_target_mask = ref_target_mask[None, None, None, ...] + mask_sum = torch.clamp(ref_target_mask.sum(), min=torch.finfo(noise_q.dtype).eps) + x_ref_attn_map = (attn_source * ref_target_mask).sum(-1) / mask_sum + x_ref_attn_map = x_ref_attn_map.permute(0, 2, 1).mean(-1) + x_ref_attn_maps.append(x_ref_attn_map) + return torch.cat(x_ref_attn_maps, dim=0) + + +def _get_attn_map_with_target( + noise_q: torch.Tensor, + key: torch.Tensor, + shape: tuple[int, int, int], + ref_target_masks: torch.Tensor, + split_num: int = 2, + cp_split_hw=None, +) -> torch.Tensor: + if cp_split_hw is not None and cp_split_hw[0] * cp_split_hw[1] > 1: + raise NotImplementedError("LongCat-Video-Avatar multi-speaker masks do not support context parallel.") + _, num_h, num_w = shape + ref_seq_len = num_h * num_w + ref_k = key[:, :ref_seq_len] + if ref_target_masks.shape[-1] != ref_seq_len: + raise ValueError(f"Expected ref_target_masks token length {ref_seq_len}, got {ref_target_masks.shape[-1]}.") + + _, seq_len, heads, _ = noise_q.shape + class_num = ref_target_masks.shape[0] + x_ref_attn_maps = torch.zeros(class_num, seq_len, dtype=noise_q.dtype, device=noise_q.device) + split_num = max(1, min(split_num, heads)) + for noise_q_chunk, ref_k_chunk in zip( + torch.chunk(noise_q.contiguous(), split_num, dim=2), + torch.chunk(ref_k.contiguous(), split_num, dim=2), + ): + x_ref_attn_maps += _calculate_x_ref_attn_map(noise_q_chunk, ref_k_chunk, ref_target_masks) + return x_ref_attn_maps / split_num + + +class RotaryPositionalEmbedding(nn.Module): + _MAX_FREQS_LRU_CACHE_SIZE = 32 + + def __init__(self, head_dim, cp_split_hw=None): + super().__init__() + self.head_dim = head_dim + assert self.head_dim % 8 == 0, "Dim must be a multiply of 8 for 3D RoPE." + self.cp_split_hw = tuple(cp_split_hw or (1, 1)) + self.base = 10000 + self.freqs_dict = OrderedDict() + + def register_grid_size(self, grid_size, key_name, frame_index=None, num_ref_latents=None): + if key_name in self.freqs_dict: + self.freqs_dict.move_to_end(key_name) + return + self.freqs_dict[key_name] = self.precompute_freqs_cis_3d(grid_size, frame_index, num_ref_latents) + if len(self.freqs_dict) > self._MAX_FREQS_LRU_CACHE_SIZE: + self.freqs_dict.popitem(last=False) + + def precompute_freqs_cis_3d(self, grid_size, frame_index=None, num_ref_latents=None): + num_frames, height, width = grid_size + dim_t = self.head_dim - 4 * (self.head_dim // 6) + dim_h = 2 * (self.head_dim // 6) + dim_w = 2 * (self.head_dim // 6) + freqs_t = 1.0 / (self.base ** (torch.arange(0, dim_t, 2)[: (dim_t // 2)].float() / dim_t)) + freqs_h = 1.0 / (self.base ** (torch.arange(0, dim_h, 2)[: (dim_h // 2)].float() / dim_h)) + freqs_w = 1.0 / (self.base ** (torch.arange(0, dim_w, 2)[: (dim_w // 2)].float() / dim_w)) + if frame_index is not None and num_ref_latents is not None: + grid_t = torch.cat( + [ + torch.tensor([frame_index], dtype=torch.float32), + torch.arange(0, num_frames - num_ref_latents, dtype=torch.float32), + ], + dim=0, + ) + else: + grid_t = torch.from_numpy(np.linspace(0, num_frames, num_frames, endpoint=False, dtype=np.float32)).float() + grid_h = torch.from_numpy(np.linspace(0, height, height, endpoint=False, dtype=np.float32)).float() + grid_w = torch.from_numpy(np.linspace(0, width, width, endpoint=False, dtype=np.float32)).float() + freqs_t = torch.einsum("..., f -> ... f", grid_t, freqs_t) + freqs_h = torch.einsum("..., f -> ... f", grid_h, freqs_h) + freqs_w = torch.einsum("..., f -> ... f", grid_w, freqs_w) + freqs_t = repeat(freqs_t, "... n -> ... (n r)", r=2) + freqs_h = repeat(freqs_h, "... n -> ... (n r)", r=2) + freqs_w = repeat(freqs_w, "... n -> ... (n r)", r=2) + freqs = broadcast_concat( + (freqs_t[:, None, None, :], freqs_h[None, :, None, :], freqs_w[None, None, :, :]), dim=-1 + ) + freqs = rearrange(freqs, "T H W D -> (T H W) D") + if self.cp_split_hw[0] * self.cp_split_hw[1] > 1: + raise NotImplementedError("LongCat-Video-Avatar context parallel RoPE is not supported in native MVP.") + return freqs + + def forward(self, q, k, grid_size, frame_index=None, num_ref_latents=None): + key_name = (tuple(grid_size), frame_index, num_ref_latents) + self.register_grid_size(grid_size, key_name, frame_index, num_ref_latents) + + freqs_cis = self.freqs_dict[key_name].to(q.device) + q_, k_ = q.float(), k.float() + freqs_cis = freqs_cis.float().to(q.device) + cos, sin = freqs_cis.cos(), freqs_cis.sin() + cos, sin = rearrange(cos, "n d -> 1 1 n d"), rearrange(sin, "n d -> 1 1 n d") + q_ = (q_ * cos) + (rotate_half(q_) * sin) + k_ = (k_ * cos) + (rotate_half(k_) * sin) + return q_.type_as(q), k_.type_as(k) + + +class Attention(nn.Module): + """LongCat self-attention projections with vLLM-Omni attention execution.""" + + def __init__( + self, + dim: int, + num_heads: int, + enable_flashattn3: bool = False, + enable_flashattn2: bool = False, + enable_xformers: bool = False, + enable_bsa: bool = False, + bsa_params: dict | None = None, + cp_split_hw: tuple[int, int] | list[int] | None = None, + ) -> None: + super().__init__() + if dim % num_heads != 0: + raise ValueError("dim must be divisible by num_heads") + if enable_bsa: + raise NotImplementedError("LongCat-Video-Avatar block-sparse attention is not supported in native MVP.") + if cp_split_hw not in (None, [1, 1], (1, 1)): + raise NotImplementedError("LongCat-Video-Avatar context parallel attention is not supported in native MVP.") + + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim**-0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=True) + self.q_norm = RMSNorm_FP32(self.head_dim, eps=1e-6) + self.k_norm = RMSNorm_FP32(self.head_dim, eps=1e-6) + self.proj = nn.Linear(dim, dim) + + self.rope_3d = RotaryPositionalEmbedding(self.head_dim, cp_split_hw=(1, 1)) + self.attn = VllmAttention( + num_heads=num_heads, + head_size=self.head_dim, + softmax_scale=self.scale, + causal=False, + role="longcat_video_avatar.self", + role_category="self", + qkv_layout="BSND", + ) + + def _process_attn(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + if q.shape[2] == 0: + return q + q = q.transpose(1, 2).contiguous() + k = k.transpose(1, 2).contiguous() + v = v.transpose(1, 2).contiguous() + x = self.attn(q, k, v) + return x.transpose(1, 2).contiguous() + + def forward( + self, + x: torch.Tensor, + shape: tuple[int, int, int] | None = None, + num_cond_latents: int | None = None, + return_kv: bool = False, + num_ref_latents=None, + ref_img_index=None, + mask_frame_range=None, + ref_target_masks=None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]: + bsz, seq_len, channels = x.shape + qkv = self.qkv(x) + qkv = qkv.view(bsz, seq_len, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + q, k = self.q_norm(q), self.k_norm(k) + + if return_kv: + k_cache, v_cache = k.clone(), v.clone() + + q, k = self.rope_3d(q, k, shape, ref_img_index, num_ref_latents) + + num_cond_latents_thw = None + if num_cond_latents is not None and num_cond_latents > 0: + if shape is None: + raise ValueError("shape is required when num_cond_latents is set.") + tokens_per_frame = seq_len // shape[0] + num_cond_latents_thw = num_cond_latents * tokens_per_frame + if num_cond_latents == 1: + q_cond = q[:, :, :num_cond_latents_thw].contiguous() + k_cond = k[:, :, :num_cond_latents_thw].contiguous() + v_cond = v[:, :, :num_cond_latents_thw].contiguous() + x_cond = self._process_attn(q_cond, k_cond, v_cond) + q_noise = q[:, :, num_cond_latents_thw:].contiguous() + x_noise = self._process_attn(q_noise, k, v) + x = torch.cat([x_cond, x_noise], dim=2).contiguous() + else: + if num_ref_latents is None or ref_img_index is None: + raise ValueError("Video continuation requires num_ref_latents and ref_img_index.") + num_ref_latents_thw = num_ref_latents * tokens_per_frame + q_ref = q[:, :, :num_ref_latents_thw].contiguous() + k_ref = k[:, :, :num_ref_latents_thw].contiguous() + v_ref = v[:, :, :num_ref_latents_thw].contiguous() + q_cond = q[:, :, num_ref_latents_thw:num_cond_latents_thw].contiguous() + k_cond = k[:, :, num_ref_latents_thw:num_cond_latents_thw].contiguous() + v_cond = v[:, :, num_ref_latents_thw:num_cond_latents_thw].contiguous() + x_ref = self._process_attn(q_ref, k_ref, v_ref) + x_cond = self._process_attn(q_cond, k_cond, v_cond) + if num_cond_latents == shape[0]: + x = torch.cat([x_ref, x_cond], dim=2).contiguous() + else: + q_noise = q[:, :, num_cond_latents_thw:].contiguous() + start_noise, end_noise, num_noisy_frames = 0, 0, shape[0] - num_cond_latents + if mask_frame_range is not None and mask_frame_range > 0: + start_noise = ref_img_index - mask_frame_range - num_cond_latents + num_ref_latents + end_noise = ref_img_index + mask_frame_range - num_cond_latents + num_ref_latents + 1 + + if start_noise >= 0 and end_noise > start_noise and end_noise <= num_noisy_frames: + start_pos = start_noise * tokens_per_frame + end_pos = end_noise * tokens_per_frame + q_noise_front = q_noise[:, :, :start_pos].contiguous() + q_noise_maskref = q_noise[:, :, start_pos:end_pos].contiguous() + q_noise_back = q_noise[:, :, end_pos:].contiguous() + k_non_ref = k[:, :, num_ref_latents_thw:].contiguous() + v_non_ref = v[:, :, num_ref_latents_thw:].contiguous() + x_noise_front = self._process_attn(q_noise_front, k, v) + x_noise_maskref = self._process_attn(q_noise_maskref, k_non_ref, v_non_ref) + x_noise_back = self._process_attn(q_noise_back, k, v) + x_noise = torch.cat([x_noise_front, x_noise_maskref, x_noise_back], dim=2).contiguous() + else: + x_noise = self._process_attn(q_noise, k, v) + x = torch.cat([x_ref, x_cond, x_noise], dim=2).contiguous() + else: + x = self._process_attn(q, k, v) + + x = x.transpose(1, 2).reshape(bsz, seq_len, channels) + x = self.proj(x) + x_ref_attn_map = None + if ref_target_masks is not None: + if shape is None or num_cond_latents_thw is None or num_cond_latents_thw <= 0: + raise ValueError("Multi-speaker Avatar masks require image/video condition latents.") + x_ref_attn_map = _get_attn_map_with_target( + q.permute(0, 2, 1, 3)[:, num_cond_latents_thw:].type_as(x), + k.permute(0, 2, 1, 3).type_as(x), + shape, + ref_target_masks, + cp_split_hw=self.rope_3d.cp_split_hw, + ) + if return_kv: + return x, (k_cache, v_cache), x_ref_attn_map + return x, x_ref_attn_map + + def forward_with_kv_cache( + self, + x: torch.Tensor, + shape: tuple[int, int, int] | None = None, + num_cond_latents: int | None = None, + kv_cache=None, + num_ref_latents=None, + ref_img_index=None, + mask_frame_range=None, + ref_target_masks=None, + ) -> torch.Tensor: + if shape is None: + raise ValueError("shape is required for LongCat-Video-Avatar KV-cache attention.") + if kv_cache is None: + raise ValueError("kv_cache is required for forward_with_kv_cache.") + + bsz, seq_len, channels = x.shape + qkv = self.qkv(x) + qkv = qkv.view(bsz, seq_len, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + q, k = self.q_norm(q), self.k_norm(k) + + k_cache, v_cache = kv_cache + k_cache = k_cache.to(device=x.device, dtype=k.dtype) + v_cache = v_cache.to(device=x.device, dtype=v.dtype) + if k_cache.shape[0] == 1 and bsz > 1: + k_cache = k_cache.repeat(bsz, 1, 1, 1) + v_cache = v_cache.repeat(bsz, 1, 1, 1) + if k_cache.shape[0] != bsz: + raise ValueError(f"KV cache batch {k_cache.shape[0]} must be 1 or match input batch {bsz}.") + + k_full = torch.cat([k_cache, k], dim=2).contiguous() + v_full = torch.cat([v_cache, v], dim=2).contiguous() + if num_cond_latents is not None and num_cond_latents > 0: + q_padding = torch.cat([torch.empty_like(k_cache), q], dim=2).contiguous() + q_padding, k_full = self.rope_3d( + q_padding, + k_full, + (shape[0] + num_cond_latents, shape[1], shape[2]), + ref_img_index, + num_ref_latents, + ) + q = q_padding[:, :, -seq_len:].contiguous() + + tokens_per_frame = seq_len // shape[0] + start_noise, end_noise, num_noisy_frames = 0, 0, shape[0] + if mask_frame_range is not None and mask_frame_range > 0: + start_noise = ref_img_index - mask_frame_range - num_cond_latents + num_ref_latents + end_noise = ref_img_index + mask_frame_range - num_cond_latents + num_ref_latents + 1 + + if start_noise >= 0 and end_noise > start_noise and end_noise <= num_noisy_frames: + num_ref_latents_thw = (num_ref_latents or 0) * tokens_per_frame + start_pos = start_noise * tokens_per_frame + end_pos = end_noise * tokens_per_frame + q_noise_front = q[:, :, :start_pos].contiguous() + q_noise_maskref = q[:, :, start_pos:end_pos].contiguous() + q_noise_back = q[:, :, end_pos:].contiguous() + k_non_ref = k_full[:, :, num_ref_latents_thw:].contiguous() + v_non_ref = v_full[:, :, num_ref_latents_thw:].contiguous() + x_noise_front = self._process_attn(q_noise_front, k_full, v_full) + x_noise_maskref = self._process_attn(q_noise_maskref, k_non_ref, v_non_ref) + x_noise_back = self._process_attn(q_noise_back, k_full, v_full) + x_attn = torch.cat([x_noise_front, x_noise_maskref, x_noise_back], dim=2).contiguous() + else: + x_attn = self._process_attn(q, k_full, v_full) + + x_out = x_attn.transpose(1, 2).reshape(bsz, seq_len, channels) + x_out = self.proj(x_out) + + x_ref_attn_map = None + if ref_target_masks is not None: + x_ref_attn_map = _get_attn_map_with_target( + q.permute(0, 2, 1, 3).type_as(x_out), + k_full.permute(0, 2, 1, 3).type_as(x_out), + shape, + ref_target_masks, + cp_split_hw=self.rope_3d.cp_split_hw, + ) + return x_out, x_ref_attn_map + + +class MultiHeadCrossAttention(nn.Module): + """LongCat cross-attention projections with vLLM-Omni attention execution.""" + + def __init__( + self, + dim: int, + num_heads: int, + enable_flashattn3: bool = False, + enable_flashattn2: bool = False, + enable_xformers: bool = False, + ): + super().__init__() + if dim % num_heads != 0: + raise ValueError("dim must be divisible by num_heads") + + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim**-0.5 + + self.q_linear = nn.Linear(dim, dim) + self.kv_linear = nn.Linear(dim, dim * 2) + self.proj = nn.Linear(dim, dim) + + self.q_norm = RMSNorm_FP32(self.head_dim, eps=1e-6) + self.k_norm = RMSNorm_FP32(self.head_dim, eps=1e-6) + self.attn = VllmAttention( + num_heads=num_heads, + head_size=self.head_dim, + softmax_scale=self.scale, + causal=False, + role="longcat_video_avatar.cross", + role_category="cross", + qkv_layout="BSND", + skip_sequence_parallel=True, + disable_kv_quant=True, + ) + + def _unpack_condition( + self, + cond: torch.Tensor, + kv_seqlen: list[int], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if cond.shape[0] != 1: + return cond, None + + batch = len(kv_seqlen) + max_len = max(kv_seqlen) + if len(set(kv_seqlen)) == 1: + return cond.view(batch, max_len, self.dim), None + + padded = cond.new_zeros(batch, max_len, self.dim) + mask = torch.zeros(batch, max_len, dtype=torch.bool, device=cond.device) + offset = 0 + flat = cond.squeeze(0) + for idx, seq_len in enumerate(kv_seqlen): + padded[idx, :seq_len] = flat[offset : offset + seq_len] + mask[idx, :seq_len] = True + offset += seq_len + return padded, mask + + def _process_cross_attn( + self, + x: torch.Tensor, + cond: torch.Tensor, + kv_seqlen: list[int], + ) -> torch.Tensor: + bsz, seq_len, channels = x.shape + if channels != self.dim: + raise ValueError(f"Expected x channel dim {self.dim}, got {channels}.") + + cond, attn_mask = self._unpack_condition(cond, kv_seqlen) + q = self.q_linear(x).view(bsz, seq_len, self.num_heads, self.head_dim) + kv = self.kv_linear(cond).view(bsz, cond.shape[1], 2, self.num_heads, self.head_dim) + k, v = kv.unbind(2) + q, k = self.q_norm(q), self.k_norm(k) + + metadata = AttentionMetadata(attn_mask=attn_mask) if attn_mask is not None else None + x = self.attn(q, k, v, metadata) + x = x.reshape(bsz, seq_len, channels) + return self.proj(x) + + def forward( + self, + x: torch.Tensor, + cond: torch.Tensor, + kv_seqlen: list[int], + num_cond_latents: int | None = None, + shape: tuple[int, int, int] | None = None, + ) -> torch.Tensor: + if num_cond_latents is None or num_cond_latents == 0: + return self._process_cross_attn(x, cond, kv_seqlen) + + bsz, seq_len, channels = x.shape + if shape is None: + raise ValueError("shape is required when num_cond_latents is set.") + num_cond_latents_thw = num_cond_latents * (seq_len // shape[0]) + x_noise = x[:, num_cond_latents_thw:] + output_noise = self._process_cross_attn(x_noise, cond, kv_seqlen) + return torch.cat( + [ + torch.zeros( + (bsz, num_cond_latents_thw, channels), + dtype=output_noise.dtype, + device=output_noise.device, + ), + output_noise, + ], + dim=1, + ).contiguous() + + +class AvatarSelfAttention(Attention): + """Avatar self-attention with vLLM-Omni attention execution.""" + + def forward( + self, + x: torch.Tensor, + shape: tuple[int, int, int] | None = None, + num_cond_latents: int | None = None, + return_kv: bool = False, + num_ref_latents=None, + ref_img_index=None, + mask_frame_range=None, + ref_target_masks=None, + ): + out = super().forward( + x, + shape=shape, + num_cond_latents=num_cond_latents, + return_kv=return_kv, + num_ref_latents=num_ref_latents, + ref_img_index=ref_img_index, + mask_frame_range=mask_frame_range, + ref_target_masks=ref_target_masks, + ) + if return_kv: + x_out, kv_cache, x_ref_attn_map = out + return x_out, kv_cache, x_ref_attn_map + return out + + +class SingleStreamAudioAttention(nn.Module): + """Avatar audio cross-attention using vLLM-Omni attention backend.""" + + def __init__( + self, + dim: int, + encoder_hidden_states_dim: int, + num_heads: int, + qkv_bias: bool, + qk_norm: bool, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + eps: float = 1e-6, + class_range: int = 24, + class_interval: int = 4, + cp_split_hw=None, + enable_flashattn3: bool = False, + enable_flashattn2: bool = False, + enable_xformers: bool = False, + ) -> None: + super().__init__() + if dim % num_heads != 0: + raise ValueError("dim must be divisible by num_heads") + if cp_split_hw not in (None, [1, 1], (1, 1)): + raise NotImplementedError("LongCat-Video-Avatar context parallel is not supported in native MVP.") + if attn_drop != 0.0 or proj_drop != 0.0: + raise NotImplementedError("LongCat-Video-Avatar native MVP only supports zero attention/proj dropout.") + + self.dim = dim + self.encoder_hidden_states_dim = encoder_hidden_states_dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim**-0.5 + self.cp_split_hw = cp_split_hw + + self.q_linear = nn.Linear(dim, dim, bias=qkv_bias) + self.q_norm = RMSNorm_FP32(self.head_dim, eps=eps) if qk_norm else nn.Identity() + self.kv_linear = nn.Linear(encoder_hidden_states_dim, dim * 2, bias=qkv_bias) + self.k_norm = RMSNorm_FP32(self.head_dim, eps=eps) if qk_norm else nn.Identity() + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Identity() + self.class_interval = class_interval + self.class_range = class_range + self.rope_h1 = (0, self.class_interval) + self.rope_h2 = (self.class_range - self.class_interval, self.class_range) + self.rope_bak = int(self.class_range // 2) + self.rope_1d = RotaryPositionalEmbedding1D(self.head_dim) + + self.attn = VllmAttention( + num_heads=num_heads, + head_size=self.head_dim, + softmax_scale=self.scale, + causal=False, + role="longcat_video.avatar_audio", + role_category="cross", + qkv_layout="BSND", + skip_sequence_parallel=True, + disable_kv_quant=True, + ) + + def _process_cross_attn( + self, + x: torch.Tensor, + cond: torch.Tensor, + frames_num: int, + x_ref_attn_map=None, + human_num=None, + ) -> torch.Tensor: + out_dtype = x.dtype + x = rearrange(x, "B (N_t S) C -> (B N_t) S C", N_t=frames_num) + batch, seq_len, channels = x.shape + q = self.q_linear(x).view(batch, seq_len, self.num_heads, self.head_dim) + q = self.q_norm(q) + + if x_ref_attn_map is not None: + if x_ref_attn_map.shape[0] < 2: + raise ValueError("Multi-speaker Avatar requires at least two speaker attention maps.") + human1 = normalize_and_scale( + x_ref_attn_map[0], + (x_ref_attn_map[0].min(), x_ref_attn_map[0].max()), + (self.rope_h1[0], self.rope_h1[1]), + ) + human2 = normalize_and_scale( + x_ref_attn_map[1], + (x_ref_attn_map[1].min(), x_ref_attn_map[1].max()), + (self.rope_h2[0], self.rope_h2[1]), + ) + background_pos = self.rope_bak if x_ref_attn_map.shape[0] <= 3 else 100 + background = torch.full( + (x_ref_attn_map.size(1),), + background_pos, + dtype=human1.dtype, + device=human1.device, + ) + max_indices = x_ref_attn_map.argmax(dim=0).clamp(max=2) + normalized_map = torch.stack([human1, human2, background], dim=1) + normalized_pos = normalized_map[torch.arange(x_ref_attn_map.size(1), device=q.device), max_indices] + q = rearrange(q, "(B N_t) S H D -> B H (N_t S) D", N_t=frames_num) + q = self.rope_1d(q, normalized_pos) + q = rearrange(q, "B H (N_t S) D -> (B N_t) S H D", N_t=frames_num) + + kv = self.kv_linear(cond).view(batch, cond.shape[1], 2, self.num_heads, self.head_dim) + k, v = kv.unbind(2) + k = self.k_norm(k) + + if x_ref_attn_map is not None: + per_frame = torch.zeros(cond.shape[1], dtype=k.dtype, device=k.device) + human1_pos = (self.rope_h1[0] + self.rope_h1[1]) / 2 + human2_pos = (self.rope_h2[0] + self.rope_h2[1]) / 2 + if human_num is not None and human_num > 2: + background_pos = self.rope_bak if x_ref_attn_map.shape[0] <= 3 else 100 + tokens_per_human = max(1, per_frame.size(0) // human_num) + per_frame[:tokens_per_human] = human1_pos + per_frame[tokens_per_human : 2 * tokens_per_human] = human2_pos + per_frame[2 * tokens_per_human :] = background_pos + else: + midpoint = per_frame.size(0) // 2 + per_frame[:midpoint] = human1_pos + per_frame[midpoint:] = human2_pos + encoder_pos = torch.cat([per_frame] * frames_num, dim=0) + k = rearrange(k, "(B N_t) S H D -> B H (N_t S) D", N_t=frames_num) + k = self.rope_1d(k, encoder_pos) + k = rearrange(k, "B H (N_t S) D -> (B N_t) S H D", N_t=frames_num) + + x = self.attn(q, k, v) + x = self.proj(x.reshape(batch, seq_len, channels)) + x = self.proj_drop(x) + x = rearrange(x, "(B N_t) S C -> B (N_t S) C", N_t=frames_num) + return x.to(out_dtype) + + def forward( + self, + x: torch.Tensor, + cond: torch.Tensor, + shape: tuple[int, int, int] | None = None, + num_cond_latents: int | None = None, + x_ref_attn_map=None, + human_num=None, + ): + if shape is None: + raise ValueError("shape is required for LongCat-Video-Avatar audio attention.") + batch, seq_len, channels = x.shape + if num_cond_latents is None or num_cond_latents == 0: + return None, self._process_cross_attn(x, cond, shape[0], x_ref_attn_map, human_num) + + num_cond_latents_thw = num_cond_latents * (seq_len // shape[0]) + x_noise = x[:, num_cond_latents_thw:] + cond = rearrange(cond, "(B N_t) M C -> B N_t M C", B=batch) + cond = cond[:, num_cond_latents:] + cond = rearrange(cond, "B N_t M C -> (B N_t) M C") + frames_num = shape[0] - num_cond_latents + output_noise = self._process_cross_attn(x_noise, cond, frames_num, x_ref_attn_map, human_num) + output_cond = torch.zeros( + (batch, num_cond_latents_thw, channels), + dtype=output_noise.dtype, + device=output_noise.device, + ) + return output_cond, output_noise + + +class AudioProjModel(nn.Module): + """Project Whisper/Wav2Vec frame windows into Avatar audio context tokens.""" + + def __init__( + self, + seq_len: int = 5, + seq_len_vf: int = 12, + blocks: int = 12, + channels: int = 768, + intermediate_dim: int = 512, + output_dim: int = 768, + context_tokens: int = 32, + norm_output_audio: bool = True, + enable_compile: bool = False, + ): + super().__init__() + self.seq_len = seq_len + self.blocks = blocks + self.channels = channels + self.input_dim = seq_len * blocks * channels + self.input_dim_vf = seq_len_vf * blocks * channels + self.intermediate_dim = intermediate_dim + self.context_tokens = context_tokens + self.output_dim = output_dim + self.enable_compile = enable_compile + self.flops = 0.0 + + self.proj1 = nn.Linear(self.input_dim, intermediate_dim) + self.proj1_vf = nn.Linear(self.input_dim_vf, intermediate_dim) + self.proj2 = nn.Linear(intermediate_dim, intermediate_dim) + self.proj3 = nn.Linear(intermediate_dim, context_tokens * output_dim) + self.norm = nn.LayerNorm(output_dim) if norm_output_audio else nn.Identity() + + def forward(self, audio_embeds: torch.Tensor, audio_embeds_vf: torch.Tensor) -> torch.Tensor: + video_length = audio_embeds.shape[1] + audio_embeds_vf.shape[1] + batch, _, _, _, _ = audio_embeds.shape + + audio_embeds = rearrange(audio_embeds, "bz f w b c -> (bz f) w b c") + batch_size, window_size, blocks, channels = audio_embeds.shape + audio_embeds = audio_embeds.view(batch_size, window_size * blocks * channels) + + audio_embeds_vf = rearrange(audio_embeds_vf, "bz f w b c -> (bz f) w b c") + batch_size_vf, window_size_vf, blocks_vf, channels_vf = audio_embeds_vf.shape + audio_embeds_vf = audio_embeds_vf.view(batch_size_vf, window_size_vf * blocks_vf * channels_vf) + + audio_embeds = torch.relu(self.proj1(audio_embeds)) + audio_embeds_vf = torch.relu(self.proj1_vf(audio_embeds_vf)) + + audio_embeds = rearrange(audio_embeds, "(bz f) c -> bz f c", bz=batch) + audio_embeds_vf = rearrange(audio_embeds_vf, "(bz f) c -> bz f c", bz=batch) + audio_embeds_c = torch.concat([audio_embeds, audio_embeds_vf], dim=1) + + audio_embeds_c = audio_embeds_c.reshape(audio_embeds_c.shape[0] * audio_embeds_c.shape[1], -1) + audio_embeds_c = torch.relu(self.proj2(audio_embeds_c)) + context_tokens = self.proj3(audio_embeds_c).reshape(-1, self.context_tokens, self.output_dim) + context_tokens = self.norm(context_tokens) + return rearrange(context_tokens, "(bz f) m c -> bz f m c", bz=batch, f=video_length) + + +class QuantizedLinear(nn.Module): + """INT8 weight-only quantized linear layer with per-channel scales.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = False): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.register_buffer("weight_int8", torch.zeros(out_features, in_features, dtype=torch.int8)) + self.register_buffer("weight_scale", torch.zeros(out_features, dtype=torch.float32)) + if bias: + self.register_buffer("bias", torch.zeros(out_features, dtype=torch.bfloat16)) + else: + self.bias = None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + compute_dtype = x.dtype + weight = self.weight_int8.to(compute_dtype) * self.weight_scale.to(compute_dtype).unsqueeze(1) + bias = self.bias.to(compute_dtype) if self.bias is not None else None + return F.linear(x, weight, bias) + + +DEFAULT_SKIP_PATTERNS = {"final_layer.linear"} + + +def _read_config(config_path: Path) -> dict: + with config_path.open("r", encoding="utf-8") as f: + config = json.load(f) + signature = inspect.signature(LongCatVideoAvatarTransformer3DModel.__init__) + allowed_keys = set(signature.parameters.keys()) + return {key: value for key, value in config.items() if key in allowed_keys} + + +def replace_linear_with_quantized(model: nn.Module, skip_patterns: set[str] | None = None) -> None: + skip_patterns = skip_patterns or DEFAULT_SKIP_PATTERNS + modules_to_replace = {} + for name, module in model.named_modules(): + if isinstance(module, nn.Linear) and not any(pattern in name for pattern in skip_patterns): + modules_to_replace[name] = QuantizedLinear( + module.in_features, + module.out_features, + bias=module.bias is not None, + ) + + for name, qlinear in modules_to_replace.items(): + parent = model + parts = name.split(".") + for part in parts[:-1]: + parent = getattr(parent, part) + setattr(parent, parts[-1], qlinear) + + +class LoRAUPParallel(nn.Module): + def __init__(self, blocks): + super().__init__() + self.blocks = nn.ModuleList(blocks) + + def forward(self, x): + if x.shape[-1] % len(self.blocks) != 0: + raise ValueError("LoRAUPParallel input dim must be divisible by number of blocks.") + xs = torch.chunk(x, len(self.blocks), dim=-1) + return torch.cat([block(xs[idx]) for idx, block in enumerate(self.blocks)], dim=-1) + + +class LoRAModule(nn.Module): + def __init__( + self, + lora_name: str, + org_module: nn.Module, + multiplier: float = 1.0, + lora_dim: int = 4, + alpha: float = 1, + n_separate: int = 1, + ): + super().__init__() + self.lora_name = lora_name + if org_module.__class__.__name__ not in ("Linear", "QuantizedLinear"): + raise TypeError(f"LoRA target must be Linear or QuantizedLinear, got {org_module.__class__.__name__}") + + in_dim = org_module.in_features + out_dim = org_module.out_features + self.lora_dim = lora_dim + if n_separate > 1: + if out_dim % n_separate != 0: + raise ValueError("LoRA separate blocks must evenly divide output dimension.") + self.lora_down = nn.Linear(in_dim, n_separate * self.lora_dim, bias=False) + self.lora_up = LoRAUPParallel( + [nn.Linear(self.lora_dim, out_dim // n_separate, bias=False) for _ in range(n_separate)] + ) + else: + self.lora_down = nn.Linear(in_dim, self.lora_dim, bias=False) + self.lora_up = nn.Linear(self.lora_dim, out_dim, bias=False) + + if isinstance(alpha, torch.Tensor): + alpha = float(alpha.detach().float().cpu().numpy()) + alpha = self.lora_dim if alpha is None or alpha == 0 else alpha + self.register_buffer("alpha_scale", torch.tensor(alpha / self.lora_dim)) + nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5)) + if n_separate > 1: + for block in self.lora_up.blocks: + nn.init.zeros_(block.weight) + else: + nn.init.zeros_(self.lora_up.weight) + self.multiplier = multiplier + self.use_lora = True + + +class LoRANetwork(nn.Module): + def __init__( + self, + model: nn.Module, + lora_state_dict: dict[str, torch.Tensor], + multiplier: float = 1.0, + lora_dim: int = 128, + alpha: float = 64, + ): + super().__init__() + self.multiplier = multiplier + self.lora_dim = lora_dim + self.alpha = alpha + self.loras: list[LoRAModule] = [] + + lora_names = {key.split(".lora_down.weight")[0] for key in lora_state_dict if key.endswith("lora_down.weight")} + for lora_name in sorted(lora_names): + module_name = lora_name.replace("lora___lorahyphen___", "").replace("___lorahyphen___", ".") + try: + module = model + for part in module_name.split("."): + module = getattr(module, part) + except AttributeError: + continue + if module.__class__.__name__ not in ("Linear", "QuantizedLinear"): + continue + prefix = lora_name + ".lora_up.blocks" + n_blocks = sum(1 for key in lora_state_dict if key.startswith(prefix)) + lora = LoRAModule( + lora_name, + module, + multiplier=multiplier, + lora_dim=lora_dim, + alpha=alpha, + n_separate=n_blocks if n_blocks > 0 else 1, + ) + self.loras.append(lora) + self.add_module(lora_name, lora) + + +def create_lora_network( + transformer: nn.Module, + lora_state_dict: dict[str, torch.Tensor], + multiplier: float, + network_dim: int | None, + network_alpha: float | None, +) -> LoRANetwork: + network = LoRANetwork( + transformer, + lora_state_dict, + multiplier=multiplier, + lora_dim=network_dim or 128, + alpha=network_alpha or 64, + ) + network.load_state_dict(lora_state_dict, strict=True) + return network + + +class LongCatAvatarSingleStreamBlock(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_ratio: int, + adaln_tembed_dim: int, + enable_flashattn3: bool = False, + enable_flashattn2: bool = False, + enable_xformers: bool = False, + enable_bsa: bool = False, + bsa_params=None, + cp_split_hw=None, + output_dim: int = 768, + audio_prenorm: bool = True, + class_range: int = 24, + class_interval: int = 4, + ): + super().__init__() + self.hidden_size = hidden_size + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(adaln_tembed_dim, 6 * hidden_size, bias=True)) + self.audio_adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(adaln_tembed_dim, 3 * hidden_size, bias=True), + ) + + self.mod_norm_attn = LayerNorm_FP32(hidden_size, eps=1e-6, elementwise_affine=False) + self.mod_norm_ffn = LayerNorm_FP32(hidden_size, eps=1e-6, elementwise_affine=False) + self.pre_crs_attn_norm = LayerNorm_FP32(hidden_size, eps=1e-6, elementwise_affine=True) + self.pre_video_crs_attn_norm = LayerNorm_FP32(hidden_size, eps=1e-6, elementwise_affine=True) + self.pre_audio_crs_attn_norm = ( + LayerNorm_FP32(output_dim, eps=1e-6, elementwise_affine=True) if audio_prenorm else nn.Identity() + ) + + self.attn = AvatarSelfAttention( + dim=hidden_size, + num_heads=num_heads, + enable_flashattn3=enable_flashattn3, + enable_flashattn2=enable_flashattn2, + enable_xformers=enable_xformers, + enable_bsa=enable_bsa, + bsa_params=bsa_params, + cp_split_hw=cp_split_hw, + ) + self.cross_attn = MultiHeadCrossAttention( + dim=hidden_size, + num_heads=num_heads, + enable_flashattn3=enable_flashattn3, + enable_flashattn2=enable_flashattn2, + enable_xformers=enable_xformers, + ) + self.audio_cross_attn = SingleStreamAudioAttention( + dim=hidden_size, + encoder_hidden_states_dim=output_dim, + num_heads=num_heads, + qk_norm=True, + qkv_bias=True, + class_range=class_range, + class_interval=class_interval, + cp_split_hw=cp_split_hw, + enable_flashattn3=enable_flashattn3, + enable_flashattn2=enable_flashattn2, + enable_xformers=enable_xformers, + ) + self.ffn = FeedForwardSwiGLU(dim=hidden_size, hidden_dim=int(hidden_size * mlp_ratio)) + + def forward( + self, + x, + y, + t, + y_seqlen, + latent_shape, + num_cond_latents=None, + return_kv=False, + kv_cache=None, + skip_crs_attn=False, + num_ref_latents=None, + audio_hidden_states=None, + ref_img_index=None, + mask_frame_range=None, + token_ref_target_masks=None, + human_num=None, + ): + x_dtype = x.dtype + batch, seq_len, channels = x.shape + num_frames, _, _ = latent_shape + num_cond = int(num_cond_latents or 0) + + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + silu_linear_fp32(self.adaLN_modulation, t).unsqueeze(2).chunk(6, dim=-1) + ) + x_m = modulate_fp32(self.mod_norm_attn, x.view(batch, num_frames, -1, channels), shift_msa, scale_msa) + x_m = x_m.view(batch, seq_len, channels) + if kv_cache is not None: + attn_outputs = self.attn.forward_with_kv_cache( + x_m, + shape=latent_shape, + num_cond_latents=num_cond_latents, + kv_cache=kv_cache, + num_ref_latents=num_ref_latents, + ref_img_index=ref_img_index, + mask_frame_range=mask_frame_range, + ref_target_masks=token_ref_target_masks, + ) + else: + attn_outputs = self.attn( + x_m, + shape=latent_shape, + num_cond_latents=num_cond_latents, + return_kv=return_kv, + num_ref_latents=num_ref_latents, + ref_img_index=ref_img_index, + mask_frame_range=mask_frame_range, + ref_target_masks=token_ref_target_masks, + ) + if return_kv: + x_s, new_kv_cache, x_ref_attn_map = attn_outputs + else: + x_s, x_ref_attn_map = attn_outputs + x = x + (gate_msa * x_s.view(batch, -1, seq_len // num_frames, channels)).view(batch, -1, channels) + x = x.to(x_dtype) + + if not skip_crs_attn: + cross_num_cond_latents = None if kv_cache is not None else num_cond_latents + x = x + self.cross_attn( + self.pre_crs_attn_norm(x), + y, + y_seqlen, + num_cond_latents=cross_num_cond_latents, + shape=latent_shape, + ) + + audio_num_cond = 0 if kv_cache is not None else num_cond + audio_shift_mca, audio_scale_mca, audio_gate_mca = ( + silu_linear_fp32(self.audio_adaLN_modulation, t[:, audio_num_cond:]).unsqueeze(2).chunk(3, dim=-1) + ) + audio_output_cond, audio_output_noise = self.audio_cross_attn( + self.pre_video_crs_attn_norm(x), + self.pre_audio_crs_attn_norm(audio_hidden_states), + shape=latent_shape, + num_cond_latents=audio_num_cond, + x_ref_attn_map=x_ref_attn_map, + human_num=human_num, + ) + audio_output_noise = modulate_fp32( + self.mod_norm_attn, + audio_output_noise.view(batch, num_frames - audio_num_cond, -1, channels), + audio_shift_mca, + audio_scale_mca, + ).view(batch, -1, channels) + audio_add_x = ( + audio_gate_mca * audio_output_noise.view(batch, num_frames - audio_num_cond, -1, channels) + ).view(batch, -1, channels) + if audio_output_cond is not None: + audio_add_x = torch.cat([audio_output_cond, audio_add_x], dim=1).contiguous() + x = (x + audio_add_x).to(x_dtype) + + x_m = modulate_fp32(self.mod_norm_ffn, x.view(batch, -1, seq_len // num_frames, channels), shift_mlp, scale_mlp) + x_m = x_m.view(batch, -1, channels) + x_s = self.ffn(x_m) + x = x + (gate_mlp * x_s.view(batch, -1, seq_len // num_frames, channels)).view(batch, -1, channels) + if return_kv: + return x.to(x_dtype), new_kv_cache + return x.to(x_dtype) + + +class LongCatVideoAvatarTransformer3DModel(nn.Module): + _supports_gradient_checkpointing = True + + def __init__( + self, + in_channels: int = 16, + out_channels: int = 16, + hidden_size: int = 4096, + depth: int = 48, + num_heads: int = 32, + caption_channels: int = 4096, + mlp_ratio: int = 4, + adaln_tembed_dim: int = 512, + frequency_embedding_size: int = 256, + patch_size: tuple[int, int, int] | list[int] = (1, 2, 2), + enable_flashattn3: bool = False, + enable_flashattn2: bool = False, + enable_xformers: bool = False, + enable_bsa: bool = False, + bsa_params: dict | None = None, + cp_split_hw: tuple[int, int] | list[int] | None = None, + text_tokens_zero_pad: bool = False, + audio_window: int = 5, + audio_block: int = 12, + audio_channel: int = 768, + intermediate_dim: int = 512, + output_dim: int = 768, + context_tokens: int = 32, + vae_scale: int = 4, + audio_prenorm: bool = False, + class_range: int = 24, + class_interval: int = 4, + ) -> None: + super().__init__() + if enable_bsa: + raise NotImplementedError("LongCat-Video-Avatar BSA is not supported in native MVP.") + if cp_split_hw not in (None, [1, 1], (1, 1)): + raise NotImplementedError("LongCat-Video-Avatar context parallel is not supported in native MVP.") + + self.patch_size = tuple(patch_size) + self.in_channels = in_channels + self.out_channels = out_channels + self.cp_split_hw = (1, 1) + self.vae_scale = vae_scale + self.audio_window = audio_window + self.text_tokens_zero_pad = text_tokens_zero_pad + self.gradient_checkpointing = False + self.lora_dict = {} + self.active_loras = [] + self.config = SimpleNamespace( + in_channels=in_channels, + out_channels=out_channels, + hidden_size=hidden_size, + depth=depth, + num_heads=num_heads, + caption_channels=caption_channels, + mlp_ratio=mlp_ratio, + adaln_tembed_dim=adaln_tembed_dim, + frequency_embedding_size=frequency_embedding_size, + patch_size=self.patch_size, + enable_flashattn3=enable_flashattn3, + enable_flashattn2=enable_flashattn2, + enable_xformers=enable_xformers, + enable_bsa=enable_bsa, + bsa_params=bsa_params, + cp_split_hw=None, + text_tokens_zero_pad=text_tokens_zero_pad, + audio_window=audio_window, + audio_block=audio_block, + audio_channel=audio_channel, + intermediate_dim=intermediate_dim, + output_dim=output_dim, + context_tokens=context_tokens, + vae_scale=vae_scale, + audio_prenorm=audio_prenorm, + class_range=class_range, + class_interval=class_interval, + ) + + self.x_embedder = PatchEmbed3D(self.patch_size, in_channels, hidden_size) + self.t_embedder = TimestepEmbedder( + t_embed_dim=adaln_tembed_dim, + frequency_embedding_size=frequency_embedding_size, + ) + self.y_embedder = CaptionEmbedder(in_channels=caption_channels, hidden_size=hidden_size) + self.blocks = nn.ModuleList( + [ + LongCatAvatarSingleStreamBlock( + hidden_size=hidden_size, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + adaln_tembed_dim=adaln_tembed_dim, + enable_flashattn3=enable_flashattn3, + enable_flashattn2=enable_flashattn2, + enable_xformers=enable_xformers, + enable_bsa=enable_bsa, + bsa_params=bsa_params, + cp_split_hw=self.cp_split_hw, + output_dim=output_dim, + audio_prenorm=audio_prenorm, + class_range=class_range, + class_interval=class_interval, + ) + for _ in range(depth) + ] + ) + self.audio_proj = AudioProjModel( + seq_len=audio_window, + seq_len_vf=audio_window + vae_scale - 1, + blocks=audio_block, + channels=audio_channel, + intermediate_dim=intermediate_dim, + output_dim=output_dim, + context_tokens=context_tokens, + ) + self.final_layer = FinalLayer_FP32(hidden_size, np.prod(self.patch_size), out_channels, adaln_tembed_dim) + + @property + def dtype(self) -> torch.dtype: + return self.x_embedder.proj.weight.dtype + + def load_lora(self, lora_path, lora_key, multiplier=1.0, lora_network_dim=128, lora_network_alpha=64): + lora_state_dict = load_file(lora_path, device="cpu") + self.lora_dict[lora_key] = create_lora_network( + self, + lora_state_dict, + multiplier=multiplier, + network_dim=lora_network_dim, + network_alpha=lora_network_alpha, + ) + + def enable_loras(self, lora_key_list=None): + lora_key_list = lora_key_list or [] + self.disable_all_loras() + module_loras = {} + model_device = next(self.parameters()).device + model_dtype = next(self.parameters()).dtype + for lora_key in lora_key_list: + if lora_key not in self.lora_dict: + continue + for lora in self.lora_dict[lora_key].loras: + lora.to(model_device, dtype=model_dtype, non_blocking=True) + module_name = lora.lora_name.replace("lora___lorahyphen___", "").replace("___lorahyphen___", ".") + module_loras.setdefault(module_name, []).append(lora) + self.active_loras.append(lora_key) + + # This mutates module.forward and is expected to run only during + # single-threaded model initialization, before requests are served. + for module_name, loras in module_loras.items(): + module = self._get_module_by_name(module_name) + if not hasattr(module, "org_forward"): + module.org_forward = module.forward + module.forward = self._create_multi_lora_forward(module, loras) + + def _create_multi_lora_forward(self, module, loras): + def multi_lora_forward(x, *args, **kwargs): + weight_dtype = x.dtype + org_output = module.org_forward(x, *args, **kwargs) + total_lora_output = 0 + for lora in loras: + if lora.use_lora: + lx = lora.lora_down(x.to(lora.lora_down.weight.dtype)) + lx = lora.lora_up(lx) + total_lora_output = total_lora_output + lx.to(weight_dtype) * lora.multiplier * lora.alpha_scale + return org_output + total_lora_output + + return multi_lora_forward + + def _get_module_by_name(self, module_name): + module = self + for part in module_name.split("."): + module = getattr(module, part) + return module + + def disable_all_loras(self): + for module in self.modules(): + if hasattr(module, "org_forward"): + module.forward = module.org_forward + delattr(module, "org_forward") + for lora_network in self.lora_dict.values(): + for lora in lora_network.loras: + lora.to("cpu") + self.active_loras.clear() + + def enable_bsa(self): + raise NotImplementedError("LongCat-Video-Avatar BSA is not supported in native MVP.") + + def disable_bsa(self): + return None + + def forward( + self, + hidden_states, + timestep, + encoder_hidden_states, + encoder_attention_mask=None, + num_cond_latents=0, + return_kv=False, + kv_cache_dict=None, + skip_crs_attn=False, + offload_kv_cache=False, + audio_embs=None, + num_ref_latents=None, + ref_img_index=None, + mask_frame_range=None, + ref_target_masks=None, + ): + if audio_embs is None: + raise ValueError("LongCat-Video-Avatar requires audio_embs.") + kv_cache_dict = kv_cache_dict or {} + + batch, _, num_frames, height, width = hidden_states.shape + n_t = num_frames // self.patch_size[0] + n_h = height // self.patch_size[1] + n_w = width // self.patch_size[2] + if self.patch_size[0] != 1: + raise NotImplementedError("LongCat-Video-Avatar native MVP requires temporal patch size 1.") + if len(timestep.shape) == 1: + timestep = timestep.unsqueeze(1).expand(-1, n_t) + + dtype = self.x_embedder.proj.weight.dtype + hidden_states = self.x_embedder(hidden_states.to(dtype)) + timestep = timestep.to(dtype) + encoder_hidden_states = encoder_hidden_states.to(dtype) + t = self.t_embedder(timestep.float().flatten(), dtype=torch.float32).reshape(batch, n_t, -1) + encoder_hidden_states = self.y_embedder(encoder_hidden_states) + + audio_cond = audio_embs.to(device=hidden_states.device, dtype=hidden_states.dtype) + first_frame_audio_emb_s = audio_cond[:, :1, ...] + latter_frame_audio_emb = audio_cond[:, 1:, ...] + latter_frame_audio_emb = rearrange(latter_frame_audio_emb, "b (n_t n) w s c -> b n_t n w s c", n=self.vae_scale) + middle_index = self.audio_window // 2 + latter_first_frame_audio_emb = latter_frame_audio_emb[:, :, :1, : middle_index + 1, ...] + latter_first_frame_audio_emb = rearrange(latter_first_frame_audio_emb, "b n_t n w s c -> b n_t (n w) s c") + latter_last_frame_audio_emb = latter_frame_audio_emb[:, :, -1:, middle_index:, ...] + latter_last_frame_audio_emb = rearrange(latter_last_frame_audio_emb, "b n_t n w s c -> b n_t (n w) s c") + latter_middle_frame_audio_emb = latter_frame_audio_emb[:, :, 1:-1, middle_index : middle_index + 1, ...] + latter_middle_frame_audio_emb = rearrange(latter_middle_frame_audio_emb, "b n_t n w s c -> b n_t (n w) s c") + latter_frame_audio_emb_s = torch.concat( + [latter_first_frame_audio_emb, latter_middle_frame_audio_emb, latter_last_frame_audio_emb], + dim=2, + ) + audio_hidden_states = self.audio_proj(first_frame_audio_emb_s, latter_frame_audio_emb_s) + if num_ref_latents is not None and num_ref_latents > 0: + audio_start_ref = audio_hidden_states[:, [0], :, :] + audio_hidden_states = torch.cat([audio_start_ref, audio_hidden_states], dim=1).contiguous() + audio_hidden_states = audio_hidden_states[:, -n_t:] + + human_num = None + if ref_target_masks is not None: + if batch != 1: + raise NotImplementedError("LongCat-Video-Avatar multi-speaker AI2V supports batch size 1.") + human_num = len(audio_hidden_states) + audio_hidden_states = torch.concat(audio_hidden_states.split(1), dim=2).squeeze(0) + else: + audio_hidden_states = rearrange(audio_hidden_states, "b t n c -> (b t) n c") + + token_ref_target_masks = None + if ref_target_masks is not None: + ref_target_masks = ref_target_masks.unsqueeze(0).to(device=hidden_states.device, dtype=torch.float32) + token_ref_target_masks = F.interpolate(ref_target_masks, size=(n_h, n_w), mode="nearest") + token_ref_target_masks = token_ref_target_masks.squeeze(0) + token_ref_target_masks = (token_ref_target_masks > 0).view(token_ref_target_masks.shape[0], -1) + token_ref_target_masks = token_ref_target_masks.to(dtype) + + if self.text_tokens_zero_pad and encoder_attention_mask is not None: + encoder_hidden_states = encoder_hidden_states * encoder_attention_mask[:, None, :, None] + encoder_attention_mask = (encoder_attention_mask * 0 + 1).to(encoder_attention_mask.dtype) + + if encoder_attention_mask is not None: + encoder_attention_mask = encoder_attention_mask.squeeze(1).squeeze(1) + encoder_hidden_states = ( + encoder_hidden_states.squeeze(1) + .masked_select(encoder_attention_mask.unsqueeze(-1) != 0) + .view(1, -1, hidden_states.shape[-1]) + ) + y_seqlens = encoder_attention_mask.sum(dim=1).tolist() + else: + y_seqlens = [encoder_hidden_states.shape[2]] * encoder_hidden_states.shape[0] + encoder_hidden_states = encoder_hidden_states.squeeze(1).view(1, -1, hidden_states.shape[-1]) + + kv_cache_dict_ret = {} + for idx, block in enumerate(self.blocks): + block_outputs = block( + hidden_states, + encoder_hidden_states, + t, + y_seqlens, + (n_t, n_h, n_w), + num_cond_latents, + return_kv, + kv_cache_dict.get(idx), + skip_crs_attn, + num_ref_latents, + audio_hidden_states, + ref_img_index, + mask_frame_range, + token_ref_target_masks, + human_num, + ) + if return_kv: + hidden_states, kv_cache = block_outputs + if offload_kv_cache: + kv_cache_dict_ret[idx] = (kv_cache[0].cpu(), kv_cache[1].cpu()) + else: + kv_cache_dict_ret[idx] = (kv_cache[0].contiguous(), kv_cache[1].contiguous()) + else: + hidden_states = block_outputs + + hidden_states = self.final_layer(hidden_states, t, (n_t, n_h, n_w)) + hidden_states = self.unpatchify(hidden_states, n_t, n_h, n_w) + hidden_states = hidden_states.to(torch.float32) + if return_kv: + return hidden_states, kv_cache_dict_ret + return hidden_states + + def unpatchify(self, x, n_t, n_h, n_w): + t_p, h_p, w_p = self.patch_size + return rearrange( + x, + "B (N_t N_h N_w) (T_p H_p W_p C_out) -> B C_out (N_t T_p) (N_h H_p) (N_w W_p)", + N_t=n_t, + N_h=n_h, + N_w=n_w, + T_p=t_p, + H_p=h_p, + W_p=w_p, + C_out=self.out_channels, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + params_dict = dict(self.named_parameters()) + params_dict.update(dict(self.named_buffers())) + + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +def create_quantized_avatar_dit(checkpoint_dir: str | os.PathLike[str], subfolder: str = "base_model_int8", **kwargs): + quantized_dir = Path(checkpoint_dir) / subfolder + config = _read_config(quantized_dir / "config.json") + config.update(kwargs) + model = LongCatVideoAvatarTransformer3DModel(**config) + replace_linear_with_quantized(model) + + for module in model.modules(): + if isinstance(module, QuantizedLinear): + continue + for param in module.parameters(recurse=False): + if param.dtype == torch.float32: + param.data = param.data.to(torch.bfloat16) + return model.eval() + + +def create_full_precision_avatar_dit(checkpoint_dir: str | os.PathLike[str], subfolder: str = "base_model", **kwargs): + model_dir = Path(checkpoint_dir) / subfolder + config = _read_config(model_dir / "config.json") + config.update(kwargs) + model = LongCatVideoAvatarTransformer3DModel(**config) + return model.eval() diff --git a/vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py b/vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py new file mode 100644 index 000000000..33b0e2a45 --- /dev/null +++ b/vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py @@ -0,0 +1,1929 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import html +import json +import math +import os +import re +import shutil +import tempfile +import time +from collections.abc import Iterable +from contextlib import nullcontext +from pathlib import Path +from typing import Any, ClassVar + +import numpy as np +import torch +import torch.nn.functional as F +from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler +from diffusers.video_processor import VideoProcessor +from PIL import Image +from torch import nn +from transformers import AutoFeatureExtractor, AutoTokenizer, UMT5EncoderModel, WhisperModel +from vllm.logger import init_logger +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan +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 SupportAudioInput, SupportImageInput +from vllm_omni.diffusion.models.longcat_video.longcat_video_avatar_transformer import ( + create_full_precision_avatar_dit, + create_quantized_avatar_dit, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniTextPrompt +from vllm_omni.platforms import current_omni_platform + +logger = init_logger(__name__) + +_DEFAULT_NEGATIVE_PROMPT = ( + "Close-up, Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, " + "static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, " + "poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, " + "messy background, three legs, many people in the background, walking backwards" +) +_DEFAULT_AT2V_RATIO = 480 / 832 + +# Whisper uses a 30-second context window: 30s * 16kHz raw samples, +# represented as 3000 mel frames after feature extraction. +_WHISPER_AUDIO_SAMPLE_RATE = 16000 +_WHISPER_AUDIO_CHUNK_SAMPLES = 30 * _WHISPER_AUDIO_SAMPLE_RATE +_WHISPER_ENCODER_CHUNK_FRAMES = 3000 + +# Mirrors official LongCat-Video `longcat_video/utils/bukcet_config.py` +# for scale_factor_spatial in {16, 32}. Used for Avatar condition sizing +# and AT2V defaults without importing the official repository at runtime. +_LONGCAT_480P_BUCKETS = { + "0.26": ([320, 1216], 1), + "0.31": ([352, 1120], 1), + "0.38": ([384, 1024], 1), + "0.43": ([416, 960], 1), + "0.52": ([448, 864], 1), + "0.58": ([480, 832], 1), + "0.67": ([512, 768], 1), + "0.74": ([544, 736], 1), + "0.86": ([576, 672], 1), + "0.95": ([608, 640], 1), + "1.05": ([640, 608], 1), + "1.17": ([672, 576], 1), + "1.29": ([704, 544], 1), + "1.35": ([736, 544], 1), + "1.50": ([768, 512], 1), + "1.67": ([800, 480], 1), + "1.73": ([832, 480], 1), + "2.00": ([896, 448], 1), + "2.31": ([960, 416], 1), + "2.58": ([992, 384], 1), + "2.75": ([1056, 384], 1), + "3.09": ([1088, 352], 1), + "3.70": ([1184, 320], 1), + "3.80": ([1216, 320], 1), + "3.90": ([1248, 320], 1), + "4.00": ([1280, 320], 1), +} + +_LONGCAT_720P_BUCKETS = { + "0.25": ([480, 1920], 1), + "0.29": ([512, 1792], 1), + "0.32": ([544, 1696], 1), + "0.36": ([576, 1600], 1), + "0.40": ([608, 1504], 1), + "0.49": ([672, 1376], 1), + "0.54": ([704, 1312], 1), + "0.59": ([736, 1248], 1), + "0.69": ([800, 1152], 1), + "0.74": ([832, 1120], 1), + "0.82": ([864, 1056], 1), + "0.88": ([896, 1024], 1), + "0.94": ([928, 992], 1), + "1.00": ([960, 960], 1), + "1.07": ([992, 928], 1), + "1.14": ([1024, 896], 1), + "1.22": ([1056, 864], 1), + "1.31": ([1088, 832], 1), + "1.35": ([1120, 832], 1), + "1.44": ([1152, 800], 1), + "1.70": ([1248, 736], 1), + "2.00": ([1344, 672], 1), + "2.05": ([1376, 672], 1), + "2.47": ([1504, 608], 1), + "2.53": ([1536, 608], 1), + "2.83": ([1632, 576], 1), + "3.06": ([1664, 544], 1), + "3.12": ([1696, 544], 1), + "3.62": ([1856, 512], 1), + "3.93": ([1888, 480], 1), + "4.00": ([1920, 480], 1), +} + + +def _get_longcat_bucket_config(resolution: str) -> dict[str, tuple[list[int], int]]: + if resolution == "480p": + return _LONGCAT_480P_BUCKETS + if resolution == "720p": + return _LONGCAT_720P_BUCKETS + raise ValueError(f"Unsupported LongCat-Video-Avatar resolution {resolution!r}. Expected '480p' or '720p'.") + + +def _default_at2v_shape(resolution: str) -> tuple[int, int]: + bucket_config = _get_longcat_bucket_config(resolution) + closest_bucket = sorted(bucket_config.keys(), key=lambda key: abs(float(key) - _DEFAULT_AT2V_RATIO))[0] + target_h, target_w = bucket_config[closest_bucket][0] + return target_h, target_w + + +def _avatar_model_allow_patterns(use_int8: bool) -> list[str]: + weight_subfolder = "base_model_int8" if use_int8 else "base_model" + return [ + "scheduler/*", + f"{weight_subfolder}/*", + "lora/dmd_lora.safetensors", + "whisper-large-v3/config.json", + "whisper-large-v3/generation_config.json", + "whisper-large-v3/model.safetensors", + "whisper-large-v3/preprocessor_config.json", + "vocal_separator/*", + "config.json", + "model_index.json", + ] + + +def _adjust_num_frames(num_frames: int, temporal_scale: int = 4) -> int: + if num_frames % temporal_scale != 1: + num_frames = num_frames // temporal_scale * temporal_scale + 1 + return max(num_frames, 1) + + +def _resolve_num_segments( + value: Any, + *, + audio_duration: float, + num_frames: int, + num_cond_frames: int, + save_fps: int, +) -> int: + if value is None: + return 1 + if isinstance(value, str) and value.lower() == "auto": + first_segment_duration = num_frames / save_fps + stride_duration = (num_frames - num_cond_frames) / save_fps + if audio_duration <= first_segment_duration or stride_duration <= 0: + return 1 + return 1 + math.ceil((audio_duration - first_segment_duration) / stride_duration) + return max(1, int(value)) + + +def _requested_num_segments(value: Any) -> int: + if value is None: + return 1 + if isinstance(value, str) and value.lower() == "auto": + return 0 + return max(1, int(value)) + + +def _requested_audio_duration( + *, + auto_num_segments: bool, + num_segments: int, + num_frames: int, + num_cond_frames: int, + save_fps: int, +) -> float: + if auto_num_segments: + return 0.0 + return num_frames / save_fps + (num_segments - 1) * (num_frames - num_cond_frames) / save_fps + + +def _load_json(path: Path) -> dict[str, Any]: + with path.open("r", encoding="utf-8") as f: + return json.load(f) + + +def _infer_asset_root_from_path(path: Path | None) -> Path | None: + if path is None: + return None + resolved = path.resolve() if path.exists() else path + parts = resolved.parts + if "assets" not in parts: + return resolved.parent + asset_idx = parts.index("assets") + if asset_idx == 0: + return None + return Path(*parts[:asset_idx]) + + +def _resolve_path_from_roots( + path: str | os.PathLike[str] | None, + *, + asset_roots: tuple[Path, ...], +) -> str | None: + if path is None: + return None + resolved = Path(path) + if resolved.is_absolute(): + return str(resolved) + for asset_root in asset_roots: + candidate = asset_root / resolved + if candidate.exists(): + return str(candidate) + return str(resolved) + + +def _speaker_sort_key(name: str) -> tuple[int, str]: + match = re.search(r"(\d+)$", str(name)) + return (int(match.group(1)) if match else 10_000, str(name)) + + +def _normalize_audio_paths( + audio_field: Any, + *, + asset_roots: tuple[Path, ...], +) -> dict[str, str]: + if audio_field is None: + return {} + if isinstance(audio_field, dict): + return { + str(speaker): resolved + for speaker, path in sorted(audio_field.items(), key=lambda item: _speaker_sort_key(str(item[0]))) + if (resolved := _resolve_path_from_roots(path, asset_roots=asset_roots)) is not None + } + if isinstance(audio_field, list | tuple): + audio_paths: dict[str, str] = {} + for idx, item in enumerate(audio_field, start=1): + if isinstance(item, dict): + audio_paths.update(_normalize_audio_paths(item, asset_roots=asset_roots)) + continue + resolved = _resolve_path_from_roots(item, asset_roots=asset_roots) + if resolved is not None: + audio_paths[f"person{idx}"] = resolved + return dict(sorted(audio_paths.items(), key=lambda item: _speaker_sort_key(item[0]))) + if isinstance(audio_field, str | os.PathLike): + resolved = _resolve_path_from_roots(audio_field, asset_roots=asset_roots) + return {"person1": resolved} if resolved is not None else {} + raise TypeError(f"Unsupported LongCat-Video-Avatar audio input type {type(audio_field)}.") + + +def _prepare_multi_speaker_audio_arrays( + speech_arrays: list[np.ndarray], + *, + audio_type: str, + generate_duration: float, + sample_rate: int, +) -> list[np.ndarray]: + if not speech_arrays: + raise ValueError("Multi-speaker Avatar requires at least one audio track.") + + arrays = [np.asarray(array, dtype=np.float32) for array in speech_arrays] + min_samples = max(1, math.ceil(generate_duration * sample_rate)) + audio_type = audio_type.lower() + + if audio_type == "para": + target_len = max(min_samples, *(len(array) for array in arrays)) + return [np.pad(array, (0, max(0, target_len - len(array)))) for array in arrays] + + if audio_type == "add": + source_total = sum(len(array) for array in arrays) + target_len = max(min_samples, source_total) + outputs = [] + offset = 0 + for array in arrays: + output = np.zeros(target_len, dtype=np.float32) + output[offset : offset + len(array)] = array + outputs.append(output) + offset += len(array) + return outputs + + raise NotImplementedError(f"Unsupported LongCat-Video-Avatar multi-speaker audio_type {audio_type!r}.") + + +def _bbox_to_mask(bbox: list[int] | tuple[int, int, int, int], height: int, width: int) -> torch.Tensor: + if len(bbox) != 4: + raise ValueError(f"Expected bbox [y_min, x_min, y_max, x_max], got {bbox!r}.") + y_min, x_min, y_max, x_max = [int(value) for value in bbox] + y_min = max(0, min(height, y_min)) + y_max = max(0, min(height, y_max)) + x_min = max(0, min(width, x_min)) + x_max = max(0, min(width, x_max)) + mask = torch.zeros((height, width), dtype=torch.float32) + mask[y_min:y_max, x_min:x_max] = 1 + return mask + + +def _normalize_other_bboxes(raw_bboxes: Any) -> list[list[int]]: + if raw_bboxes is None: + return [] + if ( + isinstance(raw_bboxes, list | tuple) + and raw_bboxes + and all(isinstance(value, int | float) for value in raw_bboxes) + ): + if len(raw_bboxes) % 4 != 0: + raise ValueError("bbox['others'] must contain a multiple of four values.") + return [list(raw_bboxes[idx : idx + 4]) for idx in range(0, len(raw_bboxes), 4)] + if isinstance(raw_bboxes, list | tuple): + return [list(bbox) for bbox in raw_bboxes] + raise TypeError("bbox['others'] must be a flat bbox list or a list of bboxes.") + + +def _build_multi_speaker_ref_target_masks( + image: Image.Image, + bbox: dict[str, Any] | None, +) -> tuple[torch.Tensor, bool]: + src_width, src_height = image.size + bbox = bbox or {} + person1_bbox = bbox.get("person1") + person2_bbox = bbox.get("person2") + + if person1_bbox is None and person2_bbox is None: + face_scale = 0.1 + split_width = src_width // 2 + y_min = int(src_height * face_scale) + y_max = int(src_height * (1 - face_scale)) + person1_bbox = [ + y_min, + int(split_width * face_scale), + y_max, + int(split_width * (1 - face_scale)), + ] + person2_bbox = [ + y_min, + int(split_width * face_scale + split_width), + y_max, + int(split_width * (1 - face_scale) + split_width), + ] + elif person1_bbox is None or person2_bbox is None: + raise NotImplementedError("Multi-speaker Avatar requires both person1 and person2 bboxes when bbox is set.") + + human_mask1 = _bbox_to_mask(person1_bbox, src_height, src_width) + human_mask2 = _bbox_to_mask(person2_bbox, src_height, src_width) + background_mask = torch.where(human_mask1 + human_mask2 > 0, torch.zeros(()), torch.ones(())) + total_masks = [human_mask1, human_mask2, background_mask] + + other_bboxes = _normalize_other_bboxes(bbox.get("others")) + for other_bbox in other_bboxes: + total_masks.append(_bbox_to_mask(other_bbox, src_height, src_width)) + + return torch.stack(total_masks, dim=0), bool(other_bboxes) + + +def _ensure_local_dir(model: str | os.PathLike[str], allow_patterns: list[str] | None = None) -> Path: + model_path = Path(model) + if model_path.exists(): + return model_path + from huggingface_hub import snapshot_download + + return Path(snapshot_download(str(model), allow_patterns=allow_patterns)) + + +def _ensure_avatar_model_index(model_dir: Path, use_int8: bool) -> None: + model_index_path = model_dir / "model_index.json" + model_index = {} + if model_index_path.exists(): + with model_index_path.open("r", encoding="utf-8") as f: + model_index = json.load(f) + model_index.update( + { + "_class_name": "LongCatVideoAvatarPipeline", + "_diffusers_version": "0.38.0", + "model_name": "LongCat-Video-Avatar-1.5", + } + ) + model_index_path.write_text(json.dumps(model_index, indent=2, sort_keys=True), encoding="utf-8") + + source_subfolder = "base_model_int8" if use_int8 else "base_model" + source_config = model_dir / source_subfolder / "config.json" + target_config = model_dir / "transformer" / "config.json" + if source_config.exists(): + target_config.parent.mkdir(parents=True, exist_ok=True) + target_config.write_text(source_config.read_text(encoding="utf-8"), encoding="utf-8") + + +def prepare_longcat_video_avatar_model_for_omni(model: str, use_int8: bool = True) -> str: + model_dir = _ensure_local_dir(model, allow_patterns=_avatar_model_allow_patterns(use_int8)) + _ensure_avatar_model_index(model_dir, use_int8) + return str(model_dir) + + +def _load_image(raw_image: Any) -> Image.Image: + if isinstance(raw_image, list): + if not raw_image: + raise ValueError("LongCat-Video-Avatar received an empty image list.") + raw_image = raw_image[0] + if isinstance(raw_image, Image.Image): + return raw_image.convert("RGB") + if isinstance(raw_image, np.ndarray): + return Image.fromarray(raw_image).convert("RGB") + if isinstance(raw_image, torch.Tensor): + tensor = raw_image.detach().cpu() + if tensor.dim() == 4 and tensor.shape[0] == 1: + tensor = tensor[0] + if tensor.dim() == 3 and tensor.shape[0] in (1, 3, 4): + tensor = tensor.permute(1, 2, 0) + array = tensor.numpy() + if np.issubdtype(array.dtype, np.floating): + array = (np.clip(array, 0.0, 1.0) * 255).round().astype("uint8") + return Image.fromarray(array).convert("RGB") + if isinstance(raw_image, str | os.PathLike): + return Image.open(raw_image).convert("RGB") + raise TypeError( + f"Unsupported LongCat-Video-Avatar image input type {type(raw_image)}. " + "Pass a PIL image, numpy array, torch tensor, or file path." + ) + + +def _as_bool(value: Any, default: bool) -> bool: + if value is None: + return default + if isinstance(value, bool): + return value + if isinstance(value, str): + return value.strip().lower() not in {"0", "false", "no", "off"} + return bool(value) + + +def _video_tensor_to_pil_frames(video: Any) -> list[Image.Image]: + if isinstance(video, torch.Tensor): + array = video.detach().cpu().numpy() + else: + array = np.asarray(video) + if array.ndim == 5 and array.shape[0] == 1: + array = array[0] + if array.ndim == 4 and array.shape[1] in (1, 3, 4) and array.shape[-1] not in (1, 3, 4): + array = np.transpose(array, (0, 2, 3, 1)) + if array.ndim != 4: + raise ValueError(f"Expected video frames with shape [T,H,W,C], got {array.shape}.") + if np.issubdtype(array.dtype, np.floating): + array = (np.clip(array, 0.0, 1.0) * 255).round().astype("uint8") + else: + array = np.clip(array, 0, 255).astype("uint8") + return [Image.fromarray(frame).convert("RGB") for frame in array] + + +def _retrieve_latents( + encoder_output: torch.Tensor, + generator: torch.Generator | None = None, + sample_mode: str = "sample", +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + if hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + if hasattr(encoder_output, "latents"): + return encoder_output.latents + raise AttributeError("Could not access latents of provided encoder_output") + + +def get_longcat_video_avatar_post_process_func(od_config: OmniDiffusionConfig): + def post_process_func(video: Any, sampling_params=None): + fps = 25 + if sampling_params is not None: + fps = int( + sampling_params.extra_args.get("save_fps") + or sampling_params.extra_args.get("fps") + or sampling_params.fps + or 25 + ) + return { + "video": [_video_tensor_to_pil_frames(video)], + "custom_output": {"fps": fps}, + "fps": fps, + } + + return post_process_func + + +def get_longcat_video_avatar_pre_process_func(od_config: OmniDiffusionConfig): + def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: + prompt = request.prompt + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt) + prompt.setdefault("multi_modal_data", {}) + prompt.setdefault("additional_information", {}) + request.prompt = prompt + return request + + return pre_process_func + + +class LongCatVideoAvatarPipeline(nn.Module, SupportImageInput, SupportAudioInput): + """Native LongCat-Video-Avatar A2V/AI2V pipeline.""" + + support_image_input: ClassVar[bool] = True + support_audio_input: ClassVar[bool] = True + dummy_run_num_frames: ClassVar[int] = 1 + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__() + self.od_config = od_config + self.device = get_local_device() + additional_config = getattr(od_config, "additional_config", {}) or {} + self.model_type = str(additional_config.get("model_type") or "avatar-v1.5") + if self.model_type != "avatar-v1.5": + raise NotImplementedError("LongCat-Video-Avatar native MVP only supports avatar-v1.5.") + self.use_distill = _as_bool(additional_config.get("use_distill"), True) + self.use_int8 = _as_bool(additional_config.get("use_int8"), True) + self.build_components_on_gpu = _as_bool(additional_config.get("build_components_on_gpu"), False) + self.resolution = str(additional_config.get("resolution") or "480p") + self.model_dir = _ensure_local_dir( + od_config.model, + allow_patterns=_avatar_model_allow_patterns(self.use_int8), + ) + self.save_fps = 25 + self.audio_stride = 1 + self.default_num_frames = 93 + self.video_processor = VideoProcessor(vae_scale_factor=8) + dit_subfolder = "base_model_int8" if self.use_int8 else "base_model" + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=str(self.model_dir), + subfolder=dit_subfolder, + revision=None, + prefix="transformer.", + fall_back_to_pt=False, + ) + ] + + self._load_components() + + def _base_model_dir(self) -> Path: + configured = (getattr(self.od_config, "additional_config", {}) or {}).get("base_model_dir") + if configured: + return _ensure_local_dir( + configured, + allow_patterns=["tokenizer/*", "text_encoder/*", "vae/*", "config.json", "model_index.json"], + ) + sibling = self.model_dir.parent / "LongCat-Video" + if sibling.exists(): + return sibling + return _ensure_local_dir( + "meituan-longcat/LongCat-Video", + allow_patterns=["tokenizer/*", "text_encoder/*", "vae/*", "config.json", "model_index.json"], + ) + + def _load_components(self) -> None: + try: + from audio_separator.separator import Separator + except ImportError as exc: + raise ImportError( + "LongCat-Video-Avatar requires audio-separator. " + "Install the optional extra with `pip install 'vllm-omni[longcat-video-avatar]'`." + ) from exc + + dtype = getattr(self.od_config, "dtype", torch.bfloat16) + base_dir = self._base_model_dir() + self.tokenizer = AutoTokenizer.from_pretrained(str(base_dir), subfolder="tokenizer") + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(str(self.model_dir), subfolder="scheduler") + + # vLLM initializes native pipelines under the target device context. + # Build LongCat's large components on CPU first by default to avoid + # transient full-precision CUDA allocations before int8 DiT replacement. + build_context = nullcontext() if self.build_components_on_gpu else torch.device("cpu") + enable_distill_lora = False + with build_context: + self.text_encoder = UMT5EncoderModel.from_pretrained( + str(base_dir), subfolder="text_encoder", torch_dtype=dtype + ) + self.vae = DistributedAutoencoderKLWan.from_pretrained(str(base_dir), subfolder="vae", torch_dtype=dtype) + + if self.use_int8: + self.transformer = create_quantized_avatar_dit( + str(self.model_dir), + subfolder="base_model_int8", + cp_split_hw=[1, 1], + ) + else: + self.transformer = create_full_precision_avatar_dit( + str(self.model_dir), + subfolder="base_model", + cp_split_hw=[1, 1], + ) + self.transformer = self.transformer.to(dtype=dtype) + + if self.use_distill: + lora_path = self.model_dir / "lora" / "dmd_lora.safetensors" + if lora_path.exists(): + self.transformer.load_lora( + str(lora_path), + "dmd", + multiplier=1.0, + lora_network_dim=128, + lora_network_alpha=64, + ) + enable_distill_lora = True + + audio_model_path = self.model_dir / "whisper-large-v3" + self.audio_encoder = WhisperModel.from_pretrained(str(audio_model_path)).eval() + + self.text_encoder = self.text_encoder.to(self.device) + self.vae = self.vae.to(self.device) + self.transformer = self.transformer.to(device=self.device) + if enable_distill_lora: + self.transformer.enable_loras(["dmd"]) + self.audio_encoder = self.audio_encoder.to(self.device) + self.audio_encoder.requires_grad_(False) + self.audio_feature_extractor = AutoFeatureExtractor.from_pretrained(str(audio_model_path)) + + separator_path = self.model_dir / "vocal_separator" / "Kim_Vocal_2.onnx" + self._audio_temp_dir = tempfile.TemporaryDirectory(prefix="longcat_avatar_audio_") + self.audio_temp_dir = Path(self._audio_temp_dir.name) + self.vocal_separator = Separator( + output_dir=self.audio_temp_dir / "vocals", + output_single_stem="vocals", + model_file_dir=str(separator_path.parent), + ) + self.vocal_separator.load_model(separator_path.name) + + def to(self, *args, **kwargs): + super().to(*args, **kwargs) + # .to() is overloaded (device / dtype / tensor); only treat args[0] as + # a device when it actually is one, otherwise keep self.device. + candidate = kwargs.get("device", args[0] if args else None) + device = torch.device(candidate) if isinstance(candidate, str | torch.device | int) else self.device + self.device = device + for name in ("text_encoder", "vae", "transformer", "audio_encoder"): + module = getattr(self, name, None) + if module is not None: + module.to(device) + return self + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load DiT weights using vLLM's diffusion weight loader.""" + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) + + @staticmethod + def _prompt_clean(text: str) -> str: + return re.sub(r"\s+", " ", html.unescape(html.unescape(text))).strip() + + def _get_t5_prompt_embeds( + self, + prompt: str | list[str], + num_videos_per_prompt: int, + max_sequence_length: int, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt = [self._prompt_clean(p) for p in prompt] + batch_size = len(prompt) + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids.to(device) + mask = text_inputs.attention_mask.to(device) + prompt_embeds = self.text_encoder(text_input_ids, mask).last_hidden_state + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, 1, seq_len, -1) + mask = mask.repeat_interleave(num_videos_per_prompt, dim=0) + return prompt_embeds, mask + + def encode_prompt( + self, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + max_sequence_length: int = 512, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + device = device or self.device + dtype = dtype or self.transformer.dtype + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + prompt_embeds, prompt_attention_mask = self._get_t5_prompt_embeds( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + if not do_classifier_free_guidance: + return prompt_embeds, prompt_attention_mask, None, None + + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_embeds, negative_prompt_attention_mask = self._get_t5_prompt_embeds( + prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + @property + def do_classifier_free_guidance(self): + return self._text_guidance_scale > 1.0 or self._audio_guidance_scale > 1.0 + + @property + def num_timesteps(self): + return 1000 + + @property + def num_distill_sample_steps(self): + return 8 + + def get_timesteps_sigmas(self, sampling_steps: int, use_distill: bool = False) -> torch.Tensor: + if use_distill: + distill_indices = torch.arange(1, self.num_distill_sample_steps + 1, dtype=torch.float32) + distill_indices = (distill_indices * (self.num_timesteps // self.num_distill_sample_steps)).round().long() + distill_indices = self.num_timesteps - distill_indices + sigmas = torch.flip(torch.linspace(0, 1, self.num_timesteps), [0]) + sigmas = torch.flip(sigmas[distill_indices], [0]).float() + else: + sigmas = torch.linspace(1, 0.001, sampling_steps) + return sigmas.to(torch.float32) + + def _update_kv_cache_dict(self, kv_cache_dict): + self.kv_cache_dict = kv_cache_dict + + def _cache_clean_latents( + self, + cond_latents: torch.Tensor, + max_sequence_length: int, + *, + offload_kv_cache: bool, + dtype: torch.dtype, + audio_embs: torch.Tensor, + num_cond_latents: int, + num_ref_latents: int, + ref_img_index: int, + ) -> None: + timestep = torch.zeros(cond_latents.shape[0], cond_latents.shape[2], device=self.device, dtype=dtype) + empty_embeds = torch.zeros( + [cond_latents.shape[0], 1, max_sequence_length, self.text_encoder.config.d_model], + device=self.device, + dtype=dtype, + ) + _, kv_cache_dict = self.transformer( + hidden_states=cond_latents, + timestep=timestep, + encoder_hidden_states=empty_embeds, + num_cond_latents=num_cond_latents, + return_kv=True, + skip_crs_attn=True, + offload_kv_cache=offload_kv_cache, + audio_embs=audio_embs, + num_ref_latents=num_ref_latents, + ref_img_index=ref_img_index, + ) + self._update_kv_cache_dict(kv_cache_dict) + + def _get_kv_cache_dict(self): + return getattr(self, "kv_cache_dict", {}) + + def _clear_cache(self) -> None: + self.kv_cache_dict = {} + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + + def normalize_latents(self, latents): + latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, 1) + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1) + latents_std = latents_std.to(latents.device, latents.dtype) + return (latents - latents_mean) * latents_std + + def denormalize_latents(self, latents): + latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, 1) + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1) + latents_std = latents_std.to(latents.device, latents.dtype) + return latents / latents_std + latents_mean + + def prepare_latents( + self, + image: torch.Tensor | None, + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + num_frames: int, + num_cond_frames: int, + dtype: torch.dtype, + device: torch.device, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + video: torch.Tensor | None = None, + need_encode: bool = True, + ) -> torch.Tensor: + if image is not None and video is not None: + raise ValueError("Cannot provide both image and video conditioning.") + if latents is None: + num_latent_frames = (num_frames - 1) // self.vae.config.scale_factor_temporal + 1 + shape = ( + batch_size, + num_channels_latents, + num_latent_frames, + int(height) // self.vae.config.scale_factor_spatial, + int(width) // self.vae.config.scale_factor_spatial, + ) + latents = torch.randn(shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device=device, dtype=dtype) + + if image is not None or video is not None: + condition_data = image if image is not None else video + cond_latents = [] + num_cond_latents = 1 + (num_cond_frames - 1) // self.vae.config.scale_factor_temporal + if need_encode: + for idx in range(batch_size): + gen = generator[idx] if isinstance(generator, list) else generator + if image is not None: + encoded_input = condition_data[idx].unsqueeze(0).unsqueeze(2) + else: + encoded_input = condition_data[idx][:, -num_cond_frames:].unsqueeze(0) + latent = _retrieve_latents(self.vae.encode(encoded_input), gen, sample_mode="argmax") + cond_latents.append(latent) + cond_latents = torch.cat(cond_latents, dim=0).to(dtype) + cond_latents = self.normalize_latents(cond_latents) + else: + cond_latents = condition_data[:, :, -num_cond_latents:].to(device=device, dtype=dtype) + latents[:, :, :num_cond_latents] = cond_latents + return latents + + def _resolve_request_inputs(self, req: OmniDiffusionRequest) -> dict[str, Any]: + first_prompt = req.prompts[0] if req.prompts else {"prompt": ""} + if isinstance(first_prompt, str): + prompt_text = first_prompt + mm_data: dict[str, Any] = {} + info: dict[str, Any] = {} + negative_prompt = None + else: + prompt_text = first_prompt.get("prompt") or "" + mm_data = first_prompt.get("multi_modal_data") or {} + info = first_prompt.get("additional_information") or {} + negative_prompt = first_prompt.get("negative_prompt") + + extra = req.sampling_params.extra_args or {} + input_json = extra.get("input_json") or info.get("input_json") + input_asset_root = None + if input_json: + input_json_path = Path(input_json).expanduser() + sample = _load_json(input_json_path) + input_asset_root = _infer_asset_root_from_path(input_json_path) + else: + sample = {} + asset_roots = (input_asset_root,) if input_asset_root is not None else () + prompt = prompt_text or sample.get("prompt") or "" + negative_prompt = negative_prompt or extra.get("negative_prompt") or _DEFAULT_NEGATIVE_PROMPT + + audio_field = mm_data.get("audio") + if audio_field is None: + audio_field = extra.get("audio_path") or info.get("audio_path") or sample.get("cond_audio") + audio_paths = _normalize_audio_paths(audio_field, asset_roots=asset_roots) + if not audio_paths: + raise ValueError("LongCat-Video-Avatar requires multi_modal_data['audio'] or extra_args['audio_path'].") + audio_path = audio_paths[sorted(audio_paths, key=_speaker_sort_key)[0]] + + image_input = mm_data.get("image") + image_path = extra.get("image_path") or info.get("image_path") or sample.get("cond_image") + if image_input is None and image_path is not None: + image_input = _resolve_path_from_roots(image_path, asset_roots=asset_roots) + + inferred_stage = "ai2v" if image_input is not None else "at2v" + stage = str(extra.get("stage") or extra.get("stage_1") or info.get("stage") or inferred_stage).lower() + resolution = str(extra.get("resolution") or info.get("resolution") or self.resolution) + bbox = extra.get("bbox") or info.get("bbox") or sample.get("bbox") or {} + if isinstance(bbox, str): + bbox = json.loads(bbox) + return { + "prompt": prompt, + "negative_prompt": negative_prompt, + "audio_path": audio_path, + "audio_paths": audio_paths, + "audio_type": str(extra.get("audio_type") or info.get("audio_type") or sample.get("audio_type") or "para"), + "bbox": bbox, + "image": image_input, + "is_multi_speaker": len(audio_paths) > 1, + "stage": stage, + "resolution": resolution, + } + + def _extract_vocal(self, audio_path: str) -> str: + fd, path = tempfile.mkstemp(prefix="longcat_avatar_vocal_", suffix=".wav") + os.close(fd) + target_path = Path(path) + outputs = self.vocal_separator.separate(audio_path) + if not outputs: + target_path.unlink(missing_ok=True) + raise RuntimeError("Audio separation produced no vocal stem.") + default_vocal_path = (self.audio_temp_dir / "vocals" / outputs[0]).resolve() + shutil.move(str(default_vocal_path), str(target_path)) + return str(target_path) + + def _loudness_norm(self, audio_array, sr=16000, lufs=-23, threshold=100): + try: + import pyloudnorm as pyln + except ImportError as exc: + raise ImportError( + "LongCat-Video-Avatar requires pyloudnorm. " + "Install the optional extra with `pip install 'vllm-omni[longcat-video-avatar]'`." + ) from exc + + meter = pyln.Meter(sr) + loudness = meter.integrated_loudness(audio_array) + if abs(loudness) > threshold: + return audio_array + return pyln.normalize.loudness(audio_array, loudness, lufs) + + @torch.no_grad() + def _get_audio_embedding_whisper(self, speech_array, fps=25, sample_rate=16000): + enc_fps = 50 + audio_duration = len(speech_array) / sample_rate + video_length = int(audio_duration * fps) + speech_array = self._loudness_norm(speech_array, sample_rate) + + mel_chunks = [] + for idx in range(0, len(speech_array), _WHISPER_AUDIO_CHUNK_SAMPLES): + mel = self.audio_feature_extractor( + speech_array[idx : idx + _WHISPER_AUDIO_CHUNK_SAMPLES], + sampling_rate=sample_rate, + return_tensors="pt", + ).input_features + mel_chunks.append(mel) + audio_features = torch.cat(mel_chunks, dim=-1).to(self.audio_encoder.dtype) + + enc_chunks = [] + for idx in range(0, audio_features.shape[-1], _WHISPER_ENCODER_CHUNK_FRAMES): + chunk_hs = self.audio_encoder.encoder( + audio_features[:, :, idx : idx + _WHISPER_ENCODER_CHUNK_FRAMES].to(self.device), + output_hidden_states=True, + ).hidden_states + enc_chunks.append(torch.stack(chunk_hs, dim=2)) + audio_prompts = torch.cat(enc_chunks, dim=1)[:, : video_length * 2] + + def interpolate(features, input_fps, output_fps, output_len): + features = features.transpose(1, 2) + output_features = F.interpolate(features, size=output_len, align_corners=True, mode="linear") + return output_features.transpose(1, 2) + + feat0 = interpolate(audio_prompts[:, :, 0:8].mean(dim=2), enc_fps, fps, video_length) + feat1 = interpolate(audio_prompts[:, :, 8:16].mean(dim=2), enc_fps, fps, video_length) + feat2 = interpolate(audio_prompts[:, :, 16:24].mean(dim=2), enc_fps, fps, video_length) + feat3 = interpolate(audio_prompts[:, :, 24:32].mean(dim=2), enc_fps, fps, video_length) + feat4 = interpolate(audio_prompts[:, :, 32], enc_fps, fps, video_length) + return torch.stack([feat0, feat1, feat2, feat3, feat4], dim=2)[0] + + def _audio_embedding_from_array( + self, + speech_array: np.ndarray, + *, + sample_rate: int, + num_frames: int, + save_fps: int, + ) -> torch.Tensor: + full_audio_emb = self._full_audio_embedding_from_array( + speech_array, + sample_rate=sample_rate, + min_duration=num_frames / save_fps, + save_fps=save_fps, + ) + return self._slice_full_audio_embedding(full_audio_emb, audio_start_idx=0, num_frames=num_frames) + + def _full_audio_embedding_from_array( + self, + speech_array: np.ndarray, + *, + sample_rate: int, + min_duration: float, + save_fps: int, + ) -> torch.Tensor: + source_duration = len(speech_array) / sample_rate + added_sample_nums = math.ceil((min_duration - source_duration) * sample_rate) + if added_sample_nums > 0: + speech_array = np.append(speech_array, np.zeros(added_sample_nums, dtype=np.float32)) + + full_audio_emb = self._get_audio_embedding_whisper( + speech_array, + fps=save_fps * self.audio_stride, + sample_rate=sample_rate, + ) + if torch.isnan(full_audio_emb).any(): + raise ValueError("Audio embedding contains NaN values.") + return full_audio_emb.to(self.device) + + def _slice_full_audio_embedding( + self, + full_audio_emb: torch.Tensor, + *, + audio_start_idx: int, + num_frames: int, + ) -> torch.Tensor: + indices = torch.arange(5, device=full_audio_emb.device) - 2 + audio_end_idx = audio_start_idx + self.audio_stride * num_frames + center_indices = torch.arange(audio_start_idx, audio_end_idx, self.audio_stride, device=full_audio_emb.device) + center_indices = center_indices.unsqueeze(1) + indices.unsqueeze(0) + center_indices = torch.clamp(center_indices, min=0, max=full_audio_emb.shape[0] - 1) + return full_audio_emb[center_indices][None, ...].to(self.device) + + def _audio_embedding(self, audio_path: str, num_frames: int, save_fps: int) -> torch.Tensor: + full_audio_emb, _ = self._audio_full_embedding( + audio_path, + min_duration=num_frames / save_fps, + save_fps=save_fps, + ) + return self._slice_full_audio_embedding(full_audio_emb, audio_start_idx=0, num_frames=num_frames) + + def _audio_full_embedding( + self, audio_path: str, *, min_duration: float, save_fps: int + ) -> tuple[torch.Tensor, float]: + import soundfile as sf + + temp_vocal_path = self._extract_vocal(audio_path) + try: + speech_array, sample_rate = self._load_audio(temp_vocal_path, target_sr=16000, soundfile_module=sf) + effective_duration = max(len(speech_array) / sample_rate, min_duration) + full_audio_emb = self._full_audio_embedding_from_array( + speech_array, + sample_rate=sample_rate, + min_duration=effective_duration, + save_fps=save_fps, + ) + return full_audio_emb, effective_duration + finally: + Path(temp_vocal_path).unlink(missing_ok=True) + + def _multi_audio_embedding( + self, + audio_paths: dict[str, str], + *, + audio_type: str, + num_frames: int, + save_fps: int, + include_background_silent_audio: bool, + ) -> torch.Tensor: + import soundfile as sf + + temp_vocal_paths: list[str] = [] + try: + speech_arrays = [] + sample_rate = 16000 + for _, audio_path in sorted(audio_paths.items(), key=lambda item: _speaker_sort_key(item[0])): + temp_vocal_path = self._extract_vocal(audio_path) + temp_vocal_paths.append(temp_vocal_path) + speech_array, sample_rate = self._load_audio(temp_vocal_path, target_sr=16000, soundfile_module=sf) + speech_arrays.append(speech_array) + + aligned_arrays = _prepare_multi_speaker_audio_arrays( + speech_arrays, + audio_type=audio_type, + generate_duration=num_frames / save_fps, + sample_rate=sample_rate, + ) + if include_background_silent_audio: + aligned_arrays.append(np.zeros_like(aligned_arrays[0], dtype=np.float32)) + + audio_embs = [ + self._audio_embedding_from_array( + speech_array, + sample_rate=sample_rate, + num_frames=num_frames, + save_fps=save_fps, + ) + for speech_array in aligned_arrays + ] + return torch.cat(audio_embs, dim=0) + finally: + for temp_vocal_path in temp_vocal_paths: + Path(temp_vocal_path).unlink(missing_ok=True) + + def _multi_audio_full_embeddings( + self, + audio_paths: dict[str, str], + *, + audio_type: str, + min_duration: float, + save_fps: int, + include_background_silent_audio: bool, + ) -> tuple[list[torch.Tensor], float]: + import soundfile as sf + + temp_vocal_paths: list[str] = [] + try: + speech_arrays = [] + sample_rate = 16000 + for _, audio_path in sorted(audio_paths.items(), key=lambda item: _speaker_sort_key(item[0])): + temp_vocal_path = self._extract_vocal(audio_path) + temp_vocal_paths.append(temp_vocal_path) + speech_array, sample_rate = self._load_audio(temp_vocal_path, target_sr=16000, soundfile_module=sf) + speech_arrays.append(speech_array) + + aligned_arrays = _prepare_multi_speaker_audio_arrays( + speech_arrays, + audio_type=audio_type, + generate_duration=min_duration, + sample_rate=sample_rate, + ) + if include_background_silent_audio: + aligned_arrays.append(np.zeros_like(aligned_arrays[0], dtype=np.float32)) + + effective_duration = len(aligned_arrays[0]) / sample_rate + full_embs = [ + self._full_audio_embedding_from_array( + speech_array, + sample_rate=sample_rate, + min_duration=effective_duration, + save_fps=save_fps, + ) + for speech_array in aligned_arrays + ] + return full_embs, effective_duration + finally: + for temp_vocal_path in temp_vocal_paths: + Path(temp_vocal_path).unlink(missing_ok=True) + + def _slice_multi_audio_embeddings( + self, + full_audio_embs: list[torch.Tensor], + *, + audio_start_idx: int, + num_frames: int, + ) -> torch.Tensor: + return torch.cat( + [ + self._slice_full_audio_embedding( + full_audio_emb, + audio_start_idx=audio_start_idx, + num_frames=num_frames, + ) + for full_audio_emb in full_audio_embs + ], + dim=0, + ) + + def _prepare_audio_embeddings( + self, + inputs: dict[str, Any], + *, + is_multi_speaker: bool, + include_background_silent_audio: bool, + needs_full_audio: bool, + auto_num_segments: bool, + num_segments_value: Any, + num_segments: int, + num_frames: int, + num_cond_frames: int, + save_fps: int, + ) -> tuple[torch.Tensor, list[torch.Tensor] | None, int]: + if not needs_full_audio: + if is_multi_speaker: + audio_emb = self._multi_audio_embedding( + inputs["audio_paths"], + audio_type=inputs["audio_type"], + num_frames=num_frames, + save_fps=save_fps, + include_background_silent_audio=include_background_silent_audio, + ) + else: + audio_emb = self._audio_embedding(inputs["audio_path"], num_frames, save_fps) + return audio_emb, None, num_segments + + requested_duration = _requested_audio_duration( + auto_num_segments=auto_num_segments, + num_segments=num_segments, + num_frames=num_frames, + num_cond_frames=num_cond_frames, + save_fps=save_fps, + ) + if is_multi_speaker: + full_audio_embs, effective_audio_duration = self._multi_audio_full_embeddings( + inputs["audio_paths"], + audio_type=inputs["audio_type"], + save_fps=save_fps, + min_duration=requested_duration, + include_background_silent_audio=include_background_silent_audio, + ) + else: + full_audio_emb, effective_audio_duration = self._audio_full_embedding( + inputs["audio_path"], + min_duration=requested_duration, + save_fps=save_fps, + ) + full_audio_embs = [full_audio_emb] + + num_segments = _resolve_num_segments( + num_segments_value, + audio_duration=effective_audio_duration, + num_frames=num_frames, + num_cond_frames=num_cond_frames, + save_fps=save_fps, + ) + audio_emb = self._slice_multi_audio_embeddings( + full_audio_embs, + audio_start_idx=0, + num_frames=num_frames, + ) + return audio_emb, full_audio_embs, num_segments + + @staticmethod + def _load_audio(audio_path: str, target_sr: int, soundfile_module) -> tuple[np.ndarray, int]: + speech_array, sr = soundfile_module.read(audio_path, dtype="float32", always_2d=True) + speech = torch.from_numpy(speech_array.T).float() + if speech.shape[0] > 1: + speech = speech.mean(dim=0, keepdim=True) + if sr != target_sr: + try: + import torchaudio.functional as taF + + speech = taF.resample(speech, sr, target_sr) + except ImportError: + target_len = max(1, int(round(speech.shape[-1] * float(target_sr) / float(sr)))) + speech = F.interpolate(speech.unsqueeze(0), size=target_len, mode="linear", align_corners=False)[0] + sr = target_sr + return speech.squeeze(0).cpu().numpy(), sr + + def _condition_shape(self, image: Image.Image, resolution: str) -> tuple[int, int]: + bucket_config = _get_longcat_bucket_config(resolution) + ratio = image.height / image.width + closest_bucket = sorted(bucket_config.keys(), key=lambda key: abs(float(key) - ratio))[0] + target_h, target_w = bucket_config[closest_bucket][0] + return target_h, target_w + + def _preprocess_image(self, image: Image.Image, height: int, width: int, resize_mode: str = "crop") -> torch.Tensor: + try: + return self.video_processor.preprocess(image, height=height, width=width, resize_mode=resize_mode) + except TypeError: + return self.video_processor.preprocess(image, height=height, width=width) + + def _preprocess_video_frames( + self, + video: list[Image.Image], + height: int, + width: int, + resize_mode: str = "crop", + ) -> torch.Tensor: + if hasattr(self.video_processor, "preprocess_video"): + try: + return self.video_processor.preprocess_video( + video, + height=height, + width=width, + resize_mode=resize_mode, + ) + except TypeError: + return self.video_processor.preprocess_video(video, height=height, width=width) + + frame_tensors = [ + self._preprocess_image(frame.convert("RGB"), height, width, resize_mode=resize_mode) for frame in video + ] + return torch.stack([frame.squeeze(0) for frame in frame_tensors], dim=1).unsqueeze(0) + + def _resize_and_centercrop_tensor( + self, + tensor: torch.Tensor, + height: int, + width: int, + resize_mode: str, + ) -> torch.Tensor: + if tensor.shape[-2:] == (height, width): + return tensor + + tensor = tensor.float().unsqueeze(1) + if resize_mode == "crop": + src_height, src_width = tensor.shape[-2:] + scale = max(height / src_height, width / src_width) + resized_height = max(height, int(math.ceil(src_height * scale))) + resized_width = max(width, int(math.ceil(src_width * scale))) + tensor = F.interpolate(tensor, size=(resized_height, resized_width), mode="nearest") + top = max(0, (resized_height - height) // 2) + left = max(0, (resized_width - width) // 2) + tensor = tensor[..., top : top + height, left : left + width] + else: + tensor = F.interpolate(tensor, size=(height, width), mode="nearest") + return tensor.squeeze(1) + + def _decode_latents(self, latents: torch.Tensor, output_type: str = "np"): + if output_type == "latent": + return latents + latents = self.denormalize_latents(latents.to(self.vae.dtype)) + output_video = self.vae.decode(latents, return_dict=False)[0] + return self.video_processor.postprocess_video(output_video, output_type=output_type) + + def _generate_at2v( + self, + prompt: str, + negative_prompt: str, + audio_emb: torch.Tensor, + height: int, + width: int, + num_frames: int, + steps: int, + text_guidance_scale: float, + audio_guidance_scale: float, + use_distill: bool, + generator, + latents: torch.Tensor | None, + max_sequence_length: int, + return_latents: bool = False, + ): + self._text_guidance_scale = text_guidance_scale + self._audio_guidance_scale = audio_guidance_scale + device = self.device + dtype = self.transformer.dtype + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + audio_cond_embs = audio_emb + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + audio_uncond_embs = torch.zeros_like(audio_cond_embs) + audio_cond_embs = torch.cat([audio_cond_embs, audio_cond_embs], dim=0) + else: + audio_uncond_embs = None + + sigmas = self.get_timesteps_sigmas(steps, use_distill=use_distill) + self.scheduler.set_timesteps(steps, sigmas=sigmas, device=device) + timesteps = self.scheduler.timesteps + latents = self.prepare_latents( + image=None, + batch_size=1, + num_channels_latents=self.transformer.config.in_channels, + height=height, + width=width, + num_frames=num_frames, + num_cond_frames=0, + dtype=torch.float32, + device=device, + generator=generator, + latents=latents, + ) + + with torch.no_grad(): + for t in timesteps: + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + latent_model_input = latent_model_input.to(dtype) + timestep = t.expand(latent_model_input.shape[0]).to(dtype) + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_attention_mask, + audio_embs=audio_cond_embs, + ) + if self.do_classifier_free_guidance: + timestep_uncond = t.expand(latents.shape[0]).to(dtype) + noise_pred_uncond = self.transformer( + hidden_states=latents, + timestep=timestep_uncond, + encoder_hidden_states=negative_prompt_embeds, + encoder_attention_mask=negative_prompt_attention_mask, + audio_embs=audio_uncond_embs, + ) + noise_pred_uncond_text, noise_pred_cond = noise_pred.chunk(2) + noise_pred = ( + noise_pred_uncond + + text_guidance_scale * (noise_pred_cond - noise_pred_uncond_text) + + audio_guidance_scale * (noise_pred_uncond_text - noise_pred_uncond) + ) + latents = self.scheduler.step(-noise_pred, t, latents, return_dict=False)[0] + output = self._decode_latents(latents, output_type="np") + if return_latents: + return output, latents.clone() + return output + + def _generate_ai2v( + self, + image: Image.Image, + prompt: str, + negative_prompt: str, + audio_emb: torch.Tensor, + resolution: str, + num_frames: int, + steps: int, + text_guidance_scale: float, + audio_guidance_scale: float, + use_distill: bool, + generator, + latents: torch.Tensor | None, + max_sequence_length: int, + resize_mode: str, + ref_target_masks: torch.Tensor | None = None, + ): + height, width = self._condition_shape(image, resolution) + self._text_guidance_scale = text_guidance_scale + self._audio_guidance_scale = audio_guidance_scale + device = self.device + dtype = self.transformer.dtype + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + audio_cond_embs = audio_emb + audio_num = audio_cond_embs.shape[0] + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + audio_uncond_embs = torch.zeros_like(audio_cond_embs) + audio_cond_embs = torch.cat([audio_cond_embs, audio_cond_embs], dim=0) + else: + audio_uncond_embs = None + + sigmas = self.get_timesteps_sigmas(steps, use_distill=use_distill) + self.scheduler.set_timesteps(steps, sigmas=sigmas, device=device) + timesteps = self.scheduler.timesteps + image_tensor = self._preprocess_image(image, height, width, resize_mode=resize_mode).to( + device=device, dtype=prompt_embeds.dtype + ) + latents = self.prepare_latents( + image=image_tensor, + batch_size=1, + num_channels_latents=self.transformer.config.in_channels, + height=height, + width=width, + num_frames=num_frames, + num_cond_frames=1, + dtype=torch.float32, + device=device, + generator=generator, + latents=latents, + ) + if ref_target_masks is not None: + ref_target_masks = self._resize_and_centercrop_tensor(ref_target_masks, height, width, resize_mode) + ref_target_masks = ref_target_masks.to(device=device, dtype=torch.float32) + + with torch.no_grad(): + for t in timesteps: + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + latent_model_input = latent_model_input.to(dtype) + timestep = t.expand(latent_model_input.shape[0]).to(dtype) + timestep = timestep.unsqueeze(-1).repeat(1, latent_model_input.shape[2]) + timestep[:, :1] = 0 + if self.do_classifier_free_guidance and ref_target_masks is not None: + noise_pred_uncond_text = self.transformer( + hidden_states=latent_model_input[:1], + timestep=timestep[:1], + encoder_hidden_states=prompt_embeds[:1], + encoder_attention_mask=prompt_attention_mask[:1], + num_cond_latents=1, + audio_embs=audio_cond_embs[:audio_num], + ref_target_masks=ref_target_masks, + ) + noise_pred_cond = self.transformer( + hidden_states=latent_model_input[1:], + timestep=timestep[1:], + encoder_hidden_states=prompt_embeds[1:], + encoder_attention_mask=prompt_attention_mask[1:], + num_cond_latents=1, + audio_embs=audio_cond_embs[audio_num : 2 * audio_num], + ref_target_masks=ref_target_masks, + ) + noise_pred = torch.cat([noise_pred_uncond_text, noise_pred_cond]) + else: + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_attention_mask, + num_cond_latents=1, + audio_embs=audio_cond_embs, + ref_target_masks=ref_target_masks, + ) + if self.do_classifier_free_guidance: + timestep_uncond = t.expand(latents.shape[0]).to(dtype) + timestep_uncond = timestep_uncond.unsqueeze(-1).repeat(1, latent_model_input.shape[2]) + timestep_uncond[:, :1] = 0 + noise_pred_uncond = self.transformer( + hidden_states=latents, + timestep=timestep_uncond, + encoder_hidden_states=negative_prompt_embeds, + encoder_attention_mask=negative_prompt_attention_mask, + num_cond_latents=1, + audio_embs=audio_uncond_embs, + ref_target_masks=ref_target_masks, + ) + noise_pred_uncond_text, noise_pred_cond = noise_pred.chunk(2) + noise_pred = ( + noise_pred_uncond + + text_guidance_scale * (noise_pred_cond - noise_pred_uncond_text) + + audio_guidance_scale * (noise_pred_uncond_text - noise_pred_uncond) + ) + latents[:, :, 1:] = self.scheduler.step( + -noise_pred[:, :, 1:], + t, + latents[:, :, 1:], + return_dict=False, + )[0] + return self._decode_latents(latents, output_type="np"), latents.clone(), height, width + + def _generate_avc( + self, + video: list[Image.Image], + video_latent: torch.Tensor, + prompt: str, + negative_prompt: str, + audio_emb: torch.Tensor, + height: int, + width: int, + num_frames: int, + num_cond_frames: int, + steps: int, + text_guidance_scale: float, + audio_guidance_scale: float, + use_distill: bool, + generator, + latents: torch.Tensor | None, + max_sequence_length: int, + resize_mode: str, + ref_latent: torch.Tensor | None, + ref_img_index: int, + mask_frame_range: int, + ref_target_masks: torch.Tensor | None = None, + use_kv_cache: bool = True, + offload_kv_cache: bool = False, + ): + self._text_guidance_scale = text_guidance_scale + self._audio_guidance_scale = audio_guidance_scale + device = self.device + dtype = self.transformer.dtype + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + audio_cond_embs = audio_emb + audio_num = audio_cond_embs.shape[0] + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + audio_uncond_embs = torch.zeros_like(audio_cond_embs) + audio_cond_embs = torch.cat([audio_cond_embs, audio_cond_embs], dim=0) + else: + audio_uncond_embs = None + + sigmas = self.get_timesteps_sigmas(steps, use_distill=use_distill) + self.scheduler.set_timesteps(steps, sigmas=sigmas, device=device) + timesteps = self.scheduler.timesteps + + video_tensor = self._preprocess_video_frames(video, height, width, resize_mode=resize_mode).to( + device=device, + dtype=prompt_embeds.dtype, + ) + cond_videos = video_tensor[:, :, -num_cond_frames:] + cond_videos_latents = _retrieve_latents(self.vae.encode(cond_videos), generator, sample_mode="argmax") + cond_videos_latents = self.normalize_latents(cond_videos_latents) + + latents = self.prepare_latents( + image=None, + video=video_latent, + batch_size=1, + num_channels_latents=self.transformer.config.in_channels, + height=height, + width=width, + num_frames=num_frames, + num_cond_frames=num_cond_frames, + dtype=torch.float32, + device=device, + generator=generator, + latents=latents, + need_encode=False, + ) + num_cond_latents = 1 + (num_cond_frames - 1) // self.vae.config.scale_factor_temporal + latents[:, :, :num_cond_latents] = cond_videos_latents + + if ref_target_masks is not None: + ref_target_masks = self._resize_and_centercrop_tensor(ref_target_masks, height, width, resize_mode) + ref_target_masks = ref_target_masks.to(device=device, dtype=torch.float32) + + num_ref_latents = 0 + if ref_latent is not None: + num_cond_latents += 1 + num_ref_latents = 1 + latents = torch.cat([ref_latent.to(device=device, dtype=latents.dtype), latents], dim=2) + + if use_kv_cache: + cond_latents = latents[:, :, :num_cond_latents] + self._cache_clean_latents( + cond_latents, + max_sequence_length, + offload_kv_cache=offload_kv_cache, + dtype=dtype, + audio_embs=audio_emb, + num_cond_latents=num_cond_latents, + num_ref_latents=num_ref_latents, + ref_img_index=ref_img_index, + ) + kv_cache_dict = self._get_kv_cache_dict() + latents = latents[:, :, num_cond_latents:] + else: + cond_latents = None + kv_cache_dict = {} + + with torch.no_grad(): + for t in timesteps: + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + latent_model_input = latent_model_input.to(dtype) + timestep = t.expand(latent_model_input.shape[0]).to(dtype) + timestep = timestep.unsqueeze(-1).repeat(1, latent_model_input.shape[2]) + if not use_kv_cache: + timestep[:, :num_cond_latents] = 0 + + if self.do_classifier_free_guidance and ref_target_masks is not None: + noise_pred_uncond_text = self.transformer( + hidden_states=latent_model_input[:1], + timestep=timestep[:1], + encoder_hidden_states=prompt_embeds[:1], + encoder_attention_mask=prompt_attention_mask[:1], + num_cond_latents=num_cond_latents, + kv_cache_dict=kv_cache_dict, + audio_embs=audio_cond_embs[:audio_num], + num_ref_latents=num_ref_latents, + ref_img_index=ref_img_index, + mask_frame_range=mask_frame_range, + ref_target_masks=ref_target_masks, + ) + noise_pred_cond = self.transformer( + hidden_states=latent_model_input[1:], + timestep=timestep[1:], + encoder_hidden_states=prompt_embeds[1:], + encoder_attention_mask=prompt_attention_mask[1:], + num_cond_latents=num_cond_latents, + kv_cache_dict=kv_cache_dict, + audio_embs=audio_cond_embs[audio_num : 2 * audio_num], + num_ref_latents=num_ref_latents, + ref_img_index=ref_img_index, + mask_frame_range=mask_frame_range, + ref_target_masks=ref_target_masks, + ) + noise_pred = torch.cat([noise_pred_uncond_text, noise_pred_cond]) + else: + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_attention_mask, + num_cond_latents=num_cond_latents, + kv_cache_dict=kv_cache_dict, + audio_embs=audio_cond_embs, + num_ref_latents=num_ref_latents, + ref_img_index=ref_img_index, + mask_frame_range=mask_frame_range, + ref_target_masks=ref_target_masks, + ) + + if self.do_classifier_free_guidance: + timestep_uncond = t.expand(latents.shape[0]).to(dtype) + timestep_uncond = timestep_uncond.unsqueeze(-1).repeat(1, latent_model_input.shape[2]) + if not use_kv_cache: + timestep_uncond[:, :num_cond_latents] = 0 + noise_pred_uncond = self.transformer( + hidden_states=latents, + timestep=timestep_uncond, + encoder_hidden_states=negative_prompt_embeds, + encoder_attention_mask=negative_prompt_attention_mask, + num_cond_latents=num_cond_latents, + kv_cache_dict=kv_cache_dict, + audio_embs=audio_uncond_embs, + num_ref_latents=num_ref_latents, + ref_img_index=ref_img_index, + mask_frame_range=mask_frame_range, + ref_target_masks=ref_target_masks, + ) + noise_pred_uncond_text, noise_pred_cond = noise_pred.chunk(2) + noise_pred = ( + noise_pred_uncond + + text_guidance_scale * (noise_pred_cond - noise_pred_uncond_text) + + audio_guidance_scale * (noise_pred_uncond_text - noise_pred_uncond) + ) + + if use_kv_cache: + latents = self.scheduler.step(-noise_pred, t, latents, return_dict=False)[0] + else: + latents[:, :, num_cond_latents:] = self.scheduler.step( + -noise_pred[:, :, num_cond_latents:], + t, + latents[:, :, num_cond_latents:], + return_dict=False, + )[0] + + if use_kv_cache: + latents = torch.cat([cond_latents, latents], dim=2) + self._clear_cache() + + if ref_latent is not None: + latents = latents[:, :, num_ref_latents:] + num_cond_latents -= num_ref_latents + latents[:, :, :num_cond_latents] = cond_videos_latents + + return self._decode_latents(latents, output_type="np"), latents.clone() + + def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + if req.is_dummy_run(): + height = req.sampling_params.height or 512 + width = req.sampling_params.width or 512 + frame = torch.zeros((1, height, width, 3), dtype=torch.uint8) + return DiffusionOutput(output=frame) + + started = time.perf_counter() + inputs = self._resolve_request_inputs(req) + extra = req.sampling_params.extra_args or {} + stage = inputs["stage"] + if stage not in {"at2v", "ai2v"}: + raise ValueError(f"LongCat-Video-Avatar stage must be 'at2v' or 'ai2v', got {stage!r}.") + resolution = inputs["resolution"] + if resolution not in {"480p", "720p"}: + raise ValueError(f"Unsupported LongCat-Video-Avatar resolution {resolution!r}.") + + num_frames = _adjust_num_frames(int(extra.get("num_frames") or req.sampling_params.num_frames or 93)) + save_fps = int(extra.get("save_fps") or req.sampling_params.fps or self.save_fps) + num_cond_frames = int(extra.get("num_cond_frames") or 13) + use_distill = _as_bool(extra.get("use_distill"), self.use_distill) + steps = int(extra.get("num_inference_steps") or req.sampling_params.num_inference_steps or 50) + text_guidance_scale = float(extra.get("text_guidance_scale") or req.sampling_params.guidance_scale or 4.0) + audio_guidance_scale = float(extra.get("audio_guidance_scale") or req.sampling_params.guidance_scale_2 or 4.0) + num_segments_value = extra.get("num_segments") + resize_mode = str(extra.get("resize_mode") or "crop") + use_kv_cache = _as_bool(extra.get("use_kv_cache"), True) + offload_kv_cache = _as_bool(extra.get("offload_kv_cache"), False) + ref_img_index = int(extra.get("ref_img_index") or 10) + mask_frame_range = int(extra.get("mask_frame_range") or 3) + if use_distill: + steps = 8 + text_guidance_scale = 1.0 + audio_guidance_scale = 1.0 + + generator = req.sampling_params.generator + if generator is None: + seed = req.sampling_params.seed if req.sampling_params.seed is not None else 42 + generator = torch.Generator(device=self.device).manual_seed(seed) + latents = req.sampling_params.latents + max_sequence_length = req.sampling_params.max_sequence_length or int(extra.get("max_sequence_length") or 512) + + image = None + ref_target_masks = None + include_background_silent_audio = False + explicit_num_segments = _requested_num_segments(num_segments_value) + auto_num_segments = isinstance(num_segments_value, str) and num_segments_value.lower() == "auto" + needs_full_audio = auto_num_segments or explicit_num_segments > 1 + num_segments = explicit_num_segments or 1 + if inputs["is_multi_speaker"]: + if stage != "ai2v": + raise NotImplementedError( + "LongCat-Video-Avatar multi-speaker generation requires AI2V " + "because speaker bounding boxes and masks are defined on a reference image." + ) + image = _load_image(inputs["image"]) + ref_target_masks, include_background_silent_audio = _build_multi_speaker_ref_target_masks( + image, + inputs["bbox"], + ) + + audio_emb, full_audio_embs, num_segments = self._prepare_audio_embeddings( + inputs, + is_multi_speaker=inputs["is_multi_speaker"], + include_background_silent_audio=include_background_silent_audio, + needs_full_audio=needs_full_audio, + auto_num_segments=auto_num_segments, + num_segments_value=num_segments_value, + num_segments=num_segments, + num_frames=num_frames, + num_cond_frames=num_cond_frames, + save_fps=save_fps, + ) + + if num_segments > 1 and num_cond_frames >= num_frames: + raise ValueError( + "LongCat-Video-Avatar AVC continuation requires num_cond_frames to be smaller than num_frames." + ) + + latent = None + if stage == "at2v": + default_height, default_width = _default_at2v_shape(resolution) + height = int(extra.get("height") or req.sampling_params.height or default_height) + width = int(extra.get("width") or req.sampling_params.width or default_width) + output = self._generate_at2v( + prompt=inputs["prompt"], + negative_prompt=inputs["negative_prompt"], + audio_emb=audio_emb, + height=height, + width=width, + num_frames=num_frames, + steps=steps, + text_guidance_scale=text_guidance_scale, + audio_guidance_scale=audio_guidance_scale, + use_distill=use_distill, + generator=generator, + latents=latents, + max_sequence_length=max_sequence_length, + return_latents=num_segments > 1, + ) + if num_segments > 1: + output, latent = output + else: + image = image or _load_image(inputs["image"]) + output, latent, height, width = self._generate_ai2v( + image=image, + prompt=inputs["prompt"], + negative_prompt=inputs["negative_prompt"], + audio_emb=audio_emb, + resolution=resolution, + num_frames=num_frames, + steps=steps, + text_guidance_scale=text_guidance_scale, + audio_guidance_scale=audio_guidance_scale, + use_distill=use_distill, + generator=generator, + latents=latents, + max_sequence_length=max_sequence_length, + resize_mode=resize_mode, + ref_target_masks=ref_target_masks, + ) + + if num_segments > 1: + if full_audio_embs is None or latent is None: + raise RuntimeError("LongCat-Video-Avatar continuation requires full audio embeddings and latents.") + segment_frames = output[0] + current_video = [ + Image.fromarray((np.clip(frame, 0.0, 1.0) * 255).round().astype("uint8")) for frame in segment_frames + ] + all_generated_frames = list(current_video) + ref_latent = latent[:, :, :1].clone() + audio_start_idx = 0 + for _segment_idx in range(1, num_segments): + audio_start_idx += self.audio_stride * (num_frames - num_cond_frames) + segment_audio_emb = self._slice_multi_audio_embeddings( + full_audio_embs, + audio_start_idx=audio_start_idx, + num_frames=num_frames, + ) + segment_output, latent = self._generate_avc( + video=current_video, + video_latent=latent, + prompt=inputs["prompt"], + negative_prompt=inputs["negative_prompt"], + audio_emb=segment_audio_emb, + height=height, + width=width, + num_frames=num_frames, + num_cond_frames=num_cond_frames, + steps=steps, + text_guidance_scale=text_guidance_scale, + audio_guidance_scale=audio_guidance_scale, + use_distill=use_distill, + generator=generator, + latents=None, + max_sequence_length=max_sequence_length, + resize_mode=resize_mode, + ref_latent=ref_latent, + ref_img_index=ref_img_index, + mask_frame_range=mask_frame_range, + ref_target_masks=ref_target_masks, + use_kv_cache=use_kv_cache, + offload_kv_cache=offload_kv_cache, + ) + new_video = [ + Image.fromarray((np.clip(frame, 0.0, 1.0) * 255).round().astype("uint8")) + for frame in segment_output[0] + ] + all_generated_frames.extend(new_video[num_cond_frames:]) + current_video = new_video + frames = np.asarray([np.asarray(frame) for frame in all_generated_frames], dtype=np.uint8) + else: + frames = (np.clip(output[0], 0.0, 1.0) * 255).round().astype("uint8") + return DiffusionOutput( + output=torch.from_numpy(frames), + stage_durations={"avatar_generate_s": time.perf_counter() - started}, + ) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 3a6aeb7c2..75ef7391a 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -126,6 +126,11 @@ _DIFFUSION_MODELS = { "pipeline_longcat_image", "LongCatImagePipeline", ), + "LongCatVideoAvatarPipeline": ( + "longcat_video", + "pipeline_longcat_video_avatar", + "LongCatVideoAvatarPipeline", + ), "BagelPipeline": ( "bagel", "pipeline_bagel", @@ -525,6 +530,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "WanT2VDMD2Pipeline": "get_wan22_post_process_func", "WanI2VDMD2Pipeline": "get_wan22_i2v_post_process_func", "LongCatImagePipeline": "get_longcat_image_post_process_func", + "LongCatVideoAvatarPipeline": "get_longcat_video_avatar_post_process_func", "BagelPipeline": "get_bagel_post_process_func", "LancePipeline": "get_lance_post_process_func", "MingImagePipeline": "get_ming_image_post_process_func", @@ -575,6 +581,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { "QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func", "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_pre_process_func", "LongCatImageEditPipeline": "get_longcat_image_edit_pre_process_func", + "LongCatVideoAvatarPipeline": "get_longcat_video_avatar_pre_process_func", "QwenImageLayeredPipeline": "get_qwen_image_layered_pre_process_func", "WanPipeline": "get_wan22_pre_process_func", "WanVACEPipeline": "get_wan22_vace_pre_process_func",