diff --git a/docs/assets/WeChat.jpg b/docs/assets/WeChat.jpg index 3d049b546..c5e88d616 100644 Binary files a/docs/assets/WeChat.jpg and b/docs/assets/WeChat.jpg differ diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 980488852..804655d00 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -30,6 +30,7 @@ th { |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | | `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | +| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | |`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | |`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | diff --git a/examples/offline_inference/helios/README.md b/examples/offline_inference/helios/README.md new file mode 100644 index 000000000..2c7ba525e --- /dev/null +++ b/examples/offline_inference/helios/README.md @@ -0,0 +1,193 @@ +# Helios Video Generation + +Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants: + +| Variant | Description | Key Features | +|---------|-------------|--------------| +| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` | +| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support | +| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` | + +## Setup + +Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. + +## Run Examples + +Get into the example folder: +```bash +cd examples/offline_inference/helios +``` + +### Text-to-Video (T2V) + +**Helios-Base** (Stage 1 only): +```bash +python end2end.py \ + --model BestWishYsh/Helios-Base \ + --sample-type t2v \ + --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \ + --guidance-scale 5.0 \ + --output helios_t2v_base.mp4 +``` + +**Helios-Mid** (Stage 2 + CFG-Zero*): +```bash +python end2end.py \ + --model BestWishYsh/Helios-Mid \ + --sample-type t2v \ + --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \ + --guidance-scale 5.0 \ + --is-enable-stage2 \ + --pyramid-num-inference-steps-list 20 20 20 \ + --use-cfg-zero-star --use-zero-init --zero-steps 1 \ + --output helios_t2v_mid.mp4 +``` + +**Helios-Distilled** (Stage 2 pyramid + DMD): +```bash +python end2end.py \ + --model BestWishYsh/Helios-Distilled \ + --sample-type t2v \ + --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \ + --num-frames 240 \ + --guidance-scale 1.0 \ + --is-enable-stage2 \ + --pyramid-num-inference-steps-list 2 2 2 \ + --is-amplify-first-chunk \ + --output helios_t2v_distilled.mp4 +``` + +### Image-to-Video (I2V) + +**Helios-Base**: +```bash +python end2end.py \ + --model BestWishYsh/Helios-Base \ + --sample-type i2v \ + --image-path /path/to/image.jpg \ + --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \ + --guidance-scale 5.0 \ + --output helios_i2v_base.mp4 +``` + +**Helios-Mid** (Stage 2 + CFG-Zero*): +```bash +python end2end.py \ + --model BestWishYsh/Helios-Mid \ + --sample-type i2v \ + --image-path /path/to/image.jpg \ + --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \ + --guidance-scale 5.0 \ + --is-enable-stage2 \ + --pyramid-num-inference-steps-list 20 20 20 \ + --use-cfg-zero-star --use-zero-init --zero-steps 1 \ + --output helios_i2v_mid.mp4 +``` + +**Helios-Distilled**: +```bash +python end2end.py \ + --model BestWishYsh/Helios-Distilled \ + --sample-type i2v \ + --image-path /path/to/image.jpg \ + --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \ + --num-frames 240 \ + --guidance-scale 1.0 \ + --is-enable-stage2 \ + --pyramid-num-inference-steps-list 2 2 2 \ + --is-amplify-first-chunk \ + --output helios_i2v_distilled.mp4 +``` + +### Video-to-Video (V2V) + +**Helios-Base**: +```bash +python end2end.py \ + --model BestWishYsh/Helios-Base \ + --sample-type v2v \ + --video-path /path/to/video.mp4 \ + --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \ + --guidance-scale 5.0 \ + --output helios_v2v_base.mp4 +``` + +**Helios-Mid** (Stage 2 + CFG-Zero*): +```bash +python end2end.py \ + --model BestWishYsh/Helios-Mid \ + --sample-type v2v \ + --video-path /path/to/video.mp4 \ + --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \ + --guidance-scale 5.0 \ + --is-enable-stage2 \ + --pyramid-num-inference-steps-list 20 20 20 \ + --use-cfg-zero-star --use-zero-init --zero-steps 1 \ + --output helios_v2v_mid.mp4 +``` + +**Helios-Distilled**: +```bash +python end2end.py \ + --model BestWishYsh/Helios-Distilled \ + --sample-type v2v \ + --video-path /path/to/video.mp4 \ + --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \ + --num-frames 240 \ + --guidance-scale 1.0 \ + --is-enable-stage2 \ + --pyramid-num-inference-steps-list 2 2 2 \ + --is-amplify-first-chunk \ + --output helios_v2v_distilled.mp4 +``` + +## Common Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path | +| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` | +| `--prompt` | — | Text prompt describing the video | +| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) | +| `--image-path` | — | Input image (required for `i2v`) | +| `--video-path` | — | Input video (required for `v2v`) | +| `--height` | `384` | Video height in pixels | +| `--width` | `640` | Video width in pixels | +| `--num-frames` | `99` | Number of output frames | +| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) | +| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) | +| `--seed` | `42` | Random seed | +| `--fps` | `16` | Output video frame rate | +| `--output` | `helios_output.mp4` | Output file path | + +### Stage 2 / Pyramid Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising | +| `--pyramid-num-stages` | `3` | Number of pyramid stages | +| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage | +| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) | + +### CFG-Zero* Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) | +| `--use-zero-init` | off | Zero init for first steps | +| `--zero-steps` | `1` | Number of zero-init steps | + +### Memory & Parallelism + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `--vae-use-slicing` | off | Enable VAE slicing | +| `--vae-use-tiling` | off | Enable VAE tiling | +| `--enforce-eager` | off | Disable torch.compile | +| `--enable-cpu-offload` | off | CPU offloading | +| `--enable-layerwise-offload` | off | Layerwise offloading | +| `--ulysses-degree` | `1` | Ulysses SP degree | +| `--ring-degree` | `1` | Ring SP degree | +| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) | +| `--tensor-parallel-size` | `1` | Tensor parallelism size | diff --git a/examples/offline_inference/helios/end2end.py b/examples/offline_inference/helios/end2end.py new file mode 100644 index 000000000..450c65cc4 --- /dev/null +++ b/examples/offline_inference/helios/end2end.py @@ -0,0 +1,343 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Helios video generation example supporting T2V, I2V, and V2V. + +Usage (T2V, Helios-Base, Stage 1 only): + python end2end.py \ + --model /path/to/Helios-Base --sample-type t2v \ + --prompt "A serene lakeside sunrise with mist over the water." \ + --height 384 --width 640 --num-frames 99 \ + --num-inference-steps 50 --guidance-scale 5.0 + +Usage (I2V, Helios-Base): + python end2end.py \ + --model /path/to/Helios-Base --sample-type i2v \ + --image-path /path/to/image.jpg \ + --prompt "Description of desired animation." \ + --guidance-scale 5.0 + +Usage (V2V, Helios-Base): + python end2end.py \ + --model /path/to/Helios-Base --sample-type v2v \ + --video-path /path/to/video.mp4 \ + --prompt "Description of desired transformation." \ + --guidance-scale 5.0 + +Usage (Helios-Mid, Stage 2 + CFG-Zero*): + python end2end.py \ + --model /path/to/Helios-Mid --sample-type t2v \ + --prompt "A serene lakeside sunrise with mist over the water." \ + --guidance-scale 5.0 \ + --is-enable-stage2 \ + --pyramid-num-inference-steps-list 20 20 20 \ + --use-cfg-zero-star --use-zero-init --zero-steps 1 + +Usage (Helios-Distilled, Stage 2 pyramid + DMD): + python end2end.py \ + --model /path/to/Helios-Distilled --sample-type t2v \ + --prompt "A serene lakeside sunrise with mist over the water." \ + --guidance-scale 1.0 \ + --is-enable-stage2 \ + --pyramid-num-inference-steps-list 2 2 2 \ + --is-amplify-first-chunk +""" + +import argparse +import time +from pathlib import Path + +import numpy as np +import torch + +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.platforms import current_omni_platform + + +def load_image_as_tensor(image_path: str, height: int, width: int) -> torch.Tensor: + """Load an image and return a (1, C, H, W) tensor normalized to [-1, 1].""" + from PIL import Image + from torchvision import transforms + + image = Image.open(image_path).convert("RGB").resize((width, height)) + transform = transforms.Compose( + [ + transforms.ToTensor(), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), + ] + ) + return transform(image).unsqueeze(0) + + +def load_video_as_tensor(video_path: str, height: int, width: int) -> torch.Tensor: + """Load a video and return a (1, C, T, H, W) tensor normalized to [-1, 1].""" + from torchvision import transforms + from torchvision.io import read_video + + video_frames, _, _ = read_video(video_path, output_format="TCHW") + video_frames = video_frames.float() / 255.0 + + transform = transforms.Compose( + [ + transforms.Resize((height, width), antialias=True), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), + ] + ) + video_frames = torch.stack([transform(f) for f in video_frames]) + return video_frames.permute(1, 0, 2, 3).unsqueeze(0) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Generate a video with Helios (T2V / I2V / V2V).") + parser.add_argument( + "--model", + default="BestWishYsh/Helios-Base", + help="Helios model ID or local path (e.g. Helios-Base, Helios-Mid, Helios-Distilled).", + ) + parser.add_argument( + "--sample-type", + choices=["t2v", "i2v", "v2v"], + default="t2v", + help="Generation mode: t2v (text-to-video), i2v (image-to-video), v2v (video-to-video).", + ) + parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.") + parser.add_argument( + "--negative-prompt", + default=( + "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, " + "images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, " + "incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, " + "misshapen limbs, fused fingers, still picture, messy background, three legs, many people " + "in the background, walking backwards" + ), + help="Negative prompt.", + ) + parser.add_argument("--seed", type=int, default=42, help="Random seed.") + parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.") + parser.add_argument("--height", type=int, default=384, help="Video height.") + parser.add_argument("--width", type=int, default=640, help="Video width.") + parser.add_argument("--num-frames", type=int, default=99, help="Number of video frames.") + parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps (Stage 1 only).") + parser.add_argument("--output", type=str, default="helios_output.mp4", help="Output video path.") + parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.") + + # I2V / V2V + parser.add_argument("--image-path", type=str, default=None, help="Input image path for I2V mode.") + parser.add_argument("--video-path", type=str, default=None, help="Input video path for V2V mode.") + + # Stage 2 (pyramid multi-stage denoising) + parser.add_argument( + "--is-enable-stage2", + action="store_true", + help="Enable pyramid multi-stage denoising (Stage 2). Required for Helios-Distilled.", + ) + parser.add_argument( + "--pyramid-num-stages", + type=int, + default=3, + help="Number of pyramid stages for Stage 2.", + ) + parser.add_argument( + "--pyramid-num-inference-steps-list", + type=int, + nargs="+", + default=[10, 10, 10], + help="Inference steps per pyramid stage.", + ) + + # DMD + parser.add_argument( + "--is-amplify-first-chunk", + action="store_true", + help="Enable DMD amplification for the first chunk (Helios-Distilled).", + ) + + # CFG Zero Star + parser.add_argument( + "--use-cfg-zero-star", + action="store_true", + help="Enable CFG Zero Star guidance (recommended for Helios-Mid).", + ) + parser.add_argument( + "--use-zero-init", + action="store_true", + help="Use zero initialization for the first denoising steps with CFG-Zero*.", + ) + parser.add_argument( + "--zero-steps", + type=int, + default=1, + help="Number of initial denoising steps using zero prediction (default: 1).", + ) + + # Memory & parallelism + parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.") + parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.") + parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.") + parser.add_argument("--enable-cpu-offload", action="store_true", help="Enable CPU offloading.") + parser.add_argument("--enable-layerwise-offload", action="store_true", help="Enable layerwise offloading.") + parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.") + parser.add_argument("--ring-degree", type=int, default=1, help="Ring SP degree.") + parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.") + parser.add_argument("--tensor-parallel-size", type=int, default=1, help="Tensor parallelism size.") + + return parser.parse_args() + + +def main(): + args = parse_args() + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) + + parallel_config = DiffusionParallelConfig( + ulysses_degree=args.ulysses_degree, + ring_degree=args.ring_degree, + cfg_parallel_size=args.cfg_parallel_size, + tensor_parallel_size=args.tensor_parallel_size, + ) + + omni = Omni( + model=args.model, + enable_layerwise_offload=args.enable_layerwise_offload, + vae_use_slicing=args.vae_use_slicing, + vae_use_tiling=args.vae_use_tiling, + enable_cpu_offload=args.enable_cpu_offload, + parallel_config=parallel_config, + enforce_eager=args.enforce_eager, + ) + + # Validate I2V / V2V arguments + if args.sample_type == "i2v" and not args.image_path: + raise ValueError("--image-path is required for I2V mode.") + if args.sample_type == "v2v" and not args.video_path: + raise ValueError("--video-path is required for V2V mode.") + + # Build extra_args for Helios-specific parameters + extra_args = {} + if args.is_enable_stage2: + extra_args["is_enable_stage2"] = True + extra_args["pyramid_num_stages"] = args.pyramid_num_stages + extra_args["pyramid_num_inference_steps_list"] = args.pyramid_num_inference_steps_list + if args.is_amplify_first_chunk: + extra_args["is_amplify_first_chunk"] = True + if args.use_cfg_zero_star: + extra_args["use_cfg_zero_star"] = True + if args.use_zero_init: + extra_args["use_zero_init"] = True + extra_args["zero_steps"] = args.zero_steps + + if args.sample_type == "i2v": + image_tensor = load_image_as_tensor(args.image_path, args.height, args.width) + extra_args["image"] = image_tensor + elif args.sample_type == "v2v": + video_tensor = load_video_as_tensor(args.video_path, args.height, args.width) + extra_args["video"] = video_tensor + + # Print generation configuration + print(f"\n{'=' * 60}") + print("Helios Generation Configuration:") + print(f" Model: {args.model}") + print(f" Sample type: {args.sample_type.upper()}") + if args.sample_type == "i2v": + print(f" Image: {args.image_path}") + elif args.sample_type == "v2v": + print(f" Video: {args.video_path}") + print(f" Prompt: {args.prompt[:80]}{'...' if len(args.prompt) > 80 else ''}") + print(f" Video size: {args.width}x{args.height}, {args.num_frames} frames") + print(f" Inference steps: {args.num_inference_steps}") + print(f" Guidance scale: {args.guidance_scale}") + if args.is_enable_stage2: + print(f" Stage 2: enabled (stages={args.pyramid_num_stages}, steps={args.pyramid_num_inference_steps_list})") + if args.is_amplify_first_chunk: + print(" DMD amplify first chunk: enabled") + if args.use_cfg_zero_star: + print(f" CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})") + else: + if args.use_cfg_zero_star: + print(f" CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})") + print(" Stage 2: disabled (Stage 1 only)") + print(f"{'=' * 60}\n") + + generation_start = time.perf_counter() + frames = omni.generate( + { + "prompt": args.prompt, + "negative_prompt": args.negative_prompt, + }, + OmniDiffusionSamplingParams( + height=args.height, + width=args.width, + generator=generator, + guidance_scale=args.guidance_scale, + num_inference_steps=args.num_inference_steps, + num_frames=args.num_frames, + extra_args=extra_args, + ), + ) + generation_end = time.perf_counter() + generation_time = generation_end - generation_start + + print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)") + + # Extract video frames from OmniRequestOutput + if isinstance(frames, list) and len(frames) > 0: + first_item = frames[0] + + if hasattr(first_item, "final_output_type"): + if first_item.final_output_type != "image": + raise ValueError( + f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation." + ) + + if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output: + if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0: + inner_output = first_item.request_output[0] + if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"): + frames = inner_output.images[0] if inner_output.images else None + if frames is None: + raise ValueError("No video frames found in output.") + elif hasattr(first_item, "images") and first_item.images: + frames = first_item.images + else: + raise ValueError("No video frames found in OmniRequestOutput.") + + # Unwrap batch list from postprocess_video: [numpy(T,H,W,C)] -> numpy(T,H,W,C) + if isinstance(frames, list) and len(frames) > 0 and isinstance(frames[0], np.ndarray): + frames = frames[0] + + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + try: + from diffusers.utils import export_to_video + except ImportError: + raise ImportError("diffusers is required for export_to_video.") + + if isinstance(frames, torch.Tensor): + video_tensor = frames.detach().cpu() + if video_tensor.dim() == 5: + if video_tensor.shape[1] in (3, 4): + video_tensor = video_tensor[0].permute(1, 2, 3, 0) + else: + video_tensor = video_tensor[0] + elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4): + video_tensor = video_tensor.permute(1, 2, 3, 0) + if video_tensor.is_floating_point(): + video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5 + video_array = video_tensor.float().numpy() + elif isinstance(frames, np.ndarray): + video_array = frames + else: + video_array = frames + + if isinstance(video_array, np.ndarray) and video_array.ndim == 5: + video_array = video_array[0] + + export_to_video(video_array, str(output_path), fps=args.fps) + print(f"Saved generated video to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index ea967a7b0..fe4e654e8 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -182,7 +182,11 @@ def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch monkeypatch.setattr(manager_mod, "replace_submodule", _fake_replace_submodule) pipeline = torch.nn.Module() - pipeline.packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]} + pipeline.stacked_params_mapping = [ + (".to_qkv.", ".to_q.", "q"), + (".to_qkv.", ".to_k.", "k"), + (".to_qkv.", ".to_v.", "v"), + ] pipeline.transformer = torch.nn.Module() pipeline.transformer.to_qkv = _FakeLinearBase() @@ -194,7 +198,7 @@ def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch ) # Treat the dummy layer as a packed 3-slice projection so the manager uses - # `packed_modules_mapping` to decide replacement based on target_modules. + # `stacked_params_mapping` to decide replacement based on target_modules. monkeypatch.setattr(manager, "_get_packed_modules_list", lambda _module: ["q", "k", "v"]) peft_helper = type("_PH", (), {"r": 1, "target_modules": ["to_q"]})() @@ -254,7 +258,11 @@ def test_lora_manager_activates_fused_lora_on_packed_layer(): def test_lora_manager_activates_packed_lora_from_sublayers(): pipeline = torch.nn.Module() - pipeline.packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]} + pipeline.stacked_params_mapping = [ + (".to_qkv", ".to_q", "q"), + (".to_qkv", ".to_k", "k"), + (".to_qkv", ".to_v", "v"), + ] manager = DiffusionLoRAManager( pipeline=pipeline, device=torch.device("cpu"), diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 7af4e078f..10dbd3a3b 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -349,13 +349,13 @@ class TestTTSMethods: req = OpenAICreateSpeechRequest(input="Hello", language="InvalidLang") assert "Invalid language" in speech_server._validate_tts_request(req) - # When no speakers loaded, any voice is accepted (unconstrained) + # CustomVoice on model with no speakers -> rejected req = OpenAICreateSpeechRequest(input="Hello", voice="Invalid") - assert speech_server._validate_tts_request(req) is None + assert "does not support CustomVoice" in speech_server._validate_tts_request(req) - # Valid request - req = OpenAICreateSpeechRequest(input="Hello", voice="Vivian") - assert speech_server._validate_tts_request(req) is None + # CustomVoice without voice on model with no speakers -> also rejected + req = OpenAICreateSpeechRequest(input="Hello") + assert "does not support CustomVoice" in speech_server._validate_tts_request(req) def test_validate_tts_request_task_types(self, speech_server): """Test task-specific validation.""" @@ -367,9 +367,43 @@ class TestTTSMethods: req = OpenAICreateSpeechRequest(input="Hello", task_type="VoiceDesign") assert "instructions" in speech_server._validate_tts_request(req) - # ref_text only for Base + # ref_text without task_type auto-infers Base, then fails on missing ref_audio req = OpenAICreateSpeechRequest(input="Hello", ref_text="text") - assert "Base task" in speech_server._validate_tts_request(req) + assert "ref_audio" in speech_server._validate_tts_request(req) + + def test_validate_tts_request_auto_infer_base(self, speech_server): + """Test auto-inference of Base task when ref_audio/ref_text is provided.""" + # ref_audio without task_type -> infers Base, requires non-empty ref_text + req = OpenAICreateSpeechRequest(input="Hello", ref_audio="data:audio/wav;base64,abc") + result = speech_server._validate_tts_request(req) + assert "ref_text" in result + assert req.task_type == "Base" + + # ref_text without task_type -> infers Base, requires ref_audio + req = OpenAICreateSpeechRequest(input="Hello", ref_text="some text") + result = speech_server._validate_tts_request(req) + assert "ref_audio" in result + assert req.task_type == "Base" + + def test_validate_tts_request_base_empty_ref_text(self, speech_server): + """Empty ref_text on Base task returns 400 instead of crashing engine.""" + req = OpenAICreateSpeechRequest( + input="Hello", task_type="Base", ref_audio="data:audio/wav;base64,abc", ref_text="" + ) + result = speech_server._validate_tts_request(req) + assert "non-empty 'ref_text'" in result + + # x_vector_only_mode bypasses ref_text requirement + req = OpenAICreateSpeechRequest( + input="Hello", task_type="Base", ref_audio="data:audio/wav;base64,abc", ref_text="", x_vector_only_mode=True + ) + assert speech_server._validate_tts_request(req) is None + + def test_validate_tts_request_customvoice_no_speakers(self, speech_server): + """CustomVoice on a model with no speakers returns 400 instead of crashing engine.""" + req = OpenAICreateSpeechRequest(input="Hello", task_type="CustomVoice") + result = speech_server._validate_tts_request(req) + assert "does not support CustomVoice" in result def test_build_tts_params(self, speech_server): """Test TTS parameter building.""" diff --git a/tools/nightly/send_nightly_perf_email.py b/tools/nightly/send_nightly_perf_email.py index af11cc73a..9b4bf7ece 100644 --- a/tools/nightly/send_nightly_perf_email.py +++ b/tools/nightly/send_nightly_perf_email.py @@ -13,6 +13,7 @@ import logging import os import smtplib import sys +from email.header import Header from email.mime.application import MIMEApplication from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText @@ -123,7 +124,7 @@ def _send_mail( subject = _build_subject(prefix=prefix, date_str=date_str) msg = MIMEMultipart() - msg["Subject"] = subject + msg["Subject"] = Header(subject, "utf-8") msg["From"] = sender msg["To"] = ", ".join(recipients) msg.attach(MIMEText(body, "plain", "utf-8")) diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 04fdaca18..a397e6085 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -22,6 +22,8 @@ from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapt ) from vllm_omni.outputs import OmniModelRunnerOutput +logger = init_logger(__name__) + class OmniGenerationScheduler(VLLMScheduler): def __init__(self, *args, **kwargs): @@ -307,7 +309,7 @@ class OmniGenerationScheduler(VLLMScheduler): except Exception: # If anything goes wrong, leave the original output unchanged - init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData") + logger.exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData") return scheduler_output @@ -462,6 +464,10 @@ class OmniGenerationScheduler(VLLMScheduler): prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) if new_token_ids or pooler_output is not None or kv_transfer_params or stopped: # Add EngineCoreOutput for this Request. + num_cached = request.num_cached_tokens + if num_cached < 0: + logger.warning("Negative num_cached_tokens (%d) for request %s, clamping to 0", num_cached, req_id) + num_cached = 0 outputs[request.client_index].append( EngineCoreOutput( request_id=req_id, @@ -474,7 +480,7 @@ class OmniGenerationScheduler(VLLMScheduler): events=request.take_events(), kv_transfer_params=kv_transfer_params, trace_headers=request.trace_headers, - num_cached_tokens=request.num_cached_tokens, + num_cached_tokens=num_cached, num_external_computed_tokens=request.num_external_computed_tokens, routed_experts=routed_experts, num_nans_in_logits=request.num_nans_in_logits, @@ -496,6 +502,12 @@ class OmniGenerationScheduler(VLLMScheduler): requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids] self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR) for request in requests: + num_cached = request.num_cached_tokens + if num_cached < 0: + logger.warning( + "Negative num_cached_tokens (%d) for request %s, clamping to 0", num_cached, request.request_id + ) + num_cached = 0 outputs[request.client_index].append( EngineCoreOutput( request_id=request.request_id, @@ -503,7 +515,7 @@ class OmniGenerationScheduler(VLLMScheduler): finish_reason=request.get_finished_reason(), events=request.take_events(), trace_headers=request.trace_headers, - num_cached_tokens=request.num_cached_tokens, + num_cached_tokens=num_cached, ) ) if self.chunk_transfer_adapter is not None: diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 3cccc732e..0dd5ba986 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -132,17 +132,43 @@ class DiffusionLoRAManager: def _compute_packed_modules_mapping(self) -> dict[str, list[str]]: """Collect packed->sublayer mappings from the diffusion model. - vLLM models declare `packed_modules_mapping` on the model class. For - diffusion pipelines, we attach the same mapping on the transformer - module(s) that implement packed (fused) projections, so LoRA loading can - accept checkpoints trained against the logical sub-projections. + Diffusion models often use packed (fused) projections like `to_qkv` or + `w13`, while LoRA checkpoints are typically saved against the logical + sub-projections (e.g. `to_q`/`to_k`/`to_v`, `w1`/`w3`). Many diffusion + model implementations already define these relationships in + `load_weights()` via `stacked_params_mapping`. To avoid duplicating the + mapping in multiple places, we derive packed→sublayer mappings from the + model's `stacked_params_mapping`. """ + + def _derive_from_stacked_params_mapping(stacked: object) -> dict[str, list[str]]: + if not isinstance(stacked, (list, tuple)): + return {} + derived: dict[str, list[str]] = {} + for item in stacked: + if not isinstance(item, (list, tuple)) or len(item) < 2: + continue + packed_suffix, sub_suffix = item[0], item[1] + if not isinstance(packed_suffix, str) or not packed_suffix: + continue + if not isinstance(sub_suffix, str) or not sub_suffix: + continue + # The mapping strings are usually suffix patterns (e.g. ".to_qkv"), + # but some models scope them under submodules (e.g. ".attn1.to_qkv"). + # For LoRA we only care about the leaf module names. + packed_name = packed_suffix.strip(".").split(".")[-1] + sub_name = sub_suffix.strip(".").split(".")[-1] + existing = derived.get(packed_name) + if existing is None: + derived[packed_name] = [sub_name] + elif sub_name not in existing: + existing.append(sub_name) + return derived + mapping: dict[str, list[str]] = {} for module in self.pipeline.modules(): - packed = getattr(module, "packed_modules_mapping", None) - if not isinstance(packed, dict): - continue - for packed_name, sub_names in packed.items(): + derived = _derive_from_stacked_params_mapping(getattr(module, "stacked_params_mapping", None)) + for packed_name, sub_names in derived.items(): if not isinstance(packed_name, str) or not packed_name: continue if not isinstance(sub_names, (list, tuple)) or not all(isinstance(s, str) for s in sub_names): @@ -156,7 +182,7 @@ class DiffusionLoRAManager: mapping[packed_name] = sub_names_list elif existing != sub_names_list: logger.warning( - "Conflicting packed_modules_mapping for %s: %s vs %s; using %s", + "Conflicting packed module mapping for %s: %s vs %s; using %s", packed_name, existing, sub_names_list, @@ -171,7 +197,7 @@ class DiffusionLoRAManager: return None if len(sub_suffixes) != n_slices: logger.warning( - "packed_modules_mapping[%s] has %d slices but layer expects %d; skipping sublayer lookup", + "Packed module mapping[%s] has %d slices but layer expects %d; skipping sublayer lookup", packed_module_suffix, len(sub_suffixes), n_slices, diff --git a/vllm_omni/diffusion/lora/utils.py b/vllm_omni/diffusion/lora/utils.py index 5f1baea34..465edc99c 100644 --- a/vllm_omni/diffusion/lora/utils.py +++ b/vllm_omni/diffusion/lora/utils.py @@ -39,9 +39,10 @@ def _expand_expected_modules_for_packed_layers( `supported_modules`, but the sublayer names are not. Expanding the set ensures these sublayer keys are not dropped when loading a LoRA checkpoint. - The packed→sublayer mapping is model-specific (see each diffusion model's - `packed_modules_mapping`) so new packed layers are added alongside the model - implementation rather than hard-coded in the LoRA framework. + The packed→sublayer mapping is model-specific and is derived from each + diffusion model's `stacked_params_mapping` (used by `load_weights()`), so + new packed layers are added alongside the model implementation rather than + hard-coded in the LoRA framework. """ expanded = set(supported_modules) if not packed_modules_mapping: diff --git a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py index 5ee5ee440..2499b9c16 100644 --- a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py +++ b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py @@ -583,10 +583,6 @@ class Flux2Transformer2DModel(nn.Module): """ _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"] - packed_modules_mapping = { - "to_qkv": ["to_q", "to_k", "to_v"], - "add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"], - } def __init__( self, @@ -762,6 +758,8 @@ class Flux2Transformer2DModel(nn.Module): (".add_kv_proj", ".add_k_proj", "k"), (".add_kv_proj", ".add_v_proj", "v"), ] + # Expose packed shard mappings for LoRA handling of fused projections. + self.stacked_params_mapping = stacked_params_mapping params_dict = dict(self.named_parameters()) diff --git a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py index ecc418d4a..cf6ec8f03 100644 --- a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py +++ b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py @@ -544,9 +544,7 @@ class GlmImageTransformer2DModel(CachedTransformer): `od_config.tf_model_config`. """ - packed_modules_mapping = { - "to_qkv": ["to_q", "to_k", "to_v"], - } + _repeated_blocks = ["GlmImageTransformerBlock"] def __init__( self, @@ -717,6 +715,8 @@ class GlmImageTransformer2DModel(CachedTransformer): (".to_qkv", ".to_k", "k"), (".to_qkv", ".to_v", "v"), ] + # Expose packed shard mappings for LoRA handling of fused projections. + self.stacked_params_mapping = stacked_params_mapping params_dict = dict(self.named_parameters()) diff --git a/vllm_omni/diffusion/models/helios/__init__.py b/vllm_omni/diffusion/models/helios/__init__.py new file mode 100644 index 000000000..2852ccdab --- /dev/null +++ b/vllm_omni/diffusion/models/helios/__init__.py @@ -0,0 +1,19 @@ +from .helios_transformer import HeliosTransformer3DModel +from .pipeline_helios import ( + HeliosPipeline, + create_transformer_from_config, + get_helios_post_process_func, + get_helios_pre_process_func, + load_transformer_config, +) +from .scheduling_helios import HeliosScheduler + +__all__ = [ + "HeliosPipeline", + "HeliosTransformer3DModel", + "HeliosScheduler", + "get_helios_post_process_func", + "get_helios_pre_process_func", + "load_transformer_config", + "create_transformer_from_config", +] diff --git a/vllm_omni/diffusion/models/helios/helios_transformer.py b/vllm_omni/diffusion/models/helios/helios_transformer.py new file mode 100644 index 000000000..17b9ec37e --- /dev/null +++ b/vllm_omni/diffusion/models/helios/helios_transformer.py @@ -0,0 +1,929 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from Helios (https://github.com/BestWishYsh/Helios) + +import math +from collections.abc import Iterable +from functools import lru_cache +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.normalization import FP32LayerNorm +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.conv import Conv3dLayer +from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, + SequenceParallelOutput, +) + +logger = init_logger(__name__) + + +def pad_for_3d_conv(x, kernel_size): + b, c, t, h, w = x.shape + pt, ph, pw = kernel_size + pad_t = (pt - (t % pt)) % pt + pad_h = (ph - (h % ph)) % ph + pad_w = (pw - (w % pw)) % pw + return F.pad(x, (0, pad_w, 0, pad_h, 0, pad_t), mode="replicate") + + +def center_down_sample_3d(x, kernel_size): + return F.avg_pool3d(x, kernel_size, stride=kernel_size) + + +def apply_rotary_emb_helios( + hidden_states: torch.Tensor, + freqs_cis: torch.Tensor, +) -> torch.Tensor: + """Apply Helios-style rotary embeddings. + + freqs_cis contains [cos_t, cos_y, cos_x, sin_t, sin_y, sin_x] concatenated + along the last dimension, with shape [B, seq, D*2] where D = DT+DY+DX. + hidden_states has shape [B, seq, H, head_dim]. + """ + x_1, x_2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) + cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1) + out = torch.empty_like(hidden_states) + out[..., 0::2] = x_1 * cos[..., 0::2] - x_2 * sin[..., 1::2] + out[..., 1::2] = x_1 * sin[..., 1::2] + x_2 * cos[..., 0::2] + return out.type_as(hidden_states) + + +class DistributedRMSNorm(nn.Module): + """RMSNorm that computes global RMS across tensor parallel ranks.""" + + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(hidden_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + tp_size = get_tensor_model_parallel_world_size() + + input_dtype = x.dtype + x_float = x.float() + local_sum_sq = (x_float**2).sum(dim=-1, keepdim=True) + local_count = x.shape[-1] + + if tp_size > 1: + global_sum_sq = local_sum_sq.clone() + tensor_model_parallel_all_reduce(global_sum_sq) + global_count = local_count * tp_size + else: + global_sum_sq = local_sum_sq + global_count = local_count + + rms = torch.sqrt(global_sum_sq / global_count + self.eps) + + output = (x_float / rms) * self.weight.float() + return output.to(input_dtype) + + +class ColumnParallelGELU(nn.Module): + """Column parallel linear with GELU activation.""" + + def __init__(self, dim_in: int, dim_out: int, *, approximate: str = "tanh", bias: bool = True): + super().__init__() + self.proj = ColumnParallelLinear( + dim_in, + dim_out, + bias=bias, + gather_output=False, + return_bias=False, + ) + self.approximate = approximate + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.proj(x) + return F.gelu(x, approximate=self.approximate) + + +class HeliosFeedForward(nn.Module): + """TP-enabled FeedForward network for Helios.""" + + def __init__( + self, + dim: int, + inner_dim: int, + dim_out: int | None = None, + bias: bool = True, + ) -> None: + super().__init__() + dim_out = dim_out or dim + + self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias) + self.net_1 = nn.Identity() + self.net_2 = RowParallelLinear( + inner_dim, + dim_out, + bias=bias, + input_is_parallel=True, + return_bias=False, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.net_0(hidden_states) + hidden_states = self.net_1(hidden_states) + hidden_states = self.net_2(hidden_states) + return hidden_states + + +class HeliosRotaryPosEmbed(nn.Module): + """Helios-style 3D rotary position embeddings using explicit frame indices.""" + + def __init__(self, rope_dim: tuple[int, int, int], theta: float): + super().__init__() + self.DT, self.DY, self.DX = rope_dim + self.theta = theta + self.register_buffer("freqs_base_t", self._get_freqs_base(self.DT), persistent=False) + self.register_buffer("freqs_base_y", self._get_freqs_base(self.DY), persistent=False) + self.register_buffer("freqs_base_x", self._get_freqs_base(self.DX), persistent=False) + + def _get_freqs_base(self, dim): + return 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32)[: (dim // 2)] / dim)) + + @torch.no_grad() + def get_frequency_batched(self, freqs_base, pos): + freqs = torch.einsum("d,bthw->dbthw", freqs_base, pos) + freqs = freqs.repeat_interleave(2, dim=0) + return freqs.cos(), freqs.sin() + + @torch.no_grad() + @lru_cache(maxsize=32) + def _get_spatial_meshgrid(self, height, width, device_str): + device = torch.device(device_str) + grid_y_coords = torch.arange(height, device=device, dtype=torch.float32) + grid_x_coords = torch.arange(width, device=device, dtype=torch.float32) + grid_y, grid_x = torch.meshgrid(grid_y_coords, grid_x_coords, indexing="ij") + return grid_y, grid_x + + @torch.no_grad() + def forward(self, frame_indices, height, width, device): + batch_size = frame_indices.shape[0] + num_frames = frame_indices.shape[1] + + frame_indices = frame_indices.to(device=device, dtype=torch.float32) + grid_y, grid_x = self._get_spatial_meshgrid(height, width, str(device)) + + grid_t = frame_indices[:, :, None, None].expand(batch_size, num_frames, height, width) + grid_y_batch = grid_y[None, None, :, :].expand(batch_size, num_frames, -1, -1) + grid_x_batch = grid_x[None, None, :, :].expand(batch_size, num_frames, -1, -1) + + freqs_cos_t, freqs_sin_t = self.get_frequency_batched(self.freqs_base_t, grid_t) + freqs_cos_y, freqs_sin_y = self.get_frequency_batched(self.freqs_base_y, grid_y_batch) + freqs_cos_x, freqs_sin_x = self.get_frequency_batched(self.freqs_base_x, grid_x_batch) + + result = torch.cat( + [freqs_cos_t, freqs_cos_y, freqs_cos_x, freqs_sin_t, freqs_sin_y, freqs_sin_x], + dim=0, + ) + + return result.permute(1, 0, 2, 3, 4) + + +class HeliosTimeTextEmbedding(nn.Module): + """Combined time and text condition embeddings for Helios.""" + + def __init__( + self, + dim: int, + time_freq_dim: int, + time_proj_dim: int, + text_embed_dim: int, + ): + super().__init__() + + self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim, time_embed_dim=dim) + self.act_fn = nn.SiLU() + self.time_proj = nn.Linear(dim, time_proj_dim) + self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim, act_fn="gelu_tanh") + + def forward( + self, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + is_return_encoder_hidden_states: bool = True, + ): + timestep = self.timesteps_proj(timestep) + + time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype + if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8: + timestep = timestep.to(time_embedder_dtype) + temb = self.time_embedder(timestep).type_as(encoder_hidden_states) + timestep_proj = self.time_proj(self.act_fn(temb)) + + if encoder_hidden_states is not None and is_return_encoder_hidden_states: + encoder_hidden_states = self.text_embedder(encoder_hidden_states) + + return temb, timestep_proj, encoder_hidden_states + + +class HeliosOutputNorm(nn.Module): + """Output normalization that extracts only original_context_length tokens.""" + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + self.norm = FP32LayerNorm(dim, eps, elementwise_affine=False) + + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor, original_context_length: int): + temb = temb[:, -original_context_length:, :] + shift, scale = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2) + shift, scale = shift.squeeze(2).to(hidden_states.device), scale.squeeze(2).to(hidden_states.device) + hidden_states = hidden_states[:, -original_context_length:, :] + hidden_states = (self.norm(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states) + return hidden_states + + +class HeliosSelfAttention(nn.Module): + """Optimized self-attention for Helios with history amplification support.""" + + def __init__( + self, + dim: int, + num_heads: int, + head_dim: int, + eps: float = 1e-5, + dropout: float = 0.0, + is_amplify_history: bool = False, + history_scale_mode: str = "per_head", + ): + super().__init__() + + self.dim = dim + self.num_heads = num_heads + self.head_dim = head_dim + self.inner_dim = num_heads * head_dim + + self.to_qkv = QKVParallelLinear( + hidden_size=dim, + head_size=head_dim, + total_num_heads=num_heads, + bias=True, + ) + + self.num_heads = self.to_qkv.num_heads + self.num_kv_heads = self.to_qkv.num_kv_heads + self.tp_inner_dim = self.num_heads * head_dim + + self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + + self.to_out = RowParallelLinear( + self.inner_dim, + dim, + bias=True, + input_is_parallel=True, + return_bias=False, + ) + self.dropout = nn.Dropout(dropout) + + self.attn = Attention( + num_heads=self.num_heads, + head_size=head_dim, + num_kv_heads=self.num_kv_heads, + softmax_scale=1.0 / (head_dim**0.5), + causal=False, + ) + + self.is_amplify_history = is_amplify_history + if is_amplify_history: + tp_size = get_tensor_model_parallel_world_size() + local_heads = num_heads // tp_size + if history_scale_mode == "scalar": + self.history_key_scale = nn.Parameter(torch.ones(1)) + elif history_scale_mode == "per_head": + self.history_key_scale = nn.Parameter(torch.ones(local_heads)) + else: + raise ValueError(f"Unknown history_scale_mode: {history_scale_mode}") + self.history_scale_mode = history_scale_mode + self.max_scale = 10.0 + + def forward( + self, + hidden_states: torch.Tensor, + rotary_emb: torch.Tensor | None = None, + original_context_length: int | None = None, + ) -> torch.Tensor: + qkv, _ = self.to_qkv(hidden_states) + + q_size = self.num_heads * self.head_dim + kv_size = self.num_kv_heads * self.head_dim + query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1) + + query = self.norm_q(query) + key = self.norm_k(key) + + query = query.unflatten(2, (self.num_heads, self.head_dim)) + key = key.unflatten(2, (self.num_kv_heads, self.head_dim)) + value = value.unflatten(2, (self.num_kv_heads, self.head_dim)) + + if rotary_emb is not None: + query = apply_rotary_emb_helios(query, rotary_emb) + key = apply_rotary_emb_helios(key, rotary_emb) + + if self.is_amplify_history and original_context_length is not None: + history_seq_len = hidden_states.shape[1] - original_context_length + if history_seq_len > 0: + scale_key = 1.0 + torch.sigmoid(self.history_key_scale) * (self.max_scale - 1.0) + if self.history_scale_mode == "per_head": + scale_key = scale_key.view(1, 1, -1, 1) + key = torch.cat( + [key[:, :history_seq_len] * scale_key, key[:, history_seq_len:]], + dim=1, + ) + + hidden_states = self.attn(query, key, value) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.type_as(query) + + hidden_states = self.to_out(hidden_states) + hidden_states = self.dropout(hidden_states) + + return hidden_states + + +class HeliosCrossAttention(nn.Module): + """Optimized cross-attention for Helios.""" + + def __init__( + self, + dim: int, + num_heads: int, + head_dim: int, + eps: float = 1e-5, + dropout: float = 0.0, + ): + super().__init__() + + self.dim = dim + self.num_heads = num_heads + self.head_dim = head_dim + self.inner_dim = num_heads * head_dim + + self.to_q = ColumnParallelLinear( + dim, + self.inner_dim, + bias=True, + gather_output=False, + return_bias=False, + ) + self.to_k = ColumnParallelLinear( + dim, + self.inner_dim, + bias=True, + gather_output=False, + return_bias=False, + ) + self.to_v = ColumnParallelLinear( + dim, + self.inner_dim, + bias=True, + gather_output=False, + return_bias=False, + ) + + tp_size = get_tensor_model_parallel_world_size() + self.num_heads = num_heads // tp_size + self.tp_inner_dim = self.num_heads * head_dim + + self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + + self.to_out = RowParallelLinear( + self.inner_dim, + dim, + bias=True, + input_is_parallel=True, + return_bias=False, + ) + self.dropout = nn.Dropout(dropout) + + self.attn = Attention( + num_heads=self.num_heads, + head_size=head_dim, + num_kv_heads=self.num_heads, + softmax_scale=1.0 / (head_dim**0.5), + causal=False, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + ) -> torch.Tensor: + query = self.to_q(hidden_states) + query = self.norm_q(query) + + key = self.to_k(encoder_hidden_states) + value = self.to_v(encoder_hidden_states) + key = self.norm_k(key) + + query = query.unflatten(2, (self.num_heads, self.head_dim)) + key = key.unflatten(2, (self.num_heads, self.head_dim)) + value = value.unflatten(2, (self.num_heads, self.head_dim)) + + hidden_states = self.attn(query, key, value) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.type_as(query) + + hidden_states = self.to_out(hidden_states) + hidden_states = self.dropout(hidden_states) + + return hidden_states + + +class HeliosTransformerBlock(nn.Module): + """Transformer block with guidance cross-attention and history support.""" + + def __init__( + self, + dim: int, + ffn_dim: int, + num_heads: int, + eps: float = 1e-6, + cross_attn_norm: bool = False, + guidance_cross_attn: bool = False, + is_amplify_history: bool = False, + history_scale_mode: str = "per_head", + ): + super().__init__() + + head_dim = dim // num_heads + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.attn1 = HeliosSelfAttention( + dim=dim, + num_heads=num_heads, + head_dim=head_dim, + eps=eps, + is_amplify_history=is_amplify_history, + history_scale_mode=history_scale_mode, + ) + + # 2. Cross-attention + self.attn2 = HeliosCrossAttention( + dim=dim, + num_heads=num_heads, + head_dim=head_dim, + eps=eps, + ) + self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + + # 3. Feed-forward + self.ffn = HeliosFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim) + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + # 4. Guidance cross-attention flag + self.guidance_cross_attn = guidance_cross_attn + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + rotary_emb: torch.Tensor, + original_context_length: int | None = None, + ) -> torch.Tensor: + if temb.ndim == 4: + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + self.scale_shift_table.unsqueeze(0) + temb.float() + ).chunk(6, dim=2) + shift_msa = shift_msa.squeeze(2) + scale_msa = scale_msa.squeeze(2) + gate_msa = gate_msa.squeeze(2) + c_shift_msa = c_shift_msa.squeeze(2) + c_scale_msa = c_scale_msa.squeeze(2) + c_gate_msa = c_gate_msa.squeeze(2) + else: + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + self.scale_shift_table + temb.float() + ).chunk(6, dim=1) + + # 1. Self-attention + norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states) + attn_output = self.attn1(norm_hidden_states, rotary_emb, original_context_length) + hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states) + + # 2. Cross-attention (with optional guidance: only current chunk attends to text) + if self.guidance_cross_attn and original_context_length is not None: + history_seq_len = hidden_states.shape[1] - original_context_length + + history_hidden_states, current_hidden_states = ( + hidden_states[:, :history_seq_len], + hidden_states[:, history_seq_len:], + ) + norm_hidden_states = self.norm2(current_hidden_states.float()).type_as(current_hidden_states) + attn_output = self.attn2(norm_hidden_states, encoder_hidden_states) + current_hidden_states = current_hidden_states + attn_output + hidden_states = torch.cat([history_hidden_states, current_hidden_states], dim=1) + else: + norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states) + attn_output = self.attn2(norm_hidden_states, encoder_hidden_states) + hidden_states = hidden_states + attn_output + + # 3. Feed-forward + norm_hidden_states = (self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as( + hidden_states + ) + ff_output = self.ffn(norm_hidden_states) + hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + + return hidden_states + + +class HeliosTransformer3DModel(nn.Module): + """Optimized Helios Transformer model for video generation using vLLM layers. + + Helios extends the Wan2.2 architecture with multi-term memory patches, + guidance cross-attention, and chunked video generation support. + """ + + _repeated_blocks = ["HeliosTransformerBlock"] + _layerwise_offload_blocks_attr = "blocks" + packed_modules_mapping = { + "to_qkv": ["to_q", "to_k", "to_v"], + } + + @staticmethod + def _is_transformer_block(name: str, module) -> bool: + return "blocks" in name and name.split(".")[-1].isdigit() + + _hsdp_shard_conditions = [_is_transformer_block] + + _sp_plan = { + "rope": { + 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), + 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), + }, + "blocks.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True), + }, + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + + def __init__( + self, + patch_size: tuple[int, ...] = (1, 2, 2), + num_attention_heads: int = 40, + attention_head_dim: int = 128, + in_channels: int = 16, + out_channels: int = 16, + text_dim: int = 4096, + freq_dim: int = 256, + ffn_dim: int = 13824, + num_layers: int = 40, + cross_attn_norm: bool = True, + qk_norm: str | None = "rms_norm_across_heads", + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + rope_dim: tuple[int, ...] = (44, 42, 42), + rope_theta: float = 10000.0, + guidance_cross_attn: bool = True, + zero_history_timestep: bool = True, + has_multi_term_memory_patch: bool = True, + is_amplify_history: bool = False, + history_scale_mode: str = "per_head", + ) -> None: + super().__init__() + + inner_dim = num_attention_heads * attention_head_dim + out_channels = out_channels or in_channels + + self.config = type( + "Config", + (), + { + "patch_size": patch_size, + "num_attention_heads": num_attention_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "text_dim": text_dim, + "freq_dim": freq_dim, + "ffn_dim": ffn_dim, + "num_layers": num_layers, + "cross_attn_norm": cross_attn_norm, + "qk_norm": qk_norm, + "eps": eps, + "added_kv_proj_dim": added_kv_proj_dim, + "rope_dim": rope_dim, + "rope_theta": rope_theta, + "guidance_cross_attn": guidance_cross_attn, + "zero_history_timestep": zero_history_timestep, + "has_multi_term_memory_patch": has_multi_term_memory_patch, + "is_amplify_history": is_amplify_history, + "history_scale_mode": history_scale_mode, + }, + )() + + self.inner_dim = inner_dim + self.zero_history_timestep = zero_history_timestep + self.has_multi_term_memory_patch = has_multi_term_memory_patch + + # 1. Patch & position embedding + self.rope = HeliosRotaryPosEmbed(rope_dim=rope_dim, theta=rope_theta) + self.patch_embedding = Conv3dLayer( + in_channels=in_channels, + out_channels=inner_dim, + kernel_size=patch_size, + stride=patch_size, + ) + + # 2. Multi-term memory patches + if has_multi_term_memory_patch: + self.patch_short = Conv3dLayer( + in_channels=in_channels, + out_channels=inner_dim, + kernel_size=(1, 2, 2), + stride=(1, 2, 2), + ) + self.patch_mid = Conv3dLayer( + in_channels=in_channels, + out_channels=inner_dim, + kernel_size=(2, 4, 4), + stride=(2, 4, 4), + ) + self.patch_long = Conv3dLayer( + in_channels=in_channels, + out_channels=inner_dim, + kernel_size=(4, 8, 8), + stride=(4, 8, 8), + ) + + # 3. Condition embeddings + self.condition_embedder = HeliosTimeTextEmbedding( + dim=inner_dim, + time_freq_dim=freq_dim, + time_proj_dim=inner_dim * 6, + text_embed_dim=text_dim, + ) + + # 4. Transformer blocks + self.blocks = nn.ModuleList( + [ + HeliosTransformerBlock( + inner_dim, + ffn_dim, + num_attention_heads, + eps, + cross_attn_norm, + guidance_cross_attn=guidance_cross_attn, + is_amplify_history=is_amplify_history, + history_scale_mode=history_scale_mode, + ) + for _ in range(num_layers) + ] + ) + + # 5. Output norm & projection + self.norm_out = HeliosOutputNorm(inner_dim, eps) + self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) + + @property + def dtype(self) -> torch.dtype: + return next(self.parameters()).dtype + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.LongTensor, + encoder_hidden_states: torch.Tensor, + indices_hidden_states: torch.Tensor | None = None, + indices_latents_history_short: torch.Tensor | None = None, + indices_latents_history_mid: torch.Tensor | None = None, + indices_latents_history_long: torch.Tensor | None = None, + latents_history_short: torch.Tensor | None = None, + latents_history_mid: torch.Tensor | None = None, + latents_history_long: torch.Tensor | None = None, + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + ) -> torch.Tensor | Transformer2DModelOutput: + batch_size = hidden_states.shape[0] + p_t, p_h, p_w = self.config.patch_size + + # 1. Process noisy latents + hidden_states = self.patch_embedding(hidden_states) + _, _, post_patch_num_frames, post_patch_height, post_patch_width = hidden_states.shape + + if indices_hidden_states is None: + indices_hidden_states = torch.arange(0, post_patch_num_frames).unsqueeze(0).expand(batch_size, -1) + + hidden_states = hidden_states.flatten(2).transpose(1, 2) + rotary_emb = self.rope( + frame_indices=indices_hidden_states, + height=post_patch_height, + width=post_patch_width, + device=hidden_states.device, + ) + rotary_emb = rotary_emb.flatten(2).transpose(1, 2) + original_context_length = hidden_states.shape[1] + + # 2. Process short history latents + if latents_history_short is not None and indices_latents_history_short is not None: + latents_history_short = latents_history_short.to(hidden_states) + latents_history_short = self.patch_short(latents_history_short) + _, _, _, H1, W1 = latents_history_short.shape + latents_history_short = latents_history_short.flatten(2).transpose(1, 2) + + rotary_emb_history_short = self.rope( + frame_indices=indices_latents_history_short, + height=H1, + width=W1, + device=latents_history_short.device, + ) + rotary_emb_history_short = rotary_emb_history_short.flatten(2).transpose(1, 2) + + hidden_states = torch.cat([latents_history_short, hidden_states], dim=1) + rotary_emb = torch.cat([rotary_emb_history_short, rotary_emb], dim=1) + + # 3. Process mid history latents + if latents_history_mid is not None and indices_latents_history_mid is not None: + latents_history_mid = latents_history_mid.to(hidden_states) + latents_history_mid = pad_for_3d_conv(latents_history_mid, (2, 4, 4)) + latents_history_mid = self.patch_mid(latents_history_mid) + latents_history_mid = latents_history_mid.flatten(2).transpose(1, 2) + + rotary_emb_history_mid = self.rope( + frame_indices=indices_latents_history_mid, + height=H1, + width=W1, + device=latents_history_mid.device, + ) + rotary_emb_history_mid = pad_for_3d_conv(rotary_emb_history_mid, (2, 2, 2)) + rotary_emb_history_mid = center_down_sample_3d(rotary_emb_history_mid, (2, 2, 2)) + rotary_emb_history_mid = rotary_emb_history_mid.flatten(2).transpose(1, 2) + + hidden_states = torch.cat([latents_history_mid, hidden_states], dim=1) + rotary_emb = torch.cat([rotary_emb_history_mid, rotary_emb], dim=1) + + # 4. Process long history latents + if latents_history_long is not None and indices_latents_history_long is not None: + latents_history_long = latents_history_long.to(hidden_states) + latents_history_long = pad_for_3d_conv(latents_history_long, (4, 8, 8)) + latents_history_long = self.patch_long(latents_history_long) + latents_history_long = latents_history_long.flatten(2).transpose(1, 2) + + rotary_emb_history_long = self.rope( + frame_indices=indices_latents_history_long, + height=H1, + width=W1, + device=latents_history_long.device, + ) + rotary_emb_history_long = pad_for_3d_conv(rotary_emb_history_long, (4, 4, 4)) + rotary_emb_history_long = center_down_sample_3d(rotary_emb_history_long, (4, 4, 4)) + rotary_emb_history_long = rotary_emb_history_long.flatten(2).transpose(1, 2) + + hidden_states = torch.cat([latents_history_long, hidden_states], dim=1) + rotary_emb = torch.cat([rotary_emb_history_long, rotary_emb], dim=1) + + history_context_length = hidden_states.shape[1] - original_context_length + + # 5. Compute timestep embeddings + if indices_hidden_states is not None and self.zero_history_timestep: + timestep_t0 = torch.zeros((1), dtype=timestep.dtype, device=timestep.device) + temb_t0, timestep_proj_t0, _ = self.condition_embedder( + timestep_t0, encoder_hidden_states, is_return_encoder_hidden_states=False + ) + temb_t0 = temb_t0.unsqueeze(1).expand(batch_size, history_context_length, -1) + timestep_proj_t0 = ( + timestep_proj_t0.unflatten(-1, (6, -1)) + .view(1, 6, 1, -1) + .expand(batch_size, -1, history_context_length, -1) + ) + + temb, timestep_proj, encoder_hidden_states = self.condition_embedder(timestep, encoder_hidden_states) + timestep_proj = timestep_proj.unflatten(-1, (6, -1)) + + if indices_hidden_states is not None and not self.zero_history_timestep: + main_repeat_size = hidden_states.shape[1] + else: + main_repeat_size = original_context_length + temb = temb.view(batch_size, 1, -1).expand(batch_size, main_repeat_size, -1) + timestep_proj = timestep_proj.view(batch_size, 6, 1, -1).expand(batch_size, 6, main_repeat_size, -1) + + if indices_hidden_states is not None and self.zero_history_timestep: + temb = torch.cat([temb_t0, temb], dim=1) + timestep_proj = torch.cat([timestep_proj_t0, timestep_proj], dim=2) + + if timestep_proj.ndim == 4: + timestep_proj = timestep_proj.permute(0, 2, 1, 3) + + # 6. Transformer blocks + hidden_states = hidden_states.contiguous() + encoder_hidden_states = encoder_hidden_states.contiguous() + rotary_emb = rotary_emb.contiguous() + + for block in self.blocks: + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep_proj, + rotary_emb, + original_context_length, + ) + + # 7. Output normalization + hidden_states = self.norm_out(hidden_states, temb, original_context_length) + hidden_states = self.proj_out(hidden_states) + + # 8. Unpatchify + hidden_states = hidden_states.reshape( + batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1 + ) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights with QKV fusion, FFN remapping, and TP norm sharding.""" + tp_rank = get_tensor_model_parallel_rank() + tp_size = get_tensor_model_parallel_world_size() + + stacked_params_mapping = [ + (".attn1.to_qkv", ".attn1.to_q", "q"), + (".attn1.to_qkv", ".attn1.to_k", "k"), + (".attn1.to_qkv", ".attn1.to_v", "v"), + ] + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + original_name = name + lookup_name = name + + # Handle QKV fusion for self-attention + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in original_name: + continue + lookup_name = original_name.replace(weight_name, param_name) + param = params_dict[lookup_name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # FFN remapping: net.0 -> net_0, net.2 -> net_2 + if ".ffn.net.0." in lookup_name: + lookup_name = lookup_name.replace(".ffn.net.0.", ".ffn.net_0.") + elif ".ffn.net.2." in lookup_name: + lookup_name = lookup_name.replace(".ffn.net.2.", ".ffn.net_2.") + + # Output projection: to_out.0 -> to_out + if ".to_out.0." in lookup_name: + lookup_name = lookup_name.replace(".to_out.0.", ".to_out.") + + if lookup_name not in params_dict: + logger.warning(f"Skipping weight {original_name} -> {lookup_name}") + continue + + param = params_dict[lookup_name] + + # TP sharding for RMSNorm weights + if tp_size > 1 and any( + norm_name in lookup_name + for norm_name in [ + ".attn1.norm_q.", + ".attn1.norm_k.", + ".attn2.norm_q.", + ".attn2.norm_k.", + ] + ): + shard_size = loaded_weight.shape[0] // tp_size + loaded_weight = loaded_weight[tp_rank * shard_size : (tp_rank + 1) * shard_size] + + # TP sharding for history_key_scale (per_head mode) + if tp_size > 1 and ".history_key_scale" in lookup_name and loaded_weight.dim() == 1: + shard_size = loaded_weight.shape[0] // tp_size + if shard_size > 0: + loaded_weight = loaded_weight[tp_rank * shard_size : (tp_rank + 1) * shard_size] + + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + loaded_params.add(original_name) + loaded_params.add(lookup_name) + + return loaded_params diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py new file mode 100644 index 000000000..0a89d2d02 --- /dev/null +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -0,0 +1,1064 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from Helios (https://github.com/BestWishYsh/Helios) + +from __future__ import annotations + +import json +import logging +import math +import os +from collections.abc import Iterable +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F +from diffusers import AutoencoderKLWan +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from transformers import AutoTokenizer, UMT5EncoderModel +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel +from vllm_omni.diffusion.models.helios.scheduling_helios import HeliosScheduler +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.platforms import current_omni_platform + +logger = logging.getLogger(__name__) + + +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +def optimized_scale(positive_flat, negative_flat): + dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) + squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8 + st_star = dot_product / squared_norm + return st_star + + +def load_json_config(model_path: str, subfolder: str, filename: str, local_files_only: bool = True) -> dict: + """Load a JSON config file from a local path or HuggingFace Hub repo.""" + if local_files_only: + config_path = os.path.join(model_path, subfolder, filename) + if os.path.exists(config_path): + with open(config_path) as f: + return json.load(f) + else: + try: + from huggingface_hub import hf_hub_download + + config_path = hf_hub_download( + repo_id=model_path, + filename=f"{subfolder}/{filename}", + ) + with open(config_path) as f: + return json.load(f) + except Exception: + pass + return {} + + +def load_transformer_config(model_path: str, subfolder: str = "transformer", local_files_only: bool = True) -> dict: + return load_json_config(model_path, subfolder, "config.json", local_files_only) + + +def create_transformer_from_config(config: dict) -> HeliosTransformer3DModel: + kwargs = {} + + key_map = [ + "patch_size", + "num_attention_heads", + "attention_head_dim", + "in_channels", + "out_channels", + "text_dim", + "freq_dim", + "ffn_dim", + "num_layers", + "cross_attn_norm", + "qk_norm", + "eps", + "added_kv_proj_dim", + "rope_dim", + "rope_theta", + "guidance_cross_attn", + "zero_history_timestep", + "has_multi_term_memory_patch", + "is_amplify_history", + "history_scale_mode", + ] + for key in key_map: + if key in config: + val = config[key] + if key in ("patch_size", "rope_dim") and isinstance(val, list): + val = tuple(val) + kwargs[key] = val + + return HeliosTransformer3DModel(**kwargs) + + +def get_helios_post_process_func( + od_config: OmniDiffusionConfig, +): + from diffusers.video_processor import VideoProcessor + + video_processor = VideoProcessor(vae_scale_factor=8) + + def post_process_func( + video: torch.Tensor, + output_type: str = "np", + ): + if output_type == "latent": + return video + return video_processor.postprocess_video(video, output_type=output_type) + + return post_process_func + + +def get_helios_pre_process_func( + od_config: OmniDiffusionConfig, +): + def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: + return request + + return pre_process_func + + +class HeliosPipeline(nn.Module, CFGParallelMixin): + """Helios text-to-video / image-to-video / video-to-video pipeline for vllm-omni. + + Supports T2V, I2V (with image input), and V2V (with video input). + Implements chunked video generation with multi-term memory history context. + """ + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + + self.device = get_local_device() + dtype = getattr(od_config, "dtype", torch.bfloat16) + + model = od_config.model + local_files_only = os.path.exists(model) + + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ) + ] + + self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + self.text_encoder = UMT5EncoderModel.from_pretrained( + model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only + ).to(self.device) + self.vae = AutoencoderKLWan.from_pretrained( + model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only + ).to(self.device) + + transformer_config = load_transformer_config(model, "transformer", local_files_only) + self.transformer = create_transformer_from_config(transformer_config) + + # Read scheduler config to determine scheduler type + sched_cfg = load_json_config(model, "scheduler", "scheduler_config.json", local_files_only) + scheduler_kwargs = {} + passthrough_keys = [ + "num_train_timesteps", + "shift", + "stages", + "stage_range", + "gamma", + "thresholding", + "prediction_type", + "solver_order", + "predict_x0", + "solver_type", + "lower_order_final", + "disable_corrector", + "use_flow_sigmas", + "scheduler_type", + "use_dynamic_shifting", + "time_shift_type", + ] + for key in passthrough_keys: + if key in sched_cfg: + scheduler_kwargs[key] = sched_cfg[key] + + self.scheduler = HeliosScheduler(**scheduler_kwargs) + + self.is_distilled = scheduler_kwargs.get("scheduler_type") == "dmd" + + self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4 + self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8 + + self._guidance_scale = None + self._num_timesteps = None + self._current_timestep = None + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale is not None and self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + def forward( + self, + req: OmniDiffusionRequest, + prompt: str | None = None, + negative_prompt: str | None = None, + height: int = 384, + width: int = 640, + num_inference_steps: int = 50, + guidance_scale: float = 5.0, + frame_num: int = 132, + output_type: str | None = "np", + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + attention_kwargs: dict | None = None, + # Helios-specific + history_sizes: list | None = None, + num_latent_frames_per_chunk: int = 9, + keep_first_frame: bool = True, + # I2V + image: torch.Tensor | None = None, + add_noise_to_image_latents: bool = True, + image_noise_sigma_min: float = 0.111, + image_noise_sigma_max: float = 0.135, + # V2V + video: torch.Tensor | None = None, + add_noise_to_video_latents: bool = True, + video_noise_sigma_min: float = 0.111, + video_noise_sigma_max: float = 0.135, + # Stage 2 (pyramid multi-stage denoising) + is_enable_stage2: bool = False, + pyramid_num_stages: int = 3, + pyramid_num_inference_steps_list: list | None = None, + is_skip_first_chunk: bool = False, + # DMD + is_amplify_first_chunk: bool = False, + # CFG Zero Star + use_cfg_zero_star: bool = False, + use_zero_init: bool = True, + zero_steps: int = 1, + **kwargs, + ) -> DiffusionOutput: + if pyramid_num_inference_steps_list is None: + pyramid_num_inference_steps_list = [10, 10, 10] + if history_sizes is None: + history_sizes = [16, 2, 1] + + # Read Helios-specific params from extra_args + extra = getattr(req.sampling_params, "extra_args", {}) or {} + is_enable_stage2 = extra.get("is_enable_stage2", is_enable_stage2) + pyramid_num_stages = extra.get("pyramid_num_stages", pyramid_num_stages) + pyramid_num_inference_steps_list = extra.get( + "pyramid_num_inference_steps_list", pyramid_num_inference_steps_list + ) + is_amplify_first_chunk = extra.get("is_amplify_first_chunk", is_amplify_first_chunk) + use_cfg_zero_star = extra.get("use_cfg_zero_star", use_cfg_zero_star) + use_zero_init = extra.get("use_zero_init", use_zero_init) + zero_steps = extra.get("zero_steps", zero_steps) + is_skip_first_chunk = extra.get("is_skip_first_chunk", is_skip_first_chunk) + + image = extra.get("image", image) + video = extra.get("video", video) + add_noise_to_image_latents = extra.get("add_noise_to_image_latents", add_noise_to_image_latents) + image_noise_sigma_min = extra.get("image_noise_sigma_min", image_noise_sigma_min) + image_noise_sigma_max = extra.get("image_noise_sigma_max", image_noise_sigma_max) + add_noise_to_video_latents = extra.get("add_noise_to_video_latents", add_noise_to_video_latents) + video_noise_sigma_min = extra.get("video_noise_sigma_min", video_noise_sigma_min) + video_noise_sigma_max = extra.get("video_noise_sigma_max", video_noise_sigma_max) + + if image is not None and video is not None: + raise ValueError("image and video cannot be provided simultaneously") + + if len(req.prompts) > 1: + raise ValueError("This model only supports a single prompt, not a batched request.") + if len(req.prompts) == 1: + prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt") + negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt") + + if prompt is None and prompt_embeds is None: + raise ValueError("Prompt or prompt_embeds is required for Helios generation.") + + height = req.sampling_params.height or height + width = req.sampling_params.width or width + num_frames = req.sampling_params.num_frames if req.sampling_params.num_frames else frame_num + num_steps = req.sampling_params.num_inference_steps or num_inference_steps + + if req.sampling_params.guidance_scale_provided: + guidance_scale = req.sampling_params.guidance_scale + + self._guidance_scale = guidance_scale + + height = (height // 16) * 16 + width = (width // 16) * 16 + num_frames = max(num_frames, 1) + + device = self.device + dtype = self.transformer.dtype + + if generator is None: + generator = req.sampling_params.generator + if generator is None and req.sampling_params.seed is not None: + generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) + + # Encode prompts + if prompt_embeds is None: + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, + max_sequence_length=req.sampling_params.max_sequence_length or 226, + device=device, + dtype=dtype, + ) + else: + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + if negative_prompt_embeds is not None: + negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + + batch_size = prompt_embeds.shape[0] + + history_sizes = sorted(history_sizes, reverse=True) + + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(self.vae.device, self.vae.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + self.vae.device, self.vae.dtype + ) + + # Prepare I2V image latents + fake_image_latents = None + image_latents = None + if image is not None: + image_latents, fake_image_latents = self.prepare_image_latents( + image, + latents_mean=latents_mean, + latents_std=latents_std, + num_latent_frames_per_chunk=num_latent_frames_per_chunk, + dtype=torch.float32, + device=device, + generator=generator, + ) + + if image_latents is not None and add_noise_to_image_latents: + image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min) + + image_noise_sigma_min + ) + image_latents = ( + image_noise_sigma * randn_tensor(image_latents.shape, generator=generator, device=device) + + (1 - image_noise_sigma) * image_latents + ) + fake_image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (video_noise_sigma_max - video_noise_sigma_min) + + video_noise_sigma_min + ) + fake_image_latents = ( + fake_image_noise_sigma * randn_tensor(fake_image_latents.shape, generator=generator, device=device) + + (1 - fake_image_noise_sigma) * fake_image_latents + ) + + # Prepare V2V video latents + video_latents = None + if video is not None: + image_latents, video_latents = self.prepare_video_latents( + video, + latents_mean=latents_mean, + latents_std=latents_std, + num_latent_frames_per_chunk=num_latent_frames_per_chunk, + dtype=torch.float32, + device=device, + generator=generator, + ) + + if video_latents is not None and add_noise_to_video_latents: + image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min) + + image_noise_sigma_min + ) + image_latents = ( + image_noise_sigma * randn_tensor(image_latents.shape, generator=generator, device=device) + + (1 - image_noise_sigma) * image_latents + ) + + noisy_latents_chunks = [] + num_latent_chunks = video_latents.shape[2] // num_latent_frames_per_chunk + for i in range(num_latent_chunks): + chunk_start = i * num_latent_frames_per_chunk + chunk_end = chunk_start + num_latent_frames_per_chunk + latent_chunk = video_latents[:, :, chunk_start:chunk_end, :, :] + chunk_frames = latent_chunk.shape[2] + frame_sigmas = ( + torch.rand(chunk_frames, device=device, generator=generator) + * (video_noise_sigma_max - video_noise_sigma_min) + + video_noise_sigma_min + ) + frame_sigmas = frame_sigmas.view(1, 1, chunk_frames, 1, 1) + noisy_chunk = ( + frame_sigmas * randn_tensor(latent_chunk.shape, generator=generator, device=device) + + (1 - frame_sigmas) * latent_chunk + ) + noisy_latents_chunks.append(noisy_chunk) + video_latents = torch.cat(noisy_latents_chunks, dim=2) + + # Prepare latent variables + num_channels_latents = self.transformer.config.in_channels + window_num_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 + num_latent_chunk = max(1, (num_frames + window_num_frames - 1) // window_num_frames) + num_history_latent_frames = sum(history_sizes) + history_video = None + total_generated_latent_frames = 0 + + if not keep_first_frame: + history_sizes[-1] = history_sizes[-1] + 1 + + history_latents = torch.zeros( + batch_size, + num_channels_latents, + num_history_latent_frames, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + device=device, + dtype=torch.float32, + ) + + if fake_image_latents is not None: + history_latents = torch.cat([history_latents[:, :, :-1, :, :], fake_image_latents], dim=2) + total_generated_latent_frames += 1 + + if video_latents is not None: + history_frames = history_latents.shape[2] + video_frames = video_latents.shape[2] + if video_frames < history_frames: + keep_frames = history_frames - video_frames + history_latents = torch.cat([history_latents[:, :, :keep_frames, :, :], video_latents], dim=2) + else: + history_latents = video_latents + total_generated_latent_frames += video_latents.shape[2] + + # Prepare frame indices + if keep_first_frame: + indices = torch.arange(0, sum([1, *history_sizes, num_latent_frames_per_chunk])) + ( + indices_prefix, + indices_latents_history_long, + indices_latents_history_mid, + indices_latents_history_1x, + indices_hidden_states, + ) = indices.split([1, *history_sizes, num_latent_frames_per_chunk], dim=0) + indices_latents_history_short = torch.cat([indices_prefix, indices_latents_history_1x], dim=0) + else: + indices = torch.arange(0, sum([*history_sizes, num_latent_frames_per_chunk])) + ( + indices_latents_history_long, + indices_latents_history_mid, + indices_latents_history_short, + indices_hidden_states, + ) = indices.split([*history_sizes, num_latent_frames_per_chunk], dim=0) + + indices_hidden_states = indices_hidden_states.unsqueeze(0) + indices_latents_history_short = indices_latents_history_short.unsqueeze(0) + indices_latents_history_mid = indices_latents_history_mid.unsqueeze(0) + indices_latents_history_long = indices_latents_history_long.unsqueeze(0) + + if attention_kwargs is None: + attention_kwargs = {} + + vae_dtype = self.vae.dtype + + # Chunked denoising loop + for k in range(num_latent_chunk): + is_first_chunk = k == 0 + is_second_chunk = k == 1 + + if keep_first_frame: + latents_history_long, latents_history_mid, latents_history_1x = history_latents[ + :, :, -num_history_latent_frames: + ].split(history_sizes, dim=2) + if image_latents is None and is_first_chunk: + latents_prefix = torch.zeros( + ( + batch_size, + num_channels_latents, + 1, + latents_history_1x.shape[-2], + latents_history_1x.shape[-1], + ), + device=latents_history_1x.device, + dtype=latents_history_1x.dtype, + ) + else: + latents_prefix = image_latents + latents_history_short = torch.cat([latents_prefix, latents_history_1x], dim=2) + else: + latents_history_long, latents_history_mid, latents_history_short = history_latents[ + :, :, -num_history_latent_frames: + ].split(history_sizes, dim=2) + + # Prepare noise for this chunk + latents = self.prepare_latents( + batch_size, + num_channels_latents, + height, + width, + window_num_frames, + dtype=torch.float32, + device=device, + generator=generator, + latents=None, + ) + + if not is_enable_stage2: + # Stage 1 only: single-stage denoising + patch_size = self.transformer.config.patch_size + image_seq_len = (latents.shape[-1] * latents.shape[-2] * latents.shape[-3]) // ( + patch_size[0] * patch_size[1] * patch_size[2] + ) + sigmas = np.linspace(0.999, 0.0, num_steps + 1)[:-1] + mu = calculate_shift(image_seq_len) + self.scheduler.set_timesteps(num_steps, device=device, sigmas=sigmas, mu=mu) + timesteps = self.scheduler.timesteps + self._num_timesteps = len(timesteps) + + latents = self._stage1_sample( + latents=latents, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + timesteps=timesteps, + guidance_scale=guidance_scale, + indices_hidden_states=indices_hidden_states, + indices_latents_history_short=indices_latents_history_short, + indices_latents_history_mid=indices_latents_history_mid, + indices_latents_history_long=indices_latents_history_long, + latents_history_short=latents_history_short, + latents_history_mid=latents_history_mid, + latents_history_long=latents_history_long, + attention_kwargs=attention_kwargs, + transformer_dtype=dtype, + use_cfg_zero_star=use_cfg_zero_star, + use_zero_init=use_zero_init, + zero_steps=zero_steps, + ) + else: + # Stage 2: pyramid multi-stage denoising + latents = self._stage2_sample( + latents=latents, + pyramid_num_stages=pyramid_num_stages, + pyramid_num_inference_steps_list=pyramid_num_inference_steps_list, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + guidance_scale=guidance_scale, + indices_hidden_states=indices_hidden_states, + indices_latents_history_short=indices_latents_history_short, + indices_latents_history_mid=indices_latents_history_mid, + indices_latents_history_long=indices_latents_history_long, + latents_history_short=latents_history_short, + latents_history_mid=latents_history_mid, + latents_history_long=latents_history_long, + attention_kwargs=attention_kwargs, + transformer_dtype=dtype, + is_amplify_first_chunk=is_amplify_first_chunk and is_first_chunk, + use_cfg_zero_star=use_cfg_zero_star, + use_zero_init=use_zero_init, + zero_steps=zero_steps, + device=device, + ) + + if keep_first_frame and ( + (is_first_chunk and image_latents is None) or (is_skip_first_chunk and is_second_chunk) + ): + image_latents = latents[:, :, 0:1, :, :] + + total_generated_latent_frames += latents.shape[2] + history_latents = torch.cat([history_latents, latents], dim=2) + real_history_latents = history_latents[:, :, -total_generated_latent_frames:] + index_slice = ( + slice(None), + slice(None), + slice(-num_latent_frames_per_chunk, None), + ) + + current_latents = real_history_latents[index_slice].to(vae_dtype) / latents_std + latents_mean + current_video = self.vae.decode(current_latents, return_dict=False)[0] + + if history_video is None: + history_video = current_video + else: + history_video = torch.cat([history_video, current_video], dim=2) + + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + self._current_timestep = None + + if output_type == "latent": + output = real_history_latents + else: + output = history_video + + return DiffusionOutput(output=output) + + def _stage1_sample( + self, + latents: torch.Tensor, + prompt_embeds: torch.Tensor, + negative_prompt_embeds: torch.Tensor | None, + timesteps: torch.Tensor, + guidance_scale: float, + indices_hidden_states: torch.Tensor, + indices_latents_history_short: torch.Tensor, + indices_latents_history_mid: torch.Tensor, + indices_latents_history_long: torch.Tensor, + latents_history_short: torch.Tensor, + latents_history_mid: torch.Tensor, + latents_history_long: torch.Tensor, + attention_kwargs: dict, + transformer_dtype: torch.dtype, + use_cfg_zero_star: bool = False, + use_zero_init: bool = True, + zero_steps: int = 1, + ) -> torch.Tensor: + """Single-stage denoising loop for one chunk.""" + batch_size = latents.shape[0] + do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None + + for i, t in enumerate(timesteps): + self._current_timestep = t + timestep = t.expand(batch_size) + + transformer_kwargs = { + "hidden_states": latents.to(transformer_dtype), + "timestep": timestep, + "indices_hidden_states": indices_hidden_states, + "indices_latents_history_short": indices_latents_history_short, + "indices_latents_history_mid": indices_latents_history_mid, + "indices_latents_history_long": indices_latents_history_long, + "latents_history_short": latents_history_short.to(transformer_dtype), + "latents_history_mid": latents_history_mid.to(transformer_dtype), + "latents_history_long": latents_history_long.to(transformer_dtype), + "attention_kwargs": attention_kwargs, + "return_dict": False, + } + + if use_cfg_zero_star and do_true_cfg: + noise_pred = self.transformer( + encoder_hidden_states=prompt_embeds, + **transformer_kwargs, + )[0] + + noise_uncond = self.transformer( + encoder_hidden_states=negative_prompt_embeds, + **transformer_kwargs, + )[0] + + positive_flat = noise_pred.view(batch_size, -1) + negative_flat = noise_uncond.view(batch_size, -1) + alpha_cfg = optimized_scale(positive_flat, negative_flat) + alpha_cfg = alpha_cfg.view(batch_size, *([1] * (len(noise_pred.shape) - 1))) + alpha_cfg = alpha_cfg.to(noise_pred.dtype) + + if (i <= zero_steps) and use_zero_init: + noise_pred = noise_pred * 0.0 + else: + noise_pred = noise_uncond * alpha_cfg + guidance_scale * (noise_pred - noise_uncond * alpha_cfg) + else: + positive_kwargs = { + "encoder_hidden_states": prompt_embeds, + **transformer_kwargs, + } + if do_true_cfg: + negative_kwargs = { + "encoder_hidden_states": negative_prompt_embeds, + **transformer_kwargs, + } + else: + negative_kwargs = None + + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=guidance_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) + + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + + return latents + + def _stage2_sample( + self, + latents: torch.Tensor, + pyramid_num_stages: int, + pyramid_num_inference_steps_list: list[int], + prompt_embeds: torch.Tensor, + negative_prompt_embeds: torch.Tensor | None, + guidance_scale: float, + indices_hidden_states: torch.Tensor, + indices_latents_history_short: torch.Tensor, + indices_latents_history_mid: torch.Tensor, + indices_latents_history_long: torch.Tensor, + latents_history_short: torch.Tensor, + latents_history_mid: torch.Tensor, + latents_history_long: torch.Tensor, + attention_kwargs: dict, + transformer_dtype: torch.dtype, + is_amplify_first_chunk: bool = False, + use_cfg_zero_star: bool = False, + use_zero_init: bool = True, + zero_steps: int = 1, + device: torch.device | None = None, + ) -> torch.Tensor: + """Pyramid multi-stage denoising for one chunk.""" + batch_size, num_channel, num_frames_lat, height, width = latents.shape + + # Downsample latents to the smallest pyramid level + latents_flat = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames_lat, num_channel, height, width) + for _ in range(pyramid_num_stages - 1): + height //= 2 + width //= 2 + latents_flat = F.interpolate(latents_flat, size=(height, width), mode="bilinear") * 2 + latents = latents_flat.reshape(batch_size, num_frames_lat, num_channel, height, width).permute(0, 2, 1, 3, 4) + + start_point_list = None + if self.is_distilled: + start_point_list = [latents] + + do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None + + for i_s in range(pyramid_num_stages): + patch_size = self.transformer.config.patch_size + image_seq_len = (latents.shape[-1] * latents.shape[-2] * latents.shape[-3]) // ( + patch_size[0] * patch_size[1] * patch_size[2] + ) + mu = calculate_shift(image_seq_len) + self.scheduler.set_timesteps( + pyramid_num_inference_steps_list[i_s], + stage_index=i_s, + device=device, + mu=mu, + is_amplify_first_chunk=is_amplify_first_chunk, + ) + timesteps = self.scheduler.timesteps + + if i_s > 0: + # Upsample latents to next pyramid level + height *= 2 + width *= 2 + num_frames_cur = latents.shape[2] + latents_flat = latents.permute(0, 2, 1, 3, 4).reshape( + batch_size * num_frames_cur, num_channel, height // 2, width // 2 + ) + latents_flat = F.interpolate(latents_flat, size=(height, width), mode="nearest") + latents = latents_flat.reshape(batch_size, num_frames_cur, num_channel, height, width).permute( + 0, 2, 1, 3, 4 + ) + + # Add block noise to fix artifacts between stages + ori_sigma = 1 - self.scheduler.ori_start_sigmas[i_s] + gamma = self.scheduler.config.gamma + alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma) + beta = alpha * (1 - ori_sigma) / math.sqrt(gamma) + + noise = self.sample_block_noise(batch_size, num_channel, latents.shape[2], height, width, patch_size) + noise = noise.to(device=device, dtype=transformer_dtype) + latents = alpha * latents + beta * noise + + if self.is_distilled and start_point_list is not None: + start_point_list.append(latents) + + for idx, t in enumerate(timesteps): + self._current_timestep = t + timestep = t.expand(latents.shape[0]).to(torch.int64) + + transformer_kwargs = { + "hidden_states": latents.to(transformer_dtype), + "timestep": timestep, + "indices_hidden_states": indices_hidden_states, + "indices_latents_history_short": indices_latents_history_short, + "indices_latents_history_mid": indices_latents_history_mid, + "indices_latents_history_long": indices_latents_history_long, + "latents_history_short": latents_history_short.to(transformer_dtype), + "latents_history_mid": latents_history_mid.to(transformer_dtype), + "latents_history_long": latents_history_long.to(transformer_dtype), + "attention_kwargs": attention_kwargs, + "return_dict": False, + } + + noise_pred = self.transformer( + encoder_hidden_states=prompt_embeds, + **transformer_kwargs, + )[0] + + if do_true_cfg: + noise_uncond = self.transformer( + encoder_hidden_states=negative_prompt_embeds, + **transformer_kwargs, + )[0] + + if use_cfg_zero_star: + positive_flat = noise_pred.view(batch_size, -1) + negative_flat = noise_uncond.view(batch_size, -1) + alpha_cfg = optimized_scale(positive_flat, negative_flat) + alpha_cfg = alpha_cfg.view(batch_size, *([1] * (len(noise_pred.shape) - 1))) + alpha_cfg = alpha_cfg.to(noise_pred.dtype) + + if (i_s == 0 and idx <= zero_steps) and use_zero_init: + noise_pred = noise_pred * 0.0 + else: + noise_pred = noise_uncond * alpha_cfg + guidance_scale * ( + noise_pred - noise_uncond * alpha_cfg + ) + else: + noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond) + + latents = self.scheduler.step( + noise_pred, + t, + latents, + return_dict=False, + cur_sampling_step=idx, + dmd_noisy_tensor=start_point_list[i_s] if start_point_list is not None else None, + dmd_sigmas=self.scheduler.sigmas, + dmd_timesteps=self.scheduler.timesteps, + all_timesteps=timesteps, + )[0] + + return latents + + def sample_block_noise(self, batch_size, channel, num_frames, height, width, patch_size=(1, 2, 2)): + gamma = self.scheduler.config.gamma + _, ph, pw = patch_size + block_size = ph * pw + + cov = torch.eye(block_size) * (1 + gamma) - torch.ones(block_size, block_size) * gamma + dist = torch.distributions.MultivariateNormal(torch.zeros(block_size, device=cov.device), covariance_matrix=cov) + block_number = batch_size * channel * num_frames * (height // ph) * (width // pw) + + noise = dist.sample((block_number,)) + noise = noise.view(batch_size, channel, num_frames, height // ph, width // pw, ph, pw) + noise = noise.permute(0, 1, 2, 3, 5, 4, 6).reshape(batch_size, channel, num_frames, height, width) + return noise + + def predict_noise(self, **kwargs: Any) -> torch.Tensor: + return self.transformer(**kwargs)[0] + + def encode_prompt( + self, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + max_sequence_length: int = 226, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + device = device or self.device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt_clean = [self._prompt_clean(p) for p in prompt] + batch_size = len(prompt_clean) + + text_inputs = self.tokenizer( + prompt_clean, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + ids, mask = text_inputs.input_ids, text_inputs.attention_mask + seq_lens = mask.gt(0).sum(dim=1).long() + + prompt_embeds = self.text_encoder(ids.to(device), mask.to(device)).last_hidden_state + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0 + ) + + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + negative_prompt_embeds = None + if do_classifier_free_guidance: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + neg_text_inputs = self.tokenizer( + [self._prompt_clean(p) for p in negative_prompt], + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + ids_neg, mask_neg = neg_text_inputs.input_ids, neg_text_inputs.attention_mask + seq_lens_neg = mask_neg.gt(0).sum(dim=1).long() + negative_prompt_embeds = self.text_encoder(ids_neg.to(device), mask_neg.to(device)).last_hidden_state + negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device) + negative_prompt_embeds = [u[:v] for u, v in zip(negative_prompt_embeds, seq_lens_neg)] + negative_prompt_embeds = torch.stack( + [ + torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) + for u in negative_prompt_embeds + ], + dim=0, + ) + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + return prompt_embeds, negative_prompt_embeds + + @staticmethod + def _prompt_clean(text: str) -> str: + return " ".join(text.strip().split()) + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + num_frames: int, + dtype: torch.dtype | None, + device: torch.device | None, + generator: torch.Generator | list[torch.Generator] | None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + shape = ( + batch_size, + num_channels_latents, + num_latent_frames, + int(height) // self.vae_scale_factor_spatial, + int(width) // self.vae_scale_factor_spatial, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError(f"Generator list length {len(generator)} does not match batch size {batch_size}.") + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + def prepare_image_latents( + self, + image: torch.Tensor, + latents_mean: torch.Tensor, + latents_std: torch.Tensor, + num_latent_frames_per_chunk: int, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Encode a single image into VAE latent space for I2V generation. + + Returns (image_latents, fake_image_latents) where fake_image_latents + is the last-frame latent of a repeated-frame video, used to seed the + history buffer for the first denoising chunk. + """ + device = device or self.device + image = image.unsqueeze(2).to(device=device, dtype=self.vae.dtype) + latents = self.vae.encode(image).latent_dist.sample(generator=generator) + latents = (latents - latents_mean) * latents_std + + min_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 + fake_video = image.repeat(1, 1, min_frames, 1, 1) + fake_latents_full = self.vae.encode(fake_video).latent_dist.sample(generator=generator) + fake_latents_full = (fake_latents_full - latents_mean) * latents_std + fake_latents = fake_latents_full[:, :, -1:, :, :] + + return latents.to(device=device, dtype=dtype), fake_latents.to(device=device, dtype=dtype) + + def prepare_video_latents( + self, + video: torch.Tensor, + latents_mean: torch.Tensor, + latents_std: torch.Tensor, + num_latent_frames_per_chunk: int, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Encode a video into VAE latent space for V2V generation. + + Returns (first_frame_latent, video_latents) where first_frame_latent + is used as the image prefix, and video_latents fills the history buffer. + """ + device = device or self.device + video = video.to(device=device, dtype=self.vae.dtype) + + num_frames = video.shape[2] + min_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 + num_chunks = num_frames // min_frames + if num_chunks == 0: + raise ValueError( + f"Video must have at least {min_frames} frames (got {num_frames}). " + f"Required: (num_latent_frames_per_chunk - 1) * {self.vae_scale_factor_temporal} + 1" + ) + total_valid_frames = num_chunks * min_frames + start_frame = num_frames - total_valid_frames + + first_frame = video[:, :, 0:1, :, :] + first_frame_latent = self.vae.encode(first_frame).latent_dist.sample(generator=generator) + first_frame_latent = (first_frame_latent - latents_mean) * latents_std + + latents_chunks = [] + for i in range(num_chunks): + chunk_start = start_frame + i * min_frames + chunk_end = chunk_start + min_frames + video_chunk = video[:, :, chunk_start:chunk_end, :, :] + chunk_latents = self.vae.encode(video_chunk).latent_dist.sample(generator=generator) + chunk_latents = (chunk_latents - latents_mean) * latents_std + latents_chunks.append(chunk_latents) + video_latents = torch.cat(latents_chunks, dim=2) + + return first_frame_latent.to(device=device, dtype=dtype), video_latents.to(device=device, dtype=dtype) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/models/helios/scheduling_helios.py b/vllm_omni/diffusion/models/helios/scheduling_helios.py new file mode 100644 index 000000000..847dea2d0 --- /dev/null +++ b/vllm_omni/diffusion/models/helios/scheduling_helios.py @@ -0,0 +1,800 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from Helios (https://github.com/BestWishYsh/Helios) + +import math +from dataclasses import dataclass +from typing import Literal + +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class HeliosSchedulerOutput(BaseOutput): + prev_sample: torch.FloatTensor + model_outputs: torch.FloatTensor | None = None + last_sample: torch.FloatTensor | None = None + this_order: int | None = None + + +class HeliosScheduler(SchedulerMixin, ConfigMixin): + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + stages: int = 3, + stage_range: list = [0, 1 / 3, 2 / 3, 1], + gamma: float = 1 / 3, + thresholding: bool = False, + prediction_type: str = "flow_prediction", + solver_order: int = 2, + predict_x0: bool = True, + solver_type: str = "bh2", + lower_order_final: bool = True, + disable_corrector: list[int] = [], + solver_p: SchedulerMixin = None, + use_flow_sigmas: bool = True, + scheduler_type: str = "unipc", + use_dynamic_shifting: bool = False, + time_shift_type: Literal["exponential", "linear"] = "linear", + ): + self.timestep_ratios = {} + self.timesteps_per_stage = {} + self.sigmas_per_stage = {} + self.start_sigmas = {} + self.end_sigmas = {} + self.ori_start_sigmas = {} + + self.init_sigmas_for_each_stage() + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + self.gamma = gamma + + if solver_type not in ["bh1", "bh2"]: + if solver_type in ["midpoint", "heun", "logrho"]: + self.register_to_config(solver_type="bh2") + else: + raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}") + + self.predict_x0 = predict_x0 + self.model_outputs = [None] * solver_order + self.timestep_list = [None] * solver_order + self.lower_order_nums = 0 + self.disable_corrector = disable_corrector + self.solver_p = solver_p + self.last_sample = None + self._step_index = None + self._begin_index = None + + def init_sigmas(self): + num_train_timesteps = self.config.num_train_timesteps + shift = self.config.shift + + alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps + 1) + sigmas = 1.0 - alphas + sigmas = np.flip(shift * sigmas / (1 + (shift - 1) * sigmas))[:-1].copy() + sigmas = torch.from_numpy(sigmas) + timesteps = (sigmas * num_train_timesteps).clone() + + self._step_index = None + self._begin_index = None + self.timesteps = timesteps + self.sigmas = sigmas.to("cpu") + + def init_sigmas_for_each_stage(self): + self.init_sigmas() + + stage_distance = [] + stages = self.config.stages + training_steps = self.config.num_train_timesteps + stage_range = self.config.stage_range + + for i_s in range(stages): + start_indice = int(stage_range[i_s] * training_steps) + start_indice = max(start_indice, 0) + end_indice = int(stage_range[i_s + 1] * training_steps) + end_indice = min(end_indice, training_steps) + start_sigma = self.sigmas[start_indice].item() + end_sigma = self.sigmas[end_indice].item() if end_indice < training_steps else 0.0 + self.ori_start_sigmas[i_s] = start_sigma + + if i_s != 0: + ori_sigma = 1 - start_sigma + gamma = self.config.gamma + corrected_sigma = (1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma)) * ori_sigma + start_sigma = 1 - corrected_sigma + + stage_distance.append(start_sigma - end_sigma) + self.start_sigmas[i_s] = start_sigma + self.end_sigmas[i_s] = end_sigma + + tot_distance = sum(stage_distance) + for i_s in range(stages): + if i_s == 0: + start_ratio = 0.0 + else: + start_ratio = sum(stage_distance[:i_s]) / tot_distance + if i_s == stages - 1: + end_ratio = 0.9999999999999999 + else: + end_ratio = sum(stage_distance[: i_s + 1]) / tot_distance + + self.timestep_ratios[i_s] = (start_ratio, end_ratio) + + for i_s in range(stages): + timestep_ratio = self.timestep_ratios[i_s] + timestep_max = min(self.timesteps[int(timestep_ratio[0] * training_steps)], 999) + timestep_min = self.timesteps[min(int(timestep_ratio[1] * training_steps), training_steps - 1)] + timesteps = np.linspace(timestep_max, timestep_min, training_steps + 1) + self.timesteps_per_stage[i_s] = ( + timesteps[:-1] if isinstance(timesteps, torch.Tensor) else torch.from_numpy(timesteps[:-1]) + ) + stage_sigmas = np.linspace(0.999, 0, training_steps + 1) + self.sigmas_per_stage[i_s] = torch.from_numpy(stage_sigmas[:-1]) + + @property + def step_index(self): + return self._step_index + + @property + def begin_index(self): + return self._begin_index + + def set_begin_index(self, begin_index: int = 0): + self._begin_index = begin_index + + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def set_timesteps( + self, + num_inference_steps: int, + stage_index: int | None = None, + device: str | torch.device = None, + sigmas: bool | None = None, + mu: bool | None = None, + is_amplify_first_chunk: bool = False, + ): + if self.config.scheduler_type == "dmd": + if is_amplify_first_chunk: + num_inference_steps = num_inference_steps * 2 + 1 + else: + num_inference_steps = num_inference_steps + 1 + + self.num_inference_steps = num_inference_steps + self.init_sigmas() + + if self.config.stages == 1 or stage_index is None: + if sigmas is None: + sigmas = np.linspace(1, 1 / self.config.num_train_timesteps, num_inference_steps + 1)[:-1].astype( + np.float32 + ) + if self.config.shift != 1.0: + assert not self.config.use_dynamic_shifting + sigmas = self.time_shift(self.config.shift, 1.0, sigmas) + timesteps = (sigmas * self.config.num_train_timesteps).copy() + sigmas = torch.from_numpy(sigmas) + else: + stage_timesteps = self.timesteps_per_stage[stage_index] + timesteps = np.linspace( + stage_timesteps[0].item(), + stage_timesteps[-1].item(), + num_inference_steps, + ) + + stage_sigmas = self.sigmas_per_stage[stage_index] + ratios = np.linspace(stage_sigmas[0].item(), stage_sigmas[-1].item(), num_inference_steps) + sigmas = torch.from_numpy(ratios) + + self.timesteps = torch.from_numpy(timesteps).to(device=device) + self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device) + + self._step_index = None + self.reset_scheduler_history() + + if self.config.scheduler_type == "dmd": + self.timesteps = self.timesteps[:-1] + self.sigmas = torch.cat([self.sigmas[:-2], self.sigmas[-1:]]) + + if self.config.use_dynamic_shifting: + assert self.config.shift == 1.0 + self.sigmas = self.time_shift(mu, 1.0, self.sigmas) + if self.config.stages == 1 or stage_index is None: + self.timesteps = self.sigmas[:-1] * self.config.num_train_timesteps + else: + self.timesteps = self.timesteps_per_stage[stage_index].min() + self.sigmas[:-1] * ( + self.timesteps_per_stage[stage_index].max() - self.timesteps_per_stage[stage_index].min() + ) + + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + if self.config.time_shift_type == "exponential": + return self._time_shift_exponential(mu, sigma, t) + elif self.config.time_shift_type == "linear": + return self._time_shift_linear(mu, sigma, t) + + def _time_shift_exponential(self, mu, sigma, t): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + def _time_shift_linear(self, mu, sigma, t): + return mu / (mu + (1 / t - 1) ** sigma) + + # ---------------------------------- Euler ---------------------------------- + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + pos = 1 if len(indices) > 1 else 0 + return indices[pos].item() + + def _init_step_index(self, timestep): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step_euler( + self, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor = None, + sample: torch.FloatTensor = None, + generator: torch.Generator | None = None, + sigma: torch.FloatTensor | None = None, + sigma_next: torch.FloatTensor | None = None, + return_dict: bool = True, + ) -> HeliosSchedulerOutput | tuple: + assert (sigma is None) == (sigma_next is None) + + if sigma is None and sigma_next is None: + if ( + isinstance(timestep, int) + or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor) + ): + raise ValueError( + "Passing integer indices as timesteps to EulerDiscreteScheduler.step() is not supported. " + "Pass one of scheduler.timesteps as a timestep." + ) + + if self.step_index is None: + self._step_index = 0 + + sample = sample.to(torch.float32) + + if sigma is None and sigma_next is None: + sigma = self.sigmas[self.step_index] + sigma_next = self.sigmas[self.step_index + 1] + + prev_sample = sample + (sigma_next - sigma) * model_output + prev_sample = prev_sample.to(model_output.dtype) + + self._step_index += 1 + + if not return_dict: + return (prev_sample,) + + return HeliosSchedulerOutput(prev_sample=prev_sample) + + # ---------------------------------- UniPC ---------------------------------- + def _sigma_to_alpha_sigma_t(self, sigma): + if self.config.use_flow_sigmas: + alpha_t = 1 - sigma + sigma_t = torch.clamp(sigma, min=1e-8) + else: + alpha_t = 1 / ((sigma**2 + 1) ** 0.5) + sigma_t = sigma * alpha_t + + return alpha_t, sigma_t + + def convert_model_output( + self, + model_output: torch.Tensor, + *args, + sample: torch.Tensor = None, + sigma: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError("missing `sample` as a required keyword argument") + + flag = False + if sigma is None: + flag = True + sigma = self.sigmas[self.step_index] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + + if self.predict_x0: + if self.config.prediction_type == "epsilon": + x0_pred = (sample - sigma_t * model_output) / alpha_t + elif self.config.prediction_type == "sample": + x0_pred = model_output + elif self.config.prediction_type == "v_prediction": + x0_pred = alpha_t * sample - sigma_t * model_output + elif self.config.prediction_type == "flow_prediction": + if flag: + sigma_t = self.sigmas[self.step_index] + else: + sigma_t = sigma + x0_pred = sample - sigma_t * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, " + "`v_prediction`, or `flow_prediction`." + ) + + if self.config.thresholding: + x0_pred = self._threshold_sample(x0_pred) + + return x0_pred + else: + if self.config.prediction_type == "epsilon": + return model_output + elif self.config.prediction_type == "sample": + epsilon = (sample - alpha_t * model_output) / sigma_t + return epsilon + elif self.config.prediction_type == "v_prediction": + epsilon = alpha_t * model_output + sigma_t * sample + return epsilon + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" + " `v_prediction`." + ) + + def multistep_uni_p_bh_update( + self, + model_output: torch.Tensor, + *args, + sample: torch.Tensor = None, + order: int = None, + sigma: torch.Tensor = None, + sigma_next: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError("missing `sample` as a required keyword argument") + if order is None: + if len(args) > 2: + order = args[2] + else: + raise ValueError("missing `order` as a required keyword argument") + + model_output_list = self.model_outputs + + s0 = self.timestep_list[-1] + m0 = model_output_list[-1] + x = sample + + if self.solver_p: + x_t = self.solver_p.step(model_output, s0, x).prev_sample + return x_t + + if sigma_next is None and sigma is None: + sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index] + else: + sigma_t, sigma_s0 = sigma_next, sigma + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + device = sample.device + + rks = [] + D1s = [] + for i in range(1, order): + si = self.step_index - i + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(1.0) + rks = torch.tensor(rks, device=device, dtype=torch.float32) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + hh = hh.float() + h_phi_1 = torch.expm1(hh) + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R) + b = torch.tensor(b, device=device, dtype=torch.float32) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) + if order == 2: + rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device) + else: + rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype) + else: + D1s = None + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) + else: + pred_res = 0 + x_t = x_t_ - alpha_t * B_h * pred_res + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) + else: + pred_res = 0 + x_t = x_t_ - sigma_t * B_h * pred_res + + x_t = x_t.to(x.dtype) + return x_t + + def multistep_uni_c_bh_update( + self, + this_model_output: torch.Tensor, + *args, + last_sample: torch.Tensor = None, + this_sample: torch.Tensor = None, + order: int = None, + sigma_before: torch.Tensor = None, + sigma: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + if last_sample is None: + if len(args) > 1: + last_sample = args[1] + else: + raise ValueError("missing `last_sample` as a required keyword argument") + if this_sample is None: + if len(args) > 2: + this_sample = args[2] + else: + raise ValueError("missing `this_sample` as a required keyword argument") + if order is None: + if len(args) > 3: + order = args[3] + else: + raise ValueError("missing `order` as a required keyword argument") + + model_output_list = self.model_outputs + + m0 = model_output_list[-1] + x = last_sample + x_t = this_sample + model_t = this_model_output + + if sigma_before is None and sigma is None: + sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[self.step_index - 1] + else: + sigma_t, sigma_s0 = sigma, sigma_before + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + device = this_sample.device + + rks = [] + D1s = [] + for i in range(1, order): + si = self.step_index - (i + 1) + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(1.0) + rks = torch.tensor(rks, device=device, dtype=torch.float32) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + hh = hh.float() + h_phi_1 = torch.expm1(hh) + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R) + b = torch.tensor(b, device=device, dtype=torch.float32) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) + else: + D1s = None + + if order == 1: + rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device) + else: + rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype) + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) + x_t = x_t.to(x.dtype) + return x_t + + def step_unipc( + self, + model_output: torch.Tensor, + timestep: int | torch.Tensor = None, + sample: torch.Tensor = None, + return_dict: bool = True, + model_outputs: list = None, + timestep_list: list = None, + sigma_before: torch.Tensor = None, + sigma: torch.Tensor = None, + sigma_next: torch.Tensor = None, + cus_step_index: int = None, + cus_lower_order_num: int = None, + cus_this_order: int = None, + cus_last_sample: torch.Tensor = None, + ) -> HeliosSchedulerOutput | tuple: + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + if cus_step_index is None: + if self.step_index is None: + self._step_index = 0 + else: + self._step_index = cus_step_index + + if cus_lower_order_num is not None: + self.lower_order_nums = cus_lower_order_num + + if cus_this_order is not None: + self.this_order = cus_this_order + + if cus_last_sample is not None: + self.last_sample = cus_last_sample + + use_corrector = ( + self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None + ) + + model_output_convert = self.convert_model_output(model_output, sample=sample, sigma=sigma) + + if model_outputs is not None and timestep_list is not None: + self.model_outputs = model_outputs[:-1] + self.timestep_list = timestep_list[:-1] + + if use_corrector: + sample = self.multistep_uni_c_bh_update( + this_model_output=model_output_convert, + last_sample=self.last_sample, + this_sample=sample, + order=self.this_order, + sigma_before=sigma_before, + sigma=sigma, + ) + + if model_outputs is not None and timestep_list is not None: + model_outputs[-1] = model_output_convert + self.model_outputs = model_outputs[1:] + self.timestep_list = timestep_list[1:] + else: + for i in range(self.config.solver_order - 1): + self.model_outputs[i] = self.model_outputs[i + 1] + self.timestep_list[i] = self.timestep_list[i + 1] + self.model_outputs[-1] = model_output_convert + self.timestep_list[-1] = timestep + + if self.config.lower_order_final: + this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index) + else: + this_order = self.config.solver_order + self.this_order = min(this_order, self.lower_order_nums + 1) + assert self.this_order > 0 + + self.last_sample = sample + prev_sample = self.multistep_uni_p_bh_update( + model_output=model_output, + sample=sample, + order=self.this_order, + sigma=sigma, + sigma_next=sigma_next, + ) + + if cus_lower_order_num is None: + if self.lower_order_nums < self.config.solver_order: + self.lower_order_nums += 1 + + if cus_step_index is None: + self._step_index += 1 + + if not return_dict: + return (prev_sample, model_outputs, self.last_sample, self.this_order) + + return HeliosSchedulerOutput( + prev_sample=prev_sample, + model_outputs=model_outputs, + last_sample=self.last_sample, + this_order=self.this_order, + ) + + # ---------------------------------- DMD ---------------------------------- + def add_noise(self, original_samples, noise, timestep, sigmas, timesteps): + sigmas = sigmas.to(noise.device) + timesteps = timesteps.to(noise.device) + timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1) + sample = (1 - sigma) * original_samples + sigma * noise + return sample.type_as(noise) + + def convert_flow_pred_to_x0(self, flow_pred, xt, timestep, sigmas, timesteps): + original_dtype = flow_pred.dtype + device = flow_pred.device + flow_pred, xt, sigmas, timesteps = (x.double().to(device) for x in (flow_pred, xt, sigmas, timesteps)) + + timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1) + x0_pred = xt - sigma_t * flow_pred + return x0_pred.to(original_dtype) + + def step_dmd( + self, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor = None, + sample: torch.FloatTensor = None, + generator: torch.Generator | None = None, + return_dict: bool = True, + cur_sampling_step: int = 0, + dmd_noisy_tensor: torch.FloatTensor | None = None, + dmd_sigmas: torch.FloatTensor | None = None, + dmd_timesteps: torch.FloatTensor | None = None, + all_timesteps: torch.FloatTensor | None = None, + ): + pred_image_or_video = self.convert_flow_pred_to_x0( + flow_pred=model_output, + xt=sample, + timestep=torch.full((model_output.shape[0],), timestep, dtype=torch.long, device=model_output.device), + sigmas=dmd_sigmas, + timesteps=dmd_timesteps, + ) + if cur_sampling_step < len(all_timesteps) - 1: + prev_sample = self.add_noise( + pred_image_or_video, + dmd_noisy_tensor, + torch.full( + (model_output.shape[0],), + all_timesteps[cur_sampling_step + 1], + dtype=torch.long, + device=model_output.device, + ), + sigmas=dmd_sigmas, + timesteps=dmd_timesteps, + ) + else: + prev_sample = pred_image_or_video + + if not return_dict: + return (prev_sample,) + + return HeliosSchedulerOutput(prev_sample=prev_sample) + + # ---------------------------------- Merge ---------------------------------- + def step( + self, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor = None, + sample: torch.FloatTensor = None, + generator: torch.Generator | None = None, + return_dict: bool = True, + cur_sampling_step: int = 0, + dmd_noisy_tensor: torch.FloatTensor | None = None, + dmd_sigmas: torch.FloatTensor | None = None, + dmd_timesteps: torch.FloatTensor | None = None, + all_timesteps: torch.FloatTensor | None = None, + ) -> HeliosSchedulerOutput | tuple: + if self.config.scheduler_type == "euler": + return self.step_euler( + model_output=model_output, + timestep=timestep, + sample=sample, + generator=generator, + return_dict=return_dict, + ) + elif self.config.scheduler_type == "unipc": + return self.step_unipc( + model_output=model_output, + timestep=timestep, + sample=sample, + return_dict=return_dict, + ) + elif self.config.scheduler_type == "dmd": + return self.step_dmd( + model_output=model_output, + timestep=timestep, + sample=sample, + generator=generator, + return_dict=return_dict, + cur_sampling_step=cur_sampling_step, + dmd_noisy_tensor=dmd_noisy_tensor, + dmd_sigmas=dmd_sigmas, + dmd_timesteps=dmd_timesteps, + all_timesteps=all_timesteps, + ) + else: + raise NotImplementedError + + def reset_scheduler_history(self): + self.model_outputs = [None] * self.config.solver_order + self.timestep_list = [None] * self.config.solver_order + self.lower_order_nums = 0 + self.disable_corrector = self.config.disable_corrector + self.solver_p = self.config.solver_p + self.last_sample = None + self._step_index = None + self._begin_index = None + + def __len__(self): + return self.config.num_train_timesteps diff --git a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py index 329fc648b..4334ef62e 100644 --- a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py +++ b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py @@ -509,10 +509,7 @@ class LongCatImageTransformer2DModel(nn.Module): Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig. """ - packed_modules_mapping = { - "to_qkv": ["to_q", "to_k", "to_v"], - "add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"], - } + _repeated_blocks = ["LongCatImageTransformerBlock", "LongCatImageSingleTransformerBlock"] def __init__( self, @@ -714,6 +711,8 @@ class LongCatImageTransformer2DModel(nn.Module): (".add_kv_proj", ".add_k_proj", "k"), (".add_kv_proj", ".add_v_proj", "v"), ] + # Expose packed shard mappings for LoRA handling of fused projections. + self.stacked_params_mapping = stacked_params_mapping params_dict = dict(self.named_parameters()) diff --git a/vllm_omni/diffusion/models/ovis_image/ovis_image_transformer.py b/vllm_omni/diffusion/models/ovis_image/ovis_image_transformer.py index ae6cf6b0c..bd2a3b483 100644 --- a/vllm_omni/diffusion/models/ovis_image/ovis_image_transformer.py +++ b/vllm_omni/diffusion/models/ovis_image/ovis_image_transformer.py @@ -366,10 +366,6 @@ class OvisImageTransformer2DModel(nn.Module): """ _repeated_blocks = ["OvisImageTransformerBlock", "OvisImageSingleTransformerBlock"] - packed_modules_mapping = { - "to_qkv": ["to_q", "to_k", "to_v"], - "add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"], - } def __init__( self, @@ -518,6 +514,8 @@ class OvisImageTransformer2DModel(nn.Module): (".add_kv_proj", ".add_k_proj", "k"), (".add_kv_proj", ".add_v_proj", "v"), ] + # Expose packed shard mappings for LoRA handling of fused projections. + self.stacked_params_mapping = stacked_params_mapping params_dict = dict(self.named_parameters()) diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index 2d8d49eee..8da6ffdb2 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -1082,6 +1082,8 @@ class QwenImageTransformer2DModel(CachedTransformer): (".add_kv_proj", ".add_k_proj", "k"), (".add_kv_proj", ".add_v_proj", "v"), ] + # Expose packed shard mappings for LoRA handling of fused projections. + self.stacked_params_mapping = stacked_params_mapping params_dict = dict(self.named_parameters()) diff --git a/vllm_omni/diffusion/models/sd3/sd3_transformer.py b/vllm_omni/diffusion/models/sd3/sd3_transformer.py index 071716657..c4a04febb 100644 --- a/vllm_omni/diffusion/models/sd3/sd3_transformer.py +++ b/vllm_omni/diffusion/models/sd3/sd3_transformer.py @@ -381,10 +381,6 @@ class SD3Transformer2DModel(nn.Module): """ _repeated_blocks = ["SD3TransformerBlock"] - packed_modules_mapping = { - "to_qkv": ["to_q", "to_k", "to_v"], - "add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"], - } def __init__( self, @@ -518,6 +514,8 @@ class SD3Transformer2DModel(nn.Module): (".add_kv_proj", ".add_k_proj", "k"), (".add_kv_proj", ".add_v_proj", "v"), ] + # Expose packed shard mappings for LoRA handling of fused projections. + self.stacked_params_mapping = stacked_params_mapping params_dict = dict(self.named_parameters()) diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index 231a001f7..2ff604159 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -965,13 +965,17 @@ class WanTransformer3DModel(nn.Module): """ tp_rank = get_tensor_model_parallel_rank() tp_size = get_tensor_model_parallel_world_size() - # Stacked params mapping for self-attention QKV fusion + # Stacked params mapping for self-attention QKV fusion. + # Format: (param_name, shard_name, shard_id) + # Note: Only fuse attn1 (self-attention), NOT attn2 (cross-attention) stacked_params_mapping = [ # self-attention QKV fusion (".attn1.to_qkv", ".attn1.to_q", "q"), (".attn1.to_qkv", ".attn1.to_k", "k"), (".attn1.to_qkv", ".attn1.to_v", "v"), ] + # Expose packed shard mappings for LoRA handling of fused projections. + self.stacked_params_mapping = stacked_params_mapping # Remap scale_shift_table to new module location weight_name_remapping = { diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index 83d0ecb9b..3e61610b1 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -579,10 +579,6 @@ class ZImageTransformer2DModel(CachedTransformer): """ _repeated_blocks = ["ZImageTransformerBlock"] - packed_modules_mapping = { - "to_qkv": ["to_q", "to_k", "to_v"], - "w13": ["w1", "w3"], - } # Sequence Parallelism for Z-Image (following diffusers' _cp_plan pattern) # Similar to how Wan uses `rope` module's split_output to shard rotary embeddings, @@ -985,6 +981,8 @@ class ZImageTransformer2DModel(CachedTransformer): (".w13", ".w1", 0), (".w13", ".w3", 1), ] + # Expose packed shard mappings for LoRA handling of fused projections. + self.stacked_params_mapping = stacked_params_mapping params_dict = dict(self.named_parameters()) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index c806dbae9..37bc316cc 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -111,6 +111,16 @@ _DIFFUSION_MODELS = { "pipeline_omnigen2", "OmniGen2Pipeline", ), + "HeliosPipeline": ( + "helios", + "pipeline_helios", + "HeliosPipeline", + ), + "HeliosPyramidPipeline": ( + "helios", + "pipeline_helios", + "HeliosPipeline", + ), } @@ -275,6 +285,8 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "NextStep11Pipeline": "get_nextstep11_post_process_func", "FluxPipeline": "get_flux_post_process_func", "OmniGen2Pipeline": "get_omnigen2_post_process_func", + "HeliosPipeline": "get_helios_post_process_func", + "HeliosPyramidPipeline": "get_helios_post_process_func", } _DIFFUSION_PRE_PROCESS_FUNCS = { @@ -289,6 +301,8 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { "WanPipeline": "get_wan22_pre_process_func", "WanImageToVideoPipeline": "get_wan22_i2v_pre_process_func", "OmniGen2Pipeline": "get_omnigen2_pre_process_func", + "HeliosPipeline": "get_helios_pre_process_func", + "HeliosPyramidPipeline": "get_helios_pre_process_func", } diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index b50e50d1e..16a91822c 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -120,6 +120,9 @@ class OmniRequestState(RequestState): # When the audio tensor shape is inconsistent, torch.cat will fail. # We need to use torch.cat in -1 dimension. continue + elif k == "sr": + # Sample rate is a constant scalar, keep last value. + self.mm_accumulated[k] = v[-1] else: self.mm_accumulated[k] = torch.cat(v, dim=0) except Exception: diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 196cba0dd..986e06ace 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -235,6 +235,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _validate_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate TTS request parameters. Returns error message or None.""" + # Infer Base task when ref_audio or ref_text is provided without explicit task_type. + if request.task_type is None and (request.ref_audio is not None or request.ref_text is not None): + request.task_type = "Base" task_type = request.task_type or "CustomVoice" # Normalize voice to lowercase for case-insensitive matching @@ -250,8 +253,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return f"Invalid language '{request.language}'. Supported: {', '.join(sorted(_TTS_LANGUAGES))}" # Validate speaker for CustomVoice task - if task_type == "CustomVoice" and request.voice is not None: - if self.supported_speakers and request.voice not in self.supported_speakers: + if task_type == "CustomVoice": + if not self.supported_speakers: + return ( + "This model does not support CustomVoice task (no speakers configured). " + "Use task_type='Base' with ref_audio/ref_text for voice cloning, " + "or use a CustomVoice model." + ) + if request.voice is not None and request.voice not in self.supported_speakers: return f"Invalid speaker '{request.voice}'. Supported: {', '.join(sorted(self.supported_speakers))}" # Validate Base task requirements @@ -261,6 +270,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Validate ref_audio format if not (request.ref_audio.startswith(("http://", "https://")) or request.ref_audio.startswith("data:")): return "ref_audio must be a URL (http/https) or base64 data URL (data:...)" + # In-context voice cloning (default) requires non-empty ref_text. + # x_vector_only_mode skips in-context and only uses speaker embedding. + if not request.x_vector_only_mode: + if not request.ref_text or not request.ref_text.strip(): + return ( + "Base task requires non-empty 'ref_text' (transcript of " + "the reference audio) unless 'x_vector_only_mode' is enabled" + ) # Validate cross-parameter dependencies if task_type != "Base": diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py index 28469a8d1..a712a2180 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py @@ -396,7 +396,8 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): bsz = int(inputs_embeds.shape[0]) qlen = 2 # Flatten to token-major. - hs = inputs_embeds.to(dtype=torch.bfloat16).reshape(bsz * qlen, -1) + model_dtype = next(self.model.parameters()).dtype + hs = inputs_embeds.to(dtype=model_dtype).reshape(bsz * qlen, -1) hs = self.small_to_mtp_projection(hs) query_lens = torch.full((bsz,), qlen, dtype=torch.int32) @@ -497,11 +498,15 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): for step in range(1, num_groups): # Sample or argmax from logits. if do_sample and temperature > 0: - scaled = logits / temperature + scaled = logits.float() / temperature if top_k > 0: topk_vals, _ = scaled.topk(top_k, dim=-1) scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) probs = torch.softmax(scaled, dim=-1) + # Guard against NaN/Inf from lower-precision dtypes. + probs = probs.clamp(min=0.0) + row_sums = probs.sum(dim=-1, keepdim=True) + probs = torch.where(row_sums > 0, probs / row_sums, torch.ones_like(probs) / probs.shape[-1]) next_ids = torch.multinomial(probs, num_samples=1) # [B, 1] else: next_ids = logits.argmax(dim=-1, keepdim=True) # [B, 1] diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index 1bdcd566e..e80f1785f 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -638,12 +638,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): def _get_tokenizer(self): if self._tokenizer is None: - self._tokenizer = AutoTokenizer.from_pretrained( - self.model_path, - trust_remote_code=True, - fix_mistral_regex=True, - use_fast=True, - ) + import transformers + + kwargs = dict(trust_remote_code=True, use_fast=True) + if transformers.__version__ < "5": + kwargs["fix_mistral_regex"] = True + self._tokenizer = AutoTokenizer.from_pretrained(self.model_path, **kwargs) self._tokenizer.padding_side = "left" return self._tokenizer diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index bbec6a525..5719d1024 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py @@ -48,6 +48,22 @@ from .configuration_qwen3_tts_tokenizer_v2 import ( logger = logging.get_logger(__name__) +def _default_rope_init(config, device=None, seq_len=None, layer_type=None): + """Vanilla sinusoidal RoPE (no scaling). + + transformers>=5.0 removed the 'default' entry from + ``ROPE_INIT_FUNCTIONS`` (see ``modeling_rope_utils.py`` in + huggingface/transformers), but the speech tokenizer config still + declares ``rope_type="default"``. This reimplements the original + behaviour. + """ + head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + inv_freq = 1.0 / ( + config.rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim) + ) + return inv_freq, 1.0 + + @dataclass @auto_docstring class Qwen3TTSTokenizerV2EncoderOutput(ModelOutput): @@ -256,7 +272,15 @@ class Qwen3TTSTokenizerV2DecoderRotatoryEmbedding(nn.Module): self.original_max_seq_len = config.max_position_embeddings self.config = config - self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + if self.rope_type in ROPE_INIT_FUNCTIONS: + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + elif self.rope_type == "default": + self.rope_init_fn = _default_rope_init + else: + raise ValueError( + f"Unsupported rope_type '{self.rope_type}'. Expected one of {list(ROPE_INIT_FUNCTIONS)} or 'default'." + ) inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 7734590d7..569ee60f9 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -7,7 +7,6 @@ from vllm.config import CUDAGraphMode from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import set_forward_context from vllm.logger import init_logger -from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding from vllm.model_executor.models.interfaces import supports_mrope from vllm.model_executor.models.interfaces_base import VllmModelForPooling from vllm.sampling_params import SamplingType @@ -20,6 +19,7 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices +from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding from vllm_omni.model_executor.models.output_templates import OmniOutput if TYPE_CHECKING: