Merge remote-tracking branch 'remotes/origin/main' into dev/rebase-v0169
This commit is contained in:
tzhouam committed
commit
395eab8984
59 files changed
+11285
-22
No files matched your search
@@ -1,6 +1,6 @@
|
||||
<p align="center">
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
|
||||
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
|
||||
<img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
|
||||
</picture>
|
||||
</p>
|
||||
@@ -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).
|
||||
|
||||
+15
-3
@@ -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
|
||||
|
||||
@@ -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` |
|
||||
|
||||
@@ -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.<img>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).
|
||||
@@ -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.<img>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}<img>{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)
|
||||
@@ -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
|
||||
```
|
||||
@@ -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()
|
||||
@@ -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 ``<base>_0<ext>``, ``<base>_1<ext>``, …
|
||||
"""
|
||||
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()
|
||||
@@ -203,3 +203,5 @@ extend-ignore-identifiers-re = [
|
||||
".*[Oo]no_[Aa]nna.*",
|
||||
".*cann.*",
|
||||
]
|
||||
[tool.typos.default.extend-words]
|
||||
ue = "ue"
|
||||
@@ -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
|
||||
@@ -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: <repo_root>/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()
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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]
|
||||
@@ -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}")
|
||||
|
||||
@@ -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": {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
@@ -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.",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
Whitespace-only changes.
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -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
|
||||
@@ -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))
|
||||
@@ -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
|
||||
@@ -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,
|
||||
},
|
||||
}
|
||||
@@ -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!")
|
||||
@@ -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}")
|
||||
@@ -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 = {}
|
||||
@@ -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]",
|
||||
"<strong>",
|
||||
"</strong>",
|
||||
"[noise]",
|
||||
"[laughter]",
|
||||
"[cough]",
|
||||
"[clucking]",
|
||||
"[accent]",
|
||||
"[quick_breath]",
|
||||
"<laughter>",
|
||||
"</laughter>",
|
||||
"[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]",
|
||||
"<strong>",
|
||||
"</strong>",
|
||||
"[noise]",
|
||||
"[laughter]",
|
||||
"[cough]",
|
||||
"[clucking]",
|
||||
"[accent]",
|
||||
"[quick_breath]",
|
||||
"<laughter>",
|
||||
"</laughter>",
|
||||
"[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
|
||||
@@ -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
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
File diff suppressed because it is too large.
Load diff
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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}")
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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(
|
||||
{
|
||||
**{
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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]
|
||||
@@ -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
|
||||
Whitespace-only changes.
@@ -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|>",
|
||||
"<tool_call>",
|
||||
"</tool_call>",
|
||||
"<|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
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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:
|
||||
|
||||
@@ -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"]
|
||||
|
||||
Reference in new issue
Block a user