Separate FastVideoArgs from SamplingParam
This commit is contained in:
@@ -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()):
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
from fastvideo.v1.configs.sample.base import SamplingParam
|
||||
|
||||
__all__ = [
|
||||
"SamplingParam"
|
||||
]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
}
|
||||
@@ -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!
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user