# SPDX-License-Identifier: Apache-2.0 """ VideoGenerator module for FastVideo. This module provides a consolidated interface for generating videos using diffusion models. """ import math import os import time from copy import deepcopy from typing import Any import imageio import numpy as np import torch import torchvision from einops import rearrange from fastvideo.configs.sample import SamplingParam from fastvideo.fastvideo_args import FastVideoArgs 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 logger = init_logger(__name__) 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): """ Initialize the video generator. Args: fastvideo_args: The inference arguments executor_class: The executor class to use for inference """ self.fastvideo_args = fastvideo_args self.executor = executor_class(fastvideo_args) @classmethod def from_pretrained(cls, model_path: str, device: str | None = None, torch_dtype: torch.dtype | None = None, **kwargs) -> "VideoGenerator": """ Create a video generator from a pretrained model. Args: model_path: Path or identifier for the pretrained model device: Device to load the model on (e.g., "cuda", "cuda:0", "cpu") torch_dtype: Data type for model weights (e.g., torch.float16) 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 """ # If users also provide some kwargs, it will override the FastVideoArgs and PipelineConfig. kwargs['model_path'] = model_path fastvideo_args = FastVideoArgs.from_kwargs(**kwargs) return cls.from_fastvideo_args(fastvideo_args) @classmethod def from_fastvideo_args(cls, fastvideo_args: FastVideoArgs) -> "VideoGenerator": """ Create a video generator with the specified arguments. Args: fastvideo_args: The inference arguments 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 ) def generate_video( self, prompt: str | None = None, sampling_param: SamplingParam | None = None, **kwargs, ) -> dict[str, Any] | list[np.ndarray] | 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) output_video_name: Name of the video file to save. Default is the first 100 characters of the prompt. save_video: Whether to save the video to disk return_frames: Whether to return the raw frames 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: Either the output dictionary, list of frames, or list of results for batch processing """ # Handle batch processing from text file if self.fastvideo_args.prompt_txt is not None: prompt_txt_path = self.fastvideo_args.prompt_txt if 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) if sampling_param is not None: original_output_video_name = sampling_param.output_video_name else: original_output_video_name = None 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 if sampling_param is not None and original_output_video_name is not None: sampling_param.output_video_name = original_output_video_name + f"_{i}" result = self._generate_single_video( batch_prompt, sampling_param, **kwargs) # Add prompt info to result if isinstance(result, dict): 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: raise ValueError("Either prompt or prompt_txt must be provided") return self._generate_single_video(prompt, sampling_param, **kwargs) def _generate_single_video( self, prompt: str, sampling_param: SamplingParam | None = None, **kwargs, ) -> dict[str, Any] | list[np.ndarray]: """Internal method for single video generation""" # Create a copy of inference args to avoid modifying the original fastvideo_args = self.fastvideo_args pipeline_config = fastvideo_args.pipeline_config # Validate inputs if not isinstance(prompt, str): raise TypeError( f"`prompt` must be a string, but got {type(prompt)}") prompt = prompt.strip() if sampling_param is None: sampling_param = SamplingParam.from_pretrained( fastvideo_args.model_path) else: sampling_param = deepcopy(sampling_param) kwargs["prompt"] = prompt sampling_param.update(kwargs) # 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}") temporal_scale_factor = pipeline_config.vae_config.arch_config.temporal_compression_ratio num_frames = sampling_param.num_frames num_gpus = fastvideo_args.num_gpus use_temporal_scaling_frames = pipeline_config.vae_config.use_temporal_scaling_frames # Adjust number of frames based on number of GPUs if use_temporal_scaling_frames: orig_latent_num_frames = (num_frames - 1) // temporal_scale_factor + 1 else: # stepvideo only orig_latent_num_frames = sampling_param.num_frames // 17 * 3 if orig_latent_num_frames % fastvideo_args.num_gpus != 0: # Adjust latent frames to be divisible by number of GPUs if sampling_param.num_frames_round_down: # Ensure we have at least 1 batch per GPU new_latent_num_frames = max( 1, (orig_latent_num_frames // num_gpus)) * num_gpus else: new_latent_num_frames = math.ceil( orig_latent_num_frames / num_gpus) * num_gpus if use_temporal_scaling_frames: # Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor new_num_frames = (new_latent_num_frames - 1) * temporal_scale_factor + 1 else: # stepvideo only # Find the least common multiple of 3 and num_gpus divisor = math.lcm(3, num_gpus) # Round up to the nearest multiple of this LCM new_latent_num_frames = ( (new_latent_num_frames + divisor - 1) // divisor) * divisor # Convert back to actual frames using the StepVideo formula new_num_frames = new_latent_num_frames // 3 * 17 logger.info( "Adjusting number of frames from %s to %s based on number of GPUs (%s)", sampling_param.num_frames, new_num_frames, fastvideo_args.num_gpus) sampling_param.num_frames = new_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: {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: {sampling_param.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, extra={}, ) # Use prompt[:100] for video name if batch.output_video_name is None: batch.output_video_name = prompt[:100] # Run inference start_time = time.perf_counter() output_batch = self.executor.execute_forward(batch, fastvideo_args) samples = output_batch.output logging_info = output_batch.logging_info gen_time = time.perf_counter() - start_time logger.info("Generated successfully in %.2f seconds", gen_time) # Process outputs videos = rearrange(samples, "b c t h w -> t b c h w") frames = [] for x in videos: x = torchvision.utils.make_grid(x, nrow=6) x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) frames.append((x * 255).numpy().astype(np.uint8)) # Save video if requested if batch.save_video: output_path = batch.output_path if output_path: os.makedirs(output_path, exist_ok=True) video_path = os.path.join(output_path, f"{batch.output_video_name}.mp4") imageio.mimsave(video_path, frames, fps=batch.fps, format="mp4") logger.info("Saved video to %s", video_path) else: logger.warning("No output path provided, video not saved") if batch.return_frames: return frames else: return { "samples": samples, "frames": frames, "prompts": prompt, "size": (target_height, target_width, batch.num_frames), "generation_time": gen_time, "logging_info": logging_info, "trajectory": output_batch.trajectory_latents, "trajectory_timesteps": output_batch.trajectory_timesteps, "trajectory_decoded": output_batch.trajectory_decoded, } def set_lora_adapter(self, lora_nickname: str, lora_path: str | None = None) -> None: self.executor.set_lora_adapter(lora_nickname, lora_path) 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