Boyang ZHONG
dc72a35f4e
For n > 1, the image endpoints returned one preview grid instead of n images. The endpoints now run one model batch, save each decoded sample as its own file, and return one URL or base64 item for each image. This applies to generations and edits. A keyword-only return_samples field (default False) is added to OutputConfig, SamplingParam and ForwardBatch. Batched image requests set it and do not build preview frames. The worker keeps the sample tensor for these requests, and non-output ranks clear the field. Defaults and positional arguments do not change.
1553 lines
66 KiB
Python
1553 lines
66 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
"""
|
||
VideoGenerator module for FastVideo.
|
||
|
||
This module provides a consolidated interface for generating videos using
|
||
diffusion models.
|
||
"""
|
||
|
||
import os
|
||
import re
|
||
import shutil
|
||
import subprocess
|
||
import threading
|
||
import time
|
||
import tempfile
|
||
import types
|
||
import warnings
|
||
from collections.abc import Callable, Mapping
|
||
from contextlib import suppress
|
||
from copy import deepcopy
|
||
from functools import partial
|
||
from typing import Any, cast, TypeVar
|
||
|
||
import imageio
|
||
import numpy as np
|
||
import torch
|
||
import torchvision
|
||
from einops import rearrange
|
||
|
||
import fastvideo.envs as envs
|
||
from fastvideo.api.compat import (
|
||
REQUEST_BATCH_EXTRA_PASSTHROUGH_FIELDS,
|
||
expand_request_prompt_batch,
|
||
generator_config_to_fastvideo_args,
|
||
legacy_from_pretrained_to_config,
|
||
legacy_generate_call_to_request,
|
||
load_generator_config_from_file,
|
||
normalize_generation_request,
|
||
normalize_generator_config,
|
||
request_to_batch_extra,
|
||
validate_request_batch_extra,
|
||
request_to_pipeline_overrides,
|
||
request_to_sampling_param,
|
||
)
|
||
from fastvideo.api.results import (
|
||
GenerationResult,
|
||
VideoFinalEvent,
|
||
VideoProgressEvent,
|
||
)
|
||
from fastvideo.api.schema import (
|
||
GenerationRequest,
|
||
GeneratorConfig,
|
||
InputConfig,
|
||
OutputConfig,
|
||
SamplingConfig,
|
||
)
|
||
from fastvideo.api.sampling_param import SamplingParam
|
||
from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
|
||
from fastvideo.logger import init_logger
|
||
from fastvideo.pipelines import ForwardBatch
|
||
from fastvideo.utils import align_to, allocate_cpu_tensor_with_pin_fallback, pixels_to_uint8, shallow_asdict
|
||
from fastvideo.worker.executor import Executor
|
||
|
||
fcntl: types.ModuleType | None
|
||
try:
|
||
import fcntl
|
||
except ImportError:
|
||
fcntl = None
|
||
|
||
logger = init_logger(__name__)
|
||
_FFMPEG_ENCODER_OPTION_CACHE: dict[tuple[str, str, str], bool] = {}
|
||
_T = TypeVar("_T")
|
||
|
||
_BATCH_EXTRA_PASSTHROUGH_KEYS = tuple(REQUEST_BATCH_EXTRA_PASSTHROUGH_FIELDS)
|
||
|
||
_FROM_PRETRAINED_CONVENIENCE_KWARGS = frozenset({
|
||
"num_gpus",
|
||
"revision",
|
||
"trust_remote_code",
|
||
"distributed_executor_backend",
|
||
"tp_size",
|
||
"sp_size",
|
||
"hsdp_replicate_dim",
|
||
"hsdp_shard_dim",
|
||
"dist_timeout",
|
||
"use_fsdp_inference",
|
||
"disable_autocast",
|
||
"enable_stage_verification",
|
||
"dit_cpu_offload",
|
||
"dit_layerwise_offload",
|
||
"text_encoder_cpu_offload",
|
||
"image_encoder_cpu_offload",
|
||
"vae_cpu_offload",
|
||
"pin_cpu_memory",
|
||
"enable_torch_compile",
|
||
"torch_compile_kwargs",
|
||
"lora_path",
|
||
"lora_strength",
|
||
"output_type",
|
||
"nvfp4_fa4",
|
||
"fa4_pv_mode",
|
||
})
|
||
|
||
|
||
def _infer_latent_batch_size(batch: ForwardBatch) -> int:
|
||
if isinstance(batch.prompt, list):
|
||
latent_batch_size = len(batch.prompt)
|
||
elif batch.prompt is not None:
|
||
latent_batch_size = 1
|
||
elif batch.prompt_embeds is not None and len(batch.prompt_embeds) > 0:
|
||
latent_batch_size = batch.prompt_embeds[0].shape[0]
|
||
else:
|
||
raise ValueError("Cannot infer batch size from batch; no prompt or prompt_embeds found")
|
||
latent_batch_size *= batch.num_videos_per_prompt
|
||
return latent_batch_size
|
||
|
||
|
||
def _resolve_output_size(
|
||
samples: torch.Tensor,
|
||
fallback: tuple[int, int, int],
|
||
*,
|
||
pixel_output: bool,
|
||
) -> tuple[int, int, int]:
|
||
"""Report the final decoded video's `(height, width, frames)`.
|
||
|
||
Refiner stages can produce a different resolution from the base request, so
|
||
pixel outputs use the final `[batch, channels, frames, height, width]` tensor.
|
||
Latent and audio outputs keep the requested fallback because their tensor
|
||
dimensions do not describe decoded pixels.
|
||
"""
|
||
if pixel_output and samples.ndim == 5:
|
||
return (int(samples.shape[-2]), int(samples.shape[-1]), int(samples.shape[-3]))
|
||
return fallback
|
||
|
||
|
||
def _quantize_video_frames_to_uint8(src: torch.Tensor) -> torch.Tensor:
|
||
"""`[b, c, t, h, w]` decoded output -> CPU uint8 `[t, b, c, h, w]`.
|
||
|
||
Quantizes on ``src``'s own device (typically CUDA) before the device->host copy: a full fp32 D->H copy scales with
|
||
resolution x frames x batch, while casting to uint8 first makes the transfer 4x smaller and moves the elementwise
|
||
work onto the GPU. ``clamp_()`` also fixes a latent overflow bug (VAE output slightly outside [0, 1] wrapped mod
|
||
256 in the old unclamped cast). Equivalence with a CPU quantize is SSIM-gated, not bit-exact (float->uint8 differs
|
||
<=1 LSB CPU vs GPU). uint8 input is already quantized by the worker (MiniMax-H3 decode stage) and passes through
|
||
untouched (see ``pixels_to_uint8``).
|
||
|
||
A large DiT run with ``dit_cpu_offload=False`` can leave a GPU with only a few hundred MB free once denoising and
|
||
decode are done, in a *different* process than the one that allocates ``src`` (the worker keeps the model
|
||
resident); the tiny elementwise allocation this needs can then be the single straw that OOMs. That failure is
|
||
recoverable -- the data itself is already on the device (``src`` is IPC-shared, not a fresh allocation) and only
|
||
needs a device->host copy, which does not need a device allocation -- so on ``torch.cuda.OutOfMemoryError`` this
|
||
falls back to copying ``src`` to CPU first and quantizing there instead of crashing the whole generation.
|
||
"""
|
||
try:
|
||
vid_u8 = pixels_to_uint8(src)
|
||
return rearrange(vid_u8, "b c t h w -> t b c h w").cpu()
|
||
except torch.cuda.OutOfMemoryError:
|
||
logger.warning("GPU out of memory quantizing the decoded video (the DiT/VAE likely still hold the device "
|
||
"resident); falling back to a CPU quantize. Pass dit_cpu_offload=True / "
|
||
"vae_cpu_offload=True to avoid this, at some speed cost.")
|
||
src_cpu = src.detach().to("cpu")
|
||
vid_u8 = pixels_to_uint8(src_cpu if src_cpu.dtype == torch.uint8 else src_cpu.float())
|
||
return rearrange(vid_u8, "b c t h w -> t b c h w")
|
||
|
||
|
||
def _validate_request_stage_overrides(model_path: str, request: GenerationRequest) -> None:
|
||
"""Validate typed stage overrides against the model's registered preset."""
|
||
if not request.stage_overrides:
|
||
return
|
||
from fastvideo.api.presets import validate_preset_selection
|
||
from fastvideo.registry import get_preset_selection
|
||
preset_name, model_family = get_preset_selection(model_path)
|
||
if preset_name is None or model_family is None:
|
||
raise ValueError(f"Model {model_path!r} has no preset for stage override validation")
|
||
validate_preset_selection(
|
||
preset_name,
|
||
model_family,
|
||
stage_overrides=request.stage_overrides,
|
||
)
|
||
|
||
|
||
class VideoGenerator:
|
||
"""
|
||
A unified class for generating videos using diffusion models.
|
||
|
||
This class provides a simple interface for video generation with rich
|
||
customization options, similar to popular frameworks like HF Diffusers.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
fastvideo_args: FastVideoArgs,
|
||
executor_class: type[Executor],
|
||
log_stats: bool,
|
||
*,
|
||
log_queue=None,
|
||
):
|
||
"""
|
||
Initialize the video generator.
|
||
|
||
Args:
|
||
fastvideo_args: The inference arguments
|
||
executor_class: The executor class to use for inference
|
||
log_stats: Whether to log statistics
|
||
log_queue: Optional multiprocessing.Queue to forward worker logs to
|
||
"""
|
||
self.config: GeneratorConfig | None = None
|
||
self.fastvideo_args = fastvideo_args
|
||
self.executor = executor_class(fastvideo_args, log_queue=log_queue)
|
||
|
||
@classmethod
|
||
def from_pretrained(
|
||
cls,
|
||
model_path: str | GeneratorConfig | Mapping[str, Any] | None = None,
|
||
**kwargs,
|
||
) -> "VideoGenerator":
|
||
"""
|
||
Create a video generator from a pretrained model.
|
||
|
||
Args:
|
||
model_path: Path or identifier for the pretrained model
|
||
pipeline_config: Pipeline config to use for inference
|
||
**kwargs: Additional arguments to customize model loading, set any FastVideoArgs or PipelineConfig attributes here.
|
||
|
||
Returns:
|
||
The created video generator
|
||
|
||
Priority level: Default pipeline config < User's pipeline config < User's kwargs
|
||
|
||
Stable convenience kwargs remain supported here for common engine and
|
||
offload settings. Advanced model- or pipeline-specific options should
|
||
move to VideoGenerator.from_config(...).
|
||
"""
|
||
log_queue = kwargs.pop("log_queue", None)
|
||
if kwargs.pop("nvfp4_fa4", False):
|
||
import os
|
||
os.environ["FASTVIDEO_NVFP4_FA4"] = "1"
|
||
envs.setdefault_external("CUTE_DSL_ENABLE_TVM_FFI", "1")
|
||
fa4_pv_mode = kwargs.pop("fa4_pv_mode", None)
|
||
if fa4_pv_mode is not None:
|
||
# Same env bridge as nvfp4_fa4: model code constructs attention
|
||
# impls with fixed literals, so the knob has no kwarg path of its
|
||
# own (and FastVideoArgs would drop it).
|
||
envs.FASTVIDEO_FA4_PV_MODE.set(str(fa4_pv_mode))
|
||
typed_config = kwargs.pop("config", None)
|
||
if typed_config is not None:
|
||
if model_path is not None:
|
||
raise TypeError("Pass either model_path or config to from_pretrained, not both")
|
||
if kwargs:
|
||
unexpected = ", ".join(sorted(kwargs))
|
||
raise TypeError(f"Unexpected keyword arguments with config: {unexpected}")
|
||
return cls.from_config(typed_config, log_queue=log_queue)
|
||
|
||
if isinstance(model_path, GeneratorConfig | Mapping):
|
||
if kwargs:
|
||
unexpected = ", ".join(sorted(kwargs))
|
||
raise TypeError(f"Unexpected keyword arguments with typed config: {unexpected}")
|
||
return cls.from_config(model_path, log_queue=log_queue)
|
||
|
||
if model_path is None:
|
||
raise TypeError("model_path or config is required")
|
||
|
||
legacy_only_kwargs = sorted(set(kwargs) - _FROM_PRETRAINED_CONVENIENCE_KWARGS)
|
||
if legacy_only_kwargs:
|
||
warnings.warn(
|
||
"VideoGenerator.from_pretrained(...) received legacy-only kwargs "
|
||
f"({', '.join(legacy_only_kwargs)}); prefer VideoGenerator.from_config(...) "
|
||
"for advanced configuration.",
|
||
DeprecationWarning,
|
||
stacklevel=2,
|
||
)
|
||
return cls.from_config(
|
||
legacy_from_pretrained_to_config(model_path, kwargs),
|
||
log_queue=log_queue,
|
||
)
|
||
|
||
@classmethod
|
||
def from_config(
|
||
cls,
|
||
config: GeneratorConfig | Mapping[str, Any],
|
||
*,
|
||
log_queue=None,
|
||
) -> "VideoGenerator":
|
||
normalized = normalize_generator_config(config)
|
||
fastvideo_args = generator_config_to_fastvideo_args(normalized)
|
||
generator = cls.from_fastvideo_args(fastvideo_args, log_queue=log_queue)
|
||
generator.config = normalized
|
||
return generator
|
||
|
||
@classmethod
|
||
def from_file(
|
||
cls,
|
||
path: str,
|
||
overrides: list[str] | Mapping[str, Any] | None = None,
|
||
*,
|
||
log_queue=None,
|
||
) -> "VideoGenerator":
|
||
return cls.from_config(
|
||
load_generator_config_from_file(path, overrides=overrides),
|
||
log_queue=log_queue,
|
||
)
|
||
|
||
@classmethod
|
||
def from_fastvideo_args(
|
||
cls,
|
||
fastvideo_args: FastVideoArgs,
|
||
*,
|
||
log_queue=None,
|
||
) -> "VideoGenerator":
|
||
"""
|
||
Create a video generator with the specified arguments.
|
||
|
||
Args:
|
||
fastvideo_args: The inference arguments
|
||
log_queue: Optional multiprocessing.Queue to forward worker logs to
|
||
|
||
Returns:
|
||
The created video generator
|
||
"""
|
||
# Initialize distributed environment if needed
|
||
# initialize_distributed_and_parallelism(fastvideo_args)
|
||
|
||
executor_class = Executor.get_class(fastvideo_args, allow_external_launcher=True)
|
||
return cls(
|
||
fastvideo_args=fastvideo_args,
|
||
executor_class=executor_class,
|
||
log_stats=False, # TODO: implement
|
||
log_queue=log_queue,
|
||
)
|
||
|
||
def _run_on_output_rank(
|
||
self,
|
||
operation: Callable[[], _T],
|
||
*,
|
||
description: str,
|
||
) -> _T:
|
||
"""Run rank-local setup on rank zero and broadcast its result/error."""
|
||
if not self.executor.uses_spmd_execution:
|
||
return operation()
|
||
|
||
payload: tuple[bool, _T | str] | None = None
|
||
source_error: Exception | None = None
|
||
if self.executor.is_output_rank:
|
||
try:
|
||
payload = (True, operation())
|
||
except Exception as error:
|
||
source_error = error
|
||
payload = (False, f"{type(error).__name__}: {error}")
|
||
|
||
synchronized = self.executor.broadcast_from_output_rank(payload)
|
||
if not synchronized[0]:
|
||
propagated = RuntimeError(f"External-launcher {description} failed on output rank: {synchronized[1]}")
|
||
if source_error is not None:
|
||
raise propagated from source_error
|
||
raise propagated
|
||
return cast(_T, synchronized[1])
|
||
|
||
def _synchronize_request_prompt(self, request: GenerationRequest) -> GenerationRequest:
|
||
"""Use rank zero's prompt selection for an offline SPMD request."""
|
||
if not self.executor.uses_spmd_execution:
|
||
return request
|
||
synchronized = deepcopy(request)
|
||
prompt_and_path = self.executor.broadcast_from_output_rank((
|
||
request.prompt, request.inputs.prompt_path) if self.executor.is_output_rank else None)
|
||
synchronized.prompt, synchronized.inputs.prompt_path = prompt_and_path
|
||
return synchronized
|
||
|
||
def generate(
|
||
self,
|
||
request: GenerationRequest | Mapping[str, Any],
|
||
*,
|
||
log_queue=None,
|
||
) -> GenerationResult | list[GenerationResult]:
|
||
"""
|
||
Generate video or image outputs from a typed inference request.
|
||
|
||
Args:
|
||
request: A `GenerationRequest` instance or a mapping that can be
|
||
parsed into one. This is the primary public inference
|
||
entrypoint for the typed API.
|
||
log_queue: Optional multiprocessing.Queue to forward worker logs to
|
||
during this request.
|
||
|
||
Returns:
|
||
A `GenerationResult` for single-request generation, or a list of
|
||
`GenerationResult` objects when the request expands into multiple
|
||
prompts.
|
||
"""
|
||
normalized_request = self._synchronize_request_prompt(normalize_generation_request(request))
|
||
if log_queue:
|
||
self.executor.set_log_queue(log_queue)
|
||
|
||
try:
|
||
return self._generate_request_impl(normalized_request)
|
||
finally:
|
||
if log_queue:
|
||
self.executor.clear_log_queue()
|
||
|
||
async def generate_async(
|
||
self,
|
||
request: GenerationRequest | Mapping[str, Any],
|
||
*,
|
||
log_queue=None,
|
||
):
|
||
"""Async generation that yields typed :class:`VideoEvent`s.
|
||
|
||
Three consumers share this substrate:
|
||
|
||
* Streaming server (:mod:`fastvideo.entrypoints.streaming`) —
|
||
pipes :class:`VideoPartialEvent` frames into fMP4.
|
||
* Stateless OpenAI server — ignores progress events, forwards
|
||
:class:`VideoFinalEvent` as the HTTP response body.
|
||
* Dynamo native backend
|
||
(``components/src/dynamo/fastvideo/``) — wraps each event as
|
||
an ``NvVideosResponse`` chunk.
|
||
|
||
The aggregated code path shipped here yields a single
|
||
:class:`VideoProgressEvent` at start and one
|
||
:class:`VideoFinalEvent` at end. Future work will thread
|
||
per-step progress events through the pipeline's denoise loop
|
||
so streaming consumers don't have to wait on a materialized
|
||
final.
|
||
"""
|
||
import asyncio
|
||
|
||
normalized = self._synchronize_request_prompt(normalize_generation_request(request))
|
||
total_steps = max(1, normalized.sampling.num_inference_steps)
|
||
yield VideoProgressEvent(step=0, total_steps=total_steps, stage="denoise")
|
||
|
||
if log_queue:
|
||
self.executor.set_log_queue(log_queue)
|
||
try:
|
||
result = await asyncio.to_thread(self._generate_request_impl, normalized)
|
||
finally:
|
||
if log_queue:
|
||
self.executor.clear_log_queue()
|
||
|
||
if isinstance(result, list):
|
||
# Prompt-batch expansion — emit one Final per sub-result.
|
||
for sub in result:
|
||
yield await asyncio.to_thread(_final_event_from_result, sub)
|
||
return
|
||
yield await asyncio.to_thread(_final_event_from_result, result)
|
||
|
||
@staticmethod
|
||
def default_health_check_request() -> GenerationRequest:
|
||
"""Return the minimal typed request Dynamo uses for probes.
|
||
|
||
256x256, 8 frames, 1 inference step -- fast enough to be a
|
||
viable liveness check, non-trivial enough to exercise the
|
||
DiT -> VAE -> decode path. Consumers adapt this shape to their
|
||
transport's health-check payload (see
|
||
``docs/design/server_contracts/dynamo.md``).
|
||
"""
|
||
return GenerationRequest(
|
||
prompt="health check",
|
||
inputs=InputConfig(),
|
||
sampling=SamplingConfig(
|
||
num_frames=8,
|
||
height=256,
|
||
width=256,
|
||
fps=24,
|
||
num_inference_steps=1,
|
||
guidance_scale=1.0,
|
||
),
|
||
output=OutputConfig(save_video=False, return_frames=False),
|
||
)
|
||
|
||
def generate_video(
|
||
self,
|
||
prompt: str | None = None,
|
||
sampling_param: SamplingParam | None = None,
|
||
# Action control inputs (Matrix-Game)
|
||
mouse_cond: torch.Tensor | None = None,
|
||
keyboard_cond: torch.Tensor | None = None,
|
||
grid_sizes: tuple[int, int, int] | list[int] | torch.Tensor
|
||
| None = None,
|
||
**kwargs,
|
||
) -> dict[str, Any] | list[dict[str, Any]]:
|
||
"""
|
||
Generate a video based on the given prompt.
|
||
|
||
Args:
|
||
prompt: The prompt to use for generation (optional if prompt_txt is provided)
|
||
negative_prompt: The negative prompt to use (overrides the one in fastvideo_args)
|
||
output_path: Path to save the video (overrides the one in fastvideo_args)
|
||
prompt_path: Path to prompt file
|
||
save_video: Whether to save the video to disk
|
||
return_frames: Whether to include raw frames in the result dict
|
||
num_inference_steps: Number of denoising steps (overrides fastvideo_args)
|
||
guidance_scale: Classifier-free guidance scale (overrides fastvideo_args)
|
||
num_frames: Number of frames to generate (overrides fastvideo_args)
|
||
height: Height of generated video (overrides fastvideo_args)
|
||
width: Width of generated video (overrides fastvideo_args)
|
||
fps: Frames per second for saved video (overrides fastvideo_args)
|
||
seed: Random seed for generation (overrides fastvideo_args)
|
||
callback: Callback function called after each step
|
||
callback_steps: Number of steps between each callback
|
||
|
||
Returns:
|
||
A metadata dictionary for single-prompt generation, or a list of
|
||
metadata dictionaries for prompt-file batch generation.
|
||
"""
|
||
log_queue = kwargs.pop("log_queue", None)
|
||
warnings.warn(
|
||
"VideoGenerator.generate_video(...) is deprecated; use "
|
||
"VideoGenerator.generate(request=...) instead.",
|
||
DeprecationWarning,
|
||
stacklevel=2,
|
||
)
|
||
if log_queue:
|
||
self.executor.set_log_queue(log_queue)
|
||
|
||
try:
|
||
extra_overrides: dict[str, Any] = {}
|
||
for _ek in _BATCH_EXTRA_PASSTHROUGH_KEYS:
|
||
if _ek in kwargs:
|
||
extra_overrides[_ek] = kwargs.pop(_ek)
|
||
|
||
request = legacy_generate_call_to_request(
|
||
prompt,
|
||
sampling_param,
|
||
mouse_cond=mouse_cond,
|
||
keyboard_cond=keyboard_cond,
|
||
grid_sizes=grid_sizes,
|
||
legacy_kwargs=kwargs,
|
||
)
|
||
|
||
fastvideo_args = self.fastvideo_args
|
||
pipeline_overrides = request_to_pipeline_overrides(request)
|
||
if pipeline_overrides:
|
||
fastvideo_args = deepcopy(self.fastvideo_args)
|
||
for key, value in pipeline_overrides.items():
|
||
if not hasattr(fastvideo_args.pipeline_config, key):
|
||
raise ValueError(f"Request field {key!r} is not supported by pipeline config overrides")
|
||
setattr(fastvideo_args.pipeline_config, key, deepcopy(value))
|
||
|
||
resolved_sampling_param = fastvideo_args.pipeline_config.apply_request_constraints(
|
||
request,
|
||
request_to_sampling_param(request, model_path=self.fastvideo_args.model_path),
|
||
)
|
||
return self._generate_video_impl(
|
||
prompt=request.prompt,
|
||
sampling_param=resolved_sampling_param,
|
||
fastvideo_args=fastvideo_args,
|
||
**extra_overrides,
|
||
)
|
||
finally:
|
||
if log_queue:
|
||
self.executor.clear_log_queue()
|
||
|
||
def _generate_request_impl(
|
||
self,
|
||
request: GenerationRequest,
|
||
) -> GenerationResult | list[GenerationResult]:
|
||
_validate_request_stage_overrides(self.fastvideo_args.model_path, request)
|
||
if isinstance(request.prompt, list):
|
||
if request.inputs.prompt_path is not None:
|
||
raise ValueError("request.prompt list cannot be combined with request.inputs.prompt_path")
|
||
results: list[GenerationResult] = []
|
||
for index, single_request in enumerate(expand_request_prompt_batch(request)):
|
||
prompt = single_request.prompt
|
||
wrapped = self._generate_single_request(single_request)
|
||
if isinstance(wrapped, list):
|
||
results.extend(wrapped)
|
||
continue
|
||
wrapped.prompt_index = index
|
||
if wrapped.prompt is None:
|
||
wrapped.prompt = prompt
|
||
results.append(wrapped)
|
||
return results
|
||
|
||
return self._generate_single_request(request)
|
||
|
||
def _generate_single_request(
|
||
self,
|
||
request: GenerationRequest,
|
||
) -> GenerationResult | list[GenerationResult]:
|
||
fastvideo_args = self.fastvideo_args
|
||
pipeline_overrides = request_to_pipeline_overrides(request)
|
||
if pipeline_overrides:
|
||
fastvideo_args = deepcopy(self.fastvideo_args)
|
||
for key, value in pipeline_overrides.items():
|
||
if not hasattr(fastvideo_args.pipeline_config, key):
|
||
raise ValueError(f"Request field {key!r} is not supported by pipeline config overrides")
|
||
setattr(fastvideo_args.pipeline_config, key, deepcopy(value))
|
||
|
||
sampling_param = fastvideo_args.pipeline_config.apply_request_constraints(
|
||
request,
|
||
request_to_sampling_param(request, model_path=self.fastvideo_args.model_path),
|
||
)
|
||
batch_extra = request_to_batch_extra(request)
|
||
result = self._generate_video_impl(
|
||
prompt=request.prompt,
|
||
sampling_param=sampling_param,
|
||
fastvideo_args=fastvideo_args,
|
||
**batch_extra,
|
||
)
|
||
return self._wrap_legacy_result(result)
|
||
|
||
def _generate_video_impl(
|
||
self,
|
||
prompt: str | None = None,
|
||
sampling_param: SamplingParam | None = None,
|
||
mouse_cond: torch.Tensor | None = None,
|
||
keyboard_cond: torch.Tensor | None = None,
|
||
grid_sizes: tuple[int, int, int] | list[int] | torch.Tensor
|
||
| None = None,
|
||
fastvideo_args: FastVideoArgs | None = None,
|
||
**kwargs,
|
||
) -> dict[str, Any] | list[np.ndarray] | list[dict[str, Any]]:
|
||
"""Internal implementation of generate_video."""
|
||
if fastvideo_args is None:
|
||
fastvideo_args = self.fastvideo_args
|
||
|
||
validate_request_batch_extra(kwargs, model_path=fastvideo_args.model_path)
|
||
|
||
# Handle batch processing from text file
|
||
if sampling_param is None:
|
||
sampling_param = SamplingParam.from_pretrained(fastvideo_args.model_path)
|
||
|
||
if self.executor.uses_spmd_execution:
|
||
prompt = self.executor.broadcast_from_output_rank(prompt if self.executor.is_output_rank else None)
|
||
|
||
# Add action control inputs to kwargs if provided
|
||
if mouse_cond is not None:
|
||
kwargs['mouse_cond'] = mouse_cond
|
||
if keyboard_cond is not None:
|
||
kwargs['keyboard_cond'] = keyboard_cond
|
||
if grid_sizes is not None:
|
||
kwargs['grid_sizes'] = grid_sizes
|
||
|
||
extra_overrides: dict[str, Any] = {}
|
||
for _ek in _BATCH_EXTRA_PASSTHROUGH_KEYS:
|
||
if _ek in kwargs:
|
||
extra_overrides[_ek] = kwargs.pop(_ek)
|
||
|
||
prompt_embeds = kwargs.pop("prompt_embeds", None)
|
||
sampling_param.update(kwargs)
|
||
kwargs["_extra_overrides"] = extra_overrides
|
||
|
||
prompt_txt_path = sampling_param.prompt_path or fastvideo_args.prompt_txt
|
||
if self.executor.uses_spmd_execution:
|
||
prompt_txt_path = self.executor.broadcast_from_output_rank(
|
||
prompt_txt_path if self.executor.is_output_rank else None)
|
||
|
||
if prompt_txt_path is not None:
|
||
|
||
def _load_prompt_file() -> tuple[str, list[str]]:
|
||
if not prompt_txt_path or not os.path.exists(prompt_txt_path):
|
||
raise FileNotFoundError(f"Prompt text file not found: {prompt_txt_path}")
|
||
with open(prompt_txt_path, encoding='utf-8') as f:
|
||
prompts = [line.strip() for line in f if line.strip()]
|
||
if not prompts:
|
||
raise ValueError(f"No prompts found in file: {prompt_txt_path}")
|
||
return prompt_txt_path, prompts
|
||
|
||
prompt_txt_path, prompts = self._run_on_output_rank(
|
||
_load_prompt_file,
|
||
description="prompt-file setup",
|
||
)
|
||
|
||
logger.info("Found %d prompts in %s", len(prompts), prompt_txt_path)
|
||
|
||
results = []
|
||
for i, batch_prompt in enumerate(prompts):
|
||
logger.info("Processing prompt %d/%d: %s...", i + 1, len(prompts), batch_prompt[:100])
|
||
try:
|
||
# Generate video for this prompt using the same logic below
|
||
output_path = self._run_on_output_rank(
|
||
partial(self._prepare_output_path, sampling_param.output_path, batch_prompt),
|
||
description="output-path setup",
|
||
)
|
||
kwargs["output_path"] = output_path
|
||
result = self._generate_single_video(
|
||
prompt=batch_prompt,
|
||
sampling_param=sampling_param,
|
||
fastvideo_args=fastvideo_args,
|
||
**kwargs,
|
||
)
|
||
|
||
# Add prompt info to result
|
||
result["prompt_index"] = i
|
||
result["prompt"] = batch_prompt
|
||
|
||
results.append(result)
|
||
logger.info("Successfully generated video for prompt %d", i + 1)
|
||
|
||
except Exception as e:
|
||
logger.error("Failed to generate video for prompt %d: %s", i + 1, e)
|
||
if self.executor.uses_spmd_execution:
|
||
raise
|
||
continue
|
||
|
||
logger.info("Completed batch processing. Generated %d videos successfully.", len(results))
|
||
return results
|
||
|
||
# Single prompt generation (original behavior)
|
||
output_name_hint: str | None = None
|
||
if prompt is None:
|
||
if fastvideo_args.workload_type is WorkloadType.V2A:
|
||
# Video semantics are sufficient conditioning for V2A models;
|
||
# model-specific text stages interpret the empty string using
|
||
# their native tokenizer/empty-prompt contract.
|
||
prompt = ""
|
||
elif getattr(fastvideo_args.pipeline_config, "prompt_optional", False):
|
||
# Video-to-video pipelines without a text tower (Kandinsky6 SR)
|
||
# are conditioned by the input video alone; name the output
|
||
# after that video instead of the (empty) prompt.
|
||
prompt = ""
|
||
if isinstance(sampling_param.video_path, str) and sampling_param.video_path:
|
||
output_name_hint = os.path.splitext(os.path.basename(sampling_param.video_path))[0]
|
||
else:
|
||
raise ValueError("Either prompt or prompt_txt must be provided")
|
||
output_path = self._run_on_output_rank(
|
||
lambda: self._prepare_output_path(sampling_param.output_path, output_name_hint
|
||
if output_name_hint else prompt),
|
||
description="output-path setup",
|
||
)
|
||
kwargs["output_path"] = output_path
|
||
if prompt_embeds is not None:
|
||
kwargs["prompt_embeds"] = prompt_embeds
|
||
return self._generate_single_video(
|
||
prompt=prompt,
|
||
sampling_param=sampling_param,
|
||
fastvideo_args=fastvideo_args,
|
||
**kwargs,
|
||
)
|
||
|
||
def _is_image_workload(self) -> bool:
|
||
"""Return True when the workload produces a single image (t2i, i2i …)."""
|
||
args = getattr(self, "fastvideo_args", None)
|
||
if args is None:
|
||
return False
|
||
return args.workload_type.value.endswith("2i")
|
||
|
||
def _is_audio_workload(self) -> bool:
|
||
"""Return True when the workload produces standalone audio."""
|
||
args = getattr(self, "fastvideo_args", None)
|
||
if args is None:
|
||
return False
|
||
return args.workload_type.value.endswith("2a")
|
||
|
||
def _prepare_output_path(
|
||
self,
|
||
output_path: str,
|
||
prompt: str,
|
||
) -> str:
|
||
"""Build a unique, sanitized output file path.
|
||
|
||
Image workloads preserve explicit PNG, JPEG, or WebP extensions and
|
||
otherwise default to ``.png``. Video workloads use ``.mp4``.
|
||
|
||
- If ``output_path`` already carries the correct extension, treat it
|
||
as a file path.
|
||
- Otherwise, treat ``output_path`` as a directory and derive the
|
||
filename from the prompt.
|
||
- Invalid filename characters are removed; if the name changes, a
|
||
warning is logged.
|
||
- If the target path already exists, a numeric suffix is appended.
|
||
"""
|
||
base_path, extension = os.path.splitext(output_path)
|
||
extension_lower = extension.lower()
|
||
is_directory = os.path.isdir(output_path)
|
||
if self._is_image_workload():
|
||
target_ext = extension_lower if not is_directory and extension_lower in {".png", ".jpg", ".jpeg", ".webp"
|
||
} else ".png"
|
||
elif self._is_audio_workload():
|
||
target_ext = ".wav"
|
||
else:
|
||
target_ext = ".mp4"
|
||
|
||
def _sanitize_filename_component(name: str) -> str:
|
||
# Remove characters invalid on common filesystems, strip spaces/dots
|
||
sanitized = re.sub(r'[\\/:*?"<>|]', '', name)
|
||
sanitized = sanitized.strip().strip('.')
|
||
sanitized = re.sub(r'\s+', ' ', sanitized)
|
||
return sanitized or "output"
|
||
|
||
if extension_lower == target_ext and not is_directory:
|
||
output_dir = os.path.dirname(output_path)
|
||
base_name = os.path.basename(base_path) # filename without extension
|
||
sanitized_base = _sanitize_filename_component(base_name)
|
||
if sanitized_base != base_name:
|
||
logger.warning(
|
||
"The output name '%s' contained invalid characters. "
|
||
"It has been renamed to '%s%s'",
|
||
os.path.basename(output_path),
|
||
sanitized_base,
|
||
target_ext,
|
||
)
|
||
out_name = f"{sanitized_base}{target_ext}"
|
||
else:
|
||
# Treat as directory; inform if an unexpected extension was
|
||
# provided.
|
||
if extension:
|
||
logger.info(
|
||
"Output path '%s' has extension '%s' which does not "
|
||
"match the target '%s'; treating it as a directory",
|
||
output_path,
|
||
extension,
|
||
target_ext,
|
||
)
|
||
output_dir = output_path
|
||
prompt_component = _sanitize_filename_component(prompt[:100])
|
||
out_name = f"{prompt_component}{target_ext}"
|
||
|
||
if output_dir:
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
new_output_path = os.path.join(output_dir, out_name)
|
||
counter = 1
|
||
while os.path.exists(new_output_path):
|
||
name_part, ext_part = os.path.splitext(out_name)
|
||
new_name = f"{name_part}_{counter}{ext_part}"
|
||
new_output_path = os.path.join(output_dir, new_name)
|
||
counter += 1
|
||
return new_output_path
|
||
|
||
def _generate_single_video(
|
||
self,
|
||
prompt: str,
|
||
sampling_param: SamplingParam | None = None,
|
||
fastvideo_args: FastVideoArgs | None = None,
|
||
**kwargs,
|
||
) -> dict[str, Any]:
|
||
"""Internal method for single video generation"""
|
||
# Generation boundary: drop any stale cancellation left over from a
|
||
# previous run so it can never silently skip this one, and open the
|
||
# window in which executor.interrupt() is honored.
|
||
begin_generation = getattr(self.executor, "begin_generation", None)
|
||
if begin_generation is not None:
|
||
begin_generation()
|
||
if fastvideo_args is None:
|
||
fastvideo_args = self.fastvideo_args
|
||
|
||
# Validate inputs
|
||
if not isinstance(prompt, str):
|
||
raise TypeError(f"`prompt` must be a string, but got {type(prompt)}")
|
||
prompt = prompt.strip()
|
||
sampling_param = deepcopy(sampling_param)
|
||
output_path = kwargs["output_path"]
|
||
prompt_embeds = kwargs.get("prompt_embeds")
|
||
sampling_param.prompt = prompt
|
||
# Process negative prompt
|
||
if sampling_param.negative_prompt is not None:
|
||
sampling_param.negative_prompt = sampling_param.negative_prompt.strip()
|
||
|
||
# Validate dimensions
|
||
if (sampling_param.height <= 0 or sampling_param.width <= 0 or sampling_param.num_frames <= 0):
|
||
raise ValueError(f"Height, width, and num_frames must be positive integers, got "
|
||
f"height={sampling_param.height}, width={sampling_param.width}, "
|
||
f"num_frames={sampling_param.num_frames}")
|
||
|
||
# Calculate sizes
|
||
target_height = align_to(sampling_param.height, 16)
|
||
target_width = align_to(sampling_param.width, 16)
|
||
|
||
# Calculate latent sizes
|
||
latents_size = [(sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.width // 8]
|
||
n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
|
||
|
||
# Log parameters
|
||
debug_str = f"""
|
||
height: {target_height}
|
||
width: {target_width}
|
||
video_length: {sampling_param.num_frames}
|
||
prompt: {sampling_param.prompt}
|
||
image_path: {sampling_param.image_path}
|
||
neg_prompt: {sampling_param.negative_prompt}
|
||
seed: {sampling_param.seed}
|
||
infer_steps: {sampling_param.num_inference_steps}
|
||
num_videos_per_prompt: {sampling_param.num_videos_per_prompt}
|
||
guidance_scale: {sampling_param.guidance_scale}
|
||
n_tokens: {n_tokens}
|
||
flow_shift: {fastvideo_args.pipeline_config.flow_shift}
|
||
embedded_guidance_scale: {fastvideo_args.pipeline_config.embedded_cfg_scale}
|
||
save_video: {sampling_param.save_video}
|
||
output_path: {output_path}
|
||
""" # type: ignore[attr-defined]
|
||
logger.info(debug_str)
|
||
|
||
# Prepare batch
|
||
batch = ForwardBatch(
|
||
**shallow_asdict(sampling_param),
|
||
eta=0.0,
|
||
n_tokens=n_tokens,
|
||
VSA_sparsity=fastvideo_args.VSA_sparsity,
|
||
)
|
||
# Allow precomputed prompt_embeds (e.g. from diffusers) to skip text encoding
|
||
if prompt_embeds is not None:
|
||
batch.prompt_embeds = (list(prompt_embeds) if isinstance(prompt_embeds, list | tuple) else [prompt_embeds])
|
||
|
||
extra_overrides = kwargs.pop("_extra_overrides", {})
|
||
for _ek, _ev in extra_overrides.items():
|
||
batch.extra[_ek] = _ev
|
||
|
||
# External-launcher mode is SPMD: every rank runs the same pipeline
|
||
# collectives, while only world rank 0 owns user-facing outputs. Keep
|
||
# non-output ranks in the forward pass but skip their decoded-tensor
|
||
# host copy and disk write.
|
||
if not self.executor.is_output_rank:
|
||
batch.save_video = False
|
||
batch.return_frames = False
|
||
batch.return_samples = False
|
||
batch.return_trajectory_latents = False
|
||
batch.return_trajectory_decoded = False
|
||
batch.return_continuation_state = False
|
||
|
||
# Run inference
|
||
start_time = time.perf_counter()
|
||
|
||
# Execute forward pass in a new thread for non-blocking tensor
|
||
# allocation. Capture thread exceptions so we can surface the true
|
||
# failure in the main thread instead of later hitting None outputs.
|
||
result_container = {"output_batch": ForwardBatch(data_type=batch.data_type)}
|
||
thread_error: dict[str, BaseException | None] = {"error": None}
|
||
thread_error_traceback: dict[str, str] = {"traceback": ""}
|
||
|
||
def execute_forward_thread():
|
||
import traceback
|
||
try:
|
||
result_container["output_batch"] = self.executor.execute_forward(batch, fastvideo_args)
|
||
except BaseException as error: # noqa: BLE001
|
||
thread_error["error"] = error
|
||
thread_error_traceback["traceback"] = traceback.format_exc()
|
||
|
||
thread = threading.Thread(target=execute_forward_thread)
|
||
thread.start()
|
||
latent_batch_size = _infer_latent_batch_size(batch)
|
||
is_latent_output = fastvideo_args.output_type == "latent"
|
||
needs_frame_output = batch.return_frames or (batch.save_video and not is_latent_output)
|
||
# A populated ``samples`` has exactly one consumer — the result
|
||
# dict (``"samples": samples if needs_samples_out else None``).
|
||
# Post-decode frame building reads ``output_batch.output``
|
||
# directly (the GPU ``vid_u8`` path), not ``samples``. So when
|
||
# neither frames nor samples are requested, the pinned fp32 allocation
|
||
# and D->H copy are dead weight — the CLI generate flow (``save_video=True``,
|
||
# ``return_frames=False``) hits this on every call.
|
||
# ``output_type == "latent"`` keeps its existing branch (shape
|
||
# mismatch falls through to ``.cpu()`` below) for callers that
|
||
# *do* ask for the latent samples via ``return_frames=True``.
|
||
# ``skip_pixel_prealloc`` also gates the slow-path warning.
|
||
needs_samples_out = batch.return_frames or batch.return_samples
|
||
# ``output_shape_from_input`` pipelines (Kandinsky6 SR) decide the output geometry from the input video, so the
|
||
# request's height / width / num_frames say nothing about the buffer to pre-allocate.
|
||
output_shape_from_input = bool(getattr(fastvideo_args.pipeline_config, "output_shape_from_input", False))
|
||
skip_pixel_prealloc = is_latent_output or not needs_samples_out or output_shape_from_input
|
||
if skip_pixel_prealloc:
|
||
samples = torch.empty(0, device='cpu')
|
||
else:
|
||
samples = allocate_cpu_tensor_with_pin_fallback(
|
||
(latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width),
|
||
pin_memory=fastvideo_args.pin_cpu_memory)
|
||
try:
|
||
thread.join()
|
||
except BaseException: # noqa: BLE001
|
||
# Ctrl-C lands in the main thread while the in-process forward pass
|
||
# keeps running in `execute_forward_thread`; flag the pipeline so the
|
||
# denoise loop stops instead of running to completion. No-op for
|
||
# executors without an interrupt hook (e.g. MultiprocExecutor, whose
|
||
# worker processes already receive SIGINT directly).
|
||
try:
|
||
interrupt = getattr(self.executor, "interrupt", None)
|
||
if interrupt is not None:
|
||
interrupt()
|
||
except Exception:
|
||
logger.debug("executor.interrupt() failed while aborting forward thread", exc_info=True)
|
||
raise
|
||
|
||
if thread_error["error"] is not None:
|
||
raise RuntimeError("Forward execution thread failed.\n"
|
||
f"{thread_error_traceback['traceback']}") from thread_error["error"]
|
||
|
||
output_batch = result_container["output_batch"]
|
||
if output_batch.output is None:
|
||
raise RuntimeError("Forward execution returned no output tensor. "
|
||
"This usually means the executor/pipeline failed earlier.")
|
||
|
||
audio_only = bool(output_batch.extra.get("audio_only"))
|
||
if not needs_samples_out:
|
||
# Nothing downstream reads ``samples`` (the result dict
|
||
# returns None when neither output is requested); keep the empty
|
||
# placeholder allocated above and skip the fp32 D->H copy
|
||
# entirely.
|
||
pass
|
||
elif audio_only:
|
||
# Audio-only return-frames requests expose the small placeholder
|
||
# without warning about expected video shape.
|
||
samples = output_batch.output.cpu()
|
||
elif output_batch.output.shape == samples.shape:
|
||
samples.copy_(output_batch.output)
|
||
if output_batch.output.dtype == torch.uint8:
|
||
# A worker may hand back uint8 pixels (the MiniMax-H3 decode
|
||
# stage does); keep ``samples`` on its [0, 1] float contract.
|
||
samples.div_(255)
|
||
else:
|
||
if not skip_pixel_prealloc:
|
||
logger.warning("Output shape %s does not match expected shape %s; use slow path",
|
||
output_batch.output.shape, samples.shape)
|
||
samples = output_batch.output.cpu()
|
||
if samples.dtype == torch.uint8:
|
||
samples = samples.float().div_(255)
|
||
logging_info = output_batch.logging_info
|
||
|
||
gen_time = time.perf_counter() - start_time
|
||
logger.info("Generated successfully in %.2f seconds", gen_time)
|
||
|
||
# Three mutually-exclusive output modes determine whether (a) we
|
||
# build an RGB frame buffer and (b) what file we write to disk:
|
||
#
|
||
# 1. `output_type == "latent"` — VAE is bypassed in DecodingStage
|
||
# and `samples` holds raw latents (arbitrary channel count).
|
||
# The RGB grid / uint8 / mp4 / png pipeline below cannot
|
||
# consume those, so we skip it entirely and let callers work
|
||
# with the latent tensor directly via `result["samples"]`.
|
||
# 2. Audio-only workload — `samples` is a 1×3×1×8×8 placeholder
|
||
# no caller will use; skip the grid loop and save a `.wav`.
|
||
# 3. Pixel video / image — the historical happy path.
|
||
# `GenerationResult.size` describes the produced media, not only the
|
||
# base-stage request. Refiner pipelines can change the final pixel
|
||
# dimensions, so derive this result metadata from the decoded output.
|
||
# Read the geometry from `output_batch.output` (a shape-only access,
|
||
# no D->H copy): when `return_frames=False` the `samples` mirror
|
||
# stays an empty placeholder and no longer carries the decoded
|
||
# shape. Metadata-only calls keep the request fallback and never
|
||
# inspect the (possibly dropped) worker output.
|
||
output_size = _resolve_output_size(
|
||
output_batch.output if needs_frame_output else samples,
|
||
(target_height, target_width, batch.num_frames),
|
||
pixel_output=not is_latent_output and not audio_only,
|
||
)
|
||
|
||
postprocess_start = time.perf_counter()
|
||
frames: list[np.ndarray] | None
|
||
if is_latent_output or audio_only:
|
||
frames = [] if audio_only and batch.return_frames else None
|
||
elif not needs_frame_output:
|
||
frames = None
|
||
else:
|
||
# `samples` above is just the pinned-CPU mirror of `output_batch.output` (`samples.copy_(output)` or
|
||
# `output.cpu()`) with no intervening preprocessing, so reading `output_batch.output` here is the same
|
||
# data. See `_quantize_video_frames_to_uint8` for why this quantizes on-device and its OOM fallback.
|
||
vid_u8 = _quantize_video_frames_to_uint8(output_batch.output)
|
||
frames = [
|
||
torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy() for x in vid_u8
|
||
]
|
||
postprocess_time = time.perf_counter() - postprocess_start
|
||
logger.info("PostDecodeFrameProcessStage completed in %.3f s", postprocess_time)
|
||
if logging_info is not None:
|
||
logging_info.add_stage_execution_time("PostDecodeFrameProcessStage", postprocess_time)
|
||
|
||
save_to_disk = batch.save_video and not is_latent_output
|
||
save_video_time = 0.0
|
||
audio_mux_time = 0.0
|
||
if save_to_disk:
|
||
if audio_only:
|
||
# Audio-only workload: write a standalone .wav rather than
|
||
# muxing the audio into a placeholder mp4 (which forces
|
||
# ffmpeg to round 8x8 placeholder frames up to 16x16).
|
||
output_path = self._rewrite_extension(output_path, ".wav")
|
||
save_start = time.perf_counter()
|
||
self._write_pcm_wav(
|
||
output_path,
|
||
output_batch.extra["audio"],
|
||
int(output_batch.extra["audio_sample_rate"]),
|
||
)
|
||
save_video_time = time.perf_counter() - save_start
|
||
logger.info("Saved audio to %s", output_path)
|
||
elif self._is_image_workload():
|
||
# Image workloads (t2i, i2i, …): save the first frame as PNG.
|
||
assert frames is not None # implied by save_to_disk and not audio_only
|
||
save_start = time.perf_counter()
|
||
imageio.imwrite(output_path, frames[0])
|
||
save_video_time = time.perf_counter() - save_start
|
||
logger.info("Saved image to %s", output_path)
|
||
else:
|
||
assert frames is not None # implied by save_to_disk and not audio_only
|
||
# ``batch`` is the request-side batch. A pipeline that derives the frame rate from its input (Kandinsky6
|
||
# SR follows the source clip) mutates a worker-side copy, so it reports the rate via ``extra``.
|
||
save_fps = int(output_batch.extra.get("output_fps") or batch.fps)
|
||
audio = output_batch.extra.get("audio")
|
||
audio_sample_rate = output_batch.extra.get("audio_sample_rate")
|
||
if audio is not None and audio_sample_rate is not None:
|
||
# Single-pass save path: encode video+audio once, avoiding
|
||
# second-pass remux overhead. AudioMuxStage remains 0.0 on
|
||
# success because audio is already present in the saved MP4.
|
||
save_start = time.perf_counter()
|
||
save_ok = self._save_video_with_audio_ffmpeg_pipe(
|
||
output_path=output_path,
|
||
frames=frames,
|
||
fps=save_fps,
|
||
audio=audio,
|
||
sample_rate=int(audio_sample_rate),
|
||
)
|
||
if not save_ok:
|
||
logger.warning("ffmpeg pipe save failed; trying PyAV single-pass save.")
|
||
save_ok = self._save_video_with_audio_single_pass(
|
||
output_path=output_path,
|
||
frames=frames,
|
||
fps=save_fps,
|
||
audio=audio,
|
||
sample_rate=int(audio_sample_rate),
|
||
)
|
||
save_video_time = time.perf_counter() - save_start
|
||
if save_ok:
|
||
audio_mux_time = 0.0
|
||
else:
|
||
logger.warning("Single-pass save failed; falling back to two-step save/mux.")
|
||
save_start = time.perf_counter()
|
||
imageio.mimsave(output_path, frames, fps=save_fps, format="mp4")
|
||
save_video_time = time.perf_counter() - save_start
|
||
mux_start = time.perf_counter()
|
||
mux_ok = self._mux_audio(output_path, audio, int(audio_sample_rate))
|
||
audio_mux_time = time.perf_counter() - mux_start
|
||
if not mux_ok:
|
||
logger.warning("Audio mux failed; saved video without audio.")
|
||
else:
|
||
save_start = time.perf_counter()
|
||
imageio.mimsave(output_path, frames, fps=save_fps, format="mp4")
|
||
save_video_time = time.perf_counter() - save_start
|
||
audio_mux_time = 0.0
|
||
logger.info("Saved video to %s", output_path)
|
||
|
||
logger.info("VideoSaveStage completed in %.3f s", save_video_time)
|
||
if logging_info is not None:
|
||
logging_info.add_stage_execution_time("VideoSaveStage", save_video_time)
|
||
logger.info("AudioMuxStage completed in %.3f s", audio_mux_time)
|
||
if logging_info is not None:
|
||
logging_info.add_stage_execution_time("AudioMuxStage", audio_mux_time)
|
||
|
||
e2e_time = time.perf_counter() - start_time
|
||
logger.info("End-to-end latency: %.2f seconds", e2e_time)
|
||
|
||
result: dict[str, Any] = {
|
||
"prompts": prompt,
|
||
"samples": samples if needs_samples_out else None,
|
||
"frames": frames if batch.return_frames else None,
|
||
# Audio is the primary output for audio workloads — return it
|
||
# whenever the pipeline produced one, regardless of
|
||
# `return_frames` (which gates the video-shaped buffers).
|
||
"audio": output_batch.extra.get("audio") if self.executor.is_output_rank else None,
|
||
"audio_sample_rate": output_batch.extra.get("audio_sample_rate") if self.executor.is_output_rank else None,
|
||
"ltx2_audio_latents":
|
||
output_batch.extra.get("ltx2_audio_latents") if self.executor.is_output_rank else None,
|
||
"size": output_size,
|
||
"generation_time": gen_time,
|
||
"e2e_latency": e2e_time,
|
||
"logging_info": logging_info,
|
||
"trajectory": output_batch.trajectory_latents if self.executor.is_output_rank else None,
|
||
"trajectory_timesteps": output_batch.trajectory_timesteps if self.executor.is_output_rank else None,
|
||
"trajectory_decoded": output_batch.trajectory_decoded if self.executor.is_output_rank else None,
|
||
"video_path": output_path if save_to_disk else None,
|
||
"peak_memory_mb": output_batch.extra.get("peak_memory_mb"),
|
||
}
|
||
|
||
return result
|
||
|
||
@staticmethod
|
||
def _wrap_legacy_result(
|
||
result: dict[str, Any] | list[dict[str, Any]], ) -> GenerationResult | list[GenerationResult]:
|
||
if isinstance(result, list):
|
||
return [GenerationResult.from_legacy_result(item) for item in result]
|
||
return GenerationResult.from_legacy_result(result)
|
||
|
||
@staticmethod
|
||
def _unwrap_typed_result(
|
||
result: GenerationResult | list[GenerationResult], ) -> dict[str, Any] | list[dict[str, Any]]:
|
||
if isinstance(result, list):
|
||
return [item.to_legacy_dict() for item in result]
|
||
return result.to_legacy_dict()
|
||
|
||
@staticmethod
|
||
def _rewrite_extension(path: str, new_ext: str) -> str:
|
||
root, old_ext = os.path.splitext(path)
|
||
new_path = root + new_ext
|
||
if old_ext and old_ext.lower() != new_ext.lower():
|
||
logger.info("Rewriting output extension %s -> %s.", old_ext, new_ext)
|
||
return new_path
|
||
|
||
@staticmethod
|
||
def _audio_to_int16(audio: torch.Tensor | np.ndarray, ) -> tuple[np.ndarray, int]:
|
||
"""Normalize `[samples]` / `[samples, channels]` / `[channels,
|
||
samples]` audio in roughly [-1, 1] to a `(int16 [samples,
|
||
channels], num_channels)` pair. Raises `ValueError` for shapes
|
||
we can't classify.
|
||
"""
|
||
if torch.is_tensor(audio):
|
||
audio_np = audio.detach().cpu().float().numpy()
|
||
else:
|
||
audio_np = np.asarray(audio, dtype=np.float32)
|
||
if audio_np.ndim == 1:
|
||
audio_np = audio_np[:, None]
|
||
elif audio_np.ndim == 2:
|
||
if audio_np.shape[0] <= 8 and audio_np.shape[1] > audio_np.shape[0]:
|
||
audio_np = audio_np.T
|
||
else:
|
||
raise ValueError(f"Unexpected audio shape {audio_np.shape}.")
|
||
audio_np = np.clip(audio_np, -1.0, 1.0)
|
||
audio_int16 = (audio_np * 32767.0).astype(np.int16)
|
||
return audio_int16, audio_int16.shape[1]
|
||
|
||
@classmethod
|
||
def _write_pcm_wav(
|
||
cls,
|
||
wav_path: str,
|
||
audio: torch.Tensor | np.ndarray,
|
||
sample_rate: int,
|
||
) -> int:
|
||
"""Write 16-bit PCM WAV; returns the channel count."""
|
||
import wave
|
||
audio_int16, num_channels = cls._audio_to_int16(audio)
|
||
with wave.open(wav_path, "wb") as f:
|
||
f.setnchannels(num_channels)
|
||
f.setsampwidth(2)
|
||
f.setframerate(sample_rate)
|
||
f.writeframes(audio_int16.tobytes())
|
||
return num_channels
|
||
|
||
@classmethod
|
||
def _save_video_with_audio_single_pass(
|
||
cls,
|
||
output_path: str,
|
||
frames: list[np.ndarray],
|
||
fps: int,
|
||
audio: torch.Tensor | np.ndarray,
|
||
sample_rate: int,
|
||
) -> bool:
|
||
"""Encode video+audio into MP4 in one pass using PyAV."""
|
||
try:
|
||
import av
|
||
except ImportError:
|
||
logger.warning("PyAV not installed; cannot use single-pass save.")
|
||
return False
|
||
|
||
if not frames:
|
||
return False
|
||
|
||
output = None
|
||
try:
|
||
audio_int16, num_channels = cls._audio_to_int16(audio)
|
||
layout = "stereo" if num_channels == 2 else "mono"
|
||
output = av.open(output_path, mode="w")
|
||
video_stream = output.add_stream("libx264", rate=fps)
|
||
video_stream.width = int(frames[0].shape[1])
|
||
video_stream.height = int(frames[0].shape[0])
|
||
video_stream.pix_fmt = "yuv420p"
|
||
video_stream.options = {
|
||
"preset": "ultrafast",
|
||
"tune": "zerolatency",
|
||
}
|
||
|
||
audio_stream = output.add_stream("aac", rate=sample_rate, layout=layout)
|
||
|
||
for frame_np in frames:
|
||
vframe = av.VideoFrame.from_ndarray(np.ascontiguousarray(frame_np), format="rgb24")
|
||
for packet in video_stream.encode(vframe):
|
||
output.mux(packet)
|
||
for packet in video_stream.encode():
|
||
output.mux(packet)
|
||
|
||
# AAC commonly uses 1024 samples/frame. Pad the tail frame to avoid
|
||
# shape errors in PyAV/FFmpeg.
|
||
chunk_size = 1024
|
||
for start in range(0, audio_int16.shape[0], chunk_size):
|
||
chunk = audio_int16[start:start + chunk_size]
|
||
if chunk.shape[0] < chunk_size:
|
||
pad = np.zeros((chunk_size - chunk.shape[0], chunk.shape[1]), dtype=chunk.dtype)
|
||
chunk = np.concatenate([chunk, pad], axis=0)
|
||
chunk_planar = np.ascontiguousarray(chunk.T)
|
||
aframe = av.AudioFrame.from_ndarray(chunk_planar, format="s16p", layout=layout)
|
||
aframe.sample_rate = sample_rate
|
||
for packet in audio_stream.encode(aframe):
|
||
output.mux(packet)
|
||
for packet in audio_stream.encode():
|
||
output.mux(packet)
|
||
|
||
output.close()
|
||
return True
|
||
except Exception as e:
|
||
logger.warning("Single-pass video+audio save failed: %s", e)
|
||
if output is not None:
|
||
with suppress(Exception):
|
||
output.close()
|
||
return False
|
||
|
||
@staticmethod
|
||
def _ffmpeg_encoder_supports_option(ffmpeg_bin: str, codec: str, option_name: str) -> bool:
|
||
"""Best-effort check whether ffmpeg encoder supports an option."""
|
||
cache_key = (ffmpeg_bin, codec, option_name)
|
||
cached = _FFMPEG_ENCODER_OPTION_CACHE.get(cache_key)
|
||
if cached is not None:
|
||
return cached
|
||
|
||
try:
|
||
result = subprocess.run(
|
||
[
|
||
ffmpeg_bin,
|
||
"-hide_banner",
|
||
"-h",
|
||
f"encoder={codec}",
|
||
],
|
||
check=False,
|
||
capture_output=True,
|
||
text=True,
|
||
)
|
||
if result.returncode != 0:
|
||
_FFMPEG_ENCODER_OPTION_CACHE[cache_key] = False
|
||
return False
|
||
haystack = f"{result.stdout}\n{result.stderr}"
|
||
supported = f"-{option_name}" in haystack
|
||
_FFMPEG_ENCODER_OPTION_CACHE[cache_key] = supported
|
||
return supported
|
||
except Exception:
|
||
_FFMPEG_ENCODER_OPTION_CACHE[cache_key] = False
|
||
return False
|
||
|
||
@classmethod
|
||
def _save_video_with_audio_ffmpeg_pipe(
|
||
cls,
|
||
output_path: str,
|
||
frames: list[np.ndarray],
|
||
fps: int,
|
||
audio: torch.Tensor | np.ndarray,
|
||
sample_rate: int,
|
||
) -> bool:
|
||
"""Encode video+audio using ffmpeg via rawvideo stdin + WAV input."""
|
||
ffmpeg_bin = shutil.which(envs.FASTVIDEO_FFMPEG_BIN.get())
|
||
if ffmpeg_bin is None:
|
||
logger.warning("ffmpeg not found; cannot use ffmpeg pipe save.")
|
||
return False
|
||
|
||
if not frames:
|
||
return False
|
||
|
||
height = int(frames[0].shape[0])
|
||
width = int(frames[0].shape[1])
|
||
codec = envs.FASTVIDEO_VIDEO_CODEC.get()
|
||
|
||
try:
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
wav_path = os.path.join(tmpdir, "audio.wav")
|
||
try:
|
||
cls._write_pcm_wav(wav_path, audio, sample_rate)
|
||
except ValueError as e:
|
||
logger.warning("Unexpected audio tensor for ffmpeg pipe save: %s", e)
|
||
return False
|
||
|
||
cmd = [
|
||
ffmpeg_bin,
|
||
"-hide_banner",
|
||
"-loglevel",
|
||
"error",
|
||
"-y",
|
||
"-f",
|
||
"rawvideo",
|
||
"-pix_fmt",
|
||
"rgb24",
|
||
"-s:v",
|
||
f"{width}x{height}",
|
||
"-r",
|
||
str(fps),
|
||
"-i",
|
||
"pipe:0",
|
||
"-i",
|
||
wav_path,
|
||
"-c:v",
|
||
codec,
|
||
]
|
||
|
||
if codec.endswith("_nvenc"):
|
||
nvenc_options = [
|
||
("preset", envs.FASTVIDEO_NVENC_PRESET.get()),
|
||
("tune", envs.FASTVIDEO_NVENC_TUNE.get()),
|
||
("rc", envs.FASTVIDEO_NVENC_RC.get()),
|
||
("qp", envs.FASTVIDEO_NVENC_QP.get()),
|
||
("bf", envs.FASTVIDEO_NVENC_BF.get()),
|
||
]
|
||
for option_name, option_value in nvenc_options:
|
||
if cls._ffmpeg_encoder_supports_option(ffmpeg_bin, codec, option_name):
|
||
cmd += [f"-{option_name}", option_value]
|
||
else:
|
||
cmd += ["-preset", envs.FASTVIDEO_X264_PRESET.get()]
|
||
|
||
cmd += [
|
||
"-c:a",
|
||
"aac",
|
||
"-pix_fmt",
|
||
envs.FASTVIDEO_OUTPUT_PIX_FMT.get(),
|
||
# Audio and video decoders may produce different durations.
|
||
# Preserve every video frame and pad/trim audio to match.
|
||
# No -frames:v: ffmpeg 4.4 closes every stream once the video
|
||
# reaches that count, which cuts the padded audio short.
|
||
"-af",
|
||
"apad",
|
||
"-t",
|
||
str(len(frames) / fps),
|
||
"-movflags",
|
||
"+faststart",
|
||
output_path,
|
||
]
|
||
|
||
proc = subprocess.Popen(cmd, stdin=subprocess.PIPE)
|
||
if proc.stdin is None:
|
||
proc.kill()
|
||
proc.wait()
|
||
logger.warning("ffmpeg stdin unavailable.")
|
||
return False
|
||
if fcntl is not None and hasattr(fcntl, "F_SETPIPE_SZ"):
|
||
with suppress(OSError):
|
||
fcntl.fcntl(proc.stdin.fileno(), fcntl.F_SETPIPE_SZ, 1048576)
|
||
|
||
try:
|
||
ts_start = time.perf_counter()
|
||
for frame in frames:
|
||
proc.stdin.write(frame.tobytes())
|
||
proc.stdin.close()
|
||
logger.info("Wrote frames to ffmpeg stdin in %.3f s", time.perf_counter() - ts_start)
|
||
rc = proc.wait()
|
||
if rc != 0:
|
||
logger.warning("ffmpeg pipe save failed with return code %d", rc)
|
||
return False
|
||
return True
|
||
except Exception:
|
||
proc.kill()
|
||
proc.wait()
|
||
raise
|
||
except Exception as e:
|
||
logger.warning("ffmpeg pipe save failed: %s", e)
|
||
return False
|
||
|
||
@classmethod
|
||
def _mux_audio(
|
||
cls,
|
||
video_path: str,
|
||
audio: torch.Tensor | np.ndarray,
|
||
sample_rate: int,
|
||
) -> bool:
|
||
"""Mux audio into video using PyAV."""
|
||
try:
|
||
import av
|
||
except ImportError:
|
||
logger.warning("PyAV not installed; cannot mux audio. "
|
||
"Install with: uv pip install av")
|
||
return False
|
||
|
||
try:
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
out_path = os.path.join(tmpdir, "muxed.mp4")
|
||
wav_path = os.path.join(tmpdir, "audio.wav")
|
||
|
||
num_channels = cls._write_pcm_wav(wav_path, audio, sample_rate)
|
||
layout = "stereo" if num_channels == 2 else "mono"
|
||
|
||
# Open input video and audio
|
||
input_video = av.open(video_path)
|
||
input_audio = av.open(wav_path)
|
||
|
||
# Create output with both streams
|
||
output = av.open(out_path, mode="w")
|
||
|
||
# Add video stream (copy codec from input)
|
||
in_video_stream = input_video.streams.video[0]
|
||
out_video_stream = output.add_stream(
|
||
codec_name=in_video_stream.codec_context.name,
|
||
rate=in_video_stream.average_rate,
|
||
)
|
||
out_video_stream.width = in_video_stream.width
|
||
out_video_stream.height = in_video_stream.height
|
||
out_video_stream.pix_fmt = in_video_stream.pix_fmt
|
||
|
||
# Add audio stream (AAC)
|
||
out_audio_stream = output.add_stream("aac", rate=sample_rate)
|
||
out_audio_stream.layout = layout
|
||
|
||
# Remux video (decode and re-encode to be safe)
|
||
for frame in input_video.decode(video=0):
|
||
for packet in out_video_stream.encode(frame):
|
||
output.mux(packet)
|
||
for packet in out_video_stream.encode():
|
||
output.mux(packet)
|
||
|
||
# Encode audio
|
||
for frame in input_audio.decode(audio=0):
|
||
frame.pts = None # Let encoder assign PTS
|
||
for packet in out_audio_stream.encode(frame):
|
||
output.mux(packet)
|
||
for packet in out_audio_stream.encode():
|
||
output.mux(packet)
|
||
|
||
input_video.close()
|
||
input_audio.close()
|
||
output.close()
|
||
shutil.move(out_path, video_path)
|
||
return True
|
||
except Exception as e:
|
||
logger.warning("Audio mux failed: %s", e)
|
||
return False
|
||
|
||
def set_lora_adapter(self,
|
||
lora_nickname: str,
|
||
lora_path: str | None = None,
|
||
strength: float = 1.0,
|
||
accumulate: bool = False) -> None:
|
||
self.executor.set_lora_adapter(lora_nickname, lora_path, strength=strength, accumulate=accumulate)
|
||
|
||
def unmerge_lora_weights(self) -> None:
|
||
"""
|
||
Use unmerged weights for inference to produce videos that align with
|
||
validation videos generated during training.
|
||
"""
|
||
self.executor.unmerge_lora_weights()
|
||
|
||
def merge_lora_weights(self) -> None:
|
||
self.executor.merge_lora_weights()
|
||
|
||
def shutdown(self):
|
||
"""
|
||
Shutdown the video generator.
|
||
"""
|
||
self.executor.shutdown()
|
||
del self.executor
|
||
|
||
|
||
def _final_event_from_result(result: GenerationResult) -> VideoFinalEvent:
|
||
"""Build a :class:`VideoFinalEvent` from a terminal result.
|
||
|
||
Streaming consumers prefer ``frames`` for the MSE-backed path;
|
||
server-side consumers prefer encoded ``video_bytes``. We carry both
|
||
— whichever the pipeline actually produced — and attach the full
|
||
:class:`GenerationResult` so callers that want the complete object
|
||
don't have to keep a second reference.
|
||
"""
|
||
video_bytes: bytes | None = None
|
||
if result.video_path and os.path.isfile(result.video_path):
|
||
try:
|
||
with open(result.video_path, "rb") as f:
|
||
video_bytes = f.read()
|
||
except OSError:
|
||
video_bytes = None
|
||
metadata = {
|
||
"generation_time": result.generation_time,
|
||
"peak_memory_mb": result.peak_memory_mb,
|
||
"video_path": result.video_path,
|
||
}
|
||
return VideoFinalEvent(
|
||
video_bytes=video_bytes,
|
||
tensor=result.samples,
|
||
frames=result.frames,
|
||
metadata=metadata,
|
||
continuation_state=result.state,
|
||
result=result,
|
||
)
|