# 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 Mapping from contextlib import suppress from copy import deepcopy from typing import Any import imageio import numpy as np import torch import torchvision from einops import rearrange from fastvideo.api.compat import ( 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_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, 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] = {} _BATCH_EXTRA_PASSTHROUGH_KEYS: tuple[str, ...] = ( "ltx2_audio_latents", "ltx2_audio_clean_latent", "ltx2_audio_denoise_mask", "audio_num_frames", "video_position_offset_sec", # MiniMax-H3 VSA per-request knobs (read by the H3 denoising stage; # sparsity itself flows through the existing ForwardBatch.VSA_sparsity) "vsa_mode", "vsa_dense_first_n_steps", "vsa_dense_layers", ) _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", "output_type", "nvfp4_fa4", }) 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 _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" os.environ.setdefault("CUTE_DSL_ENABLE_TVM_FFI", "1") 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) return cls( fastvideo_args=fastvideo_args, executor_class=executor_class, log_stats=False, # TODO: implement log_queue=log_queue, ) 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 = 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 = 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 = 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 = request_to_sampling_param( request, model_path=self.fastvideo_args.model_path, ) result = self._generate_video_impl( prompt=request.prompt, sampling_param=sampling_param, fastvideo_args=fastvideo_args, ) 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 # Handle batch processing from text file if sampling_param is None: sampling_param = SamplingParam.from_pretrained(fastvideo_args.model_path) # 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 if fastvideo_args.prompt_txt is not None or sampling_param.prompt_path is not None: prompt_txt_path = sampling_param.prompt_path or fastvideo_args.prompt_txt if not prompt_txt_path or not os.path.exists(prompt_txt_path): raise FileNotFoundError(f"Prompt text file not found: {prompt_txt_path}") # Read prompts from file 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}") 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._prepare_output_path(sampling_param.output_path, batch_prompt) 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) continue logger.info("Completed batch processing. Generated %d videos successfully.", len(results)) return results # Single prompt generation (original behavior) 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 = "" else: raise ValueError("Either prompt or prompt_txt must be provided") output_path = self._prepare_output_path(sampling_param.output_path, prompt) 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. The file extension is chosen automatically based on the workload type: ``.png`` for image workloads (``t2i``, ``i2i``, …) and ``.mp4`` for video workloads. - 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. """ if self._is_image_workload(): target_ext = ".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" base_path, extension = os.path.splitext(output_path) extension_lower = extension.lower() if extension_lower == target_ext: 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""" 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 # 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 batch.return_frames else None``). # Post-decode frame building reads ``output_batch.output`` # directly (the GPU ``vid_u8`` path), not ``samples``. So when # ``return_frames=False`` the pinned fp32 alloc + 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 skip_pixel_prealloc = is_latent_output or not needs_samples_out if skip_pixel_prealloc: samples = torch.empty(0, device='cpu') else: samples = torch.empty( (latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width), device='cpu', pin_memory=fastvideo_args.pin_cpu_memory) thread.join() 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 ``return_frames=False``); 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) 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() 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: # Quantize on the source device (typically CUDA) BEFORE the # device->host copy. `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. The old path # paid a full fp32 video D->H copy (which scales with # resolution x frames x batch) and then a single-threaded # per-frame CPU *255/cast loop. Casting to uint8 on-device # makes the transfer 4x smaller, ships it in a single copy, # 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 is SSIM-gated, not bit-exact: float->uint8 # differs <=1 LSB CPU vs GPU.) src = output_batch.output vid_u8 = (src * 255).clamp_(0, 255).to(torch.uint8) vid_u8 = rearrange(vid_u8, "b c t h w -> t b c h w").cpu() 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 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=batch.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=batch.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=batch.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=batch.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 batch.return_frames 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"), "audio_sample_rate": output_batch.extra.get("audio_sample_rate"), "ltx2_audio_latents": output_batch.extra.get("ltx2_audio_latents"), "size": output_size, "generation_time": gen_time, "e2e_latency": e2e_time, "logging_info": logging_info, "trajectory": output_batch.trajectory_latents, "trajectory_timesteps": output_batch.trajectory_timesteps, "trajectory_decoded": output_batch.trajectory_decoded, "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(os.getenv("FASTVIDEO_FFMPEG_BIN", "ffmpeg")) 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 = os.getenv("FASTVIDEO_VIDEO_CODEC", "libx264") 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", os.getenv("FASTVIDEO_NVENC_PRESET", "p1")), ("tune", os.getenv("FASTVIDEO_NVENC_TUNE", "ull")), ("rc", os.getenv("FASTVIDEO_NVENC_RC", "constqp")), ("qp", os.getenv("FASTVIDEO_NVENC_QP", "28")), ("bf", os.getenv("FASTVIDEO_NVENC_BF", "0")), ] 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", os.getenv("FASTVIDEO_X264_PRESET", "ultrafast")] cmd += [ "-c:a", "aac", "-pix_fmt", os.getenv("FASTVIDEO_OUTPUT_PIX_FMT", "yuv420p"), "-shortest", "-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, )