diff --git a/README.md b/README.md index 35b0594b7..d20a90e12 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@

- + vllm-omni

@@ -17,6 +17,7 @@ Easy, fast, and cheap omni-modality model serving for everyone *Latest News* 🔥 +- [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more. - [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation. - [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results. - [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm). diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index f75bd568b..3d59f9b04 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -9,8 +9,7 @@ ARG PYTHON_VERSION=3.12 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu" RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \ - echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \ - add-apt-repository -y ppa:kobuk-team/intel-graphics + echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list RUN apt clean && apt-get update -y && \ apt-get install -y --no-install-recommends --fix-missing \ @@ -31,9 +30,22 @@ RUN apt clean && apt-get update -y && \ python3-pip RUN apt update && apt upgrade -y && \ - apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing intel-ocloc && \ apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3 +# Install UMD with fixed version +RUN mkdir neo && \ + cd neo && \ + wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \ + wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \ + wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \ + wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \ + wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \ + wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \ + wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \ + dpkg -i *.deb && \ + cd .. && \ + rm -rf neo + ENV PATH="/root/.local/bin:$PATH" ENV VIRTUAL_ENV="/opt/venv" ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 804655d00..98c746e09 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -35,6 +35,8 @@ th { |`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | |`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | +|`CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | +|`MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | |`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | |`OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | diff --git a/examples/offline_inference/hunyuan_image3/README.md b/examples/offline_inference/hunyuan_image3/README.md new file mode 100644 index 000000000..da28a44d9 --- /dev/null +++ b/examples/offline_inference/hunyuan_image3/README.md @@ -0,0 +1,25 @@ +# HunyuanImage-3.0 Image-to-Text Inference + +This example demonstrates how to run HunyuanImage-3.0 Image-to-Text with the vLLM-Omni. + +## Local CLI Usage + +Download the example image: + +```bash +wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg +``` + +Run example: + +```bash +python image_to_text.py \ + --image cherry_blossom.jpg \ + --prompt "<|startoftext|>You are an assistant that understands images and outputs text.Describe the content of the picture." +``` + +Key arguments: + +- `--model`: Model used. Default is: tencent/HunyuanImage-3.0-Instruct (Optional). +- `--image`: Path to input image (required). +- `--prompt`: Text description used to guide image understanding (required). diff --git a/examples/offline_inference/hunyuan_image3/image_to_text.py b/examples/offline_inference/hunyuan_image3/image_to_text.py new file mode 100644 index 000000000..dbae04315 --- /dev/null +++ b/examples/offline_inference/hunyuan_image3/image_to_text.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import argparse +import os + +from PIL import Image + +from vllm_omni.entrypoints.omni import Omni + +""" +The tencent/HunyuanImage-3.0-Instruct base model is built on the Hunyuan v1 architecture, specifically the tencent/Hunyuan-A13B-Instruct model. It utilizes two tokenizer delimiter templates: + +1) Pretrained template (default for gen_text mode), which concatenates system, image + tokens, and user question WITHOUT role delimiters: +"<|startoftext|>{system_prompt}{image_tokens}{user_question}" + + Example (before image token expansion): +"<|startoftext|>You are an assistant that understands images and outputs text.Describe the content of the picture." + +2) Instruct template, which uses explicit role prefixes and separators. +""" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Generate text from image using HunyuanImage-3.0-Instruct.") + parser.add_argument( + "--model", + default="tencent/HunyuanImage-3.0-Instruct", + help="Model name or local path.", + ) + parser.add_argument( + "--image", + type=str, + required=True, + help="Path to input image file (PNG, JPG, etc.).", + ) + parser.add_argument( + "--prompt", + type=str, + required=True, + help="Pretrain template prompt: <|startoftext|>{system}{question}", + ) + return parser.parse_args() + + +def load_image(image_path: str) -> Image.Image: + """Load an image from file path.""" + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") + return Image.open(image_path).convert("RGB") + + +def main(args: argparse.Namespace) -> None: + omni = Omni(model=args.model) + + prompt_dict = { + "prompt": args.prompt, + "modalities": ["text"], + } + + # Add image input if provided + if args.image: + if not os.path.exists(args.image): + raise FileNotFoundError(f"Input image not found: {args.image}") + + input_image = load_image(args.image) + prompt_dict["multi_modal_data"] = {"image": input_image} + + prompts = [prompt_dict] + omni_outputs = omni.generate(prompts=prompts) + + prompt_text = omni_outputs[0].request_output[0].prompt + generated_text = omni_outputs[0].request_output[0].outputs[0].text + print(f"Prompt: {prompt_text}") + print(f"Text: {generated_text}") + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/examples/offline_inference/mammothmodal2_preview/README.md b/examples/offline_inference/mammothmodal2_preview/README.md new file mode 100644 index 000000000..44b20c76a --- /dev/null +++ b/examples/offline_inference/mammothmodal2_preview/README.md @@ -0,0 +1,32 @@ +# MammothModa2-Preview + +## Run examples (MammothModa2-Preview) + +Download model +```bash +hf bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview +``` + +### Text-to-Image (T2I) + +```bash +python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py \ + --model ./MammothModa2-Preview \ + --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \ + --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \ + --height 1024 \ + --width 1024 \ + --num-inference-steps 50 \ + --text-guidance-scale 4.0 \ + --out output.png +``` + +### Image Summary + +```bash +python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \ + --model ./MammothModa2-Preview \ + --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \ + --question "Summarize this image." \ + --image ./image.png +``` diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py new file mode 100644 index 000000000..e29f907f2 --- /dev/null +++ b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py @@ -0,0 +1,115 @@ +""" +Offline inference example: MammothModa2 image summarization (single AR stage). + +Example: + uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \ + --model path/to/MammothModa2-Preview \ + --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \ + --image /path/to/input.jpg \ + --question "Please summarize the content of this image." +""" + +from __future__ import annotations + +import argparse +import os + +from PIL import Image +from vllm import SamplingParams +from vllm.multimodal.image import convert_image_mode + +from vllm_omni import Omni + +DEFAULT_SYSTEM = "You are a helpful assistant." +DEFAULT_QUESTION = "Please summarize the content of this image." + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).") + parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.") + parser.add_argument( + "--stage-config", type=str, required=True, help="Path to stage config yaml (single-stage AR->text)." + ) + parser.add_argument("--image", type=str, required=True, help="Path to input image.") + parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.") + parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.") + parser.add_argument( + "--max-tokens", + type=int, + default=512, + help="Max new tokens to generate.", + ) + parser.add_argument("--temperature", type=float, default=0.2) + parser.add_argument("--top-p", type=float, default=0.9) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--trust-remote-code", action="store_true") + return parser.parse_args() + + +def build_prompt(system: str, question: str) -> str: + return ( + f"<|im_start|>system\n{system}<|im_end|>\n" + "<|im_start|>user\n" + "<|vision_start|><|image_pad|><|vision_end|>" + f"{question}<|im_end|>\n" + "<|im_start|>assistant\n" + ) + + +def main() -> None: + args = parse_args() + + if not os.path.exists(args.image): + raise FileNotFoundError(f"Image file not found: {args.image}") + + pil_image = Image.open(args.image) + image_data = convert_image_mode(pil_image, "RGB") + prompt = build_prompt(args.system, args.question) + + omni = Omni( + model=args.model, + stage_configs_path=args.stage_config, + trust_remote_code=args.trust_remote_code, + ) + try: + sp = SamplingParams( + temperature=float(args.temperature), + top_p=float(args.top_p), + top_k=-1, + max_tokens=int(args.max_tokens), + seed=int(args.seed), + detokenize=True, + ) + # NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None]. + # Consume it inside the try block so the worker isn't closed early. + outputs = list( + omni.generate( + [ + { + "prompt": prompt, + "multi_modal_data": {"image": image_data}, + "additional_information": {"omni_task": ["chat"]}, + } + ], + [sp], + ) + ) + finally: + omni.close() + + lines: list[str] = [] + for stage_outputs in outputs: + req_outputs = getattr(stage_outputs, "request_output", stage_outputs) + req_outputs = req_outputs if isinstance(req_outputs, list) else [req_outputs] + for ro in req_outputs: + text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro) + lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n") + lines.append("answer:\n") + lines.append(text.strip() + "\n") + lines.append("\n") + + print("\n".join(lines)) + + +if __name__ == "__main__": + main() diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py new file mode 100644 index 000000000..1f10fb288 --- /dev/null +++ b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py @@ -0,0 +1,251 @@ +""" +Offline inference example for MammothModa2 Text-to-Image (T2I) generation. +This script uses the vllm_omni.Omni pipeline with a multi-stage configuration. + +Workflow: +1. Stage 0 (AR): Generates visual tokens and their corresponding hidden states. +2. Stage 1 (DiT): Consumes the hidden states as conditions to perform diffusion + and VAE decoding to produce the final image. + +Example Usage: + uv run python examples/offline_inference/run_mammothmoda2_t2i.py \ + --model path/to/MammothModa2-Preview \ + --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \ + --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \ + --out output.png +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +from pathlib import Path +from typing import NamedTuple + +import torch +from PIL import Image +from vllm.sampling_params import SamplingParams + +from vllm_omni import Omni + +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- +_PATCH_SIZE = 16 # AR image grid patch size (pixels per token) + + +class T2IGenConfig(NamedTuple): + eol_token_id: int + visual_token_start_id: int + visual_token_end_id: int + top_k: int # AR sampling top-k (covers the full visual generation vocabulary) + # Qwen2.5-VL special vision tokens: <|image_pad|>, <|video_pad|>, <|vision_start|>, <|vision_end|> + visual_ids: list[int] + + +def load_t2i_generation_config(model_dir: str) -> T2IGenConfig: + """Load T2I token IDs from t2i_generation_config.json and config.json. + + Supports both local directory paths and HuggingFace Hub model IDs. + """ + model_path = Path(model_dir) + + def _read_json(filename: str) -> dict: + local = model_path / filename + if local.exists(): + with local.open(encoding="utf-8") as f: + return json.load(f) + # Fall back to HuggingFace Hub when model_dir is a repo ID. + try: + from huggingface_hub import hf_hub_download + except ImportError as exc: + raise ImportError( + "huggingface_hub is required to load configs from HF Hub. Install it with: pip install huggingface_hub" + ) from exc + cached = hf_hub_download(repo_id=model_dir, filename=filename) + with open(cached, encoding="utf-8") as f: + return json.load(f) + + gen_cfg = _read_json("t2i_generation_config.json") + llm_cfg = _read_json("config.json").get("llm_config", {}) + + return T2IGenConfig( + eol_token_id=int(gen_cfg["eol_token_id"]), + visual_token_start_id=int(gen_cfg["visual_token_start_id"]), + visual_token_end_id=int(gen_cfg["visual_token_end_id"]), + top_k=int(gen_cfg["top_k"]), + visual_ids=[ + int(llm_cfg["image_token_id"]), + int(llm_cfg["video_token_id"]), + int(llm_cfg["vision_start_token_id"]), + int(llm_cfg["vision_end_token_id"]), + ], + ) + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description="Run MammothModa2 T2I (AR -> DiT) with vLLM-Omni.") + p.add_argument("--model", type=str, required=True, help="Path to the model directory.") + p.add_argument("--stage-config", type=str, required=True, help="Path to the multi-stage YAML configuration.") + p.add_argument( + "--prompt", + type=str, + action="append", + default=None, + help=( + "Text prompt for image generation. Can be provided multiple times " + "to generate multiple images with shared height/width/CFG settings." + ), + ) + p.add_argument("--height", type=int, default=1024, help="Output image height (must be a multiple of 16).") + p.add_argument("--width", type=int, default=1024, help="Output image width (must be a multiple of 16).") + p.add_argument("--num-inference-steps", type=int, default=50, help="Number of diffusion steps for the DiT stage.") + p.add_argument( + "--text-guidance-scale", type=float, default=9.0, help="Classifier-Free Guidance (CFG) scale for DiT." + ) + p.add_argument( + "--cfg-range", + type=float, + nargs=2, + default=(0.0, 1.0), + help="Relative step range [start, end] where CFG is active.", + ) + p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.") + p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.") + args = p.parse_args() + if not args.prompt: + args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."] + return args + + +def tensor_to_pil(image: torch.Tensor) -> Image.Image: + """Convert a normalized torch tensor [-1, 1] to a PIL Image.""" + if image.ndim == 4: + image = image[0] + image = image.detach().to("cpu") + image = (image / 2 + 0.5).clamp(0, 1) + image = (image * 255).to(torch.uint8) + image = image.permute(1, 2, 0).contiguous().numpy() + return Image.fromarray(image) + + +def _format_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str: + """Build the AR-stage prompt string including the image grid header.""" + return ( + "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n" + f"<|im_start|>user\n{user_prompt}<|im_end|>\n" + "<|im_start|>assistant\n" + f"<|image start|>{ar_width}*{ar_height}<|image token|>" + ) + + +def _collect_images(outputs: list) -> list[torch.Tensor]: + """Extract all image tensors produced by the final (DiT) stage.""" + images: list[torch.Tensor] = [] + for out in outputs: + ro_list = getattr(out, "request_output", out) + if not isinstance(ro_list, list): + ro_list = [ro_list] + for ro_item in ro_list: + for completion in getattr(ro_item, "outputs", None) or []: + mm = getattr(completion, "multimodal_output", None) + if not isinstance(mm, dict) or "image" not in mm: + raise RuntimeError(f"Missing image in multimodal output: {mm}") + payload = mm["image"] + for tensor in payload if isinstance(payload, list) else [payload]: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"Expected image tensor, got {type(tensor)}") + images.append(tensor) + return images + + +def _save_images(images: list[torch.Tensor], out_path: str) -> list[str]: + """Save image tensors to disk. + + Single image: written to *out_path* exactly. + Multiple images: suffixed as ``_0``, ``_1``, … + """ + if not images: + raise RuntimeError("No images to save.") + base, ext = os.path.splitext(out_path) + ext = ext or ".png" + paths = [] + for i, tensor in enumerate(images): + path = out_path if len(images) == 1 else f"{base}_{i}{ext}" + tensor_to_pil(tensor).save(path) + paths.append(path) + return paths + + +def main() -> None: + args = parse_args() + os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True) + + if args.height <= 0 or args.width <= 0: + raise ValueError(f"Height and width must be positive, got {args.height}x{args.width}") + if args.height % _PATCH_SIZE != 0 or args.width % _PATCH_SIZE != 0: + raise ValueError(f"Height and width must be multiples of {_PATCH_SIZE}, got {args.height}x{args.width}") + + ar_height = args.height // _PATCH_SIZE + ar_width = args.width // _PATCH_SIZE + gen_cfg = load_t2i_generation_config(args.model) + expected_grid_tokens = ar_height * (ar_width + 1) + + logger.info("Initializing Omni pipeline...") + omni = Omni(model=args.model, stage_configs_path=args.stage_config, trust_remote_code=args.trust_remote_code) + try: + ar_sampling = SamplingParams( + temperature=1.0, + top_p=1.0, + top_k=gen_cfg.top_k, + max_tokens=max(1, expected_grid_tokens + 1), # +1 for hidden state of eoi + detokenize=False, + ) + dit_sampling = SamplingParams( + temperature=0.0, + top_p=1.0, + top_k=-1, + max_tokens=1, + detokenize=False, + ) + + additional_information = { + "omni_task": ["t2i"], + "ar_width": [ar_width], + "ar_height": [ar_height], + "eol_token_id": [gen_cfg.eol_token_id], + "visual_token_start_id": [gen_cfg.visual_token_start_id], + "visual_token_end_id": [gen_cfg.visual_token_end_id], + "image_height": [args.height], + "image_width": [args.width], + "num_inference_steps": [args.num_inference_steps], + "text_guidance_scale": [args.text_guidance_scale], + "cfg_range": [args.cfg_range[0], args.cfg_range[1]], + "visual_ids": gen_cfg.visual_ids, + } + inputs = [ + { + "prompt": _format_prompt(p, ar_width, ar_height), + "additional_information": dict(additional_information), + } + for p in args.prompt + ] + + logger.info("Starting generation...") + # omni.generate() returns a Generator; consume it to run the full pipeline. + outputs = list(omni.generate(inputs, [ar_sampling, dit_sampling])) + + logger.info("Post-processing and saving image(s)...") + for path in _save_images(_collect_images(outputs), args.out): + logger.info(f"Saved: {path}") + finally: + omni.close() + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 95e43d871..99194c648 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -203,3 +203,5 @@ extend-ignore-identifiers-re = [ ".*[Oo]no_[Aa]nna.*", ".*cann.*", ] +[tool.typos.default.extend-words] +ue = "ue" diff --git a/requirements/common.txt b/requirements/common.txt index 7e5f5f9b2..548f8f0bf 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -12,6 +12,8 @@ torchsde>=0.2.6 openai-whisper>=20250625 imageio[ffmpeg]>=2.37.2 sox>=1.5.0 +x-transformers>=2.12.2 +einops>=0.8.1 prettytable>=3.8.0 aenum==3.1.16 pyzmq>=25.0.0 diff --git a/tests/e2e/offline_inference/test_mammoth_moda2.py b/tests/e2e/offline_inference/test_mammoth_moda2.py new file mode 100644 index 000000000..5293b5ed1 --- /dev/null +++ b/tests/e2e/offline_inference/test_mammoth_moda2.py @@ -0,0 +1,215 @@ +""" +End-to-end test for MammothModa2 text-to-image generation. + +Verifies that the AR->DiT pipeline produces an image tensor whose pixel values +match a golden reference. + +Model weights: env var ``MAMMOTHMODA2_MODEL_PATH`` + (default: /MammothModa2-Preview) +Stage config: env var ``MAMMOTHMODA2_T2I_STAGE_CONFIG`` + (default: vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml) + +Golden pixel file: ``tests/e2e/offline_inference/fixtures/mammoth_moda2_t2i_golden.json`` + Regenerate with: ``UPDATE_GOLDEN=1 pytest tests/e2e/offline_inference/test_mammoth_moda2.py`` +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +import pytest +import torch +from vllm.sampling_params import SamplingParams + +from tests.utils import hardware_test + +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- +_IMAGE_TOKEN_ID = 151655 # "<|image_pad|>" +_VIDEO_TOKEN_ID = 151656 # "<|video_pad|>" +_VISION_START_TOKEN_ID = 151652 # "<|vision_start|>" +_VISION_END_TOKEN_ID = 151653 # "<|vision_end|>" +_AR_PATCH_SIZE = 16 + +_STAGE_CONFIGS_DIR = Path(__file__).resolve().parents[3] / "vllm_omni" / "model_executor" / "stage_configs" +MODEL_PATH = os.environ.get( + "MAMMOTHMODA2_MODEL_PATH", + str(Path(__file__).resolve().parents[3] / "MammothModa2-Preview"), +) +T2I_STAGE_CONFIG = os.environ.get( + "MAMMOTHMODA2_T2I_STAGE_CONFIG", + str(_STAGE_CONFIGS_DIR / "mammoth_moda2.yaml"), +) + +# Golden pixel reference file. Set UPDATE_GOLDEN=1 to regenerate. +_GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json" +# Fixed sampling coordinates: (channel, row_fraction, col_fraction) +# Covers corners, centre, and mid-edges across all 3 channels. +_PIXEL_SAMPLE_COORDS = [ + (0, 0.0, 0.0), + (0, 0.5, 0.5), + (0, 1.0, 1.0), + (0, 0.25, 0.75), + (1, 0.0, 1.0), + (1, 0.5, 0.0), + (1, 0.75, 0.25), + (1, 1.0, 0.5), + (2, 0.0, 0.5), + (2, 0.5, 1.0), + (2, 0.75, 0.75), + (2, 1.0, 0.0), +] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _load_t2i_gen_config(model_dir: str) -> dict: + cfg_path = Path(model_dir) / "t2i_generation_config.json" + if not cfg_path.exists(): + pytest.skip(f"t2i_generation_config.json not found at {cfg_path}") + with cfg_path.open() as f: + return json.load(f) + + +def _format_t2i_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str: + return ( + "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n" + f"<|im_start|>user\n{user_prompt}<|im_end|>\n" + "<|im_start|>assistant\n" + f"<|image start|>{ar_width}*{ar_height}<|image token|>" + ) + + +def _sample_pixels(img_tensor: torch.Tensor) -> list[float]: + """Sample pixel values at fixed fractional coordinates from a (C, H, W) tensor.""" + t = img_tensor.float().clamp(0.0, 1.0) + if t.ndim == 4: + t = t[0] # unbatch + C, H, W = t.shape + values = [] + for c, rh, rw in _PIXEL_SAMPLE_COORDS: + ri = min(int(rh * (H - 1)), H - 1) + ci = min(int(rw * (W - 1)), W - 1) + values.append(round(float(t[c, ri, ci]), 6)) + return values + + +# --------------------------------------------------------------------------- +# End-to-end test +# --------------------------------------------------------------------------- +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4"}) +def test_mammothmoda2_t2i_e2e(): + """ + End-to-end text-to-image generation with MammothModa2 (AR -> DiT). + + Verifies: + - Omni pipeline initialises with the two-stage YAML config. + - DiT stage outputs an image tensor with the correct shape. + - A fixed set of pixel values matches a golden reference + (regenerate with ``UPDATE_GOLDEN=1``). + """ + from vllm_omni import Omni + + if not Path(MODEL_PATH).exists(): + pytest.skip(f"Model weights not found at {MODEL_PATH}") + if not Path(T2I_STAGE_CONFIG).exists(): + pytest.skip(f"Stage config not found at {T2I_STAGE_CONFIG}") + + gen_cfg = _load_t2i_gen_config(MODEL_PATH) + eol_token_id = int(gen_cfg["eol_token_id"]) + visual_start = int(gen_cfg["visual_token_start_id"]) + visual_end = int(gen_cfg["visual_token_end_id"]) + + height, width = 256, 256 # small for CI speed + ar_height, ar_width = height // _AR_PATCH_SIZE, width // _AR_PATCH_SIZE + expected_grid_tokens = ar_height * (ar_width + 1) + + prompt_text = "A cat sitting on a laptop keyboard" + formatted_prompt = _format_t2i_prompt(prompt_text, ar_width, ar_height) + + omni = Omni(model=MODEL_PATH, stage_configs_path=T2I_STAGE_CONFIG, trust_remote_code=True) + try: + # Greedy / deterministic sampling so pixel values are reproducible. + ar_sampling = SamplingParams( + temperature=0.0, + top_k=1, + max_tokens=max(1, expected_grid_tokens + 1), + detokenize=False, + ) + dit_sampling = SamplingParams(temperature=0.0, max_tokens=1, detokenize=False) + + outputs = list( + omni.generate( + [ + { + "prompt": formatted_prompt, + "additional_information": { + "omni_task": ["t2i"], + "ar_width": [ar_width], + "ar_height": [ar_height], + "eol_token_id": [eol_token_id], + "visual_token_start_id": [visual_start], + "visual_token_end_id": [visual_end], + "image_height": [height], + "image_width": [width], + "num_inference_steps": [2], + "text_guidance_scale": [1.0], + "cfg_range": [0.0, 1.0], + "visual_ids": [ + _IMAGE_TOKEN_ID, + _VIDEO_TOKEN_ID, + _VISION_START_TOKEN_ID, + _VISION_END_TOKEN_ID, + ], + }, + } + ], + [ar_sampling, dit_sampling], + ) + ) + + assert len(outputs) > 0, "Pipeline produced no outputs" + + # Find the image tensor produced by the DiT stage and compare pixels. + found_image = False + for out in outputs: + ro_list = getattr(out, "request_output", out) + if not isinstance(ro_list, list): + ro_list = [ro_list] + for ro in ro_list: + completion_outputs = getattr(ro, "outputs", None) + if not isinstance(completion_outputs, list): + continue + for completion in completion_outputs: + mm = getattr(completion, "multimodal_output", None) + if not (isinstance(mm, dict) and "image" in mm): + continue + img_list = mm["image"] if isinstance(mm["image"], list) else [mm["image"]] + for img_tensor in img_list: + assert isinstance(img_tensor, torch.Tensor), f"Expected image tensor, got {type(img_tensor)}" + assert img_tensor.ndim in (3, 4), f"Expected 3D or 4D image tensor, got {img_tensor.ndim}D" + + sampled = _sample_pixels(img_tensor) + + if os.environ.get("UPDATE_GOLDEN"): + _GOLDEN_T2I_PATH.parent.mkdir(parents=True, exist_ok=True) + _GOLDEN_T2I_PATH.write_text(json.dumps({"pixels": sampled}, indent=2)) + print(f"\nGolden file written to {_GOLDEN_T2I_PATH}") + elif _GOLDEN_T2I_PATH.exists(): + golden = json.loads(_GOLDEN_T2I_PATH.read_text())["pixels"] + for i, (got, exp) in enumerate(zip(sampled, golden)): + assert abs(got - exp) < 1e-4, f"Pixel {i} mismatch: got {got}, expected {exp}" + + found_image = True + + assert found_image, "No image tensor found in pipeline output" + finally: + omni.close() diff --git a/tests/entrypoints/test_omni_diffusion.py b/tests/entrypoints/test_omni_diffusion.py index 29e6dcf6a..90b0fd053 100644 --- a/tests/entrypoints/test_omni_diffusion.py +++ b/tests/entrypoints/test_omni_diffusion.py @@ -4,6 +4,7 @@ from queue import Empty, Queue from typing import Any import pytest +import torch from pytest_mock import MockerFixture from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK @@ -1340,3 +1341,91 @@ def test_close_sends_shutdown_signal(monkeypatch: pytest.MonkeyPatch, mocker: Mo # Verify stop_stage_worker was called (process should be set) assert omni.stage_list[0]._proc is not None + + +# Tests below are for diffusion dtype normalization fix from the following: +# https://github.com/vllm-project/vllm-omni/pull/1391 +# In the future we should ensure dtypes are parsed in less hacky way and make +# these tests more atomic. +@pytest.mark.parametrize("dtype", ["float16", torch.float16]) +def test_dtype_normalization_valid_types( + monkeypatch, dtype: str | torch.dtype, mocker: MockerFixture, fake_stage_config +): + """Ensure Diffusion Config builder coerces valid types correctly.""" + + def _fake_loader(model: str, base_engine_args=None): + # Return not stage configs to fall back to the diffusion cfg builder + return [] + + import sys + + for module_name in [ + "vllm_omni.entrypoints.utils", + "vllm_omni.entrypoints.omni", + "vllm_omni.entrypoints.omni_stage", + ]: + if module_name in sys.modules: + del sys.modules[module_name] + + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) + _setup_ipc_mocks(monkeypatch) + _setup_log_mocks(monkeypatch) + + from vllm_omni.entrypoints.omni import Omni + + monkeypatch.setattr( + "vllm_omni.entrypoints.utils.load_stage_configs_from_model", + _fake_loader, + raising=False, + ) + + omni = Omni(model="any", init_timeout=1, dtype=dtype) + + # Dtype parsing being checked is on the diffusion path + assert len(omni.stage_configs) == 1 + assert omni.stage_configs[0]["stage_type"] == "diffusion" + + # Regardless of whether a str / dtype is passed, it resolves correctly + engine_args = omni.stage_configs[0]["engine_args"] + assert "dtype" in engine_args + assert isinstance(engine_args["dtype"], str) + assert engine_args["dtype"] == "float16" + + +def test_dtype_normalization_invalid_types(monkeypatch, mocker: MockerFixture, fake_stage_config): + """Ensure Diffusion Config builder correctly handles bad dtype overrides.""" + + def _fake_loader(model: str, base_engine_args=None): + # Return not stage configs to fall back to the diffusion cfg builder + return [] + + class NotATorchDtype: + pass + + import sys + + for module_name in [ + "vllm_omni.entrypoints.utils", + "vllm_omni.entrypoints.omni", + "vllm_omni.entrypoints.omni_stage", + ]: + if module_name in sys.modules: + del sys.modules[module_name] + + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) + _setup_ipc_mocks(monkeypatch) + _setup_log_mocks(monkeypatch) + + from vllm_omni.entrypoints.omni import Omni + + monkeypatch.setattr( + "vllm_omni.entrypoints.utils.load_stage_configs_from_model", + _fake_loader, + raising=False, + ) + + # Raise TypeError if we get an unrecognized type + with pytest.raises(TypeError): + Omni(model="any", init_timeout=1, dtype=NotATorchDtype) diff --git a/tests/model_executor/test_cosyvoice3_components.py b/tests/model_executor/test_cosyvoice3_components.py new file mode 100644 index 000000000..3b1471365 --- /dev/null +++ b/tests/model_executor/test_cosyvoice3_components.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for CosyVoice3 components.""" + +import pytest +import torch +import torch.nn as nn + + +class TestPreLookaheadLayer: + """Tests for PreLookaheadLayer.""" + + @pytest.fixture + def layer(self): + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer + + return PreLookaheadLayer(in_channels=512, channels=512, pre_lookahead_len=3) + + def test_forward_shape(self, layer): + """Test that output shape matches input shape.""" + batch, seq_len, channels = 2, 10, 512 + x = torch.randn(batch, seq_len, channels) + + out = layer(x) + + assert out.shape == x.shape + + def test_forward_with_context(self, layer): + """Test forward with context for streaming.""" + batch, seq_len, channels = 1, 10, 512 + x = torch.randn(batch, seq_len, channels) + context = torch.randn(batch, 3, channels) # pre_lookahead_len=3 + + layer.eval() + out = layer(x, context=context) + + assert out.shape == x.shape + + def test_residual_connection(self, layer): + """Test that residual connection is applied.""" + batch, seq_len, channels = 1, 5, 512 + x = torch.zeros(batch, seq_len, channels) + + # With zero input, output should also be close to zero due to residual + out = layer(x) + + # Output should be close to input (residual) plus conv output + assert out.shape == x.shape + + +class TestDiTAttention: + """Tests for DiTAttention with diffusion backend.""" + + @pytest.fixture + def attention(self): + from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiTAttention + + return DiTAttention(dim=512, heads=8, dim_head=64, dropout=0.0) + + def test_forward_shape(self, attention): + """Test attention output shape.""" + batch, seq_len, dim = 2, 16, 512 + x = torch.randn(batch, seq_len, dim) + + out = attention(x) + + assert out.shape == x.shape + + def test_forward_with_mask(self, attention): + """Test attention with mask.""" + batch, seq_len, dim = 2, 16, 512 + x = torch.randn(batch, seq_len, dim) + mask = torch.ones(batch, seq_len, dtype=torch.bool) + mask[:, -3:] = False # Mask last 3 positions + + out = attention(x, mask=mask) + + assert out.shape == x.shape + # Masked positions should be zero + assert torch.allclose(out[:, -3:], torch.zeros_like(out[:, -3:])) + + def test_qkv_projections(self, attention): + """Test that Q/K/V projections exist and have correct dimensions.""" + assert hasattr(attention, "to_q") + assert hasattr(attention, "to_k") + assert hasattr(attention, "to_v") + assert attention.to_q.out_features == 512 # heads * dim_head + assert attention.to_k.out_features == 512 + assert attention.to_v.out_features == 512 + + +class TestDiTBlock: + """Tests for DiTBlock.""" + + @pytest.fixture + def block(self): + from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiTBlock + + return DiTBlock(dim=512, heads=8, dim_head=64, ff_mult=4, dropout=0.0) + + def test_forward_shape(self, block): + """Test block output shape.""" + batch, seq_len, dim = 2, 16, 512 + x = torch.randn(batch, seq_len, dim) + t = torch.randn(batch, dim) # Timestep embedding + + out = block(x, t) + + assert out.shape == x.shape + + def test_adalayernorm_modulation(self, block): + """Test that AdaLayerNorm modulates based on timestep.""" + batch, seq_len, dim = 1, 8, 512 + x = torch.randn(batch, seq_len, dim) + t1 = torch.zeros(batch, dim) + t2 = torch.ones(batch, dim) + + out1 = block(x, t1) + out2 = block(x, t2) + + # Different timesteps should produce different outputs + assert not torch.allclose(out1, out2) + + +class TestDiT: + """Tests for the full DiT model.""" + + @pytest.fixture + def dit(self): + from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT + + return DiT( + dim=256, + depth=2, + heads=4, + dim_head=64, + dropout=0.0, + ff_mult=2, + mel_dim=80, + mu_dim=80, + spk_dim=80, + long_skip_connection=True, + ) + + def test_forward_shape(self, dit): + """Test DiT forward output shape.""" + batch, mel_dim, seq_len = 1, 80, 32 + x = torch.randn(batch, mel_dim, seq_len) + mask = torch.ones(batch, 1, seq_len) + mu = torch.randn(batch, mel_dim, seq_len) + t = torch.tensor([0.5]) + spks = torch.randn(batch, 80) + cond = torch.randn(batch, mel_dim, seq_len) + + out = dit(x, mask, mu, t, spks=spks, cond=cond) + + assert out.shape == (batch, mel_dim, seq_len) + + def test_timestep_embedding(self, dit): + """Test that different timesteps produce different outputs.""" + batch, mel_dim, seq_len = 1, 80, 16 + x = torch.randn(batch, mel_dim, seq_len) + mask = torch.ones(batch, 1, seq_len) + mu = torch.randn(batch, mel_dim, seq_len) + spks = torch.randn(batch, 80) + cond = torch.randn(batch, mel_dim, seq_len) + + out1 = dit(x, mask, mu, torch.tensor([0.0]), spks=spks, cond=cond) + out2 = dit(x, mask, mu, torch.tensor([1.0]), spks=spks, cond=cond) + + assert not torch.allclose(out1, out2) + + +class TestCFM: + """Tests for Conditional Flow Matching classes.""" + + @pytest.fixture + def dummy_estimator(self): + """Create a dummy estimator for testing.""" + + class DummyEstimator(nn.Module): + def __init__(self, mel_dim=80): + super().__init__() + self.mel_dim = mel_dim + + def forward(self, x, mask, mu, t, spks=None, cond=None): + return torch.zeros_like(x) + + return DummyEstimator() + + def test_causal_conditional_cfm_forward(self, dummy_estimator): + """Test CausalConditionalCFM forward pass.""" + from omegaconf import DictConfig + + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import CausalConditionalCFM + + cfm_params = DictConfig( + { + "sigma_min": 1e-6, + "solver": "euler", + "t_scheduler": "cosine", + "training_cfg_rate": 0.2, + "inference_cfg_rate": 0.7, + } + ) + + cfm = CausalConditionalCFM( + in_channels=80, + cfm_params=cfm_params, + n_spks=1, + spk_emb_dim=80, + estimator=dummy_estimator, + ) + + batch, mel_dim, seq_len = 1, 80, 32 + mu = torch.randn(batch, mel_dim, seq_len) + mask = torch.ones(batch, 1, seq_len) + spks = torch.randn(batch, 80) + cond = torch.randn(batch, mel_dim, seq_len) + + out, _ = cfm(mu, mask, n_timesteps=2, spks=spks, cond=cond) + + assert out.shape == mu.shape + + +class TestSDPAFallback: + """Test SDPA fallback for float32 inputs.""" + + def test_float32_uses_sdpa(self): + """Test that float32 inputs use SDPA fallback.""" + from vllm_omni.diffusion.attention.layer import Attention + + attn = Attention( + num_heads=8, + head_size=64, + causal=False, + softmax_scale=1.0 / 8.0, + ) + + batch, seq_len, heads, dim = 1, 16, 8, 64 + q = torch.randn(batch, seq_len, heads, dim, dtype=torch.float32) + k = torch.randn(batch, seq_len, heads, dim, dtype=torch.float32) + v = torch.randn(batch, seq_len, heads, dim, dtype=torch.float32) + + # Should not raise error - SDPA fallback handles float32 + out = attn(q, k, v) + + assert out.shape == (batch, seq_len, heads, dim) + assert out.dtype == torch.float32 diff --git a/tests/model_executor/test_cosyvoice3_utils.py b/tests/model_executor/test_cosyvoice3_utils.py new file mode 100644 index 000000000..828bb2b14 --- /dev/null +++ b/tests/model_executor/test_cosyvoice3_utils.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for CosyVoice3 utility functions.""" + +import pytest +import torch + + +class TestMakePadMask: + """Tests for make_pad_mask utility.""" + + def test_basic_mask_creation(self): + """Test basic padding mask creation.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + + lengths = torch.tensor([5, 3, 2]) + mask = make_pad_mask(lengths) + + # Expected: [[F,F,F,F,F], [F,F,F,T,T], [F,F,T,T,T]] + expected = torch.tensor( + [ + [False, False, False, False, False], + [False, False, False, True, True], + [False, False, True, True, True], + ] + ) + assert torch.equal(mask, expected) + + def test_single_sequence(self): + """Test mask for single sequence.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + + lengths = torch.tensor([3]) + mask = make_pad_mask(lengths) + + expected = torch.tensor([[False, False, False]]) + assert torch.equal(mask, expected) + + def test_explicit_max_len(self): + """Test with explicit max_len parameter.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + + lengths = torch.tensor([2, 3]) + mask = make_pad_mask(lengths, max_len=5) + + expected = torch.tensor( + [ + [False, False, True, True, True], + [False, False, False, True, True], + ] + ) + assert torch.equal(mask, expected) + + def test_zero_max_len_uses_max_length(self): + """Test that max_len=0 falls back to max(lengths).""" + from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + + lengths = torch.tensor([2, 4]) + mask = make_pad_mask(lengths, max_len=0) + + # Should use max(lengths) = 4 + assert mask.shape == (2, 4) + + def test_full_length_sequence(self): + """Test sequence that fills the entire length (no padding).""" + from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + + lengths = torch.tensor([5]) + mask = make_pad_mask(lengths, max_len=5) + + # No positions should be masked + assert not mask.any() + + def test_empty_sequence(self): + """Test sequence with zero length (all padding).""" + from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + + lengths = torch.tensor([0]) + mask = make_pad_mask(lengths, max_len=3) + + # All positions should be masked + assert mask.all() + + def test_device_preservation(self): + """Test that mask is created on the same device as input.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + + lengths = torch.tensor([3, 2]) + mask = make_pad_mask(lengths) + + assert mask.device == lengths.device + + +class TestDynamicRangeCompression: + """Tests for dynamic_range_compression_torch utility.""" + + def test_basic_compression(self): + """Test basic compression with default parameters.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import dynamic_range_compression_torch + + x = torch.tensor([1.0, 10.0, 100.0]) + result = dynamic_range_compression_torch(x) + + # Result should be log(x * 1) = log(x) + expected = torch.log(x) + assert torch.allclose(result, expected) + + def test_clipping_small_values(self): + """Test that small values are clipped.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import dynamic_range_compression_torch + + x = torch.tensor([0.0, 1e-10, 1e-6]) + clip_val = 1e-5 + result = dynamic_range_compression_torch(x, clip_val=clip_val) + + # All values below clip_val should be clipped + expected = torch.log(torch.tensor([clip_val, clip_val, clip_val])) + assert torch.allclose(result, expected) + + def test_scaling_factor(self): + """Test with different scaling factor c.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import dynamic_range_compression_torch + + x = torch.tensor([1.0, 10.0]) + c = 2.0 + result = dynamic_range_compression_torch(x, c=c) + + expected = torch.log(x * c) + assert torch.allclose(result, expected) + + +class TestSpectralNormalize: + """Tests for spectral_normalize_torch utility.""" + + def test_is_wrapper_for_compression(self): + """Test that spectral_normalize is a wrapper for dynamic_range_compression.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import ( + dynamic_range_compression_torch, + spectral_normalize_torch, + ) + + magnitudes = torch.tensor([1.0, 10.0, 100.0]) + result = spectral_normalize_torch(magnitudes) + expected = dynamic_range_compression_torch(magnitudes) + + assert torch.allclose(result, expected) + + +class TestExactDiv: + """Tests for exact_div utility.""" + + def test_exact_division(self): + """Test exact division with valid inputs.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import exact_div + + assert exact_div(10, 2) == 5 + assert exact_div(100, 25) == 4 + assert exact_div(0, 5) == 0 + + def test_non_exact_division_raises(self): + """Test that non-exact division raises AssertionError.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import exact_div + + with pytest.raises(AssertionError): + exact_div(10, 3) + + with pytest.raises(AssertionError): + exact_div(7, 2) + + +class TestConcatTextWithPromptIds: + """Tests for concat_text_with_prompt_ids utility.""" + + def test_basic_concat(self): + """Test basic concatenation.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import concat_text_with_prompt_ids + + text = torch.tensor([[5, 6, 7]], dtype=torch.int64) + text_len = torch.tensor([3], dtype=torch.int32) + prompt_text = torch.tensor([[1, 2]], dtype=torch.int64) + prompt_text_len = torch.tensor([2], dtype=torch.int32) + + result, result_len = concat_text_with_prompt_ids(text, text_len, prompt_text, prompt_text_len) + + assert result.tolist() == [[1, 2, 5, 6, 7]] + assert result_len.item() == 5 + + def test_empty_prompt(self): + """Test with empty prompt.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import concat_text_with_prompt_ids + + text = torch.tensor([[1, 2, 3]], dtype=torch.int64) + text_len = torch.tensor([3], dtype=torch.int32) + prompt_text = torch.tensor([[]], dtype=torch.int64).reshape(1, 0) + prompt_text_len = torch.tensor([0], dtype=torch.int32) + + result, result_len = concat_text_with_prompt_ids(text, text_len, prompt_text, prompt_text_len) + + assert result.tolist() == [[1, 2, 3]] + assert result_len.item() == 3 + + def test_batch_concat(self): + """Test batched concatenation.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import concat_text_with_prompt_ids + + text = torch.tensor([[5, 6], [7, 8]], dtype=torch.int64) + text_len = torch.tensor([2, 2], dtype=torch.int32) + prompt_text = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64) + prompt_text_len = torch.tensor([2, 2], dtype=torch.int32) + + result, result_len = concat_text_with_prompt_ids(text, text_len, prompt_text, prompt_text_len) + + assert result.tolist() == [[1, 2, 5, 6], [3, 4, 7, 8]] + assert result_len.tolist() == [4, 4] + + def test_prompt_comes_first(self): + """Test that prompt tokens come before text tokens.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import concat_text_with_prompt_ids + + text = torch.tensor([[100, 200]], dtype=torch.int64) + text_len = torch.tensor([2], dtype=torch.int32) + prompt_text = torch.tensor([[1, 2, 3]], dtype=torch.int64) + prompt_text_len = torch.tensor([3], dtype=torch.int32) + + result, _ = concat_text_with_prompt_ids(text, text_len, prompt_text, prompt_text_len) + + # First 3 tokens should be prompt + assert result[0, :3].tolist() == [1, 2, 3] + # Last 2 tokens should be text + assert result[0, 3:].tolist() == [100, 200] diff --git a/tests/perf/scripts/run_benchmark.py b/tests/perf/scripts/run_benchmark.py index 113c7e826..eb2ac5359 100644 --- a/tests/perf/scripts/run_benchmark.py +++ b/tests/perf/scripts/run_benchmark.py @@ -188,6 +188,10 @@ def benchmark_params(request, omni_server): if param_index >= len(all_params): raise ValueError(f"No benchmark parameters found for index {param_index} in test: {test_name}") + if all_params[param_index]["dataset_name"] == "random-mm": + # TODO: Due to known issues, skip the random-mm dataset. + pytest.skip("Skipping parameter for random-mm dataset.") + current = param_index + 1 total = len(all_params) print(f"\n Running benchmark {current}/{total} for {test_name}") diff --git a/tests/perf/tests/test.json b/tests/perf/tests/test.json index 32bbf1b5a..c2cc801b2 100644 --- a/tests/perf/tests/test.json +++ b/tests/perf/tests/test.json @@ -18,8 +18,8 @@ 4, 10 ], - "random_input_len": 2500, - "random_output_len": 900, + "random_input_len": 100, + "random_output_len": 100, "ignore_eos": true, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { @@ -102,8 +102,8 @@ 4, 10 ], - "random_input_len": 2500, - "random_output_len": 900, + "random_input_len": 100, + "random_output_len": 100, "ignore_eos": true, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py index 25ea24b19..c9d7fcf2c 100644 --- a/vllm_omni/__init__.py +++ b/vllm_omni/__init__.py @@ -20,6 +20,8 @@ except ModuleNotFoundError as exc: # pragma: no cover - optional dependency # Allow importing vllm_omni without vllm (e.g., documentation builds) patch = None # type: ignore +# Register custom configs (AutoConfig, AutoTokenizer) as early as possible. +from vllm_omni.transformers_utils import configs as _configs # noqa: F401, E402 from .config import OmniModelConfig from .entrypoints.async_omni import AsyncOmni diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index 9a644e1d9..c2389cc15 100644 --- a/vllm_omni/diffusion/layers/adalayernorm.py +++ b/vllm_omni/diffusion/layers/adalayernorm.py @@ -1,3 +1,5 @@ +from importlib.util import find_spec + import torch import torch.nn as nn from vllm.logger import init_logger @@ -6,6 +8,8 @@ from vllm_omni.diffusion.layers.custom_op import CustomOp logger = init_logger(__name__) +_HAS_MINDIESD = find_spec("mindiesd") is not None + class AdaLayerNorm(CustomOp): """ @@ -83,10 +87,21 @@ class AdaLayerNorm(CustomOp): ) -> torch.Tensor: shift_result, scale_result, gate_result = self.preprocess(mod_params, index) + if _HAS_MINDIESD: + try: + from mindiesd import layernorm_scale_shift + + output = layernorm_scale_shift(self.layernorm, x, scale_result, shift_result, fused=True) + + return output, gate_result + except ImportError as e: + logger.warning_once(f"mindiesd import failed, falling back to torch_npu: {e}") + import torch_npu - output = torch_npu.npu_layer_norm_eval( - x, normalized_shape=[self.hidden_size], weight=(1 + scale_result), bias=shift_result, eps=self.eps + output = ( + torch_npu.npu_layer_norm_eval(x, normalized_shape=[self.hidden_size], eps=self.eps) * (1 + scale_result) + + shift_result ) return output, gate_result diff --git a/vllm_omni/diffusion/models/cosyvoice3_audio/__init__.py b/vllm_omni/diffusion/models/cosyvoice3_audio/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/diffusion/models/cosyvoice3_audio/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py new file mode 100644 index 000000000..9c5af8f33 --- /dev/null +++ b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py @@ -0,0 +1,446 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from https://github.com/FunAudioLLM/CosyVoice/tree/main/cosyvoice/flow/DiT +# Refactored to use vllm_omni diffusion infrastructure for optimized attention backends + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F +from diffusers.models.normalization import AdaLayerNormZero +from einops import repeat +from torch import nn +from vllm.logger import init_logger +from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb + +from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention + +logger = init_logger(__name__) + +""" +in notation: +b - batch +n - sequence +nt - text sequence +nw - raw wave length +d - dimension +""" + + +def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0): + """Precompute rotary embedding frequencies.""" + theta *= theta_rescale_factor ** (dim / (dim - 2)) + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + t = torch.arange(end, device=freqs.device) + freqs = torch.outer(t, freqs).float() + freqs_cos = torch.cos(freqs) # real part + freqs_sin = torch.sin(freqs) # imaginary part + return torch.cat([freqs_cos, freqs_sin], dim=-1) + + +def get_pos_embed_indices(start, length, max_pos, scale=1.0): + """Get position embedding indices.""" + scale = scale * torch.ones_like(start, dtype=torch.float32) + pos = ( + start.unsqueeze(1) + + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long() + ) + pos = torch.where(pos < max_pos, pos, max_pos - 1) + return pos + + +class FeedForward(nn.Module): + """Feed-forward network with GELU activation.""" + + def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none"): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + activation = nn.GELU(approximate=approximate) + project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation) + self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)) + + def forward(self, x): + return self.ff(x) + + +class DiTAttention(nn.Module): + """Attention module using diffusion infrastructure for optimized backends. + + This replaces the original Attention class to leverage FlashAttention, + SageAttention, or SDPA backends automatically. + """ + + def __init__( + self, + dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + ): + super().__init__() + self.dim = dim + self.heads = heads + self.dim_head = dim_head + self.inner_dim = dim_head * heads + self.dropout = dropout + self.scale = 1.0 / math.sqrt(dim_head) + + # Q/K/V projections + self.to_q = nn.Linear(dim, self.inner_dim) + self.to_k = nn.Linear(dim, self.inner_dim) + self.to_v = nn.Linear(dim, self.inner_dim) + + # Output projection + self.to_out = nn.Sequential(nn.Linear(self.inner_dim, dim), nn.Dropout(dropout)) + + # Diffusion attention backend (Flash/Sage/SDPA) + self.attn = DiffusionAttention( + num_heads=heads, + head_size=dim_head, + softmax_scale=self.scale, + causal=False, + ) + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + rope=None, + ) -> torch.Tensor: + batch_size, seq_len = x.shape[0], x.shape[1] + + # Project to Q, K, V + query = self.to_q(x) + key = self.to_k(x) + value = self.to_v(x) + + # Apply rotary position embedding + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + # Reshape for attention: (batch, seq, heads, head_dim) + query = query.view(batch_size, seq_len, self.heads, self.dim_head) + key = key.view(batch_size, seq_len, self.heads, self.dim_head) + value = value.view(batch_size, seq_len, self.heads, self.dim_head) + + # Use diffusion attention backend + # The diffusion Attention layer expects (batch, seq, heads, head_dim) + out = self.attn(query, key, value, attn_metadata=None) + + # Reshape back: (batch, seq, dim) + out = out.view(batch_size, seq_len, self.inner_dim) + out = out.to(query.dtype) + + # Output projection + out = self.to_out(out) + + # Apply mask if provided + if mask is not None: + if mask.dim() == 2: + mask = mask.unsqueeze(-1) + elif mask.dim() == 4: + # (batch, heads, seq, seq) -> use last dim + mask = mask[:, 0, -1].unsqueeze(-1) + out = out.masked_fill(~mask.bool(), 0.0) + + return out + + +class DiTBlock(nn.Module): + """DiT block with AdaLayerNorm modulation.""" + + def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1): + super().__init__() + + self.attn_norm = AdaLayerNormZero(dim) + self.attn = DiTAttention( + dim=dim, + heads=heads, + dim_head=dim_head, + dropout=dropout, + ) + + self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh") + + def forward(self, x, t, mask=None, rope=None): + # pre-norm & modulation for attention input + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) + + # attention + attn_output = self.attn(x=norm, mask=mask, rope=rope) + + # process attention output for input x + x = x + gate_msa.unsqueeze(1) * attn_output + + ff_norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ff_output = self.ff(ff_norm) + x = x + gate_mlp.unsqueeze(1) * ff_output + + return x + + +class CausalConvPositionEmbedding(nn.Module): + """Causal convolutional position embedding.""" + + def __init__(self, dim, kernel_size=31, groups=16): + super().__init__() + assert kernel_size % 2 != 0 + self.kernel_size = kernel_size + self.conv1 = nn.Sequential( + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=0), + nn.Mish(), + ) + self.conv2 = nn.Sequential( + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=0), + nn.Mish(), + ) + + def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None): + if mask is not None: + mask = mask[..., None] + x = x.masked_fill(~mask, 0.0) + + x = x.permute(0, 2, 1) + x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) + x = self.conv1(x) + x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) + x = self.conv2(x) + out = x.permute(0, 2, 1) + + if mask is not None: + out = out.masked_fill(~mask, 0.0) + + return out + + +class AdaLayerNormZero_Final(nn.Module): + """AdaLayerNormZero for final layer - returns only modulated x.""" + + def __init__(self, dim): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 2) + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x, emb): + emb = self.linear(self.silu(emb)) + scale, shift = torch.chunk(emb, 2, dim=1) + x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return x + + +class GRN(nn.Module): + """Global Response Normalization layer.""" + + def __init__(self, dim): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, dim)) + + def forward(self, x): + Gx = torch.norm(x, p=2, dim=1, keepdim=True) + Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) + return self.gamma * (x * Nx) + self.beta + x + + +class ConvNeXtV2Block(nn.Module): + """ConvNeXt-V2 Block.""" + + def __init__(self, dim: int, intermediate_dim: int, dilation: int = 1): + super().__init__() + padding = (dilation * (7 - 1)) // 2 + self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation) + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, intermediate_dim) + self.act = nn.GELU() + self.grn = GRN(intermediate_dim) + self.pwconv2 = nn.Linear(intermediate_dim, dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = x.transpose(1, 2) + x = self.dwconv(x) + x = x.transpose(1, 2) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + return residual + x + + +class SinusPositionEmbedding(nn.Module): + """Sinusoidal position embedding.""" + + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x, scale=1000): + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) + emb = scale * x.unsqueeze(1) * emb.unsqueeze(0) + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + +class TimestepEmbedding(nn.Module): + """Timestep embedding with MLP.""" + + def __init__(self, dim, freq_embed_dim=256): + super().__init__() + self.time_embed = SinusPositionEmbedding(freq_embed_dim) + self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + + def forward(self, timestep: torch.Tensor): + time_hidden = self.time_embed(timestep) + time_hidden = time_hidden.to(timestep.dtype) + time = self.time_mlp(time_hidden) + return time + + +class TextEmbedding(nn.Module): + """Text embedding with optional ConvNeXt modeling.""" + + def __init__(self, text_num_embeds, text_dim, conv_layers=0, conv_mult=2): + super().__init__() + self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim) + + if conv_layers > 0: + self.extra_modeling = True + self.precompute_max_pos = 4096 + self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False) + self.text_blocks = nn.Sequential( + *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)] + ) + else: + self.extra_modeling = False + + def forward(self, text: torch.Tensor, seq_len, drop_text=False): + batch, text_len = text.shape[0], text.shape[1] + text = text + 1 + text = text[:, :seq_len] + text = F.pad(text, (0, seq_len - text_len), value=0) + + if drop_text: + text = torch.zeros_like(text) + + text = self.text_embed(text) + + if self.extra_modeling: + batch_start = torch.zeros((batch,), dtype=torch.long) + pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos) + text_pos_embed = self.freqs_cis[pos_idx] + text = text + text_pos_embed + text = self.text_blocks(text) + + return text + + +class InputEmbedding(nn.Module): + """Input embedding combining noised audio, condition, text, and speaker.""" + + def __init__(self, mel_dim, text_dim, out_dim, spk_dim=None): + super().__init__() + spk_dim = 0 if spk_dim is None else spk_dim + self.spk_dim = spk_dim + self.proj = nn.Linear(mel_dim * 2 + text_dim + spk_dim, out_dim) + self.conv_pos_embed = CausalConvPositionEmbedding(dim=out_dim) + + def forward(self, x, cond, text_embed, spks): + to_cat = [x, cond, text_embed] + if self.spk_dim > 0: + spks = repeat(spks, "b c -> b t c", t=x.shape[1]) + to_cat.append(spks) + + x = self.proj(torch.cat(to_cat, dim=-1)) + x = self.conv_pos_embed(x) + x + return x + + +class DiT(nn.Module): + """Diffusion Transformer backbone using optimized attention backends. + + This is a drop-in replacement for the original DiT that uses the + vllm_omni diffusion infrastructure for FlashAttention/SageAttention/SDPA. + """ + + def __init__( + self, + *, + dim, + depth=8, + heads=8, + dim_head=64, + dropout=0.1, + ff_mult=4, + mel_dim=80, + mu_dim=None, + long_skip_connection=False, + spk_dim=None, + out_channels=None, + static_chunk_size=50, + num_decoding_left_chunks=2, + ): + super().__init__() + + self.time_embed = TimestepEmbedding(dim) + if mu_dim is None: + mu_dim = mel_dim + self.input_embed = InputEmbedding(mel_dim, mu_dim, dim, spk_dim) + + self.rotary_embed = RotaryEmbedding(dim_head) + + self.dim = dim + self.depth = depth + + self.transformer_blocks = nn.ModuleList( + [DiTBlock(dim=dim, heads=heads, dim_head=dim_head, ff_mult=ff_mult, dropout=dropout) for _ in range(depth)] + ) + self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None + + self.norm_out = AdaLayerNormZero_Final(dim) + self.proj_out = nn.Linear(dim, mel_dim) + self.out_channels = out_channels + self.static_chunk_size = static_chunk_size + self.num_decoding_left_chunks = num_decoding_left_chunks + + def forward(self, x, mask, mu, t, spks=None, cond=None): + x = x.transpose(1, 2) + mu = mu.transpose(1, 2) + cond = cond.transpose(1, 2) + spks = spks.unsqueeze(dim=1) + batch, seq_len = x.shape[0], x.shape[1] + if t.ndim == 0: + t = t.repeat(batch) + + # t: conditioning time, c: context (text + masked cond audio), x: noised input audio + t = self.time_embed(t) + x = self.input_embed(x, cond, mu, spks.squeeze(1)) + + rope = self.rotary_embed.forward_from_seq_len(seq_len) + + if self.long_skip_connection is not None: + residual = x + + attn_mask = mask.bool().repeat(1, x.size(1), 1).unsqueeze(dim=1) + + for block in self.transformer_blocks: + x = block(x, t, mask=attn_mask.bool(), rope=rope) + + if self.long_skip_connection is not None: + x = self.long_skip_connection(torch.cat((x, residual), dim=-1)) + + x = self.norm_out(x, t) + output = self.proj_out(x).transpose(1, 2) + return output diff --git a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py new file mode 100644 index 000000000..ed42e24ca --- /dev/null +++ b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py @@ -0,0 +1,809 @@ +import math + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.attention_processor import Attention +from diffusers.models.embeddings import TimestepEmbedding, Timesteps +from diffusers.models.modeling_utils import ModelMixin +from einops import rearrange, repeat +from torch import nn +from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm + +from .rope_real import RotaryPosEmbedReal + +try: + from transformers.modeling_flash_attention_utils import ( # type: ignore + flash_attn_varlen_func, # pyright: ignore[reportAttributeAccessIssue] + is_flash_attn_available, + ) +except Exception: # pragma: no cover - best-effort compatibility + flash_attn_varlen_func = None # type: ignore[assignment] + + def is_flash_attn_available() -> bool: # type: ignore[override] + return False + + +from .rope_real import apply_real_rotary_emb + +_HAS_FLASH_ATTN_VARLEN = bool(is_flash_attn_available()) and flash_attn_varlen_func is not None + + +class LuminaRMSNormZero(nn.Module): + """ + Norm layer adaptive RMS normalization zero. + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + """ + + def __init__( + self, + embedding_dim: int, + norm_eps: float, + norm_elementwise_affine: bool, + ): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear( + min(embedding_dim, 1024), + 4 * embedding_dim, + bias=True, + ) + + self.norm = Qwen2RMSNorm(embedding_dim, eps=norm_eps) + + def forward( + self, + x: torch.Tensor, + emb: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + emb = self.linear(self.silu(emb)) + scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1) + x = self.norm(x) * (1 + scale_msa[:, None]) + return x, gate_msa, scale_mlp, gate_mlp + + +class LuminaFeedForward(nn.Module): + def __init__( + self, + dim: int, + inner_dim: int, + multiple_of: int | None = 256, + ffn_dim_multiplier: float | None = None, + ): + super().__init__() + + # custom hidden_size factor multiplier + if ffn_dim_multiplier is not None: + inner_dim = int(ffn_dim_multiplier * inner_dim) + inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of) + + self.linear_1 = nn.Linear( + dim, + inner_dim, + bias=False, + ) + self.linear_2 = nn.Linear( + inner_dim, + dim, + bias=False, + ) + self.linear_3 = nn.Linear( + dim, + inner_dim, + bias=False, + ) + + def swiglu(self, x, y): + return F.silu(x.float(), inplace=False).to(x.dtype) * y + + def forward(self, x): + h1, h2 = self.linear_1(x), self.linear_3(x) + return self.linear_2(self.swiglu(h1, h2)) + + +class LuminaLayerNormContinuous(nn.Module): + def __init__( + self, + embedding_dim: int, + conditioning_embedding_dim: int, + # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters + # because the output is immediately scaled and shifted by the projected conditioning embeddings. + # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters. + # However, this is how it was implemented in the original code, and it's rather likely you should + # set `elementwise_affine` to False. + elementwise_affine=True, + eps=1e-5, + bias=True, + norm_type="layer_norm", + out_dim: int | None = None, + ): + super().__init__() + + self.silu = nn.SiLU() + self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias) + + if norm_type == "layer_norm": + self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias) + elif norm_type == "rms_norm": + self.norm = Qwen2RMSNorm(embedding_dim, eps=eps) + else: + raise ValueError(f"unknown norm_type {norm_type}") + + self.linear_2 = None + if out_dim is not None: + self.linear_2 = nn.Linear(embedding_dim, out_dim, bias=bias) + + def forward( + self, + x: torch.Tensor, + conditioning_embedding: torch.Tensor, + ) -> torch.Tensor: + scale = self.linear_1(self.silu(conditioning_embedding).to(x.dtype)) + x = self.norm(x) * (1 + scale)[:, None, :] + + if self.linear_2 is not None: + x = self.linear_2(x) + + return x + + +class Lumina2CombinedTimestepCaptionEmbedding(nn.Module): + def __init__( + self, + hidden_size: int = 2520, + text_feat_dim: int = 3584, + frequency_embedding_size: int = 256, + norm_eps: float = 1e-5, + timestep_scale: float = 1.0, + ) -> None: + super().__init__() + + self.time_proj = Timesteps( + num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0.0, scale=timestep_scale + ) + + self.timestep_embedder = TimestepEmbedding( + in_channels=frequency_embedding_size, time_embed_dim=min(hidden_size, 1024) + ) + + self.caption_embedder = nn.Sequential( + Qwen2RMSNorm(text_feat_dim, eps=norm_eps), + nn.Linear(text_feat_dim, hidden_size, bias=True), + ) + + def forward( + self, + timestep: torch.Tensor, + text_hidden_states: torch.Tensor, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + timestep_proj = self.time_proj(timestep).to(dtype=dtype) + time_embed = self.timestep_embedder(timestep_proj) + caption_embed = self.caption_embedder(text_hidden_states) + return time_embed, caption_embed + + +class SimpleQFormerImageRefiner(nn.Module): + def __init__( + self, + hidden_size: int, + num_queries: int = 128, + num_layers: int = 2, + num_heads: int | None = None, + dropout: float = 0.0, + norm_eps: float = 1e-5, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.num_queries = num_queries + # ensure num_heads divides hidden_size + if num_heads is None: + num_heads = max(1, hidden_size // 128) + self.num_heads = self._choose_valid_num_heads(hidden_size, num_heads) + self.input_proj = nn.Sequential( + Qwen2RMSNorm(hidden_size, eps=norm_eps), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + + # Learnable query embeddings + scale = hidden_size**-0.5 + self.query = nn.Parameter(scale * torch.randn(1, num_queries, hidden_size)) + + # Decoder layers + self.layers = nn.ModuleList() + for _ in range(num_layers): + self.layers.append( + nn.ModuleDict( + dict( + ln_q1=Qwen2RMSNorm(hidden_size, eps=norm_eps), + self_attn=nn.MultiheadAttention( + embed_dim=hidden_size, num_heads=self.num_heads, dropout=dropout, batch_first=True + ), + ln_q2=Qwen2RMSNorm(hidden_size, eps=norm_eps), + cross_attn=nn.MultiheadAttention( + embed_dim=hidden_size, num_heads=self.num_heads, dropout=dropout, batch_first=True + ), + ln_ffn=Qwen2RMSNorm(hidden_size, eps=norm_eps), + ffn=LuminaFeedForward(dim=hidden_size, inner_dim=4 * hidden_size), + ) + ) + ) + + @staticmethod + def _choose_valid_num_heads(hidden_size: int, proposed_heads: int, preferred_head_dim: int = 128) -> int: + """Pick a number of heads that divides hidden_size, close to proposed or preferred.""" + # If proposed is valid, use it + if proposed_heads > 0 and hidden_size % proposed_heads == 0: + return proposed_heads + # target based on preferred head dim + target = max(1, round(hidden_size / preferred_head_dim)) + # collect divisors up to 128 heads (more than enough) + max_heads_cap = min(128, hidden_size) + divisors = [d for d in range(1, max_heads_cap + 1) if hidden_size % d == 0] + # choose closest to target + best = min(divisors, key=lambda d: (abs(d - target), -d)) + return best + + def forward(self, x: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + """ + Args: + x: Tensor of shape (batch, seq_len, input_dim) + Returns: + Tensor of shape (batch, num_queries, hidden_size) + """ + batch, _, _ = x.shape + kv = self.input_proj(x) + q = self.query.repeat(batch, 1, 1).to(kv.dtype) + + for layer in self.layers: + # Self-attention on queries + q_norm = layer["ln_q1"](q) + attn_out, _ = layer["self_attn"](q_norm, q_norm, q_norm, need_weights=False) + q = q + attn_out + + # Cross-attention: queries attend to inputs + q_norm = layer["ln_q2"](q) + cross_out, _ = layer["cross_attn"](q_norm, kv, kv, need_weights=False, key_padding_mask=attention_mask) + q = q + cross_out + + # Feed-forward + q = q + layer["ffn"](layer["ln_ffn"](q)) + + return q + + +class AttnProcessor: + def __init__(self) -> None: + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("AttnProcessor requires PyTorch 2.0+ (F.scaled_dot_product_attention).") + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, + base_sequence_length: int | None = None, + ) -> torch.Tensor: + batch_size, sequence_length, _ = hidden_states.shape + + # Get Query-Key-Value Pair + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + query_dim = query.shape[-1] + inner_dim = key.shape[-1] + head_dim = query_dim // attn.heads + dtype = query.dtype + + # Get key-value heads + kv_heads = inner_dim // head_dim + + # Reshape tensors for attention computation + query = query.view(batch_size, -1, attn.heads, head_dim) + key = key.view(batch_size, -1, kv_heads, head_dim) + value = value.view(batch_size, -1, kv_heads, head_dim) + + # Apply Query-Key normalization + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + # Apply Rotary Position Embeddings + if image_rotary_emb is not None: + query = apply_real_rotary_emb(query, image_rotary_emb[0], image_rotary_emb[1]) + key = apply_real_rotary_emb(key, image_rotary_emb[0], image_rotary_emb[1]) + + query, key = query.to(dtype), key.to(dtype) + + # Calculate attention scale + if base_sequence_length is not None: + softmax_scale = math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale + else: + softmax_scale = attn.scale + + if _HAS_FLASH_ATTN_VARLEN and attention_mask is not None and hidden_states.is_cuda: + # Flash-Attn varlen expects packed tokens + cu_seqlens. Here we only need + # the self-attention case (q/k/v share the same padding mask). + attention_mask = attention_mask.to(torch.bool) + seqlens = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen = int(seqlens.max().item()) + cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)) + + query_states = query.reshape(batch_size * sequence_length, attn.heads, head_dim)[indices] + key_states = key.reshape(batch_size * sequence_length, kv_heads, head_dim)[indices] + value_states = value.reshape(batch_size * sequence_length, kv_heads, head_dim)[indices] + + if kv_heads < attn.heads: + key_states = repeat(key_states, "l h c -> l (h k) c", k=attn.heads // kv_heads) + value_states = repeat(value_states, "l h c -> l (h k) c", k=attn.heads // kv_heads) + + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + dropout_p=0.0, + causal=False, + softmax_scale=softmax_scale, + ) + + out = torch.zeros( + (batch_size * sequence_length, attn.heads, head_dim), + device=attn_output_unpad.device, + dtype=attn_output_unpad.dtype, + ) + out[indices] = attn_output_unpad + hidden_states = out.view(batch_size, sequence_length, attn.heads, head_dim).flatten(-2) + hidden_states = hidden_states.type_as(query) + else: + # PyTorch SDPA path. + attn_mask = None + if attention_mask is not None: + attention_mask = attention_mask.to(torch.bool) + attn_mask = attention_mask.view(batch_size, 1, 1, -1) + + query = query.transpose(1, 2) # [B, H, S, D] + key = key.transpose(1, 2) # [B, H_kv, S, D] + value = value.transpose(1, 2) + + if kv_heads < attn.heads: + key = key.repeat_interleave(attn.heads // kv_heads, dim=1) + value = value.repeat_interleave(attn.heads // kv_heads, dim=1) + + hidden_states = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attn_mask, + dropout_p=0.0, + is_causal=False, + scale=softmax_scale, + ) + + if attention_mask is not None: + # Keep padding tokens consistent with the flash-varlen path (zero output). + hidden_states = hidden_states * attention_mask[:, None, :, None] + + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states = hidden_states.type_as(query) + + # Apply output projection + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + + return hidden_states + + +class TransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + num_kv_heads: int, + multiple_of: int, + ffn_dim_multiplier: float, + norm_eps: float, + modulation: bool = True, + ) -> None: + """Initialize the transformer block.""" + super().__init__() + self.head_dim = dim // num_attention_heads + self.modulation = modulation + + processor = AttnProcessor() + + # Initialize attention layer + self.attn = Attention( + query_dim=dim, + cross_attention_dim=None, + dim_head=dim // num_attention_heads, + qk_norm=None, + heads=num_attention_heads, + kv_heads=num_kv_heads, + eps=1e-5, + bias=False, + out_bias=False, + processor=processor, + ) + # 显式使用 transformers 的 Qwen2RMSNorm，避免依赖 diffusers 内部创建的 `RMSNorm` 再做递归替换。 + self.attn.norm_q = Qwen2RMSNorm(self.head_dim, eps=1e-5) + self.attn.norm_k = Qwen2RMSNorm(self.head_dim, eps=1e-5) + + # Initialize feed-forward network + self.feed_forward = LuminaFeedForward( + dim=dim, inner_dim=4 * dim, multiple_of=multiple_of, ffn_dim_multiplier=ffn_dim_multiplier + ) + + # Initialize normalization layers + if modulation: + self.norm1 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True) + else: + self.norm1 = Qwen2RMSNorm(dim, eps=norm_eps) + + self.ffn_norm1 = Qwen2RMSNorm(dim, eps=norm_eps) + self.norm2 = Qwen2RMSNorm(dim, eps=norm_eps) + self.ffn_norm2 = Qwen2RMSNorm(dim, eps=norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + image_rotary_emb: torch.Tensor, + temb: torch.Tensor | None = None, + ) -> torch.Tensor: + if self.modulation: + if temb is None: + raise ValueError("temb must be provided when modulation is enabled") + + norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb) + attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + hidden_states = hidden_states + gate_msa.unsqueeze(1).tanh() * self.norm2(attn_output) + mlp_output = self.feed_forward(self.ffn_norm1(hidden_states) * (1 + scale_mlp.unsqueeze(1))) + hidden_states = hidden_states + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(mlp_output) + else: + norm_hidden_states = self.norm1(hidden_states) + attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + hidden_states = hidden_states + self.norm2(attn_output) + mlp_output = self.feed_forward(self.ffn_norm1(hidden_states)) + hidden_states = hidden_states + self.ffn_norm2(mlp_output) + + return hidden_states + + +class Transformer2DModel(ModelMixin, ConfigMixin): + """MammothModa2 DiT transformer""" + + @register_to_config + def __init__( + self, + patch_size: int = 2, + in_channels: int = 16, + out_channels: int | None = None, + hidden_size: int = 2304, + num_layers: int = 26, + num_refiner_layers: int = 2, + num_attention_heads: int = 24, + num_kv_heads: int = 8, + multiple_of: int = 256, + ffn_dim_multiplier: float | None = None, + norm_eps: float = 1e-5, + axes_dim_rope: tuple[int, int, int] = (32, 32, 32), + axes_lens: tuple[int, int, int] = (300, 512, 512), + text_feat_dim: int = 1024, + timestep_scale: float = 1.0, + ) -> None: + """Initialize the transformer model.""" + super().__init__() + self.hidden_size = hidden_size + + # Validate configuration + if (hidden_size // num_attention_heads) != sum(axes_dim_rope): + raise ValueError( + f"hidden_size // num_attention_heads ({hidden_size // num_attention_heads}) " + f"must equal sum(axes_dim_rope) ({sum(axes_dim_rope)})" + ) + + self.out_channels = out_channels or in_channels + # Initialize embeddings + self.rope_embedder = RotaryPosEmbedReal( + theta=10000, + axes_dim=axes_dim_rope, + axes_lens=axes_lens, + patch_size=patch_size, + ) + + self.x_embedder = nn.Linear( + in_features=patch_size * patch_size * in_channels, + out_features=hidden_size, + ) + + self.ref_image_patch_embedder = nn.Linear( + in_features=patch_size * patch_size * in_channels, + out_features=hidden_size, + ) + + self.time_caption_embed = Lumina2CombinedTimestepCaptionEmbedding( + hidden_size=hidden_size, + text_feat_dim=text_feat_dim, + norm_eps=norm_eps, + timestep_scale=timestep_scale, + ) + + # Initialize transformer blocks + self.noise_refiner = nn.ModuleList( + [ + TransformerBlock( + hidden_size, + num_attention_heads, + num_kv_heads, + multiple_of, + ffn_dim_multiplier, + norm_eps, + modulation=True, + ) + for _ in range(num_refiner_layers) + ] + ) + + self.ref_image_refiner = nn.ModuleList( + [ + TransformerBlock( + hidden_size, + num_attention_heads, + num_kv_heads, + multiple_of, + ffn_dim_multiplier, + norm_eps, + modulation=True, + ) + for _ in range(num_refiner_layers) + ] + ) + + self.context_refiner = nn.ModuleList( + [ + TransformerBlock( + hidden_size, + num_attention_heads, + num_kv_heads, + multiple_of, + ffn_dim_multiplier, + norm_eps, + modulation=False, + ) + for _ in range(num_refiner_layers) + ] + ) + + # 3. Transformer blocks + self.layers = nn.ModuleList( + [ + TransformerBlock( + hidden_size, + num_attention_heads, + num_kv_heads, + multiple_of, + ffn_dim_multiplier, + norm_eps, + modulation=True, + ) + for _ in range(num_layers) + ] + ) + + # 4. Output norm & projection + self.norm_out = LuminaLayerNormContinuous( + embedding_dim=hidden_size, + conditioning_embedding_dim=min(hidden_size, 1024), + elementwise_affine=False, + eps=1e-6, + bias=True, + out_dim=patch_size * patch_size * self.out_channels, + ) + + # Add learnable embeddings to distinguish different images + self.image_index_embedding = nn.Parameter(torch.randn(5, hidden_size)) # support max 5 ref images + + def _validate_inputs( + self, + hidden_states: torch.Tensor, + text_hidden_states: torch.Tensor, + text_attention_mask: torch.Tensor, + ref_image_hidden_states: list[list[torch.Tensor]] | None, + return_dict: bool, + ) -> tuple[int, int, int]: + if return_dict: + raise ValueError("return_dict=True is not supported in vLLM inference.") + if ref_image_hidden_states is not None: + raise ValueError("ref_image_hidden_states is not supported in vLLM inference.") + if hidden_states.ndim != 4: + raise ValueError(f"Expected hidden_states to be 4D [B,C,H,W], got shape={tuple(hidden_states.shape)}") + + batch_size, _channels, height, width = hidden_states.shape + if batch_size != text_hidden_states.shape[0] or batch_size != text_attention_mask.shape[0]: + raise ValueError( + "Batch size mismatch: " + f"hidden_states={batch_size}, text_hidden_states={text_hidden_states.shape[0]}, " + f"text_attention_mask={text_attention_mask.shape[0]}" + ) + + p = self.config.patch_size + if height % p != 0 or width % p != 0: + raise ValueError(f"Input latent H/W must be divisible by patch_size={p}, got {height}x{width}") + return batch_size, height, width + + def _prepare_embeddings( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + text_hidden_states: torch.Tensor, + text_attention_mask: torch.Tensor, + freqs_cis: torch.Tensor, + batch_size: int, + height: int, + width: int, + ): + device = hidden_states.device + p = self.config.patch_size + + temb, text_hidden_states = self.time_caption_embed(timestep, text_hidden_states, hidden_states.dtype) + + img_tokens = rearrange(hidden_states, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=p, p2=p) + img_tokens = self.x_embedder(img_tokens) + + img_len = (height // p) * (width // p) + img_mask = torch.ones((batch_size, img_len), dtype=torch.bool, device=device) + l_effective_img_len = [img_len for _ in range(batch_size)] + img_sizes = [(height, width) for _ in range(batch_size)] + + l_effective_ref_img_len = [[] for _ in range(batch_size)] + ref_img_sizes = [None for _ in range(batch_size)] + + ( + context_rotary_emb, + _ref_img_rotary_emb, + noise_rotary_emb, + rotary_emb, + encoder_seq_lengths, + seq_lengths, + ) = self.rope_embedder( + freqs_cis, + text_attention_mask, + l_effective_ref_img_len, + l_effective_img_len, + ref_img_sizes, + img_sizes, + device, + ) + + return ( + temb, + text_hidden_states, + img_tokens, + img_mask, + img_len, + context_rotary_emb, + noise_rotary_emb, + rotary_emb, + encoder_seq_lengths, + seq_lengths, + ) + + def _apply_refiners( + self, + text_hidden_states: torch.Tensor, + text_attention_mask: torch.Tensor, + context_rotary_emb: torch.Tensor, + img_tokens: torch.Tensor, + img_mask: torch.Tensor, + noise_rotary_emb: torch.Tensor, + temb: torch.Tensor, + ): + for layer in self.context_refiner: + text_hidden_states = layer(text_hidden_states, text_attention_mask, context_rotary_emb) + + for layer in self.noise_refiner: + img_tokens = layer(img_tokens, img_mask, noise_rotary_emb, temb) + + return text_hidden_states, img_tokens + + def _apply_transformer_layers(self, hidden_states, attention_mask, rotary_emb, temb): + for layer in self.layers: + hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb) + return hidden_states + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + text_hidden_states: torch.Tensor, + freqs_cis: torch.Tensor, + text_attention_mask: torch.Tensor, + ref_image_hidden_states: list[list[torch.Tensor]] | None = None, + return_dict: bool = False, + ) -> torch.Tensor: + batch_size, height, width = self._validate_inputs( + hidden_states, text_hidden_states, text_attention_mask, ref_image_hidden_states, return_dict + ) + + ( + temb, + text_hidden_states, + img_tokens, + img_mask, + img_len, + context_rotary_emb, + noise_rotary_emb, + rotary_emb, + encoder_seq_lengths, + seq_lengths, + ) = self._prepare_embeddings( + hidden_states, + timestep, + text_hidden_states, + text_attention_mask, + freqs_cis, + batch_size, + height, + width, + ) + + text_hidden_states, img_tokens = self._apply_refiners( + text_hidden_states, + text_attention_mask, + context_rotary_emb, + img_tokens, + img_mask, + noise_rotary_emb, + temb, + ) + + max_seq_len = max(seq_lengths) + attention_mask = hidden_states.new_zeros(batch_size, max_seq_len, dtype=torch.bool) + joint_hidden_states = hidden_states.new_zeros(batch_size, max_seq_len, self.config.hidden_size) + for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)): + attention_mask[i, :seq_len] = True + joint_hidden_states[i, :encoder_seq_len] = text_hidden_states[i, :encoder_seq_len] + joint_hidden_states[i, encoder_seq_len : encoder_seq_len + img_len] = img_tokens[i, :img_len] + + hidden_states = self._apply_transformer_layers(joint_hidden_states, attention_mask, rotary_emb, temb) + + hidden_states = self.norm_out(hidden_states, temb) + + p = self.config.patch_size + img_hidden_states = torch.stack( + [ + hidden_states[i, encoder_seq_len : encoder_seq_len + img_len] + for i, encoder_seq_len in enumerate(encoder_seq_lengths) + ], + dim=0, + ) + output = rearrange( + img_hidden_states, + "b (h w) (p1 p2 c) -> b c (h p1) (w p2)", + h=height // p, + w=width // p, + p1=p, + p2=p, + ) + return output diff --git a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py new file mode 100644 index 000000000..e55338963 --- /dev/null +++ b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py @@ -0,0 +1,298 @@ +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import torch +from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm +from vllm.config import VllmConfig +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config + +from .mammothmoda2_dit_model import SimpleQFormerImageRefiner, Transformer2DModel +from .rope_real import RotaryPosEmbedReal +from .schedulers import FlowMatchEulerDiscreteScheduler + + +class MammothModa2DiTPipeline(nn.Module): + """ + MammothModa2 DiT + VAE generation stage (non-autoregressive). + + This stage expects "image condition token hidden states" from the upstream AR stage, + and outputs image tensors via diffusion transformer + VAE decode. + + """ + + have_multimodal_outputs = True + + # Load only gen_* weights; ignore llm_model.* to prevent loading the entire LLM backbone in the DiT stage. + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "llm_model.": None, + } + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + del prefix + + hf_config = vllm_config.model_config.hf_config + if not isinstance(hf_config, Mammothmoda2Config): + raise TypeError(f"Expected Mammothmoda2Config, got {type(hf_config)}") + + self.config = hf_config + + # --- Build DiT / VAE modules (names must match checkpoint keys) --- + if self.config.gen_vae_config is None or self.config.gen_dit_config is None: + raise ValueError("Mammothmoda2Config.gen_vae_config / gen_dit_config must not be None") + + self.gen_vae = AutoencoderKL.from_config(self.config.gen_vae_config) + self.gen_transformer = Transformer2DModel.from_config(self.config.gen_dit_config) + + llm_hidden_size = int(getattr(self.config.llm_config, "hidden_size", 0) or 0) + if llm_hidden_size <= 0: + raise ValueError("Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.hidden_size") + self._reinit_caption_embedder(llm_hidden_size) + + # Optional: image condition refiner (Q-Former) + if self.config.gen_image_condition_refiner_config is not None: + self.gen_image_condition_refiner = SimpleQFormerImageRefiner( + hidden_size=llm_hidden_size, + **self.config.gen_image_condition_refiner_config, + ) + else: + self.gen_image_condition_refiner = None + + # Precompute rotary freqs for diffusion transformer + # IMPORTANT: follow upstream mammothmoda: use top-level `config.gen_axes_*` + # (the checkpoint's `gen_dit_config.axes_lens` can be as small as 1024, + # which is insufficient for vLLM dummy-run/cudagraph warmup). + self.gen_freqs_cis = RotaryPosEmbedReal.get_freqs_real( + tuple(self.config.gen_axes_dim_rope), + tuple(self.config.gen_axes_lens), + theta=10000, + ) + + # vLLM PP interface compatibility + self.make_empty_intermediate_tensors = lambda: None + + self._llm_hidden_size = llm_hidden_size + + def _reinit_caption_embedder(self, in_features: int) -> None: + # Align with upstream Mammothmoda2Model's `reinit_caption_embedder`: + # Use Qwen2RMSNorm(in_features) + Linear(in_features -> out_features). + out_features = int(getattr(self.gen_transformer, "hidden_size", 0) or self.gen_transformer.config.hidden_size) + self.gen_transformer.time_caption_embed.caption_embedder = nn.Sequential( + Qwen2RMSNorm(in_features, eps=1e-5), + nn.Linear(in_features, out_features, bias=True), + ) + + def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]: + if num_reqs <= 0: + raise ValueError(f"num_reqs must be positive, got {num_reqs}") + if num_reqs > 1: + raise NotImplementedError( + f"get_dummy_runtime_additional_information does not support num_reqs > 1, got {num_reqs}" + ) + text_prompt_embeds = torch.zeros((1, self._llm_hidden_size), dtype=torch.float32) + image_prompt_embeds = torch.zeros((1, self._llm_hidden_size), dtype=torch.float32) + negative_prompt_embeds = torch.zeros((0, self._llm_hidden_size), dtype=torch.float32) + info = { + "text_prompt_embeds": text_prompt_embeds, + "image_prompt_embeds": image_prompt_embeds, + "negative_prompt_embeds": negative_prompt_embeds, + "negative_prompt_attention_mask": [], + "image_height": [512], + "image_width": [512], + "text_guidance_scale": [1.0], + "cfg_range": [0.0, 1.0], + "num_inference_steps": [1], + } + return [info for _ in range(num_reqs)] + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + # DiT stage does not consume token embeddings; return a dummy tensor. + try: + dtype = next(self.parameters()).dtype + except StopIteration: + dtype = torch.float32 + return torch.zeros( + (input_ids.numel(), self._llm_hidden_size), + device=input_ids.device, + dtype=dtype, + ) + + @torch.inference_mode() + def forward( + self, + *, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, # noqa: ARG002 + ) -> OmniOutput: + runtime_addi = kwargs.get("runtime_additional_information", None) + info = runtime_addi[0] + text_cond = info["text_prompt_embeds"] + image_cond = info["image_prompt_embeds"] + negative_cond = info.get("negative_prompt_embeds") + negative_attention_mask = info.get("negative_prompt_attention_mask") + image_hw = info["image_height"][0], info["image_width"][0] + text_guidance_scale = info["text_guidance_scale"][0] + cfg_range = info["cfg_range"][0], info["cfg_range"][1] + num_inference_steps = info["num_inference_steps"][0] + + # Move to model device/dtype. + model_device = next(self.parameters()).device + if self.gen_image_condition_refiner is not None: + target_dtype = next(self.gen_image_condition_refiner.parameters()).dtype + else: + target_dtype = next(self.gen_transformer.parameters()).dtype + + def _ensure_2d(x: torch.Tensor, name: str) -> torch.Tensor: + if x.ndim == 3 and x.shape[0] == 1: + x = x[0] + if x.ndim != 2: + raise ValueError(f"Expected {name} to be 2D [T,H], got shape={tuple(x.shape)}") + return x + + text_cond = _ensure_2d(text_cond, "text_prompt_embeds") + image_cond = _ensure_2d(image_cond, "image_prompt_embeds") + text_cond = text_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous() + image_cond = image_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous() + + text_embeds = text_cond.unsqueeze(0) # [1, T_text, H] + text_attention_mask = torch.ones( + (1, text_embeds.shape[1]), + dtype=torch.bool, + device=text_embeds.device, + ) + + image_embeds = image_cond.unsqueeze(0) # [1, T_img, H] + image_attention_mask = torch.ones( + (1, image_embeds.shape[1]), + dtype=torch.bool, + device=image_embeds.device, + ) + + # Apply optional refiner ONLY on image condition tokens. + if self.gen_image_condition_refiner is not None and image_embeds.shape[1] > 0: + image_embeds = self.gen_image_condition_refiner(image_embeds, ~image_attention_mask.bool()) + image_attention_mask = torch.ones( + image_embeds.shape[:2], + dtype=torch.bool, + device=image_embeds.device, + ) + + prompt_embeds = torch.cat([text_embeds, image_embeds], dim=1) + prompt_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1) + + # Prepare negative prompt (for CFG). If none provided, fall back to unconditional. + negative_prompt_embeds = None + negative_prompt_attention_mask = None + if text_guidance_scale > 1.0: + if negative_cond is not None: + negative_cond = _ensure_2d(negative_cond, "negative_prompt_embeds") + negative_prompt_embeds = ( + negative_cond.to(device=model_device, dtype=target_dtype, non_blocking=True) + .contiguous() + .unsqueeze(0) + ) + if isinstance(negative_attention_mask, torch.Tensor): + neg_mask = negative_attention_mask + elif isinstance(negative_attention_mask, list): + neg_mask = torch.tensor(negative_attention_mask, dtype=torch.bool) + else: + neg_mask = None + if neg_mask is None: + negative_prompt_attention_mask = torch.ones( + (1, negative_prompt_embeds.shape[1]), + dtype=torch.bool, + device=negative_prompt_embeds.device, + ) + else: + neg_mask = neg_mask.to(device=negative_prompt_embeds.device, dtype=torch.bool) + if neg_mask.ndim == 1: + neg_mask = neg_mask.unsqueeze(0) + negative_prompt_attention_mask = neg_mask + else: + hidden_size = int(prompt_embeds.shape[-1]) + negative_prompt_embeds = torch.zeros( + (1, 0, hidden_size), + dtype=target_dtype, + device=prompt_embeds.device, + ) + negative_prompt_attention_mask = torch.zeros( + (1, 0), + dtype=torch.bool, + device=prompt_embeds.device, + ) + + # Output image size (px), passed from stage input processor. + height, width = image_hw + if height <= 0 or width <= 0: + raise ValueError(f"Invalid image size: {height}x{width}") + if height % 16 != 0 or width % 16 != 0: + raise ValueError(f"Image size must be multiples of 16, got {height}x{width}") + vae_scale_factor = 16 + + latent_channels = int(self.gen_transformer.config.in_channels) + shape = (1, latent_channels, 2 * height // vae_scale_factor, 2 * width // vae_scale_factor) + latents = randn_tensor(shape, device=prompt_embeds.device, dtype=prompt_embeds.dtype) + + scheduler = FlowMatchEulerDiscreteScheduler() + + scheduler.set_timesteps( + num_inference_steps=num_inference_steps, + device=prompt_embeds.device, + num_tokens=latents.shape[-2] * latents.shape[-1], + ) + + # Run diffusion loop (CFG supported when text_guidance_scale > 1.0) + total_steps = max(1, len(scheduler.timesteps)) + for i, t in enumerate(scheduler.timesteps): + timestep = t.expand(latents.shape[0]).to(latents.dtype) + model_pred = self.gen_transformer( + hidden_states=latents, + timestep=timestep, + text_hidden_states=prompt_embeds, + text_attention_mask=prompt_attention_mask, + ref_image_hidden_states=None, + freqs_cis=self.gen_freqs_cis, + ) + guidance_scale = text_guidance_scale if cfg_range[0] <= i / total_steps <= cfg_range[1] else 1.0 + if guidance_scale > 1.0 and negative_prompt_embeds is not None: + model_pred_uncond = self.gen_transformer( + hidden_states=latents, + timestep=timestep, + text_hidden_states=negative_prompt_embeds, + text_attention_mask=negative_prompt_attention_mask, + ref_image_hidden_states=None, + freqs_cis=self.gen_freqs_cis, + ) + model_pred = model_pred_uncond + guidance_scale * (model_pred - model_pred_uncond) + latents = scheduler.step(model_pred, t, latents, return_dict=False)[0] + latents = latents.to(dtype=prompt_embeds.dtype) + + # VAE decode + if self.gen_vae.config.scaling_factor is not None: + latents = latents / self.gen_vae.config.scaling_factor + if self.gen_vae.config.shift_factor is not None: + latents = latents + self.gen_vae.config.shift_factor + image = self.gen_vae.decode(latents, return_dict=False)[0] + + return OmniOutput( + text_hidden_states=inputs_embeds, # placeholder, not used by runner + multimodal_outputs=image, + intermediate_tensors=None, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: # noqa: ARG002 + return None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm_omni/diffusion/models/mammoth_moda2/rope_real.py b/vllm_omni/diffusion/models/mammoth_moda2/rope_real.py new file mode 100644 index 000000000..d16181a69 --- /dev/null +++ b/vllm_omni/diffusion/models/mammoth_moda2/rope_real.py @@ -0,0 +1,250 @@ +# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import torch +from diffusers.models.embeddings import get_1d_rotary_pos_embed as _get_1d_rotary_pos_embed +from einops import repeat +from torch import nn + + +def apply_real_rotary_emb(x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor) -> torch.Tensor: + """ + Apply real-valued rotary embeddings to input tensor. + + Args: + x: Input tensor of shape [..., seq_len, num_heads, dim] or [..., seq_len, dim] + freqs_cos: Cosine frequencies of shape [batch, seq_len, dim] or [seq_len, dim] + freqs_sin: Sine frequencies of shape [batch, seq_len, dim] or [seq_len, dim] + + Returns: + Tensor with rotary embeddings applied + """ + # x: [batch, seq_len, num_heads, dim] or [batch, seq_len, dim] + # freqs_cos: [batch, seq_len, dim] or [seq_len, dim] + # freqs_sin: [batch, seq_len, dim] or [seq_len, dim] + + x_shape = x.shape + if len(x_shape) == 4: + batch, seq_len, num_heads, dim = x_shape + x_reshaped = x.view(batch, seq_len, num_heads, dim // 2, 2) + elif len(x_shape) == 3: + batch, seq_len, dim = x_shape + num_heads = None + x_reshaped = x.view(batch, seq_len, dim // 2, 2) + else: + raise ValueError(f"Unsupported x shape: {x.shape}") + + # freqs_cos/sin: [batch, seq_len, dim] or [seq_len, dim] + # Expand freqs_cos/sin to [batch, seq_len, dim] if needed + if freqs_cos.dim() == 2: + # [seq_len, dim] -> [1, seq_len, dim] + freqs_cos = freqs_cos.unsqueeze(0) + freqs_sin = freqs_sin.unsqueeze(0) + if freqs_cos.shape[0] == 1 and batch > 1: + freqs_cos = freqs_cos.expand(batch, -1, -1) + freqs_sin = freqs_sin.expand(batch, -1, -1) + + # Reshape freqs to [batch, seq_len, dim//2, 2] + freqs_cos_reshaped = freqs_cos.view(batch, seq_len, dim // 2, 2) + freqs_sin_reshaped = freqs_sin.view(batch, seq_len, dim // 2, 2) + + cos_1 = freqs_cos_reshaped[..., 0] # [batch, seq_len, dim//2] + cos_2 = freqs_cos_reshaped[..., 1] # [batch, seq_len, dim//2] + sin_1 = freqs_sin_reshaped[..., 0] # [batch, seq_len, dim//2] + sin_2 = freqs_sin_reshaped[..., 1] # [batch, seq_len, dim//2] + + # Broadcast cos/sin to match x_reshaped + if len(x_shape) == 4: + # [batch, seq_len, 1, dim//2] + cos_1 = cos_1.unsqueeze(2) + cos_2 = cos_2.unsqueeze(2) + sin_1 = sin_1.unsqueeze(2) + sin_2 = sin_2.unsqueeze(2) + + x1 = x_reshaped[..., 0] # [..., seq_len, num_heads, dim//2] or [..., seq_len, dim//2] + x2 = x_reshaped[..., 1] # same + + out1 = x1 * cos_1 - x2 * sin_1 + out2 = x1 * sin_2 + x2 * cos_2 + + out = torch.stack([out1, out2], dim=-1) + return out.view(*x_shape) + + +def get_1d_rotary_pos_embed_real( + dim: int, + pos: np.ndarray | int, + theta: float = 10000.0, + linear_factor: float = 1.0, + ntk_factor: float = 1.0, + freqs_dtype: torch.dtype = torch.float32, +): + freqs_cos, freqs_sin = _get_1d_rotary_pos_embed( + dim, + pos, + theta=theta, + use_real=True, + linear_factor=linear_factor, + ntk_factor=ntk_factor, + repeat_interleave_real=True, + freqs_dtype=freqs_dtype, + ) + return freqs_cos, freqs_sin + + +class RotaryPosEmbedReal(nn.Module): + def __init__( + self, theta: int, axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], patch_size: int = 2 + ): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + self.axes_lens = axes_lens + self.patch_size = patch_size + + @staticmethod + def get_freqs_real( + axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], theta: int + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + freqs_real = [] + freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64 + for i, (d, e) in enumerate(zip(axes_dim, axes_lens)): + cos_emb, sin_emb = get_1d_rotary_pos_embed_real(d, e, theta=theta, freqs_dtype=freqs_dtype) + freqs_real.append((cos_emb, sin_emb)) + return freqs_real + + def _get_freqs_real( + self, freqs_real: list[tuple[torch.Tensor, torch.Tensor]], ids: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + device = ids.device + if ids.device.type == "mps": + ids = ids.to("cpu") + + cos_result = [] + sin_result = [] + for i in range(len(self.axes_dim)): + freqs_cos, freqs_sin = freqs_real[i] + freqs_cos = freqs_cos.to(ids.device) + freqs_sin = freqs_sin.to(ids.device) + index = ids[:, :, i : i + 1].repeat(1, 1, freqs_cos.shape[-1]).to(torch.int64) + cos_result.append(torch.gather(freqs_cos.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index)) + sin_result.append(torch.gather(freqs_sin.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index)) + + combined_cos = torch.cat(cos_result, dim=-1).to(device) + combined_sin = torch.cat(sin_result, dim=-1).to(device) + return combined_cos, combined_sin + + def forward( + self, freqs_real, attention_mask, l_effective_ref_img_len, l_effective_img_len, ref_img_sizes, img_sizes, device + ): + batch_size = len(attention_mask) + p = self.patch_size + + encoder_seq_len = attention_mask.shape[1] + l_effective_cap_len = attention_mask.sum(dim=1).tolist() + l_effective_cap_len = [int(len) for len in l_effective_cap_len] + seq_lengths = [ + int(cap_len + sum(ref_img_len) + img_len) + for cap_len, ref_img_len, img_len in zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len) + ] + + max_seq_len = max(seq_lengths) + max_ref_img_len = max([int(sum(ref_img_len)) for ref_img_len in l_effective_ref_img_len]) + max_img_len = int(max(l_effective_img_len)) + + # Create position IDs + position_ids = torch.zeros(batch_size, int(max_seq_len), 3, dtype=torch.int32, device=device) + + for i, (cap_seq_len, seq_len) in enumerate(zip(l_effective_cap_len, seq_lengths)): + # add text position ids + position_ids[i, :cap_seq_len] = repeat( + torch.arange(cap_seq_len, dtype=torch.int32, device=device), "l -> l 3" + ) + pe_shift = cap_seq_len + pe_shift_len = cap_seq_len + + if ref_img_sizes[i] is not None: + for ref_img_size, ref_img_len in zip(ref_img_sizes[i], l_effective_ref_img_len[i]): + H, W = ref_img_size + ref_H_tokens, ref_W_tokens = H // p, W // p + assert ref_H_tokens * ref_W_tokens == ref_img_len + # add image position ids + + row_ids = repeat( + torch.arange(ref_H_tokens, dtype=torch.int32, device=device), "h -> h w", w=ref_W_tokens + ).flatten() + col_ids = repeat( + torch.arange(ref_W_tokens, dtype=torch.int32, device=device), "w -> h w", h=ref_H_tokens + ).flatten() + position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 0] = pe_shift + position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 1] = row_ids + position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 2] = col_ids + + pe_shift += max(ref_H_tokens, ref_W_tokens) + pe_shift_len += ref_img_len + + H, W = img_sizes[i] + H_tokens, W_tokens = H // p, W // p + assert H_tokens * W_tokens == l_effective_img_len[i] + + row_ids = repeat(torch.arange(H_tokens, dtype=torch.int32, device=device), "h -> h w", w=W_tokens).flatten() + col_ids = repeat(torch.arange(W_tokens, dtype=torch.int32, device=device), "w -> h w", h=H_tokens).flatten() + + assert pe_shift_len + l_effective_img_len[i] == seq_len + position_ids[i, pe_shift_len:seq_len, 0] = pe_shift + position_ids[i, pe_shift_len:seq_len, 1] = row_ids + position_ids[i, pe_shift_len:seq_len, 2] = col_ids + + # Get combined rotary embeddings (real version) + freqs_cos, freqs_sin = self._get_freqs_real(freqs_real, position_ids) + + # create separate rotary embeddings for captions and images + cap_freqs_cos = torch.zeros( + batch_size, encoder_seq_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype + ) + cap_freqs_sin = torch.zeros( + batch_size, encoder_seq_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype + ) + ref_img_freqs_cos = torch.zeros( + batch_size, max_ref_img_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype + ) + ref_img_freqs_sin = torch.zeros( + batch_size, max_ref_img_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype + ) + img_freqs_cos = torch.zeros(batch_size, max_img_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype) + img_freqs_sin = torch.zeros(batch_size, max_img_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype) + + for i, (cap_seq_len, ref_img_len, img_len, seq_len) in enumerate( + zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len, seq_lengths) + ): + cap_freqs_cos[i, :cap_seq_len] = freqs_cos[i, :cap_seq_len] + cap_freqs_sin[i, :cap_seq_len] = freqs_sin[i, :cap_seq_len] + ref_img_freqs_cos[i, : sum(ref_img_len)] = freqs_cos[i, cap_seq_len : cap_seq_len + sum(ref_img_len)] + ref_img_freqs_sin[i, : sum(ref_img_len)] = freqs_sin[i, cap_seq_len : cap_seq_len + sum(ref_img_len)] + img_freqs_cos[i, :img_len] = freqs_cos[ + i, cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len + ] + img_freqs_sin[i, :img_len] = freqs_sin[ + i, cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len + ] + + return ( + (cap_freqs_cos, cap_freqs_sin), + (ref_img_freqs_cos, ref_img_freqs_sin), + (img_freqs_cos, img_freqs_sin), + (freqs_cos, freqs_sin), + l_effective_cap_len, + seq_lengths, + ) diff --git a/vllm_omni/diffusion/models/mammoth_moda2/schedulers.py b/vllm_omni/diffusion/models/mammoth_moda2/schedulers.py new file mode 100644 index 000000000..ac73301ce --- /dev/null +++ b/vllm_omni/diffusion/models/mammoth_moda2/schedulers.py @@ -0,0 +1,137 @@ +# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# SPDX-License-Identifier: Apache-2.0 +# +# This file has been modified by ByteDance Ltd. on 2025-09-30. +# +# Original file was released under Apache-2.0, with the full license text +# available at https://www.apache.org/licenses/LICENSE-2.0 +# +# This modified file is released under the same license. +# +# --- Upstream header preserved below --- +# +# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import torch + + +@dataclass +class FlowMatchEulerDiscreteSchedulerOutput: + prev_sample: torch.FloatTensor + + +class FlowMatchEulerDiscreteScheduler: + order = 1 + + def __init__(self, num_train_timesteps: int = 1000, dynamic_time_shift: bool = True): + self.num_train_timesteps = int(num_train_timesteps) + self.dynamic_time_shift = bool(dynamic_time_shift) + + timesteps = torch.linspace(0, 1, self.num_train_timesteps + 1, dtype=torch.float32)[:-1] + self.timesteps = timesteps + self._timesteps = torch.cat([timesteps, torch.ones(1, dtype=timesteps.dtype)]) + + self._step_index: int | None = None + self._begin_index: int | None = None + + @property + def step_index(self) -> int | None: + return self._step_index + + @property + def begin_index(self) -> int | None: + return self._begin_index + + def set_begin_index(self, begin_index: int = 0) -> None: + self._begin_index = int(begin_index) + + def index_for_timestep(self, timestep: torch.Tensor, schedule_timesteps: torch.Tensor | None = None) -> int: + schedule_timesteps = self._timesteps if schedule_timesteps is None else schedule_timesteps + indices = (schedule_timesteps == timestep).nonzero() + pos = 1 if len(indices) > 1 else 0 + return int(indices[pos].item()) + + def set_timesteps( + self, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[float] | None = None, + num_tokens: int | None = None, + ) -> None: + if timesteps is None: + if num_inference_steps is None: + raise ValueError("`num_inference_steps` must be provided when `timesteps` is None.") + timesteps_np = np.linspace(0, 1, num_inference_steps + 1, dtype=np.float32)[:-1] + if self.dynamic_time_shift and num_tokens is not None: + m = np.sqrt(float(num_tokens)) / 40.0 + timesteps_np = timesteps_np / (m - m * timesteps_np + timesteps_np) + else: + timesteps_np = np.asarray(timesteps, dtype=np.float32) + + timesteps_t = torch.from_numpy(timesteps_np).to(dtype=torch.float32, device=device) + self.timesteps = timesteps_t + self._timesteps = torch.cat([timesteps_t, torch.ones(1, device=timesteps_t.device, dtype=timesteps_t.dtype)]) + + self._step_index = None + self._begin_index = None + + def _init_step_index(self, timestep: float | torch.Tensor) -> None: + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + else: + timestep = torch.tensor(timestep, device=self.timesteps.device, dtype=self.timesteps.dtype) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step( + self, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor, + sample: torch.FloatTensor, + generator: torch.Generator | None = None, # noqa: ARG002 - kept for API compatibility + return_dict: bool = True, + ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor]: + if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)): + raise ValueError( + "Integer timesteps (e.g. a loop index from enumerate) are not supported; " + "pass a value from `scheduler.timesteps` instead.", + ) + + if self.step_index is None: + self._init_step_index(timestep) + assert self._step_index is not None + + sample_fp32 = sample.to(torch.float32) + t = self._timesteps[self._step_index] + t_next = self._timesteps[self._step_index + 1] + + prev_sample = sample_fp32 + (t_next - t) * model_output + prev_sample = prev_sample.to(model_output.dtype) + + self._step_index += 1 + + if not return_dict: + return (prev_sample,) + return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample) + + def __len__(self) -> int: + return self.num_train_timesteps diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 263ce3126..105629f34 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -68,6 +68,21 @@ class DiffusionModelRunner: # Initialize KV cache manager for connector management self.kv_transfer_manager = OmniKVTransferManager.from_od_config(od_config) + def _compile_transformer(self, attr_name: str) -> None: + """Compile a transformer attribute on the pipeline with torch.compile.""" + model = getattr(self.pipeline, attr_name, None) + if model is None: + return + try: + setattr(self.pipeline, attr_name, regionally_compile(model, dynamic=True)) + logger.info("Model runner: %s compiled with torch.compile.", attr_name) + except Exception as e: + logger.warning( + "Model runner: torch.compile for %s failed: %s. Using eager mode.", + attr_name, + e, + ) + def load_model( self, memory_pool_context_fn: callable | None = None, @@ -131,14 +146,8 @@ class DiffusionModelRunner: # Apply torch.compile if not in eager mode if not self.od_config.enforce_eager: if current_omni_platform.supports_torch_inductor(): - try: - self.pipeline.transformer = regionally_compile( - self.pipeline.transformer, - dynamic=True, - ) - logger.info("Model runner: Model compiled with torch.compile.") - except Exception as e: - logger.warning(f"Model runner: torch.compile failed with error: {e}. Using eager mode.") + self._compile_transformer("transformer") + self._compile_transformer("transformer_2") else: logger.warning( "Model runner: Platform %s does not support torch inductor, skipping torch.compile.", diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 1cf66759b..7a75d3563 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -16,6 +16,7 @@ def _register_omni_hf_configs() -> None: try: from transformers import AutoConfig + from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( Qwen3TTSConfig, ) @@ -25,6 +26,7 @@ def _register_omni_hf_configs() -> None: try: AutoConfig.register("qwen3_tts", Qwen3TTSConfig) + AutoConfig.register("cosyvoice3", CosyVoice3Config) except ValueError: # Already registered elsewhere; ignore. return diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 58dd88eb0..9ecc5f762 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -13,6 +13,7 @@ from typing import Any, Literal, overload import huggingface_hub import msgspec.msgpack +import torch import zmq from omegaconf import OmegaConf from tqdm.auto import tqdm @@ -217,8 +218,11 @@ class OmniBase: # We temporally create a default config for diffusion stage. # In the future, we should merge the default config with the user-provided config. # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error. - if "dtype" in kwargs: - kwargs["dtype"] = str(kwargs["dtype"]) + if "dtype" in kwargs and not isinstance(kwargs["dtype"], str): + if not isinstance(kwargs["dtype"], torch.dtype): + raise TypeError(f"Provided dtype must be a string or torch.dtype, got {type(kwargs['dtype']).__name__}") + kwargs["dtype"] = str(kwargs["dtype"]).removeprefix("torch.") + cache_backend = kwargs.get("cache_backend", "none") cache_config = self._normalize_cache_config(cache_backend, kwargs.get("cache_config", None)) # TODO: hack, calculate devices based on parallel config. diff --git a/vllm_omni/model_executor/models/cosyvoice3/__init__.py b/vllm_omni/model_executor/models/cosyvoice3/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/__init__.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py new file mode 100644 index 000000000..7281cd81f --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py @@ -0,0 +1,322 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/FunAudioLLM/CosyVoice/tree/main/cosyvoice/flow +"""Conditional Flow Matching (CFM) classes for audio generation.""" + +from abc import ABC + +import torch +import torch.nn as nn +from omegaconf import DictConfig +from torch.nn import functional as F +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + +logger = init_logger(__name__) + + +class BASECFM(torch.nn.Module, ABC): + def __init__( + self, + n_feats, + cfm_params, + n_spks=1, + spk_emb_dim=128, + ): + super().__init__() + self.n_feats = n_feats + self.n_spks = n_spks + self.spk_emb_dim = spk_emb_dim + self.solver = cfm_params.solver + if hasattr(cfm_params, "sigma_min"): + self.sigma_min = cfm_params.sigma_min + else: + self.sigma_min = 1e-4 + + self.estimator = None + + +class ConditionalCFM(BASECFM): + def __init__(self, in_channels, cfm_params, n_spks=1, spk_emb_dim=64, estimator: torch.nn.Module = None): + super().__init__( + n_feats=in_channels, + cfm_params=cfm_params, + n_spks=n_spks, + spk_emb_dim=spk_emb_dim, + ) + self.t_scheduler = cfm_params.t_scheduler + self.training_cfg_rate = cfm_params.training_cfg_rate + self.inference_cfg_rate = cfm_params.inference_cfg_rate + in_channels = in_channels + (spk_emb_dim if n_spks > 0 else 0) + # Just change the architecture of the estimator here + self.estimator = estimator + + @torch.inference_mode() + def forward( + self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, prompt_len=0, cache=torch.zeros(1, 80, 0, 2) + ): + """Forward diffusion + + Args: + mu (torch.Tensor): output of encoder + shape: (batch_size, n_feats, mel_timesteps) + mask (torch.Tensor): output_mask + shape: (batch_size, 1, mel_timesteps) + n_timesteps (int): number of diffusion steps + temperature (float, optional): temperature for scaling noise. Defaults to 1.0. + spks (torch.Tensor, optional): speaker ids. Defaults to None. + shape: (batch_size, spk_emb_dim) + cond (Optional[Any], optional): Not used but kept for future purposes + + Returns: + sample (torch.Tensor): generated mel-spectrogram + shape: (batch_size, n_feats, mel_timesteps) + """ + + z = torch.randn_like(mu).to(mu.device).to(mu.dtype) * temperature + cache_size = cache.shape[2] + # fix prompt and overlap part mu and z + if cache_size != 0: + z[:, :, :cache_size] = cache[:, :, :, 0] + mu[:, :, :cache_size] = cache[:, :, :, 1] + z_cache = torch.concat([z[:, :, :prompt_len], z[:, :, -34:]], dim=2) + mu_cache = torch.concat([mu[:, :, :prompt_len], mu[:, :, -34:]], dim=2) + cache = torch.stack([z_cache, mu_cache], dim=-1) + + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + if self.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), cache + + def solve_euler(self, x, t_span, mu, mask, spks, cond): + """ + Fixed euler solver for ODEs. + Args: + x (torch.Tensor): random noise + t_span (torch.Tensor): n_timesteps interpolated + shape: (n_timesteps + 1,) + mu (torch.Tensor): output of encoder + shape: (batch_size, n_feats, mel_timesteps) + mask (torch.Tensor): output_mask + shape: (batch_size, 1, mel_timesteps) + spks (torch.Tensor, optional): speaker ids. Defaults to None. + shape: (batch_size, spk_emb_dim) + cond (Optional[Any], optional): Not used but kept for future purposes + """ + t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0] + t = t.unsqueeze(dim=0) + + sol = [] + + # Do not use concat, it may cause memory format changed and trt infer with wrong results! + # NOTE when flow run in amp mode, x.dtype is float32, which cause nan in trt fp16 + # inference, so set dtype=spks.dtype + x_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) + mask_in = torch.zeros([2, 1, x.size(2)], device=x.device, dtype=spks.dtype) + mu_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) + t_in = torch.zeros([2], device=x.device, dtype=spks.dtype) + spks_in = torch.zeros([2, 80], device=x.device, dtype=spks.dtype) + cond_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) + for step in range(1, len(t_span)): + # Classifier-Free Guidance inference introduced in VoiceBox + x_in[:] = x + mask_in[:] = mask + mu_in[0] = mu + t_in[:] = t.unsqueeze(0) + spks_in[0] = spks + cond_in[0] = cond + dphi_dt = self.forward_estimator(x_in, mask_in, mu_in, t_in, spks_in, cond_in) + dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [x.size(0), x.size(0)], dim=0) + dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt + x = x + dt * dphi_dt + t = t + dt + sol.append(x) + if step < len(t_span) - 1: + dt = t_span[step + 1] - t + + return sol[-1].float() + + def forward_estimator(self, x, mask, mu, t, spks, cond): + if isinstance(self.estimator, torch.nn.Module): + return self.estimator(x, mask, mu, t, spks, cond) + else: + [estimator, stream], trt_engine = self.estimator.acquire_estimator() + # NOTE need to synchronize when switching stream + torch.cuda.current_stream().synchronize() + with stream: + estimator.set_input_shape("x", (2, 80, x.size(2))) + estimator.set_input_shape("mask", (2, 1, x.size(2))) + estimator.set_input_shape("mu", (2, 80, x.size(2))) + estimator.set_input_shape("t", (2,)) + estimator.set_input_shape("spks", (2, 80)) + estimator.set_input_shape("cond", (2, 80, x.size(2))) + data_ptrs = [ + x.contiguous().data_ptr(), + mask.contiguous().data_ptr(), + mu.contiguous().data_ptr(), + t.contiguous().data_ptr(), + spks.contiguous().data_ptr(), + cond.contiguous().data_ptr(), + x.data_ptr(), + ] + for i, j in enumerate(data_ptrs): + estimator.set_tensor_address(trt_engine.get_tensor_name(i), j) + # run trt engine + assert estimator.execute_async_v3(torch.cuda.current_stream().cuda_stream) is True + torch.cuda.current_stream().synchronize() + self.estimator.release_estimator(estimator, stream) + return x + + +class CausalConditionalCFM(ConditionalCFM): + def __init__(self, in_channels, cfm_params, n_spks=1, spk_emb_dim=64, estimator: torch.nn.Module = None): + super().__init__(in_channels, cfm_params, n_spks, spk_emb_dim, estimator) + + @torch.inference_mode() + def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None): + """Forward diffusion + + Args: + mu (torch.Tensor): output of encoder + shape: (batch_size, n_feats, mel_timesteps) + mask (torch.Tensor): output_mask + shape: (batch_size, 1, mel_timesteps) + n_timesteps (int): number of diffusion steps + temperature (float, optional): temperature for scaling noise. Defaults to 1.0. + spks (torch.Tensor, optional): speaker ids. Defaults to None. + shape: (batch_size, spk_emb_dim) + cond (Optional[Any], optional): Not used but kept for future purposes + + Returns: + sample (torch.Tensor): generated mel-spectrogram + shape: (batch_size, n_feats, mel_timesteps) + """ + + z = ( + torch.randn( + (mu.size(0), mu.size(1), mu.size(2)), + device=mu.device, + dtype=mu.dtype, + ) + * temperature + ) + + # fix prompt and overlap part mu and z + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + + if self.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + + return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), None + + +class CausalMaskedDiffWithDiT(torch.nn.Module): + def __init__( + self, + input_size: int = 512, + output_size: int = 80, + spk_embed_dim: int = 192, + output_type: str = "mel", + vocab_size: int = 4096, + input_frame_rate: int = 50, + only_mask_loss: bool = True, + token_mel_ratio: int = 2, + pre_lookahead_len: int = 3, + pre_lookahead_layer: torch.nn.Module = None, + decoder: torch.nn.Module = None, + decoder_conf: dict = { + "in_channels": 240, + "out_channel": 80, + "spk_emb_dim": 80, + "n_spks": 1, + "cfm_params": DictConfig( + { + "sigma_min": 1e-06, + "solver": "euler", + "t_scheduler": "cosine", + "training_cfg_rate": 0.2, + "inference_cfg_rate": 0.7, + "reg_loss_type": "l1", + } + ), + "decoder_params": { + "channels": [256, 256], + "dropout": 0.0, + "attention_head_dim": 64, + "n_blocks": 4, + "num_mid_blocks": 12, + "num_heads": 8, + "act_fn": "gelu", + }, + }, + ): + super().__init__() + self.input_size = input_size + self.output_size = output_size + self.decoder_conf = decoder_conf + self.vocab_size = vocab_size + self.output_type = output_type + self.input_frame_rate = input_frame_rate + logger.info(f"input frame rate={self.input_frame_rate}") + self.input_embedding = nn.Embedding(vocab_size, input_size) + self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, output_size) + self.pre_lookahead_len = pre_lookahead_len + self.pre_lookahead_layer = pre_lookahead_layer + self.decoder = decoder + self.only_mask_loss = only_mask_loss + self.token_mel_ratio = token_mel_ratio + + @torch.inference_mode() + def inference( + self, + token, + token_len, + prompt_token, + prompt_token_len, + prompt_feat, + prompt_feat_len, + embedding, + finalize, + ): + assert token.shape[0] == 1 + # xvec projection + + embedding = F.normalize(embedding, dim=1) + + embedding = self.spk_embed_affine_layer(embedding) + + # concat text and prompt_text + token, token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + token_len + mask = (~make_pad_mask(token_len)).unsqueeze(-1).to(embedding) + token = self.input_embedding(torch.clamp(token, min=0)) * mask + # text encode + if finalize is True: + h = self.pre_lookahead_layer(token) + else: + h = self.pre_lookahead_layer( + token[:, : -self.pre_lookahead_len], context=token[:, -self.pre_lookahead_len :] + ) + h = h.repeat_interleave(self.token_mel_ratio, dim=1) + + mel_len1, mel_len2 = prompt_feat.shape[1], h.shape[1] - prompt_feat.shape[1] + + # get conditions + conds = torch.zeros([1, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype) + + conds[:, :mel_len1] = prompt_feat + + conds = conds.transpose(1, 2) + + mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h) + feat, _ = self.decoder( + mu=h.transpose(1, 2).contiguous(), + mask=mask.unsqueeze(1), + spks=embedding, + cond=conds, + n_timesteps=10, + ) + + feat = feat[:, :, mel_len1:] + assert feat.shape[2] == mel_len2 + return feat.float(), None diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py new file mode 100644 index 000000000..bd31443cc --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py @@ -0,0 +1,951 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/FunAudioLLM/CosyVoice/tree/main/cosyvoice/hifigan +# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu) + + +"""HIFI-GAN""" + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from scipy.signal import get_window +from torch import pow, sin +from torch.nn import Conv1d, ConvTranspose1d +from torch.nn.utils import remove_weight_norm + +try: + from torch.nn.utils.parametrizations import weight_norm +except ImportError: + from torch.nn.utils import weight_norm +from torch.distributions.uniform import Uniform +from torch.nn import Parameter + + +# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license. +# LICENSE is in incl_licenses directory. +class Snake(nn.Module): + """ + Implementation of a sine-based periodic activation function + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter + References: + - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://arxiv.org/abs/2006.08195 + Examples: + >>> a1 = snake(256) + >>> x = torch.randn(256) + >>> x = a1(x) + """ + + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + """ + Initialization. + INPUT: + - in_features: shape of the input + - alpha: trainable parameter + alpha is initialized to 1 by default, higher values = higher-frequency. + alpha will be trained along with the rest of your model. + """ + super().__init__() + self.in_features = in_features + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + """ + Forward pass of the function. + Applies the function to the input elementwise. + Snake ∶= x + 1/a * sin^2 (xa) + """ + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + if self.alpha_logscale: + alpha = torch.exp(alpha) + x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + return x + + +"""hifigan based generator implementation. + +This code is modified from https://github.com/jik876/hifi-gan + ,https://github.com/kan-bayashi/ParallelWaveGAN and + https://github.com/NVIDIA/BigVGAN + +""" + + +def init_weights(m, mean=0.0, std=0.01): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + m.weight.data.normal_(mean, std) + + +def get_padding(kernel_size, dilation=1): + return int((kernel_size * dilation - dilation) / 2) + + +class ResBlock(torch.nn.Module): + """Residual block module in HiFiGAN/BigVGAN.""" + + def __init__( + self, + channels: int = 512, + kernel_size: int = 3, + dilations: list[int] = [1, 3, 5], + causal: bool = False, + ): + super().__init__() + self.causal = causal + self.convs1 = nn.ModuleList() + self.convs2 = nn.ModuleList() + + for dilation in dilations: + self.convs1.append( + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation, + padding=get_padding(kernel_size, dilation), + ) + if causal is False + else CausalConv1d(channels, channels, kernel_size, 1, dilation=dilation, causal_type="left") + ) + ) + self.convs2.append( + weight_norm( + Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)) + if causal is False + else CausalConv1d(channels, channels, kernel_size, 1, dilation=1, causal_type="left") + ) + ) + self.convs1.apply(init_weights) + self.convs2.apply(init_weights) + self.activations1 = nn.ModuleList([Snake(channels, alpha_logscale=False) for _ in range(len(self.convs1))]) + self.activations2 = nn.ModuleList([Snake(channels, alpha_logscale=False) for _ in range(len(self.convs2))]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for idx in range(len(self.convs1)): + xt = self.activations1[idx](x) + xt = self.convs1[idx](xt) + xt = self.activations2[idx](xt) + xt = self.convs2[idx](xt) + x = xt + x + return x + + def remove_weight_norm(self): + for idx in range(len(self.convs1)): + remove_weight_norm(self.convs1[idx]) + remove_weight_norm(self.convs2[idx]) + + +class SineGen(torch.nn.Module): + """Definition of sine generator + SineGen(samp_rate, harmonic_num = 0, + sine_amp = 0.1, noise_std = 0.003, + voiced_threshold = 0, + flag_for_pulse=False) + samp_rate: sampling rate in Hz + harmonic_num: number of harmonic overtones (default 0) + sine_amp: amplitude of sine-wavefrom (default 0.1) + noise_std: std of Gaussian noise (default 0.003) + voiced_thoreshold: F0 threshold for U/V classification (default 0) + flag_for_pulse: this SinGen is used inside PulseGen (default False) + Note: when flag_for_pulse is True, the first time step of a voiced + segment is always sin(np.pi) or cos(0) + """ + + def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voiced_threshold=0): + super().__init__() + self.sine_amp = sine_amp + self.noise_std = noise_std + self.harmonic_num = harmonic_num + self.sampling_rate = samp_rate + self.voiced_threshold = voiced_threshold + + def _f02uv(self, f0): + # generate uv signal + uv = (f0 > self.voiced_threshold).type(torch.float32) + return uv + + @torch.no_grad() + def forward(self, f0): + """sine_tensor, uv = forward(f0) + input F0: tensor(batchsize=1, dim=1, length) + f0 for unvoiced steps should be 0 + output sine_tensor: tensor(batchsize=1, length, dim) + output uv: tensor(batchsize=1, length, 1) + """ + f0 = f0.transpose(1, 2) + F_mat = torch.zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1))).to(f0.device) + for i in range(self.harmonic_num + 1): + F_mat[:, i : i + 1, :] = f0 * (i + 1) / self.sampling_rate + + theta_mat = 2 * np.pi * (torch.cumsum(F_mat, dim=-1) % 1) + u_dist = Uniform(low=-np.pi, high=np.pi) + phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device) + phase_vec[:, 0, :] = 0 + + # generate sine waveforms + sine_waves = self.sine_amp * torch.sin(theta_mat + phase_vec) + + # generate uv signal + uv = self._f02uv(f0) + + # noise: for unvoiced should be similar to sine_amp + # std = self.sine_amp/3 -> max value ~ self.sine_amp + # . for voiced regions is self.noise_std + noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 + noise = noise_amp * torch.randn_like(sine_waves) + + # first: set the unvoiced part to 0 by uv + # then: additive noise + sine_waves = sine_waves * uv + noise + return sine_waves.transpose(1, 2), uv.transpose(1, 2), noise + + +class SineGen2(torch.nn.Module): + """Definition of sine generator + SineGen(samp_rate, harmonic_num = 0, + sine_amp = 0.1, noise_std = 0.003, + voiced_threshold = 0, + flag_for_pulse=False) + samp_rate: sampling rate in Hz + harmonic_num: number of harmonic overtones (default 0) + sine_amp: amplitude of sine-wavefrom (default 0.1) + noise_std: std of Gaussian noise (default 0.003) + voiced_thoreshold: F0 threshold for U/V classification (default 0) + flag_for_pulse: this SinGen is used inside PulseGen (default False) + Note: when flag_for_pulse is True, the first time step of a voiced + segment is always sin(np.pi) or cos(0) + """ + + def __init__( + self, + samp_rate, + upsample_scale, + harmonic_num=0, + sine_amp=0.1, + noise_std=0.003, + voiced_threshold=0, + flag_for_pulse=False, + causal=False, + ): + super().__init__() + self.sine_amp = sine_amp + self.noise_std = noise_std + self.harmonic_num = harmonic_num + self.dim = self.harmonic_num + 1 + self.sampling_rate = samp_rate + self.voiced_threshold = voiced_threshold + self.flag_for_pulse = flag_for_pulse + self.upsample_scale = upsample_scale + self.causal = causal + if causal is True: + self.rand_ini = torch.rand(1, 9) + self.rand_ini[:, 0] = 0 + self.sine_waves = torch.rand(1, 300 * 24000, 9) + + def _f02uv(self, f0): + # generate uv signal + uv = (f0 > self.voiced_threshold).type(torch.float32) + return uv + + def _f02sine(self, f0_values): + """f0_values: (batchsize, length, dim) + where dim indicates fundamental tone and overtones + """ + # convert to F0 in rad. The integer part n can be ignored + # because 2 * np.pi * n doesn't affect phase + rad_values = (f0_values / self.sampling_rate) % 1 + + # initial phase noise (no noise for fundamental component) + if self.training is False and self.causal is True: + rad_values[:, 0, :] = rad_values[:, 0, :] + self.rand_ini.to(rad_values.device) + else: + rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device) + rand_ini[:, 0] = 0 + rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini + + # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad) + if not self.flag_for_pulse: + rad_values = torch.nn.functional.interpolate( + rad_values.transpose(1, 2), scale_factor=1 / self.upsample_scale, mode="linear" + ).transpose(1, 2) + + phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi + phase = torch.nn.functional.interpolate( + phase.transpose(1, 2) * self.upsample_scale, + scale_factor=self.upsample_scale, + mode="nearest" if self.causal is True else "linear", + ).transpose(1, 2) + sines = torch.sin(phase) + else: + # If necessary, make sure that the first time step of every + # voiced segments is sin(pi) or cos(0) + # This is used for pulse-train generation + + # identify the last time step in unvoiced segments + uv = self._f02uv(f0_values) + uv_1 = torch.roll(uv, shifts=-1, dims=1) + uv_1[:, -1, :] = 1 + u_loc = (uv < 1) * (uv_1 > 0) + + # get the instantanouse phase + tmp_cumsum = torch.cumsum(rad_values, dim=1) + # different batch needs to be processed differently + for idx in range(f0_values.shape[0]): + temp_sum = tmp_cumsum[idx, u_loc[idx, :, 0], :] + temp_sum[1:, :] = temp_sum[1:, :] - temp_sum[0:-1, :] + # stores the accumulation of i.phase within + # each voiced segments + tmp_cumsum[idx, :, :] = 0 + tmp_cumsum[idx, u_loc[idx, :, 0], :] = temp_sum + + # rad_values - tmp_cumsum: remove the accumulation of i.phase + # within the previous voiced segment. + i_phase = torch.cumsum(rad_values - tmp_cumsum, dim=1) + + # get the sines + sines = torch.cos(i_phase * 2 * np.pi) + return sines + + def forward(self, f0): + """sine_tensor, uv = forward(f0) + input F0: tensor(batchsize=1, length, dim=1) + f0 for unvoiced steps should be 0 + output sine_tensor: tensor(batchsize=1, length, dim) + output uv: tensor(batchsize=1, length, 1) + """ + # fundamental component + fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device)) + + # generate sine waveforms + sine_waves = self._f02sine(fn) * self.sine_amp + + # generate uv signal + uv = self._f02uv(f0) + + # noise: for unvoiced should be similar to sine_amp + # std = self.sine_amp/3 -> max value ~ self.sine_amp + # . for voiced regions is self.noise_std + noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 + if self.training is False and self.causal is True: + noise = noise_amp * self.sine_waves[:, : sine_waves.shape[1]].to(sine_waves.device) + else: + noise = noise_amp * torch.randn_like(sine_waves) + + # first: set the unvoiced part to 0 by uv + # then: additive noise + sine_waves = sine_waves * uv + noise + return sine_waves, uv, noise + + +class SourceModuleHnNSF(torch.nn.Module): + """SourceModule for hn-nsf + SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1, + add_noise_std=0.003, voiced_threshold=0) + sampling_rate: sampling_rate in Hz + harmonic_num: number of harmonic above F0 (default: 0) + sine_amp: amplitude of sine source signal (default: 0.1) + add_noise_std: std of additive Gaussian noise (default: 0.003) + note that amplitude of noise in unvoiced is decided + by sine_amp + voiced_threshold: threshold to set U/V given F0 (default: 0) + Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) + F0_sampled (batchsize, length, 1) + Sine_source (batchsize, length, 1) + noise_source (batchsize, length 1) + uv (batchsize, length, 1) + """ + + def __init__( + self, + sampling_rate, + upsample_scale, + harmonic_num=0, + sine_amp=0.1, + add_noise_std=0.003, + voiced_threshold=0, + sinegen_type="1", + causal=False, + ): + super().__init__() + + self.sine_amp = sine_amp + self.noise_std = add_noise_std + + # to produce sine waveforms + if sinegen_type == "1": + self.l_sin_gen = SineGen(sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold) + else: + self.l_sin_gen = SineGen2( + sampling_rate, upsample_scale, harmonic_num, sine_amp, add_noise_std, voiced_threshold, causal=causal + ) + + # to merge source harmonics into a single excitation + self.l_linear = torch.nn.Linear(harmonic_num + 1, 1) + self.l_tanh = torch.nn.Tanh() + self.causal = causal + if causal is True: + self.uv = torch.rand(1, 300 * 24000, 1) + + def forward(self, x): + """ + Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) + F0_sampled (batchsize, length, 1) + Sine_source (batchsize, length, 1) + noise_source (batchsize, length 1) + """ + # source for harmonic branch + with torch.no_grad(): + sine_wavs, uv, _ = self.l_sin_gen(x) + sine_merge = self.l_tanh(self.l_linear(sine_wavs)) + + # source for noise branch, in the same shape as uv + if self.training is False and self.causal is True: + noise = self.uv[:, : uv.shape[1]] * self.sine_amp / 3 + else: + noise = torch.randn_like(uv) * self.sine_amp / 3 + return sine_merge, noise, uv + + +class HiFTGenerator(nn.Module): + """ + HiFTNet Generator: Neural Source Filter + ISTFTNet + https://arxiv.org/abs/2309.09493 + """ + + def __init__( + self, + in_channels: int = 80, + base_channels: int = 512, + nb_harmonics: int = 8, + sampling_rate: int = 22050, + nsf_alpha: float = 0.1, + nsf_sigma: float = 0.003, + nsf_voiced_threshold: float = 10, + upsample_rates: list[int] = [8, 8], + upsample_kernel_sizes: list[int] = [16, 16], + istft_params: dict[str, int] = {"n_fft": 16, "hop_len": 4}, + resblock_kernel_sizes: list[int] = [3, 7, 11], + resblock_dilation_sizes: list[list[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + source_resblock_kernel_sizes: list[int] = [7, 11], + source_resblock_dilation_sizes: list[list[int]] = [[1, 3, 5], [1, 3, 5]], + lrelu_slope: float = 0.1, + audio_limit: float = 0.99, + f0_predictor: torch.nn.Module = None, + ): + super().__init__() + + self.out_channels = 1 + self.nb_harmonics = nb_harmonics + self.sampling_rate = sampling_rate + self.istft_params = istft_params + self.lrelu_slope = lrelu_slope + self.audio_limit = audio_limit + + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + # NOTE in CosyVoice2, we use the original SineGen implementation + self.m_source = SourceModuleHnNSF( + sampling_rate=sampling_rate, + upsample_scale=np.prod(upsample_rates) * istft_params["hop_len"], + harmonic_num=nb_harmonics, + sine_amp=nsf_alpha, + add_noise_std=nsf_sigma, + voiced_threshold=nsf_voiced_threshold, + sinegen_type="1" if self.sampling_rate == 22050 else "2", + causal=False, + ) + self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * istft_params["hop_len"]) + + self.conv_pre = weight_norm(Conv1d(in_channels, base_channels, 7, 1, padding=3)) + + # Up + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): + self.ups.append( + weight_norm( + ConvTranspose1d( + base_channels // (2**i), + base_channels // (2 ** (i + 1)), + k, + u, + padding=(k - u) // 2, + ) + ) + ) + + # Down + self.source_downs = nn.ModuleList() + self.source_resblocks = nn.ModuleList() + downsample_rates = [1] + upsample_rates[::-1][:-1] + downsample_cum_rates = np.cumprod(downsample_rates) + for i, (u, k, d) in enumerate( + zip(downsample_cum_rates[::-1], source_resblock_kernel_sizes, source_resblock_dilation_sizes) + ): + if u == 1: + self.source_downs.append(Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), 1, 1)) + else: + self.source_downs.append( + Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), u * 2, u, padding=(u // 2)) + ) + + self.source_resblocks.append(ResBlock(base_channels // (2 ** (i + 1)), k, d)) + + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = base_channels // (2 ** (i + 1)) + for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)): + self.resblocks.append(ResBlock(ch, k, d)) + + self.conv_post = weight_norm(Conv1d(ch, istft_params["n_fft"] + 2, 7, 1, padding=3)) + self.ups.apply(init_weights) + self.conv_post.apply(init_weights) + self.reflection_pad = nn.ReflectionPad1d((1, 0)) + self.stft_window = torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32)) + self.f0_predictor = f0_predictor + + def remove_weight_norm(self): + for layer in self.ups: + remove_weight_norm(layer) + for block in self.resblocks: + block.remove_weight_norm() + remove_weight_norm(self.conv_pre) + remove_weight_norm(self.conv_post) + self.m_source.remove_weight_norm() + for layer in self.source_downs: + remove_weight_norm(layer) + for block in self.source_resblocks: + block.remove_weight_norm() + + def _stft(self, x): + spec = torch.stft( + x, + self.istft_params["n_fft"], + self.istft_params["hop_len"], + self.istft_params["n_fft"], + window=self.stft_window.to(x.device), + return_complex=True, + ) + spec = torch.view_as_real(spec) # [B, F, TT, 2] + return spec[..., 0], spec[..., 1] + + def _istft(self, magnitude, phase): + magnitude = torch.clip(magnitude, max=1e2) + real = magnitude * torch.cos(phase) + img = magnitude * torch.sin(phase) + inverse_transform = torch.istft( + torch.complex(real, img), + self.istft_params["n_fft"], + self.istft_params["hop_len"], + self.istft_params["n_fft"], + window=self.stft_window.to(magnitude.device), + ) + return inverse_transform + + def decode(self, x: torch.Tensor, s: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor: + s_stft_real, s_stft_imag = self._stft(s.squeeze(1)) + s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1) + + x = self.conv_pre(x) + for i in range(self.num_upsamples): + x = F.leaky_relu(x, self.lrelu_slope) + x = self.ups[i](x) + + if i == self.num_upsamples - 1: + x = self.reflection_pad(x) + + # fusion + si = self.source_downs[i](s_stft) + si = self.source_resblocks[i](si) + x = x + si + + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs / self.num_kernels + + x = F.leaky_relu(x) + x = self.conv_post(x) + magnitude = torch.exp(x[:, : self.istft_params["n_fft"] // 2 + 1, :]) + phase = torch.sin(x[:, self.istft_params["n_fft"] // 2 + 1 :, :]) # actually, sin is redundancy + + x = self._istft(magnitude, phase) + x = torch.clamp(x, -self.audio_limit, self.audio_limit) + return x + + def forward( + self, + batch: dict, + device: torch.device, + ) -> dict[str, torch.Tensor | None]: + speech_feat = batch["speech_feat"].transpose(1, 2).to(device) + # mel->f0 + f0 = self.f0_predictor(speech_feat) + # f0->source + s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t + s, _, _ = self.m_source(s) + s = s.transpose(1, 2) + # mel+source->speech + generated_speech = self.decode(x=speech_feat, s=s) + return generated_speech, f0 + + @torch.inference_mode() + def inference(self, speech_feat: torch.Tensor, cache_source: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor: + # mel->f0 + f0 = self.f0_predictor(speech_feat) + # f0->source + s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t + s, _, _ = self.m_source(s) + s = s.transpose(1, 2) + # use cache_source to avoid glitch + if cache_source.shape[2] != 0: + s[:, :, : cache_source.shape[2]] = cache_source + generated_speech = self.decode(x=speech_feat, s=s) + return generated_speech, s + + +class CausalHiFTGenerator(HiFTGenerator): + """ + HiFTNet Generator: Neural Source Filter + ISTFTNet + https://arxiv.org/abs/2309.09493 + """ + + def __init__( + self, + in_channels: int = 80, + base_channels: int = 512, + nb_harmonics: int = 8, + sampling_rate: int = 22050, + nsf_alpha: float = 0.1, + nsf_sigma: float = 0.003, + nsf_voiced_threshold: float = 10, + upsample_rates: list[int] = [8, 8], + upsample_kernel_sizes: list[int] = [16, 16], + istft_params: dict[str, int] = {"n_fft": 16, "hop_len": 4}, + resblock_kernel_sizes: list[int] = [3, 7, 11], + resblock_dilation_sizes: list[list[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + source_resblock_kernel_sizes: list[int] = [7, 11], + source_resblock_dilation_sizes: list[list[int]] = [[1, 3, 5], [1, 3, 5]], + lrelu_slope: float = 0.1, + audio_limit: float = 0.99, + conv_pre_look_right: int = 4, + f0_predictor: torch.nn.Module = None, + ): + torch.nn.Module.__init__(self) + + self.out_channels = 1 + self.nb_harmonics = nb_harmonics + self.sampling_rate = sampling_rate + self.istft_params = istft_params + self.lrelu_slope = lrelu_slope + self.audio_limit = audio_limit + + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + self.m_source = SourceModuleHnNSF( + sampling_rate=sampling_rate, + upsample_scale=np.prod(upsample_rates) * istft_params["hop_len"], + harmonic_num=nb_harmonics, + sine_amp=nsf_alpha, + add_noise_std=nsf_sigma, + voiced_threshold=nsf_voiced_threshold, + sinegen_type="1" if self.sampling_rate == 22050 else "2", + causal=True, + ) + self.upsample_rates = upsample_rates + self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * istft_params["hop_len"]) + + self.conv_pre = weight_norm( + CausalConv1d(in_channels, base_channels, conv_pre_look_right + 1, 1, causal_type="right") + ) + + # Up + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): + self.ups.append( + weight_norm( + CausalConv1dUpsample( + base_channels // (2**i), + base_channels // (2 ** (i + 1)), + k, + u, + ) + ) + ) + + # Down + self.source_downs = nn.ModuleList() + self.source_resblocks = nn.ModuleList() + downsample_rates = [1] + upsample_rates[::-1][:-1] + downsample_cum_rates = np.cumprod(downsample_rates) + for i, (u, k, d) in enumerate( + zip(downsample_cum_rates[::-1], source_resblock_kernel_sizes, source_resblock_dilation_sizes) + ): + if u == 1: + self.source_downs.append( + CausalConv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), 1, 1, causal_type="left") + ) + else: + self.source_downs.append( + CausalConv1dDownSample(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), u * 2, u) + ) + + self.source_resblocks.append(ResBlock(base_channels // (2 ** (i + 1)), k, d, causal=True)) + + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = base_channels // (2 ** (i + 1)) + for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)): + self.resblocks.append(ResBlock(ch, k, d, causal=True)) + + self.conv_post = weight_norm(CausalConv1d(ch, istft_params["n_fft"] + 2, 7, 1, causal_type="left")) + self.ups.apply(init_weights) + self.conv_post.apply(init_weights) + self.reflection_pad = nn.ReflectionPad1d((1, 0)) + self.stft_window = torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32)) + self.conv_pre_look_right = conv_pre_look_right + self.f0_predictor = f0_predictor + + def decode(self, x: torch.Tensor, s: torch.Tensor = torch.zeros(1, 1, 0), finalize: bool = True) -> torch.Tensor: + s_stft_real, s_stft_imag = self._stft(s.squeeze(1)) + if finalize is True: + x = self.conv_pre(x) + else: + x = self.conv_pre(x[:, :, : -self.conv_pre_look_right], x[:, :, -self.conv_pre_look_right :]) + s_stft_real = s_stft_real[:, :, : -int(np.prod(self.upsample_rates) * self.conv_pre_look_right)] + s_stft_imag = s_stft_imag[:, :, : -int(np.prod(self.upsample_rates) * self.conv_pre_look_right)] + s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1) + + for i in range(self.num_upsamples): + x = F.leaky_relu(x, self.lrelu_slope) + x = self.ups[i](x) + + if i == self.num_upsamples - 1: + x = self.reflection_pad(x) + + # fusion + si = self.source_downs[i](s_stft) + si = self.source_resblocks[i](si) + x = x + si + + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs / self.num_kernels + + x = F.leaky_relu(x) + x = self.conv_post(x) + magnitude = torch.exp(x[:, : self.istft_params["n_fft"] // 2 + 1, :]) + phase = torch.sin(x[:, self.istft_params["n_fft"] // 2 + 1 :, :]) # actually, sin is redundancy + + x = self._istft(magnitude, phase) + if finalize is False: + x = x[:, : -int(np.prod(self.upsample_rates) * self.istft_params["hop_len"])] + x = torch.clamp(x, -self.audio_limit, self.audio_limit) + return x + + @torch.inference_mode() + def inference(self, speech_feat: torch.Tensor, finalize: bool = True) -> torch.Tensor: + # mel->f0 NOTE f0_predictor precision is crucial for causal inference, move + # self.f0_predictor to cpu if necessary + self.f0_predictor.to("cpu") + f0 = self.f0_predictor(speech_feat.cpu(), finalize=finalize).to(speech_feat) + # f0->source + s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t + s, _, _ = self.m_source(s) + s = s.transpose(1, 2) + if finalize is True: + generated_speech = self.decode(x=speech_feat, s=s, finalize=finalize) + else: + generated_speech = self.decode( + x=speech_feat[:, :, : -self.f0_predictor.condnet[0].causal_padding], s=s, finalize=finalize + ) + return generated_speech, s + + +class CausalConv1dUpsample(torch.nn.Conv1d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + device=None, + dtype=None, + ) -> None: + super().__init__( + in_channels, + out_channels, + kernel_size, + 1, + padding=0, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=padding_mode, + device=device, + dtype=dtype, + ) + assert dilation == 1 + self.causal_padding = kernel_size - 1 + self.upsample = torch.nn.Upsample(scale_factor=stride, mode="nearest") + + def forward(self, x: torch.Tensor, cache: torch.Tensor = torch.zeros(0, 0, 0)) -> tuple[torch.Tensor, torch.Tensor]: + x = self.upsample(x) + input_timestep = x.shape[2] + if cache.size(2) == 0: + x = F.pad(x, (self.causal_padding, 0), value=0.0) + else: + assert cache.size(2) == self.causal_padding + x = torch.concat([cache, x], dim=2) + x = super().forward(x) + assert input_timestep == x.shape[2] + return x + + +class CausalConv1dDownSample(torch.nn.Conv1d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + device=None, + dtype=None, + ) -> None: + super().__init__( + in_channels, + out_channels, + kernel_size, + stride, + padding=0, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=padding_mode, + device=device, + dtype=dtype, + ) + assert stride != 1 and dilation == 1 + assert kernel_size % stride == 0 + self.causal_padding = stride - 1 + + def forward(self, x: torch.Tensor, cache: torch.Tensor = torch.zeros(0, 0, 0)) -> tuple[torch.Tensor, torch.Tensor]: + if cache.size(2) == 0: + x = F.pad(x, (self.causal_padding, 0), value=0.0) + else: + assert cache.size(2) == self.causal_padding + x = torch.concat([cache, x], dim=2) + x = super().forward(x) + return x + + +# NOTE(Xiang Lyu) causal conv module used in convolution-based vocoder +class CausalConv1d(torch.nn.Conv1d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + causal_type: str = "left", + device=None, + dtype=None, + ) -> None: + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=padding_mode, + device=device, + dtype=dtype, + ) + assert stride == 1 + self.causal_padding = int((kernel_size * dilation - dilation) / 2) * 2 + (kernel_size + 1) % 2 + assert causal_type in ["left", "right"] + self.causal_type = causal_type + + def forward(self, x: torch.Tensor, cache: torch.Tensor = torch.zeros(0, 0, 0)) -> tuple[torch.Tensor]: + input_timestep = x.shape[2] + if cache.size(2) == 0: + cache = torch.zeros(x.shape[0], x.shape[1], self.causal_padding).to(x) + assert cache.size(2) == self.causal_padding + if self.causal_type == "left": + x = torch.concat([cache, x], dim=2) + else: + x = torch.concat([x, cache], dim=2) + x = super().forward(x) + assert x.shape[2] == input_timestep + return x + + +class CausalConvRNNF0Predictor(nn.Module): + def __init__(self, num_class: int = 1, in_channels: int = 80, cond_channels: int = 512): + super().__init__() + + self.num_class = num_class + self.condnet = nn.Sequential( + weight_norm(CausalConv1d(in_channels, cond_channels, kernel_size=4, causal_type="right")), + nn.ELU(), + weight_norm(CausalConv1d(cond_channels, cond_channels, kernel_size=3, causal_type="left")), + nn.ELU(), + weight_norm(CausalConv1d(cond_channels, cond_channels, kernel_size=3, causal_type="left")), + nn.ELU(), + weight_norm(CausalConv1d(cond_channels, cond_channels, kernel_size=3, causal_type="left")), + nn.ELU(), + weight_norm(CausalConv1d(cond_channels, cond_channels, kernel_size=3, causal_type="left")), + nn.ELU(), + ) + self.classifier = nn.Linear(in_features=cond_channels, out_features=self.num_class) + + def forward(self, x: torch.Tensor, finalize: bool = True) -> torch.Tensor: + if finalize is True: + x = self.condnet[0](x) + else: + x = self.condnet[0](x[:, :, : -self.condnet[0].causal_padding], x[:, :, -self.condnet[0].causal_padding :]) + for i in range(1, len(self.condnet)): + x = self.condnet[i](x) + x = x.transpose(1, 2) + return torch.abs(self.classifier(x).squeeze(-1)) diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/layers.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/layers.py new file mode 100644 index 000000000..40d6bb54e --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/layers.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Neural network layers for CosyVoice3 Code2Wav.""" + +import torch +import torch.nn as nn +from torch.nn import functional as F + + +class PreLookaheadLayer(nn.Module): + """Pre-lookahead convolutional layer for causal processing.""" + + def __init__(self, in_channels: int, channels: int, pre_lookahead_len: int = 1): + super().__init__() + self.in_channels = in_channels + self.channels = channels + self.pre_lookahead_len = pre_lookahead_len + self.conv1 = nn.Conv1d( + in_channels, + channels, + kernel_size=pre_lookahead_len + 1, + stride=1, + padding=0, + ) + self.conv2 = nn.Conv1d( + channels, + in_channels, + kernel_size=3, + stride=1, + padding=0, + ) + + def forward(self, inputs: torch.Tensor, context: torch.Tensor = torch.zeros(0, 0, 0)) -> torch.Tensor: + """ + Args: + inputs: (batch_size, seq_len, channels) + context: Optional context tensor for streaming + """ + outputs = inputs.transpose(1, 2).contiguous() + context = context.transpose(1, 2).contiguous() + # look ahead + if context.size(2) == 0: + outputs = F.pad(outputs, (0, self.pre_lookahead_len), mode="constant", value=0.0) + else: + assert self.training is False, "you have passed context, make sure that you are running inference mode" + assert context.size(2) == self.pre_lookahead_len + outputs = F.pad( + torch.concat([outputs, context], dim=2), + (0, self.pre_lookahead_len - context.size(2)), + mode="constant", + value=0.0, + ) + outputs = F.leaky_relu(self.conv1(outputs)) + # outputs + outputs = F.pad(outputs, (self.conv2.kernel_size[0] - 1, 0), mode="constant", value=0.0) + outputs = self.conv2(outputs) + outputs = outputs.transpose(1, 2).contiguous() + + # residual connection + outputs = outputs + inputs + return outputs diff --git a/vllm_omni/model_executor/models/cosyvoice3/config.py b/vllm_omni/model_executor/models/cosyvoice3/config.py new file mode 100644 index 000000000..0c9a28997 --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/config.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from transformers.configuration_utils import PretrainedConfig + + +class CosyVoice3Config(PretrainedConfig): + model_type = "cosyvoice3" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.sample_rate = 24000 + self.llm_input_size = 896 + self.llm_output_size = 896 + self.hidden_size = self.llm_output_size + self.num_attention_heads = 14 + self.num_hidden_layers = 24 + self.spk_embed_dim = 192 + self.token_frame_rate = 25 + self.token_mel_ratio = 2 + self.vocab_size = 151923 + self.min_token_text_ratio = 2 + self.max_token_text_ratio = 20 + self.allowed_special = "all" + self.skip_special_tokens = True + self.target_sr = 24000 + self.feat_extractor = { + "n_fft": 1920, + "num_mels": 80, + "sampling_rate": self.sample_rate, + "hop_size": 480, + "win_size": 1920, + "fmin": 0, + "fmax": None, + "center": False, + } + self.speculative_config = None + self.qwen_pretrain_path = "CosyVoice-BlankEN" + self.campplus_onxx_path = "campplus.onnx" + self.speech_tokenizer_path = "speech_tokenizer_v3.onnx" + self.spk2info_path = "spk2info.pt" + self.version = "cosyvoice3" + self.llm = { + "llm_input_size": self.llm_input_size, + "llm_output_size": self.llm_output_size, + "speech_token_size": 6561, + "eos_token_id": 6561 + 1, + "length_normalized_loss": True, + "lsm_weight": 0, + "mix_ratio": [5, 15], + "llm": { + "pretrain_path": self.qwen_pretrain_path, + }, + "sampling": { + "top_p": 0.8, + "top_k": 25, + "win_size": 10, + "tau_r": 0.1, + }, + "spk_embed_dim": self.spk_embed_dim, + } + self.flow = { + "input_size": 80, + "output_size": 80, + "spk_embed_dim": self.spk_embed_dim, + "output_type": "mel", + "vocab_size": 6561, + "input_frame_rate": self.token_frame_rate, + "only_mask_loss": True, + "token_mel_ratio": self.token_mel_ratio, + "pre_lookahead_len": 3, + "pre_lookahead_layer": { + "in_channels": 80, + "channels": 1024, + "pre_lookahead_len": 3, + }, + "decoder": { + "in_channels": 240, + "n_spks": 1, + "spk_emb_dim": 80, + "cfm_params": { + "sigma_min": 1e-06, + "solver": "euler", + "t_scheduler": "cosine", + "training_cfg_rate": 0.2, + "inference_cfg_rate": 0.7, + "reg_loss_type": "l1", + }, + "estimator": { + "dim": 1024, + "depth": 22, + "heads": 16, + "dim_head": 64, + "ff_mult": 2, + "mel_dim": 80, + "mu_dim": 80, + "spk_dim": 80, + "out_channels": 80, + "static_chunk_size": self.token_frame_rate * self.token_mel_ratio, + "num_decoding_left_chunks": -1, + }, + }, + } + self.hift = { + "in_channels": 80, + "base_channels": 512, + "nb_harmonics": 8, + "sampling_rate": self.sample_rate, + "nsf_alpha": 0.1, + "nsf_sigma": 0.003, + "nsf_voiced_threshold": 10, + "upsample_rates": [8, 5, 3], + "upsample_kernel_sizes": [16, 11, 7], + "istft_params": { + "n_fft": 16, + "hop_len": 4, + }, + "resblock_kernel_sizes": [3, 7, 11], + "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "source_resblock_kernel_sizes": [7, 7, 11], + "source_resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "lrelu_slope": 0.1, + "audio_limit": 0.99, + "conv_pre_look_right": 4, + "f0_predictor": { + "num_class": 1, + "in_channels": 80, + "cond_channels": 512, + }, + } diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py new file mode 100644 index 000000000..87c5f323a --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -0,0 +1,468 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os +from collections.abc import Iterable, Mapping, Sequence +from functools import partial + +import numpy as np +import torch +import torch.nn as nn +from transformers.feature_extraction_utils import BatchFeature +from vllm.config import VllmConfig +from vllm.config.multimodal import BaseDummyOptions +from vllm.logger import init_logger +from vllm.model_executor.models.interfaces import SupportsMultiModal +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalDataDict, MultiModalFieldConfig, MultiModalKwargsItems +from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + ProcessorInputs, + PromptIndexTargets, + PromptInsertion, + PromptUpdate, +) +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config +from vllm_omni.model_executor.models.cosyvoice3.utils import ( + concat_text_with_prompt_ids, + extract_speech_feat, + extract_speech_token, + extract_spk_embedding, + extract_text_token, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + + +class CosyVoice3MultiModalProcessingInfo(BaseProcessingInfo): + def get_hf_config(self): + """If the config is not already present pass it + as a class and it will try to find it in your + model directory just copy the config class there also. + """ + return self.ctx.get_hf_config(CosyVoice3Config) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + """How many audio can you pass. I think I should keep it as 1 + For now I have kept it None. + """ + return {"audio": None} + + def get_data_parser(self): + return MultiModalDataParser( + target_sr=self.ctx.get_hf_config().target_sr, + expected_hidden_size=self._get_expected_hidden_size(), + ) + + +class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModalProcessingInfo]): + def _ensure_cached_runtime_components(self, model_dir: str, config: CosyVoice3Config) -> None: + cached_model_dir = getattr(self, "_cached_model_dir", None) + if cached_model_dir == model_dir: + return + + import onnxruntime + + from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer + from vllm_omni.model_executor.models.cosyvoice3.utils import mel_spectrogram + + option = onnxruntime.SessionOptions() + option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + option.intra_op_num_threads = 1 + + self.tokenizer = get_qwen_tokenizer( + token_path=os.path.join(model_dir, config.qwen_pretrain_path), + skip_special_tokens=config.skip_special_tokens, + version=config.version, + ) + self.speech_tokenizer = onnxruntime.InferenceSession( + os.path.join(model_dir, config.speech_tokenizer_path), + sess_options=option, + providers=["CUDAExecutionProvider" if torch.cuda.is_available() else "CPUExecutionProvider"], + ) + self.feat_extractor = partial(mel_spectrogram, **getattr(config, "feat_extractor", {})) + self.campplus_session = onnxruntime.InferenceSession( + os.path.join(model_dir, config.campplus_onxx_path), + sess_options=option, + providers=["CPUExecutionProvider"], + ) + self._cached_model_dir = model_dir + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + """ + apply-> cached_apply_hf_processor -> apply_hf_processor_mm -> + _call_hf_processor. + _call_hf_processor takes input prompt and mm_data and returns + token ids and tensors + """ + config = self.info.ctx.get_hf_config() + model_dir = self.info.ctx.model_config.model + self._ensure_cached_runtime_components(model_dir, config) + + audio = mm_data.get("audio", None) + + if audio is None: + audio = mm_data.get("audios") + if audio is not None: + audio = audio[0], config.target_sr + + text_token, text_token_len = extract_text_token(prompt, self.tokenizer, config.allowed_special) + if audio is None: + # Text-only path for profiling/cache + return BatchFeature({"input_ids": text_token, "input_len": [text_token_len]}) + + prompt_text = mm_kwargs.get("prompt_text") + + if not isinstance(prompt_text, str): + raise ValueError(f"prompt text is None : {prompt_text}") + + prompt_text_token, prompt_text_token_len = extract_text_token( + prompt_text, self.tokenizer, config.allowed_special + ) + + input_ids, input_len = concat_text_with_prompt_ids( + text_token, + text_token_len, + prompt_text_token, + prompt_text_token_len, + ) + logger.debug( + "cosyvoice _call_hf_processor: prompt_text_token=%s text_token=%s input_ids=%s " + "prompt_text_len=%s text_len=%s input_len=%s", + prompt_text_token.tolist(), + text_token.tolist(), + input_ids.tolist(), + int(prompt_text_token_len), + int(text_token_len), + int(input_len), + ) + device = "cpu" + + speech_token, speech_token_len = extract_speech_token(audio, self.speech_tokenizer, device) + speech_feat, speech_feat_len = extract_speech_feat(audio, self.feat_extractor, device) + + if config.sample_rate == 24000: + token_len = min(int(speech_feat.shape[1] / 2), speech_token.shape[1]) + speech_feat, speech_feat_len[:] = speech_feat[:, : 2 * token_len], 2 * token_len + speech_token, speech_token_len[:] = speech_token[:, :token_len], token_len + + embedding = extract_spk_embedding(audio, self.campplus_session, device) + + ft = BatchFeature( + { + "input_ids": input_ids, + "speech_feat": speech_feat, + "speech_token": speech_token, + "speech_token_len": [speech_token_len], + "embedding": embedding, + } + ) + + return ft + + def _get_mm_fields_config( + self, + hf_inputs: "BatchFeature", + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + return { + "speech_feat": MultiModalFieldConfig.batched("audio"), + "speech_token": MultiModalFieldConfig.batched("audio"), + "speech_token_len": MultiModalFieldConfig.batched("audio"), + "embedding": MultiModalFieldConfig.batched("audio"), + } + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return False + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + def insertion_end(item_idx): + # TODO: Think if this can be done better + # sos + task + audio token ... ideally this needs to be split into + # two start and end but somehow I couldn't pass two of these + # wutg target .start() and .end() + token_len = out_mm_kwargs["audio"][0]["speech_token_len"].data[0].item() + return [1] * (1 + 1 + token_len) + + return [ + PromptInsertion( + modality="audio", + target=PromptIndexTargets.start(), + insertion=insertion_end, + ), + ] + + +class CosyVoice3DummyInputsBuilder(BaseDummyInputsBuilder[CosyVoice3MultiModalProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return "Hello, this is a test of the CosyVoice3 system capability." + + def get_dummy_mm_data( + self, seq_len: int, mm_counts: Mapping[str, int], mm_options: Mapping[str, BaseDummyOptions] | None = None + ) -> MultiModalDataDict: + num_audios = mm_counts.get("audio") + max_prompt_seconds = 30 + prompt_sample_rate = 24000 + target_audio_length = max_prompt_seconds * prompt_sample_rate + + audio_overrides = mm_options.get("audio") if mm_options else None + mm_data = { + "audio": ( + self._get_dummy_audios( + length=target_audio_length, + num_audios=num_audios, + overrides=audio_overrides, + )[0], + 24000, + ), + } + return mm_data + + def get_dummy_processor_inputs( + self, seq_len: int, mm_counts: Mapping[str, int], mm_options: Mapping[str, BaseDummyOptions] | None = None + ) -> ProcessorInputs: + inputs = super().get_dummy_processor_inputs(seq_len, mm_counts, mm_options) + inputs.hf_processor_mm_kwargs = {"prompt_text": "Testing my voices. Why should I not?"} + return inputs + + +@MULTIMODAL_REGISTRY.register_processor( + CosyVoice3MultiModalProcessor, + info=CosyVoice3MultiModalProcessingInfo, + dummy_inputs=CosyVoice3DummyInputsBuilder, +) +class CosyVoice3Model( + nn.Module, + SupportsMultiModal, +): + supports_multimodal_raw_input_only = True + supports_multimodal = True + requires_raw_input_tokens = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_config + self.have_multimodal_outputs = True + self.model_stage = vllm_config.model_config.model_stage + self.model_dir = vllm_config.model_config.model + self.model = None + if self.model_stage == "talker": + # Initialize talker stage (text to speech tokens) + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_talker import CosyVoice3LM, VLLMQwen2Encoder + + llm_vllm_config = self._create_llm_vllm_config(vllm_config) + llm = VLLMQwen2Encoder(vllm_config=llm_vllm_config, prefix="model") + self.talker = CosyVoice3LM( + llm_input_size=self.config.llm["llm_input_size"], + llm_output_size=self.config.llm["llm_output_size"], + speech_token_size=self.config.llm["speech_token_size"], + llm=llm, + length_normalized_loss=self.config.llm["length_normalized_loss"], + lsm_weight=self.config.llm["lsm_weight"], + mix_ratio=self.config.llm["mix_ratio"], + ) + # KV cache is now managed externally by vLLM's PagedAttention + # No need for self.llm_cache + self.model = self.talker + elif self.model_stage == "code2wav": + # Initialize code2wav stage (flow matching + vocoder) + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + + self.code2wav = CosyVoice3Code2Wav(self.config) + self.model = self.code2wav.flow_model + self.hift = self.code2wav.hift + + # Expose streaming parameters + self.token_overlap_len = self.code2wav.token_overlap_len + self.mel_overlap_len = self.code2wav.mel_overlap_len + self.mel_window = self.code2wav.mel_window + self.mel_cache_len = self.code2wav.mel_cache_len + self.source_cache_len = self.code2wav.source_cache_len + self.speech_window = self.code2wav.speech_window + else: + raise ValueError(f"Model stage not supported {self.model_stage}") + + def _create_llm_vllm_config(self, parent_config: VllmConfig) -> VllmConfig: + """Create VllmConfig for the inner Qwen2 LLM. + + This creates a modified VllmConfig with the Qwen2 HF config loaded from + the pretrained model directory. The cache config is inherited from the parent + to enable PagedAttention with the same memory configuration. + """ + from transformers import Qwen2Config + + qwen_config_path = os.path.join(self.model_dir, self.config.llm["llm"]["pretrain_path"]) + qwen_hf_config = Qwen2Config.from_pretrained(qwen_config_path) + + # Use parent's cache config - critical for PagedAttention to work correctly + return parent_config.with_hf_config(qwen_hf_config, architectures=["Qwen2Model"]) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput) -> torch.Tensor | None: + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if self.model_stage == "talker": + logits = self.model.llm_decoder(hidden_states) + vocab_size = self.config.vocab_size + pad_size = vocab_size - logits.size(-1) + pad_shape = logits.shape[:-1] + (pad_size,) + pad = logits.new_full(pad_shape, float("-inf")) + eos_token_val = logits[..., self.config.llm["eos_token_id"]].clone() + logits[..., -200:] = float("-inf") + logits[..., self.config.llm["eos_token_id"]] = eos_token_val + logits = torch.cat([logits, pad], dim=-1) + return logits + else: + raise RuntimeError(f"compute_logits is only valid for {self.model_stage}.") + + def embed_multimodal(self, **kwargs: object) -> torch.Tensor: + if self.model_stage == "talker": + speech_token = kwargs["speech_token"] + speech_token_emb = self.model.speech_embedding(speech_token) + return speech_token_emb + else: + raise RuntimeError(f"embed_multimodal is only valid for {self.model_stage}.") + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + is_multimodal=None, + ) -> torch.Tensor: + if self.model_stage == "talker": + if is_multimodal is not None and any(is_multimodal): + embed_tokens = self.model.llm.model.embed_tokens(input_ids) + sos = self.model.speech_embedding.weight[self.model.sos].reshape(1, -1) + task_id = self.model.speech_embedding.weight[self.model.task_id].reshape(1, -1) + prompt_speech_token_emb = multimodal_embeddings[0] + pstoken_len = prompt_speech_token_emb.shape[0] # Get length from tensor shape + embed_tokens = torch.cat( + [sos, embed_tokens[2 + pstoken_len :], task_id, prompt_speech_token_emb], dim=0 + ) + else: + embed_tokens = self.model.speech_embedding.weight[input_ids] + return embed_tokens + elif self.model_stage == "code2wav": + assert input_ids.dim() == 1 + hidden = int(self.config.hidden_size) + return torch.zeros( + (input_ids.shape[0], hidden), + ) + else: + raise RuntimeError(f"embed_input_ids is not valid for {self.model_stage}.") + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + additional_information: dict[str, object] | None = None, + **kwargs: object, + ) -> OmniOutput: + if self.model_stage == "talker": + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + + # [total_tokens, hidden] + hidden_states = self.model.llm(inputs_embeds, positions) + + multimodal_outputs = {} + + if "speech_token" in kwargs: + # Wrap in lists to pass through gpu_ar_model_runner shape filtering + multimodal_outputs = { + "speech_token": [kwargs.get("speech_token")], + "speech_feat": [kwargs.get("speech_feat")], + "embedding": [kwargs.get("embedding")], + } + + return OmniOutput(text_hidden_states=hidden_states, multimodal_outputs=multimodal_outputs) + elif self.model_stage == "code2wav": + runtime_info = kwargs.get("runtime_additional_information", []) + if not runtime_info: + length = 30 * 24000 + audio = np.zeros((length,)) + return OmniOutput(text_hidden_states=None, multimodal_outputs={"audio": audio}) + + # Remove the last eos token and add batch dimension + token = input_ids[..., :-1].unsqueeze(0) + + # Generate audio using code2wav + tts_speech = self.code2wav( + token=token, + prompt_token=runtime_info[0]["speech_token"][:1], + prompt_feat=runtime_info[0]["speech_feat"][:1], + embedding=runtime_info[0]["embedding"][:1], + n_timesteps=10, + ) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"audio": tts_speech}, + ) + else: + raise ValueError(f"Unsupported model_stage: {self.model_stage}") + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + if self.model_stage == "talker": + # Load weights for text to speech LM stage using vLLM's weight loading + llm_weight_path = os.path.join(self.model_dir, "llm.pt") + device = next(self.parameters()).device + checkpoint = torch.load(llm_weight_path, map_location=device) + + # 1. Load Qwen2 model weights into vLLM's Qwen2Model + # The checkpoint has prefix "llm.model.model." for the transformer weights + # vLLM's Qwen2Model expects just the model structure without extra prefixes + qwen_weights = [] + for name, weight in checkpoint.items(): + if name.startswith("llm.model.model."): + # Strip prefix: llm.model.model.X -> X (for vLLM's Qwen2Model) + vllm_name = name.replace("llm.model.model.", "") + qwen_weights.append((vllm_name, weight)) + + # Use vLLM's built-in load_weights which handles stacked params + # (q_proj+k_proj+v_proj -> qkv_proj, gate_proj+up_proj -> gate_up_proj) + self.model.llm.model.load_weights(iter(qwen_weights)) + + # 2. Load CosyVoice3LM-specific weights (speech_embedding, llm_decoder) + speech_emb_state = { + k.replace("speech_embedding.", ""): v + for k, v in checkpoint.items() + if k.startswith("speech_embedding.") + } + self.model.speech_embedding.load_state_dict(speech_emb_state) + + llm_decoder_state = { + k.replace("llm_decoder.", ""): v for k, v in checkpoint.items() if k.startswith("llm_decoder.") + } + self.model.llm_decoder.load_state_dict(llm_decoder_state) + + self.model.to(device).eval() + elif self.model_stage == "code2wav": + # Load weights for code2wav stage (flow + hift) + device = next(self.parameters()).device + self.code2wav.load_weights(self.model_dir, device) + else: + raise ValueError(f"{self.model_stage} not supported yet!") diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py new file mode 100644 index 000000000..f5e0d04a8 --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +CosyVoice3 Code2Wav Stage - Converts speech tokens to audio waveforms. + +This module contains the code2wav (token-to-waveform) stage which uses: +1. DiT (Diffusion Transformer) with optimized attention backends +2. CFM (Conditional Flow Matching) for mel spectrogram generation +3. HiFiGAN vocoder for waveform synthesis +""" + +from __future__ import annotations + +import numpy as np +import torch +import torch.nn as nn +from omegaconf import DictConfig +from torch.nn import functional as F +from vllm.logger import init_logger + +from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import ( + CausalConditionalCFM, + CausalMaskedDiffWithDiT, +) +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( + CausalConvRNNF0Predictor, + CausalHiFTGenerator, +) +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer +from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config +from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + +logger = init_logger(__name__) + + +class CosyVoice3Code2Wav(nn.Module): + """CosyVoice3 Code2Wav stage for token-to-waveform conversion. + + This class encapsulates: + - Flow matching decoder with DiT backbone (using diffusion attention) + - HiFiGAN vocoder for mel-to-waveform conversion + """ + + def __init__(self, config: CosyVoice3Config): + super().__init__() + self.config = config + + # Build flow matching components + pre_lookahead_layer = PreLookaheadLayer(**config.flow["pre_lookahead_layer"]) + + decoder_cfg = config.flow["decoder"] + cfm_params = DictConfig(decoder_cfg["cfm_params"]) + + # DiT estimator using diffusion attention (Flash/Sage/SDPA backends) + estimator = DiT(**decoder_cfg["estimator"]) + + decoder = CausalConditionalCFM( + in_channels=decoder_cfg["in_channels"], + estimator=estimator, + cfm_params=cfm_params, + n_spks=decoder_cfg["n_spks"], + spk_emb_dim=decoder_cfg["spk_emb_dim"], + ) + + self.flow_model = CausalMaskedDiffWithDiT( + input_size=config.flow["input_size"], + output_size=config.flow["output_size"], + spk_embed_dim=config.flow["spk_embed_dim"], + output_type=config.flow["output_type"], + vocab_size=config.flow["vocab_size"], + input_frame_rate=config.flow["input_frame_rate"], + only_mask_loss=config.flow["only_mask_loss"], + token_mel_ratio=config.flow["token_mel_ratio"], + pre_lookahead_len=config.flow["pre_lookahead_len"], + pre_lookahead_layer=pre_lookahead_layer, + decoder=decoder, + ) + + # Build HiFiGAN vocoder + f0_predictor = CausalConvRNNF0Predictor( + num_class=config.hift["f0_predictor"]["num_class"], + in_channels=config.hift["f0_predictor"]["in_channels"], + cond_channels=config.hift["f0_predictor"]["cond_channels"], + ) + + self.hift = CausalHiFTGenerator( + in_channels=config.hift["in_channels"], + base_channels=config.hift["base_channels"], + nb_harmonics=config.hift["nb_harmonics"], + sampling_rate=config.hift["sampling_rate"], + nsf_alpha=config.hift["nsf_alpha"], + nsf_sigma=config.hift["nsf_sigma"], + nsf_voiced_threshold=config.hift["nsf_voiced_threshold"], + upsample_rates=config.hift["upsample_rates"], + upsample_kernel_sizes=config.hift["upsample_kernel_sizes"], + istft_params=config.hift["istft_params"], + resblock_kernel_sizes=config.hift["resblock_kernel_sizes"], + resblock_dilation_sizes=config.hift["resblock_dilation_sizes"], + source_resblock_kernel_sizes=config.hift["source_resblock_kernel_sizes"], + source_resblock_dilation_sizes=config.hift["source_resblock_dilation_sizes"], + lrelu_slope=config.hift["lrelu_slope"], + audio_limit=config.hift["audio_limit"], + conv_pre_look_right=config.hift["conv_pre_look_right"], + f0_predictor=f0_predictor, + ) + # Run hift in float32 to avoid dtype mismatches in internal ops + self.hift = self.hift.float() + + # Streaming/chunking parameters + self.token_overlap_len = 20 + self.mel_overlap_len = int(self.token_overlap_len / self.flow_model.input_frame_rate * 22050 / 256) + self.mel_window = np.hamming(2 * self.mel_overlap_len) + self.mel_cache_len = 20 + self.source_cache_len = int(self.mel_cache_len * 256) + self.speech_window = np.hamming(2 * self.source_cache_len) + + @property + def input_frame_rate(self) -> int: + """Input frame rate from flow model.""" + return self.flow_model.input_frame_rate + + @property + def token_mel_ratio(self) -> int: + """Token to mel ratio.""" + return self.flow_model.token_mel_ratio + + @property + def output_size(self) -> int: + """Output mel dimension.""" + return self.flow_model.output_size + + @property + def input_embedding(self) -> nn.Embedding: + """Token embedding layer.""" + return self.flow_model.input_embedding + + @property + def pre_lookahead_layer(self) -> nn.Module: + """Pre-lookahead layer.""" + return self.flow_model.pre_lookahead_layer + + @property + def decoder(self) -> nn.Module: + """Flow matching decoder.""" + return self.flow_model.decoder + + @property + def spk_embed_affine_layer(self) -> nn.Linear: + """Speaker embedding affine layer.""" + return self.flow_model.spk_embed_affine_layer + + @torch.inference_mode() + def forward( + self, + token: torch.Tensor, + prompt_token: torch.Tensor, + prompt_feat: torch.Tensor, + embedding: torch.Tensor, + n_timesteps: int = 10, + ) -> torch.Tensor: + """Generate audio waveform from speech tokens. + + Args: + token: Speech tokens from talker stage [batch, seq_len] + prompt_token: Prompt speech tokens [batch, prompt_len] + prompt_feat: Prompt mel features [batch, feat_len, mel_dim] + embedding: Speaker embedding [batch, spk_dim] + n_timesteps: Number of diffusion steps + + Returns: + Audio waveform [batch, 1, audio_len] + """ + device = token.device + dtype = next(self.flow_model.parameters()).dtype + + # Normalize and project speaker embedding + embedding = embedding.to(device=device, dtype=dtype) + embedding = F.normalize(embedding, dim=1) + embedding = self.spk_embed_affine_layer(embedding) + + # Prepare tokens + prompt_token = prompt_token.to(device=device) + token_len1, token_len2 = prompt_token.shape[1], token.shape[1] + prompt_token_len = torch.tensor([token_len1], device=device, dtype=torch.int32) + token_len = torch.tensor([token_len2], device=device, dtype=torch.int32) + + # Concatenate prompt and target tokens + full_token = torch.cat([prompt_token, token], dim=1) + full_token_len = prompt_token_len + token_len + + # Create mask + mask = (~make_pad_mask(full_token_len)).unsqueeze(-1).to(embedding) + + # Token embedding + token_emb = self.input_embedding(torch.clamp(full_token, min=0)) * mask + + # Pre-lookahead processing + h = self.pre_lookahead_layer(token_emb) + h = h.repeat_interleave(self.token_mel_ratio, dim=1) + + # Calculate mel lengths + mel_len1 = prompt_feat.shape[1] + mel_len2 = h.shape[1] - mel_len1 + + # Build conditioning + conds = torch.zeros( + [1, mel_len1 + mel_len2, self.output_size], + device=device, + dtype=h.dtype, + ) + conds[:, :mel_len1] = prompt_feat + conds = conds.transpose(1, 2) + + # Create mel mask + mel_mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h) + + # Run flow matching decoder + feat, _ = self.decoder( + mu=h.transpose(1, 2).contiguous(), + mask=mel_mask.unsqueeze(1), + spks=embedding, + cond=conds, + n_timesteps=n_timesteps, + ) + + # Extract generated portion (after prompt) + feat = feat[:, :, mel_len1:] + + # Run vocoder + hift_weight = self.hift.m_source.l_linear.weight + tts_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype) + + if tts_mel.shape[-1] == 0: + tts_speech = torch.zeros( + (tts_mel.shape[0], 1, 0), + device=tts_mel.device, + dtype=tts_mel.dtype, + ) + else: + tts_speech, _ = self.hift.inference(speech_feat=tts_mel) + + return tts_speech + + def load_weights(self, model_dir: str, device: torch.device) -> None: + """Load flow.pt and hift.pt weights. + + Args: + model_dir: Model directory containing flow.pt and hift.pt + device: Device to load weights to + """ + import os + + # Load flow weights + flow_path = os.path.join(model_dir, "flow.pt") + self.flow_model.load_state_dict(torch.load(flow_path, map_location=device), strict=True) + self.flow_model.to(device).eval() + logger.info(f"Loaded flow weights from {flow_path}") + + # Load hift weights + hift_path = os.path.join(model_dir, "hift.pt") + hift_state_dict = { + k.replace("generator.", ""): v for k, v in torch.load(hift_path, map_location=device).items() + } + self.hift.load_state_dict(hift_state_dict, strict=True) + self.hift.to(device).eval() + logger.info(f"Loaded hift weights from {hift_path}") diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_talker.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_talker.py new file mode 100644 index 000000000..aeea2f8cd --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_talker.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/FunAudioLLM/CosyVoice/blob/main/cosyvoice/llm/llm.py +from collections.abc import Callable + +import torch +from torch import nn +from vllm.config import VllmConfig +from vllm.model_executor.models.qwen2 import Qwen2Model + + +class VLLMQwen2Encoder(torch.nn.Module): + """Qwen2 encoder using vLLM's Qwen2Model with external KV cache management. + + This replaces the HuggingFace Qwen2ForCausalLM with vLLM's optimized implementation + that uses PagedAttention and external KV cache via ForwardContext. + """ + + def __init__(self, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.model = Qwen2Model(vllm_config=vllm_config, prefix=prefix) + self.hidden_size = vllm_config.model_config.hf_config.hidden_size + + def forward(self, inputs_embeds: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """Forward pass using vLLM's attention with external KV cache. + + Args: + inputs_embeds: Input embeddings [total_tokens, hidden_size] or [batch, seq, hidden] + positions: Position tensor for RoPE [total_tokens] + + Returns: + hidden_states: Output hidden states [total_tokens, hidden_size] + """ + # vLLM model expects flattened tensors [total_tokens, hidden_size] + if inputs_embeds.dim() == 3: + inputs_flat = inputs_embeds.view(-1, self.hidden_size) + else: + inputs_flat = inputs_embeds + positions_flat = positions.view(-1) + + # KV cache managed externally via ForwardContext (set by GPUARModelRunner) + # input_ids is required but ignored when inputs_embeds is provided + hidden_states = self.model( + input_ids=torch.zeros(inputs_flat.size(0), dtype=torch.long, device=inputs_flat.device), + positions=positions_flat, + intermediate_tensors=None, + inputs_embeds=inputs_flat, + ) + return hidden_states + + +class TransformerLM(torch.nn.Module): + def __init__( + self, + text_encoder_input_size: int, + llm_input_size: int, + llm_output_size: int, + text_token_size: int, + speech_token_size: int, + text_encoder: torch.nn.Module, + llm: torch.nn.Module, + length_normalized_loss: bool = True, + lsm_weight: float = 0.0, + spk_embed_dim: int = 192, + ): + super().__init__() + self.llm_input_size = llm_input_size + self.speech_token_size = speech_token_size + # 1. build text token inputs related modules + self.text_embedding = torch.nn.Embedding(text_token_size, text_encoder_input_size) + self.text_encoder = text_encoder + self.text_encoder_affine_layer = nn.Linear(self.text_encoder.output_size(), llm_input_size) + + # 2. build speech token language model related modules + self.sos = 0 + self.task_id = 1 + self.eos_token = self.speech_token_size + self.llm_embedding = torch.nn.Embedding(2, llm_input_size) + self.llm = llm + self.llm_decoder = nn.Linear(llm_output_size, speech_token_size + 1) + + # 3. [Optional] build speech token related modules + self.speech_embedding = torch.nn.Embedding(speech_token_size, llm_input_size) + self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, llm_input_size) + + +class Qwen2LM(TransformerLM): + def __init__( + self, + llm_input_size: int, + llm_output_size: int, + speech_token_size: int, + llm: torch.nn.Module, + sampling: Callable, + length_normalized_loss: bool = True, + lsm_weight: float = 0.0, + mix_ratio: list[int] = [5, 15], + ): + torch.nn.Module.__init__(self) + self.llm_input_size = llm_input_size + self.llm_output_size = llm_output_size + self.speech_token_size = speech_token_size + # 2. build speech token language model related modules + self.sos = 0 + self.task_id = 1 + self.eos_token = speech_token_size + self.fill_token = speech_token_size + 2 + + self.llm_embedding = torch.nn.Embedding(2, llm_input_size) + self.llm = llm + self.llm_decoder = nn.Linear(llm_output_size, speech_token_size + 3) + + # 3. [Optional] build speech token related modules + self.speech_embedding = torch.nn.Embedding(speech_token_size + 3, llm_input_size) + + # 4. sampling method + self.sampling = sampling + self.mix_ratio = mix_ratio + + # 5. vllm related + self.stop_token_ids = [speech_token_size + i for i in range(3)] + self.vllm_output_queue = {} + + +class CosyVoice3LM(Qwen2LM): + def __init__( + self, + llm_input_size: int, + llm_output_size: int, + speech_token_size: int, + llm: torch.nn.Module, + length_normalized_loss: bool = True, + lsm_weight: float = 0.0, + mix_ratio: list[int] = [5, 15], + ): + torch.nn.Module.__init__(self) + self.llm_input_size = llm_input_size + self.llm_output_size = llm_output_size + self.speech_token_size = speech_token_size + # 2. build speech token language model related modules + self.sos = speech_token_size + 0 + self.eos_token = speech_token_size + 1 + self.task_id = speech_token_size + 2 + self.fill_token = speech_token_size + 3 + + self.llm = llm + self.llm_decoder = nn.Linear(llm_output_size, speech_token_size + 200, bias=False) + + # 3. [Optional] build speech token related modules + self.speech_embedding = torch.nn.Embedding(speech_token_size + 200, llm_input_size) + + # 4. sampling method + # self.sampling = sampling + self.mix_ratio = mix_ratio + + # 5. vllm related + self.stop_token_ids = [speech_token_size + i for i in range(200)] + self.vllm_output_queue = {} diff --git a/vllm_omni/model_executor/models/cosyvoice3/tokenizer.py b/vllm_omni/model_executor/models/cosyvoice3/tokenizer.py new file mode 100644 index 000000000..6049df61e --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/tokenizer.py @@ -0,0 +1,356 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from functools import cache + +import torch +from transformers import AutoTokenizer + + +class CosyVoice2Tokenizer: + def __init__(self, token_path, skip_special_tokens=True): + super().__init__() + # NOTE: non-chat model, all these special tokens keep randomly initialized. + special_tokens = { + "eos_token": "<|endoftext|>", + "pad_token": "<|endoftext|>", + "additional_special_tokens": [ + "<|im_start|>", + "<|im_end|>", + "<|endofprompt|>", + "[breath]", + "", + "", + "[noise]", + "[laughter]", + "[cough]", + "[clucking]", + "[accent]", + "[quick_breath]", + "", + "", + "[hissing]", + "[sigh]", + "[vocalized-noise]", + "[lipsmack]", + "[mn]", + ], + } + self.special_tokens = special_tokens + self.tokenizer = AutoTokenizer.from_pretrained(token_path) + self.tokenizer.add_special_tokens(special_tokens) + self.skip_special_tokens = skip_special_tokens + + def encode(self, text, **kwargs): + tokens = self.tokenizer([text], return_tensors="pt") + tokens = tokens["input_ids"][0].cpu().tolist() + return tokens + + def decode(self, tokens): + tokens = torch.tensor(tokens, dtype=torch.int64) + text = self.tokenizer.batch_decode([tokens], skip_special_tokens=self.skip_special_tokens)[0] + return text + + +class CosyVoice3Tokenizer(CosyVoice2Tokenizer): + def __init__(self, token_path, skip_special_tokens=True): + # NOTE: non-chat model, all these special tokens keep randomly initialized. + special_tokens = { + "eos_token": "<|endoftext|>", + "pad_token": "<|endoftext|>", + "additional_special_tokens": [ + "<|im_start|>", + "<|im_end|>", + "<|endofprompt|>", + "[breath]", + "", + "", + "[noise]", + "[laughter]", + "[cough]", + "[clucking]", + "[accent]", + "[quick_breath]", + "", + "", + "[hissing]", + "[sigh]", + "[vocalized-noise]", + "[lipsmack]", + "[mn]", + "<|endofsystem|>", + "[AA]", + "[AA0]", + "[AA1]", + "[AA2]", + "[AE]", + "[AE0]", + "[AE1]", + "[AE2]", + "[AH]", + "[AH0]", + "[AH1]", + "[AH2]", + "[AO]", + "[AO0]", + "[AO1]", + "[AO2]", + "[AW]", + "[AW0]", + "[AW1]", + "[AW2]", + "[AY]", + "[AY0]", + "[AY1]", + "[AY2]", + "[B]", + "[CH]", + "[D]", + "[DH]", + "[EH]", + "[EH0]", + "[EH1]", + "[EH2]", + "[ER]", + "[ER0]", + "[ER1]", + "[ER2]", + "[EY]", + "[EY0]", + "[EY1]", + "[EY2]", + "[F]", + "[G]", + "[HH]", + "[IH]", + "[IH0]", + "[IH1]", + "[IH2]", + "[IT]", + "[IT0]", + "[IT1]", + "[IT2]", + "[JH]", + "[K]", + "[L]", + "[M]", + "[N]", + "[NG]", + "[OW]", + "[OW0]", + "[OW1]", + "[OW2]", + "[OY]", + "[OY0]", + "[OY1]", + "[OY2]", + "[P]", + "[R]", + "[S]", + "[SH]", + "[T]", + "[TH]", + "[UH]", + "[UH0]", + "[UH1]", + "[UH2]", + "[UW]", + "[UW0]", + "[UW1]", + "[UW2]", + "[V]", + "[W]", + "[Y]", + "[Z]", + "[ZH]", + "[a]", + "[ai]", + "[an]", + "[ang]", + "[ao]", + "[b]", + "[c]", + "[ch]", + "[d]", + "[e]", + "[ei]", + "[en]", + "[eng]", + "[f]", + "[g]", + "[h]", + "[i]", + "[ian]", + "[in]", + "[ing]", + "[iu]", + "[ià]", + "[iàn]", + "[iàng]", + "[iào]", + "[iá]", + "[ián]", + "[iáng]", + "[iáo]", + "[iè]", + "[ié]", + "[iòng]", + "[ióng]", + "[iù]", + "[iú]", + "[iā]", + "[iān]", + "[iāng]", + "[iāo]", + "[iē]", + "[iě]", + "[iōng]", + "[iū]", + "[iǎ]", + "[iǎn]", + "[iǎng]", + "[iǎo]", + "[iǒng]", + "[iǔ]", + "[j]", + "[k]", + "[l]", + "[m]", + "[n]", + "[o]", + "[ong]", + "[ou]", + "[p]", + "[q]", + "[r]", + "[s]", + "[sh]", + "[t]", + "[u]", + "[uang]", + "[ue]", + "[un]", + "[uo]", + "[uà]", + "[uài]", + "[uàn]", + "[uàng]", + "[uá]", + "[uái]", + "[uán]", + "[uáng]", + "[uè]", + "[ué]", + "[uì]", + "[uí]", + "[uò]", + "[uó]", + "[uā]", + "[uāi]", + "[uān]", + "[uāng]", + "[uē]", + "[uě]", + "[uī]", + "[uō]", + "[uǎ]", + "[uǎi]", + "[uǎn]", + "[uǎng]", + "[uǐ]", + "[uǒ]", + "[vè]", + "[w]", + "[x]", + "[y]", + "[z]", + "[zh]", + "[à]", + "[ài]", + "[àn]", + "[àng]", + "[ào]", + "[á]", + "[ái]", + "[án]", + "[áng]", + "[áo]", + "[è]", + "[èi]", + "[èn]", + "[èng]", + "[èr]", + "[é]", + "[éi]", + "[én]", + "[éng]", + "[ér]", + "[ì]", + "[ìn]", + "[ìng]", + "[í]", + "[ín]", + "[íng]", + "[ò]", + "[òng]", + "[òu]", + "[ó]", + "[óng]", + "[óu]", + "[ù]", + "[ùn]", + "[ú]", + "[ún]", + "[ā]", + "[āi]", + "[ān]", + "[āng]", + "[āo]", + "[ē]", + "[ēi]", + "[ēn]", + "[ēng]", + "[ě]", + "[ěi]", + "[ěn]", + "[ěng]", + "[ěr]", + "[ī]", + "[īn]", + "[īng]", + "[ō]", + "[ōng]", + "[ōu]", + "[ū]", + "[ūn]", + "[ǎ]", + "[ǎi]", + "[ǎn]", + "[ǎng]", + "[ǎo]", + "[ǐ]", + "[ǐn]", + "[ǐng]", + "[ǒ]", + "[ǒng]", + "[ǒu]", + "[ǔ]", + "[ǔn]", + "[ǘ]", + "[ǚ]", + "[ǜ]", + ], + } + self.special_tokens = special_tokens + self.tokenizer = AutoTokenizer.from_pretrained(token_path) + self.tokenizer.add_special_tokens(special_tokens) + self.skip_special_tokens = skip_special_tokens + + +@cache +def get_qwen_tokenizer(token_path: str, skip_special_tokens: bool, version: str = "cosyvoice3"): + if version == "cosyvoice2": + return CosyVoice2Tokenizer(token_path=token_path, skip_special_tokens=skip_special_tokens) + elif version == "cosyvoice3": + return CosyVoice3Tokenizer(token_path=token_path, skip_special_tokens=skip_special_tokens) + else: + raise ValueError diff --git a/vllm_omni/model_executor/models/cosyvoice3/utils.py b/vllm_omni/model_executor/models/cosyvoice3/utils.py new file mode 100644 index 000000000..e1310cd3b --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/utils.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os +from functools import cache, lru_cache + +import numpy as np +import torch +import torch.nn.functional as F +import torchaudio +import torchaudio.compliance.kaldi as kaldi +from librosa.filters import mel as librosa_mel_fn + +IGNORE_ID = -1 + + +def dynamic_range_compression_torch(x, c=1, clip_val=1e-5): + return torch.log(torch.clamp(x, min=clip_val) * c) + + +def spectral_normalize_torch(magnitudes): + output = dynamic_range_compression_torch(magnitudes) + return output + + +@lru_cache +def _get_mel_basis( + sampling_rate: int, + n_fft: int, + num_mels: int, + fmin: float, + fmax: float | None, + device_str: str, +) -> torch.Tensor: + mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax) + return torch.from_numpy(mel).float().to(torch.device(device_str)) + + +@lru_cache +def _get_hann_window(win_size: int, device_str: str) -> torch.Tensor: + return torch.hann_window(win_size).to(torch.device(device_str)) + + +def mel_spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False): + device_str = str(y.device) + mel = _get_mel_basis( + int(sampling_rate), + int(n_fft), + int(num_mels), + float(fmin), + fmax, + device_str, + ) + window = _get_hann_window(int(win_size), device_str) + + y = torch.nn.functional.pad( + y.unsqueeze(1), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), mode="reflect" + ) + y = y.squeeze(1) + + spec = torch.view_as_real( + torch.stft( + y, + n_fft, + hop_length=hop_size, + win_length=win_size, + window=window, + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + ) + + spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) + + spec = torch.matmul(mel, spec) + spec = spectral_normalize_torch(spec) + + return spec + + +def load_wav(wav, target_sr, min_sr=16000): + if not isinstance(wav, tuple): + speech, sample_rate = torchaudio.load(wav, backend="soundfile") + else: + speech, sample_rate = wav + if isinstance(speech, np.ndarray): + speech = torch.from_numpy(speech).float().unsqueeze(0) + + if sample_rate != target_sr: + if sample_rate < min_sr: + raise ValueError( + f"Audio sample rate {sample_rate} Hz is too low. " + f"Minimum required: {min_sr} Hz, target: {target_sr} Hz. " + f"Please provide audio with sample rate >= {min_sr} Hz." + ) + speech = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_sr)(speech) + + speech = speech.to(dtype=torch.float32) + return speech + + +def extract_speech_feat(prompt_wav, feat_extractor, device): + speech = load_wav(prompt_wav, 24000) + speech_feat = feat_extractor(speech).squeeze(dim=0).transpose(0, 1).to(device) + speech_feat = speech_feat.unsqueeze(dim=0) + speech_feat_len = torch.tensor([speech_feat.shape[1]], dtype=torch.int32).to(device) + return speech_feat, speech_feat_len + + +# Adopted from https://github.com/openai/whisper/blob/main/whisper/audio.py + + +def exact_div(x, y): + assert x % y == 0 + return x // y + + +@cache +def mel_filters(device, n_mels: int) -> torch.Tensor: + """ + load the mel filterbank matrix for projecting STFT into a Mel spectrogram. + Allows decoupling librosa dependency; saved using: + + np.savez_compressed( + "mel_filters.npz", + mel_80=librosa.filters.mel(sr=16000, n_fft=400, n_mels=80), + mel_128=librosa.filters.mel(sr=16000, n_fft=400, n_mels=128), + ) + """ + assert n_mels in {80, 128}, f"Unsupported n_mels: {n_mels}" + + filters_path = os.path.join(os.path.dirname(__file__), "assets", "mel_filters.npz") + if not os.path.exists(filters_path): + source_url = "https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz" + raise FileNotFoundError( + "Missing CosyVoice3 mel filter asset:\n" + f" {filters_path}\n" + "Download it manually from:\n" + f" {source_url}\n" + "Example:\n" + f" mkdir -p {os.path.dirname(filters_path)} && " + f"curl -L {source_url} -o {filters_path}" + ) + + with np.load(filters_path, allow_pickle=False) as f: + return torch.from_numpy(f[f"mel_{n_mels}"]).to(device) + + +def log_mel_spectrogram( + audio: str | np.ndarray | torch.Tensor, + n_mels: int = 80, + padding: int = 0, + device: str | torch.device | None = None, +): + """ + Compute the log-Mel spectrogram of + + Parameters + ---------- + audio: Union[str, np.ndarray, torch.Tensor], shape = (*) + The path to audio or either a NumPy array or Tensor containing the audio waveform in 16 kHz + + n_mels: int + The number of Mel-frequency filters, only 80 and 128 are supported + + padding: int + Number of zero samples to pad to the right + + device: Optional[Union[str, torch.device]] + If given, the audio tensor is moved to this device before STFT + + Returns + ------- + torch.Tensor, shape = (n_mels, n_frames) + A Tensor that contains the Mel spectrogram + """ + N_FFT = 400 + HOP_LENGTH = 160 + + if not torch.is_tensor(audio): + raise Exception(f"audio is not tensor {type(audio)}") + + if device is not None: + audio = audio.to(device) + if padding > 0: + audio = F.pad(audio, (0, padding)) + window = torch.hann_window(N_FFT).to(audio.device) + stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True) + magnitudes = stft[..., :-1].abs() ** 2 + + filters = mel_filters(audio.device, n_mels) + mel_spec = filters @ magnitudes + + log_spec = torch.clamp(mel_spec, min=1e-10).log10() + log_spec = torch.maximum(log_spec, log_spec.max() - 8.0) + log_spec = (log_spec + 4.0) / 4.0 + return log_spec + + +def extract_speech_token(prompt_wav, speech_tokenizer_session, device): + speech = load_wav(prompt_wav, 16000) + assert speech.shape[1] / 16000 <= 30, "do not support extract speech token for audio longer than 30s" + feat = log_mel_spectrogram(speech, n_mels=128) + speech_token = ( + speech_tokenizer_session.run( + None, + { + speech_tokenizer_session.get_inputs()[0].name: feat.detach().cpu().numpy(), + speech_tokenizer_session.get_inputs()[1].name: np.array([feat.shape[2]], dtype=np.int32), + }, + )[0] + .flatten() + .tolist() + ) + speech_token = torch.tensor([speech_token], dtype=torch.int32).to(device) + speech_token_len = torch.tensor([speech_token.shape[1]], dtype=torch.int32).to(device) + return speech_token, speech_token_len + + +def extract_spk_embedding(prompt_wav, campplus_session, device): + speech = load_wav(prompt_wav, 16000) + feat = kaldi.fbank(speech, num_mel_bins=80, dither=0, sample_frequency=16000) + feat = feat - feat.mean(dim=0, keepdim=True) + embedding = ( + campplus_session.run(None, {campplus_session.get_inputs()[0].name: feat.unsqueeze(dim=0).cpu().numpy()})[0] + .flatten() + .tolist() + ) + embedding = torch.tensor([embedding]).to(device) + return embedding + + +def extract_text_token(text, tokenizer, allowed_special): + text_token = tokenizer.encode(text, allowed_special=allowed_special) + text_token = torch.tensor([text_token], dtype=torch.int32) + text_token_len = text_token.shape[1] + return text_token, text_token_len + + +def concat_text_with_prompt_ids( + text: torch.Tensor, + text_len: torch.Tensor, + prompt_text: torch.Tensor, + prompt_text_len: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + text = torch.concat([prompt_text, text], dim=1) + text_len = text_len + prompt_text_len + return text, text_len + + +def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor: + """Make mask tensor containing indices of padded part. + + See description of make_non_pad_mask. + + Args: + lengths (torch.Tensor): Batch of lengths (B,). + Returns: + torch.Tensor: Mask tensor containing indices of padded part. + + Examples: + >>> lengths = [5, 3, 2] + >>> make_pad_mask(lengths) + masks = [[0, 0, 0, 0 ,0], + [0, 0, 0, 1, 1], + [0, 0, 1, 1, 1]] + """ + batch_size = lengths.size(0) + max_len = max_len if max_len > 0 else lengths.max().item() + seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device) + seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len) + seq_length_expand = lengths.unsqueeze(-1) + mask = seq_range_expand >= seq_length_expand + return mask diff --git a/vllm_omni/model_executor/models/hunyuan_image3/__init__.py b/vllm_omni/model_executor/models/hunyuan_image3/__init__.py new file mode 100644 index 000000000..a3618aefc --- /dev/null +++ b/vllm_omni/model_executor/models/hunyuan_image3/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .hunyuan_image3 import HunyuanImage3ForConditionalGeneration + +__all__ = ["HunyuanImage3ForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/hunyuan_image3/autoencoder_kl_3d.py b/vllm_omni/model_executor/models/hunyuan_image3/autoencoder_kl_3d.py new file mode 100644 index 000000000..e87d9b8b1 --- /dev/null +++ b/vllm_omni/model_executor/models/hunyuan_image3/autoencoder_kl_3d.py @@ -0,0 +1,934 @@ +""" +Reference code +[FLUX] https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/autoencoder.py +[DCAE] https://github.com/mit-han-lab/efficientvit/blob/master/efficientvit/models/efficientvit/dc_ae.py +""" + +import math +import random +from collections.abc import Iterable +from dataclasses import dataclass + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.utils import BaseOutput +from diffusers.utils.torch_utils import randn_tensor +from einops import rearrange +from torch import Tensor, nn + + +class DiagonalGaussianDistribution: + def __init__(self, parameters: torch.Tensor, deterministic: bool = False): + if parameters.ndim == 3: + dim = 2 # (B, L, C) + elif parameters.ndim == 5 or parameters.ndim == 4: + dim = 1 # (B, C, T, H ,W) / (B, C, H, W) + else: + raise NotImplementedError + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=dim) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like( + self.mean, device=self.parameters.device, dtype=self.parameters.dtype + ) + + def sample(self, generator: torch.Generator | None = None) -> torch.FloatTensor: + # make sure sample is on the same device as the parameters and has same dtype + sample = randn_tensor( + self.mean.shape, + generator=generator, + device=self.parameters.device, + dtype=self.parameters.dtype, + ) + x = self.mean + self.std * sample + return x + + def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + else: + reduce_dim = list(range(1, self.mean.ndim)) + if other is None: + return 0.5 * torch.sum( + torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, + dim=reduce_dim, + ) + else: + return 0.5 * torch.sum( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var + - 1.0 + - self.logvar + + other.logvar, + dim=reduce_dim, + ) + + def nll(self, sample: torch.Tensor, dims: tuple[int, ...] = [1, 2, 3]) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum( + logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, + dim=dims, + ) + + def mode(self) -> torch.Tensor: + return self.mean + + +@dataclass +class DecoderOutput(BaseOutput): + sample: torch.FloatTensor + posterior: DiagonalGaussianDistribution | None = None + + +def swish(x: Tensor) -> Tensor: + return x * torch.sigmoid(x) + + +def forward_with_checkpointing(module, *inputs, use_checkpointing=False): + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if use_checkpointing: + return torch.utils.checkpoint.checkpoint(create_custom_forward(module), *inputs, use_reentrant=False) + else: + return module(*inputs) + + +class Conv3d(nn.Conv3d): + """Perform Conv3d on patches with numerical differences from nn.Conv3d within 1e-5. + Only symmetric padding is supported. + """ + + def forward(self, input): + B, C, T, H, W = input.shape + memory_count = (C * T * H * W) * 2 / 1024**3 + if memory_count > 2: + n_split = math.ceil(memory_count / 2) + assert n_split >= 2 + chunks = torch.chunk(input, chunks=n_split, dim=-3) + padded_chunks = [] + for i in range(len(chunks)): + if self.padding[0] > 0: + padded_chunk = F.pad( + chunks[i], + (0, 0, 0, 0, self.padding[0], self.padding[0]), + mode="constant" if self.padding_mode == "zeros" else self.padding_mode, + value=0, + ) + if i > 0: + padded_chunk[:, :, : self.padding[0]] = chunks[i - 1][:, :, -self.padding[0] :] + if i < len(chunks) - 1: + padded_chunk[:, :, -self.padding[0] :] = chunks[i + 1][:, :, : self.padding[0]] + else: + padded_chunk = chunks[i] + padded_chunks.append(padded_chunk) + padding_bak = self.padding + self.padding = (0, self.padding[1], self.padding[2]) + outputs = [] + for i in range(len(padded_chunks)): + outputs.append(super().forward(padded_chunks[i])) + self.padding = padding_bak + return torch.cat(outputs, dim=-3) + else: + return super().forward(input) + + +class AttnBlock(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.in_channels = in_channels + + self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + + self.q = Conv3d(in_channels, in_channels, kernel_size=1) + self.k = Conv3d(in_channels, in_channels, kernel_size=1) + self.v = Conv3d(in_channels, in_channels, kernel_size=1) + self.proj_out = Conv3d(in_channels, in_channels, kernel_size=1) + + def attention(self, h_: Tensor) -> Tensor: + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + b, c, f, h, w = q.shape + q = rearrange(q, "b c f h w -> b 1 (f h w) c").contiguous() + k = rearrange(k, "b c f h w -> b 1 (f h w) c").contiguous() + v = rearrange(v, "b c f h w -> b 1 (f h w) c").contiguous() + h_ = nn.functional.scaled_dot_product_attention(q, k, v) + + return rearrange(h_, "b 1 (f h w) c -> b c f h w", f=f, h=h, w=w, c=c, b=b) + + def forward(self, x: Tensor) -> Tensor: + return x + self.proj_out(self.attention(x)) + + +class ResnetBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.conv1 = Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) + self.conv2 = Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if self.in_channels != self.out_channels: + self.nin_shortcut = Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x): + h = x + h = self.norm1(h) + h = swish(h) + h = self.conv1(h) + + h = self.norm2(h) + h = swish(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +class Downsample(nn.Module): + def __init__(self, in_channels: int, add_temporal_downsample: bool = True): + super().__init__() + self.add_temporal_downsample = add_temporal_downsample + stride = (2, 2, 2) if add_temporal_downsample else (1, 2, 2) # `THE`, `THAW` + # no asymmetric padding in torch conv, must do it ourselves + self.conv = Conv3d(in_channels, in_channels, kernel_size=3, stride=stride, padding=0) + + def forward(self, x: Tensor): + spatial_pad = (0, 1, 0, 1, 0, 0) # WHAT + x = nn.functional.pad(x, spatial_pad, mode="constant", value=0) + + temporal_pad = (0, 0, 0, 0, 0, 1) if self.add_temporal_downsample else (0, 0, 0, 0, 1, 1) + x = nn.functional.pad(x, temporal_pad, mode="replicate") + + x = self.conv(x) + return x + + +class DownsampleDCAE(nn.Module): + def __init__(self, in_channels: int, out_channels: int, add_temporal_downsample: bool = True): + super().__init__() + factor = 2 * 2 * 2 if add_temporal_downsample else 1 * 2 * 2 + assert out_channels % factor == 0 + self.conv = Conv3d(in_channels, out_channels // factor, kernel_size=3, stride=1, padding=1) + + self.add_temporal_downsample = add_temporal_downsample + self.group_size = factor * in_channels // out_channels + + def forward(self, x: Tensor): + r1 = 2 if self.add_temporal_downsample else 1 + h = self.conv(x) + h = rearrange(h, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + shortcut = rearrange(x, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + + B, C, T, H, W = shortcut.shape + shortcut = shortcut.view(B, h.shape[1], self.group_size, T, H, W).mean(dim=2) + return h + shortcut + + +class Upsample(nn.Module): + def __init__(self, in_channels: int, add_temporal_upsample: bool = True): + super().__init__() + self.add_temporal_upsample = add_temporal_upsample + self.scale_factor = (2, 2, 2) if add_temporal_upsample else (1, 2, 2) # `THE`, `THAW` + self.conv = Conv3d(in_channels, in_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x: Tensor): + x = nn.functional.interpolate(x, scale_factor=self.scale_factor, mode="nearest") + x = self.conv(x) + return x + + +class UpsampleDCAE(nn.Module): + def __init__(self, in_channels: int, out_channels: int, add_temporal_upsample: bool = True): + super().__init__() + factor = 2 * 2 * 2 if add_temporal_upsample else 1 * 2 * 2 + self.conv = Conv3d(in_channels, out_channels * factor, kernel_size=3, stride=1, padding=1) + + self.add_temporal_upsample = add_temporal_upsample + self.repeats = factor * out_channels // in_channels + + def forward(self, x: Tensor): + r1 = 2 if self.add_temporal_upsample else 1 + h = self.conv(x) + h = rearrange(h, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + shortcut = x.repeat_interleave(repeats=self.repeats, dim=1) + shortcut = rearrange(shortcut, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + return h + shortcut + + +class Encoder(nn.Module): + def __init__( + self, + in_channels: int, + z_channels: int, + block_out_channels: tuple[int, ...], + num_res_blocks: int, + ffactor_spatial: int, + ffactor_temporal: int, + downsample_match_channel: bool = True, + ): + super().__init__() + assert block_out_channels[-1] % (2 * z_channels) == 0 + + self.z_channels = z_channels + self.block_out_channels = block_out_channels + self.num_res_blocks = num_res_blocks + + # downsampling + self.conv_in = Conv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) + + self.down = nn.ModuleList() + block_in = block_out_channels[0] + for i_level, ch in enumerate(block_out_channels): + block = nn.ModuleList() + block_out = ch + for _ in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + down = nn.Module() + down.block = block + + add_spatial_downsample = bool(i_level < np.log2(ffactor_spatial)) + add_temporal_downsample = add_spatial_downsample and bool( + i_level >= np.log2(ffactor_spatial // ffactor_temporal) + ) + if add_spatial_downsample or add_temporal_downsample: + assert i_level < len(block_out_channels) - 1 + block_out = block_out_channels[i_level + 1] if downsample_match_channel else block_in + down.downsample = DownsampleDCAE(block_in, block_out, add_temporal_downsample) + block_in = block_out + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = Conv3d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x: Tensor) -> Tensor: + with torch.no_grad(): + use_checkpointing = bool(self.training and self.gradient_checkpointing) + + # downsampling + h = self.conv_in(x) + for i_level in range(len(self.block_out_channels)): + for i_block in range(self.num_res_blocks): + h = forward_with_checkpointing( + self.down[i_level].block[i_block], h, use_checkpointing=use_checkpointing + ) + if hasattr(self.down[i_level], "downsample"): + h = forward_with_checkpointing( + self.down[i_level].downsample, h, use_checkpointing=use_checkpointing + ) + + # middle + h = forward_with_checkpointing(self.mid.block_1, h, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing(self.mid.attn_1, h, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing(self.mid.block_2, h, use_checkpointing=use_checkpointing) + + # end + group_size = self.block_out_channels[-1] // (2 * self.z_channels) + shortcut = rearrange(h, "b (c r) f h w -> b c r f h w", r=group_size).mean(dim=2) + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + h += shortcut + return h + + +class Decoder(nn.Module): + def __init__( + self, + z_channels: int, + out_channels: int, + block_out_channels: tuple[int, ...], + num_res_blocks: int, + ffactor_spatial: int, + ffactor_temporal: int, + upsample_match_channel: bool = True, + ): + super().__init__() + assert block_out_channels[0] % z_channels == 0 + + self.z_channels = z_channels + self.block_out_channels = block_out_channels + self.num_res_blocks = num_res_blocks + + # z to block_in + block_in = block_out_channels[0] + self.conv_in = Conv3d(z_channels, block_in, kernel_size=3, stride=1, padding=1) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + # upsampling + self.up = nn.ModuleList() + for i_level, ch in enumerate(block_out_channels): + block = nn.ModuleList() + block_out = ch + for _ in range(self.num_res_blocks + 1): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + up = nn.Module() + up.block = block + + add_spatial_upsample = bool(i_level < np.log2(ffactor_spatial)) + add_temporal_upsample = bool(i_level < np.log2(ffactor_temporal)) + if add_spatial_upsample or add_temporal_upsample: + assert i_level < len(block_out_channels) - 1 + block_out = block_out_channels[i_level + 1] if upsample_match_channel else block_in + up.upsample = UpsampleDCAE(block_in, block_out, add_temporal_upsample) + block_in = block_out + self.up.append(up) + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = Conv3d(block_in, out_channels, kernel_size=3, stride=1, padding=1) + + self.gradient_checkpointing = False + + def forward(self, z: Tensor) -> Tensor: + with torch.no_grad(): + use_checkpointing = bool(self.training and self.gradient_checkpointing) + # z to block_in + repeats = self.block_out_channels[0] // (self.z_channels) + h = self.conv_in(z) + z.repeat_interleave(repeats=repeats, dim=1) + # middle + h = forward_with_checkpointing(self.mid.block_1, h, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing(self.mid.attn_1, h, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing(self.mid.block_2, h, use_checkpointing=use_checkpointing) + # upsampling + for i_level in range(len(self.block_out_channels)): + for i_block in range(self.num_res_blocks + 1): + h = forward_with_checkpointing( + self.up[i_level].block[i_block], h, use_checkpointing=use_checkpointing + ) + if hasattr(self.up[i_level], "upsample"): + h = forward_with_checkpointing(self.up[i_level].upsample, h, use_checkpointing=use_checkpointing) + # end + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + return h + + +class AutoencoderKLConv3D(ModelMixin, ConfigMixin): + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int, + out_channels: int, + latent_channels: int, + block_out_channels: tuple[int, ...], + layers_per_block: int, + ffactor_spatial: int, + ffactor_temporal: int, + sample_size: int, + sample_tsize: int, + scaling_factor: float = None, + shift_factor: float | None = None, + downsample_match_channel: bool = True, + upsample_match_channel: bool = True, + only_encoder: bool = False, + only_decoder: bool = False, + ): + super().__init__() + self.ffactor_spatial = ffactor_spatial + self.ffactor_temporal = ffactor_temporal + self.scaling_factor = scaling_factor + self.shift_factor = shift_factor + + if not only_decoder: + self.encoder = Encoder( + in_channels=in_channels, + z_channels=latent_channels, + block_out_channels=block_out_channels, + num_res_blocks=layers_per_block, + ffactor_spatial=ffactor_spatial, + ffactor_temporal=ffactor_temporal, + downsample_match_channel=downsample_match_channel, + ) + if not only_encoder: + self.decoder = Decoder( + z_channels=latent_channels, + out_channels=out_channels, + block_out_channels=list(reversed(block_out_channels)), + num_res_blocks=layers_per_block, + ffactor_spatial=ffactor_spatial, + ffactor_temporal=ffactor_temporal, + upsample_match_channel=upsample_match_channel, + ) + + self.use_slicing = False + self.slicing_bsz = 1 + self.use_spatial_tiling = False + self.use_temporal_tiling = False + self.use_tiling_during_training = False + + # only relevant if vae tiling is enabled + self.tile_sample_min_size = sample_size + self.tile_latent_min_size = sample_size // ffactor_spatial + self.tile_sample_min_tsize = sample_tsize + self.tile_latent_min_tsize = sample_tsize // ffactor_temporal + self.tile_overlap_factor = 0.125 + + self.use_compile = False + + self.empty_cache = torch.empty(0, device="cuda") + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (Encoder, Decoder)): + module.gradient_checkpointing = value + + def enable_tiling_during_training(self, use_tiling: bool = True): + self.use_tiling_during_training = use_tiling + + def disable_tiling_during_training(self): + self.enable_tiling_during_training(False) + + def enable_temporal_tiling(self, use_tiling: bool = True): + self.use_temporal_tiling = use_tiling + + def disable_temporal_tiling(self): + self.enable_temporal_tiling(False) + + def enable_spatial_tiling(self, use_tiling: bool = True): + self.use_spatial_tiling = use_tiling + + def disable_spatial_tiling(self): + self.enable_spatial_tiling(False) + + def enable_tiling(self, use_tiling: bool = True): + self.enable_spatial_tiling(use_tiling) + + def disable_tiling(self): + self.disable_spatial_tiling() + + def enable_slicing(self): + self.use_slicing = True + + def disable_slicing(self): + self.use_slicing = False + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( + x / blend_extent + ) + return b + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( + y / blend_extent + ) + return b + + def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * ( + x / blend_extent + ) + return b + + def spatial_tiled_encode(self, x: torch.Tensor): + B, C, T, H, W = x.shape + overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) # 256 * (1 - 0.25) = 192 + blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) # 8 * 0.25 = 2 + row_limit = self.tile_latent_min_size - blend_extent # 8 - 2 = 6 + + rows = [] + for i in range(0, H, overlap_size): + row = [] + for j in range(0, W, overlap_size): + tile = x[:, :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size] + tile = self.encoder(tile) + row.append(tile) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + moments = torch.cat(result_rows, dim=-2) + return moments + + def temporal_tiled_encode(self, x: torch.Tensor): + B, C, T, H, W = x.shape + overlap_size = int(self.tile_sample_min_tsize * (1 - self.tile_overlap_factor)) # 64 * (1 - 0.25) = 48 + blend_extent = int(self.tile_latent_min_tsize * self.tile_overlap_factor) # 8 * 0.25 = 2 + t_limit = self.tile_latent_min_tsize - blend_extent # 8 - 2 = 6 + + row = [] + for i in range(0, T, overlap_size): + tile = x[:, :, i : i + self.tile_sample_min_tsize, :, :] + if self.use_spatial_tiling and ( + tile.shape[-1] > self.tile_sample_min_size or tile.shape[-2] > self.tile_sample_min_size + ): + tile = self.spatial_tiled_encode(tile) + else: + tile = self.encoder(tile) + row.append(tile) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + moments = torch.cat(result_row, dim=-3) + return moments + + def spatial_tiled_decode(self, z: torch.Tensor): + B, C, T, H, W = z.shape + overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) # 24 * (1 - 0.125) = 21 + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) # 384 * 0.125 = 48 + row_limit = self.tile_sample_min_size - blend_extent # 384 - 48 = 336 + + # Distributed/multi-GPU: no padding on input -> each rank pads decoded output to right/bottom + # -> GPU all_gather -> rank0 reassembles/fuses/crops + if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: + rank = dist.get_rank() + world_size = dist.get_world_size() + + # Count tiles + num_rows = math.ceil(H / overlap_size) + num_cols = math.ceil(W / overlap_size) + total_tiles = num_rows * num_cols + tiles_per_rank = math.ceil(total_tiles / world_size) + + # Tile indices for this rank (round-robin allocation): rank, rank+world_size, ... + my_linear_indices = list(range(rank, total_tiles, world_size)) + if my_linear_indices == []: + my_linear_indices = [0] + + decoded_tiles = [] # tiles + decoded_metas = [] # (ri, rj, pad_w, pad_h) + H_out_std = self.tile_sample_min_size + W_out_std = self.tile_sample_min_size + for lin_idx in my_linear_indices: + ri = lin_idx // num_cols + rj = lin_idx % num_cols + i = ri * overlap_size + j = rj * overlap_size + tile = z[ + :, + :, + :, + i : i + self.tile_latent_min_size, + j : j + self.tile_latent_min_size, + ] + dec = self.decoder(tile) + # Pad boundary tile outputs to right/bottom to standard size + pad_h = max(0, H_out_std - dec.shape[-2]) + pad_w = max(0, W_out_std - dec.shape[-1]) + if pad_h > 0 or pad_w > 0: + dec = F.pad(dec, (0, pad_w, 0, pad_h, 0, 0), "constant", 0) + decoded_tiles.append(dec) + decoded_metas.append(torch.tensor([ri, rj, pad_w, pad_h], device=z.device, dtype=torch.int64)) + + # Different ranks may have different counts, pad to same length + T_out = decoded_tiles[0].shape[2] if len(decoded_tiles) > 0 else (T - 1) * self.ffactor_temporal + 1 + while len(decoded_tiles) < tiles_per_rank: + decoded_tiles.append( + torch.zeros( + [1, 3, T_out, self.tile_sample_min_size, self.tile_sample_min_size], + device=z.device, + dtype=dec.dtype, + ) + ) + decoded_metas.append( + torch.tensor( + [-1, -1, self.tile_sample_min_size, self.tile_sample_min_size], + device=z.device, + dtype=torch.int64, + ) + ) + + # Perform GPU all_gather + decoded_tiles = torch.stack(decoded_tiles, dim=0) + decoded_metas = torch.stack(decoded_metas, dim=0) + + tiles_gather_list = [torch.empty_like(decoded_tiles) for _ in range(world_size)] + metas_gather_list = [torch.empty_like(decoded_metas) for _ in range(world_size)] + + dist.all_gather(tiles_gather_list, decoded_tiles) + dist.all_gather(metas_gather_list, decoded_metas) + + if rank != 0: + # Non-zero ranks return empty placeholder, results are only valid on rank0 + return torch.empty(0, device=z.device) + + # rank0: Reconstruct tile grid based on (ri, rj) metadata; skip placeholder items where (ri, rj) == (-1, -1) + rows = [[None for _ in range(num_cols)] for _ in range(num_rows)] + for r in range(world_size): + gathered_tiles_r = tiles_gather_list[r] # [tiles_per_rank, B, C, T, H, W] + gathered_metas_r = metas_gather_list[r] # [tiles_per_rank, 4]，elements: (ri, rj, pad_w, pad_h) + for k in range(gathered_tiles_r.shape[0]): + ri = int(gathered_metas_r[k][0]) + rj = int(gathered_metas_r[k][1]) + if ri < 0 or rj < 0: + continue + if ri < num_rows and rj < num_cols: + # Remove padding + pad_w = int(gathered_metas_r[k][2]) + pad_h = int(gathered_metas_r[k][3]) + h_end = None if pad_h == 0 else -pad_h + w_end = None if pad_w == 0 else -pad_w + rows[ri][rj] = gathered_tiles_r[k][:, :, :, :h_end, :w_end] + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if tile is None: + continue + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + dec = torch.cat(result_rows, dim=-2) + return dec + + # Single GPU: original sequential logic + rows = [] + for i in range(0, H, overlap_size): + row = [] + for j in range(0, W, overlap_size): + tile = z[ + :, + :, + :, + i : i + self.tile_latent_min_size, + j : j + self.tile_latent_min_size, + ] + decoded = self.decoder(tile) + row.append(decoded) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=-2) + return dec + + def temporal_tiled_decode(self, z: torch.Tensor): + B, C, T, H, W = z.shape + overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor)) # 8 * (1 - 0.25) = 6 + blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor) # 64 * 0.25 = 16 + t_limit = self.tile_sample_min_tsize - blend_extent # 64 - 16 = 48 + assert 0 < overlap_size < self.tile_latent_min_tsize + + row = [] + for i in range(0, T, overlap_size): + tile = z[:, :, i : i + self.tile_latent_min_tsize, :, :] + if self.use_spatial_tiling and ( + tile.shape[-1] > self.tile_latent_min_size or tile.shape[-2] > self.tile_latent_min_size + ): + decoded = self.spatial_tiled_decode(tile) + else: + decoded = self.decoder(tile) + row.append(decoded) + + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + dec = torch.cat(result_row, dim=-3) + return dec + + def encode(self, x: Tensor, return_dict: bool = True): + def _encode(x): + if self.use_temporal_tiling and x.shape[-3] > self.tile_sample_min_tsize: + return self.temporal_tiled_encode(x) + if self.use_spatial_tiling and ( + x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size + ): + return self.spatial_tiled_encode(x) + + if self.use_compile: + + @torch.compile + def encoder(x): + return self.encoder(x) + + return encoder(x) + return self.encoder(x) + + if len(x.shape) != 5: # (B, C, T, H, W) + x = x[:, :, None] + assert len(x.shape) == 5 # (B, C, T, H, W) + if x.shape[2] == 1: + x = x.expand(-1, -1, self.ffactor_temporal, -1, -1) + else: + assert x.shape[2] != self.ffactor_temporal and x.shape[2] % self.ffactor_temporal == 0 + + if self.use_slicing and x.shape[0] > 1: + if self.slicing_bsz == 1: + encoded_slices = [_encode(x_slice) for x_slice in x.split(1)] + else: + sections = [self.slicing_bsz] * (x.shape[0] // self.slicing_bsz) + if x.shape[0] % self.slicing_bsz != 0: + sections.append(x.shape[0] % self.slicing_bsz) + encoded_slices = [_encode(x_slice) for x_slice in x.split(sections)] + h = torch.cat(encoded_slices) + else: + h = _encode(x) + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + + return AutoencoderKLOutput(latent_dist=posterior) + + def decode(self, z: Tensor, return_dict: bool = True, generator=None): + def _decode(z): + if self.use_temporal_tiling and z.shape[-3] > self.tile_latent_min_tsize: + return self.temporal_tiled_decode(z) + if self.use_spatial_tiling and ( + z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size + ): + return self.spatial_tiled_decode(z) + return self.decoder(z) + + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [_decode(z_slice) for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = _decode(z) + if torch.distributed.is_initialized(): + if torch.distributed.get_rank() != 0: + return self.empty_cache + + if z.shape[-3] == 1: + decoded = decoded[:, :, -1:] + if not return_dict: + return (decoded,) + + return DecoderOutput(sample=decoded) + + def decode_dist(self, z: Tensor, return_dict: bool = True, generator=None): + z = z.cuda() + self.use_spatial_tiling = True + decoded = self.decode(z) + self.use_spatial_tiling = False + return decoded + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_posterior: bool = True, + return_dict: bool = True, + ): + posterior = self.encode(sample).latent_dist + z = posterior.sample() if sample_posterior else posterior.mode() + dec = self.decode(z).sample + return DecoderOutput(sample=dec, posterior=posterior) if return_dict else (dec, posterior) + + def random_reset_tiling(self, x: torch.Tensor): + if x.shape[-3] == 1: + self.disable_spatial_tiling() + self.disable_temporal_tiling() + return + + # Tiling has many restrictions on input_shape and sample_size, arbitrary input_shape + # and sample_size may not meet conditions, so use fixed values here + min_sample_size = int(1 / self.tile_overlap_factor) * self.ffactor_spatial + min_sample_tsize = int(1 / self.tile_overlap_factor) * self.ffactor_temporal + sample_size = random.choice([None, 1 * min_sample_size, 2 * min_sample_size, 3 * min_sample_size]) + if sample_size is None: + self.disable_spatial_tiling() + else: + self.tile_sample_min_size = sample_size + self.tile_latent_min_size = sample_size // self.ffactor_spatial + self.enable_spatial_tiling() + + sample_tsize = random.choice([None, 1 * min_sample_tsize, 2 * min_sample_tsize, 3 * min_sample_tsize]) + if sample_tsize is None: + self.disable_temporal_tiling() + else: + self.tile_sample_min_tsize = sample_tsize + self.tile_latent_min_tsize = sample_tsize // self.ffactor_temporal + self.enable_temporal_tiling() + + +def load_weights(model, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + def update_state_dict(state_dict: dict[str, torch.Tensor], name, weight): + if name not in state_dict: + raise ValueError(f"Unexpected weight {name}") + + model_tensor = state_dict[name] + if model_tensor.shape != weight.shape: + raise ValueError( + f"Shape mismatch for weight {name}: " + f"model tensor shape {model_tensor.shape} vs. " + f"loaded tensor shape {weight.shape}" + ) + if isinstance(weight, torch.Tensor): + model_tensor.data.copy_(weight.data) + else: + raise ValueError(f"Unsupported tensor type in load_weights for {name}: {type(weight)}") + + loaded_params = set() + for name, load_tensor in weights.items(): + updated = True + name = name.replace("vae.", "") + if name in model.state_dict(): + update_state_dict(model.state_dict(), name, load_tensor) + else: + updated = False + + if updated: + loaded_params.add(name) + + return loaded_params diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py new file mode 100644 index 000000000..b015679c2 --- /dev/null +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -0,0 +1,1561 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math +import typing +from collections.abc import Callable, Iterable, Mapping, Sequence +from typing import Any, Literal, TypeAlias + +import numpy as np +import regex as re +import torch +from einops import rearrange +from PIL import Image +from torch import nn +from torchvision import transforms +from transformers import PretrainedConfig, Siglip2ImageProcessorFast +from transformers.feature_extraction_utils import BatchFeature +from transformers.image_utils import ImageInput +from transformers.tokenization_utils_base import PreTokenizedInput, TextInput +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.config.multimodal import BaseDummyOptions +from vllm.distributed import get_pp_group +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import SharedFusedMoE +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + DEFAULT_VOCAB_PADDING_SIZE, + ParallelLMHead, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.hunyuan_v1 import ( + HunYuanModel, + _get_cla_factor, + _is_moe, +) +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsLoRA, + SupportsMRoPE, + SupportsMultiModal, + SupportsPP, + _require_is_multimodal, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + _merge_multimodal_embeddings, + is_pp_missing_parameter, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.image import rgba_to_rgb +from vllm.multimodal.inputs import ( + MultiModalDataDict, + MultiModalFeatureSpec, + MultiModalFieldConfig, + MultiModalKwargsItems, +) +from vllm.multimodal.parse import ( + MultiModalDataItems, +) +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.tokenizer import get_tokenizer +from vllm.utils.tensor_schema import TensorSchema +from vllm.v1.sample.metadata import SamplingMetadata + +from vllm_omni.model_executor.models.hunyuan_image3.autoencoder_kl_3d import AutoencoderKLConv3D +from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import LightProjector, Siglip2VisionTransformer + +logger = init_logger(__name__) + + +@support_torch_compile +class HunyuanModel(HunYuanModel): + def _split_qkv_weight(self, qkv: torch.Tensor): + num_attention_heads = self.config.num_attention_heads + num_kv_heads = getattr(self.config, "num_key_value_heads", self.config.num_attention_heads) + num_key_value_groups = num_attention_heads // num_kv_heads + hidden_size = self.config.hidden_size + + if hasattr(self.config, "head_dim"): + attention_head_dim = self.config.head_dim + elif hasattr(self.config, "attention_head_dim"): + attention_head_dim = self.config.attention_head_dim + else: + attention_head_dim = self.config.hidden_size // num_attention_heads + + qkv = qkv.reshape(num_kv_heads, num_key_value_groups + 2, attention_head_dim, hidden_size) + q, k, v = torch.split(qkv, (num_key_value_groups, 1, 1), dim=1) + q = q.reshape(-1, hidden_size) + k = k.reshape(-1, hidden_size) + v = v.reshape(-1, hidden_size) + return torch.concat((q, k, v)) + + def get_expert_mapping(self) -> tuple[list[tuple[str, str, int, str]], dict[str, tuple[str, int, int]]]: + if _is_moe(self.config): + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + fused_moe_expert_mapping = SharedFusedMoE.make_expert_params_mapping( + model=self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + num_redundant_experts=self.num_redundant_experts, + ) + expert_weights_remapping = { + "gate_proj": ("gate_and_up_proj", 1, 2), + "up_proj": ("gate_and_up_proj", 0, 2), + } + return fused_moe_expert_mapping, expert_weights_remapping + else: + return [], {} + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + cla_factor = _get_cla_factor(self.config) + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + + num_attention_heads = self.config.num_attention_heads + num_kv_heads = getattr(self.config, "num_key_value_heads", self.config.num_attention_heads) + split_params_mapping = [ + (".gate_up_proj", ".gate_and_up_proj", 2, [(1, 1), (0, 1)], None), + ( + ".qkv_proj", + ".qkv_proj", + num_attention_heads + num_kv_heads * 2, + [("q", num_attention_heads), ("k", num_kv_heads), ("v", num_kv_heads)], + self._split_qkv_weight, + ), + ] + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + expert_params_mapping, expert_weights_remapping = self.get_expert_mapping() + + # List of unexpected keywords in weight names + unexpected_keywords = [ + "vae", + "vision_aligner", + "vision_model", + "final_layer", + "patch_embed", + "timestep_emb", + "time_embed", + "time_embed_2", + "guidance_emb", + "timestep_r_emb", + ] + + def contains_unexpected_keyword(name, keywords): + for keyword in keywords: + if keyword in name: + return True + return False + + for name, loaded_weight in weights: + if contains_unexpected_keyword(name, unexpected_keywords): + continue + + if "rotary_emb.inv_freq" in name: + continue + + if "gate_proj_bias" in name: + name = name.replace("gate_proj_bias", "gate_proj.bias") + if "up_proj_bias" in name: + name = name.replace("up_proj_bias", "up_proj.bias") + + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + + # With tie_word_embeddings, we can skip lm_head.weight + # The weight might appear unnecessarily in the files if the model is + # processed with quantization, LoRA, fine-tuning, etc. + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + + if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)): + # Loading kv cache scales for compressed-tensors quantization + param = params_dict[scale_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + loaded_weight = loaded_weight[0] + weight_loader(param, loaded_weight) + continue + + is_found = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + + if "mlp.experts" in name: + continue + + # cross layer only have q_proj, skip qkv pack + if weight_name == ".q_proj": + match = re.search(r"layers\.\d+", name) + if match: + layer_id = int(match.group(0).split(".")[-1]) + if cla_factor > 1 and layer_id % cla_factor != 0: + continue + + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(name) + is_found = True + break + + if is_found: + continue + + for ( + param_name, + weight_name, + den, + split_param, + func, + ) in split_params_mapping: + if weight_name not in name: + continue + + if "mlp.experts" in name: + continue + + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if is_pp_missing_parameter(name, self): + continue + + assert loaded_weight.shape[0] % den == 0 + units = loaded_weight.shape[0] // den + + param = params_dict[name] + weight_loader = param.weight_loader + offset = 0 + for shard_id, num in split_param: + new_offset = offset + num * units + if func: + weight_loader(param, func(loaded_weight)[offset:new_offset], shard_id) + else: + weight_loader(param, loaded_weight[offset:new_offset], shard_id) + offset = new_offset + + loaded_params.add(name) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + is_expert_weight = False + is_found = False + found_num = 0 + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + offset = 0 + den = 1 + for ( + mapped_weight_substr, + origin_weight_info, + ) in expert_weights_remapping.items(): + if mapped_weight_substr in weight_name: + origin_weight_name, offset, den = origin_weight_info + weight_name = weight_name.replace(mapped_weight_substr, origin_weight_name) + break + + if weight_name not in name: + continue + + # this is an expert weight and should not be + # attempted to load as other weights later + is_expert_weight = True + + # Do not modify `name` since the loop may continue here + # Instead, create a new variable + name_mapped = name.replace(weight_name, param_name) + found_num += 1 + + if is_pp_missing_parameter(name_mapped, self): + continue + + param = params_dict[name_mapped] + # We should ask the weight loader to return success or not + # here since otherwise we may skip experts with other + # available replicas. + weight_loader = typing.cast(Callable[..., bool], param.weight_loader) + assert loaded_weight.shape[0] % den == 0 + units = loaded_weight.shape[0] // den + + success = weight_loader( + param, + loaded_weight[offset * units : offset * units + units], + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + loaded_params.add(name_mapped) + is_found = True + if found_num == den: + break + + if is_found: + continue + + if is_expert_weight: + # We've checked that this is an expert weight + # However it's not mapped locally to this rank + # So we simply skip it + continue + + # Remapping the name of FP8 kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + if is_pp_missing_parameter(name, self): + continue + + if "mlp.gate.wg." in name: + name = name.replace("wg.", "") + if name == "ln_f.weight": + name = "norm.weight" + if name == "wte.weight": + name = "embed_tokens.weight" + 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 conv_nd(dims, *args, **kwargs): + """ + Create a 1D, 2D, or 3D convolution module. + """ + if dims == 1: + return nn.Conv1d(*args, **kwargs) + elif dims == 2: + return nn.Conv2d(*args, **kwargs) + elif dims == 3: + return nn.Conv3d(*args, **kwargs) + raise ValueError(f"unsupported dimensions: {dims}") + + +def linear(*args, **kwargs): + """ + Create a linear module. + """ + return nn.Linear(*args, **kwargs) + + +def avg_pool_nd(dims, *args, **kwargs): + """ + Create a 1D, 2D, or 3D average pooling module. + """ + if dims == 1: + return nn.AvgPool1d(*args, **kwargs) + elif dims == 2: + return nn.AvgPool2d(*args, **kwargs) + elif dims == 3: + return nn.AvgPool3d(*args, **kwargs) + raise ValueError(f"unsupported dimensions: {dims}") + + +def zero_module(module): + """ + Zero out the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().zero_() + return module + + +def normalization(channels, **kwargs): + """ + Make a standard normalization layer. + """ + return nn.GroupNorm(32, channels, **kwargs) + + +class Downsample(nn.Module): + """ + A downsampling layer with an optional convolution. + + :param channels: channels in the inputs and outputs. + :param use_conv: a bool determining if a convolution is applied. + :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then + downsampling occurs in the inner-two dimensions. + """ + + def __init__(self, channels, use_conv, dims=2, out_channels=None, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.dims = dims + stride = 2 if dims != 3 else (1, 2, 2) + if use_conv: + self.op = conv_nd(dims, self.channels, self.out_channels, 3, stride=stride, padding=1, **factory_kwargs) + else: + assert self.channels == self.out_channels + self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride) + + def forward(self, x): + assert x.shape[1] == self.channels + return self.op(x) + + +class ResBlock(nn.Module): + """ + A residual block that can optionally change the number of channels. + + :param in_channels: the number of input channels. + :param emb_channels: the number of timestep embedding channels. + :param dropout: the rate of dropout. + :param out_channels: if specified, the number of out channels. + :param use_conv: if True and out_channels is specified, use a spatial + convolution instead of a smaller 1x1 convolution to change the + channels in the skip connection. + :param dims: determines if the signal is 1D, 2D, or 3D. + :param down: if True, use this block for downsampling. + """ + + def __init__( + self, + in_channels, + emb_channels, + out_channels=None, + dropout=0.0, + use_conv=False, + dims=2, + down=False, + device=None, + dtype=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.in_channels = in_channels + self.dropout = dropout + self.out_channels = out_channels or self.in_channels + self.use_conv = use_conv + + self.in_layers = nn.Sequential( + normalization(self.in_channels, **factory_kwargs), + nn.SiLU(), + conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs), + ) + + self.down = down + + if down: + self.h_upd = Downsample(self.in_channels, False, dims, **factory_kwargs) + self.x_upd = Downsample(self.in_channels, False, dims, **factory_kwargs) + else: + self.h_upd = self.x_upd = nn.Identity() + + self.emb_layers = nn.Sequential(nn.SiLU(), linear(emb_channels, 2 * self.out_channels, **factory_kwargs)) + + self.out_layers = nn.Sequential( + normalization(self.out_channels, **factory_kwargs), + nn.SiLU(), + nn.Dropout(p=dropout), + zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1, **factory_kwargs)), + ) + + if self.out_channels == self.in_channels: + self.skip_connection = nn.Identity() + elif use_conv: + self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs) + else: + self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 1, **factory_kwargs) + + def forward(self, x, emb): + if self.down: + in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] + h = in_rest(x) + h = self.h_upd(h) + x = self.x_upd(x) + h = in_conv(h) + else: + h = self.in_layers(x) + + emb_out = self.emb_layers(emb) + while len(emb_out.shape) < len(h.shape): + emb_out = emb_out[..., None] + + # Adaptive Group Normalization + out_norm, out_rest = self.out_layers[0], self.out_layers[1:] + scale, shift = torch.chunk(emb_out, 2, dim=1) + h = out_norm(h) * (1.0 + scale) + shift + h = out_rest(h) + + return self.skip_connection(x) + h + + +class UNetDown(nn.Module): + """ + patch_size: one of [1, 2 ,4 ,8] + in_channels: vae latent dim + hidden_channels: hidden dim for reducing parameters + out_channels: transformer model dim + """ + + def __init__( + self, patch_size, in_channels, emb_channels, hidden_channels, out_channels, dropout=0.0, device=None, dtype=None + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + + self.patch_size = patch_size + assert self.patch_size in [1, 2, 4, 8] + + self.model = nn.ModuleList( + [ + nn.Conv2d( + in_channels=in_channels, out_channels=hidden_channels, kernel_size=3, padding=1, **factory_kwargs + ) + ] + ) + + if self.patch_size == 1: + self.model.append( + ResBlock( + in_channels=hidden_channels, + emb_channels=emb_channels, + out_channels=out_channels, + dropout=dropout, + **factory_kwargs, + ) + ) + else: + for i in range(self.patch_size // 2): + self.model.append( + ResBlock( + in_channels=hidden_channels, + emb_channels=emb_channels, + out_channels=hidden_channels if (i + 1) * 2 != self.patch_size else out_channels, + dropout=dropout, + down=True, + **factory_kwargs, + ) + ) + + def forward(self, x, t): + assert x.shape[2] % self.patch_size == 0 and x.shape[3] % self.patch_size == 0 + for module in self.model: + if isinstance(module, ResBlock): + x = module(x, t) + else: + x = module(x) + _, _, token_h, token_w = x.shape + x = rearrange(x, "b c h w -> b (h w) c") + return x, token_h, token_w + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__( + self, + hidden_size, + act_layer=nn.GELU, + frequency_embedding_size=256, + max_period=10000, + out_size=None, + dtype=None, + device=None, + ): + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.max_period = max_period + if out_size is None: + out_size = hidden_size + + self.mlp = nn.Sequential( + ColumnParallelLinear( + frequency_embedding_size, hidden_size, bias=True, gather_output=False, return_bias=False + ), + act_layer(), + RowParallelLinear(hidden_size, out_size, bias=True, input_is_parallel=True, return_bias=False), + ) + + nn.init.normal_(self.mlp[0].weight, std=0.02) + nn.init.normal_(self.mlp[2].weight, std=0.02) + + @staticmethod + def _timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + + Args: + t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional. + dim (int): the dimension of the output. + max_period (int): controls the minimum frequency of the embeddings. + + Returns: + embedding (torch.Tensor): An (N, D) Tensor of positional embeddings. + """ + # Ensure t is a tensor and get device + if not isinstance(t, torch.Tensor): + t = torch.tensor(t, dtype=torch.float32) + device = t.device + + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=device) / half + ) + # Ensure t is 2D: (N, 1) for broadcasting + if t.ndim == 0: + t = t.unsqueeze(0) + if t.ndim == 1: + t = t[:, None] + args = t.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): + t_freq = self._timestep_embedding(t, self.frequency_embedding_size, self.max_period).type( + self.mlp[0].weight.dtype + ) + t_emb = self.mlp(t_freq) + return t_emb + + +class HunyuanImage3PixelInputs(TensorSchema): + type: Literal["pixel_values"] + + pixel_values: dict[str, torch.Tensor] + + +class HunyuanImage3Processor: + """Image processor for Hunyuan Image 3.0 model.""" + + class Resolution: + def __init__(self, size, *args): + if isinstance(size, str): + if "x" in size: + size = size.split("x") + size = (int(size[0]), int(size[1])) + else: + size = int(size) + if len(args) > 0: + size = (size, args[0]) + if isinstance(size, int): + size = (size, size) + + self.h = self.height = size[0] + self.w = self.width = size[1] + self.r = self.ratio = self.height / self.width + + def __getitem__(self, idx): + if idx == 0: + return self.h + elif idx == 1: + return self.w + else: + raise IndexError(f"Index {idx} out of range") + + def __str__(self): + return f"{self.h}x{self.w}" + + class ResolutionGroup: + """Group of resolutions for image processing.""" + + def __init__(self, base_size=None, step=None, align=1): + self.align = align + self.base_size = base_size + assert base_size % align == 0, f"base_size {base_size} is not divisible by align {align}" + if base_size is not None and not isinstance(base_size, int): + raise ValueError(f"base_size must be None or int, but got {type(base_size)}") + if step is None: + step = base_size // 16 + if step is not None and step > base_size // 2: + raise ValueError(f"step must be smaller than base_size // 2, but got {step} > {base_size // 2}") + + self.step = step + self.data = self._calc_by_step() + + self.ratio = np.array([x.ratio for x in self.data]) + self.attr = ["" for _ in range(len(self.data))] + self.prefix_space = 0 + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return self.data[idx] + + def _calc_by_step(self): + assert self.align <= self.step, f"align {self.align} must be smaller than step {self.step}" + + min_height = self.base_size // 2 + min_width = self.base_size // 2 + max_height = self.base_size * 2 + max_width = self.base_size * 2 + + resolutions = [HunyuanImage3Processor.Resolution(self.base_size, self.base_size)] + + cur_height, cur_width = self.base_size, self.base_size + while True: + if cur_height >= max_height and cur_width <= min_width: + break + + cur_height = min(cur_height + self.step, max_height) + cur_width = max(cur_width - self.step, min_width) + resolutions.append( + HunyuanImage3Processor.Resolution( + cur_height // self.align * self.align, cur_width // self.align * self.align + ) + ) + + cur_height, cur_width = self.base_size, self.base_size + while True: + if cur_height <= min_height and cur_width >= max_width: + break + + cur_height = max(cur_height - self.step, min_height) + cur_width = min(cur_width + self.step, max_width) + resolutions.append( + HunyuanImage3Processor.Resolution( + cur_height // self.align * self.align, cur_width // self.align * self.align + ) + ) + + resolutions = sorted(resolutions, key=lambda x: x.ratio) + + return resolutions + + def get_target_size(self, width, height): + ratio = height / width + idx = np.argmin(np.abs(self.ratio - ratio)) + reso = self.data[idx] + return reso.w, reso.h + + def get_base_size_and_ratio_index(self, width, height): + ratio = height / width + idx = np.argmin(np.abs(self.ratio - ratio)) + return self.base_size, idx + + def __init__(self, tokenizer, hf_config, **kwargs: object): + self.tokenizer = tokenizer + self.hf_config = hf_config + self.reso_group = self.ResolutionGroup(base_size=hf_config.image_base_size) + self.vision_encoder_processor = Siglip2ImageProcessorFast.from_dict(hf_config.vit_processor) + self.vae_processor = transforms.Compose( + [ + transforms.ToTensor(), + transforms.Normalize([0.5], [0.5]), # transform to [-1, 1] + ] + ) + + def __call__( + self, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput], + images: ImageInput, + **kwargs, + ): + """ + Main method to prepare for the model one or several sequences(s) and image(s). + """ + if images is not None: + image_infors = self.process_image(images) + else: + image_infors = None + + text_inputs = self.tokenizer(text, return_tensors="pt", **kwargs) if text is not None else None + + if image_infors and text_inputs is not None: + # Combine text and image inputs into BatchFeature + combined = dict(text_inputs) + combined.update(image_infors) + return BatchFeature(combined) + elif image_infors: + return BatchFeature(image_infors) + elif text_inputs is not None: + return BatchFeature(dict(text_inputs)) + else: + return BatchFeature({}) + + def process_image(self, image_input: ImageInput): + if isinstance(image_input, Image.Image): + images = [image_input] + elif isinstance(image_input, list): + images = image_input + else: + raise TypeError(f"Unsupported image type: {type(image_input)}.") + + torch_dtype = getattr(self.hf_config, "torch_dtype", torch.bfloat16) + + batch_data = [] + for image in images: + current_info = {} + + if self.hf_config.vit["num_channels"] == 3 and image.mode == "RGBA": + image = rgba_to_rgb(image, (255, 255, 255)) + + # VIT processing + vit_pixel_values = self.vision_encoder_processor(image) + # shape: (seq_len, num_channels * patch_size * patch_size) + current_info["vit_pixel_values"] = vit_pixel_values["pixel_values"].squeeze(0) + # shape: (seq_len, ) + current_info["vit_pixel_attention_mask"] = vit_pixel_values["pixel_attention_mask"].squeeze(0) + # shape: (2, ) + current_info["vit_spatial_shapes"] = vit_pixel_values["spatial_shapes"].squeeze(0) + + # VAE processing + image_width, image_height = self.reso_group.get_target_size(image.width, image.height) + resized_image = self._resize_and_crop(image, (image_width, image_height)) + vae_pixel_values = self.vae_processor(resized_image) + token_height = image_height // (self.hf_config.vae_downsample_factor[0] * self.hf_config.patch_size) + token_width = image_width // (self.hf_config.vae_downsample_factor[1] * self.hf_config.patch_size) + current_info["vae_pixel_values"] = vae_pixel_values.squeeze(0).to(dtype=torch_dtype) + current_info["vae_token_grid_hw"] = torch.tensor([token_height, token_width]) + + # size + base_size, ratio_index = self.reso_group.get_base_size_and_ratio_index(image_width, image_height) + current_info["base_size"] = torch.tensor(base_size) + current_info["ratio_index"] = torch.tensor(ratio_index) + + batch_data.append(current_info) + + # Stack the tensors in the list into a batch dimension (B, ...) + final_image_info = {} + if len(batch_data) > 0: + for key in batch_data[0].keys(): + final_image_info[key] = torch.stack([d[key] for d in batch_data], dim=0) + + if final_image_info: + shapes_info = {k: tuple(v.shape) for k, v in final_image_info.items()} + logger.info(f"Successfully processed {len(images)} image(s). Final tensor shapes: {shapes_info}") + + return final_image_info + + def _resize_and_crop(self, image: Image.Image, target_size: tuple[int, int]) -> Image.Image: + tw, th = target_size + w, h = image.size + + tr = th / tw + r = h / w + + # resize + if r < tr: + resize_height = th + resize_width = int(round(th / h * w)) + else: + resize_width = tw + resize_height = int(round(tw / w * h)) + + image = image.resize((resize_width, resize_height), resample=Image.Resampling.LANCZOS) + + # center crop + crop_top = int(round((resize_height - th) / 2.0)) + crop_left = int(round((resize_width - tw) / 2.0)) + + image = image.crop((crop_left, crop_top, crop_left + tw, crop_top + th)) + return image + + +class HunyuanImage3ProcessingInfo(BaseProcessingInfo): + """Processing information for HunyuanImage3 model.""" + + def get_image_processor(self, **kwargs: object) -> HunyuanImage3Processor: + tokenizer = self.get_tokenizer() + hf_config = self.get_hf_config() + + return HunyuanImage3Processor( + hf_config=hf_config, + tokenizer=tokenizer, + **kwargs, + ) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"image": None} + + +class HunyuanImage3DummyInputsBuilder(BaseDummyInputsBuilder[HunyuanImage3ProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + """Generate dummy text with image placeholders.""" + num_images = mm_counts.get("image", 0) + return "" * num_images + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions] | None = None, + ) -> MultiModalDataDict: + num_images = mm_counts.get("image", 0) + + hf_config = self.info.get_hf_config() + image_size = hf_config.image_base_size + image_overrides = mm_options.get("image") if mm_options else None + + return { + "image": self._get_dummy_images( + width=image_size, + height=image_size, + num_images=num_images, + overrides=image_overrides, + ), + } + + +class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3ProcessingInfo]): + """Multimodal processor for HunyuanImage3 model.""" + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + image_processor = self.info.get_image_processor(**mm_kwargs) + images = mm_data.get("images", []) + logger.debug(f"process image count: {len(images)}") + batch_feature = image_processor(prompt, images, **tok_kwargs) + return batch_feature + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return False + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + config = {} + + if "vit_pixel_values" in hf_inputs: + config["vit_pixel_values"] = MultiModalFieldConfig.batched("image") + if "vit_pixel_attention_mask" in hf_inputs: + config["vit_pixel_attention_mask"] = MultiModalFieldConfig.batched("image") + if "vit_spatial_shapes" in hf_inputs: + config["vit_spatial_shapes"] = MultiModalFieldConfig.batched("image") + if "vae_pixel_values" in hf_inputs: + config["vae_pixel_values"] = MultiModalFieldConfig.batched("image") + if "vae_token_grid_hw" in hf_inputs: + config["vae_token_grid_hw"] = MultiModalFieldConfig.batched("image") + if "base_size" in hf_inputs: + config["base_size"] = MultiModalFieldConfig.batched("image") + if "ratio_index" in hf_inputs: + config["ratio_index"] = MultiModalFieldConfig.batched("image") + return config + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + """Replace image placeholders with the correct number of tokens.""" + + tokenizer = self.info.get_tokenizer() + img_token_id = tokenizer.convert_tokens_to_ids("") + if img_token_id is None: + raise ValueError("Image token '' not found in tokenizer vocabulary") + joint_img_sep_token_id = tokenizer.convert_tokens_to_ids("") + if joint_img_sep_token_id is None: + raise ValueError("Joint image separator token '' not found in tokenizer vocabulary") + boi_token_id = tokenizer.convert_tokens_to_ids("") + if boi_token_id is None: + raise ValueError("Beginning of image token '' not found in tokenizer vocabulary") + eoi_token_id = tokenizer.convert_tokens_to_ids("") + if eoi_token_id is None: + raise ValueError("End of image token '' not found in tokenizer vocabulary") + + out_mm_data = out_mm_kwargs.get_data() + vit_spatial_shapes = out_mm_data.get("vit_spatial_shapes") + vae_token_grid_hw = out_mm_data.get("vae_token_grid_hw") + base_size = out_mm_data.get("base_size") + ratio_index = out_mm_data.get("ratio_index") + + def get_replacement_image(item_idx: int) -> PromptUpdateDetails: + _vit_token_grid_hw = vit_spatial_shapes.tolist()[item_idx] + _vae_token_grid_hw = vae_token_grid_hw.tolist()[item_idx] + _base_size = base_size.tolist()[item_idx] + _ratio_index = ratio_index.tolist()[item_idx] + + timestep_token_num = 1 + vae_token_num = _vae_token_grid_hw[0] * _vae_token_grid_hw[1] + hf_config = self.info.get_hf_config() + vit_token_num = hf_config.vit_processor.get("max_num_patches", 729) + + base_size_token_id = tokenizer.convert_tokens_to_ids(f"") + if base_size_token_id is None: + raise ValueError(f"Base size token '' not found in tokenizer vocabulary") + ratio_token_id = tokenizer.convert_tokens_to_ids(f"") + if ratio_token_id is None: + raise ValueError(f"Ratio token '' not found in tokenizer vocabulary") + + replacement = ( + [boi_token_id] + + [base_size_token_id] + + [ratio_token_id] + + [img_token_id] * timestep_token_num + + [img_token_id] * vae_token_num + + [joint_img_sep_token_id] + + [img_token_id] * vit_token_num + + [eoi_token_id] + ) + logger.debug(f"actual replacement token count: {timestep_token_num + vae_token_num + vit_token_num}") + return PromptUpdateDetails.select_token_id(replacement, embed_token_id=img_token_id) + + return [ + PromptReplacement(modality="image", target=[img_token_id], replacement=get_replacement_image), + ] + + +@MULTIMODAL_REGISTRY.register_processor( + HunyuanImage3MultiModalProcessor, + info=HunyuanImage3ProcessingInfo, + dummy_inputs=HunyuanImage3DummyInputsBuilder, +) +class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA, SupportsPP, SupportsMRoPE): + """ + HunyuanImage3.0 model for conditional image generation. + + This is the main entry point for HunyuanImage3.0 in vLLM. It wraps: + - HunyuanModel + - VAE Encoder (AutoencoderKLConv3D + TimestepEmbedder + UNetDown) + - ViT Encoder (Siglip2VisionTransformer + LightProjector) + - LM Head for token prediction + + Supports: + - Text-to-Text and Image-to-Text generation + - Tensor Parallelism + """ + + HunyuanImage3Inputs: TypeAlias = HunyuanImage3PixelInputs + + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + + # Use mRoPE to preserve 2D positional encoding for image tokens. + if isinstance(config.rope_parameters, dict): + config.rope_parameters["rope_type"] = "default" + head_dim = getattr( + config, + "head_dim", + getattr(config, "attention_head_dim", config.hidden_size // config.num_attention_heads), + ) + config.rope_parameters["mrope_section"] = [0, head_dim // 4, head_dim // 4] + + self.config = config + self.quant_config = quant_config + self.vllm_config = vllm_config + self.model = HunyuanModel(vllm_config=vllm_config, prefix="model") + if get_pp_group().is_last_rank: + self.unpadded_vocab_size = config.vocab_size + self.lm_head = ParallelLMHead( + self.unpadded_vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + padding_size=DEFAULT_VOCAB_PADDING_SIZE, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + if config.tie_word_embeddings: + self.lm_head.weight = self.model.embed_tokens.weight + + logit_scale = getattr(config, "logit_scale", 1.0) + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size, config.vocab_size, logit_scale) + else: + self.lm_head = PPMissingLayer() + + # vae + self.vae = AutoencoderKLConv3D.from_config(config.vae) + self.patch_embed = UNetDown( + patch_size=config.patch_size, + emb_channels=config.hidden_size, + in_channels=config.vae["latent_channels"], + hidden_channels=config.patch_embed_hidden_dim, + out_channels=config.hidden_size, + ) + + # Used when converting VAE-encoded latent space (latents) to token embeddings. + self.time_embed = TimestepEmbedder(hidden_size=config.hidden_size) + + # vision + self.vision_model = Siglip2VisionTransformer(config.vit) + self.vision_aligner = LightProjector(config.vit_aligner) + + # Used to embed timestep information into the input sequence. + self.timestep_emb = TimestepEmbedder(hidden_size=config.hidden_size) + + tokenizer = get_tokenizer(vllm_config.model_config.tokenizer) + self._mrope_img_token_id = tokenizer.convert_tokens_to_ids("") + self._mrope_boi_token_id = tokenizer.convert_tokens_to_ids("") + self._mrope_eoi_token_id = tokenizer.convert_tokens_to_ids("") + self._mrope_joint_img_sep_token_id = tokenizer.convert_tokens_to_ids("") + self._mrope_max_num_patches = config.vit_processor.get("max_num_patches", 729) + + def _parse_and_validate_image_input( + self, + **kwargs: dict[str, Any], + ) -> HunyuanImage3Inputs | None: + """ + Parse and validate image input from kwargs. + """ + vit_pixel_values = kwargs.pop("vit_pixel_values", None) + vit_pixel_attention_mask = kwargs.pop("vit_pixel_attention_mask", None) + vit_spatial_shapes = kwargs.pop("vit_spatial_shapes", None) + vae_pixel_values = kwargs.pop("vae_pixel_values", None) + vae_token_grid_hw = kwargs.pop("vae_token_grid_hw", None) + + if vit_pixel_values is None or vae_pixel_values is None: + return None + + return HunyuanImage3PixelInputs( + type="pixel_values", + pixel_values={ + "vit_pixel_values": vit_pixel_values, + "vit_pixel_attention_mask": vit_pixel_attention_mask, + "vit_spatial_shapes": vit_spatial_shapes, + "vae_pixel_values": vae_pixel_values, + "vae_token_grid_hw": vae_token_grid_hw, + }, + ) + + def _vae_encode( + self, + images: torch.Tensor, + cfg_factor: int = 1, + ) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + """ + Encode images through VAE encoder. + """ + config = self.vae.config + + vae_encode_result = self.vae.encode(images) + latents = vae_encode_result.latent_dist.sample() + + # Apply shift and scaling factors if present + if hasattr(config, "shift_factor") and config.shift_factor: + latents.sub_(config.shift_factor) + if hasattr(config, "scaling_factor") and config.scaling_factor: + latents.mul_(config.scaling_factor) + + # Handle temporal dimension if present + # from (B, C, T, H, W) to (B, C, H, W) + if latents.ndim == 5: + # Check if T dimension is 1, then squeeze it + # This matches HuggingFace implementation: assert latents.shape[2] == 1, then squeeze(2) + if latents.shape[2] == 1: + latents = latents.squeeze(2) + else: + # If T > 1, take the first frame (shouldn't happen for conditional images) + logger.warning(f"T({latents.shape[2]}) > 1, should not happen for conditional images") + latents = latents[:, :, 0, :, :] + + # Always use t=0 to declare it is a clean conditional image + t = torch.zeros((latents.shape[0],), device=latents.device, dtype=latents.dtype) + + # Apply cfg_factor if needed + if cfg_factor > 1: + t = t.repeat(cfg_factor) + latents = latents.repeat(cfg_factor, 1, 1, 1) + + return t, latents + + def _vit_encode( + self, + pixel_values: torch.Tensor, + vit_attention_mask: torch.Tensor, + vit_spatial_shapes: torch.Tensor, + ) -> torch.Tensor | None: + """ + Encode pixel_values through ViT encoder (vision_model and vision_aligner). + """ + # Handle empty batch + if pixel_values.shape[0] == 0: + return None + + vision_output = self.vision_model( + pixel_values, attention_mask=vit_attention_mask, spatial_shapes=vit_spatial_shapes + ) + image_embed = vision_output.last_hidden_state + image_embed = self.vision_aligner(image_embed) + return image_embed + + def _timestep_encode( + self, + timestep: torch.Tensor, + ) -> torch.Tensor: + """Encode timestep into embedding representation. + Args: + timestep: Timestep tensor, shape (token_size,). + Returns: + Timestep embeddings, shape (token_size, hidden_size). + """ + if timestep.ndim == 0: + t_flat = timestep.unsqueeze(0) + else: + t_flat = timestep.reshape(-1) + + t_emb = self.timestep_emb(t_flat) + return t_emb + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + """Get multimodal embeddings from input.""" + image_input = self._parse_and_validate_image_input(**kwargs) + if image_input is None: + return [] + + vae_cfg_factor = 1 + pixel_values = image_input["pixel_values"] + vit_pixel_values = pixel_values["vit_pixel_values"] + vit_pixel_attention_mask = pixel_values["vit_pixel_attention_mask"] + vit_spatial_shapes = pixel_values["vit_spatial_shapes"] + vae_pixel_values = pixel_values["vae_pixel_values"] + + # Perform ViT encoding + vit_embeddings = self._vit_encode(vit_pixel_values, vit_pixel_attention_mask, vit_spatial_shapes) + + # Perform VAE encoding + t, latents = self._vae_encode(vae_pixel_values, vae_cfg_factor) + + # Process VAE latents through patch_embed to convert to token embeddings + # VAE latents are in (B, C, H, W) format, need to be converted to (B, seq_len, hidden_size) + vae_token_embeddings = [] + batch_size = latents.shape[0] + for i in range(batch_size): + t_i = t[i] + latents_i = latents[i : i + 1] # Shape: (1, C, H, W) + + # Time embedding for VAE processing + t_emb = self.time_embed(t_i) + + # Process VAE latent through patch_embed + # Input: (1, C, H, W) -> Output: (1, seq_len, hidden_size) + vae_tokens, _, _ = self.patch_embed(latents_i, t_emb) + vae_token_embeddings.append(vae_tokens) + + assert vit_embeddings is not None and vit_embeddings.shape[0] == len(vae_token_embeddings), ( + f"Number of ViT embeddings ({vit_embeddings.shape[0]}) does not match " + f"number of VAE token embeddings ({len(vae_token_embeddings)}). " + "Each image should have both VAE and ViT embeddings." + ) + + # Order per image: timestep -> VAE tokens -> ViT tokens + combined_embeddings: list[torch.Tensor] = [] + num_images = len(vae_token_embeddings) + for img_idx in range(num_images): + # 1. Timestep embedding + timestep = torch.zeros((1,)).to(vit_embeddings.device).to(vit_embeddings.dtype) + timestep_emb = self._timestep_encode(timestep) + + # 2. VAE image token embeddings + vae_token_embed = vae_token_embeddings[img_idx] + # Remove batch dimension if present: (B, seq_len, hidden_size) -> (seq_len, hidden_size) + if vae_token_embed.ndim == 3: + vae_token_embed = vae_token_embed.squeeze(0) + + # 3. ViT image embeddings + vit_embed = vit_embeddings[img_idx] + + stacked_embed = torch.cat([timestep_emb, vae_token_embed, vit_embed], dim=0) + combined_embeddings.append(stacked_embed) + + return combined_embeddings + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + handle_oov_mm_token: bool = False, + ) -> torch.Tensor: + """Embed input IDs with optional multimodal embeddings.""" + # Get text embeddings + inputs_embeds = self.model.embed_input_ids(input_ids) + + # If no multimodal embeddings, return text embeddings + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + # Merge multimodal embeddings with text embeddings + merged_embeds = _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=_require_is_multimodal(is_multimodal), + ) + return merged_embeds + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + sampling_metadata: SamplingMetadata | None = None, + logits_index: int | None = None, + sampler=None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + model_input_ids = None if inputs_embeds is not None else input_ids + model_output = self.model(model_input_ids, positions, intermediate_tensors, inputs_embeds) + return model_output + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + return logits + + def make_empty_intermediate_tensors( + self, batch_size: int, dtype: torch.dtype, device: torch.device + ) -> IntermediateTensors: + return IntermediateTensors( + { + "hidden_states": torch.zeros((batch_size, self.config.hidden_size), dtype=dtype, device=device), + "residual": torch.zeros((batch_size, self.config.hidden_size), dtype=dtype, device=device), + } + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else [] + unexpected_keywords = [ + "final_layer", + "time_embed_2", + "guidance_emb", + "timestep_r_emb", + ] + skip_prefixes.extend(unexpected_keywords) + loader = AutoWeightsLoader( + self, + skip_prefixes=skip_prefixes, + ) + + loaded_params = loader.load_weights(weights) + return loaded_params + + def get_language_model(self) -> torch.nn.Module: + return self.model + + def get_mrope_input_positions( + self, + input_tokens: list[int], + mm_features: list[MultiModalFeatureSpec] | None = None, + *, + hf_config: PretrainedConfig, + image_grid_thw: list[list[int]] | torch.Tensor, + video_grid_thw: list[list[int]] | torch.Tensor, + second_per_grid_ts: list[float] | None = None, + context_len: int = 0, + seq_len: int | None = None, + audio_feature_lengths: torch.Tensor | None = None, + use_audio_in_video: bool = False, + ) -> tuple[torch.Tensor, int]: + """Compute mRoPE positions for HunyuanImage-3. + + Maps the original model's build_2d_rope logic into vLLM's 3-dim + mRoPE position tensor [3, seq_len] where dim-0 is temporal (unused, + kept equal to 1D), dim-1 is height, and dim-2 is width. + + For text tokens and auxiliary image tokens (timestep, ViT): + All three dims get the same flat 1D position id. + For VAE image tokens: + dim-0 (T): flat 1D position id at the image start + dim-1 (H): 2D y-position using build_2d_rope centering + dim-2 (W): 2D x-position using build_2d_rope centering + """ + + # Extract per-image VAE grid dims from mm_features + vae_grids: list[tuple[int, int]] = [] + if mm_features is not None: + for mm_feature in mm_features: + mm_item = mm_feature.data + if mm_item is None: + continue + mm_input = mm_item.get_data() + vae_hw = mm_input.get("vae_token_grid_hw") + if vae_hw is not None: + grid = vae_hw.tolist() + vae_grids.append((int(grid[0]), int(grid[1]))) + + # Identify image token ids (cached from __init__) + img_token_id = self._mrope_img_token_id + boi_token_id = self._mrope_boi_token_id + eoi_token_id = self._mrope_eoi_token_id + joint_img_sep_token_id = self._mrope_joint_img_sep_token_id + max_num_patches = self._mrope_max_num_patches + + # Build position arrays + t_pos = [] # temporal (same as 1D for this model) + h_pos = [] # height + w_pos = [] # width + + pos = 0 # current 1D position counter + image_idx = 0 + i = 0 + n = len(input_tokens) + + while i < n: + tok = input_tokens[i] + + if tok == boi_token_id: + # Found start of image block. + # Structure: *timestep *vae + # *vit + # Assign a flat position + t_pos.append(pos) + h_pos.append(pos) + w_pos.append(pos) + pos += 1 + i += 1 + + # token + if i < n: + t_pos.append(pos) + h_pos.append(pos) + w_pos.append(pos) + pos += 1 + i += 1 + + # token + if i < n: + t_pos.append(pos) + h_pos.append(pos) + w_pos.append(pos) + pos += 1 + i += 1 + + # Timestep token (1 token) + if i < n and input_tokens[i] == img_token_id: + t_pos.append(pos) + h_pos.append(pos) + w_pos.append(pos) + pos += 1 + i += 1 + + # VAE tokens: get grid dims + if image_idx < len(vae_grids): + vae_h, vae_w = vae_grids[image_idx] + else: + vae_h, vae_w = 0, 0 + image_idx += 1 + + # Assign 2D positions to VAE tokens using build_2d_rope + # centering logic + L = pos # position at start of VAE region + wh = vae_w * vae_h + beta_y = L + (wh - vae_h) / 2 + beta_x = L + (wh - vae_w) / 2 + + for row in range(vae_h): + for col in range(vae_w): + if i < n and input_tokens[i] == img_token_id: + t_pos.append(L) # temporal stays flat + h_pos.append(int(beta_y + row)) + w_pos.append(int(beta_x + col)) + i += 1 + + pos = L + wh # advance past VAE region + + # token + if i < n and input_tokens[i] == joint_img_sep_token_id: + t_pos.append(pos) + h_pos.append(pos) + w_pos.append(pos) + pos += 1 + i += 1 + + # ViT tokens (max_num_patches tokens) — flat 1D + vit_consumed = 0 + while i < n and input_tokens[i] == img_token_id and vit_consumed < max_num_patches: + t_pos.append(pos) + h_pos.append(pos) + w_pos.append(pos) + pos += 1 + i += 1 + vit_consumed += 1 + + # token + if i < n and input_tokens[i] == eoi_token_id: + t_pos.append(pos) + h_pos.append(pos) + w_pos.append(pos) + pos += 1 + i += 1 + + else: + # Regular text token — flat 1D position + t_pos.append(pos) + h_pos.append(pos) + w_pos.append(pos) + pos += 1 + i += 1 + + llm_positions = torch.tensor([t_pos, h_pos, w_pos], dtype=torch.long) + mrope_position_delta = llm_positions.max() + 1 - len(input_tokens) + + if seq_len is not None: + llm_positions = llm_positions[:, context_len:seq_len] + + return llm_positions, mrope_position_delta diff --git a/vllm_omni/model_executor/models/hunyuan_image3/siglip2.py b/vllm_omni/model_executor/models/hunyuan_image3/siglip2.py new file mode 100644 index 000000000..e608baeaa --- /dev/null +++ b/vllm_omni/model_executor/models/hunyuan_image3/siglip2.py @@ -0,0 +1,561 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/LICENSE +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import warnings + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.activations import ACT2FN +from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask +from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling + + +class Config: + def __init__(self, config): + if config is not None: + for key, value in config.items(): + setattr(self, key, value) + + def __getitem__(self, key): + return getattr(self, key, None) + + def __setitem__(self, key, value): + return setattr(self, key, value) + + +class Siglip2VisionEmbeddings(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.patch_size = config.patch_size + + self.patch_embedding = nn.Linear( + in_features=config.num_channels * self.patch_size * self.patch_size, + out_features=self.embed_dim, + ) + + self.num_patches = config.num_patches + self.position_embedding_size = int(self.num_patches**0.5) + self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim) + + @staticmethod + def resize_positional_embeddings( + positional_embeddings: torch.Tensor, + spatial_shapes: torch.LongTensor, + max_length: int, + ) -> torch.Tensor: + """ + Resize positional embeddings to image-specific size and pad to a fixed size. + Args: + positional_embeddings (`torch.Tensor`): + Position embeddings of shape (height, width, embed_dim) + spatial_shapes (`torch.LongTensor`): + Spatial shapes of shape (batch_size, 2) to resize the positional embeddings to + max_length (`int`): + Maximum length of the positional embeddings to pad resized positional embeddings to + Returns: + `torch.Tensor`: Embeddings of shape (batch_size, max_length, embed_dim) + """ + batch_size = spatial_shapes.shape[0] + embed_dim = positional_embeddings.shape[-1] + source_dtype = positional_embeddings.dtype + + resulted_positional_embeddings = torch.empty( + (batch_size, max_length, embed_dim), + device=positional_embeddings.device, + dtype=source_dtype, + ) + + # (height, width, embed_dim) -> (1, embed_dim, height, width) for interpolation + positional_embeddings = positional_embeddings.permute(2, 0, 1).unsqueeze(0) + + # Upcast to float32 on CPU because antialias is not supported for bfloat16/float16 on CPU + if positional_embeddings.device.type == "cpu": + positional_embeddings = positional_embeddings.to(torch.float32) + + for i in range(batch_size): + # (1, dim, height, width) -> (1, dim, target_height, target_width) + height, width = spatial_shapes[i] + resized_embeddings = F.interpolate( + positional_embeddings, + size=(height, width), + mode="bilinear", + align_corners=False, + antialias=True, + ) + + # (1, dim, target_height, target_width) -> (target_height * target_width, dim) + resized_embeddings = resized_embeddings.reshape(embed_dim, height * width).transpose(0, 1) + + # Cast to original dtype + resized_embeddings = resized_embeddings.to(source_dtype) + + resulted_positional_embeddings[i, : height * width] = resized_embeddings + resulted_positional_embeddings[i, height * width :] = resized_embeddings[0] + + return resulted_positional_embeddings + + def forward(self, pixel_values: torch.FloatTensor, spatial_shapes: torch.LongTensor) -> torch.Tensor: + """ + Args: + pixel_values (`torch.FloatTensor`): + Pixel values of shape (batch_size, max_num_patches, num_channels * patch_size * patch_size) + spatial_shapes (`List[Tuple[int, int]]`): + Spatial shapes of shape (batch_size, 2) to resize the positional embeddings to + """ + + # Apply patch embeddings to already patchified pixel values + target_dtype = self.patch_embedding.weight.dtype + patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) + + # Get positional resized and padded positional embeddings + positional_embeddings = self.position_embedding.weight.reshape( + self.position_embedding_size, self.position_embedding_size, -1 + ) + resized_positional_embeddings = self.resize_positional_embeddings( + positional_embeddings, spatial_shapes, max_length=pixel_values.shape[1] + ) + + # Add positional embeddings to patch embeddings + embeddings = patch_embeds + resized_positional_embeddings + return embeddings + + +class Siglip2Attention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {self.num_heads})." + ) + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Input shape: Batch x Time x Channel""" + + batch_size, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2) + + k_v_seq_len = key_states.shape[-2] + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scale + + if attn_weights.size() != (batch_size, self.num_heads, q_len, k_v_seq_len): + raise ValueError( + f"Attention weights should be of size {(batch_size, self.num_heads, q_len, k_v_seq_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: + if attention_mask.size() != (batch_size, 1, q_len, k_v_seq_len): + raise ValueError( + f"Attention mask should be of size {(batch_size, 1, q_len, k_v_seq_len)}, " + f"but is {attention_mask.size()}" + ) + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (batch_size, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(batch_size, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim) + + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights + + +class Siglip2SdpaAttention(Siglip2Attention): + """ + Siglip2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from + `Siglip2Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt + to SDPA API. + """ + + is_causal = False + + # Adapted from Siglip2Attention.forward and transformers.models.llama.modeling_llama.LlamaSdpaAttention.forward + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if output_attentions: + # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` + # once this is implemented. + warnings.warn( + "Siglip2Model is using Siglip2SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` " + "does not support `output_attentions=True`. Falling back to the manual attention implementation, " + "but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. " + 'This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + return super().forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + ) + + batch_size, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2) + + # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with + # custom attn_mask, + # Reference: https://github.com/pytorch/pytorch/issues/112577. + if query_states.device.type == "cuda" and attention_mask is not None: + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an + # inline conditional assignment in SDPA to support both torch.compile's dynamic shapes and full graph options. + # An inline conditional prevents dynamic shapes from compiling. + is_causal = True if self.is_causal and q_len > 1 else False + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=attention_mask, + dropout_p=self.dropout if self.training else 0.0, + is_causal=is_causal, + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(batch_size, q_len, self.embed_dim) + + attn_output = self.out_proj(attn_output) + + return attn_output, None + + +class Siglip2MLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +class Siglip2EncoderLayer(nn.Module): + def __init__(self, config): + super().__init__() + self.embed_dim = config.hidden_size + self.self_attn = Siglip2Attention(config=config) + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.mlp = Siglip2MLP(config) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + + # Ignore copy + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + output_attentions: bool | None = False, + ) -> tuple[torch.FloatTensor]: + """ + Args: + hidden_states (`torch.FloatTensor`): + Input to the layer of shape `(batch, seq_len, embed_dim)`. + attention_mask (`torch.FloatTensor`): + Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very + large negative values. + output_attentions (`bool`, *optional*, defaults to `False`): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states, attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + return outputs + + +class Siglip2Encoder(nn.Module): + """ + Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a + [`Siglip2EncoderLayer`]. + Args: + config: Siglip2Config + """ + + def __init__(self, config): + super().__init__() + self.config = config + self.layers = nn.ModuleList([Siglip2EncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = True + + # Ignore copy + def forward( + self, + inputs_embeds, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> tuple | BaseModelOutput: + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. + This is useful if you want more control over how to convert `input_ids` indices into associated vectors + than the model's internal embedding lookup matrix. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + [What are attention masks?](../glossary#attention-mask) + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + encoder_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + hidden_states = inputs_embeds + for layer_index, encoder_layer in enumerate(self.layers): # len(self.layers): 27 + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + output_attentions=output_attentions, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None) + return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions) + + +class Siglip2MultiheadAttentionPoolingHead(nn.Module): + """Multihead Attention Pooling.""" + + def __init__(self, config): + super().__init__() + + self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size)) + self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True) + self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.mlp = Siglip2MLP(config) + self.num_heads = config.num_attention_heads + + def forward(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor | None = None): + batch_size = hidden_state.shape[0] + probe = self.probe.repeat(batch_size, 1, 1) + + if attention_mask is not None: + target_len, source_len = probe.shape[1], hidden_state.shape[1] + attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_state.dtype, target_len) + attention_mask = attention_mask.repeat(1, self.num_heads, target_len, 1) + attention_mask = attention_mask.reshape(-1, target_len, source_len) + + hidden_state = self.attention(probe, hidden_state, hidden_state, attn_mask=attention_mask)[0] + + residual = hidden_state + hidden_state = self.layernorm(hidden_state) + hidden_state = residual + self.mlp(hidden_state) + + return hidden_state[:, 0] + + +class Siglip2VisionTransformer(nn.Module): + def __init__(self, config): + super().__init__() + config = Config(config) + self.config = config + embed_dim = config.hidden_size + + self.embeddings = Siglip2VisionEmbeddings(config) + self.encoder = Siglip2Encoder(config) + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + self.use_head = True if not hasattr(config, "vision_use_head") else config.vision_use_head + if self.use_head: + self.head = Siglip2MultiheadAttentionPoolingHead(config) + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + + def forward( + self, + pixel_values: torch.FloatTensor, + attention_mask: torch.Tensor, + spatial_shapes: torch.LongTensor, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> tuple | BaseModelOutputWithPooling: + r""" + Returns: + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + hidden_states = self.embeddings(pixel_values, spatial_shapes) + + if attention_mask is not None and not self._use_flash_attention_2: + # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len] + encoder_attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + else: + encoder_attention_mask = attention_mask + + encoder_outputs = self.encoder( + inputs_embeds=hidden_states, + attention_mask=encoder_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + last_hidden_state = encoder_outputs[0] + last_hidden_state = self.post_layernorm(last_hidden_state) + + pooler_output = self.head(last_hidden_state, attention_mask) if self.use_head else None + if not return_dict: + return (last_hidden_state, pooler_output) + encoder_outputs[1:] + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooler_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +class LightProjector(nn.Module): + def __init__(self, config): + config = Config(config) + super().__init__() + + if config.projector_type == "linear": + modules = nn.Linear(config.input_dim, config.n_embed) + + elif config.projector_type == "mlp_gelu": + modules = [nn.Linear(config.input_dim, config.n_embed)] + for _ in range(1, config.depth): + modules.append(nn.GELU()) + modules.append(nn.Linear(config.n_embed, config.n_embed)) + modules = nn.Sequential(*modules) + + else: + raise ValueError(f"Unknown projector type: {config.projector_type}") + + self.layers = modules + + def forward(self, x): + return self.layers(x) + + +__all__ = [ + "Siglip2VisionTransformer", + "LightProjector", +] diff --git a/vllm_omni/model_executor/models/mammoth_moda2/__init__.py b/vllm_omni/model_executor/models/mammoth_moda2/__init__.py new file mode 100644 index 000000000..80e7c2263 --- /dev/null +++ b/vllm_omni/model_executor/models/mammoth_moda2/__init__.py @@ -0,0 +1,17 @@ +from transformers import AutoTokenizer + +from vllm_omni.tokenizers.mammoth_moda2_tokenizer import MammothUTokenizer +from vllm_omni.transformers_utils.configs.mammoth_moda2 import ( # noqa: F401 registers AutoConfig + Mammothmoda2Config, + Mammothmoda2Qwen2_5_VLConfig, +) + +from .mammoth_moda2 import MammothModa2ARForConditionalGeneration + +# AutoConfig.register is done inside transformers_utils/configs/mammoth_moda2.py +AutoTokenizer.register(config_class=Mammothmoda2Config, slow_tokenizer_class=MammothUTokenizer) +AutoTokenizer.register(config_class=Mammothmoda2Qwen2_5_VLConfig, slow_tokenizer_class=MammothUTokenizer) + +__all__ = [ + "MammothModa2ARForConditionalGeneration", +] diff --git a/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py b/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py new file mode 100644 index 000000000..be56a14a6 --- /dev/null +++ b/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py @@ -0,0 +1,826 @@ +from __future__ import annotations + +from collections.abc import Callable, Iterable, Mapping +from itertools import islice +from typing import Any + +import torch +from torch import nn +from transformers import Qwen2Config +from transformers.models.qwen2_5_vl.processing_qwen2_5_vl import Qwen2_5_VLProcessor +from vllm.config import CacheConfig, VllmConfig +from vllm.distributed import get_pp_group +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP +from vllm.model_executor.models.qwen2 import Qwen2DecoderLayer, Qwen2MLP +from vllm.model_executor.models.qwen2_5_vl import ( + Qwen2_5_VLDummyInputsBuilder, + Qwen2_5_VLForConditionalGeneration, + Qwen2_5_VLMultiModalProcessor, + Qwen2_5_VLProcessingInfo, +) +from vllm.model_executor.models.qwen2_vl import Qwen2VLMultiModalDataParser +from vllm.model_executor.models.utils import ( + PPMissingLayer, + WeightsMapper, + init_vllm_registered_model, + is_pp_missing_parameter, + make_layers, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.config import ( + is_interleaved, + patch_rope_parameters, + set_default_rope_theta, +) + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights +from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config + + +def moe_enable(moe_type, layer_type, layer_idx) -> bool: + """Determine if MoE should be enabled for a specific layer type and index. + + Args: + moe_type (str): The MoE configuration string (e.g., "ffn", "ffn_attention-14:28"). + layer_type (str): The type of layer being checked (e.g., "ffn", "attention"). + layer_idx (int): The index of the current layer. + + Returns: + bool: True if MoE should be enabled for this layer, False otherwise. + """ + if ":" in moe_type: + # moe_type like ffn_attention-14:28 + moe_type, layers = moe_type.split("-") + start, end = [int(n) for n in layers.split(":")] + else: + start, end = 0, float("inf") + assert moe_type in ["none", "attention", "ffn", "ffn_attention"] + return layer_type in moe_type and start <= layer_idx < end + + +def moe_forward( + hidden_states: torch.Tensor, + und_expert: Callable[[torch.Tensor], torch.Tensor], + gen_expert: Callable[[torch.Tensor], torch.Tensor] | None, + gen_token_mask: torch.Tensor | None = None, +) -> torch.Tensor: + """Perform Mixture-of-Experts (MoE) routing and forward pass. + + This function routes tokens to either the understanding expert (`und_expert`) or the + generation expert (`gen_expert`) based on the `gen_token_mask`. + + Routing Logic: + - If `gen_expert` is None: All tokens go to `und_expert`. + - If `gen_token_mask` is None or all False: All tokens go to `und_expert`. + - If `gen_token_mask` is all True: All tokens go to `gen_expert`. + - Otherwise (mixed batch): + - Tokens where `gen_token_mask` is True go to `gen_expert`. + - Tokens where `gen_token_mask` is False go to `und_expert`. + - Results are concatenated and reordered to match the original input order. + + Args: + hidden_states (torch.Tensor): Input hidden states. Shape: `(batch_size, seq_len, hidden_size)` + or `(num_tokens, hidden_size)`. + und_expert (Callable): The expert module for understanding/text tokens. + Takes `(N, D)` tensor, returns `(N, D_out)`. + gen_expert (Callable | None): The expert module for generation/image tokens, or None. + If provided, takes `(N, D)` tensor, returns `(N, D_out)`. + gen_token_mask (torch.Tensor | None): Boolean mask indicating generation tokens. + Shape matches `hidden_states` (excluding feature dim). True for generation tokens. + + Returns: + torch.Tensor: The processed hidden states with the same shape as input (except potentially + different feature dimension if `D_out != D`). + """ + if gen_expert is None: + return und_expert(hidden_states) + + if gen_token_mask is None or not gen_token_mask.any(): + return und_expert(hidden_states) + if gen_token_mask.all(): + return gen_expert(hidden_states) + + if hidden_states.ndim == 2: + flat_hid = hidden_states + d_model = hidden_states.shape[-1] + total_tokens = hidden_states.shape[0] + elif hidden_states.ndim == 3: + d_model = hidden_states.shape[-1] + flat_hid = hidden_states.reshape(-1, d_model) # (B*L, D) + total_tokens = flat_hid.shape[0] + else: + raise ValueError(f"Unexpected hidden_states shape: {tuple(hidden_states.shape)}") + + # Validate before reshape to catch shape mismatches where numel() would + # coincidentally match after flattening dimensions of different sizes. + if gen_token_mask.numel() != total_tokens: # type: ignore[union-attr] + raise ValueError( + "gen_token_mask shape mismatch: " + f"mask={tuple(gen_token_mask.shape)}, hidden_states={tuple(hidden_states.shape)}" + ) + # mask: [num_tokens] or [B, L] -> flatten to [total_tokens] + flat_mask = gen_token_mask.reshape(-1) # type: ignore[union-attr] + gen_pos = torch.where(flat_mask)[0] + und_pos = torch.where(~flat_mask)[0] + permute_order = torch.cat([gen_pos, und_pos], dim=0) + inverse_order = torch.argsort(permute_order) + gen_token_num = int(flat_mask.sum().item()) + gen_hid, und_hid = flat_hid[permute_order].split([gen_token_num, total_tokens - gen_token_num], dim=0) + + # 1.1 Generation tokens (True) + gen_out = gen_expert(gen_hid) # (N_gen, D) + + # 1.2 Understanding tokens (False) + und_out = und_expert(und_hid) # (N_und, D) + out_dim = und_out.shape[-1] + + merged = torch.cat([gen_out, und_out], dim=0) + merged = merged[inverse_order] + + if hidden_states.ndim == 2: + return merged.view(total_tokens, out_dim).contiguous() + return merged.view(*hidden_states.shape[:-1], out_dim).contiguous() + + +class Mammothmoda2Processor(Qwen2_5_VLProcessor): + """Qwen2.5-VL Processor with MammothU tokenizer.""" + + tokenizer_class = ("MammothUTokenizer", None) + + +class MammothModa2ARProcessingInfo(Qwen2_5_VLProcessingInfo): + """Processes multi-modal information for MammothModa2 AR, returning the VL sub-configuration.""" + + def get_hf_config(self): + mammoth_cfg: Mammothmoda2Config = self.ctx.get_hf_config(Mammothmoda2Config) + llm_cfg = getattr(mammoth_cfg, "llm_config", None) + return llm_cfg + + def get_hf_processor(self, **kwargs: object) -> Mammothmoda2Processor: + return self.ctx.get_hf_processor( + Mammothmoda2Processor, + use_fast=kwargs.pop("use_fast", True), + **kwargs, + ) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + # MammothModa2 currently supports only image input, not video. + return {"image": None} + + def get_data_parser(self) -> Qwen2VLMultiModalDataParser: + # vLLM >=0.16 expects the parser to be provided by ProcessingInfo, + # not by BaseMultiModalProcessor._get_data_parser. + return Qwen2VLMultiModalDataParser( + spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size, + expected_hidden_size=self._get_expected_hidden_size(), + ) + + +class MammothModa2ARDummyInputsBuilder(Qwen2_5_VLDummyInputsBuilder): + """Reuse Qwen2.5-VL's dummy input generation logic.""" + + +class MammothModa2ARMultiModalProcessor(Qwen2_5_VLMultiModalProcessor): + """Reuse Qwen2.5-VL's multi-modal processing,""" + + +class Mammoth2DecoderLayer(Qwen2DecoderLayer): + def __init__( + self, + config: Qwen2Config, + layer_idx: int, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + # Must patch rope parameters on config BEFORE calling super().__init__, + # because Qwen2DecoderLayer.__init__ creates Qwen2Attention using the + # config's rope settings. Patching afterwards would leave the attention + # module initialised with wrong parameters, causing shape mismatches in + # rotary_embedding at runtime. + patch_rope_parameters(config) + set_default_rope_theta(config, default_theta=1000000) + super().__init__(config, cache_config, quant_config, prefix) + + self.moe_enable = moe_enable(config.moe_type, "ffn", layer_idx) + if self.moe_enable: + self.gen_mlp = Qwen2MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.gen_mlp", + ) + else: + self.gen_mlp = None + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + gen_token_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = moe_forward(hidden_states, self.mlp, self.gen_mlp, gen_token_mask) + return hidden_states, residual + + +class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP): + def __init__( + self, *, vllm_config: VllmConfig, prefix: str = "", decoder_layer_type: type[nn.Module] = Mammoth2DecoderLayer + ): + super().__init__() + + hf_config = vllm_config.model_config.hf_config + if hasattr(hf_config, "get_text_config"): + config = hf_config.get_text_config() + elif hasattr(hf_config, "text_config"): + config = hf_config.text_config + else: + config = hf_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + self.prefix = prefix + + if is_interleaved(vllm_config.model_config.hf_text_config): + assert config.max_window_layers == config.num_hidden_layers, ( + "Sliding window for some but all layers is not supported. " + f"This model uses sliding window but `max_window_layers` = {config.max_window_layers} " + f"is less than `num_hidden_layers` = {config.num_hidden_layers}. Please open an issue " + "to discuss this feature." + ) + + self.config = config + self.quant_config = quant_config + # NOTE: MammothModa2 supports extra generation vocabulary (for image tokens). + # Token ID range: [gen_vocab_start_index, gen_vocab_start_index + gen_vocab_size). + # vLLM sampler/processor expects "last dimension of logits == model_config.get_vocab_size()", + # so we output base+gen logits in compute_logits, and embeddings must accept these IDs. + self.extra_gen_vocab = bool(getattr(config, "extra_gen_vocab", False)) + # Starting index for generation tokens (used for gen_token_mask). + self.gen_vocab_start_index = getattr(hf_config, "gen_vocab_start_index", None) or getattr( + config, "gen_vocab_start_index", None + ) + self.gen_vocab_size = int(getattr(config, "gen_vocab_size", 0) or 0) + + self.base_vocab_size = int(self.gen_vocab_start_index) if self.extra_gen_vocab else int(config.vocab_size) + # The configuration level (hf_text_config.vocab_size) has been extended to base+gen + # by the upstream config class. Use config.vocab_size as the total vocab size. + self.total_vocab_size = int(getattr(config, "vocab_size", self.base_vocab_size)) + + if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank): + self.embed_tokens = VocabParallelEmbedding( + self.base_vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + if self.extra_gen_vocab: + if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank): + self.gen_embed_tokens = VocabParallelEmbedding( + self.gen_vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.gen_embed_tokens", + ) + else: + self.gen_embed_tokens = PPMissingLayer() + else: + self.gen_embed_tokens = None + + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + self.base_vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.lm_head", + ) + else: + self.lm_head = PPMissingLayer() + + if self.extra_gen_vocab: + if get_pp_group().is_last_rank: + self.gen_head = ParallelLMHead( + self.gen_vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.gen_head", + ) + else: + self.gen_head = PPMissingLayer() + else: + self.gen_head = None + + self.logits_processor = LogitsProcessor(self.base_vocab_size) + self.gen_logits_processor = LogitsProcessor(self.gen_vocab_size) if self.extra_gen_vocab else None + + decoder_layer_type = decoder_layer_type or Mammoth2DecoderLayer + + def _make_decoder_layer(*, prefix: str) -> nn.Module: + try: + layer_idx = int(prefix.rsplit(".", 1)[-1]) + except Exception: + layer_idx = 0 + return decoder_layer_type( + config=config, + layer_idx=layer_idx, + cache_config=cache_config, + quant_config=quant_config, + prefix=prefix, + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + _make_decoder_layer, + prefix=f"{prefix}.layers", + ) + + def _make_empty_intermediate_tensors( + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + return IntermediateTensors( + { + "hidden_states": torch.zeros((batch_size, config.hidden_size), dtype=dtype, device=device), + "residual": torch.zeros((batch_size, config.hidden_size), dtype=dtype, device=device), + "gen_token_mask": torch.zeros((batch_size,), dtype=torch.bool, device=device), + } + ) + + self.make_empty_intermediate_tensors = _make_empty_intermediate_tensors + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + @property + def model(self) -> MammothModa2Qwen2ForCausalLM: + return self + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + if not self.extra_gen_vocab or self.gen_embed_tokens is None: + return self.embed_tokens(input_ids) + + gen_mask = input_ids >= int(self.gen_vocab_start_index) + if not gen_mask.any(): + return self.embed_tokens(input_ids) + if gen_mask.all(): + gen_ids = input_ids - int(self.gen_vocab_start_index) + return self.gen_embed_tokens(gen_ids) + + flat_ids = input_ids.reshape(-1) + flat_mask = gen_mask.reshape(-1) + out = torch.empty( + (flat_ids.shape[0], self.config.hidden_size), + dtype=self.embed_tokens.weight.dtype, # type: ignore[attr-defined] + device=flat_ids.device, + ) + + base_pos = torch.where(~flat_mask)[0] + gen_pos = torch.where(flat_mask)[0] + if base_pos.numel() > 0: + out[base_pos] = self.embed_tokens(flat_ids[base_pos]) + if gen_pos.numel() > 0: + gen_ids = flat_ids[gen_pos] - int(self.gen_vocab_start_index) + out[gen_pos] = self.gen_embed_tokens(gen_ids) + return out.view(*input_ids.shape, -1).contiguous() + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.get_input_embeddings(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + assert input_ids is not None + hidden_states = self.get_input_embeddings(input_ids) + # gen_token_mask: True indicates image generation tokens, which use gen_mlp. + # In vLLM v1 path, only inputs_embeds might be provided, with input_ids set to None. + # In this case, gen tokens cannot be distinguished by ID, falling back to und_expert. + if self.gen_vocab_start_index is None or input_ids is None: + gen_token_mask = None + else: + gen_token_mask = input_ids >= self.gen_vocab_start_index + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + gen_token_mask = intermediate_tensors.tensors.get("gen_token_mask") + + for idx, layer in enumerate(islice(self.layers, self.start_layer, self.end_layer)): + hidden_states, residual = layer(positions, hidden_states, residual, gen_token_mask) + + if not get_pp_group().is_last_rank: + tensors = {"hidden_states": hidden_states, "residual": residual} + if gen_token_mask is not None: + tensors["gen_token_mask"] = gen_token_mask + return IntermediateTensors(tensors) + + hidden_states, _ = self.norm(hidden_states, residual) + + return hidden_states + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + if isinstance(self.lm_head, PPMissingLayer): + return None + base_logits = self.logits_processor(self.lm_head, hidden_states) + if not self.extra_gen_vocab: + return base_logits + if self.gen_head is None or isinstance(self.gen_head, PPMissingLayer): + return base_logits + assert self.gen_logits_processor is not None + gen_logits = self.gen_logits_processor(self.gen_head, hidden_states) + if base_logits is None or gen_logits is None: + return None + return torch.cat([base_logits, gen_logits], dim=-1) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)): + param = params_dict[scale_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0] + weight_loader(param, loaded_weight) + loaded_params.add(scale_name) + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + + name = name.replace(weight_name, param_name) + + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + if name.endswith("scale"): + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if weight_loader == default_weight_loader: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + if is_pp_missing_parameter(name, self): + 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 + + +@MULTIMODAL_REGISTRY.register_processor( + MammothModa2ARMultiModalProcessor, + info=MammothModa2ARProcessingInfo, + dummy_inputs=MammothModa2ARDummyInputsBuilder, +) +class MammothModa2ARForConditionalGeneration(Qwen2_5_VLForConditionalGeneration): + """Replaces the language backbone with MoE within the Qwen2_5_VLForConditionalGeneration multi-modal framework.""" + + have_multimodal_outputs = True + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + # Skip generation-side (DiT/VAE) weights as they do not belong to the AR stage. + "gen_image_condition_refiner.": None, + "gen_transformer.": None, + "gen_vae.": None, + # LLM backbone: checkpoint uses the llm_model.* prefix. + # Extra generation vocab (image tokens) weights: mapped separately to the vLLM language_model submodule. + "llm_model.model.language_model.gen_embed_tokens.": "language_model.gen_embed_tokens.", + "llm_model.gen_head.": "language_model.gen_head.", + "llm_model.model.language_model.": "language_model.", + "llm_model.model.visual.": "visual.", + "llm_model.lm_head.": "language_model.lm_head.", + } + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + # Switch hf_config to the AR sub-config to ensure the Qwen2.5-VL path receives the correct type. + mammoth_cfg = vllm_config.model_config.hf_config + ar_hf_config = getattr(mammoth_cfg, "llm_config", mammoth_cfg) + ar_vllm_config = vllm_config.with_hf_config(ar_hf_config, architectures=vllm_config.model_config.architectures) + # Initialize multi-modal components like the vision tower first. + super().__init__(vllm_config=ar_vllm_config, prefix=prefix) + # Replace with the custom MoE language model. + lm_hf_config = getattr( + ar_vllm_config.model_config.hf_config, "text_config", ar_vllm_config.model_config.hf_config + ) + self.language_model = init_vllm_registered_model( + vllm_config=ar_vllm_config, + prefix=maybe_prefix(prefix, "language_model"), + hf_config=lm_hf_config, + architectures=["MammothModa2Qwen2ForCausalLM"], + ) + self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors + + # -------- t2i (AR grid) token constraints -------- + # Constraint logic depends on per-step sampling_metadata + runtime_additional_information. + # These are passed by the vllm-omni runner via kwargs, so caching them in the model is sufficient. + self._last_runtime_additional_information: list[dict[str, Any]] | None = None + + def _apply_t2i_token_constraints(self, logits: torch.Tensor) -> torch.Tensor: + """Applies per-request token constraints. + + - For T2I requests: constrain AR grid tokens (force EOL at row end and + restrict intra-row sampling to visual token range). + - For non-T2I (text/understanding/chat) requests: disallow sampling + from the extra generation vocabulary (image tokens) to prevent + accidentally emitting visual-token sequences. + """ + if logits is None or not isinstance(logits, torch.Tensor): + return logits + + runtime_infos = self._last_runtime_additional_information + + if runtime_infos is None: + # There is no runtime info in dummy/profile run + return logits + + neg_inf = -float("inf") + num_reqs = int(logits.shape[0]) + for i in range(num_reqs): + runtime_info = runtime_infos[i] if isinstance(runtime_infos[i], dict) else {} + omni_task = runtime_info.get("omni_task") + if not isinstance(omni_task, list) or not omni_task or omni_task[0] != "t2i": + # Text/understanding/chat: forbid sampling from the extra gen vocab. + logits[i, self.language_model.base_vocab_size :] = neg_inf + continue + + ar_width = runtime_info["ar_width"][0] + eol_token_id = runtime_info["eol_token_id"][0] + visual_start = runtime_info["visual_token_start_id"][0] + visual_end = runtime_info["visual_token_end_id"][0] + generated_len = runtime_info["generated_len"] + + row = logits[i] + column_id = generated_len % (ar_width + 1) + if column_id == ar_width: + # End-of-row token: only allow eol. + eol_logit = row[eol_token_id].clone() + row.fill_(neg_inf) + row[eol_token_id] = eol_logit + else: + # Intra-row tokens: only allow visual tokens (explicitly forbid eol). + row[:visual_start] = neg_inf + row[visual_end + 1 :] = neg_inf + row[eol_token_id] = neg_inf + + return logits + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ): + # vllm-omni runner passes sampling_metadata and runtime_additional_information + # in each forward step. compute_logits is called immediately after + # forward, so caching here enables step-by-step dynamic token constraints. + runtime_infos = kwargs.get("runtime_additional_information") + self._last_runtime_additional_information = runtime_infos if isinstance(runtime_infos, list) else None + hidden_states = super().forward( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **kwargs, + ) + if isinstance(hidden_states, IntermediateTensors) and not get_pp_group().is_last_rank: + return hidden_states + # NOTE: gpu_model_runner._dummy_run performs hidden_states[logit_indices] after forward. + # We must ensure text_hidden_states is a torch.Tensor to avoid errors when + # indexing (which happens if it's a list/tuple). + if isinstance(hidden_states, IntermediateTensors): + text_hidden_states = hidden_states["hidden_states"] + out_intermediate_tensors = hidden_states + elif isinstance(hidden_states, list): + text_hidden_states = hidden_states[0] + out_intermediate_tensors = None + else: + text_hidden_states = hidden_states + out_intermediate_tensors = None + + return OmniOutput( + text_hidden_states=text_hidden_states, + multimodal_outputs={}, + intermediate_tensors=out_intermediate_tensors, + ) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput): + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + logits = super().compute_logits(hidden_states) + if isinstance(logits, torch.Tensor): + logits = self._apply_t2i_token_constraints(logits) + return logits + + +@MULTIMODAL_REGISTRY.register_processor( + MammothModa2ARMultiModalProcessor, + info=MammothModa2ARProcessingInfo, + dummy_inputs=MammothModa2ARDummyInputsBuilder, +) +class MammothModa2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP): + # Ensure vllm_omni/worker/gpu_model_runner.py's `extract_multimodal_outputs` follows + # the OmniOutput branch to retrieve text_hidden_states as a pure torch.Tensor, + # preventing errors in `hidden_states[logit_indices]` due to type mismatch (list/tuple). + have_multimodal_outputs = True + + multimodal_cpu_fields = {"image_grid_thw", "video_grid_thw"} + merge_by_field_config = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + # Consistent with Qwen2_5OmniForConditionalGeneration: instance-level flag. + self.have_multimodal_outputs = True + self.vllm_config = vllm_config + cfg = vllm_config.model_config.hf_config + self.model_stage = vllm_config.model_config.model_stage + self.multimodal_config = vllm_config.model_config.multimodal_config + + # For debugging/alignment with qwen2.5-omni: explicitly nullify unused stages. + self.ar = None + self.dit = None + self.vae = None + + if self.model_stage == "ar": + # AR stage: multi-modal + MoE text. + self.ar = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "ar"), + hf_config=cfg.llm_config if hasattr(cfg, "llm_config") else cfg.text_config, + architectures=["MammothModa2ARForConditionalGeneration"], + ) + self.model = self.ar + elif self.model_stage == "dit": + self.dit = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "dit"), + # NOTE: init_vllm_registered_model -> VllmConfig.with_hf_config requires a + # transformers.PretrainedConfig; however, Mammothmoda2Config.gen_dit_config + # is a dict (diffusers config). The DiT stage hf_config still uses the + # top-level Mammothmoda2Config, and the DiT module reads its own + # gen_dit_config / gen_vae_config dicts. + hf_config=cfg, + architectures=["MammothModa2DiTPipeline"], + ) + self.model = self.dit + elif self.model_stage == "vae": + # Reserved: VAEs not implemented yet; raise explicit error. + raise NotImplementedError("MammothModa2 VAE stage not implemented yet.") + else: + raise ValueError(f"Unsupported model_stage: {self.model_stage}") + + # Expose intermediate tensor factory for PP if provided by the submodule. + self.make_empty_intermediate_tensors = getattr(self.model, "make_empty_intermediate_tensors", lambda: None) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int): # noqa: ARG003 + return Qwen2_5_VLForConditionalGeneration.get_placeholder_str(modality, i) + + def get_language_model(self) -> nn.Module: + if hasattr(self.model, "get_language_model"): + return self.model.get_language_model() + return self.model + + def get_multimodal_embeddings(self, **kwargs: object): + # Backward compatibility: route through embed_multimodal. + return self.embed_multimodal(**kwargs) + + def embed_multimodal(self, **kwargs: object): + if hasattr(self.model, "embed_multimodal"): + return self.model.embed_multimodal(**kwargs) + if hasattr(self.model, "get_multimodal_embeddings"): + return self.model.get_multimodal_embeddings(**kwargs) + return [] + + def get_input_embeddings(self, input_ids: torch.Tensor, multimodal_embeddings=None) -> torch.Tensor: + if hasattr(self.model, "get_input_embeddings"): + return self.model.get_input_embeddings(input_ids, multimodal_embeddings=multimodal_embeddings) + # DiT stage does not consume token embeddings from `input_ids`; it uses + # condition embeddings passed via additional_information. + # However, vLLM's generation runner may still request token embeddings + # to populate `inputs_embeds` buffers, so we provide a dummy tensor. + if self.model_stage == "dit": + hidden_size = int(self.vllm_config.model_config.get_hidden_size()) + try: + target_dtype = next(self.model.parameters()).dtype + except StopIteration: + target_dtype = self.vllm_config.model_config.dtype + return torch.zeros( + (input_ids.numel(), hidden_size), + device=input_ids.device, + dtype=target_dtype, + ) + raise NotImplementedError("Underlying model does not implement get_input_embeddings") + + def forward(self, *args, **kwargs) -> OmniOutput | torch.Tensor: + out = self.model(*args, **kwargs) + if isinstance(out, OmniOutput): + return out + if isinstance(out, list): + out = out[0] + return OmniOutput(text_hidden_states=out, multimodal_outputs={}, intermediate_tensors=None) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, *args, **kwargs): + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hasattr(self.model, "compute_logits"): + return self.model.compute_logits(hidden_states) + return None + + def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]: + if self.model_stage != "dit": + raise RuntimeError( + f"get_dummy_runtime_additional_information only valid for dit stage, got {self.model_stage}" + ) + if self.dit is None: + raise RuntimeError("dit stage model is not initialized") + if not hasattr(self.dit, "get_dummy_runtime_additional_information"): + raise AttributeError("dit model missing get_dummy_runtime_additional_information") + return self.dit.get_dummy_runtime_additional_information(num_reqs) + + def load_weights(self, weights): + if self.model_stage == "ar": + if self.ar is None or not hasattr(self.ar, "load_weights"): + return set() + loaded = self.ar.load_weights(weights) + return add_prefix_to_loaded_weights(loaded, "ar") + if self.model_stage == "dit": + if self.dit is None or not hasattr(self.dit, "load_weights"): + return set() + loaded = self.dit.load_weights(weights) + return add_prefix_to_loaded_weights(loaded, "dit") + if self.model_stage == "vae": + return set() + raise ValueError(f"Unsupported model_stage: {self.model_stage}") diff --git a/vllm_omni/model_executor/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/model_executor/models/mammoth_moda2/pipeline_mammothmoda2_dit.py new file mode 100644 index 000000000..afe38f015 --- /dev/null +++ b/vllm_omni/model_executor/models/mammoth_moda2/pipeline_mammothmoda2_dit.py @@ -0,0 +1,15 @@ +""" +Compatibility shim. + +The MammothModa2 DiT implementation lives under `vllm_omni.diffusion` to align +with other ARDiT structured models. We keep this module path so existing +OmniModelRegistry entries (and downstream code) keep working. +""" + +from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import ( # noqa: F401 + MammothModa2DiTPipeline, +) + +__all__ = [ + "MammothModa2DiTPipeline", +] diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 2cf23bd92..496dfe305 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -1,4 +1,8 @@ -from vllm.model_executor.models.registry import _VLLM_MODELS, _LazyRegisteredModel, _ModelRegistry +from vllm.model_executor.models.registry import ( + _VLLM_MODELS, + _LazyRegisteredModel, + _ModelRegistry, +) _OMNI_MODELS = { "Qwen2_5OmniForConditionalGeneration": ( @@ -48,6 +52,36 @@ _OMNI_MODELS = { "qwen3_omni_code2wav", "Qwen3OmniMoeCode2Wav", ), + "CosyVoice3Model": ( + "cosyvoice3", + "cosyvoice3", + "CosyVoice3Model", + ), + "MammothModa2Qwen2ForCausalLM": ( + "mammoth_moda2", + "mammoth_moda2", + "MammothModa2Qwen2ForCausalLM", + ), + "MammothModa2ARForConditionalGeneration": ( + "mammoth_moda2", + "mammoth_moda2", + "MammothModa2ARForConditionalGeneration", + ), + "MammothModa2DiTPipeline": ( + "mammoth_moda2", + "pipeline_mammothmoda2_dit", + "MammothModa2DiTPipeline", + ), + "MammothModa2ForConditionalGeneration": ( + "mammoth_moda2", + "mammoth_moda2", + "MammothModa2ForConditionalGeneration", + ), + "Mammothmoda2Model": ( + "mammoth_moda2", + "mammoth_moda2", + "MammothModa2ForConditionalGeneration", + ), "Qwen3TTSForConditionalGeneration": ( "qwen3_tts", "qwen3_tts_talker", @@ -85,6 +119,11 @@ _OMNI_MODELS = { "glm_image_ar", "GlmImageForConditionalGeneration", ), + "HunyuanImage3ForCausalMM": ( + "hunyuan_image3", + "hunyuan_image3", + "HunyuanImage3ForConditionalGeneration", + ), } @@ -93,7 +132,6 @@ _VLLM_OMNI_MODELS = { **_OMNI_MODELS, } - OmniModelRegistry = _ModelRegistry( { **{ diff --git a/vllm_omni/model_executor/stage_configs/cosyvoice3.yaml b/vllm_omni/model_executor/stage_configs/cosyvoice3.yaml new file mode 100644 index 000000000..cbf2193de --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/cosyvoice3.yaml @@ -0,0 +1,44 @@ +# Stage config for running CosyVoice3 with 2 stage architecture +# Stage 0: Talker (text prompt → speech tokens) +# Stage 1: Code2Wav (flow matching → acoustic features → waveform) +# Right now I have coded up stage 2 in stage 1 only; can be split in future + +stage_args: + - stage_id: 0 + runtime: + devices: 0 + engine_args: + model_stage: talker + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + model_arch: CosyVoice3Model + trust_remote_code: true + gpu_memory_utilization: 0.4 + engine_output_type: latent # Output speech tokens for chunk aware flow matching + disable_hybrid_kv_cache_manager: true + enable_prefix_caching: false + enforce_eager: false + mm_processor_cache_gb: 0 + skip_mm_profiling: true + dtype: "float32" + + - stage_id: 1 + runtime: + devices: 0 + engine_args: + model_stage: code2wav + model_arch: CosyVoice3Model + trust_remote_code: true + worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + engine_output_type: latent + gpu_memory_utilization: 0.2 + enforce_eager: false # CUDA graphs don't work with dynamic runtime_info access + disable_hybrid_kv_cache_manager: true + enable_prefix_caching: false + skip_mm_profiling: true + dtype: "float32" + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.cosyvoice3.text2flow + final_output: true + final_output_type: audio diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml new file mode 100644 index 000000000..b4b918288 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml @@ -0,0 +1,42 @@ +# Stage config for running Hunyuan-Image3.0 with architecture of OmniLLM. +# Stage 0: AR Model (vLLM implementation) + +# The following config has been verified on 8x L40S-48G GPU. +stage_args: + - stage_id: 0 + stage_type: llm # Use llm stage type to launch OmniLLM + runtime: + process: true # Run this stage in a separate process + devices: "0,1,2,3,4,5,6,7" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device) + max_batch_size: 1 + engine_args: + model_stage: AR + model_arch: HunyuanImage3ForCausalMM + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.3 + enforce_eager: true # Now we only support eager mode + trust_remote_code: true + engine_output_type: latent + enable_prefix_caching: false + max_num_batched_tokens: 32768 + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + is_comprehension: true + final_output: true + final_output_type: text + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 2048 + seed: 42 + detokenize: True + repetition_penalty: 1.1 + +# Top-level runtime config (concise): default windows and stage edges +runtime: + enabled: true + defaults: + window_size: -1 # Simplified: trigger downstream only after full upstream completion + max_inflight: 1 # Simplified: process serially within each stage diff --git a/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml b/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml new file mode 100644 index 000000000..505548ecc --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml @@ -0,0 +1,42 @@ +# Stage config for running MammothModa2 with full AR→DiT pipeline (text-to-image). +# Stage 0: AR (multimodal understanding + latent generation) +# Stage 1: DiT (latent → image) +# +# For text/image understanding tasks (text output only), use mammoth_moda2_ar.yaml. + +stage_args: + - stage_id: 0 + runtime: + devices: "0" + max_batch_size: 100 + engine_args: + model_stage: ar + model_arch: MammothModa2ForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + max_model_len: 8192 + gpu_memory_utilization: 0.5 + enforce_eager: true + trust_remote_code: true + engine_output_type: latent + enable_prefix_caching: false + final_output: false + + - stage_id: 1 + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: dit + model_arch: MammothModa2ForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + gpu_memory_utilization: 0.3 + enforce_eager: true + trust_remote_code: true + engine_output_type: image + enable_prefix_caching: false + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit + final_output: true + final_output_type: image diff --git a/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml b/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml new file mode 100644 index 000000000..f95ed2b3d --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml @@ -0,0 +1,24 @@ +# Stage config for running MammothModa2 with single AR stage (understanding / summarization). +# Stage 0: AR (multimodal understanding + text generation) +# +# Use this config for image/video understanding tasks where only text output is needed. +# For text-to-image generation (AR → DiT), use mammoth_moda2.yaml instead. + +stage_args: + - stage_id: 0 + runtime: + devices: "0" + max_batch_size: 16 + engine_args: + model_stage: ar + model_arch: MammothModa2ForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + max_model_len: 8192 + gpu_memory_utilization: 0.5 + enforce_eager: true + trust_remote_code: true + engine_output_type: text + enable_prefix_caching: false + final_output: true + final_output_type: text diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py new file mode 100644 index 000000000..b7f21eca8 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py @@ -0,0 +1,32 @@ +from typing import Any + +from vllm.inputs import TextPrompt + +from vllm_omni.inputs.data import OmniTokensPrompt + + +def text2flow( + stage_list: list[Any], + engine_input_source: list[int], + prompt: OmniTokensPrompt | TextPrompt = None, + requires_multimodal_data: bool = True, +): + """Build stage-1 inputs by prefixing stage-0 prompt ids to its outputs.""" + source_stage_id = engine_input_source[0] + source_outputs = stage_list[source_stage_id].engine_outputs + + if not isinstance(prompt, list): + prompt = [prompt] + + source_output = source_outputs[0] + output = source_output.outputs[0] + + multi_modal_data = output.multimodal_output + if multi_modal_data is None: + raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}") + + output_ids = output.token_ids + prefix_ids = source_output.prompt_token_ids + multi_modal_data["prefix_ids"] = prefix_ids + engine_input = OmniTokensPrompt(prompt_token_ids=output_ids, additional_information=multi_modal_data) + return [engine_input] diff --git a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py new file mode 100644 index 000000000..1239aafaf --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py @@ -0,0 +1,99 @@ +"""Stage input processor for MammothModa2 (AR -> DiT).""" + +from typing import Any + +import torch +from vllm.inputs import TextPrompt + +from vllm_omni.inputs.data import OmniTokensPrompt + + +def ar2dit( + stage_list: list[Any], + engine_input_source: list[int], + prompts: OmniTokensPrompt | TextPrompt | None = None, + requires_multimodal_data: bool = False, # noqa: ARG001 — interface param, unused for ar2dit +) -> list[OmniTokensPrompt]: + """Convert AR stage outputs to DiT stage inputs.""" + + source_stage_id = engine_input_source[0] + ar_outputs = stage_list[source_stage_id].engine_outputs + + dit_inputs: list[OmniTokensPrompt] = [] + for ar_output, prompt in zip(ar_outputs, prompts): + addi_info = prompt["additional_information"] + image_height = addi_info["image_height"][0] + image_width = addi_info["image_width"][0] + text_guidance_scale = addi_info["text_guidance_scale"][0] + cfg_range = addi_info["cfg_range"] + num_inference_steps = addi_info["num_inference_steps"][0] + gen_vocab_start_index = addi_info["visual_token_start_id"][0] + # ["<|image_pad|>", "<|video_pad|>", "<|vision_start|>", "<|vision_end|>"] + visual_ids = addi_info["visual_ids"] + + prompt_token_ids = ar_output.prompt_token_ids + # exclude the last token because it has no corresponding hidden state + completion_output = ar_output.outputs[0] + gen_token_ids = completion_output.token_ids[:-1] + full_token_ids = prompt_token_ids + gen_token_ids + + mm_output = getattr(completion_output, "multimodal_output", None) + if not isinstance(mm_output, dict) or "latent" not in mm_output: + raise ValueError( + "AR stage output missing latent multimodal output. " + f"request_id={getattr(ar_output, 'request_id', None)}, " + f"completion_has_mm={hasattr(completion_output, 'multimodal_output')}" + ) + full_hidden_states = mm_output["latent"] + hidden_total = int(full_hidden_states.shape[0]) + assert hidden_total == len(prompt_token_ids) + len(gen_token_ids), ( + f"Hidden states length mismatch: expected {len(prompt_token_ids) + len(gen_token_ids)}, got {hidden_total}" + ) + + mask_device = full_hidden_states.device + full_token_ids_t = torch.tensor(full_token_ids, dtype=torch.long, device=mask_device) + attention_mask = torch.ones_like(full_token_ids_t, dtype=torch.bool) + + pos = torch.arange(full_token_ids_t.shape[0], device=mask_device) + answer_start_index = len(prompt_token_ids) + questions_mask = pos < answer_start_index + answers_mask = ~questions_mask + + gen_token_mask = full_token_ids_t >= gen_vocab_start_index + + visual_token_mask = torch.isin( + full_token_ids_t, + torch.tensor(visual_ids, dtype=torch.long, device=mask_device), + ) + + text_condition_token_mask = questions_mask & ~(visual_token_mask | gen_token_mask) & attention_mask + image_condition_token_mask = answers_mask & gen_token_mask & attention_mask + + text_condition = full_hidden_states[text_condition_token_mask] + image_condition = full_hidden_states[image_condition_token_mask] + + text_prompt_embeds = text_condition.to(dtype=torch.float32).contiguous() + image_prompt_embeds = image_condition.to(dtype=torch.float32).contiguous() + + additional_information = { + "text_prompt_embeds": text_prompt_embeds, + "text_prompt_embeds_shape": list(text_prompt_embeds.shape), + "image_prompt_embeds": image_prompt_embeds, + "image_prompt_embeds_shape": list(image_prompt_embeds.shape), + "image_height": [int(image_height)], + "image_width": [int(image_width)], + "text_guidance_scale": [float(text_guidance_scale)], + "cfg_range": [float(cfg_range[0]), float(cfg_range[1])], + "num_inference_steps": [int(num_inference_steps)], + } + + dit_inputs.append( + OmniTokensPrompt( + prompt_token_ids=[0], + additional_information=additional_information, + multi_modal_data=None, + mm_processor_kwargs=None, + ) + ) + + return dit_inputs diff --git a/vllm_omni/tokenizers/__init__.py b/vllm_omni/tokenizers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/tokenizers/mammoth_moda2_tokenizer.py b/vllm_omni/tokenizers/mammoth_moda2_tokenizer.py new file mode 100644 index 000000000..a419b2a92 --- /dev/null +++ b/vllm_omni/tokenizers/mammoth_moda2_tokenizer.py @@ -0,0 +1,401 @@ +# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# SPDX-License-Identifier: Apache-2.0 +# +# This file has been modified by ByteDance Ltd. on 2025-09-30. +# +# Original file was released under Apache-2.0, with the full license text +# available at https://www.apache.org/licenses/LICENSE-2.0 +# +# This modified file is released under the same license. +# +# --- Upstream header preserved below --- +# +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import base64 +import os +import unicodedata +from collections.abc import Collection + +import tiktoken +from loguru import logger +from transformers import AddedToken, PreTrainedTokenizer + +VOCAB_FILES_NAMES = { + "vocab_file": "mammothu.tiktoken", + "special_tokens_file": "mammothu_vision_tokens.txt", +} + +PAT_STR = ( + r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?""" + r"""[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""" +) +ENDOFTEXT = "<|endoftext|>" +IMSTART = "<|im_start|>" +IMEND = "<|im_end|>" +# as the default behavior is changed to allow special tokens in +# regular texts, the surface forms of special tokens need to be +# as different as possible to minimize the impact +QWEN_SPECIAL_TOKENS = ( + "<|object_ref_start|>", + "<|object_ref_end|>", + "<|box_start|>", + "<|box_end|>", + "<|quad_start|>", + "<|quad_end|>", + "<|vision_start|>", + "<|vision_end|>", + "<|vision_pad|>", + "<|image_pad|>", + "<|video_pad|>", + "", + "", + "<|fim_prefix|>", + "<|fim_middle|>", + "<|fim_suffix|>", + "<|fim_pad|>", + "<|repo_name|>", + "<|file_sep|>", +) + +# align to qwen2.5 tokenizer length (151846) +EXTRAS = [f"<|extra_{i}|>" for i in range(181)] # 205 - 19[len(QWEN_SPECIAL_TOKENS)] - 5 +# align to qwen2.5 embedding size (152064) +EXTRAS += [f"<|extra_margin_{i}|>" for i in range(152064 - 151846)] +# append new token in gen embedding range +EXTRAS += ["<|endofline|>", "<|endoffile|>", "<|gen_placeholder|>", "<|useless token|>", "<|beginoftext|>"] +EXTRAS = tuple(EXTRAS) +# changed to use actual index to avoid misconfiguration with vocabulary expansion +SPECIAL_START_ID = 151643 + + +def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> dict[bytes, int]: + with open(tiktoken_bpe_file, "rb") as f: + contents = f.read() + return { + base64.b64decode(token): int(rank) for token, rank in (line.split() for line in contents.splitlines() if line) + } + + +class MammothUTokenizer(PreTrainedTokenizer): + """MammothU tokenizer.""" + + vocab_files_names = VOCAB_FILES_NAMES + + def __init__( + self, + vocab_file: str, + special_tokens_file: str, + errors: str = "replace", + bos_token: str = "<|beginoftext|>", + eos_token: str = "<|endoftext|>", + pad_token: str = "<|endoftext|>", + img_token: str = "<|image token|>", + boi_token: str = "<|image start|>", + eoi_token: str = "<|image end|>", + eol_token: str = "<|endofline|>", + eof_token: str = "<|endoffile|>", + **kwargs, + ) -> None: + super().__init__(**kwargs) + + # how to handle errors in decoding UTF-8 byte sequences + # use ignore if you are in streaming inference + self.errors = errors + self.mergeable_ranks = _load_tiktoken_bpe(vocab_file) + with open(special_tokens_file) as f: + vision_tokens = [t.strip() for t in f.readlines() if len(t.strip()) > 0] + SPECIAL_TOKENS = tuple( + enumerate( + ( + ( + ENDOFTEXT, + IMSTART, + IMEND, + ) + + QWEN_SPECIAL_TOKENS + + EXTRAS + + tuple(vision_tokens) + ), + start=SPECIAL_START_ID, + ) + ) + + self.special_tokens = {token: index for index, token in SPECIAL_TOKENS} + self.special_tokens_set = set(t for _, t in SPECIAL_TOKENS) + + enc = tiktoken.Encoding( + "mammothu", + pat_str=PAT_STR, + mergeable_ranks=self.mergeable_ranks, + special_tokens=self.special_tokens, + ) + + assert len(self.mergeable_ranks) + len(self.special_tokens) == enc.n_vocab, ( + f"{len(self.mergeable_ranks) + len(self.special_tokens)} != {enc.n_vocab} in encoding" + ) + + self.decoder = {v: k for k, v in self.mergeable_ranks.items()} + self.decoder.update({v: k for k, v in self.special_tokens.items()}) + self.tokenizer = enc + self.eod_id = self.tokenizer.eot_token + self.bos_token = bos_token + self.eos_token = eos_token + self.pad_token = pad_token + self.img_token = img_token + self.boi_token = boi_token + self.eoi_token = eoi_token + self.eol_token = eol_token + self.eof_token = eof_token + self.image_content_token = "<|image_pad|>" # come from Qwen2.5-VL + self.gen_image_token = "<|gen_image_pad|>" + self.gen_image_placeholder_token = "<|gen_placeholder|>" + self.visual_tokens = ["<|image_pad|>", "<|video_pad|>", "<|vision_start|>", "<|vision_end|>"] + self.visual_tokens_ids = [self.get_vocab()[token] for token in self.visual_tokens] + + self.vision_range = (self.get_vocab()[self.boi_token], self.tokenizer.n_vocab - 1) + logger.info(f"MammothUTokeniser Vision range: {self.vision_range}") + + def __getstate__(self): + # for pickle lovers + state = self.__dict__.copy() + del state["tokenizer"] + return state + + def __setstate__(self, state): + # tokenizer is not python native; don't pass it; rebuild it + self.__dict__.update(state) + enc = tiktoken.Encoding( + "mammothu", + pat_str=PAT_STR, + mergeable_ranks=self.mergeable_ranks, + special_tokens=self.special_tokens, + ) + self.tokenizer = enc + + def __len__(self) -> int: + return self.tokenizer.n_vocab + + def get_vocab(self) -> dict[bytes | str, int]: + vocab = self.mergeable_ranks.copy() + vocab.update(self.special_tokens) + return vocab + + @property + def gen_placeholder_id(self): + return self.get_vocab()[self.gen_image_placeholder_token] + + def convert_tokens_to_ids(self, tokens: bytes | str | list[bytes | str]) -> list[int]: + if isinstance(tokens, (str, bytes)): + if tokens in self.special_tokens: + return self.special_tokens[tokens] + else: + return self.mergeable_ranks.get(tokens) + + ids = [] + for token in tokens: + if token in self.special_tokens: + ids.append(self.special_tokens[token]) + else: + ids.append(self.mergeable_ranks.get(token)) + return ids + + def _add_tokens( + self, + new_tokens: list[str] | list[AddedToken], + special_tokens: bool = False, + ) -> int: + if not special_tokens and new_tokens: + raise ValueError("Adding regular tokens is not supported") + + added_tokens = 0 + for token in new_tokens: + surface_form = token.content if isinstance(token, AddedToken) else token + if surface_form in self.special_tokens_set: + # Token already exists in our special tokens set + added_tokens += 1 + else: + logger.warning(f"Token {surface_form} is not in the predefined special tokens set and cannot be added") + + return added_tokens + + def add_special_tokens(self, special_tokens_dict: dict[str, str | AddedToken]) -> int: + """ + Add special tokens to the tokenizer and update the special tokens mapping. + Only adds tokens that are already in the special_tokens_set. + + Args: + special_tokens_dict: dictionary of special tokens to add. + The key is the token type and the value is the token to add. + + Returns: + Number of tokens added to the vocabulary. + """ + added_tokens = 0 + + for token_type, token in special_tokens_dict.items(): + if token_type == "additional_special_tokens" and isinstance(token, list): + added_tokens += self._add_tokens(token, special_tokens=True) + else: + token_value = token.content if isinstance(token, AddedToken) else token + if token_value in self.special_tokens_set: + setattr(self, token_type, token_value) + added_tokens += 1 + else: + logger.warning( + f"Token {token_value} is not in the predefined special tokens set and cannot be added" + ) + + return added_tokens + + def save_vocabulary(self, save_directory: str, **kwargs) -> tuple[str]: + """ + Save only the vocabulary of the tokenizer (vocabulary). + + Returns: + `tuple(str)`: Paths to the files saved. + """ + regular_file_path = os.path.join(save_directory, self.vocab_files_names["vocab_file"]) + with open(regular_file_path, "w", encoding="utf8") as w: + for k, v in self.mergeable_ranks.items(): + line = base64.b64encode(k).decode("utf8") + " " + str(v) + "\n" + w.write(line) + + excluded_special_tokens = set( + ( + ENDOFTEXT, + IMSTART, + IMEND, + ) + + EXTRAS + ) + special_file_path = os.path.join(save_directory, self.vocab_files_names["special_tokens_file"]) + with open(special_file_path, "w", encoding="utf8") as w: + for k in self.special_tokens: + if k not in excluded_special_tokens: + print(k, file=w) + + return (regular_file_path, special_file_path) + + def tokenize( + self, + text: str, + allowed_special: set | str = "all", + disallowed_special: Collection | str = (), + **kwargs, + ) -> list[bytes | str]: + """ + Converts a string in a sequence of tokens. + + Args: + text (`str`): + The sequence to be encoded. + allowed_special (`Literal["all"]` or `set`): + The surface forms of the tokens to be encoded as special tokens in regular texts. + Default to "all". + disallowed_special (`Literal["all"]` or `Collection`): + The surface forms of the tokens that should not be in regular texts and trigger errors. + Default to an empty tuple. + + kwargs (additional keyword arguments, *optional*): + Will be passed to the underlying model specific encode method. + + Returns: + `list[bytes|str]`: The list of tokens. + """ + tokens = [] + text = unicodedata.normalize("NFC", text) + + # this implementation takes a detour: text -> token id -> token surface forms + for t in self.tokenizer.encode(text, allowed_special=allowed_special, disallowed_special=disallowed_special): + tokens.append(self.decoder[t]) + + return tokens + + def convert_tokens_to_string(self, tokens: list[bytes | str]) -> str: + """ + Converts a sequence of tokens in a single string. + """ + text = "" + temp = b"" + for t in tokens: + if isinstance(t, str): + if temp: + text += temp.decode("utf-8", errors=self.errors) + temp = b"" + text += t + elif isinstance(t, bytes): + temp += t + else: + raise TypeError("token should only be of type types or str") + if temp: + text += temp.decode("utf-8", errors=self.errors) + return text + + @property + def vocab_size(self) -> int: + return self.tokenizer.n_vocab + + def _convert_id_to_token(self, index: int) -> bytes | str: + """Converts an id to a token, special tokens included""" + if index in self.decoder: + return self.decoder[index] + raise ValueError("unknown ids") + + def _convert_token_to_id(self, token: bytes | str) -> int: + """Converts a token to an id using the vocab, special tokens included""" + if token in self.special_tokens: + return self.special_tokens[token] + if token in self.mergeable_ranks: + return self.mergeable_ranks[token] + raise ValueError("unknown token") + + def _decode( + self, + token_ids: int | list[int], + skip_special_tokens: bool = False, + errors: str | None = None, + **kwargs, + ) -> str: + if isinstance(token_ids, int): + token_ids = [token_ids] + + if skip_special_tokens: + token_ids = [i for i in token_ids if i < self.eod_id] + + return self.tokenizer.decode(token_ids, errors=errors or self.errors) + + def bytes_to_str(self, byte_tokens: dict) -> str: + """Convert byte tokens to string representation. + + Args: + byte_tokens: A dictionary where keys are byte objects and values are integers, + or a single byte object. + + Returns: + If input is a dictionary, returns a new dictionary with byte keys converted to strings. + If input is a single byte object, returns the string representation. + """ + if isinstance(byte_tokens, dict): + return {k.decode("utf-8", errors=self.errors): v for k, v in byte_tokens.items() if isinstance(k, bytes)} + if isinstance(byte_tokens, bytes): + return byte_tokens.decode("utf-8", errors=self.errors) + return byte_tokens diff --git a/vllm_omni/transformers_utils/__init__.py b/vllm_omni/transformers_utils/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/transformers_utils/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py new file mode 100644 index 000000000..136a680a3 --- /dev/null +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Custom model configs that cannot be expressed via HuggingFace Transformers alone, +following the same pattern as vllm.transformers_utils.configs. +""" + +from __future__ import annotations + +import importlib + +_CLASS_TO_MODULE: dict[str, str] = { + "Mammothmoda2Config": "vllm_omni.transformers_utils.configs.mammoth_moda2", + "Mammothmoda2Qwen2_5_VLConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", + "Mammothmoda2Qwen2_5_VLTextConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", + "Mammothmoda2Qwen2_5_VLVisionConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", +} + +__all__ = [ + "Mammothmoda2Config", + "Mammothmoda2Qwen2_5_VLConfig", + "Mammothmoda2Qwen2_5_VLTextConfig", + "Mammothmoda2Qwen2_5_VLVisionConfig", +] + + +def __getattr__(name: str): + if name in _CLASS_TO_MODULE: + module_name = _CLASS_TO_MODULE[name] + module = importlib.import_module(module_name) + return getattr(module, name) + + raise AttributeError(f"module 'vllm_omni.transformers_utils.configs' has no attribute '{name}'") + + +def __dir__(): + return sorted(list(__all__)) + + +# Eagerly import all config modules so their AutoConfig.register() side-effects +# run as soon as `vllm_omni.transformers_utils.configs` is imported. +from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2 # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/mammoth_moda2.py b/vllm_omni/transformers_utils/configs/mammoth_moda2.py new file mode 100644 index 000000000..e8a4bd6dd --- /dev/null +++ b/vllm_omni/transformers_utils/configs/mammoth_moda2.py @@ -0,0 +1,283 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import ClassVar, Literal + +from transformers import AutoConfig, PretrainedConfig +from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import ( + Qwen2_5_VLConfig, + Qwen2_5_VLTextConfig, + Qwen2_5_VLVisionConfig, +) + +__all__ = [ + "Mammothmoda2Config", + "Mammothmoda2Qwen2_5_VLConfig", + "Mammothmoda2Qwen2_5_VLTextConfig", + "Mammothmoda2Qwen2_5_VLVisionConfig", +] + + +class Mammothmoda2Qwen2_5_VLVisionConfig(Qwen2_5_VLVisionConfig): + model_type = "mammothmoda2_qwen2_5_vl_vision" + base_config_key = "vision_config" + + def __init__( + self, + depth: int = 32, + hidden_size: int = 3584, + hidden_act: str = "silu", + intermediate_size: int = 3420, + num_heads: int = 16, + in_channels: int = 3, + patch_size: int = 14, + spatial_merge_size: int = 2, + temporal_patch_size: int = 2, + tokens_per_second: int = 4, + window_size: int = 112, + out_hidden_size: int = 3584, + fullatt_block_indexes: list[int] | None = None, + initializer_range: float = 0.02, + **kwargs, + ) -> None: + super().__init__( + depth=depth, + hidden_size=hidden_size, + hidden_act=hidden_act, + intermediate_size=intermediate_size, + num_heads=num_heads, + in_channels=in_channels, + patch_size=patch_size, + spatial_merge_size=spatial_merge_size, + temporal_patch_size=temporal_patch_size, + tokens_per_second=tokens_per_second, + window_size=window_size, + out_hidden_size=out_hidden_size, + fullatt_block_indexes=fullatt_block_indexes or [7, 15, 23, 31], + initializer_range=initializer_range, + **kwargs, + ) + + +class Mammothmoda2Qwen2_5_VLTextConfig(Qwen2_5_VLTextConfig): + model_type = "mammothmoda2_qwen2_5_vl_text" + base_config_key = "text_config" + keys_to_ignore_at_inference = ["past_key_values"] + base_model_tp_plan = Qwen2_5_VLTextConfig.base_model_tp_plan + base_model_pp_plan = Qwen2_5_VLTextConfig.base_model_pp_plan + + def __init__( + self, + vocab_size: int = 152064, + hidden_size: int = 8192, + intermediate_size: int = 29568, + num_hidden_layers: int = 80, + num_attention_heads: int = 64, + num_key_value_heads: int | None = 8, + hidden_act: str = "silu", + max_position_embeddings: int = 32768, + initializer_range: float = 0.02, + rms_norm_eps: float = 1e-05, + use_cache: bool = True, + tie_word_embeddings: bool = False, + rope_theta: float = 1000000.0, + use_sliding_window: bool = False, + sliding_window: int = 4096, + max_window_layers: int = 80, + layer_types: list[str] | None = None, + attention_dropout: float = 0.0, + rope_scaling: dict | None = None, + image_token_id: int | None = None, + video_token_id: int | None = None, + extra_gen_vocab: bool = True, + gen_vocab_size: int = 32800, + gen_vocab_start_index: int | None = None, + moe_type: str = "ffn", + **kwargs, + ) -> None: + super().__init__( + vocab_size=vocab_size, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_hidden_layers=num_hidden_layers, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + hidden_act=hidden_act, + max_position_embeddings=max_position_embeddings, + initializer_range=initializer_range, + rms_norm_eps=rms_norm_eps, + use_cache=use_cache, + tie_word_embeddings=tie_word_embeddings, + rope_theta=rope_theta, + use_sliding_window=use_sliding_window, + sliding_window=sliding_window, + max_window_layers=max_window_layers, + layer_types=layer_types, + attention_dropout=attention_dropout, + rope_scaling=rope_scaling, + **kwargs, + ) + + self.extra_gen_vocab = extra_gen_vocab + self.gen_vocab_size = gen_vocab_size + self.moe_type = moe_type + if gen_vocab_start_index is None: + self.gen_vocab_start_index = ( + self.vocab_size if self.extra_gen_vocab else self.vocab_size - self.gen_vocab_size + ) + else: + self.gen_vocab_start_index = gen_vocab_start_index + + # NOTE: vLLM V1 uses `hf_text_config.vocab_size` for sampling parameter validation + # (e.g., allowed_token_ids). Although MammothModa2's gen vocab is implemented via + # independent gen_embed/gen_head, the overall vocab size should still cover the + # gen vocab token ID range from the perspective of "output logits dimension". + if self.extra_gen_vocab: + self.vocab_size = int(self.gen_vocab_start_index) + int(self.gen_vocab_size) + + # Extra token IDs for multi-modal placeholders. + self.image_token_id = image_token_id + self.video_token_id = video_token_id + + +class Mammothmoda2Qwen2_5_VLConfig(Qwen2_5_VLConfig): + """Combined configuration: text_config + vision_config.""" + + model_type = "mammothmoda2_qwen2_5_vl" + sub_configs = { + "vision_config": Mammothmoda2Qwen2_5_VLVisionConfig, + "text_config": Mammothmoda2Qwen2_5_VLTextConfig, + } + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + text_config: dict | PretrainedConfig | None = None, + vision_config: dict | PretrainedConfig | None = None, + image_token_id: int = 151655, + video_token_id: int = 151656, + vision_start_token_id: int = 151652, + vision_end_token_id: int = 151653, + extra_gen_vocab: bool = True, + gen_vocab_size: int = 32800, + gen_vocab_start_index: int | None = None, + moe_type: str = "ffn", + **kwargs, + ) -> None: + if isinstance(vision_config, dict): + vision_config = self.sub_configs["vision_config"](**vision_config) + elif vision_config is None: + vision_config = self.sub_configs["vision_config"]() + + text_extra_kwargs = { + "extra_gen_vocab": extra_gen_vocab, + "gen_vocab_size": gen_vocab_size, + "moe_type": moe_type, + "gen_vocab_start_index": gen_vocab_start_index, + } + if isinstance(text_config, dict): + for key, val in text_extra_kwargs.items(): + text_config.setdefault(key, val) + text_config = self.sub_configs["text_config"](**text_config) + elif text_config is None: + text_config = self.sub_configs["text_config"](**{**text_extra_kwargs, **kwargs}) + elif isinstance(text_config, PretrainedConfig): + text_config = text_config + + super().__init__( + text_config=text_config, + vision_config=vision_config, + image_token_id=image_token_id, + video_token_id=video_token_id, + vision_start_token_id=vision_start_token_id, + vision_end_token_id=vision_end_token_id, + **kwargs, + ) + + if not hasattr(self, "text_config"): + self.text_config = text_config + if not hasattr(self, "vision_config"): + self.vision_config = vision_config + + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.extra_gen_vocab = getattr(self.text_config, "extra_gen_vocab", extra_gen_vocab) + self.gen_vocab_size = getattr(self.text_config, "gen_vocab_size", gen_vocab_size) + self.moe_type = getattr(self.text_config, "moe_type", moe_type) + self.gen_vocab_start_index = getattr(self.text_config, "gen_vocab_start_index", gen_vocab_start_index) + self.tokenizer_class = "MammothUTokenizer" + + +class Mammothmoda2Config(PretrainedConfig): + """Top-level MammothModa2 composition configuration""" + + model_type = "mammothmoda2" + is_composition = True + sub_configs: ClassVar = {"llm_config": AutoConfig} + + def __init__( + self, + *, + llm_config: dict | None = None, + gen_vae_config: dict | None = None, + gen_dit_config: dict | None = None, + gen_condition_mode: Literal["text", "image", "text_image"] = "image", + gen_image_condition_refiner_config: dict | None = None, + gen_axes_dim_rope: list[int] | None = None, + gen_axes_lens: list[int] | None = None, + gen_transport_config: dict | None = None, + initializer_range: float = 0.02, + architectures: list[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.llm_config = AutoConfig.for_model(**llm_config) if llm_config is not None else None + self.gen_vae_config = gen_vae_config + self.gen_dit_config = gen_dit_config + + self.gen_condition_mode = gen_condition_mode + self.gen_image_condition_refiner_config = gen_image_condition_refiner_config + self.gen_axes_dim_rope = gen_axes_dim_rope or [40, 40, 40] + self.gen_axes_lens = gen_axes_lens or [10000, 10000, 10000] + self.gen_transport_config = gen_transport_config or {} + self.initializer_range = initializer_range + self.tokenizer_class = "MammothUTokenizer" + self.architectures = ["Mammothmoda2Model"] + + def get_text_config(self, decoder: bool = False) -> PretrainedConfig: # noqa: ARG002 + return self.llm_config + + def _require_llm_config(self) -> PretrainedConfig: + if self.llm_config is None: + raise AttributeError("Mammothmoda2Config.llm_config is None") + return self.llm_config + + # ---- Proxy attrs for vLLM multimodal/mrope helpers ---- + # vllm_omni/model_executor/layers/mrope.py expects these fields on `hf_config`. + # In MammothModa2, they live in the nested `llm_config` (VL config), so we + # expose them here to make the top-level composition config compatible. + @property + def vision_config(self): + return self._require_llm_config().vision_config + + @property + def image_token_id(self) -> int: + return int(self._require_llm_config().image_token_id) + + @property + def video_token_id(self) -> int: + return int(self._require_llm_config().video_token_id) + + @property + def vision_start_token_id(self) -> int: + return int(self._require_llm_config().vision_start_token_id) + + @property + def vision_end_token_id(self) -> int: + return int(self._require_llm_config().vision_end_token_id) + + +# Register model_type -> config class for AutoConfig +AutoConfig.register(Mammothmoda2Config.model_type, Mammothmoda2Config) +AutoConfig.register(Mammothmoda2Qwen2_5_VLConfig.model_type, Mammothmoda2Qwen2_5_VLConfig) +AutoConfig.register(Mammothmoda2Qwen2_5_VLTextConfig.model_type, Mammothmoda2Qwen2_5_VLTextConfig) +AutoConfig.register(Mammothmoda2Qwen2_5_VLVisionConfig.model_type, Mammothmoda2Qwen2_5_VLVisionConfig) diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index a6f7efafc..71f173424 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -672,6 +672,16 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None + # Some generation-stage models (e.g. MammothModa2DiTPipeline) require + # model-specific runtime information (such as image size and conditioning + # embeddings) even during the dummy profiling run that vLLM uses to + # estimate KV-cache capacity. get_dummy_runtime_additional_information + # provides placeholder values of the correct shape so that the profiling + # run does not raise an error due to missing inputs. + if hasattr(self.model, "get_dummy_runtime_additional_information"): + runtime_addi = self.model.get_dummy_runtime_additional_information(num_reqs) + model_kwargs["runtime_additional_information"] = runtime_addi + if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_tokens_padded] elif self.uses_xdrope_dim > 0: diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 569ee60f9..1d70e0eb3 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -950,8 +950,18 @@ class OmniGPUModelRunner(GPUModelRunner): """Gather per-request model_intermediate_buffer in batch order.""" per_req_runtime_info = [] for req_id in self.input_batch.req_ids: + req_state = self.requests.get(req_id) + # MammothModa2 AR grid constraint: the model must emit a special + # end-of-line (EOL) token at the end of each image row. To determine + # whether the current decoding step falls on a row boundary, the + # constraint logic (see MammothModa2ARForConditionalGeneration. + # _apply_t2i_token_constraints) computes: + # column_id = generated_len % (ar_width + 1) + # and forces the EOL token when column_id == ar_width. + generated_len = len(req_state.output_token_ids) if req_state is not None else 0 info = self.model_intermediate_buffer.get(req_id, {}) if info: + info["generated_len"] = generated_len per_req_runtime_info.append(info) if "thinker_reply_part_per_request" in info: q = info["thinker_reply_part_per_request"]