From 81bc6c99430dde2f367f5851b8708c2a4e354cd6 Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Thu, 24 Apr 2025 07:53:44 +0000 Subject: [PATCH] Separate FastVideoArgs from SamplingParam --- fastvideo/v1/configs/pipelines/base.py | 32 ++--- fastvideo/v1/configs/pipelines/hunyuan.py | 2 - fastvideo/v1/configs/pipelines/registry.py | 1 - fastvideo/v1/configs/pipelines/wan.py | 10 -- fastvideo/v1/configs/sample/__init__.py | 5 + fastvideo/v1/configs/sample/base.py | 68 +++++++++ fastvideo/v1/configs/sample/hunyuan.py | 18 +++ fastvideo/v1/configs/sample/registry.py | 73 ++++++++++ fastvideo/v1/configs/sample/wan.py | 22 +++ .../v1/configs/wan_1.3B_t2v_pipeline.json | 12 +- .../v1/configs/wan_14B_i2v_480p_pipeline.json | 12 +- fastvideo/v1/entrypoints/video_generator.py | 101 +++++-------- .../v1/examples/inference/basic/basic.py | 5 + fastvideo/v1/fastvideo_args.py | 134 ------------------ fastvideo/v1/models/dits/base.py | 3 +- fastvideo/v1/models/dits/hunyuanvideo.py | 1 + fastvideo/v1/models/dits/wanvideo.py | 1 + .../v1/pipelines/composed_pipeline_base.py | 3 +- .../v1/pipelines/hunyuan/hunyuan_pipeline.py | 10 +- fastvideo/v1/pipelines/pipeline_batch_info.py | 11 +- .../pipelines/stages/clip_image_encoding.py | 4 +- .../v1/pipelines/stages/clip_text_encoding.py | 6 +- fastvideo/v1/pipelines/stages/denoising.py | 4 +- fastvideo/v1/pipelines/stages/encoding.py | 10 +- .../v1/pipelines/stages/input_validation.py | 12 +- .../v1/pipelines/stages/latent_preparation.py | 20 +-- .../v1/pipelines/stages/llama_encoding.py | 6 +- fastvideo/v1/pipelines/stages/t5_encoding.py | 6 +- .../pipelines/stages/timestep_preparation.py | 2 +- .../v1/pipelines/wan/wan_i2v_pipeline.py | 10 +- fastvideo/v1/pipelines/wan/wan_pipeline.py | 10 +- fastvideo/v1/worker/gpu_worker.py | 1 - 32 files changed, 292 insertions(+), 323 deletions(-) create mode 100644 fastvideo/v1/configs/sample/__init__.py create mode 100644 fastvideo/v1/configs/sample/base.py create mode 100644 fastvideo/v1/configs/sample/hunyuan.py create mode 100644 fastvideo/v1/configs/sample/registry.py create mode 100644 fastvideo/v1/configs/sample/wan.py diff --git a/fastvideo/v1/configs/pipelines/base.py b/fastvideo/v1/configs/pipelines/base.py index 7bae1149a..b2451e72a 100644 --- a/fastvideo/v1/configs/pipelines/base.py +++ b/fastvideo/v1/configs/pipelines/base.py @@ -4,22 +4,14 @@ import json from fastvideo.v1.configs.models import ModelConfig, VAEConfig from fastvideo.v1.utils import shallow_asdict +from fastvideo.v1.logger import init_logger + +logger = init_logger(__name__) @dataclass class BaseConfig: """Base configuration for all pipeline architectures.""" - - # Video parameters - height: int = 720 - width: int = 1280 - num_frames: int = 125 - fps: int = 24 - # Video generation parameters - num_inference_steps: int = 50 - guidance_scale: float = 1.0 - seed: int = 1024 - guidance_rescale: float = 0.0 embedded_cfg_scale: float = 6.0 flow_shift: Optional[float] = None use_cpu_offload: bool = False @@ -32,12 +24,8 @@ class BaseConfig: vae_precision: str = "fp16" vae_tiling: bool = True vae_sp: bool = True - # vae_scale_factor: Optional[int] = None # Deprecated vae_config: VAEConfig = VAEConfig() - # DiT configuration - num_channels_latents: Optional[int] = None # Deprecated - # Image encoder configuration image_encoder_precision: str = "fp32" @@ -50,7 +38,19 @@ class BaseConfig: mask_strategy_file_path: Optional[str] = None enable_torch_compile: bool = False - neg_prompt: Optional[str] = None + @classmethod + def from_pretrained(cls, + model_path: str + ) -> "BaseConfig": + from fastvideo.v1.configs.pipelines.registry import get_pipeline_config_cls_for_name + pipeline_config_cls = get_pipeline_config_cls_for_name(model_path) + if pipeline_config_cls is not None: + pipeline_config = pipeline_config_cls() + else: + logger.warning("Couldn't find an optimal sampling param for %s. Using the default sampling param.", model_path) + pipeline_config = cls() + + return pipeline_config def dump_to_json(self, file_path: str): output_dict = shallow_asdict(self) diff --git a/fastvideo/v1/configs/pipelines/hunyuan.py b/fastvideo/v1/configs/pipelines/hunyuan.py index 818675254..18b6d2568 100644 --- a/fastvideo/v1/configs/pipelines/hunyuan.py +++ b/fastvideo/v1/configs/pipelines/hunyuan.py @@ -14,7 +14,6 @@ class HunyuanConfig(BaseConfig): # Denoising stage embedded_cfg_scale: int = 6 flow_shift: int = 7 - num_inference_steps: int = 50 # Text encoding stage hidden_state_skip_layer: int = 2 @@ -40,7 +39,6 @@ class FastHunyuanConfig(HunyuanConfig): """Configuration specifically optimized for FastHunyuan weights.""" # Override HunyuanConfig defaults - num_inference_steps: int = 6 flow_shift: int = 17 # No need to re-specify guidance_scale or embedded_cfg_scale as they diff --git a/fastvideo/v1/configs/pipelines/registry.py b/fastvideo/v1/configs/pipelines/registry.py index 7c2dc44a2..84ae9797e 100644 --- a/fastvideo/v1/configs/pipelines/registry.py +++ b/fastvideo/v1/configs/pipelines/registry.py @@ -66,7 +66,6 @@ def get_pipeline_config_cls_for_name( # If no match, try to use the fallback config fallback_config = None - print(pipeline_name) # Try to determine pipeline architecture for fallback for pipeline_type, detector in PIPELINE_DETECTOR.items(): if detector(pipeline_name.lower()): diff --git a/fastvideo/v1/configs/pipelines/wan.py b/fastvideo/v1/configs/pipelines/wan.py index bd90e6526..a7e7f921b 100644 --- a/fastvideo/v1/configs/pipelines/wan.py +++ b/fastvideo/v1/configs/pipelines/wan.py @@ -15,17 +15,10 @@ class WanT2V480PConfig(BaseConfig): vae_sp: bool = False # Video parameters - height: int = 480 - width: int = 832 - num_frames: int = 81 - fps: int = 16 use_cpu_offload: bool = True # Denoising stage - guidance_scale: float = 3.0 - neg_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" flow_shift: int = 3 - num_inference_steps: int = 50 # Text encoding stage text_len: int = 512 @@ -46,9 +39,6 @@ class WanI2V480PConfig(WanT2V480PConfig): """Base configuration for Wan I2V 14B 480P pipeline architecture.""" # WanConfig-specific parameters with defaults - # Denoising stage - guidance_scale: float = 5.0 - num_inference_steps: int = 40 # Precision for each component image_encoder_precision: str = "fp32" diff --git a/fastvideo/v1/configs/sample/__init__.py b/fastvideo/v1/configs/sample/__init__.py new file mode 100644 index 000000000..15e4b1455 --- /dev/null +++ b/fastvideo/v1/configs/sample/__init__.py @@ -0,0 +1,5 @@ +from fastvideo.v1.configs.sample.base import SamplingParam + +__all__ = [ + "SamplingParam" +] \ No newline at end of file diff --git a/fastvideo/v1/configs/sample/base.py b/fastvideo/v1/configs/sample/base.py new file mode 100644 index 000000000..14dcb1b56 --- /dev/null +++ b/fastvideo/v1/configs/sample/base.py @@ -0,0 +1,68 @@ +from dataclasses import dataclass +from typing import Optional, Union, List, Dict, Any + +from fastvideo.v1.logger import init_logger + +logger = init_logger(__name__) +@dataclass +class SamplingParam: + # All fields below are copied from ForwardBatch + data_type: str = "video" + + # Image inputs + image_path: Optional[str] = None + + # Text inputs + prompt: Optional[Union[str, List[str]]] = None + negative_prompt: Optional[Union[str, List[str]]] = None + prompt_path: Optional[str] = None + output_path: str = "outputs/" + + # Batch info + num_videos_per_prompt: int = 1 + seed: int = 1024 + + # Original dimensions (before VAE scaling) + num_frames: int = 125 + height: int = 720 + width: int = 1280 + fps: int = 24 + + # Denoising parameters + num_inference_steps: int = 50 + guidance_scale: float = 1.0 + guidance_rescale: float = 0.0 + + # Misc + save_video: bool = True + return_frames: bool = False + + def __post_init__(self): + self.data_type = "video" if self.num_frames > 1 else "image" + + def check_sampling_param(self): + if self.prompt_path and not self.prompt_path.endswith(".txt"): + raise ValueError("prompt_path must be a txt file") + + def update(self, source_dict: Dict[str, Any]) -> None: + for key, value in source_dict.items(): + if hasattr(self, key): + setattr(self, key, value) + else: + logger.exception("%s has no attribute %s", type(self).__name__, key) + + self.__post_init__() + + @classmethod + def from_pretrained(cls, + model_path: str + ) -> "SamplingParam": + from fastvideo.v1.configs.sample.registry import get_sampling_param_cls_for_name + sampling_cls = get_sampling_param_cls_for_name(model_path) + if sampling_cls is not None: + sampling_param = sampling_cls() + else: + logger.warning("Couldn't find an optimal sampling param for %s. Using the default sampling param.", model_path) + sampling_param = cls() + + return sampling_param \ No newline at end of file diff --git a/fastvideo/v1/configs/sample/hunyuan.py b/fastvideo/v1/configs/sample/hunyuan.py new file mode 100644 index 000000000..01bc53cc6 --- /dev/null +++ b/fastvideo/v1/configs/sample/hunyuan.py @@ -0,0 +1,18 @@ +from dataclasses import dataclass + +from fastvideo.v1.configs.sample.base import SamplingParam + +@dataclass +class HunyuanSamplingParam(SamplingParam): + num_inference_steps: int = 50 + + num_frames: int = 125 + height: int = 720 + width: int = 1280 + fps: int = 24 + + guidance_scale: float = 1.0 + +@dataclass +class FastHunyuanSamplingParam(HunyuanSamplingParam): + num_inference_steps: int = 6 \ No newline at end of file diff --git a/fastvideo/v1/configs/sample/registry.py b/fastvideo/v1/configs/sample/registry.py new file mode 100644 index 000000000..6eb4cac87 --- /dev/null +++ b/fastvideo/v1/configs/sample/registry.py @@ -0,0 +1,73 @@ +import os +from typing import Callable, Dict, Optional, Any + +from fastvideo.v1.configs.sample.hunyuan import HunyuanSamplingParam, FastHunyuanSamplingParam +from fastvideo.v1.configs.sample.wan import WanT2V480PSamplingParam, WanI2V480PSamplingParam + +from fastvideo.v1.logger import init_logger +from fastvideo.v1.utils import (maybe_download_model_index, + verify_model_config_and_directory) + +logger = init_logger(__name__) +# Registry maps specific model weights to their config classes +SAMPLING_PARAM_REGISTRY: Dict[str, Any] = { + "FastVideo/FastHunyuan-Diffusers": FastHunyuanSamplingParam, + "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam, + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PSamplingParam, + "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PSamplingParam + # Add other specific weight variants +} + +# For determining pipeline type from model ID +SAMPLING_PARAM_DETECTOR: Dict[str, Callable[[str], bool]] = { + "hunyuan": lambda id: "hunyuan" in id.lower(), + "wanpipeline": lambda id: "wanpipeline" in id.lower(), + "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(), + # Add other pipeline architecture detectors +} + +# Fallback configs when exact match isn't found but architecture is detected +SAMPLING_FALLBACK_PARAM: Dict[str, Any] = { + "hunyuan": + HunyuanSamplingParam, # Base Hunyuan config as fallback for any Hunyuan variant + "wanpipeline": + WanT2V480PSamplingParam, # Base Wan config as fallback for any Wan variant + "wanimagetovideo": WanI2V480PSamplingParam, + # Other fallbacks by architecture +} + + +def get_sampling_param_cls_for_name( + pipeline_name_or_path: str) -> Optional[Any]: + """Get the appropriate sampling param for specific pretrained weights.""" + + if os.path.exists(pipeline_name_or_path): + config = verify_model_config_and_directory(pipeline_name_or_path) + logger.warning( + "FastVideo may not correctly identify the optimal sampling param for this model, as the local directory may have been renamed." + ) + else: + config = maybe_download_model_index(pipeline_name_or_path) + + pipeline_name = config["_class_name"] + + # First try exact match for specific weights + if pipeline_name_or_path in SAMPLING_PARAM_REGISTRY: + return SAMPLING_PARAM_REGISTRY[pipeline_name_or_path] + + # Try partial matches (for local paths that might include the weight ID) + for registered_id, config_class in SAMPLING_PARAM_REGISTRY.items(): + if registered_id in pipeline_name_or_path: + return config_class + + # If no match, try to use the fallback config + fallback_config = None + # Try to determine pipeline architecture for fallback + for pipeline_type, detector in SAMPLING_PARAM_DETECTOR.items(): + if detector(pipeline_name.lower()): + fallback_config = SAMPLING_FALLBACK_PARAM.get(pipeline_type) + break + + logger.warning("No match found for pipeline %s, using fallback sampling param %s.", + pipeline_name_or_path, fallback_config) + return fallback_config \ No newline at end of file diff --git a/fastvideo/v1/configs/sample/wan.py b/fastvideo/v1/configs/sample/wan.py new file mode 100644 index 000000000..8ae8b2861 --- /dev/null +++ b/fastvideo/v1/configs/sample/wan.py @@ -0,0 +1,22 @@ +from dataclasses import dataclass + +from fastvideo.v1.configs.sample.base import SamplingParam + +@dataclass +class WanT2V480PSamplingParam(SamplingParam): + # Video parameters + height: int = 480 + width: int = 832 + num_frames: int = 81 + fps: int = 16 + + # Denoising stage + guidance_scale: float = 3.0 + negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" + num_inference_steps: int = 50 + +@dataclass +class WanI2V480PSamplingParam(WanT2V480PSamplingParam): + # Denoising stage + guidance_scale: float = 5.0 + num_inference_steps: int = 40 \ No newline at end of file diff --git a/fastvideo/v1/configs/wan_1.3B_t2v_pipeline.json b/fastvideo/v1/configs/wan_1.3B_t2v_pipeline.json index d8f3019a4..46ba70374 100644 --- a/fastvideo/v1/configs/wan_1.3B_t2v_pipeline.json +++ b/fastvideo/v1/configs/wan_1.3B_t2v_pipeline.json @@ -1,12 +1,4 @@ { - "height": 480, - "width": 832, - "num_frames": 81, - "fps": 16, - "num_inference_steps": 50, - "guidance_scale": 3.0, - "seed": 1024, - "guidance_rescale": 0.0, "embedded_cfg_scale": 6.0, "flow_shift": 3, "use_cpu_offload": true, @@ -30,12 +22,10 @@ "use_parallel_tiling": false, "use_feature_cache": true }, - "num_channels_latents": null, "image_encoder_precision": "fp32", "text_encoder_precision": "fp32", "text_len": 512, "hidden_state_skip_layer": 0, "mask_strategy_file_path": null, - "enable_torch_compile": false, - "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" + "enable_torch_compile": false } \ No newline at end of file diff --git a/fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json b/fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json index 61d57e6ae..aba468a1f 100644 --- a/fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json +++ b/fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json @@ -1,12 +1,4 @@ { - "height": 480, - "width": 832, - "num_frames": 81, - "fps": 16, - "num_inference_steps": 40, - "guidance_scale": 5.0, - "seed": 1024, - "guidance_rescale": 0.0, "embedded_cfg_scale": 6.0, "flow_shift": 3, "use_cpu_offload": true, @@ -30,12 +22,10 @@ "use_parallel_tiling": false, "use_feature_cache": true }, - "num_channels_latents": null, "image_encoder_precision": "fp32", "text_encoder_precision": "fp32", "text_len": 512, "hidden_state_skip_layer": 0, "mask_strategy_file_path": null, - "enable_torch_compile": false, - "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" + "enable_torch_compile": false } \ No newline at end of file diff --git a/fastvideo/v1/entrypoints/video_generator.py b/fastvideo/v1/entrypoints/video_generator.py index 1843805c3..c35cfd79c 100644 --- a/fastvideo/v1/entrypoints/video_generator.py +++ b/fastvideo/v1/entrypoints/video_generator.py @@ -15,8 +15,10 @@ import numpy as np import torch import torchvision from einops import rearrange +from dataclasses import asdict from fastvideo.v1.configs.pipelines import get_pipeline_config_cls_for_name, BaseConfig +from fastvideo.v1.configs.sample import SamplingParam from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.logger import init_logger from fastvideo.v1.pipelines import ForwardBatch @@ -123,20 +125,8 @@ class VideoGenerator: def generate_video( self, prompt: str, - image_path: Optional[str] = None, - negative_prompt: Optional[str] = None, - output_path: Optional[str] = None, - save_video: bool = True, - return_frames: bool = False, - num_inference_steps: Optional[int] = None, - guidance_scale: Optional[float] = None, - num_frames: Optional[int] = None, - height: Optional[int] = None, - width: Optional[int] = None, - fps: Optional[int] = None, - seed: Optional[int] = None, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, + sampling_param: SamplingParam = None, + **kwargs, ) -> Union[Dict[str, Any], List[np.ndarray]]: """ Generate a video based on the given prompt. @@ -163,69 +153,54 @@ class VideoGenerator: # Create a copy of inference args to avoid modifying the original fastvideo_args = self.fastvideo_args - # Override parameters if provided - if image_path is not None: - fastvideo_args.image_path = image_path - if negative_prompt is not None: - fastvideo_args.neg_prompt = negative_prompt - if num_inference_steps is not None: - fastvideo_args.num_inference_steps = num_inference_steps - if guidance_scale is not None: - fastvideo_args.guidance_scale = guidance_scale - if num_frames is not None: - fastvideo_args.num_frames = num_frames - if height is not None: - fastvideo_args.height = height - if width is not None: - fastvideo_args.width = width - if fps is not None: - fastvideo_args.fps = fps - if seed is not None: - fastvideo_args.seed = seed - # 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) + kwargs["prompt"] = prompt + sampling_param.update(kwargs) + # Process negative prompt - if fastvideo_args.neg_prompt is not None: - fastvideo_args.neg_prompt = fastvideo_args.neg_prompt.strip() + if sampling_param.negative_prompt is not None: + sampling_param.negative_prompt = sampling_param.negative_prompt.strip() # Validate dimensions - if (fastvideo_args.height <= 0 or fastvideo_args.width <= 0 - or fastvideo_args.num_frames <= 0): + 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={fastvideo_args.height}, width={fastvideo_args.width}, " - f"num_frames={fastvideo_args.num_frames}") + f"height={sampling_param.height}, width={sampling_param.width}, " + f"num_frames={sampling_param.num_frames}") - if (fastvideo_args.num_frames - 1) % 4 != 0: + if (sampling_param.num_frames - 1) % fastvideo_args.vae_config.arch_config.temporal_compression_ratio != 0: raise ValueError( - f"num_frames-1 must be a multiple of 4, got {fastvideo_args.num_frames}" + f"num_frames-1 must be a multiple of {fastvideo_args.vae_config.arch_config.temporal_compression_ratio}, got {sampling_param.num_frames}" ) # Calculate sizes - target_height = align_to(fastvideo_args.height, 16) - target_width = align_to(fastvideo_args.width, 16) + target_height = align_to(sampling_param.height, 16) + target_width = align_to(sampling_param.width, 16) # Calculate latent sizes - latents_size = [(fastvideo_args.num_frames - 1) // 4 + 1, - fastvideo_args.height // 8, fastvideo_args.width // 8] + 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: {fastvideo_args.num_frames} + video_length: {sampling_param.num_frames} prompt: {prompt} - 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} + 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.flow_shift} embedded_guidance_scale: {fastvideo_args.embedded_cfg_scale}""" @@ -234,19 +209,9 @@ class VideoGenerator: # Prepare batch device = torch.device(fastvideo_args.device_str) batch = ForwardBatch( - prompt=prompt, - image_path=fastvideo_args.image_path, - 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, + **asdict(sampling_param), eta=0.0, n_tokens=n_tokens, - data_type="video" if fastvideo_args.num_frames > 1 else "image", - device=device, extra={}, ) @@ -267,23 +232,23 @@ class VideoGenerator: frames.append((x * 255).numpy().astype(np.uint8)) # Save video if requested - if save_video: - save_path = output_path or fastvideo_args.output_path + if batch.save_video: + save_path = batch.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=fastvideo_args.fps) + imageio.mimsave(video_path, frames, fps=batch.fps) logger.info("Saved video to %s", video_path) else: logger.warning("No output path provided, video not saved") - if return_frames: + if batch.return_frames: return frames else: return { "samples": samples, "prompts": prompt, "size": - (target_height, target_width, fastvideo_args.num_frames), + (target_height, target_width, batch.num_frames), "generation_time": gen_time } diff --git a/fastvideo/v1/examples/inference/basic/basic.py b/fastvideo/v1/examples/inference/basic/basic.py index db8aaf4c5..366a2f66d 100644 --- a/fastvideo/v1/examples/inference/basic/basic.py +++ b/fastvideo/v1/examples/inference/basic/basic.py @@ -1,4 +1,5 @@ from fastvideo import VideoGenerator +# from fastvideo.v1.configs.sample import SamplingParam def main(): # FastVideo will automatically use the optimal default arguments for the @@ -11,9 +12,13 @@ def main(): num_gpus=4, ) + # sampling_param = SamplingParam.from_pretrained("/workspace/data/Wan-AI/Wan2.1-I2V-14B-480P-Diffusers") + # sampling_param.num_frames = 45 + # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" # Generate videos with the same simple API, regardless of GPU count prompt = "A beautiful woman in a red dress walking down a street" video = generator.generate_video(prompt) + # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/") # Generate another video with a different prompt, without reloading the # model! diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py index c811fee4d..2845fa3a7 100644 --- a/fastvideo/v1/fastvideo_args.py +++ b/fastvideo/v1/fastvideo_args.py @@ -36,12 +36,6 @@ class FastVideoArgs: dist_timeout: Optional[int] = None # timeout for torch.distributed # Video generation parameters - height: int = 720 - width: int = 1280 - num_frames: int = 117 - num_inference_steps: int = 50 - guidance_scale: float = 1.0 - guidance_rescale: float = 0.0 embedded_cfg_scale: float = 6.0 flow_shift: Optional[float] = None @@ -57,9 +51,6 @@ class FastVideoArgs: # vae_scale_factor: Optional[int] = None # Deprecated vae_config: VAEConfig = VAEConfig() - # DiT configuration - num_channels_latents: Optional[int] = None - # Image encoder configuration image_encoder_precision: str = "fp32" @@ -72,20 +63,10 @@ class FastVideoArgs: text_encoder_precision_2: str = "fp16" text_len_2: int = 77 - # Flow Matching parameters - flow_solver: str = "euler" - denoise_type: str = "flow" # Deprecated. Will use scheduler_config.json - # STA (Spatial-Temporal Attention) parameters mask_strategy_file_path: Optional[str] = None enable_torch_compile: bool = False - # Scheduler options - scheduler_type: str = "euler" # Deprecated. Will use the param in scheduler_config.json - - neg_prompt: Optional[str] = None - num_videos: int = 1 - fps: int = 24 use_cpu_offload: bool = False disable_autocast: bool = False @@ -93,11 +74,6 @@ class FastVideoArgs: log_level: str = "info" # Inference parameters - image_path: Optional[str] = None - prompt: Optional[str] = None - prompt_path: Optional[str] = None - output_path: str = "outputs/" - seed: int = 1024 device_str: Optional[str] = None device = None @@ -177,43 +153,6 @@ class FastVideoArgs: help="Set timeout for torch.distributed initialization.", ) - # Video generation parameters - parser.add_argument( - "--height", - type=int, - default=FastVideoArgs.height, - help="Height of generated video", - ) - parser.add_argument( - "--width", - type=int, - default=FastVideoArgs.width, - help="Width of generated video", - ) - parser.add_argument( - "--num-frames", - type=int, - default=FastVideoArgs.num_frames, - help="Number of frames to generate", - ) - parser.add_argument( - "--num-inference-steps", - type=int, - default=FastVideoArgs.num_inference_steps, - help="Number of inference steps", - ) - parser.add_argument( - "--guidance-scale", - type=float, - default=FastVideoArgs.guidance_scale, - help="Guidance scale for classifier-free guidance", - ) - parser.add_argument( - "--guidance-rescale", - type=float, - default=FastVideoArgs.guidance_rescale, - help="Guidance rescale for classifier-free guidance", - ) parser.add_argument( "--embedded-cfg-scale", type=float, @@ -302,20 +241,6 @@ class FastVideoArgs: help="Maximum secondary text length", ) - # Flow Matching parameters - parser.add_argument( - "--flow-solver", - type=str, - default=FastVideoArgs.flow_solver, - help="Solver for flow matching", - ) - parser.add_argument( - "--denoise-type", - type=str, - default=FastVideoArgs.denoise_type, - help="Denoise type for noised inputs", - ) - # STA (Spatial-Temporal Attention) parameters parser.add_argument( "--mask-strategy-file-path", @@ -329,33 +254,6 @@ class FastVideoArgs: "Use torch.compile for speeding up STA inference without teacache", ) - # Scheduler options - parser.add_argument( - "--scheduler-type", - type=str, - default=FastVideoArgs.scheduler_type, - help="Type of scheduler to use", - ) - - # HunYuan specific parameters - parser.add_argument( - "--neg-prompt", - type=str, - default=FastVideoArgs.neg_prompt, - help="Negative prompt for sampling", - ) - parser.add_argument( - "--num-videos", - type=int, - default=FastVideoArgs.num_videos, - help="Number of videos to generate per prompt", - ) - parser.add_argument( - "--fps", - type=int, - default=FastVideoArgs.fps, - help="Frames per second for output video", - ) parser.add_argument( "--use-cpu-offload", action="store_true", @@ -376,36 +274,6 @@ class FastVideoArgs: help="The logging level of all loggers.", ) - # Inference parameters - prompt_group = parser.add_mutually_exclusive_group(required=True) - prompt_group.add_argument( - "--prompt", - type=str, - help="Text prompt for video generation", - ) - prompt_group.add_argument( - "--prompt-path", - type=str, - help="Path to a text file containing the prompt", - ) - - parser.add_argument("--image-path", - type=str, - help="Path to the image for I2V generation") - - parser.add_argument( - "--output-path", - type=str, - default=FastVideoArgs.output_path, - help="Directory to save generated videos", - ) - parser.add_argument( - "--seed", - type=int, - default=FastVideoArgs.seed, - help="Random seed for reproducibility", - ) - return parser @classmethod @@ -454,8 +322,6 @@ class FastVideoArgs: raise ValueError( "Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True." ) - if self.prompt_path and not self.prompt_path.endswith(".txt"): - raise ValueError("prompt_path must be a text file") _current_fastvideo_args = None diff --git a/fastvideo/v1/models/dits/base.py b/fastvideo/v1/models/dits/base.py index a17d30788..410b27ed0 100644 --- a/fastvideo/v1/models/dits/base.py +++ b/fastvideo/v1/models/dits/base.py @@ -15,6 +15,7 @@ class BaseDiT(nn.Module, ABC): _param_names_mapping: dict hidden_size: int num_attention_heads: int + num_channels_latents: int # always supports torch_sdpa _supported_attention_backends: Tuple[_Backend, ...] = (_Backend.TORCH_SDPA, ) @@ -49,7 +50,7 @@ class BaseDiT(nn.Module, ABC): pass def __post_init__(self) -> None: - required_attrs = ["hidden_size", "num_attention_heads"] + required_attrs = ["hidden_size", "num_attention_heads", "num_channels_latents"] for attr in required_attrs: if not hasattr(self, attr): raise AttributeError( diff --git a/fastvideo/v1/models/dits/hunyuanvideo.py b/fastvideo/v1/models/dits/hunyuanvideo.py index a97a104b6..fb7f99ed5 100644 --- a/fastvideo/v1/models/dits/hunyuanvideo.py +++ b/fastvideo/v1/models/dits/hunyuanvideo.py @@ -578,6 +578,7 @@ class HunyuanVideoTransformer3DModel(BaseDiT): hidden_size = attention_head_dim * num_attention_heads self.patch_size = [patch_size_t, patch_size, patch_size] self.in_channels = in_channels + self.num_channels_latents = in_channels self.out_channels = in_channels if out_channels is None else out_channels self.unpatchify_channels = self.out_channels self.guidance_embeds = guidance_embeds diff --git a/fastvideo/v1/models/dits/wanvideo.py b/fastvideo/v1/models/dits/wanvideo.py index dc5b262e2..2a2fcc515 100644 --- a/fastvideo/v1/models/dits/wanvideo.py +++ b/fastvideo/v1/models/dits/wanvideo.py @@ -420,6 +420,7 @@ class WanTransformer3DModel(BaseDiT): self.num_attention_heads = num_attention_heads self.in_channels = in_channels self.out_channels = out_channels or in_channels + self.num_channels_latents = in_channels if added_kv_proj_dim is None else out_channels self.patch_size = patch_size self.text_len = text_len diff --git a/fastvideo/v1/pipelines/composed_pipeline_base.py b/fastvideo/v1/pipelines/composed_pipeline_base.py index 84ebb50ab..beab21fae 100644 --- a/fastvideo/v1/pipelines/composed_pipeline_base.py +++ b/fastvideo/v1/pipelines/composed_pipeline_base.py @@ -114,12 +114,11 @@ class ComposedPipelineBase(ABC): """ raise NotImplementedError - @abstractmethod def initialize_pipeline(self, fastvideo_args: FastVideoArgs): """ Initialize the pipeline. """ - raise NotImplementedError + pass def load_modules(self, fastvideo_args: FastVideoArgs) -> Dict[str, Any]: """ diff --git a/fastvideo/v1/pipelines/hunyuan/hunyuan_pipeline.py b/fastvideo/v1/pipelines/hunyuan/hunyuan_pipeline.py index 3b70ef11e..7846900ef 100644 --- a/fastvideo/v1/pipelines/hunyuan/hunyuan_pipeline.py +++ b/fastvideo/v1/pipelines/hunyuan/hunyuan_pipeline.py @@ -55,7 +55,8 @@ class HunyuanVideoPipeline(ComposedPipelineBase): self.add_stage(stage_name="latent_preparation_stage", stage=LatentPreparationStage( - scheduler=self.get_module("scheduler"))) + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"))) self.add_stage(stage_name="denoising_stage", stage=DenoisingStage( @@ -65,12 +66,5 @@ class HunyuanVideoPipeline(ComposedPipelineBase): self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))) - def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - """ - Initialize the pipeline. - """ - num_channels_latents = self.get_module("transformer").in_channels - fastvideo_args.num_channels_latents = num_channels_latents - EntryClass = HunyuanVideoPipeline diff --git a/fastvideo/v1/pipelines/pipeline_batch_info.py b/fastvideo/v1/pipelines/pipeline_batch_info.py index d5a2e6723..7c5e98e6c 100644 --- a/fastvideo/v1/pipelines/pipeline_batch_info.py +++ b/fastvideo/v1/pipelines/pipeline_batch_info.py @@ -7,7 +7,7 @@ This module defines the dataclasses used to pass state between pipeline componen in a functional manner, reducing the need for explicit parameter passing. """ -from dataclasses import dataclass, field +from dataclasses import dataclass, field, fields from typing import Any, Dict, List, Optional, Union import torch @@ -36,6 +36,8 @@ class ForwardBatch: # Text inputs prompt: Optional[Union[str, List[str]]] = None negative_prompt: Optional[Union[str, List[str]]] = None + prompt_path: Optional[str] = None + output_path: str = "outputs/" # Primary encoder embeddings prompt_embeds: List[torch.Tensor] = field(default_factory=list) @@ -49,6 +51,7 @@ class ForwardBatch: # Batch info batch_size: Optional[int] = None num_videos_per_prompt: int = 1 + seed: Optional[int] = None seeds: Optional[List[int]] = None # Tracking if embeddings are already processed @@ -60,7 +63,6 @@ class ForwardBatch: image_latent: Optional[torch.Tensor] = None # Latent dimensions - num_channels_latents: Optional[int] = None height_latents: Optional[int] = None width_latents: Optional[int] = None num_frames: int = 1 # Default for image models @@ -68,6 +70,7 @@ class ForwardBatch: # Original dimensions (before VAE scaling) height: Optional[int] = None width: Optional[int] = None + fps: Optional[int] = None # Timesteps timesteps: Optional[torch.Tensor] = None @@ -95,7 +98,9 @@ class ForwardBatch: # Extra parameters that might be needed by specific pipeline implementations extra: Dict[str, Any] = field(default_factory=dict) - device: torch.device = field(default_factory=lambda: torch.device("cuda")) + # Misc + save_video: bool = True + return_frames: bool = False def __post_init__(self): """Initialize dependent fields after dataclass initialization.""" diff --git a/fastvideo/v1/pipelines/stages/clip_image_encoding.py b/fastvideo/v1/pipelines/stages/clip_image_encoding.py index bc91e6b55..6b89fa677 100644 --- a/fastvideo/v1/pipelines/stages/clip_image_encoding.py +++ b/fastvideo/v1/pipelines/stages/clip_image_encoding.py @@ -53,12 +53,12 @@ class CLIPImageEncodingStage(PipelineStage): The batch with encoded prompt embeddings. """ if fastvideo_args.use_cpu_offload: - self.image_encoder = self.image_encoder.to(batch.device) + self.image_encoder = self.image_encoder.to(fastvideo_args.device) image = load_image(batch.image_path) image_inputs = self.image_processor( - images=image, return_tensors="pt").to(batch.device) + images=image, return_tensors="pt").to(fastvideo_args.device) with set_forward_context(current_timestep=0, attn_metadata=None): image_embeds = self.image_encoder(**image_inputs) diff --git a/fastvideo/v1/pipelines/stages/clip_text_encoding.py b/fastvideo/v1/pipelines/stages/clip_text_encoding.py index 0e74630b6..3aed025ec 100644 --- a/fastvideo/v1/pipelines/stages/clip_text_encoding.py +++ b/fastvideo/v1/pipelines/stages/clip_text_encoding.py @@ -52,7 +52,7 @@ class CLIPTextEncodingStage(PipelineStage): The batch with encoded prompt embeddings. """ if fastvideo_args.use_cpu_offload: - self.text_encoder = self.text_encoder.to(batch.device) + self.text_encoder = self.text_encoder.to(fastvideo_args.device) text_inputs = self.tokenizer( batch.prompt, @@ -63,7 +63,7 @@ class CLIPTextEncodingStage(PipelineStage): ) with set_forward_context(current_timestep=0, attn_metadata=None): outputs = self.text_encoder(input_ids=text_inputs["input_ids"].to( - batch.device), ) + fastvideo_args.device), ) prompt_embeds = outputs["pooler_output"] batch.prompt_embeds.append(prompt_embeds) @@ -79,7 +79,7 @@ class CLIPTextEncodingStage(PipelineStage): with set_forward_context(current_timestep=0, attn_metadata=None): negative_outputs = self.text_encoder( input_ids=negative_text_inputs["input_ids"].to( - batch.device), ) + fastvideo_args.device), ) negative_prompt_embeds = negative_outputs["pooler_output"] assert batch.negative_prompt_embeds is not None diff --git a/fastvideo/v1/pipelines/stages/denoising.py b/fastvideo/v1/pipelines/stages/denoising.py index a66f8c7a8..ae9510074 100644 --- a/fastvideo/v1/pipelines/stages/denoising.py +++ b/fastvideo/v1/pipelines/stages/denoising.py @@ -66,7 +66,7 @@ class DenoisingStage(PipelineStage): """ # If use cpu offload, need to load the model back into gpu again if fastvideo_args.use_cpu_offload: - self.transformer = self.transformer.to(batch.device) + self.transformer = self.transformer.to(fastvideo_args.device) # Prepare extra step kwargs for scheduler extra_step_kwargs = self.prepare_extra_func_kwargs( self.scheduler.step, @@ -165,7 +165,7 @@ class DenoisingStage(PipelineStage): [fastvideo_args.embedded_cfg_scale] * latent_model_input.shape[0], dtype=torch.float32, - device=batch.device, + device=fastvideo_args.device, ).to(target_dtype) * 1000.0 if fastvideo_args.embedded_cfg_scale is not None else None) diff --git a/fastvideo/v1/pipelines/stages/encoding.py b/fastvideo/v1/pipelines/stages/encoding.py index 4fbe6d7bc..61fdd9dd6 100644 --- a/fastvideo/v1/pipelines/stages/encoding.py +++ b/fastvideo/v1/pipelines/stages/encoding.py @@ -58,16 +58,16 @@ class EncodingStage(PipelineStage): image, vae_scale_factor=self.vae.spatial_compression_ratio, height=batch.height, - width=batch.width).to(batch.device, dtype=torch.float32) + width=batch.width).to(fastvideo_args.device, dtype=torch.float32) image = image.unsqueeze(2) video_condition = torch.cat([ image, image.new_zeros(image.shape[0], image.shape[1], - fastvideo_args.num_frames - 1, batch.height, + batch.num_frames - 1, batch.height, batch.width) ], dim=2) - video_condition = video_condition.to(device=batch.device, + video_condition = video_condition.to(device=fastvideo_args.device, dtype=torch.float32) # Setup VAE precision @@ -107,9 +107,9 @@ class EncodingStage(PipelineStage): else: latent_condition = latent_condition * self.vae.scaling_factor - mask_lat_size = torch.ones(1, 1, fastvideo_args.num_frames, + mask_lat_size = torch.ones(1, 1, batch.num_frames, latent_height, latent_width) - mask_lat_size[:, :, list(range(1, fastvideo_args.num_frames))] = 0 + mask_lat_size[:, :, list(range(1, batch.num_frames))] = 0 first_frame_mask = mask_lat_size[:, :, 0:1] first_frame_mask = torch.repeat_interleave( first_frame_mask, diff --git a/fastvideo/v1/pipelines/stages/input_validation.py b/fastvideo/v1/pipelines/stages/input_validation.py index 4c3f24468..f4afad434 100644 --- a/fastvideo/v1/pipelines/stages/input_validation.py +++ b/fastvideo/v1/pipelines/stages/input_validation.py @@ -24,8 +24,8 @@ class InputValidationStage(PipelineStage): def _generate_seeds(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs): """Generate seeds for the inference""" - seed = fastvideo_args.seed - num_videos_per_prompt = fastvideo_args.num_videos + seed = batch.seed + num_videos_per_prompt = batch.num_videos_per_prompt seeds = [seed + i for i in range(num_videos_per_prompt)] batch.seeds = seeds @@ -85,12 +85,4 @@ class InputValidationStage(PipelineStage): f"Guidance scale must be positive, but got {batch.guidance_scale}" ) - # Set device if not already set - if batch.device is None: - batch.device = self.device - - # Set data type if not already set - if batch.data_type is None: - batch.data_type = fastvideo_args.precision - return batch diff --git a/fastvideo/v1/pipelines/stages/latent_preparation.py b/fastvideo/v1/pipelines/stages/latent_preparation.py index abb69d1d4..b6ea3d713 100644 --- a/fastvideo/v1/pipelines/stages/latent_preparation.py +++ b/fastvideo/v1/pipelines/stages/latent_preparation.py @@ -20,9 +20,10 @@ class LatentPreparationStage(PipelineStage): denoised during the diffusion process. """ - def __init__(self, scheduler) -> None: + def __init__(self, scheduler, transformer) -> None: super().__init__() self.scheduler = scheduler + self.transformer = transformer def forward( self, @@ -40,9 +41,10 @@ class LatentPreparationStage(PipelineStage): The batch with prepared latent variables. """ + latent_num_frames = None # Adjust video length based on VAE version if needed if hasattr(self, 'adjust_video_length'): - batch = self.adjust_video_length(batch, fastvideo_args) + latent_num_frames = self.adjust_video_length(batch, fastvideo_args) # Determine batch size if isinstance(batch.prompt, list): batch_size = len(batch.prompt) @@ -56,10 +58,10 @@ class LatentPreparationStage(PipelineStage): # Get required parameters dtype = batch.prompt_embeds[0].dtype - device = batch.device + device = fastvideo_args.device generator = batch.generator latents = batch.latents - num_frames = batch.num_frames + num_frames = latent_num_frames if latent_num_frames is not None else batch.num_frames height = batch.height width = batch.width @@ -67,12 +69,10 @@ class LatentPreparationStage(PipelineStage): if height is None or width is None: raise ValueError("Height and width must be provided") - assert fastvideo_args.num_channels_latents is not None - # Calculate latent shape shape = ( batch_size, - fastvideo_args.num_channels_latents, + self.transformer.num_channels_latents, num_frames, height // fastvideo_args.vae_config.arch_config.spatial_compression_ratio, width // fastvideo_args.vae_config.arch_config.spatial_compression_ratio, @@ -104,7 +104,7 @@ class LatentPreparationStage(PipelineStage): return batch def adjust_video_length(self, batch: ForwardBatch, - fastvideo_args: FastVideoArgs) -> ForwardBatch: + fastvideo_args: FastVideoArgs) -> int: """ Adjust video length based on VAE version. @@ -118,5 +118,5 @@ class LatentPreparationStage(PipelineStage): video_length = batch.num_frames temporal_scale_factor = fastvideo_args.vae_config.arch_config.temporal_compression_ratio # TODO - batch.num_frames = (video_length - 1) // temporal_scale_factor + 1 - return batch + latent_num_frames = (video_length - 1) // temporal_scale_factor + 1 + return latent_num_frames diff --git a/fastvideo/v1/pipelines/stages/llama_encoding.py b/fastvideo/v1/pipelines/stages/llama_encoding.py index a45fd0746..d7ef6c366 100644 --- a/fastvideo/v1/pipelines/stages/llama_encoding.py +++ b/fastvideo/v1/pipelines/stages/llama_encoding.py @@ -74,7 +74,7 @@ class LlamaEncodingStage(PipelineStage): The batch with encoded prompt embeddings. """ if fastvideo_args.use_cpu_offload: - self.text_encoder = self.text_encoder.to(batch.device) + self.text_encoder = self.text_encoder.to(fastvideo_args.device) text = prompt_template_video["template"].format(batch.prompt) text_inputs = self.tokenizer( @@ -87,7 +87,7 @@ class LlamaEncodingStage(PipelineStage): hidden_state_skip_layer = 2 with set_forward_context(current_timestep=0, attn_metadata=None): outputs = self.text_encoder( - input_ids=text_inputs["input_ids"].to(batch.device), + input_ids=text_inputs["input_ids"].to(fastvideo_args.device), output_hidden_states=hidden_state_skip_layer is not None, ) @@ -111,7 +111,7 @@ class LlamaEncodingStage(PipelineStage): with set_forward_context(current_timestep=0, attn_metadata=None): negative_outputs = self.text_encoder( input_ids=negative_text_inputs["input_ids"].to( - batch.device), + fastvideo_args.device), output_hidden_states=hidden_state_skip_layer is not None, ) diff --git a/fastvideo/v1/pipelines/stages/t5_encoding.py b/fastvideo/v1/pipelines/stages/t5_encoding.py index e90d335f0..97f250e7e 100644 --- a/fastvideo/v1/pipelines/stages/t5_encoding.py +++ b/fastvideo/v1/pipelines/stages/t5_encoding.py @@ -51,7 +51,7 @@ class T5EncodingStage(PipelineStage): The batch with encoded prompt embeddings. """ if fastvideo_args.use_cpu_offload: - self.text_encoder = self.text_encoder.to(batch.device) + self.text_encoder = self.text_encoder.to(fastvideo_args.device) text = batch.prompt text_inputs = self.tokenizer( @@ -62,7 +62,7 @@ class T5EncodingStage(PipelineStage): add_special_tokens=True, return_attention_mask=True, return_tensors="pt", - ).to(batch.device) + ).to(fastvideo_args.device) text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() with set_forward_context(current_timestep=0, attn_metadata=None): @@ -89,7 +89,7 @@ class T5EncodingStage(PipelineStage): add_special_tokens=True, return_attention_mask=True, return_tensors="pt", - ).to(batch.device) + ).to(fastvideo_args.device) text_input_ids, mask = negative_text_inputs.input_ids, negative_text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() with set_forward_context(current_timestep=0, attn_metadata=None): diff --git a/fastvideo/v1/pipelines/stages/timestep_preparation.py b/fastvideo/v1/pipelines/stages/timestep_preparation.py index cbdcaaa60..9b792109e 100644 --- a/fastvideo/v1/pipelines/stages/timestep_preparation.py +++ b/fastvideo/v1/pipelines/stages/timestep_preparation.py @@ -42,7 +42,7 @@ class TimestepPreparationStage(PipelineStage): The batch with prepared timesteps. """ scheduler = self.scheduler - device = batch.device + device = fastvideo_args.device num_inference_steps = batch.num_inference_steps timesteps = batch.timesteps sigmas = batch.sigmas diff --git a/fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py b/fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py index 4bfb2e266..c8817fe5e 100644 --- a/fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py +++ b/fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py @@ -54,7 +54,8 @@ class WanImageToVideoPipeline(ComposedPipelineBase): self.add_stage(stage_name="latent_preparation_stage", stage=LatentPreparationStage( - scheduler=self.get_module("scheduler"))) + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"))) self.add_stage(stage_name="image_latent_preparation_stage", stage=EncodingStage(vae=self.get_module("vae"))) @@ -67,12 +68,5 @@ class WanImageToVideoPipeline(ComposedPipelineBase): self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))) - def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - """ - Initialize the pipeline. - """ - num_channels_latents = self.get_module("transformer").out_channels - fastvideo_args.num_channels_latents = num_channels_latents - EntryClass = WanImageToVideoPipeline diff --git a/fastvideo/v1/pipelines/wan/wan_pipeline.py b/fastvideo/v1/pipelines/wan/wan_pipeline.py index 170e5606d..0dbf648a2 100644 --- a/fastvideo/v1/pipelines/wan/wan_pipeline.py +++ b/fastvideo/v1/pipelines/wan/wan_pipeline.py @@ -47,7 +47,8 @@ class WanPipeline(ComposedPipelineBase): self.add_stage(stage_name="latent_preparation_stage", stage=LatentPreparationStage( - scheduler=self.get_module("scheduler"))) + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"))) self.add_stage(stage_name="denoising_stage", stage=DenoisingStage( @@ -57,12 +58,5 @@ class WanPipeline(ComposedPipelineBase): self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))) - def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - """ - Initialize the pipeline. - """ - num_channels_latents = self.get_module("transformer").in_channels - fastvideo_args.num_channels_latents = num_channels_latents - EntryClass = WanPipeline diff --git a/fastvideo/v1/worker/gpu_worker.py b/fastvideo/v1/worker/gpu_worker.py index c09e4d1d5..d800af0ad 100644 --- a/fastvideo/v1/worker/gpu_worker.py +++ b/fastvideo/v1/worker/gpu_worker.py @@ -89,7 +89,6 @@ class Worker: def execute_forward(self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: - self.fastvideo_args.num_inference_steps = fastvideo_args.num_inference_steps output_batch = self.pipeline.forward(forward_batch, self.fastvideo_args) return cast(ForwardBatch, output_batch)