430 lines
17 KiB
Python
430 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import argparse
|
|
import os
|
|
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 parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Generate a video with Wan2.2 T2V.")
|
|
parser.add_argument(
|
|
"--model",
|
|
default="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
|
help="Diffusers Wan2.2 model ID or local path.",
|
|
)
|
|
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
|
|
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
|
|
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
|
parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (applied to low/high).")
|
|
parser.add_argument("--guidance-scale-high", type=float, default=None, help="Optional separate CFG for high-noise.")
|
|
parser.add_argument("--height", type=int, default=720, help="Video height.")
|
|
parser.add_argument("--width", type=int, default=1280, help="Video width.")
|
|
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames (Wan default is 81).")
|
|
parser.add_argument(
|
|
"--frame-rate",
|
|
type=float,
|
|
default=None,
|
|
help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
|
|
)
|
|
parser.add_argument("--num-inference-steps", type=int, default=40, help="Sampling steps.")
|
|
parser.add_argument(
|
|
"--boundary-ratio",
|
|
type=float,
|
|
default=0.875,
|
|
help="Boundary split ratio for low/high DiT. Default 0.875 uses both transformers for best quality. Set to 1.0 to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited).",
|
|
)
|
|
parser.add_argument(
|
|
"--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
|
|
)
|
|
parser.add_argument(
|
|
"--cache-backend",
|
|
type=str,
|
|
default=None,
|
|
choices=["cache_dit"],
|
|
help=(
|
|
"Cache backend to use for acceleration. "
|
|
"Options: 'cache_dit' (DBCache + SCM + TaylorSeer). "
|
|
"Default: None (no cache acceleration)."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--enable-cache-dit-summary",
|
|
action="store_true",
|
|
help="Enable cache-dit summary logging after diffusion forward passes.",
|
|
)
|
|
parser.add_argument("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).")
|
|
parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.")
|
|
parser.add_argument(
|
|
"--vae-use-slicing",
|
|
action="store_true",
|
|
help="Enable VAE slicing for memory optimization.",
|
|
)
|
|
parser.add_argument(
|
|
"--vae-use-tiling",
|
|
action="store_true",
|
|
help="Enable VAE tiling for memory optimization.",
|
|
)
|
|
parser.add_argument(
|
|
"--enforce-eager",
|
|
action="store_true",
|
|
help="Disable torch.compile and force eager execution.",
|
|
)
|
|
parser.add_argument(
|
|
"--enable-cpu-offload",
|
|
action="store_true",
|
|
help="Enable CPU offloading for diffusion models.",
|
|
)
|
|
parser.add_argument(
|
|
"--enable-layerwise-offload",
|
|
action="store_true",
|
|
help="Enable layerwise (blockwise) offloading on DiT modules.",
|
|
)
|
|
parser.add_argument(
|
|
"--ulysses-degree",
|
|
type=int,
|
|
default=1,
|
|
help="Number of GPUs used for ulysses sequence parallelism.",
|
|
)
|
|
parser.add_argument(
|
|
"--ring-degree",
|
|
type=int,
|
|
default=1,
|
|
help="Number of GPUs used for ring sequence parallelism.",
|
|
)
|
|
parser.add_argument(
|
|
"--cfg-parallel-size",
|
|
type=int,
|
|
default=1,
|
|
choices=[1, 2],
|
|
help="Number of GPUs used for classifier free guidance parallel size.",
|
|
)
|
|
parser.add_argument(
|
|
"--tensor-parallel-size",
|
|
type=int,
|
|
default=1,
|
|
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
|
|
)
|
|
parser.add_argument(
|
|
"--audio-sample-rate",
|
|
type=int,
|
|
default=24000,
|
|
help="Sample rate for audio output when saved (default: 24000 for LTX2).",
|
|
)
|
|
parser.add_argument(
|
|
"--vae-patch-parallel-size",
|
|
type=int,
|
|
default=1,
|
|
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
|
|
)
|
|
parser.add_argument(
|
|
"--enable-expert-parallel",
|
|
action="store_true",
|
|
help="Enable expert parallelism for MoE layers.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
|
frame_rate = args.frame_rate if args.frame_rate is not None else float(args.fps)
|
|
|
|
# Wan2.2 cache-dit tuning (from cache-dit examples and cache_alignment).
|
|
cache_config = None
|
|
if args.cache_backend == "cache_dit":
|
|
cache_config = {
|
|
# DBCache parameters [cache-dit only]
|
|
"Fn_compute_blocks": 1, # Optimized for single-transformer models
|
|
"Bn_compute_blocks": 0, # Number of backward compute blocks
|
|
"max_warmup_steps": 4, # Maximum warmup steps (works for few-step models)
|
|
"max_cached_steps": 20,
|
|
"residual_diff_threshold": 0.24, # Higher threshold for more aggressive caching
|
|
"max_continuous_cached_steps": 3, # Limit to prevent precision degradation
|
|
# TaylorSeer parameters [cache-dit only]
|
|
"enable_taylorseer": False, # Disabled by default (not suitable for few-step models)
|
|
"taylorseer_order": 1, # TaylorSeer polynomial order
|
|
# SCM (Step Computation Masking) parameters [cache-dit only]
|
|
"scm_steps_mask_policy": None, # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
|
|
"scm_steps_policy": "dynamic", # SCM steps policy: "dynamic" or "static"
|
|
}
|
|
# Configure parallel settings
|
|
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,
|
|
vae_patch_parallel_size=args.vae_patch_parallel_size,
|
|
enable_expert_parallel=args.enable_expert_parallel,
|
|
)
|
|
|
|
# Check if profiling is requested via environment variable
|
|
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
|
|
|
|
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,
|
|
boundary_ratio=args.boundary_ratio,
|
|
flow_shift=args.flow_shift,
|
|
enable_cache_dit_summary=args.enable_cache_dit_summary,
|
|
enable_cpu_offload=args.enable_cpu_offload,
|
|
parallel_config=parallel_config,
|
|
enforce_eager=args.enforce_eager,
|
|
cache_backend=args.cache_backend,
|
|
cache_config=cache_config,
|
|
)
|
|
|
|
if profiler_enabled:
|
|
print("[Profiler] Starting profiling...")
|
|
omni.start_profile()
|
|
|
|
# Print generation configuration
|
|
print(f"\n{'=' * 60}")
|
|
print("Generation Configuration:")
|
|
print(f" Model: {args.model}")
|
|
print(f" Inference steps: {args.num_inference_steps}")
|
|
print(f" Frames: {args.num_frames}")
|
|
print(
|
|
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree},"
|
|
f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size},"
|
|
f" vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}"
|
|
)
|
|
print(f" Video size: {args.width}x{args.height}")
|
|
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,
|
|
guidance_scale_2=args.guidance_scale_high,
|
|
num_inference_steps=args.num_inference_steps,
|
|
num_frames=args.num_frames,
|
|
frame_rate=frame_rate,
|
|
),
|
|
)
|
|
generation_end = time.perf_counter()
|
|
generation_time = generation_end - generation_start
|
|
|
|
# Print profiling results
|
|
print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
|
|
|
|
audio = None
|
|
if isinstance(frames, list):
|
|
frames = frames[0] if frames else None
|
|
|
|
if isinstance(frames, OmniRequestOutput):
|
|
if frames.final_output_type != "image":
|
|
raise ValueError(
|
|
f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
|
|
)
|
|
if frames.multimodal_output and "audio" in frames.multimodal_output:
|
|
audio = frames.multimodal_output["audio"]
|
|
if frames.is_pipeline_output and frames.request_output is not None:
|
|
inner_output = frames.request_output
|
|
if isinstance(inner_output, list):
|
|
inner_output = inner_output[0] if inner_output else None
|
|
if isinstance(inner_output, OmniRequestOutput):
|
|
if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
|
|
audio = inner_output.multimodal_output["audio"]
|
|
frames = inner_output
|
|
if isinstance(frames, OmniRequestOutput):
|
|
if frames.images:
|
|
if len(frames.images) == 1 and isinstance(frames.images[0], tuple) and len(frames.images[0]) == 2:
|
|
frames, audio = frames.images[0]
|
|
elif len(frames.images) == 1 and isinstance(frames.images[0], dict):
|
|
audio = frames.images[0].get("audio")
|
|
frames = frames.images[0].get("frames") or frames.images[0].get("video")
|
|
else:
|
|
frames = frames.images
|
|
else:
|
|
raise ValueError("No video frames found in OmniRequestOutput.")
|
|
|
|
if isinstance(frames, list) and frames:
|
|
first_item = frames[0]
|
|
if isinstance(first_item, tuple) and len(first_item) == 2:
|
|
frames, audio = first_item
|
|
elif isinstance(first_item, dict):
|
|
audio = first_item.get("audio")
|
|
frames = first_item.get("frames") or first_item.get("video")
|
|
elif isinstance(first_item, list):
|
|
frames = first_item
|
|
|
|
if isinstance(frames, tuple) and len(frames) == 2:
|
|
frames, audio = frames
|
|
elif isinstance(frames, dict):
|
|
audio = frames.get("audio")
|
|
frames = frames.get("frames") or frames.get("video")
|
|
|
|
if frames is None:
|
|
raise ValueError("No video frames found in output.")
|
|
|
|
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.")
|
|
|
|
def _normalize_frame(frame):
|
|
if isinstance(frame, torch.Tensor):
|
|
frame_tensor = frame.detach().cpu()
|
|
if frame_tensor.dim() == 4 and frame_tensor.shape[0] == 1:
|
|
frame_tensor = frame_tensor[0]
|
|
if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4):
|
|
frame_tensor = frame_tensor.permute(1, 2, 0)
|
|
if frame_tensor.is_floating_point():
|
|
frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5
|
|
return frame_tensor.float().numpy()
|
|
if isinstance(frame, np.ndarray):
|
|
frame_array = frame
|
|
if frame_array.ndim == 4 and frame_array.shape[0] == 1:
|
|
frame_array = frame_array[0]
|
|
if np.issubdtype(frame_array.dtype, np.integer):
|
|
frame_array = frame_array.astype(np.float32) / 255.0
|
|
return frame_array
|
|
try:
|
|
from PIL import Image
|
|
except ImportError:
|
|
Image = None
|
|
if Image is not None and isinstance(frame, Image.Image):
|
|
return np.asarray(frame).astype(np.float32) / 255.0
|
|
return frame
|
|
|
|
def _ensure_frame_list(video_array):
|
|
if isinstance(video_array, list):
|
|
if len(video_array) == 0:
|
|
return video_array
|
|
first_item = video_array[0]
|
|
if isinstance(first_item, np.ndarray):
|
|
if first_item.ndim == 5:
|
|
return list(first_item[0])
|
|
if first_item.ndim == 4:
|
|
if len(video_array) == 1:
|
|
return list(first_item)
|
|
return list(first_item)
|
|
if first_item.ndim == 3:
|
|
return video_array
|
|
return video_array
|
|
if isinstance(video_array, np.ndarray):
|
|
if video_array.ndim == 5:
|
|
return list(video_array[0])
|
|
if video_array.ndim == 4:
|
|
return list(video_array)
|
|
if video_array.ndim == 3:
|
|
return [video_array]
|
|
return video_array
|
|
|
|
# frames may be np.ndarray, torch.Tensor, or list of tensors/arrays/images
|
|
# export_to_video expects a list of frames with values in [0, 1]
|
|
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
|
|
if video_array.ndim == 5:
|
|
video_array = video_array[0]
|
|
if np.issubdtype(video_array.dtype, np.integer):
|
|
video_array = video_array.astype(np.float32) / 255.0
|
|
elif isinstance(frames, list):
|
|
if len(frames) == 0:
|
|
raise ValueError("No video frames found in output.")
|
|
video_array = [_normalize_frame(frame) for frame in frames]
|
|
else:
|
|
video_array = frames
|
|
|
|
video_array = _ensure_frame_list(video_array)
|
|
|
|
use_ltx2_export = False
|
|
if args.model and "ltx" in str(args.model).lower():
|
|
use_ltx2_export = True
|
|
if audio is not None:
|
|
use_ltx2_export = True
|
|
|
|
if use_ltx2_export:
|
|
try:
|
|
from diffusers.pipelines.ltx2.export_utils import encode_video
|
|
except ImportError:
|
|
raise ImportError("diffusers is required for LTX2 encode_video.")
|
|
|
|
if isinstance(video_array, list):
|
|
frames_np = np.stack(video_array, axis=0)
|
|
elif isinstance(video_array, np.ndarray):
|
|
frames_np = video_array
|
|
else:
|
|
frames_np = np.asarray(video_array)
|
|
|
|
frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
|
|
video_tensor = torch.from_numpy(frames_u8)
|
|
|
|
audio_out = None
|
|
if audio is not None:
|
|
if isinstance(audio, list):
|
|
audio = audio[0] if audio else None
|
|
if isinstance(audio, np.ndarray):
|
|
audio = torch.from_numpy(audio)
|
|
if isinstance(audio, torch.Tensor):
|
|
audio_out = audio
|
|
if audio_out.dim() > 1:
|
|
audio_out = audio_out[0]
|
|
audio_out = audio_out.float().cpu()
|
|
|
|
encode_video(
|
|
video_tensor,
|
|
fps=args.fps,
|
|
audio=audio_out,
|
|
audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
|
|
output_path=str(output_path),
|
|
)
|
|
else:
|
|
export_to_video(video_array, str(output_path), fps=args.fps)
|
|
print(f"Saved generated video to {output_path}")
|
|
|
|
if profiler_enabled:
|
|
print("\n[Profiler] Stopping profiler and collecting results...")
|
|
profile_results = omni.stop_profile()
|
|
if profile_results and isinstance(profile_results, dict):
|
|
traces = profile_results.get("traces", [])
|
|
print("\n" + "=" * 60)
|
|
print("PROFILING RESULTS:")
|
|
for rank, trace in enumerate(traces):
|
|
print(f"\nRank {rank}:")
|
|
if trace:
|
|
print(f" • Trace: {trace}")
|
|
if not traces:
|
|
print(" No traces collected.")
|
|
print("=" * 60)
|
|
else:
|
|
print("[Profiler] No valid profiling data returned.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|