616 lines
26 KiB
Python
616 lines
26 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Generate raw LTX video and audio outputs with the official or Omni runtime."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
|
|
def _parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--backend", choices=("official", "omni"), required=True)
|
|
parser.add_argument("--request", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--model")
|
|
parser.add_argument("--model-class-name")
|
|
parser.add_argument("--official-root", type=Path)
|
|
parser.add_argument("--checkpoint", type=Path)
|
|
parser.add_argument("--gemma-root", type=Path)
|
|
parser.add_argument("--transformer-path", type=Path)
|
|
parser.add_argument("--text-encoder-path", type=Path)
|
|
parser.add_argument("--video-vae-path", type=Path)
|
|
parser.add_argument("--audio-vae-path", type=Path)
|
|
parser.add_argument("--spatial-upsampler-path", type=Path)
|
|
parser.add_argument("--distilled-lora-path", type=Path)
|
|
parser.add_argument("--connector-model", type=Path)
|
|
parser.add_argument(
|
|
"--official-pipeline",
|
|
choices=("distilled_two_stage", "full_one_stage", "full_two_stage"),
|
|
default="distilled_two_stage",
|
|
)
|
|
parser.add_argument("--enable-layerwise-offload", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def _save_outputs(
|
|
output_dir: Path,
|
|
*,
|
|
video: torch.Tensor,
|
|
audio: torch.Tensor,
|
|
audio_sample_rate: int,
|
|
fps: float,
|
|
metadata: dict[str, Any],
|
|
) -> None:
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
video = video.detach().float().cpu().clamp(0.0, 1.0)
|
|
audio = audio.detach().float().cpu()
|
|
np.save(output_dir / "video.npy", video.numpy())
|
|
np.save(output_dir / "audio.npy", audio.numpy())
|
|
|
|
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
|
|
|
frames_u8 = video.mul(255.0).round().to(torch.uint8).numpy()
|
|
audio_array = audio.numpy()
|
|
while audio_array.ndim > 2 and audio_array.shape[0] == 1:
|
|
audio_array = audio_array[0]
|
|
(output_dir / "output.mp4").write_bytes(
|
|
mux_video_audio_bytes(
|
|
frames_u8,
|
|
audio_array,
|
|
fps=float(fps),
|
|
audio_sample_rate=int(audio_sample_rate),
|
|
video_codec_options={"preset": "ultrafast"},
|
|
)
|
|
)
|
|
|
|
frame_indices = sorted({0, video.shape[0] // 2, video.shape[0] - 1})
|
|
for index in frame_indices:
|
|
frame = video[index].mul(255.0).round().to(torch.uint8).numpy()
|
|
Image.fromarray(frame).save(output_dir / f"frame_{index:04d}.png")
|
|
|
|
metadata.update(
|
|
{
|
|
"video_shape": list(video.shape),
|
|
"audio_shape": list(audio.shape),
|
|
"audio_sample_rate": int(audio_sample_rate),
|
|
"frame_indices": frame_indices,
|
|
"mp4": "output.mp4",
|
|
}
|
|
)
|
|
(output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
|
|
|
|
|
|
def _reset_accelerator_peak_memory_stats() -> None:
|
|
if torch.accelerator.is_available():
|
|
torch.accelerator.reset_peak_memory_stats()
|
|
|
|
|
|
def _accelerator_peak_memory_stats() -> dict[str, float]:
|
|
if not torch.accelerator.is_available():
|
|
return {}
|
|
torch.accelerator.synchronize()
|
|
mib = 1024**2
|
|
return {
|
|
"peak_memory_allocated_mb": torch.accelerator.max_memory_allocated() / mib,
|
|
"peak_memory_reserved_mb": torch.accelerator.max_memory_reserved() / mib,
|
|
}
|
|
|
|
|
|
def _insert_official_paths(official_root: Path) -> None:
|
|
for relative_path in ("packages/ltx-core/src", "packages/ltx-pipelines/src"):
|
|
path = str((official_root / relative_path).resolve())
|
|
if path not in sys.path:
|
|
sys.path.insert(0, path)
|
|
|
|
|
|
def _configure_official_vocoder_determinism() -> None:
|
|
"""Make the pinned official LTX-2.5 vocoder a stable reference."""
|
|
from ltx_core.model.audio_vae.vocoder import VocoderWithBWE
|
|
|
|
original_forward = VocoderWithBWE.forward
|
|
if getattr(original_forward, "_vllm_omni_deterministic", False):
|
|
return
|
|
|
|
def deterministic_forward(self: Any, mel_spec: torch.Tensor) -> torch.Tensor:
|
|
device_type = mel_spec.device.type
|
|
if device_type != "cuda":
|
|
return original_forward(self, mel_spec)
|
|
previous = torch.backends.cudnn.deterministic
|
|
try:
|
|
torch.backends.cudnn.deterministic = True
|
|
return original_forward(self, mel_spec)
|
|
finally:
|
|
torch.backends.cudnn.deterministic = previous
|
|
|
|
setattr(deterministic_forward, "_vllm_omni_deterministic", True)
|
|
setattr(VocoderWithBWE, "forward", deterministic_forward)
|
|
|
|
|
|
def _configure_official_sdpa(pipeline: Any) -> None:
|
|
"""Pin official connector and denoiser attention to cuDNN, with MATH fallback."""
|
|
from ltx_core.loader.attention_ops import set_attention_module_op
|
|
from ltx_core.model.transformer.attention import PytorchAttention
|
|
from torch.nn.attention import SDPBackend
|
|
|
|
attention = PytorchAttention(priority=[SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH])
|
|
module_op = set_attention_module_op(
|
|
attention=attention,
|
|
masked_attention=attention,
|
|
)
|
|
transformer_owners = [
|
|
owner
|
|
for attribute in ("stage", "stage_1", "stage_2")
|
|
if (owner := getattr(pipeline, attribute, None)) is not None
|
|
]
|
|
owners_and_attributes = [
|
|
*((owner, "_transformer_builder") for owner in transformer_owners),
|
|
(pipeline.prompt_encoder, "_embeddings_processor_builder"),
|
|
]
|
|
for owner, attribute in owners_and_attributes:
|
|
builder = getattr(owner, attribute)
|
|
setattr(
|
|
owner,
|
|
attribute,
|
|
builder.with_module_ops((*builder.module_ops, module_op)),
|
|
)
|
|
|
|
|
|
def _use_diffusers_connector_weights(prompt_encoder: Any, model: Path) -> None:
|
|
"""Run official connector code with the checkpoint under test."""
|
|
from safetensors import safe_open
|
|
|
|
connector_root = model / "connectors"
|
|
shards = sorted(connector_root.glob("*.safetensors"))
|
|
if not shards:
|
|
raise ValueError(f"No connector safetensors found under {connector_root}")
|
|
|
|
original_build_embeddings_processor = prompt_encoder._build_embeddings_processor
|
|
|
|
def build_embeddings_processor():
|
|
processor = original_build_embeddings_processor()
|
|
parameters = dict(processor.named_parameters())
|
|
expected = {name for name in parameters if name.startswith(("video_connector.", "audio_connector."))}
|
|
loaded: set[str] = set()
|
|
with torch.no_grad():
|
|
for shard in shards:
|
|
with safe_open(str(shard), framework="pt", device="cpu") as weights:
|
|
for source_name in weights.keys():
|
|
if not source_name.startswith(("video_connector.", "audio_connector.")):
|
|
continue
|
|
target_name = (
|
|
source_name.replace(".transformer_blocks.", ".transformer_1d_blocks.")
|
|
.replace(".attn1.norm_q.", ".attn1.q_norm.")
|
|
.replace(".attn1.norm_k.", ".attn1.k_norm.")
|
|
)
|
|
parameter = parameters[target_name]
|
|
parameter.copy_(weights.get_tensor(source_name).to(parameter.device, parameter.dtype))
|
|
loaded.add(target_name)
|
|
if loaded != expected:
|
|
missing = sorted(expected - loaded)
|
|
unexpected = sorted(loaded - expected)
|
|
raise ValueError(f"Connector weight mismatch: missing={missing}, unexpected={unexpected}")
|
|
return processor
|
|
|
|
prompt_encoder._build_embeddings_processor = build_embeddings_processor
|
|
|
|
|
|
@torch.inference_mode()
|
|
def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
|
|
if args.official_root is None:
|
|
raise ValueError("Official backend requires --official-root")
|
|
_insert_official_paths(args.official_root)
|
|
_configure_official_vocoder_determinism()
|
|
# Keep Gemma on the same SDPA path in both subprocesses. DiT attention is
|
|
# still pinned independently to the explicit cuDNN backend below.
|
|
torch.backends.cuda.enable_cudnn_sdp(False)
|
|
_reset_accelerator_peak_memory_stats()
|
|
|
|
from ltx_pipelines.utils.args import ImageConditioningInput
|
|
from ltx_pipelines.utils.model_paths import ModelPaths
|
|
from ltx_pipelines.utils.types import OffloadMode
|
|
|
|
offload_mode = OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE
|
|
image_path = request.get("image")
|
|
images = (
|
|
[]
|
|
if image_path is None
|
|
else [
|
|
ImageConditioningInput(
|
|
path=str(image_path),
|
|
frame_idx=0,
|
|
strength=1.0,
|
|
crf=request.get("image_crf", 18),
|
|
)
|
|
]
|
|
)
|
|
|
|
if args.transformer_path is not None:
|
|
required = {
|
|
"--text-encoder-path": args.text_encoder_path,
|
|
"--video-vae-path": args.video_vae_path,
|
|
"--audio-vae-path": args.audio_vae_path,
|
|
}
|
|
if args.official_pipeline in {"distilled_two_stage", "full_two_stage"}:
|
|
required["--spatial-upsampler-path"] = args.spatial_upsampler_path
|
|
if args.official_pipeline == "full_two_stage":
|
|
required["--distilled-lora-path"] = args.distilled_lora_path
|
|
missing = [name for name, value in required.items() if value is None]
|
|
if missing:
|
|
raise ValueError(f"Official LTX-2.5 split backend is missing: {', '.join(missing)}")
|
|
|
|
model_paths = ModelPaths.from_split(
|
|
transformer_path=str(args.transformer_path),
|
|
text_encoder_path=str(args.text_encoder_path),
|
|
video_vae_path=str(args.video_vae_path),
|
|
audio_vae_path=str(args.audio_vae_path),
|
|
)
|
|
if args.official_pipeline == "distilled_two_stage":
|
|
from ltx_pipelines.distilled import DistilledPipeline
|
|
|
|
pipeline = DistilledPipeline(
|
|
model_paths=model_paths,
|
|
spatial_upsampler_path=str(args.spatial_upsampler_path),
|
|
loras=(),
|
|
offload_mode=offload_mode,
|
|
)
|
|
if args.connector_model is not None:
|
|
_use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model)
|
|
_configure_official_sdpa(pipeline)
|
|
video, audio, _, _ = pipeline(
|
|
prompt=request["prompt"],
|
|
seed=request["seed"],
|
|
height=request["height"],
|
|
width=request["width"],
|
|
num_frames=request["num_frames"],
|
|
frame_rate=request["fps"],
|
|
images=images,
|
|
)
|
|
pipeline_name = "DistilledPipeline"
|
|
elif args.official_pipeline == "full_two_stage":
|
|
from ltx_core.components.guiders import MultiModalGuiderParams
|
|
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
|
|
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
|
|
|
|
pipeline = TI2VidTwoStagesPipeline(
|
|
model_paths=model_paths,
|
|
distilled_lora=[
|
|
LoraPathStrengthAndSDOps(
|
|
path=str(args.distilled_lora_path),
|
|
strength=1.0,
|
|
sd_ops=LTXV_LORA_COMFY_RENAMING_MAP,
|
|
)
|
|
],
|
|
spatial_upsampler_path=str(args.spatial_upsampler_path),
|
|
loras=[],
|
|
offload_mode=offload_mode,
|
|
)
|
|
if args.connector_model is not None:
|
|
_use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model)
|
|
_configure_official_sdpa(pipeline)
|
|
video, audio, _, _ = pipeline(
|
|
prompt=request["prompt"],
|
|
negative_prompt=request["negative_prompt"],
|
|
seed=request["seed"],
|
|
height=request["height"],
|
|
width=request["width"],
|
|
num_frames=request["num_frames"],
|
|
frame_rate=request["fps"],
|
|
num_inference_steps=request["num_inference_steps"],
|
|
video_guider_params=MultiModalGuiderParams(
|
|
cfg_scale=request["video_cfg_scale"],
|
|
stg_scale=request["video_stg_scale"],
|
|
rescale_scale=request["video_rescale_scale"],
|
|
modality_scale=request["video_modality_scale"],
|
|
skip_step=0,
|
|
stg_blocks=request["video_stg_blocks"],
|
|
),
|
|
audio_guider_params=MultiModalGuiderParams(
|
|
cfg_scale=request["audio_cfg_scale"],
|
|
stg_scale=request["audio_stg_scale"],
|
|
rescale_scale=request["audio_rescale_scale"],
|
|
modality_scale=request["audio_modality_scale"],
|
|
skip_step=0,
|
|
stg_blocks=request["audio_stg_blocks"],
|
|
),
|
|
images=images,
|
|
max_batch_size=4,
|
|
stage_1_sigmas=torch.tensor(request["stage_1_sigmas"], dtype=torch.float32),
|
|
stage_2_sigmas=torch.tensor(request["stage_2_sigmas"], dtype=torch.float32),
|
|
)
|
|
pipeline_name = "TI2VidTwoStagesPipeline"
|
|
else:
|
|
from ltx_core.components.guiders import MultiModalGuiderParams
|
|
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
|
|
|
|
pipeline = TI2VidOneStagePipeline(
|
|
model_paths=model_paths,
|
|
loras=(),
|
|
offload_mode=offload_mode,
|
|
)
|
|
if args.connector_model is not None:
|
|
_use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model)
|
|
_configure_official_sdpa(pipeline)
|
|
request_sigmas = request.get("sigmas")
|
|
video, audio, _ = pipeline(
|
|
prompt=request["prompt"],
|
|
negative_prompt=request["negative_prompt"],
|
|
seed=request["seed"],
|
|
height=request["height"],
|
|
width=request["width"],
|
|
num_frames=request["num_frames"],
|
|
frame_rate=request["fps"],
|
|
num_inference_steps=request["num_inference_steps"],
|
|
video_guider_params=MultiModalGuiderParams(
|
|
cfg_scale=request["video_cfg_scale"],
|
|
stg_scale=request["video_stg_scale"],
|
|
rescale_scale=request["video_rescale_scale"],
|
|
modality_scale=request["video_modality_scale"],
|
|
skip_step=0,
|
|
stg_blocks=request["video_stg_blocks"],
|
|
),
|
|
audio_guider_params=MultiModalGuiderParams(
|
|
cfg_scale=request["audio_cfg_scale"],
|
|
stg_scale=request["audio_stg_scale"],
|
|
rescale_scale=request["audio_rescale_scale"],
|
|
modality_scale=request["audio_modality_scale"],
|
|
skip_step=0,
|
|
stg_blocks=request["audio_stg_blocks"],
|
|
),
|
|
images=images,
|
|
max_batch_size=4,
|
|
sigmas=None if request_sigmas is None else torch.tensor(request_sigmas, dtype=torch.float32),
|
|
)
|
|
pipeline_name = "TI2VidOneStagePipeline"
|
|
checkpoint_metadata = {
|
|
"transformer_path": str(args.transformer_path),
|
|
"video_vae_path": str(args.video_vae_path),
|
|
"connector_model": None if args.connector_model is None else str(args.connector_model),
|
|
"pipeline": pipeline_name,
|
|
}
|
|
else:
|
|
if args.checkpoint is None or args.gemma_root is None:
|
|
raise ValueError("Official monolith backend requires --checkpoint and --gemma-root")
|
|
|
|
from ltx_core.components.guiders import MultiModalGuiderParams
|
|
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
|
|
|
|
pipeline = TI2VidOneStagePipeline(
|
|
model_paths=ModelPaths.from_monolith(str(args.checkpoint), str(args.gemma_root)),
|
|
loras=(),
|
|
offload_mode=offload_mode,
|
|
)
|
|
_configure_official_sdpa(pipeline)
|
|
video, audio, _ = pipeline(
|
|
prompt=request["prompt"],
|
|
negative_prompt=request["negative_prompt"],
|
|
seed=request["seed"],
|
|
height=request["height"],
|
|
width=request["width"],
|
|
num_frames=request["num_frames"],
|
|
frame_rate=request["fps"],
|
|
num_inference_steps=request["num_inference_steps"],
|
|
video_guider_params=MultiModalGuiderParams(
|
|
cfg_scale=request["video_cfg_scale"],
|
|
stg_scale=request["video_stg_scale"],
|
|
rescale_scale=request["video_rescale_scale"],
|
|
modality_scale=request["video_modality_scale"],
|
|
skip_step=0,
|
|
stg_blocks=request["video_stg_blocks"],
|
|
),
|
|
audio_guider_params=MultiModalGuiderParams(
|
|
cfg_scale=request["audio_cfg_scale"],
|
|
stg_scale=request["audio_stg_scale"],
|
|
rescale_scale=request["audio_rescale_scale"],
|
|
modality_scale=request["audio_modality_scale"],
|
|
skip_step=0,
|
|
stg_blocks=request["audio_stg_blocks"],
|
|
),
|
|
images=images,
|
|
max_batch_size=4,
|
|
)
|
|
checkpoint_metadata = {
|
|
"checkpoint": str(args.checkpoint),
|
|
"pipeline": "TI2VidOneStagePipeline",
|
|
}
|
|
|
|
video_tensor = torch.cat([chunk.detach().cpu() for chunk in video], dim=0)
|
|
_save_outputs(
|
|
args.output_dir,
|
|
video=_canonical_video(video_tensor),
|
|
audio=audio.waveform,
|
|
audio_sample_rate=audio.sampling_rate,
|
|
fps=float(request["fps"]),
|
|
metadata={
|
|
"backend": "official",
|
|
"attention_backend": "torch_sdpa_cudnn_bf16",
|
|
"official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"),
|
|
"seed": request["seed"],
|
|
"sigmas": request.get("sigmas"),
|
|
"stage_1_sigmas": request.get("stage_1_sigmas"),
|
|
"stage_2_sigmas": request.get("stage_2_sigmas"),
|
|
**_accelerator_peak_memory_stats(),
|
|
**checkpoint_metadata,
|
|
},
|
|
)
|
|
|
|
|
|
def _omni_worker_peak_memory_mb(output: Any) -> float:
|
|
result = output[0] if isinstance(output, list) and output else output
|
|
return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
|
|
|
|
|
|
def _unwrap_omni_output(output: Any) -> tuple[Any, Any, int]:
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
|
|
audio = None
|
|
audio_sample_rate = None
|
|
frames = output[0] if isinstance(output, list) and output else output
|
|
if isinstance(frames, OmniRequestOutput):
|
|
multimodal_output = frames.multimodal_output or {}
|
|
audio = multimodal_output.get("audio")
|
|
audio_sample_rate = multimodal_output.get("audio_sample_rate")
|
|
nested_output = getattr(frames, "request_output", None)
|
|
if frames.is_pipeline_output and isinstance(nested_output, OmniRequestOutput):
|
|
frames = nested_output
|
|
multimodal_output = frames.multimodal_output or {}
|
|
audio = multimodal_output.get("audio", audio)
|
|
audio_sample_rate = multimodal_output.get("audio_sample_rate", audio_sample_rate)
|
|
if isinstance(frames, OmniRequestOutput):
|
|
if not frames.images:
|
|
raise ValueError("No video frames found in OmniRequestOutput")
|
|
frames = frames.images
|
|
|
|
if isinstance(frames, list) and len(frames) == 1:
|
|
frames = frames[0]
|
|
if isinstance(frames, tuple) and len(frames) == 2:
|
|
frames, audio = frames
|
|
elif isinstance(frames, dict):
|
|
audio = frames.get("audio", audio)
|
|
audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate)
|
|
frames = frames.get("frames", frames.get("video"))
|
|
|
|
if frames is None or audio is None or audio_sample_rate is None:
|
|
raise ValueError("Omni output did not contain video, audio, and audio_sample_rate")
|
|
return frames, audio, int(audio_sample_rate)
|
|
|
|
|
|
def _canonical_video(video: Any) -> torch.Tensor:
|
|
if isinstance(video, list):
|
|
if len(video) == 1:
|
|
return _canonical_video(video[0])
|
|
tensors = [torch.as_tensor(np.asarray(item) if not isinstance(item, torch.Tensor) else item) for item in video]
|
|
if all(item.ndim == 3 for item in tensors):
|
|
video = torch.stack(tensors)
|
|
elif all(item.ndim == 4 for item in tensors):
|
|
video = torch.cat(tensors)
|
|
else:
|
|
raise ValueError(f"Cannot combine video items with dimensions {[item.ndim for item in tensors]}")
|
|
tensor = torch.as_tensor(np.asarray(video) if not isinstance(video, torch.Tensor) else video).detach().cpu()
|
|
if tensor.ndim == 5:
|
|
tensor = tensor[0]
|
|
if tensor.ndim != 4:
|
|
raise ValueError(f"Expected a 4D video tensor, got {tuple(tensor.shape)}")
|
|
if tensor.shape[-1] in (3, 4):
|
|
tensor = tensor[..., :3]
|
|
elif tensor.shape[1] in (3, 4):
|
|
tensor = tensor[:, :3].permute(0, 2, 3, 1)
|
|
elif tensor.shape[0] in (3, 4):
|
|
tensor = tensor[:3].permute(1, 2, 3, 0)
|
|
else:
|
|
raise ValueError(f"Cannot infer video channel dimension from {tuple(tensor.shape)}")
|
|
tensor = tensor.float()
|
|
if tensor.numel() and tensor.max() > 1:
|
|
tensor = tensor / 255.0
|
|
elif tensor.numel() and tensor.min() < 0:
|
|
tensor = tensor.clamp(-1.0, 1.0).add(1.0).mul(0.5)
|
|
return tensor.clamp(0.0, 1.0)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
|
|
if args.model is None or args.model_class_name is None:
|
|
raise ValueError("Omni backend requires --model and --model-class-name")
|
|
|
|
# vLLM imports can otherwise change the process-wide Gemma dispatch.
|
|
torch.backends.cuda.enable_cudnn_sdp(False)
|
|
_reset_accelerator_peak_memory_stats()
|
|
|
|
attention_config = {"default": {"backend": "CUDNN_ATTN"}}
|
|
|
|
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.model_extras import get_extra_body_params
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(request["seed"])
|
|
omni = Omni(
|
|
model=args.model,
|
|
model_class_name=args.model_class_name,
|
|
enforce_eager=True,
|
|
enable_layerwise_offload=args.enable_layerwise_offload,
|
|
diffusion_attention_config=attention_config,
|
|
parallel_config=DiffusionParallelConfig(),
|
|
)
|
|
try:
|
|
# The explicit public class argument selects the worker pipeline. The
|
|
# parent Omni proxy can only expose the checkpoint's base class, which
|
|
# is insufficient to distinguish LTX execution modes.
|
|
model_class_name = args.model_class_name
|
|
sampling_params = OmniDiffusionSamplingParams(
|
|
height=request["height"],
|
|
width=request["width"],
|
|
generator=generator,
|
|
guidance_scale=None,
|
|
num_inference_steps=request["num_inference_steps"],
|
|
num_frames=request["num_frames"],
|
|
fps=request["fps"],
|
|
frame_rate=float(request["fps"]),
|
|
output_type="np",
|
|
)
|
|
extra_args = {
|
|
key: value for key, value in request.items() if key.startswith("video_") or key.startswith("audio_")
|
|
}
|
|
for key in ("sigmas", "stage_1_sigmas", "stage_2_sigmas", "image_crf"):
|
|
if key in request:
|
|
extra_args[key] = request[key]
|
|
apply_declared_extra_args(sampling_params, get_extra_body_params(model_class_name), extra_args)
|
|
prompt: dict[str, Any] = {"prompt": request["prompt"]}
|
|
if "negative_prompt" in request:
|
|
prompt["negative_prompt"] = request["negative_prompt"]
|
|
image_path = request.get("image")
|
|
if image_path is not None:
|
|
with Image.open(str(image_path)) as source_image:
|
|
image = source_image.convert("RGB")
|
|
image.load()
|
|
prompt["multi_modal_data"] = {"image": image}
|
|
output = omni.generate(prompt, sampling_params)
|
|
video, audio, audio_sample_rate = _unwrap_omni_output(output)
|
|
audio_tensor = torch.as_tensor(np.asarray(audio) if not isinstance(audio, torch.Tensor) else audio)
|
|
_save_outputs(
|
|
args.output_dir,
|
|
video=_canonical_video(video),
|
|
audio=audio_tensor,
|
|
audio_sample_rate=audio_sample_rate,
|
|
fps=float(request["fps"]),
|
|
metadata={
|
|
"backend": "omni",
|
|
"attention_backend": "torch_sdpa_cudnn_bf16",
|
|
"seed": request["seed"],
|
|
"sigmas": request.get("sigmas"),
|
|
"model": args.model,
|
|
"stage_1_sigmas": request.get("stage_1_sigmas"),
|
|
"stage_2_sigmas": request.get("stage_2_sigmas"),
|
|
"model_class_name": model_class_name,
|
|
"worker_peak_memory_mb": _omni_worker_peak_memory_mb(output),
|
|
**_accelerator_peak_memory_stats(),
|
|
},
|
|
)
|
|
finally:
|
|
omni.shutdown()
|
|
|
|
|
|
def main() -> None:
|
|
args = _parse_args()
|
|
request = json.loads(args.request.read_text())
|
|
if args.backend == "official":
|
|
_run_official(args, request)
|
|
else:
|
|
_run_omni(args, request)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|