From 4ccf68d71e685f1bcb7d189cbd63d9b25c14ff7a Mon Sep 17 00:00:00 2001 From: Will Lin Date: Tue, 15 Apr 2025 16:05:04 -0700 Subject: [PATCH] checkpoint --- fastvideo/v1/entrypoints/video_generator.py | 150 ++++++++++---------- fastvideo/v1/fastvideo_args.py | 6 +- fastvideo/v1/worker/executor.py | 18 +-- fastvideo/v1/worker/multiproc_executor.py | 12 +- fastvideo/v1/worker/worker.py | 26 ++-- 5 files changed, 106 insertions(+), 106 deletions(-) diff --git a/fastvideo/v1/entrypoints/video_generator.py b/fastvideo/v1/entrypoints/video_generator.py index 248283fc8..5a91ffe5c 100644 --- a/fastvideo/v1/entrypoints/video_generator.py +++ b/fastvideo/v1/entrypoints/video_generator.py @@ -16,7 +16,7 @@ import torch import torchvision from einops import rearrange -from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.logger import init_logger from fastvideo.v1.pipelines import (ForwardBatch) @@ -34,17 +34,17 @@ class VideoGenerator: customization options, similar to popular frameworks like HF Diffusers. """ - def __init__(self, inference_args: InferenceArgs, + def __init__(self, fastvideo_args: FastVideoArgs, executor_class: type[Executor], log_stats: bool): """ Initialize the video generator. Args: pipeline: The pipeline to use for inference - inference_args: The inference arguments + fastvideo_args: The inference arguments """ - self.inference_args = inference_args - self.executor = executor_class(inference_args) + self.fastvideo_args = fastvideo_args + self.executor = executor_class(fastvideo_args) @classmethod def from_pretrained(cls, @@ -64,35 +64,35 @@ class VideoGenerator: Returns: The created video generator """ - inference_args = InferenceArgs( + fastvideo_args = FastVideoArgs( model_path=model_path, device_str=device or "cuda" if torch.cuda.is_available() else "cpu", **kwargs) if torch_dtype is not None: - inference_args.dtype = torch_dtype + fastvideo_args.dtype = torch_dtype - return cls.from_inference_args(inference_args) + return cls.from_fastvideo_args(fastvideo_args) @classmethod - def from_inference_args(cls, - inference_args: InferenceArgs) -> "VideoGenerator": + def from_fastvideo_args(cls, + fastvideo_args: FastVideoArgs) -> "VideoGenerator": """ Create a video generator with the specified arguments. Args: - inference_args: The inference arguments + fastvideo_args: The inference arguments Returns: The created video generator """ # Initialize distributed environment if needed - # initialize_distributed_and_parallelism(inference_args) + # initialize_distributed_and_parallelism(fastvideo_args) - executor_class = Executor.get_class(inference_args) + executor_class = Executor.get_class(fastvideo_args) return cls( - inference_args=inference_args, + fastvideo_args=fastvideo_args, executor_class=executor_class, log_stats=False, # TODO: implement ) @@ -119,17 +119,17 @@ class VideoGenerator: Args: prompt: The prompt to use for generation - negative_prompt: The negative prompt to use (overrides the one in inference_args) - output_path: Path to save the video (overrides the one in inference_args) + 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) 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 inference_args) - guidance_scale: Classifier-free guidance scale (overrides inference_args) - num_frames: Number of frames to generate (overrides inference_args) - height: Height of generated video (overrides inference_args) - width: Width of generated video (overrides inference_args) - fps: Frames per second for saved video (overrides inference_args) - seed: Random seed for generation (overrides inference_args) + 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 @@ -137,29 +137,29 @@ class VideoGenerator: Either the output dictionary or the list of frames depending on return_frames """ # Create a copy of inference args to avoid modifying the original - inference_args = self.inference_args.copy() + fastvideo_args = self.fastvideo_args.copy() # Override parameters if provided if negative_prompt is not None: - inference_args.neg_prompt = negative_prompt + fastvideo_args.neg_prompt = negative_prompt if num_inference_steps is not None: - inference_args.num_inference_steps = num_inference_steps + fastvideo_args.num_inference_steps = num_inference_steps if guidance_scale is not None: - inference_args.guidance_scale = guidance_scale + fastvideo_args.guidance_scale = guidance_scale if num_frames is not None: - inference_args.num_frames = num_frames + fastvideo_args.num_frames = num_frames if height is not None: - inference_args.height = height + fastvideo_args.height = height if width is not None: - inference_args.width = width + fastvideo_args.width = width if fps is not None: - inference_args.fps = fps + fastvideo_args.fps = fps if seed is not None: - inference_args.seed = seed + fastvideo_args.seed = seed # Store callback info - inference_args.callback = callback - inference_args.callback_steps = callback_steps + fastvideo_args.callback = callback + fastvideo_args.callback_steps = callback_steps # Validate inputs if not isinstance(prompt, str): @@ -168,61 +168,61 @@ class VideoGenerator: prompt = prompt.strip() # Process negative prompt - if inference_args.neg_prompt is not None: - inference_args.neg_prompt = inference_args.neg_prompt.strip() + if fastvideo_args.neg_prompt is not None: + fastvideo_args.neg_prompt = fastvideo_args.neg_prompt.strip() # Validate dimensions - if (inference_args.height <= 0 or inference_args.width <= 0 - or inference_args.num_frames <= 0): + if (fastvideo_args.height <= 0 or fastvideo_args.width <= 0 + or fastvideo_args.num_frames <= 0): raise ValueError( f"Height, width, and num_frames must be positive integers, got " - f"height={inference_args.height}, width={inference_args.width}, " - f"num_frames={inference_args.num_frames}") + f"height={fastvideo_args.height}, width={fastvideo_args.width}, " + f"num_frames={fastvideo_args.num_frames}") - if (inference_args.num_frames - 1) % 4 != 0: + if (fastvideo_args.num_frames - 1) % 4 != 0: raise ValueError( - f"num_frames-1 must be a multiple of 4, got {inference_args.num_frames}" + f"num_frames-1 must be a multiple of 4, got {fastvideo_args.num_frames}" ) # Calculate sizes - target_height = align_to(inference_args.height, 16) - target_width = align_to(inference_args.width, 16) + target_height = align_to(fastvideo_args.height, 16) + target_width = align_to(fastvideo_args.width, 16) # Calculate latent sizes - latents_size = [(inference_args.num_frames - 1) // 4 + 1, - inference_args.height // 8, inference_args.width // 8] + latents_size = [(fastvideo_args.num_frames - 1) // 4 + 1, + fastvideo_args.height // 8, fastvideo_args.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: {inference_args.num_frames} + video_length: {fastvideo_args.num_frames} prompt: {prompt} - neg_prompt: {inference_args.neg_prompt} - seed: {inference_args.seed} - infer_steps: {inference_args.num_inference_steps} - num_videos_per_prompt: {inference_args.num_videos} - guidance_scale: {inference_args.guidance_scale} + neg_prompt: {fastvideo_args.neg_prompt} + seed: {fastvideo_args.seed} + infer_steps: {fastvideo_args.num_inference_steps} + num_videos_per_prompt: {fastvideo_args.num_videos} + guidance_scale: {fastvideo_args.guidance_scale} n_tokens: {n_tokens} - flow_shift: {inference_args.flow_shift} - embedded_guidance_scale: {inference_args.embedded_cfg_scale}""" + flow_shift: {fastvideo_args.flow_shift} + embedded_guidance_scale: {fastvideo_args.embedded_cfg_scale}""" logger.info(debug_str) # Prepare batch - device = torch.device(inference_args.device_str) + device = torch.device(fastvideo_args.device_str) batch = ForwardBatch( prompt=prompt, - negative_prompt=inference_args.neg_prompt, - num_videos_per_prompt=inference_args.num_videos, - height=inference_args.height, - width=inference_args.width, - num_frames=inference_args.num_frames, - num_inference_steps=inference_args.num_inference_steps, - guidance_scale=inference_args.guidance_scale, + negative_prompt=fastvideo_args.neg_prompt, + num_videos_per_prompt=fastvideo_args.num_videos, + height=fastvideo_args.height, + width=fastvideo_args.width, + num_frames=fastvideo_args.num_frames, + num_inference_steps=fastvideo_args.num_inference_steps, + guidance_scale=fastvideo_args.guidance_scale, eta=0.0, n_tokens=n_tokens, - data_type="video" if inference_args.num_frames > 1 else "image", + data_type="video" if fastvideo_args.num_frames > 1 else "image", device=device, extra={}, ) @@ -231,7 +231,7 @@ class VideoGenerator: start_time = time.time() samples = self.pipeline.forward( batch=batch, - inference_args=inference_args, + fastvideo_args=fastvideo_args, ).output gen_time = time.time() - start_time @@ -247,11 +247,11 @@ class VideoGenerator: # Save video if requested if save_video: - save_path = output_path or inference_args.output_path + save_path = output_path or fastvideo_args.output_path if save_path: os.makedirs(os.path.dirname(save_path), exist_ok=True) video_path = os.path.join(save_path, f"{prompt[:100]}.mp4") - imageio.mimsave(video_path, frames, fps=inference_args.fps) + imageio.mimsave(video_path, frames, fps=fastvideo_args.fps) logger.info(f"Saved video to {video_path}") else: logger.warning("No output path provided, video not saved") @@ -263,7 +263,7 @@ class VideoGenerator: "samples": samples, "prompts": prompt, "size": - (target_height, target_width, inference_args.num_frames), + (target_height, target_width, fastvideo_args.num_frames), "generation_time": gen_time } @@ -276,14 +276,14 @@ class VideoGenerator: Args: prompts: List of prompts to generate videos for - output_path: Path to save the videos (overrides the one in inference_args) + output_path: Path to save the videos (overrides the one in fastvideo_args) **kwargs: Additional parameters to pass to generate_video Returns: List of output dictionaries from each generation """ if output_path: - self.inference_args.output_path = output_path + self.fastvideo_args.output_path = output_path results = [] for prompt in prompts: @@ -321,13 +321,13 @@ class VideoGenerator: image_tensor = image # Add image to inference args - inference_args = self.inference_args.copy() - inference_args.init_image = image_tensor - inference_args.strength = strength + fastvideo_args = self.fastvideo_args.copy() + fastvideo_args.init_image = image_tensor + fastvideo_args.strength = strength # Generate video return self.generate_video(prompt=prompt or "", - inference_args=inference_args, + fastvideo_args=fastvideo_args, **kwargs) def to(self, device: Union[str, torch.device]) -> "VideoGenerator": @@ -341,8 +341,8 @@ class VideoGenerator: Self for chaining """ device_str = str(device) - self.inference_args.device_str = device_str - self.inference_args.device = torch.device(device_str) + self.fastvideo_args.device_str = device_str + self.fastvideo_args.device = torch.device(device_str) # Move pipeline components to device self.pipeline.to(device) diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py index 762e36caa..8461302b1 100644 --- a/fastvideo/v1/fastvideo_args.py +++ b/fastvideo/v1/fastvideo_args.py @@ -124,7 +124,7 @@ class FastVideoArgs: parser.add_argument( "--distributed-executor-backend", type=str, - default=InferenceArgs.distributed_executor_backend, + default=FastVideoArgs.distributed_executor_backend, choices=["torch"], help="Backend for distributed execution", ) @@ -438,7 +438,7 @@ class FastVideoArgs: return cls(**kwargs) - def check_inference_args(self) -> None: + def check_fastvideo_args(self) -> None: """Validate inference arguments for consistency""" if self.tp_size is None: self.tp_size = self.num_gpus @@ -477,7 +477,7 @@ def prepare_fastvideo_args(argv: List[str]) -> FastVideoArgs: FastVideoArgs.add_cli_args(parser) raw_args = parser.parse_args(argv) fastvideo_args = FastVideoArgs.from_cli_args(raw_args) - fastvideo_args.check_inference_args() + fastvideo_args.check_fastvideo_args() global _current_fastvideo_args _current_fastvideo_args = fastvideo_args return fastvideo_args diff --git a/fastvideo/v1/worker/executor.py b/fastvideo/v1/worker/executor.py index eccf3b746..139bf285d 100644 --- a/fastvideo/v1/worker/executor.py +++ b/fastvideo/v1/worker/executor.py @@ -1,30 +1,30 @@ from abc import ABC, abstractmethod import torch -from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.pipelines import ForwardBatch class Executor(ABC): - def __init__(self, inference_args: InferenceArgs): - self.inference_args = inference_args + def __init__(self, fastvideo_args: FastVideoArgs): + self.fastvideo_args = fastvideo_args @abstractmethod def _init_executor(self) -> None: raise NotImplementedError @classmethod - def get_class(cls, inference_args: InferenceArgs) -> type["Executor"]: - if inference_args.distributed_executor_backend == "torch": + def get_class(cls, fastvideo_args: FastVideoArgs) -> type["Executor"]: + if fastvideo_args.distributed_executor_backend == "torch": from fastvideo.v1.worker.torchrun_executor import TorchRunExecutor - return TorchRunExecutor(inference_args) - elif inference_args.distributed_executor_backend == "mp": + return TorchRunExecutor(fastvideo_args) + elif fastvideo_args.distributed_executor_backend == "mp": from fastvideo.v1.worker.multiproc_executor import MultiprocExecutor - return MultiprocExecutor(inference_args) + return MultiprocExecutor(fastvideo_args) else: raise ValueError( - f"Unsupported distributed executor backend: {inference_args.distributed_executor_backend}" + f"Unsupported distributed executor backend: {fastvideo_args.distributed_executor_backend}" ) def execute_forward(self, forward_batch: ForwardBatch) -> torch.Tensor: diff --git a/fastvideo/v1/worker/multiproc_executor.py b/fastvideo/v1/worker/multiproc_executor.py index 7aec82fe7..86ad54ca6 100644 --- a/fastvideo/v1/worker/multiproc_executor.py +++ b/fastvideo/v1/worker/multiproc_executor.py @@ -4,7 +4,7 @@ from typing import List from multiprocessing import BaseProcess from fastvideo.v1.worker.executor import Executor -from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.worker import PipelineWorker from fastvideo.v1.utils import _init_logger @@ -13,12 +13,12 @@ logger = _init_logger(__name__) class MultiprocExecutor(Executor): - def __init__(self, inference_args: InferenceArgs): - super().__init__(inference_args) + def __init__(self, fastvideo_args: FastVideoArgs): + super().__init__(fastvideo_args) self.workers = [] - for _ in range(self.inference_args.tp_size): - self.workers.append(PipelineWorker(self.inference_args)) + for _ in range(self.fastvideo_args.tp_size): + self.workers.append(PipelineWorker(self.fastvideo_args)) def _init_executor(self) -> None: # The child processes will send SIGUSR1 when unrecoverable @@ -34,7 +34,7 @@ class MultiprocExecutor(Executor): signal.signal(signal.SIGUSR1, sigusr1_handler) - self.world_size = self.inference_args.num_gpus + self.world_size = self.fastvideo_args.num_gpus tensor_parallel_size = self.parallel_config.tensor_parallel_size self.workers: List[BaseProcess] = [] for rank in range(self.world_size): diff --git a/fastvideo/v1/worker/worker.py b/fastvideo/v1/worker/worker.py index 2b863bfcb..cabe11def 100644 --- a/fastvideo/v1/worker/worker.py +++ b/fastvideo/v1/worker/worker.py @@ -4,7 +4,7 @@ import torch from typing import Optional -from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.logger import init_logger from fastvideo.v1.pipelines import ForwardBatch from fastvideo.v1.distributed import ( @@ -17,8 +17,8 @@ logger = init_logger(__name__) class WorkerBase(ABC): - def __init__(self, inference_args: InferenceArgs): - self.inference_args = inference_args + def __init__(self, fastvideo_args: FastVideoArgs): + self.fastvideo_args = fastvideo_args def execute_forward(self, forward_batch: ForwardBatch) -> torch.Tensor: pass @@ -29,33 +29,33 @@ class WorkerBase(ABC): class Worker(WorkerBase): - def __init__(self, inference_args: InferenceArgs): - super().__init__(inference_args) + def __init__(self, fastvideo_args: FastVideoArgs): + super().__init__(fastvideo_args) def execute_forward(self, forward_batch: ForwardBatch) -> torch.Tensor: pass def init_worker_distributed_environment( - inference_args: InferenceArgs, + fastvideo_args: FastVideoArgs, rank: int, distributed_init_method: Optional[str] = None, local_rank: int = -1, ) -> None: """Initialize distributed environment and model parallelism.""" - world_size = inference_args.num_gpus + world_size = fastvideo_args.num_gpus torch.cuda.set_device(local_rank) init_distributed_environment(world_size=world_size, rank=rank, local_rank=local_rank) device_str = f"cuda:{local_rank}" - inference_args.device_str = device_str - inference_args.device = torch.device(device_str) - assert inference_args.sp_size is not None - assert inference_args.tp_size is not None + fastvideo_args.device_str = device_str + fastvideo_args.device = torch.device(device_str) + assert fastvideo_args.sp_size is not None + assert fastvideo_args.tp_size is not None initialize_model_parallel( - sequence_model_parallel_size=inference_args.sp_size, - tensor_model_parallel_size=inference_args.tp_size, + sequence_model_parallel_size=fastvideo_args.sp_size, + tensor_model_parallel_size=fastvideo_args.tp_size, )