Separate FastVideoArgs from SamplingParam

This commit is contained in:
JerryZhou54
2025-04-30 01:48:05 +00:00
parent 80450962ef
commit 81bc6c9943
32 changed files with 292 additions and 323 deletions
+16 -16
View File
@@ -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)
@@ -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
@@ -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()):
-10
View File
@@ -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"
+5
View File
@@ -0,0 +1,5 @@
from fastvideo.v1.configs.sample.base import SamplingParam
__all__ = [
"SamplingParam"
]
+68
View File
@@ -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
+18
View File
@@ -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
+73
View File
@@ -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
+22
View File
@@ -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
@@ -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
}
@@ -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
}
+33 -68
View File
@@ -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
}
@@ -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!
-134
View File
@@ -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
+2 -1
View File
@@ -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(
+1
View File
@@ -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
+1
View File
@@ -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
@@ -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]:
"""
@@ -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
@@ -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."""
@@ -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)
@@ -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
+2 -2
View File
@@ -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)
+5 -5
View File
@@ -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,
@@ -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
@@ -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
@@ -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,
)
+3 -3
View File
@@ -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):
@@ -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
@@ -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
+2 -8
View File
@@ -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
-1
View File
@@ -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)