790 lines
30 KiB
Python
790 lines
30 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import argparse
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
|
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
|
from vllm_omni.entrypoints.omni import Omni
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
from vllm_omni.lora.request import LoRARequest
|
|
from vllm_omni.lora.utils import stable_lora_int_id
|
|
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name, get_output_tensor_range
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
_MODEL_PRESETS = {
|
|
"vace": {
|
|
"height": 480,
|
|
"width": 832,
|
|
"num_frames": 81,
|
|
"num_inference_steps": 30,
|
|
"guidance_scale": 5.0,
|
|
"fps": 16,
|
|
"output": "vace_t2v_output.mp4",
|
|
},
|
|
"wan": {
|
|
"height": 720,
|
|
"width": 1280,
|
|
"num_frames": 81,
|
|
"num_inference_steps": 40,
|
|
"guidance_scale": 4.0,
|
|
"fps": 24,
|
|
"output": "wan22_output.mp4",
|
|
},
|
|
"hunyuan": {
|
|
"height": 480,
|
|
"width": 832,
|
|
"num_frames": 121,
|
|
"num_inference_steps": 50,
|
|
"guidance_scale": 6.0,
|
|
"fps": 24,
|
|
"output": "hunyuan_video_15_output.mp4",
|
|
},
|
|
"cosmos": {
|
|
"height": 720,
|
|
"width": 1280,
|
|
"num_frames": 189,
|
|
"num_inference_steps": 35,
|
|
"guidance_scale": 6.0,
|
|
"fps": 24,
|
|
"flow_shift": 10.0,
|
|
"output": "cosmos3_t2v_output.mp4",
|
|
},
|
|
# Cosmos3-Edge is a compact checkpoint with its own native resolution / guidance /
|
|
# flow_shift. It must NOT inherit the Nano/Super "cosmos" numbers above (720p / gs 6.0 /
|
|
# flow_shift 10.0), which produce degenerate output on Edge.
|
|
"cosmos3_edge": {
|
|
"height": 480,
|
|
"width": 832,
|
|
"num_frames": 189,
|
|
"num_inference_steps": 35,
|
|
"guidance_scale": 5.0,
|
|
"fps": 24,
|
|
"flow_shift": 3.0,
|
|
"output": "cosmos3_edge_t2v_output.mp4",
|
|
},
|
|
"helios": {
|
|
"height": 384,
|
|
"width": 640,
|
|
"num_frames": 99,
|
|
"num_inference_steps": 50,
|
|
"guidance_scale": 5.0,
|
|
"fps": 16,
|
|
"output": "helios_output.mp4",
|
|
},
|
|
"lingbot": {
|
|
"model_class_name": "LingBotVideoPipeline",
|
|
"height": 192,
|
|
"width": 320,
|
|
"num_frames": 9,
|
|
"num_inference_steps": 2,
|
|
"guidance_scale": 3.0,
|
|
"fps": 24,
|
|
"flow_shift": 3.0,
|
|
"output": "lingbot_video_output.mp4",
|
|
},
|
|
"ltx2": {
|
|
"height": 512,
|
|
"width": 768,
|
|
"num_frames": 121,
|
|
"num_inference_steps": 40,
|
|
"fps": 24,
|
|
"output": "ltx2_output.mp4",
|
|
},
|
|
"ltx2_distilled": {
|
|
"height": 1024,
|
|
"width": 1536,
|
|
"num_frames": 121,
|
|
"num_inference_steps": 8,
|
|
"fps": 24,
|
|
"output": "ltx2_distilled_output.mp4",
|
|
},
|
|
"ltx23": {
|
|
"height": 512,
|
|
"width": 768,
|
|
"num_frames": 121,
|
|
"num_inference_steps": 30,
|
|
"fps": 24,
|
|
"output": "ltx23_output.mp4",
|
|
},
|
|
}
|
|
|
|
|
|
def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
|
|
model_lower = model.lower()
|
|
class_lower = (model_class_name or "").lower()
|
|
if "lingbot" in model_lower or "lingbotvideo" in class_lower:
|
|
return _MODEL_PRESETS["lingbot"]
|
|
if "ltx" in class_lower or "ltx" in model_lower:
|
|
if "distilled" in class_lower or "distilled" in model_lower:
|
|
return _MODEL_PRESETS["ltx2_distilled"]
|
|
if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
|
|
return _MODEL_PRESETS["ltx23"]
|
|
return _MODEL_PRESETS["ltx2"]
|
|
if "vace" in model_lower or "vace" in class_lower:
|
|
return _MODEL_PRESETS["vace"]
|
|
# Edge must be matched before the generic cosmos branch (its "cosmos" substring would
|
|
# otherwise pick up the Nano/Super 720p / gs 6.0 / flow_shift 10.0 preset).
|
|
if ("cosmos" in model_lower or "cosmos" in class_lower) and ("edge" in model_lower or "edge" in class_lower):
|
|
return _MODEL_PRESETS["cosmos3_edge"]
|
|
if "cosmos" in model_lower or "cosmos" in class_lower:
|
|
return _MODEL_PRESETS["cosmos"]
|
|
if "hunyuan" in model_lower or "hunyuan" in class_lower:
|
|
return _MODEL_PRESETS["hunyuan"]
|
|
if "helios" in model_lower or "helios" in class_lower:
|
|
return _MODEL_PRESETS["helios"]
|
|
return _MODEL_PRESETS["wan"]
|
|
|
|
|
|
def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]:
|
|
"""Build the canonical request envelope for the shared T2V example."""
|
|
result: dict[str, Any] = {
|
|
"prompt": prompt,
|
|
"modalities": ["video"],
|
|
}
|
|
if negative_prompt is not None:
|
|
result["negative_prompt"] = negative_prompt
|
|
return result
|
|
|
|
|
|
def _normalize_float_tensor(tensor: torch.Tensor, source_range: str) -> torch.Tensor:
|
|
"""Normalize decoded frames according to the pipeline's output contract."""
|
|
if not tensor.is_floating_point():
|
|
return tensor
|
|
if source_range == "negative_one_to_one":
|
|
return tensor.clamp(-1, 1) * 0.5 + 0.5
|
|
if source_range == "zero_to_one":
|
|
return tensor.clamp(0, 1)
|
|
raise ValueError(f"Unsupported floating-point tensor range: {source_range!r}")
|
|
|
|
|
|
def parse_profiler_config(value: str) -> dict[str, Any]:
|
|
try:
|
|
config = json.loads(value)
|
|
except json.JSONDecodeError as e:
|
|
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
|
|
if not isinstance(config, dict):
|
|
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
|
|
return config
|
|
|
|
|
|
def parse_extra_body(value: str) -> dict[str, Any]:
|
|
"""Parse a JSON object of model-specific extra_body params.
|
|
|
|
Pipeline-declared knobs (see vllm_omni/model_extras/) are merged into
|
|
OmniDiffusionSamplingParams.extra_args, so a single generic example can
|
|
drive model-specific behaviour without bespoke per-model flags.
|
|
"""
|
|
try:
|
|
body = json.loads(value)
|
|
except json.JSONDecodeError as e:
|
|
raise argparse.ArgumentTypeError(f"--extra-body must be valid JSON: {e}") from e
|
|
if not isinstance(body, dict):
|
|
raise argparse.ArgumentTypeError("--extra-body must be a JSON object")
|
|
return body
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate a video from a text prompt. "
|
|
"Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, and other text-to-video models."
|
|
)
|
|
parser.add_argument(
|
|
"--model",
|
|
default="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
|
help="Diffusers model ID or local path. "
|
|
"Examples: Wan-AI/Wan2.2-T2V-A14B-Diffusers, "
|
|
"hunyuanvideo-community/HunyuanVideo-1.5-480p_t2v",
|
|
)
|
|
parser.add_argument(
|
|
"--model-class-name",
|
|
default=None,
|
|
help="Override model class name (e.g., LTX2Pipeline).",
|
|
)
|
|
parser.add_argument(
|
|
"--deploy-config",
|
|
default=None,
|
|
help="Optional deploy config YAML to use for pipeline-backed runs.",
|
|
)
|
|
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
|
|
parser.add_argument("--negative-prompt", default=None, help="Negative prompt. Default: model-specific.")
|
|
parser.add_argument(
|
|
"--extra-body",
|
|
type=parse_extra_body,
|
|
default=None,
|
|
help="JSON dict of model-specific extra_body params (declared in vllm_omni/model_extras/), "
|
|
"merged into sampling extra_args. Unknown keys for the chosen model are dropped. "
|
|
'Cosmos3 example: \'{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false, '
|
|
'"use_resolution_template": false, "use_duration_template": false}\' '
|
|
'(add "generate_sound": true and optional "sound_duration" for synchronized audio). '
|
|
'Helios-Distilled example: \'{"is_enable_stage2": true, '
|
|
'"pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}\'.',
|
|
)
|
|
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
|
parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
|
|
parser.add_argument(
|
|
"--guidance-scale-high", type=float, default=None, help="Separate CFG for high-noise stage (Wan2.2 only)."
|
|
)
|
|
parser.add_argument("--height", type=int, default=None, help="Video height. Default: model-specific.")
|
|
parser.add_argument("--width", type=int, default=None, help="Video width. Default: model-specific.")
|
|
parser.add_argument("--num-frames", type=int, default=None, help="Number of frames. Default: model-specific.")
|
|
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=None, help="Sampling steps. Default: model-specific."
|
|
)
|
|
parser.add_argument(
|
|
"--boundary-ratio",
|
|
type=float,
|
|
default=None,
|
|
help="(Wan2.2) Boundary split ratio for low/high DiT. Default 0.875.",
|
|
)
|
|
parser.add_argument(
|
|
"--flow-shift",
|
|
type=float,
|
|
default=None,
|
|
help="Scheduler flow_shift. Wan2.2: 5.0(720p)/12.0(480p). HunyuanVideo-1.5: 5.0(480p)/9.0(720p).",
|
|
)
|
|
parser.add_argument(
|
|
"--cache-backend",
|
|
type=str,
|
|
default=None,
|
|
choices=["cache_dit"],
|
|
help="Cache backend for acceleration (Wan2.2). Default: None.",
|
|
)
|
|
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=None, help="Output path (mp4). Default: model-specific.")
|
|
parser.add_argument("--fps", type=int, default=None, 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(
|
|
"--audio-sample-rate",
|
|
type=int,
|
|
default=24000,
|
|
help="Sample rate for audio output when saved (default: 24000).",
|
|
)
|
|
parser.add_argument(
|
|
"--enable-diffusion-pipeline-profiler",
|
|
action="store_true",
|
|
help="Enable diffusion pipeline profiler to display stage durations.",
|
|
)
|
|
parser.add_argument(
|
|
"--profiler-config",
|
|
type=parse_profiler_config,
|
|
default=None,
|
|
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
|
|
)
|
|
parser.add_argument(
|
|
"--quantization",
|
|
type=str,
|
|
default=None,
|
|
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"],
|
|
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
|
|
)
|
|
|
|
# Distributed and parallel execution
|
|
parser.add_argument(
|
|
"--ulysses-mode",
|
|
type=str,
|
|
default="strict",
|
|
choices=["strict", "advanced_uaa"],
|
|
help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA).",
|
|
)
|
|
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(
|
|
"--vae-patch-parallel-size",
|
|
type=int,
|
|
default=1,
|
|
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
|
|
)
|
|
parser.add_argument(
|
|
"--pipeline-parallel-size",
|
|
type=int,
|
|
default=1,
|
|
help="Number of pipeline parallel stages.",
|
|
)
|
|
parser.add_argument(
|
|
"--enable-expert-parallel",
|
|
action="store_true",
|
|
help="Enable expert parallelism for MoE layers.",
|
|
)
|
|
parser.add_argument(
|
|
"--lora-path",
|
|
type=str,
|
|
nargs="+",
|
|
default=None,
|
|
help="Path to LoRA adapter folder (PEFT format) or concrete LoRA checkpoint files. Loaded at initialization and used for generation."
|
|
"Note: for Wan2.2 MoE models, two checkpoints positionally mapped to high-noise and low-noise modules.",
|
|
)
|
|
parser.add_argument(
|
|
"--lora-scale",
|
|
type=float,
|
|
default=1.0,
|
|
help="Scale factor for LoRA weights (default: 1.0).",
|
|
)
|
|
parser.add_argument(
|
|
"--lora-backend",
|
|
type=str,
|
|
default="peft",
|
|
choices=["peft", "distill"],
|
|
help="LoRA backend for loading LoRA adapters. Default: peft"
|
|
"'peft' loads a PEFT-format adapter folder, used e.g. for RL"
|
|
"'distill' fuses one or more concrete LoRA checkpoint files, used e.g. for distilled few-step LoRAs",
|
|
)
|
|
parser.add_argument(
|
|
"--use-hsdp",
|
|
action="store_true",
|
|
help="Enable HSDP (Hybrid Sharded Data Parallel) for diffusion models.",
|
|
)
|
|
parser.add_argument(
|
|
"--hsdp-shard-size",
|
|
type=int,
|
|
default=1,
|
|
help="Number of GPUs to shard weights across for HSDP.",
|
|
)
|
|
parser.add_argument(
|
|
"--hsdp-replicate-size",
|
|
type=int,
|
|
default=1,
|
|
help="Number of HSDP replica groups.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def _extract_peak_memory_mb(result: Any) -> float:
|
|
"""Pull worker-reported peak VRAM (MiB) generation result.
|
|
|
|
Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb.
|
|
"""
|
|
if isinstance(result, list):
|
|
result = result[0] if result else None
|
|
if result is None:
|
|
return 0.0
|
|
val = getattr(result, "peak_memory_mb", 0.0)
|
|
if not val:
|
|
inner = result
|
|
if isinstance(inner, list):
|
|
inner = inner[0] if inner else None
|
|
val = getattr(inner, "peak_memory_mb", 0.0)
|
|
try:
|
|
return float(val or 0.0)
|
|
except (TypeError, ValueError):
|
|
return 0.0
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
model_class_name = args.model_class_name
|
|
|
|
preset = _detect_preset(args.model, model_class_name)
|
|
for key, default_val in preset.items():
|
|
if getattr(args, key.replace("-", "_"), None) is None:
|
|
setattr(args, key.replace("-", "_"), default_val)
|
|
model_class_name = args.model_class_name
|
|
|
|
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
|
# Cache-dit config (Wan2.2 only)
|
|
cache_config = None
|
|
if args.cache_backend == "cache_dit":
|
|
cache_config = {
|
|
"Fn_compute_blocks": 1,
|
|
"Bn_compute_blocks": 0,
|
|
"max_warmup_steps": 4,
|
|
"max_cached_steps": 20,
|
|
"residual_diff_threshold": 0.24,
|
|
"max_continuous_cached_steps": 3,
|
|
"enable_taylorseer": False,
|
|
"taylorseer_order": 1,
|
|
"scm_steps_mask_policy": None,
|
|
"scm_steps_policy": "dynamic",
|
|
}
|
|
|
|
# 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,
|
|
pipeline_parallel_size=args.pipeline_parallel_size,
|
|
enable_expert_parallel=args.enable_expert_parallel,
|
|
)
|
|
|
|
profiler_enabled = args.profiler_config is not None
|
|
|
|
omni_kwargs = dict(
|
|
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,
|
|
model_class_name=model_class_name,
|
|
cache_backend=args.cache_backend,
|
|
cache_config=cache_config,
|
|
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
|
profiler_config=args.profiler_config,
|
|
)
|
|
if args.deploy_config:
|
|
omni_kwargs["deploy_config"] = args.deploy_config
|
|
if args.boundary_ratio is not None:
|
|
omni_kwargs["boundary_ratio"] = args.boundary_ratio
|
|
if args.flow_shift is not None:
|
|
omni_kwargs["flow_shift"] = args.flow_shift
|
|
if args.quantization is not None:
|
|
omni_kwargs["quantization"] = args.quantization
|
|
if args.cache_backend is not None:
|
|
omni_kwargs["cache_backend"] = args.cache_backend
|
|
omni_kwargs["cache_config"] = cache_config
|
|
omni_kwargs["enable_cache_dit_summary"] = args.enable_cache_dit_summary
|
|
if args.lora_path is not None:
|
|
lora_path = args.lora_path
|
|
if len(lora_path) == 1:
|
|
lora_path = lora_path[0]
|
|
omni_kwargs["lora_path"] = lora_path
|
|
omni_kwargs["lora_backend"] = args.lora_backend
|
|
|
|
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
|
|
# gate is an engine-level config (offline analog of the server's --no-guardrails).
|
|
if args.extra_body and "guardrails" in args.extra_body:
|
|
omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
|
|
|
|
omni = Omni(**omni_kwargs)
|
|
model_class_name = get_model_class_name(omni) or model_class_name
|
|
declared_extra_body_params = get_extra_body_params(model_class_name)
|
|
output_tensor_range = get_output_tensor_range(model_class_name)
|
|
|
|
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}, pipeline_parallel_size={args.pipeline_parallel_size},"
|
|
f" enable_expert_parallel={args.enable_expert_parallel}"
|
|
)
|
|
print(f" Video size: {args.width}x{args.height}")
|
|
print(f"{'=' * 60}\n")
|
|
|
|
lora_request = None
|
|
if args.lora_path and args.lora_backend == "peft":
|
|
if len(args.lora_path) != 1:
|
|
raise ValueError("Only one LoRA path is expected for PEFT backend.")
|
|
|
|
lora_path = args.lora_path[0]
|
|
lora_request_id = stable_lora_int_id(lora_path)
|
|
lora_request = LoRARequest(
|
|
lora_name=Path(lora_path).stem,
|
|
lora_int_id=lora_request_id,
|
|
lora_path=lora_path,
|
|
)
|
|
|
|
negative_prompt = args.negative_prompt
|
|
if negative_prompt is None and all(preset is not _MODEL_PRESETS[name] for name in ("lingbot", "ltx2", "ltx23")):
|
|
# Preserve the historical empty-prompt behavior for non-LTX examples.
|
|
negative_prompt = ""
|
|
prompt_dict = build_text_to_video_prompt(args.prompt, negative_prompt)
|
|
|
|
extra_args = {}
|
|
if lora_request:
|
|
extra_args["lora_request"] = lora_request
|
|
extra_args["lora_scale"] = args.lora_scale
|
|
|
|
sampling_kwargs = dict(
|
|
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,
|
|
)
|
|
if args.guidance_scale_high is not None:
|
|
sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
|
|
|
|
sampling_params = OmniDiffusionSamplingParams(**sampling_kwargs)
|
|
# Route model-specific knobs through extra_body, filtered against the model's
|
|
# declared extra_body_params. Models without a declaration only forward explicit
|
|
# --extra-body JSON (preserving the generic flags' legacy behavior).
|
|
extra_body = dict(args.extra_body or {})
|
|
if declared_extra_body_params:
|
|
apply_declared_extra_args(sampling_params, declared_extra_body_params, extra_body)
|
|
elif extra_body:
|
|
sampling_params.extra_args.update({k: v for k, v in extra_body.items() if v is not None})
|
|
|
|
generation_start = time.perf_counter()
|
|
frames = omni.generate(
|
|
prompt_dict,
|
|
sampling_params,
|
|
)
|
|
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)")
|
|
|
|
peak_mb = _extract_peak_memory_mb(frames)
|
|
if peak_mb:
|
|
print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
|
|
|
|
audio = None
|
|
audio_sample_rate = args.audio_sample_rate
|
|
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"]
|
|
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
|
|
if frames.is_pipeline_output and frames is not None:
|
|
inner_output = frames
|
|
if isinstance(inner_output, OmniRequestOutput):
|
|
if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
|
|
audio = inner_output.multimodal_output["audio"]
|
|
audio_sample_rate = inner_output.multimodal_output.get("audio_sample_rate", audio_sample_rate)
|
|
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")
|
|
audio_sample_rate = frames.images[0].get("audio_sample_rate", audio_sample_rate)
|
|
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")
|
|
audio_sample_rate = first_item.get("audio_sample_rate", audio_sample_rate)
|
|
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")
|
|
audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate)
|
|
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)
|
|
frame_tensor = _normalize_float_tensor(frame_tensor, output_tensor_range)
|
|
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)
|
|
video_tensor = _normalize_float_tensor(video_tensor, output_tensor_range)
|
|
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)
|
|
|
|
if audio is not None:
|
|
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
|
|
|
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 = (np.clip(frames_np, 0.0, 1.0) * 255).round().clip(0, 255).astype("uint8")
|
|
|
|
audio_np = audio
|
|
if isinstance(audio_np, list):
|
|
audio_np = audio_np[0] if audio_np else None
|
|
if isinstance(audio_np, torch.Tensor):
|
|
audio_np = audio_np.detach().cpu().float().numpy()
|
|
if isinstance(audio_np, np.ndarray):
|
|
audio_np = np.squeeze(audio_np).astype(np.float32)
|
|
|
|
video_bytes = mux_video_audio_bytes(
|
|
frames_u8,
|
|
audio_np,
|
|
fps=float(args.fps),
|
|
audio_sample_rate=audio_sample_rate,
|
|
)
|
|
with open(str(output_path), "wb") as f:
|
|
f.write(video_bytes)
|
|
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()
|