[Model]: Helios support (#1604)

Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: gcanlin <canlinguosdu@gmail.com>

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 |

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()