Merge remote-tracking branch 'remotes/origin/main' into dev/rebase-v0170
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
This commit is contained in:
tzhouam committed
commit
36a507e071
29 files changed
+3554
-63
No files matched your search
Binary file not shown.
|
Before Width: | Height: | Size: 127 KiB After Width: | Height: | Size: 125 KiB |
@@ -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` |
|
||||
|
||||
@@ -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 |
|
||||
@@ -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()
|
||||
@@ -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"),
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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())
|
||||
|
||||
|
||||
@@ -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())
|
||||
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
File diff suppressed because it is too large.
Load diff
@@ -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
|
||||
@@ -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())
|
||||
|
||||
|
||||
@@ -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())
|
||||
|
||||
|
||||
@@ -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())
|
||||
|
||||
|
||||
@@ -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())
|
||||
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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())
|
||||
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
+25
-1
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in new issue
Block a user