nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Pass pre-commit tests

Browse Source
This commit is contained in:
JerryZhou54
2025-04-30 01:52:07 +00:00
parent e16222a5cc
commit 60f01567d1
49 changed files with 509 additions and 431 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/v1/configs/models/__init__.py
+4 -6
View File
@@ -1,11 +1,9 @@
from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
from fastvideo.v1.configs.models.encoders.base import EncoderConfig
from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
__all__ = [
"ArchConfig", "ModelConfig",
"VAEArchConfig", "VAEConfig",
"DiTArchConfig", "DiTConfig",
"EncoderConfig"
]
"ArchConfig", "ModelConfig", "VAEArchConfig", "VAEConfig", "DiTArchConfig",
"DiTConfig", "EncoderConfig"
]
fastvideo/v1/configs/models/base.py
+17 -19
View File
@@ -1,13 +1,15 @@
from dataclasses import dataclass, fields
from typing import Dict, Any
from typing import Any, Dict
# 1. ArchConfig contains all fields from diffuser's/transformer's config.json (i.e. all fields related to the architecture of the model)
# 2. ArchConfig should be inherited & overriden by each model arch_config
# 2. ArchConfig should be inherited & overridden by each model arch_config
# 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users
@dataclass
class ArchConfig:
pass
@dataclass
class ModelConfig:
# Every model config parameter can be categorized into either ArchConfig or everything else
@@ -21,31 +23,27 @@ class ModelConfig:
# Only called if 'name' is not found in ModelConfig directly
if hasattr(self.arch_config, name):
return getattr(self.arch_config, name)
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
raise AttributeError(
f"'{type(self).__name__}' object has no attribute '{name}'")
# This should be used only when loading from transformers/diffusers
def update_model_arch(
self,
source_model_dict: Dict[str, Any]
) -> None:
def update_model_arch(self, source_model_dict: Dict[str, Any]) -> None:
arch_config = self.arch_config
valid_fields = {f.name for f in fields(arch_config)}
for key, value in source_model_dict.items():
if key in valid_fields:
setattr(arch_config, key, value)
else:
raise AttributeError(f"{type(arch_config).__name__} has no field '{key}'")
raise AttributeError(
f"{type(arch_config).__name__} has no field '{key}'")
if hasattr(arch_config, "__post_init__"):
arch_config.__post_init__()
def update_model_config(
self,
source_model_dict: Dict[str, Any]
) -> None:
assert "arch_config" not in source_model_dict.keys(), "Source model config shouldn't contain arch_config."
def update_model_config(self, source_model_dict: Dict[str, Any]) -> None:
assert "arch_config" not in source_model_dict, "Source model config shouldn't contain arch_config."
valid_fields = {f.name for f in fields(self)}
for key, value in source_model_dict.items():
@@ -54,6 +52,6 @@ class ModelConfig:
else:
print(f"{type(self).__name__} does not contain field '{key}'!")
raise AttributeError(f"Invalid field: {key}")
if hasattr(self, "__post_init__"):
self.__post_init__()
self.__post_init__()
fastvideo/v1/configs/models/dits/__init__.py
+7 -5
View File
@@ -1,7 +1,9 @@
from fastvideo.v1.configs.models.dits.hunyuanvideo import HunyuanVideoArchConfig, HunyuanVideoConfig
from fastvideo.v1.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig
from fastvideo.v1.configs.models.dits.hunyuanvideo import (
HunyuanVideoArchConfig, HunyuanVideoConfig)
from fastvideo.v1.configs.models.dits.wanvideo import (WanVideoArchConfig,
WanVideoConfig)
__all__ = [
"HunyuanVideoArchConfig", "HunyuanVideoConfig",
"WanVideoArchConfig", "WanVideoConfig"
]
"HunyuanVideoArchConfig", "HunyuanVideoConfig", "WanVideoArchConfig",
"WanVideoConfig"
]
fastvideo/v1/configs/models/dits/base.py
+5 -3
View File
@@ -1,9 +1,10 @@
from dataclasses import dataclass, field
from typing import Tuple, Optional
from typing import Optional, Tuple
from fastvideo.v1.configs.models import ArchConfig, ModelConfig
from fastvideo.v1.platforms import _Backend
from fastvideo.v1.configs.quantization import QuantizationConfig
from fastvideo.v1.platforms import _Backend
@dataclass
class DiTArchConfig(ArchConfig):
@@ -11,11 +12,12 @@ class DiTArchConfig(ArchConfig):
_param_names_mapping: dict = field(default_factory=dict)
_supported_attention_backends: Tuple[_Backend,
...] = (_Backend.TORCH_SDPA, )
hidden_size: int = 0
num_attention_heads: int = 0
num_channels_latents: int = 0
@dataclass
class DiTConfig(ModelConfig):
arch_config: DiTArchConfig = DiTArchConfig()
fastvideo/v1/configs/models/dits/hunyuanvideo.py
+124 -119
View File
@@ -1,144 +1,148 @@
from dataclasses import dataclass, field
from typing import Tuple, Optional
from typing import Optional, Tuple
import torch
from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
from fastvideo.v1.platforms import _Backend
def is_double_block(n: str, m) -> bool:
return "double" in n and str.isdigit(n.split(".")[-1])
def is_single_block(n: str, m) -> bool:
return "single" in n and str.isdigit(n.split(".")[-1])
def is_refiner_block(n: str, m) -> bool:
return "refiner" in n and str.isdigit(n.split(".")[-1])
@dataclass
class HunyuanVideoArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(default_factory=lambda: [
is_double_block,
is_single_block,
is_refiner_block
])
_supported_attention_backends: Tuple[_Backend, ...] = field(default_factory=lambda: (_Backend.SLIDING_TILE_ATTN,
_Backend.FLASH_ATTN, _Backend.TORCH_SDPA))
_param_names_mapping: dict = field(default_factory=lambda: {
# 1. context_embedder.time_text_embed submodules (specific rules, applied first):
r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
r"txt_in.t_embedder.mlp.fc_in.\1",
r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_2\.(.*)$":
r"txt_in.t_embedder.mlp.fc_out.\1",
r"^context_embedder\.proj_in\.(.*)$":
r"txt_in.input_embedder.\1",
r"^context_embedder\.time_text_embed\.text_embedder\.linear_1\.(.*)$":
r"txt_in.c_embedder.fc_in.\1",
r"^context_embedder\.time_text_embed\.text_embedder\.linear_2\.(.*)$":
r"txt_in.c_embedder.fc_out.\1",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm1\.(.*)$":
r"txt_in.refiner_blocks.\1.norm1.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm2\.(.*)$":
r"txt_in.refiner_blocks.\1.norm2.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_q\.(.*)$":
(r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 0, 3),
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_k\.(.*)$":
(r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 1, 3),
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_v\.(.*)$":
(r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 2, 3),
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$":
r"txt_in.refiner_blocks.\1.self_attn_proj.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$":
r"txt_in.refiner_blocks.\1.mlp.fc_in.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$":
r"txt_in.refiner_blocks.\1.mlp.fc_out.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm_out\.linear\.(.*)$":
r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2",
_fsdp_shard_conditions: list = field(
default_factory=lambda:
[is_double_block, is_single_block, is_refiner_block])
_supported_attention_backends: Tuple[_Backend, ...] = field(
default_factory=lambda:
(_Backend.SLIDING_TILE_ATTN, _Backend.FLASH_ATTN, _Backend.TORCH_SDPA))
# 3. x_embedder mapping:
r"^x_embedder\.proj\.(.*)$":
r"img_in.proj.\1",
_param_names_mapping: dict = field(
default_factory=lambda: {
# 1. context_embedder.time_text_embed submodules (specific rules, applied first):
r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
r"txt_in.t_embedder.mlp.fc_in.\1",
r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_2\.(.*)$":
r"txt_in.t_embedder.mlp.fc_out.\1",
r"^context_embedder\.proj_in\.(.*)$":
r"txt_in.input_embedder.\1",
r"^context_embedder\.time_text_embed\.text_embedder\.linear_1\.(.*)$":
r"txt_in.c_embedder.fc_in.\1",
r"^context_embedder\.time_text_embed\.text_embedder\.linear_2\.(.*)$":
r"txt_in.c_embedder.fc_out.\1",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm1\.(.*)$":
r"txt_in.refiner_blocks.\1.norm1.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm2\.(.*)$":
r"txt_in.refiner_blocks.\1.norm2.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_q\.(.*)$":
(r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 0, 3),
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_k\.(.*)$":
(r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 1, 3),
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_v\.(.*)$":
(r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 2, 3),
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$":
r"txt_in.refiner_blocks.\1.self_attn_proj.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$":
r"txt_in.refiner_blocks.\1.mlp.fc_in.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$":
r"txt_in.refiner_blocks.\1.mlp.fc_out.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm_out\.linear\.(.*)$":
r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2",
# 4. Top-level time_text_embed mappings:
r"^time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
r"time_in.mlp.fc_in.\1",
r"^time_text_embed\.timestep_embedder\.linear_2\.(.*)$":
r"time_in.mlp.fc_out.\1",
r"^time_text_embed\.guidance_embedder\.linear_1\.(.*)$":
r"guidance_in.mlp.fc_in.\1",
r"^time_text_embed\.guidance_embedder\.linear_2\.(.*)$":
r"guidance_in.mlp.fc_out.\1",
r"^time_text_embed\.text_embedder\.linear_1\.(.*)$":
r"vector_in.fc_in.\1",
r"^time_text_embed\.text_embedder\.linear_2\.(.*)$":
r"vector_in.fc_out.\1",
# 3. x_embedder mapping:
r"^x_embedder\.proj\.(.*)$":
r"img_in.proj.\1",
# 5. transformer_blocks mapping:
r"^transformer_blocks\.(\d+)\.norm1\.linear\.(.*)$":
r"double_blocks.\1.img_mod.linear.\2",
r"^transformer_blocks\.(\d+)\.norm1_context\.linear\.(.*)$":
r"double_blocks.\1.txt_mod.linear.\2",
r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$":
r"double_blocks.\1.img_attn_q_norm.\2",
r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$":
r"double_blocks.\1.img_attn_k_norm.\2",
r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$":
(r"double_blocks.\1.img_attn_qkv.\2", 0, 3),
r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$":
(r"double_blocks.\1.img_attn_qkv.\2", 1, 3),
r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$":
(r"double_blocks.\1.img_attn_qkv.\2", 2, 3),
r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$":
(r"double_blocks.\1.txt_attn_qkv.\2", 0, 3),
r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$":
(r"double_blocks.\1.txt_attn_qkv.\2", 1, 3),
r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$":
(r"double_blocks.\1.txt_attn_qkv.\2", 2, 3),
r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$":
r"double_blocks.\1.img_attn_proj.\2",
# Corrected: merge attn.to_add_out into the main projection.
r"^transformer_blocks\.(\d+)\.attn\.to_add_out\.(.*)$":
r"double_blocks.\1.txt_attn_proj.\2",
r"^transformer_blocks\.(\d+)\.attn\.norm_added_q\.(.*)$":
r"double_blocks.\1.txt_attn_q_norm.\2",
r"^transformer_blocks\.(\d+)\.attn\.norm_added_k\.(.*)$":
r"double_blocks.\1.txt_attn_k_norm.\2",
r"^transformer_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$":
r"double_blocks.\1.img_mlp.fc_in.\2",
r"^transformer_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$":
r"double_blocks.\1.img_mlp.fc_out.\2",
r"^transformer_blocks\.(\d+)\.ff_context\.net\.0(?:\.proj)?\.(.*)$":
r"double_blocks.\1.txt_mlp.fc_in.\2",
r"^transformer_blocks\.(\d+)\.ff_context\.net\.2(?:\.proj)?\.(.*)$":
r"double_blocks.\1.txt_mlp.fc_out.\2",
# 4. Top-level time_text_embed mappings:
r"^time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
r"time_in.mlp.fc_in.\1",
r"^time_text_embed\.timestep_embedder\.linear_2\.(.*)$":
r"time_in.mlp.fc_out.\1",
r"^time_text_embed\.guidance_embedder\.linear_1\.(.*)$":
r"guidance_in.mlp.fc_in.\1",
r"^time_text_embed\.guidance_embedder\.linear_2\.(.*)$":
r"guidance_in.mlp.fc_out.\1",
r"^time_text_embed\.text_embedder\.linear_1\.(.*)$":
r"vector_in.fc_in.\1",
r"^time_text_embed\.text_embedder\.linear_2\.(.*)$":
r"vector_in.fc_out.\1",
# 6. single_transformer_blocks mapping:
r"^single_transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$":
r"single_blocks.\1.q_norm.\2",
r"^single_transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$":
r"single_blocks.\1.k_norm.\2",
r"^single_transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$":
(r"single_blocks.\1.linear1.\2", 0, 4),
r"^single_transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$":
(r"single_blocks.\1.linear1.\2", 1, 4),
r"^single_transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$":
(r"single_blocks.\1.linear1.\2", 2, 4),
r"^single_transformer_blocks\.(\d+)\.proj_mlp\.(.*)$":
(r"single_blocks.\1.linear1.\2", 3, 4),
# Corrected: map proj_out to modulation.linear rather than a separate proj_out branch.
r"^single_transformer_blocks\.(\d+)\.proj_out\.(.*)$":
r"single_blocks.\1.linear2.\2",
r"^single_transformer_blocks\.(\d+)\.norm\.linear\.(.*)$":
r"single_blocks.\1.modulation.linear.\2",
# 5. transformer_blocks mapping:
r"^transformer_blocks\.(\d+)\.norm1\.linear\.(.*)$":
r"double_blocks.\1.img_mod.linear.\2",
r"^transformer_blocks\.(\d+)\.norm1_context\.linear\.(.*)$":
r"double_blocks.\1.txt_mod.linear.\2",
r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$":
r"double_blocks.\1.img_attn_q_norm.\2",
r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$":
r"double_blocks.\1.img_attn_k_norm.\2",
r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$":
(r"double_blocks.\1.img_attn_qkv.\2", 0, 3),
r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$":
(r"double_blocks.\1.img_attn_qkv.\2", 1, 3),
r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$":
(r"double_blocks.\1.img_attn_qkv.\2", 2, 3),
r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$":
(r"double_blocks.\1.txt_attn_qkv.\2", 0, 3),
r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$":
(r"double_blocks.\1.txt_attn_qkv.\2", 1, 3),
r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$":
(r"double_blocks.\1.txt_attn_qkv.\2", 2, 3),
r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$":
r"double_blocks.\1.img_attn_proj.\2",
# Corrected: merge attn.to_add_out into the main projection.
r"^transformer_blocks\.(\d+)\.attn\.to_add_out\.(.*)$":
r"double_blocks.\1.txt_attn_proj.\2",
r"^transformer_blocks\.(\d+)\.attn\.norm_added_q\.(.*)$":
r"double_blocks.\1.txt_attn_q_norm.\2",
r"^transformer_blocks\.(\d+)\.attn\.norm_added_k\.(.*)$":
r"double_blocks.\1.txt_attn_k_norm.\2",
r"^transformer_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$":
r"double_blocks.\1.img_mlp.fc_in.\2",
r"^transformer_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$":
r"double_blocks.\1.img_mlp.fc_out.\2",
r"^transformer_blocks\.(\d+)\.ff_context\.net\.0(?:\.proj)?\.(.*)$":
r"double_blocks.\1.txt_mlp.fc_in.\2",
r"^transformer_blocks\.(\d+)\.ff_context\.net\.2(?:\.proj)?\.(.*)$":
r"double_blocks.\1.txt_mlp.fc_out.\2",
# 6. single_transformer_blocks mapping:
r"^single_transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$":
r"single_blocks.\1.q_norm.\2",
r"^single_transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$":
r"single_blocks.\1.k_norm.\2",
r"^single_transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$":
(r"single_blocks.\1.linear1.\2", 0, 4),
r"^single_transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$":
(r"single_blocks.\1.linear1.\2", 1, 4),
r"^single_transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$":
(r"single_blocks.\1.linear1.\2", 2, 4),
r"^single_transformer_blocks\.(\d+)\.proj_mlp\.(.*)$":
(r"single_blocks.\1.linear1.\2", 3, 4),
# Corrected: map proj_out to modulation.linear rather than a separate proj_out branch.
r"^single_transformer_blocks\.(\d+)\.proj_out\.(.*)$":
r"single_blocks.\1.linear2.\2",
r"^single_transformer_blocks\.(\d+)\.norm\.linear\.(.*)$":
r"single_blocks.\1.modulation.linear.\2",
# 7. Final layers mapping:
r"^norm_out\.linear\.(.*)$":
r"final_layer.adaLN_modulation.linear.\1",
r"^proj_out\.(.*)$":
r"final_layer.linear.\1",
})
# 7. Final layers mapping:
r"^norm_out\.linear\.(.*)$":
r"final_layer.adaLN_modulation.linear.\1",
r"^proj_out\.(.*)$":
r"final_layer.linear.\1",
})
patch_size: int = 2
patch_size_t: int = 1
in_channels: int = 16
@@ -161,8 +165,9 @@ class HunyuanVideoArchConfig(DiTArchConfig):
self.hidden_size: int = self.attention_head_dim * self.num_attention_heads
self.num_channels_latents: int = self.in_channels
@dataclass
class HunyuanVideoConfig(DiTConfig):
arch_config: DiTArchConfig = HunyuanVideoArchConfig()
prefix: str = "Hunyuan"
prefix: str = "Hunyuan"
fastvideo/v1/configs/models/dits/wanvideo.py
+49 -44
View File
@@ -1,59 +1,63 @@
from dataclasses import dataclass, field
from typing import Tuple, Optional
from typing import Optional, Tuple
from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
from fastvideo.v1.platforms import _Backend
def is_blocks(n: str, m) -> bool:
return "blocks" in n and str.isdigit(n.split(".")[-1])
@dataclass
class WanVideoArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
_supported_attention_backends: Tuple[_Backend, ...] = field(default_factory=lambda: (_Backend.SLIDING_TILE_ATTN,
_Backend.FLASH_ATTN, _Backend.TORCH_SDPA))
_param_names_mapping: dict = field(default_factory=lambda: {
r"^patch_embedding\.(.*)$":
r"patch_embedding.proj.\1",
r"^condition_embedder\.text_embedder\.linear_1\.(.*)$":
r"condition_embedder.text_embedder.fc_in.\1",
r"^condition_embedder\.text_embedder\.linear_2\.(.*)$":
r"condition_embedder.text_embedder.fc_out.\1",
r"^condition_embedder\.time_embedder\.linear_1\.(.*)$":
r"condition_embedder.time_embedder.mlp.fc_in.\1",
r"^condition_embedder\.time_embedder\.linear_2\.(.*)$":
r"condition_embedder.time_embedder.mlp.fc_out.\1",
r"^condition_embedder\.time_proj\.(.*)$":
r"condition_embedder.time_modulation.linear.\1",
r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$":
r"condition_embedder.image_embedder.ff.fc_in.\1",
r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$":
r"condition_embedder.image_embedder.ff.fc_out.\1",
r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$":
r"blocks.\1.to_q.\2",
r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$":
r"blocks.\1.to_k.\2",
r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$":
r"blocks.\1.to_v.\2",
r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$":
r"blocks.\1.to_out.\2",
r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$":
r"blocks.\1.norm_q.\2",
r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$":
r"blocks.\1.norm_k.\2",
r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$":
r"blocks.\1.attn2.to_out.\2",
r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$":
r"blocks.\1.ffn.fc_in.\2",
r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$":
r"blocks.\1.ffn.fc_out.\2",
r"blocks\.(\d+)\.norm2\.(.*)$":
r"blocks.\1.self_attn_residual_norm.norm.\2",
})
_supported_attention_backends: Tuple[_Backend, ...] = field(
default_factory=lambda:
(_Backend.SLIDING_TILE_ATTN, _Backend.FLASH_ATTN, _Backend.TORCH_SDPA))
_param_names_mapping: dict = field(
default_factory=lambda: {
r"^patch_embedding\.(.*)$":
r"patch_embedding.proj.\1",
r"^condition_embedder\.text_embedder\.linear_1\.(.*)$":
r"condition_embedder.text_embedder.fc_in.\1",
r"^condition_embedder\.text_embedder\.linear_2\.(.*)$":
r"condition_embedder.text_embedder.fc_out.\1",
r"^condition_embedder\.time_embedder\.linear_1\.(.*)$":
r"condition_embedder.time_embedder.mlp.fc_in.\1",
r"^condition_embedder\.time_embedder\.linear_2\.(.*)$":
r"condition_embedder.time_embedder.mlp.fc_out.\1",
r"^condition_embedder\.time_proj\.(.*)$":
r"condition_embedder.time_modulation.linear.\1",
r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$":
r"condition_embedder.image_embedder.ff.fc_in.\1",
r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$":
r"condition_embedder.image_embedder.ff.fc_out.\1",
r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$":
r"blocks.\1.to_q.\2",
r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$":
r"blocks.\1.to_k.\2",
r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$":
r"blocks.\1.to_v.\2",
r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$":
r"blocks.\1.to_out.\2",
r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$":
r"blocks.\1.norm_q.\2",
r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$":
r"blocks.\1.norm_k.\2",
r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$":
r"blocks.\1.attn2.to_out.\2",
r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$":
r"blocks.\1.ffn.fc_in.\2",
r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$":
r"blocks.\1.ffn.fc_out.\2",
r"blocks\.(\d+)\.norm2\.(.*)$":
r"blocks.\1.self_attn_residual_norm.norm.\2",
})
patch_size: Tuple[int, int, int] = (1, 2, 2)
text_len=512
text_len = 512
num_attention_heads: int = 40
attention_head_dim: int = 128
in_channels: int = 16
@@ -74,8 +78,9 @@ class WanVideoArchConfig(DiTArchConfig):
self.hidden_size = self.num_attention_heads * self.attention_head_dim
self.num_channels_latents = self.in_channels if self.added_kv_proj_dim is None else self.out_channels
@dataclass
class WanVideoConfig(DiTConfig):
arch_config: DiTArchConfig = WanVideoArchConfig()
prefix: str = "Wan"
prefix: str = "Wan"
fastvideo/v1/configs/models/encoders/__init__.py
+10 -12
View File
@@ -1,14 +1,12 @@
from fastvideo.v1.configs.models.encoders.base import EncoderArchConfig, EncoderConfig, TextEncoderArchConfig, TextEncoderConfig, ImageEncoderArchConfig, ImageEncoderConfig
from fastvideo.v1.configs.models.encoders.clip import CLIPTextArchConfig, CLIPTextConfig, CLIPVisionArchConfig, CLIPVisionConfig
from fastvideo.v1.configs.models.encoders.llama import LlamaArchConfig, LlamaConfig
from fastvideo.v1.configs.models.encoders.t5 import T5ArchConfig, T5Config
from fastvideo.v1.configs.models.encoders.base import (EncoderConfig,
ImageEncoderConfig,
TextEncoderConfig)
from fastvideo.v1.configs.models.encoders.clip import (CLIPTextConfig,
CLIPVisionConfig)
from fastvideo.v1.configs.models.encoders.llama import LlamaConfig
from fastvideo.v1.configs.models.encoders.t5 import T5Config
__all__ = [
"EncoderArchConfig", "EncoderConfig",
"TextEncoderArchConfig", "TextEncoderConfig",
"ImageEncoderArchConfig", "ImageEncoderConfig",
"CLIPTextArchConfig", "CLIPTextConfig",
"CLIPVisionArchConfig", "CLIPVisionConfig",
"LlamaArchConfig", "LlamaConfig",
"T5ArchConfig", "T5Config"
]
"EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig",
"CLIPTextConfig", "CLIPVisionConfig", "LlamaConfig", "T5Config"
]
fastvideo/v1/configs/models/encoders/base.py
+12 -6
View File
@@ -1,18 +1,20 @@
from dataclasses import dataclass, field
from typing import Tuple, Optional, Any, List
from typing import Any, List, Optional, Tuple
from fastvideo.v1.configs.models import ArchConfig, ModelConfig
from fastvideo.v1.platforms import _Backend
from fastvideo.v1.configs.quantization import QuantizationConfig
from fastvideo.v1.platforms import _Backend
@dataclass
class EncoderArchConfig(ArchConfig):
architectures: List[str] = field(default_factory=lambda: [])
_supported_attention_backends: Tuple[_Backend,
...] = (_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
_supported_attention_backends: Tuple[_Backend, ...] = (_Backend.FLASH_ATTN,
_Backend.TORCH_SDPA)
output_hidden_states: bool = False
use_return_dict: bool = True
@dataclass
class TextEncoderArchConfig(EncoderArchConfig):
vocab_size: int = 0
@@ -28,11 +30,13 @@ class TextEncoderArchConfig(EncoderArchConfig):
scalable_attention: bool = True
tie_word_embeddings: bool = False
@dataclass
class ImageEncoderArchConfig(EncoderArchConfig):
pass
@dataclass
@dataclass
class EncoderConfig(ModelConfig):
arch_config: ArchConfig = EncoderArchConfig()
@@ -40,10 +44,12 @@ class EncoderConfig(ModelConfig):
quant_config: Optional[QuantizationConfig] = None
lora_config: Optional[Any] = None
@dataclass
class TextEncoderConfig(EncoderConfig):
arch_config: ArchConfig = TextEncoderArchConfig()
@dataclass
class ImageEncoderConfig(EncoderConfig):
arch_config: ArchConfig = ImageEncoderArchConfig()
arch_config: ArchConfig = ImageEncoderArchConfig()
fastvideo/v1/configs/models/encoders/clip.py
+13 -6
View File
@@ -1,7 +1,11 @@
from dataclasses import dataclass
from typing import Optional
from fastvideo.v1.configs.models.encoders.base import TextEncoderArchConfig, TextEncoderConfig, ImageEncoderArchConfig, ImageEncoderConfig
from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
ImageEncoderConfig,
TextEncoderArchConfig,
TextEncoderConfig)
@dataclass
class CLIPTextArchConfig(TextEncoderArchConfig):
@@ -23,6 +27,7 @@ class CLIPTextArchConfig(TextEncoderArchConfig):
eos_token_id: int = 49407
text_len: int = 77
@dataclass
class CLIPVisionArchConfig(ImageEncoderArchConfig):
hidden_size: int = 768
@@ -33,12 +38,13 @@ class CLIPVisionArchConfig(ImageEncoderArchConfig):
num_channels: int = 3
image_size: int = 224
patch_size: int = 32
hidden_act: int = "quick_gelu"
layer_norm_eps: int = 1e-5
hidden_act: str = "quick_gelu"
layer_norm_eps: float = 1e-5
dropout: float = 0.0
attention_dropout: float = 0.0
initializer_range: int = 0.02
initializer_factor: int = 1.0
initializer_range: float = 0.02
initializer_factor: float = 1.0
@dataclass
class CLIPTextConfig(TextEncoderConfig):
@@ -48,10 +54,11 @@ class CLIPTextConfig(TextEncoderConfig):
require_post_norm: Optional[bool] = None
prefix: str = "clip"
@dataclass
class CLIPVisionConfig(ImageEncoderConfig):
arch_config: ImageEncoderArchConfig = CLIPVisionArchConfig()
num_hidden_layers_override: Optional[int] = None
require_post_norm: Optional[bool] = None
prefix: str = "clip"
prefix: str = "clip"
fastvideo/v1/configs/models/encoders/llama.py
+6 -3
View File
@@ -1,7 +1,9 @@
from dataclasses import dataclass
from typing import Optional
from fastvideo.v1.configs.models.encoders.base import TextEncoderArchConfig, TextEncoderConfig
from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
TextEncoderConfig)
@dataclass
class LlamaArchConfig(TextEncoderArchConfig):
@@ -16,7 +18,7 @@ class LlamaArchConfig(TextEncoderArchConfig):
initializer_range: float = 0.02
rms_norm_eps: float = 1e-6
use_cache: bool = True
pad_token_id: Optional[int] = None
pad_token_id: int = 0
bos_token_id: int = 1
eos_token_id: int = 2
pretraining_tp: int = 1
@@ -30,8 +32,9 @@ class LlamaArchConfig(TextEncoderArchConfig):
hidden_state_skip_layer: int = 2
text_len: int = 256
@dataclass
class LlamaConfig(TextEncoderConfig):
arch_config: TextEncoderArchConfig = LlamaArchConfig()
prefix: str = "llama"
prefix: str = "llama"
fastvideo/v1/configs/models/encoders/t5.py
+5 -2
View File
@@ -1,7 +1,9 @@
from dataclasses import dataclass
from typing import Optional
from fastvideo.v1.configs.models.encoders.base import TextEncoderArchConfig, TextEncoderConfig
from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
TextEncoderConfig)
@dataclass
class T5ArchConfig(TextEncoderArchConfig):
@@ -35,8 +37,9 @@ class T5ArchConfig(TextEncoderArchConfig):
if self.feed_forward_proj == "gated-gelu":
self.dense_act_fn = "gelu_new"
@dataclass
class T5Config(TextEncoderConfig):
arch_config: TextEncoderArchConfig = T5ArchConfig()
prefix: str = "t5"
prefix: str = "t5"
fastvideo/v1/configs/models/vaes/__init__.py
+7 -5
View File
@@ -1,7 +1,9 @@
from fastvideo.v1.configs.models.vaes.hunyuanvae import HunyuanVAEConfig, HunyuanVAEArchConfig
from fastvideo.v1.configs.models.vaes.wanvae import WanVAEConfig, WanVAEArchConfig
from fastvideo.v1.configs.models.vaes.hunyuanvae import (HunyuanVAEArchConfig,
HunyuanVAEConfig)
from fastvideo.v1.configs.models.vaes.wanvae import (WanVAEArchConfig,
WanVAEConfig)
__all__ = [
"HunyuanVAEConfig", "HunyuanVAEArchConfig",
"WanVAEConfig", "WanVAEArchConfig"
]
"HunyuanVAEConfig", "HunyuanVAEArchConfig", "WanVAEConfig",
"WanVAEArchConfig"
]
fastvideo/v1/configs/models/vaes/base.py
+2
View File
@@ -5,6 +5,7 @@ import torch
from fastvideo.v1.configs.models import ArchConfig, ModelConfig
@dataclass
class VAEArchConfig(ArchConfig):
scaling_factor: Union[float, torch.tensor] = 0
@@ -12,6 +13,7 @@ class VAEArchConfig(ArchConfig):
temporal_compression_ratio: int = 4
spatial_compression_ratio: int = 8
@dataclass
class VAEConfig(ModelConfig):
arch_config: VAEArchConfig = VAEArchConfig()
fastvideo/v1/configs/models/vaes/hunyuanvae.py
+6 -3
View File
@@ -1,7 +1,8 @@
from dataclasses import dataclass
from typing import Tuple
from fastvideo.v1.configs.models.vaes.base import VAEConfig, VAEArchConfig
from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
@dataclass
class HunyuanVAEArchConfig(VAEArchConfig):
@@ -30,8 +31,10 @@ class HunyuanVAEArchConfig(VAEArchConfig):
mid_block_add_attention: bool = True
def __post_init__(self):
self.spatial_compression_ratio: int = 2**(len(self.block_out_channels)-1)
self.spatial_compression_ratio: int = 2**(len(self.block_out_channels) -
1)
@dataclass
class HunyuanVAEConfig(VAEConfig):
arch_config: VAEArchConfig = HunyuanVAEArchConfig()
arch_config: VAEArchConfig = HunyuanVAEArchConfig()
fastvideo/v1/configs/models/vaes/wanvae.py
+8 -5
View File
@@ -3,7 +3,8 @@ from typing import Tuple
import torch
from fastvideo.v1.configs.models.vaes.base import VAEConfig, VAEArchConfig
from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
@dataclass
class WanVAEArchConfig(VAEArchConfig):
@@ -52,13 +53,14 @@ class WanVAEArchConfig(VAEArchConfig):
)
temporal_compression_ratio = 4
spatial_compression_ratio = 8
def __post_init__(self):
self.scaling_factor: torch.tensor = 1.0 / torch.tensor(self.latents_std).view(
1, self.z_dim, 1, 1, 1)
self.scaling_factor: torch.tensor = 1.0 / torch.tensor(
self.latents_std).view(1, self.z_dim, 1, 1, 1)
self.shift_factor: torch.tensor = torch.tensor(self.latents_mean).view(
1, self.z_dim, 1, 1, 1)
@dataclass
class WanVAEConfig(VAEConfig):
arch_config: VAEArchConfig = WanVAEArchConfig()
@@ -69,4 +71,5 @@ class WanVAEConfig(VAEConfig):
use_parallel_tiling: bool = False
def __post_init__(self):
self.blend_num_frames = (self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames) * 2
self.blend_num_frames = (self.tile_sample_min_num_frames -
self.tile_sample_stride_num_frames) * 2
fastvideo/v1/configs/pipelines/__init__.py
+8 -4
View File
@@ -1,7 +1,11 @@
from fastvideo.v1.configs.pipelines.hunyuan import HunyuanConfig, FastHunyuanConfig
from fastvideo.v1.configs.pipelines.wan import WanT2V480PConfig, WanI2V480PConfig
from fastvideo.v1.configs.pipelines.base import BaseConfig, SlidingTileAttnConfig
from fastvideo.v1.configs.pipelines.registry import get_pipeline_config_cls_for_name
from fastvideo.v1.configs.pipelines.base import (BaseConfig,
SlidingTileAttnConfig)
from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
HunyuanConfig)
from fastvideo.v1.configs.pipelines.registry import (
get_pipeline_config_cls_for_name)
from fastvideo.v1.configs.pipelines.wan import (WanI2V480PConfig,
WanT2V480PConfig)
__all__ = [
"HunyuanConfig", "FastHunyuanConfig", "BaseConfig", "SlidingTileAttnConfig",
fastvideo/v1/configs/pipelines/base.py
+19 -17
View File
@@ -1,13 +1,15 @@
from dataclasses import dataclass, asdict, fields
from typing import Optional, Dict, Any
import json
from dataclasses import asdict, dataclass, fields
from typing import Any, Dict, Optional
from fastvideo.v1.configs.models import ModelConfig, VAEConfig, DiTConfig, EncoderConfig
from fastvideo.v1.utils import shallow_asdict
from fastvideo.v1.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
VAEConfig)
from fastvideo.v1.logger import init_logger
from fastvideo.v1.utils import shallow_asdict
logger = init_logger(__name__)
@dataclass
class BaseConfig:
"""Base configuration for all pipeline architectures."""
@@ -38,15 +40,16 @@ class BaseConfig:
enable_torch_compile: bool = False
@classmethod
def from_pretrained(cls,
model_path: str
) -> "BaseConfig":
from fastvideo.v1.configs.pipelines.registry import get_pipeline_config_cls_for_name
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)
logger.warning(
"Couldn't find an optimal sampling param for %s. Using the default sampling param.",
model_path)
pipeline_config = cls()
return pipeline_config
@@ -59,19 +62,17 @@ class BaseConfig:
# Model Arch Config should be hidden away from the users
model_dict.pop("arch_config")
output_dict[key] = model_dict
with open(file_path, "w") as f:
json.dump(output_dict, f, indent=2)
def load_from_json(self, file_path: str):
with open(file_path, "r") as f:
with open(file_path) as f:
input_pipeline_dict = json.load(f)
self.update_pipeline_config(input_pipeline_dict)
def update_pipeline_config(
self,
source_pipeline_dict: Dict[str, Any]
) -> None:
def update_pipeline_config(self, source_pipeline_dict: Dict[str,
Any]) -> None:
for f in fields(self):
key = f.name
if key in source_pipeline_dict:
@@ -83,10 +84,11 @@ class BaseConfig:
current_value.update_model_config(new_value)
else:
setattr(self, key, new_value)
if hasattr(self, "__post_init__"):
self.__post_init__()
@dataclass
class SlidingTileAttnConfig(BaseConfig):
"""Configuration for sliding tile attention."""
@@ -102,4 +104,4 @@ class SlidingTileAttnConfig(BaseConfig):
# Additional configuration specific to sliding tile attention
pad_to_square: bool = False
use_overlap_optimization: bool = True
use_overlap_optimization: bool = True
fastvideo/v1/configs/pipelines/hunyuan.py
+4 -3
View File
@@ -1,10 +1,11 @@
from dataclasses import dataclass
from fastvideo.v1.configs.pipelines.base import BaseConfig
from fastvideo.v1.configs.models import VAEConfig, DiTConfig, EncoderConfig
from fastvideo.v1.configs.models.vaes import HunyuanVAEConfig
from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
from fastvideo.v1.configs.models.dits import HunyuanVideoConfig
from fastvideo.v1.configs.models.encoders import CLIPTextConfig, LlamaConfig
from fastvideo.v1.configs.models.vaes import HunyuanVAEConfig
from fastvideo.v1.configs.pipelines.base import BaseConfig
@dataclass
class HunyuanConfig(BaseConfig):
fastvideo/v1/configs/pipelines/registry.py
+4 -3
View File
@@ -4,9 +4,10 @@ import os
from typing import Callable, Dict, Optional, Type
from fastvideo.v1.configs.pipelines.base import BaseConfig
from fastvideo.v1.configs.pipelines.hunyuan import HunyuanConfig, FastHunyuanConfig
from fastvideo.v1.configs.pipelines.wan import WanT2V480PConfig, WanI2V480PConfig
from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
HunyuanConfig)
from fastvideo.v1.configs.pipelines.wan import (WanI2V480PConfig,
WanT2V480PConfig)
from fastvideo.v1.logger import init_logger
from fastvideo.v1.utils import (maybe_download_model_index,
verify_model_config_and_directory)
fastvideo/v1/configs/pipelines/wan.py
+6 -4
View File
@@ -1,10 +1,11 @@
from dataclasses import dataclass
from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
from fastvideo.v1.configs.models.dits import WanVideoConfig
from fastvideo.v1.configs.models.encoders import CLIPVisionConfig, T5Config
from fastvideo.v1.configs.models.vaes import WanVAEConfig
from fastvideo.v1.configs.pipelines.base import BaseConfig
from fastvideo.v1.configs.models import VAEConfig, DiTConfig, EncoderConfig
from fastvideo.v1.configs.models.vaes import WanVAEConfig
from fastvideo.v1.configs.models.dits import WanVideoConfig
from fastvideo.v1.configs.models.encoders import T5Config, CLIPVisionConfig
@dataclass
class WanT2V480PConfig(BaseConfig):
@@ -38,6 +39,7 @@ class WanT2V480PConfig(BaseConfig):
self.vae_config.load_encoder = False
self.vae_config.load_decoder = True
@dataclass
class WanI2V480PConfig(WanT2V480PConfig):
"""Base configuration for Wan I2V 14B 480P pipeline architecture."""
fastvideo/v1/configs/quantization/__init__.py
+1 -3
View File
@@ -1,5 +1,3 @@
from fastvideo.v1.configs.quantization.base import QuantizationConfig
__all__ = [
"QuantizationConfig"
]
__all__ = ["QuantizationConfig"]
fastvideo/v1/configs/quantization/base.py
+2 -1
View File
@@ -1,5 +1,6 @@
from dataclasses import dataclass
@dataclass
class QuantizationConfig:
pass
pass
fastvideo/v1/configs/sample/__init__.py
+1 -3
View File
@@ -1,5 +1,3 @@
from fastvideo.v1.configs.sample.base import SamplingParam
__all__ = [
"SamplingParam"
]
__all__ = ["SamplingParam"]
fastvideo/v1/configs/sample/base.py
+17 -13
View File
@@ -1,9 +1,11 @@
from dataclasses import dataclass
from typing import Optional, Union, List, Dict, Any
from typing import Any, Dict, List, Optional, Union
from fastvideo.v1.logger import init_logger
logger = init_logger(__name__)
@dataclass
class SamplingParam:
# All fields below are copied from ForwardBatch
@@ -14,7 +16,7 @@ class SamplingParam:
# Text inputs
prompt: Optional[Union[str, List[str]]] = None
negative_prompt: Optional[Union[str, List[str]]] = None
negative_prompt: Optional[str] = None
prompt_path: Optional[str] = None
output_path: str = "outputs/"
@@ -37,32 +39,34 @@ class SamplingParam:
save_video: bool = True
return_frames: bool = False
def __post_init__(self):
def __post_init__(self) -> None:
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)
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
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()
sampling_param: SamplingParam = sampling_cls()
else:
logger.warning("Couldn't find an optimal sampling param for %s. Using the default sampling param.", model_path)
logger.warning(
"Couldn't find an optimal sampling param for %s. Using the default sampling param.",
model_path)
sampling_param = cls()
return sampling_param
return sampling_param
fastvideo/v1/configs/sample/hunyuan.py
+3 -1
View File
@@ -2,6 +2,7 @@ from dataclasses import dataclass
from fastvideo.v1.configs.sample.base import SamplingParam
@dataclass
class HunyuanSamplingParam(SamplingParam):
num_inference_steps: int = 50
@@ -13,6 +14,7 @@ class HunyuanSamplingParam(SamplingParam):
guidance_scale: float = 1.0
@dataclass
class FastHunyuanSamplingParam(HunyuanSamplingParam):
num_inference_steps: int = 6
num_inference_steps: int = 6
fastvideo/v1/configs/sample/registry.py
+9 -7
View File
@@ -1,9 +1,10 @@
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 typing import Any, Callable, Dict, Optional
from fastvideo.v1.configs.sample.hunyuan import (FastHunyuanSamplingParam,
HunyuanSamplingParam)
from fastvideo.v1.configs.sample.wan import (WanI2V480PSamplingParam,
WanT2V480PSamplingParam)
from fastvideo.v1.logger import init_logger
from fastvideo.v1.utils import (maybe_download_model_index,
verify_model_config_and_directory)
@@ -68,6 +69,7 @@ def get_sampling_param_cls_for_name(
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
logger.warning(
"No match found for pipeline %s, using fallback sampling param %s.",
pipeline_name_or_path, fallback_config)
return fallback_config
fastvideo/v1/configs/sample/wan.py
+3 -1
View File
@@ -2,6 +2,7 @@ from dataclasses import dataclass
from fastvideo.v1.configs.sample.base import SamplingParam
@dataclass
class WanT2V480PSamplingParam(SamplingParam):
# Video parameters
@@ -15,8 +16,9 @@ class WanT2V480PSamplingParam(SamplingParam):
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
num_inference_steps: int = 40
fastvideo/v1/entrypoints/video_generator.py
+15 -11
View File
@@ -8,16 +8,17 @@ diffusion models.
import os
import time
from typing import Any, Callable, Dict, List, Optional, Union
from dataclasses import asdict
from typing import Any, Dict, List, Optional, Union
import imageio
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.pipelines import (BaseConfig,
get_pipeline_config_cls_for_name)
from fastvideo.v1.configs.sample import SamplingParam
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.logger import init_logger
@@ -53,7 +54,8 @@ class VideoGenerator:
model_path: str,
device: Optional[str] = None,
torch_dtype: Optional[torch.dtype] = None,
pipeline_config: Optional[Union[str | BaseConfig]] = None,
pipeline_config: Optional[Union[str
| BaseConfig]] = None,
**kwargs) -> "VideoGenerator":
"""
Create a video generator from a pretrained model.
@@ -125,7 +127,7 @@ class VideoGenerator:
def generate_video(
self,
prompt: str,
sampling_param: SamplingParam = None,
sampling_param: Optional[SamplingParam] = None,
**kwargs,
) -> Union[Dict[str, Any], List[np.ndarray]]:
"""
@@ -160,13 +162,15 @@ class VideoGenerator:
prompt = prompt.strip()
if sampling_param is None:
sampling_param = SamplingParam.from_pretrained(fastvideo_args.model_path)
sampling_param = SamplingParam.from_pretrained(
fastvideo_args.model_path)
kwargs["prompt"] = prompt
sampling_param.update(kwargs)
# Process negative prompt
if sampling_param.negative_prompt is not None:
sampling_param.negative_prompt = sampling_param.negative_prompt.strip()
sampling_param.negative_prompt = sampling_param.negative_prompt.strip(
)
# Validate dimensions
if (sampling_param.height <= 0 or sampling_param.width <= 0
@@ -176,7 +180,9 @@ class VideoGenerator:
f"height={sampling_param.height}, width={sampling_param.width}, "
f"num_frames={sampling_param.num_frames}")
if (sampling_param.num_frames - 1) % fastvideo_args.vae_config.arch_config.temporal_compression_ratio != 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 {fastvideo_args.vae_config.arch_config.temporal_compression_ratio}, got {sampling_param.num_frames}"
)
@@ -207,7 +213,6 @@ class VideoGenerator:
logger.info(debug_str)
# Prepare batch
device = torch.device(fastvideo_args.device_str)
batch = ForwardBatch(
**asdict(sampling_param),
eta=0.0,
@@ -248,7 +253,6 @@ class VideoGenerator:
return {
"samples": samples,
"prompts": prompt,
"size":
(target_height, target_width, batch.num_frames),
"size": (target_height, target_width, batch.num_frames),
"generation_time": gen_time
}
fastvideo/v1/fastvideo_args.py
+3 -16
View File
@@ -7,11 +7,10 @@ import dataclasses
from contextlib import contextmanager
from typing import List, Optional
from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
from fastvideo.v1.logger import init_logger
from fastvideo.v1.utils import FlexibleArgumentParser
from fastvideo.v1.configs.models import VAEConfig, DiTConfig, EncoderConfig
logger = init_logger(__name__)
@@ -47,8 +46,8 @@ class FastVideoArgs:
# VAE configuration
vae_precision: str = "fp16"
vae_tiling: bool = True # Might change in between forward passes
vae_sp: bool = False # Might change in between forward passes
vae_tiling: bool = True # Might change in between forward passes
vae_sp: bool = False # Might change in between forward passes
# vae_scale_factor: Optional[int] = None # Deprecated
vae_config: VAEConfig = VAEConfig()
@@ -210,12 +209,6 @@ class FastVideoArgs:
choices=["fp32", "fp16", "bf16"],
help="Precision for text encoder",
)
parser.add_argument(
"--text-len",
type=int,
default=FastVideoArgs.text_len,
help="Maximum text length",
)
# Image encoder config
parser.add_argument(
@@ -235,12 +228,6 @@ class FastVideoArgs:
choices=["fp32", "fp16", "bf16"],
help="Precision for secondary text encoder",
)
parser.add_argument(
"--text-len-2",
type=int,
default=FastVideoArgs.text_len_2,
help="Maximum secondary text length",
)
# STA (Spatial-Temporal Attention) parameters
parser.add_argument(
fastvideo/v1/inference_engine.py
+2
View File
@@ -1,3 +1,5 @@
# type: ignore
# SPDX-License-Identifier: Apache-2.0
"""
Inference module for diffusion models.
fastvideo/v1/models/dits/base.py
+5 -2
View File
@@ -5,8 +5,9 @@ from typing import List, Optional, Tuple, Union
import torch
from torch import nn
from fastvideo.v1.platforms import _Backend
from fastvideo.v1.configs.models import DiTConfig
from fastvideo.v1.platforms import _Backend
# TODO
class BaseDiT(nn.Module, ABC):
@@ -51,7 +52,9 @@ class BaseDiT(nn.Module, ABC):
pass
def __post_init__(self) -> None:
required_attrs = ["hidden_size", "num_attention_heads", "num_channels_latents"]
required_attrs = [
"hidden_size", "num_attention_heads", "num_channels_latents"
]
for attr in required_attrs:
if not hasattr(self, attr):
raise AttributeError(
fastvideo/v1/models/dits/hunyuanvideo.py
+23 -15
View File
@@ -6,6 +6,8 @@ import torch
import torch.nn as nn
from fastvideo.v1.attention import DistributedAttention, LocalAttention
from fastvideo.v1.configs.models.dits import (HunyuanVideoArchConfig,
HunyuanVideoConfig)
from fastvideo.v1.distributed.parallel_state import (
get_sequence_model_parallel_world_size)
from fastvideo.v1.layers.layernorm import (LayerNormScaleShift, ScaleResidual,
@@ -20,7 +22,6 @@ from fastvideo.v1.layers.visual_embedding import (ModulateProjection,
unpatchify)
from fastvideo.v1.models.dits.base import BaseDiT
from fastvideo.v1.platforms import _Backend
from fastvideo.v1.configs.models.dits import HunyuanVideoArchConfig, HunyuanVideoConfig
class HunyuanRMSNorm(nn.Module):
@@ -433,16 +434,19 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
# shard single stream, double stream blocks, and refiner_blocks
_fsdp_shard_conditions = HunyuanVideoArchConfig()._fsdp_shard_conditions
_supported_attention_backends = HunyuanVideoArchConfig()._supported_attention_backends
_supported_attention_backends = HunyuanVideoArchConfig(
)._supported_attention_backends
_param_names_mapping = HunyuanVideoArchConfig()._param_names_mapping
def __init__(
self,
config: HunyuanVideoConfig
):
def __init__(self, config: HunyuanVideoConfig):
super().__init__(config=config)
arch_config: HunyuanVideoArchConfig = cast(HunyuanVideoArchConfig, config.arch_config)
self.patch_size = [arch_config.patch_size_t, self.arch_config.patch_size, self.arch_config.patch_size]
arch_config: HunyuanVideoArchConfig = cast(HunyuanVideoArchConfig,
config.arch_config)
self.arch_config = arch_config
self.patch_size = [
arch_config.patch_size_t, self.arch_config.patch_size,
self.arch_config.patch_size
]
self.in_channels = arch_config.in_channels
self.num_channels_latents = arch_config.num_channels_latents
self.out_channels = arch_config.in_channels if arch_config.out_channels is None else arch_config.out_channels
@@ -458,7 +462,8 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
pe_dim = arch_config.hidden_size // arch_config.num_attention_heads
if sum(arch_config.rope_axes_dim) != pe_dim:
raise ValueError(
f"Got {arch_config.rope_axes_dim} but expected positional dim {pe_dim}")
f"Got {arch_config.rope_axes_dim} but expected positional dim {pe_dim}"
)
self.hidden_size = arch_config.hidden_size
self.num_attention_heads = arch_config.num_attention_heads
@@ -493,10 +498,11 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
prefix=f"{config.prefix}.vector_in")
# Guidance modulation
self.guidance_in = (TimestepEmbedder(self.hidden_size,
act_layer="silu",
dtype=arch_config.dtype,
prefix=f"{config.prefix}.guidance_in")
self.guidance_in = (TimestepEmbedder(
self.hidden_size,
act_layer="silu",
dtype=arch_config.dtype,
prefix=f"{config.prefix}.guidance_in")
if self.guidance_embeds else None)
# Double blocks
@@ -507,7 +513,8 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
mlp_ratio=arch_config.mlp_ratio,
dtype=arch_config.dtype,
supported_attention_backends=self._supported_attention_backends,
prefix=f"{config.prefix}.double_blocks.{i}") for i in range(arch_config.num_layers)
prefix=f"{config.prefix}.double_blocks.{i}")
for i in range(arch_config.num_layers)
])
# Single blocks
@@ -518,7 +525,8 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
mlp_ratio=arch_config.mlp_ratio,
dtype=arch_config.dtype,
supported_attention_backends=self._supported_attention_backends,
prefix=f"{config.prefix}.single_blocks.{i+arch_config.num_layers}")
prefix=
f"{config.prefix}.single_blocks.{i+arch_config.num_layers}")
for i in range(arch_config.num_single_layers)
])
fastvideo/v1/models/dits/wanvideo.py
+9 -7
View File
@@ -7,6 +7,7 @@ import torch
import torch.nn as nn
from fastvideo.v1.attention import DistributedAttention, LocalAttention
from fastvideo.v1.configs.models.dits import WanVideoArchConfig, WanVideoConfig
from fastvideo.v1.distributed.parallel_state import (
get_sequence_model_parallel_world_size)
from fastvideo.v1.layers.layernorm import (LayerNormScaleShift, RMSNorm,
@@ -22,7 +23,6 @@ from fastvideo.v1.layers.visual_embedding import (ModulateProjection,
PatchEmbed, TimestepEmbedder)
from fastvideo.v1.models.dits.base import BaseDiT
from fastvideo.v1.platforms import _Backend
from fastvideo.v1.configs.models.dits import WanVideoArchConfig, WanVideoConfig
class WanImageEmbedding(torch.nn.Module):
@@ -352,14 +352,15 @@ class WanTransformerBlock(nn.Module):
class WanTransformer3DModel(BaseDiT):
_fsdp_shard_conditions = WanVideoArchConfig()._fsdp_shard_conditions
_supported_attention_backends = WanVideoArchConfig()._supported_attention_backends
_supported_attention_backends = WanVideoArchConfig(
)._supported_attention_backends
_param_names_mapping = WanVideoArchConfig()._param_names_mapping
def __init__(self,
config: WanVideoConfig) -> None:
def __init__(self, config: WanVideoConfig) -> None:
super().__init__(config=config)
arch_config: WanVideoArchConfig = cast(WanVideoArchConfig, config.arch_config)
arch_config: WanVideoArchConfig = cast(WanVideoArchConfig,
config.arch_config)
inner_dim = arch_config.num_attention_heads * arch_config.attention_head_dim
self.hidden_size = arch_config.hidden_size
@@ -404,8 +405,9 @@ class WanTransformer3DModel(BaseDiT):
eps=arch_config.eps,
elementwise_affine=False,
dtype=torch.float32)
self.proj_out = nn.Linear(inner_dim,
arch_config.out_channels * math.prod(arch_config.patch_size))
self.proj_out = nn.Linear(
inner_dim,
arch_config.out_channels * math.prod(arch_config.patch_size))
self.scale_shift_table = nn.Parameter(
torch.randn(1, 2, inner_dim) / inner_dim**0.5)
fastvideo/v1/models/encoders/base.py
+4 -3
View File
@@ -2,12 +2,13 @@ from typing import Tuple
from torch import nn
from fastvideo.v1.platforms import _Backend
from fastvideo.v1.configs.models.encoders import EncoderConfig
from fastvideo.v1.platforms import _Backend
class BaseEncoder(nn.Module):
_supported_attention_backends: Tuple[_Backend,
...] = EncoderConfig()._supported_attention_backends
_supported_attention_backends: Tuple[
_Backend, ...] = EncoderConfig()._supported_attention_backends
def __init__(self, config: EncoderConfig) -> None:
super().__init__()
fastvideo/v1/models/encoders/clip.py
+11 -13
View File
@@ -11,6 +11,9 @@ from transformers.modeling_outputs import BaseModelOutputWithPooling
# from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask
from fastvideo.v1.attention import LocalAttention
from fastvideo.v1.configs.models.encoders import (CLIPTextConfig,
CLIPVisionConfig)
from fastvideo.v1.configs.quantization import QuantizationConfig
from fastvideo.v1.distributed import (divide,
get_tensor_model_parallel_world_size)
from fastvideo.v1.layers.activation import get_act_fn
@@ -22,8 +25,6 @@ from fastvideo.v1.models.encoders.vision import resolve_visual_encoder_outputs
# TODO: support quantization
# from vllm.model_executor.layers.quantization import QuantizationConfig
from fastvideo.v1.models.loader.weight_utils import default_weight_loader
from fastvideo.v1.configs.models.encoders import CLIPTextConfig, CLIPVisionConfig
from fastvideo.v1.configs.quantization import QuantizationConfig
logger = init_logger(__name__)
@@ -162,13 +163,13 @@ class CLIPAttention(nn.Module):
self.tp_size = get_tensor_model_parallel_world_size()
self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
self.attn = LocalAttention(self.num_heads_per_partition,
self.head_dim,
self.num_heads_per_partition,
softmax_scale=self.scale,
causal=True,
supported_attention_backends=config.
_supported_attention_backends)
self.attn = LocalAttention(
self.num_heads_per_partition,
self.head_dim,
self.num_heads_per_partition,
softmax_scale=self.scale,
causal=True,
supported_attention_backends=config._supported_attention_backends)
def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
return tensor.view(bsz, seq_len, self.num_heads,
@@ -575,10 +576,7 @@ class CLIPVisionModel(BaseEncoder):
main_input_name = "pixel_values"
packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
def __init__(
self,
config: CLIPVisionConfig
) -> None:
def __init__(self, config: CLIPVisionConfig) -> None:
super().__init__(config)
self.vision_model = CLIPVisionTransformer(
config=config,
fastvideo/v1/models/encoders/llama.py
+20 -19
View File
@@ -31,6 +31,9 @@ from transformers.modeling_outputs import BaseModelOutputWithPast
# from vllm.model_executor.layers.quantization import QuantizationConfig
from fastvideo.v1.attention import LocalAttention
# from ..utils import (extract_layer_index)
from fastvideo.v1.configs.models.encoders import LlamaConfig
from fastvideo.v1.configs.quantization import QuantizationConfig
from fastvideo.v1.distributed import get_tensor_model_parallel_world_size
from fastvideo.v1.layers.activation import SiluAndMul
from fastvideo.v1.layers.layernorm import RMSNorm
@@ -41,9 +44,6 @@ from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding
from fastvideo.v1.models.encoders.base import BaseEncoder
from fastvideo.v1.models.loader.weight_utils import (default_weight_loader,
maybe_remap_kv_scale_name)
# from ..utils import (extract_layer_index)
from fastvideo.v1.configs.models.encoders import LlamaConfig
from fastvideo.v1.configs.quantization import QuantizationConfig
class LlamaMLP(nn.Module):
@@ -277,9 +277,10 @@ class LlamaDecoderLayer(nn.Module):
class LlamaModel(BaseEncoder):
def __init__(self,
config: LlamaConfig,
):
def __init__(
self,
config: LlamaConfig,
):
super().__init__(config)
self.config = config
@@ -306,8 +307,8 @@ class LlamaModel(BaseEncoder):
self.layers = nn.ModuleList([
LlamaDecoderLayer(config=config,
quant_config=config.quant_config,
prefix=f"{config.prefix}.layers.{i}")
quant_config=config.quant_config,
prefix=f"{config.prefix}.layers.{i}")
for i in range(config.num_hidden_layers)
])
@@ -385,17 +386,17 @@ class LlamaModel(BaseEncoder):
# Models trained using ColossalAI may include these tensors in
# the checkpoint. Skip them.
continue
if (self.quant_config is not None and
(scale_name := self.quant_config.get_cache_scale(name))):
# Loading kv cache quantization scales
param = params_dict[scale_name]
weight_loader = getattr(param, "weight_loader",
default_weight_loader)
loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
loaded_weight[0])
weight_loader(param, loaded_weight)
loaded_params.add(scale_name)
continue
# if (self.quant_config is not None and
# (scale_name := self.quant_config.get_cache_scale(name))):
# # Loading kv cache quantization scales
# param = params_dict[scale_name]
# weight_loader = getattr(param, "weight_loader",
# default_weight_loader)
# loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
# loaded_weight[0])
# weight_loader(param, loaded_weight)
# loaded_params.add(scale_name)
# continue
if "scale" in name:
# Remapping the name of FP8 kv-scale.
kv_scale_name: Optional[str] = maybe_remap_kv_scale_name(
fastvideo/v1/models/encoders/t5.py
+3 -3
View File
@@ -27,6 +27,8 @@ import torch
import torch.nn.functional as F
from torch import nn
from fastvideo.v1.configs.models.encoders import T5Config
from fastvideo.v1.configs.quantization import QuantizationConfig
from fastvideo.v1.distributed import (get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size)
from fastvideo.v1.layers.activation import get_act_fn
@@ -34,10 +36,8 @@ from fastvideo.v1.layers.layernorm import RMSNorm
from fastvideo.v1.layers.linear import (MergedColumnParallelLinear,
QKVParallelLinear, RowParallelLinear)
from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding
from fastvideo.v1.models.loader.weight_utils import default_weight_loader
from fastvideo.v1.models.encoders.base import BaseEncoder
from fastvideo.v1.configs.models.encoders import T5Config
from fastvideo.v1.configs.quantization import QuantizationConfig
from fastvideo.v1.models.loader.weight_utils import default_weight_loader
class AttentionType:
fastvideo/v1/models/loader/component_loader.py
+9 -10
View File
@@ -1,9 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
import dataclasses
from dataclasses import asdict
import json
import glob
import json
import os
import time
from abc import ABC, abstractmethod
@@ -12,13 +11,12 @@ from typing import Any, Generator, Iterable, List, Optional, Tuple, cast
import torch
import torch.nn as nn
from safetensors.torch import load_file as safetensors_load_file
from transformers import AutoImageProcessor, AutoTokenizer, PretrainedConfig
from transformers import AutoImageProcessor, AutoTokenizer
from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.logger import init_logger
from fastvideo.v1.models.hf_transformer_utils import (get_diffusers_config,
get_hf_config)
from fastvideo.v1.models.hf_transformer_utils import get_diffusers_config
from fastvideo.v1.models.loader.fsdp_load import load_fsdp_model
from fastvideo.v1.models.loader.utils import set_default_torch_dtype
from fastvideo.v1.models.loader.weight_utils import (
@@ -209,7 +207,7 @@ class TextEncoderLoader(ComponentLoader):
# revision=fastvideo_args.revision,
# model_override_args=None,
# )
with open(os.path.join(model_path, "config.json"), "r") as f:
with open(os.path.join(model_path, "config.json")) as f:
model_config = json.load(f)
model_config.pop("_name_or_path", None)
model_config.pop("transformers_version", None)
@@ -222,7 +220,7 @@ class TextEncoderLoader(ComponentLoader):
encoder_config = fastvideo_args.text_encoder_config
encoder_config.update_model_arch(model_config)
encoder_precision = fastvideo_args.text_encoder_precision
except:
except Exception:
encoder_config = fastvideo_args.text_encoder_config_2
encoder_config.update_model_arch(model_config)
encoder_precision = fastvideo_args.text_encoder_precision_2
@@ -275,7 +273,7 @@ class ImageEncoderLoader(TextEncoderLoader):
# revision=fastvideo_args.revision,
# model_override_args=None,
# )
with open(os.path.join(model_path, "config.json"), "r") as f:
with open(os.path.join(model_path, "config.json")) as f:
model_config = json.load(f)
model_config.pop("_name_or_path", None)
model_config.pop("transformers_version", None)
@@ -352,7 +350,8 @@ class VAELoader(ComponentLoader):
safetensors_list
) == 1, f"Found {len(safetensors_list)} safetensors files in {model_path}"
loaded = safetensors_load_file(safetensors_list[0])
vae.load_state_dict(loaded, strict=False) # We might only load encoder or decoder
vae.load_state_dict(
loaded, strict=False) # We might only load encoder or decoder
dtype = PRECISION_TO_TYPE[fastvideo_args.vae_precision]
vae = vae.eval().to(dtype)
@@ -373,7 +372,7 @@ class TransformerLoader(ComponentLoader):
"Only diffusers format is supported.")
config.pop("_diffusers_version")
# Config from Diffusers supercedes fastvideo's model config
# Config from Diffusers supersedes fastvideo's model config
dit_config = fastvideo_args.dit_config
dit_config.update_model_arch(config)
fastvideo/v1/models/vaes/common.py
+3 -3
View File
@@ -9,9 +9,9 @@ import torch
import torch.distributed as dist
from diffusers.utils.torch_utils import randn_tensor
from fastvideo.v1.configs.models import VAEConfig
from fastvideo.v1.distributed import (get_sequence_model_parallel_rank,
get_sequence_model_parallel_world_size)
from fastvideo.v1.configs.models import VAEConfig
class ParallelTiledVAE(ABC):
@@ -43,11 +43,11 @@ class ParallelTiledVAE(ABC):
@property
def temporal_compression_ratio(self) -> int:
return self.arch_config.temporal_compression_ratio
@property
def spatial_compression_ratio(self) -> int:
return self.arch_config.spatial_compression_ratio
@property
def scaling_factor(self) -> Union[float, torch.tensor]:
return self.arch_config.scaling_factor
fastvideo/v1/models/vaes/hunyuanvae.py
+8 -5
View File
@@ -23,9 +23,10 @@ import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint
from fastvideo.v1.configs.models.vaes import (HunyuanVAEArchConfig,
HunyuanVAEConfig)
from fastvideo.v1.layers.activation import get_act_fn
from fastvideo.v1.models.vaes.common import ParallelTiledVAE
from fastvideo.v1.configs.models.vaes import HunyuanVAEConfig, HunyuanVAEArchConfig
def prepare_causal_attention_mask(
@@ -777,14 +778,15 @@ class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE):
self,
config: HunyuanVAEConfig,
) -> None:
super().__init__()
super().__init__(config)
ParallelTiledVAE.__init__(self, config)
arch_config: HunyuanVAEArchConfig = cast(HunyuanVAEArchConfig, config.arch_config)
arch_config: HunyuanVAEArchConfig = cast(HunyuanVAEArchConfig,
config.arch_config)
# TODO(will): only pass in config. We do this by manually defining a
# config for hunyuan vae
self.block_out_channels = arch_config.block_out_channels
if config.load_encoder:
self.encoder = HunyuanVideoEncoder3D(
in_channels=arch_config.in_channels,
@@ -796,7 +798,8 @@ class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE):
act_fn=arch_config.act_fn,
double_z=True,
mid_block_add_attention=arch_config.mid_block_add_attention,
temporal_compression_ratio=arch_config.temporal_compression_ratio,
temporal_compression_ratio=arch_config.
temporal_compression_ratio,
spatial_compression_ratio=arch_config.spatial_compression_ratio,
)
self.quant_conv = nn.Conv3d(2 * arch_config.latent_channels,
fastvideo/v1/models/vaes/wanvae.py
+22 -14
View File
@@ -22,9 +22,10 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
from fastvideo.v1.configs.models.vaes import WanVAEArchConfig, WanVAEConfig
from fastvideo.v1.layers.activation import get_act_fn
from fastvideo.v1.models.vaes.common import ParallelTiledVAE, DiagonalGaussianDistribution
from fastvideo.v1.configs.models.vaes import WanVAEConfig, WanVAEArchConfig
from fastvideo.v1.models.vaes.common import (DiagonalGaussianDistribution,
ParallelTiledVAE)
CACHE_T = 2
@@ -780,32 +781,39 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
_supports_gradient_checkpointing = False
def __init__(self,
config: WanVAEConfig,
) -> None:
super().__init__()
def __init__(
self,
config: WanVAEConfig,
) -> None:
super().__init__(config)
ParallelTiledVAE.__init__(self, config)
self.arch_config = cast(WanVAEArchConfig, self.arch_config)
self.z_dim = self.arch_config.z_dim
self.temperal_downsample = list(self.arch_config.temperal_downsample)
self.temperal_upsample = list(self.arch_config.temperal_downsample)[::-1]
self.temperal_upsample = list(
self.arch_config.temperal_downsample)[::-1]
self.latents_mean = list(self.arch_config.latents_mean)
self.latents_std = list(self.arch_config.latents_std)
self.shift_factor = self.arch_config.shift_factor
if config.load_encoder:
self.encoder = WanEncoder3d(self.arch_config.base_dim, self.z_dim * 2, self.arch_config.dim_mult,
self.arch_config.num_res_blocks, self.arch_config.attn_scales,
self.temperal_downsample, self.arch_config.dropout)
self.encoder = WanEncoder3d(
self.arch_config.base_dim, self.z_dim * 2,
self.arch_config.dim_mult, self.arch_config.num_res_blocks,
self.arch_config.attn_scales, self.temperal_downsample,
self.arch_config.dropout)
self.quant_conv = WanCausalConv3d(self.z_dim * 2, self.z_dim * 2, 1)
self.post_quant_conv = WanCausalConv3d(self.z_dim, self.z_dim, 1)
if config.load_decoder:
self.decoder = WanDecoder3d(self.arch_config.base_dim, self.z_dim, self.arch_config.dim_mult,
self.arch_config.num_res_blocks, self.arch_config.attn_scales,
self.temperal_upsample, self.arch_config.dropout)
self.decoder = WanDecoder3d(self.arch_config.base_dim, self.z_dim,
self.arch_config.dim_mult,
self.arch_config.num_res_blocks,
self.arch_config.attn_scales,
self.temperal_upsample,
self.arch_config.dropout)
self.use_feature_cache = config.use_feature_cache
fastvideo/v1/pipelines/composed_pipeline_base.py
+1 -1
View File
@@ -118,7 +118,7 @@ class ComposedPipelineBase(ABC):
"""
Initialize the pipeline.
"""
pass
return
def load_modules(self, fastvideo_args: FastVideoArgs) -> Dict[str, Any]:
"""
fastvideo/v1/pipelines/pipeline_batch_info.py
+1 -1
View File
@@ -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, fields
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union
import torch
fastvideo/v1/pipelines/stages/decoding.py
+1 -1
View File
@@ -7,10 +7,10 @@ import torch
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.logger import init_logger
from fastvideo.v1.models.vaes.common import ParallelTiledVAE
from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.v1.pipelines.stages.base import PipelineStage
from fastvideo.v1.utils import PRECISION_TO_TYPE
from fastvideo.v1.models.vaes.common import ParallelTiledVAE
logger = init_logger(__name__)
fastvideo/v1/pipelines/stages/encoding.py
+6 -5
View File
@@ -9,13 +9,13 @@ import torch
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.logger import init_logger
from fastvideo.v1.models.vaes.common import ParallelTiledVAE
from fastvideo.v1.models.vision_utils import (get_default_height_width,
load_image, normalize,
numpy_to_pt, pil_to_numpy, resize)
from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.v1.pipelines.stages.base import PipelineStage
from fastvideo.v1.utils import PRECISION_TO_TYPE
from fastvideo.v1.models.vaes.common import ParallelTiledVAE
logger = init_logger(__name__)
@@ -50,6 +50,8 @@ class EncodingStage(PipelineStage):
# TODO(will): remove this once we add input/output validation for stages
if image_path is None:
raise ValueError("Image Path must be provided")
assert batch.height is not None
assert batch.width is not None
latent_height = batch.height // self.vae.spatial_compression_ratio
latent_width = batch.width // self.vae.spatial_compression_ratio
@@ -63,8 +65,7 @@ class EncodingStage(PipelineStage):
video_condition = torch.cat([
image,
image.new_zeros(image.shape[0], image.shape[1],
batch.num_frames - 1, batch.height,
batch.width)
batch.num_frames - 1, batch.height, batch.width)
],
dim=2)
video_condition = video_condition.to(device=fastvideo_args.device,
@@ -107,8 +108,8 @@ class EncodingStage(PipelineStage):
else:
latent_condition = latent_condition * self.vae.scaling_factor
mask_lat_size = torch.ones(1, 1, batch.num_frames,
latent_height, latent_width)
mask_lat_size = torch.ones(1, 1, batch.num_frames, latent_height,
latent_width)
mask_lat_size[:, :, list(range(1, batch.num_frames))] = 0
first_frame_mask = mask_lat_size[:, :, 0:1]
first_frame_mask = torch.repeat_interleave(
fastvideo/v1/pipelines/stages/input_validation.py
+1
View File
@@ -27,6 +27,7 @@ class InputValidationStage(PipelineStage):
seed = batch.seed
num_videos_per_prompt = batch.num_videos_per_prompt
assert seed is not None
seeds = [seed + i for i in range(num_videos_per_prompt)]
batch.seeds = seeds
# Peiyuan: using GPU seed will cause A100 and H100 to generate different results...
fastvideo/v1/pipelines/stages/latent_preparation.py
+4 -2
View File
@@ -74,8 +74,10 @@ class LatentPreparationStage(PipelineStage):
batch_size,
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,
height //
fastvideo_args.vae_config.arch_config.spatial_compression_ratio,
width //
fastvideo_args.vae_config.arch_config.spatial_compression_ratio,
)
# Validate generator if it's a list
fastvideo/v1/sample/v1_fastvideo_inference.py
+2
View File
@@ -1,3 +1,5 @@
# type: ignore
# SPDX-License-Identifier: Apache-2.0
import os
fastvideo/v1/utils.py
+4 -2
View File
@@ -13,7 +13,7 @@ import signal
import sys
import tempfile
import traceback
from dataclasses import is_dataclass, fields
from dataclasses import fields, is_dataclass
from functools import partial, wraps
from typing import (Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar,
Union, cast)
@@ -550,11 +550,13 @@ def run_method(obj: Any, method: Union[str, bytes, Callable], args: tuple[Any],
func = partial(method, obj) # type: ignore
return func(*args, **kwargs)
def shallow_asdict(obj):
def shallow_asdict(obj) -> Dict[str, Any]:
if not is_dataclass(obj):
raise TypeError("Expected dataclass instance")
return {f.name: getattr(obj, f.name) for f in fields(obj)}
def kill_itself_when_parent_died() -> None:
# if sys.platform == "linux":
# sigkill this process when parent worker manager dies