diff --git a/fastvideo/v1/configs/models/__init__.py b/fastvideo/v1/configs/models/__init__.py index 249540597..422763cd1 100644 --- a/fastvideo/v1/configs/models/__init__.py +++ b/fastvideo/v1/configs/models/__init__.py @@ -1,7 +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 __all__ = [ "ArchConfig", "ModelConfig", - "VAEArchConfig", "VAEConfig" + "VAEArchConfig", "VAEConfig", + "DiTArchConfig", "DiTConfig" ] \ No newline at end of file diff --git a/fastvideo/v1/configs/models/base.py b/fastvideo/v1/configs/models/base.py index 9a12655be..945f23efd 100644 --- a/fastvideo/v1/configs/models/base.py +++ b/fastvideo/v1/configs/models/base.py @@ -30,6 +30,9 @@ class ModelConfig: setattr(arch_config, key, value) else: 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, @@ -44,4 +47,7 @@ class ModelConfig: setattr(self, key, value) else: print(f"{type(self).__name__} does not contain field '{key}'!") - raise AttributeError(f"Invalid field: {key}") \ No newline at end of file + raise AttributeError(f"Invalid field: {key}") + + if hasattr(self, "__post_init__"): + self.__post_init__() \ No newline at end of file diff --git a/fastvideo/v1/configs/models/dits/__init__.py b/fastvideo/v1/configs/models/dits/__init__.py new file mode 100644 index 000000000..d7b4cd7cc --- /dev/null +++ b/fastvideo/v1/configs/models/dits/__init__.py @@ -0,0 +1,7 @@ +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" +] \ No newline at end of file diff --git a/fastvideo/v1/configs/models/dits/base.py b/fastvideo/v1/configs/models/dits/base.py new file mode 100644 index 000000000..75483d741 --- /dev/null +++ b/fastvideo/v1/configs/models/dits/base.py @@ -0,0 +1,25 @@ +from dataclasses import dataclass, field +from typing import Tuple + +from fastvideo.v1.configs.models import ArchConfig, ModelConfig +from fastvideo.v1.platforms import _Backend + +@dataclass +class DiTArchConfig(ArchConfig): + _fsdp_shard_conditions: list = field(default_factory=list) + _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() + + # FastVideoDiT-specific parameters + prefix: str = "" + + # (Todo: Wei) Add quant config and teacache config here diff --git a/fastvideo/v1/configs/models/dits/hunyuanvideo.py b/fastvideo/v1/configs/models/dits/hunyuanvideo.py new file mode 100644 index 000000000..fcc70ab1e --- /dev/null +++ b/fastvideo/v1/configs/models/dits/hunyuanvideo.py @@ -0,0 +1,168 @@ +from dataclasses import dataclass, field +from typing import Tuple, Optional + +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", + + # 3. x_embedder mapping: + r"^x_embedder\.proj\.(.*)$": + r"img_in.proj.\1", + + # 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", + + # 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", + }) + + patch_size: int = 2 + patch_size_t: int = 1 + in_channels: int = 16 + out_channels: int = 16 + num_attention_heads: int = 24 + attention_head_dim: int = 128 + mlp_ratio: float = 4.0 + num_layers: int = 20 + num_single_layers: int = 40 + num_refiner_layers: int = 2 + rope_axes_dim: Tuple[int, int, int] = (16, 56, 56) + guidance_embeds: bool = False + dtype: Optional[torch.dtype] = None + text_embed_dim: int = 4096 + pooled_projection_dim: int = 768 + rope_theta: int = 256 + qk_norm: str = "rms_norm" + + def __post_init__(self): + 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" \ No newline at end of file diff --git a/fastvideo/v1/configs/models/dits/wanvideo.py b/fastvideo/v1/configs/models/dits/wanvideo.py new file mode 100644 index 000000000..853c1d4ba --- /dev/null +++ b/fastvideo/v1/configs/models/dits/wanvideo.py @@ -0,0 +1,81 @@ +from dataclasses import dataclass, field +from typing import Tuple, Optional + +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", + }) + + patch_size: Tuple[int, int, int] = (1, 2, 2) + text_len=512 + num_attention_heads: int = 40 + attention_head_dim: int = 128 + in_channels: int = 16 + out_channels: int = 16 + text_dim: int = 4096 + freq_dim: int = 256 + ffn_dim: int = 13824 + num_layers: int = 40 + cross_attn_norm: bool = True + qk_norm: str = "rms_norm_across_heads" + eps: float = 1e-6 + image_dim: Optional[int] = None + added_kv_proj_dim: Optional[int] = None + rope_max_seq_len: int = 1024 + + def __post_init__(self): + self.out_channels = self.out_channels or self.in_channels + 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" \ No newline at end of file diff --git a/fastvideo/v1/configs/pipelines/base.py b/fastvideo/v1/configs/pipelines/base.py index b2451e72a..e9621feb2 100644 --- a/fastvideo/v1/configs/pipelines/base.py +++ b/fastvideo/v1/configs/pipelines/base.py @@ -2,7 +2,7 @@ from dataclasses import dataclass, asdict, fields from typing import Optional, Dict, Any import json -from fastvideo.v1.configs.models import ModelConfig, VAEConfig +from fastvideo.v1.configs.models import ModelConfig, VAEConfig, DiTConfig from fastvideo.v1.utils import shallow_asdict from fastvideo.v1.logger import init_logger @@ -26,6 +26,9 @@ class BaseConfig: vae_sp: bool = True vae_config: VAEConfig = VAEConfig() + # DiT configuration + dit_config: DiTConfig = DiTConfig() + # Image encoder configuration image_encoder_precision: str = "fp32" diff --git a/fastvideo/v1/configs/pipelines/hunyuan.py b/fastvideo/v1/configs/pipelines/hunyuan.py index 18b6d2568..d3e61aa14 100644 --- a/fastvideo/v1/configs/pipelines/hunyuan.py +++ b/fastvideo/v1/configs/pipelines/hunyuan.py @@ -1,14 +1,17 @@ from dataclasses import dataclass from fastvideo.v1.configs.pipelines.base import BaseConfig -from fastvideo.v1.configs.models import VAEConfig +from fastvideo.v1.configs.models import VAEConfig, DiTConfig from fastvideo.v1.configs.models.vaes import HunyuanVAEConfig +from fastvideo.v1.configs.models.dits import HunyuanVideoConfig @dataclass class HunyuanConfig(BaseConfig): """Base configuration for HunYuan pipeline architecture.""" # HunyuanConfig-specific parameters with defaults + # DiT + dit_config: DiTConfig = HunyuanVideoConfig() # VAE vae_config: VAEConfig = HunyuanVAEConfig() # Denoising stage diff --git a/fastvideo/v1/configs/pipelines/wan.py b/fastvideo/v1/configs/pipelines/wan.py index a7e7f921b..8a27bdfe2 100644 --- a/fastvideo/v1/configs/pipelines/wan.py +++ b/fastvideo/v1/configs/pipelines/wan.py @@ -1,14 +1,17 @@ from dataclasses import dataclass from fastvideo.v1.configs.pipelines.base import BaseConfig -from fastvideo.v1.configs.models import VAEConfig +from fastvideo.v1.configs.models import VAEConfig, DiTConfig from fastvideo.v1.configs.models.vaes import WanVAEConfig +from fastvideo.v1.configs.models.dits import WanVideoConfig @dataclass class WanT2V480PConfig(BaseConfig): """Base configuration for Wan T2V 1.3B pipeline architecture.""" # WanConfig-specific parameters with defaults + # DiT + dit_config: DiTConfig = WanVideoConfig() # VAE vae_config: VAEConfig = WanVAEConfig() vae_tiling: bool = False diff --git a/fastvideo/v1/configs/wan_1.3B_t2v_pipeline.json b/fastvideo/v1/configs/wan_1.3B_t2v_pipeline.json index 46ba70374..708b25653 100644 --- a/fastvideo/v1/configs/wan_1.3B_t2v_pipeline.json +++ b/fastvideo/v1/configs/wan_1.3B_t2v_pipeline.json @@ -22,6 +22,9 @@ "use_parallel_tiling": false, "use_feature_cache": true }, + "dit_config": { + "prefix": "Wan" + }, "image_encoder_precision": "fp32", "text_encoder_precision": "fp32", "text_len": 512, diff --git a/fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json b/fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json index aba468a1f..6055ded55 100644 --- a/fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json +++ b/fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json @@ -22,6 +22,9 @@ "use_parallel_tiling": false, "use_feature_cache": true }, + "dit_config": { + "prefix": "Wan" + }, "image_encoder_precision": "fp32", "text_encoder_precision": "fp32", "text_len": 512, diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py index 2845fa3a7..4398d6aaa 100644 --- a/fastvideo/v1/fastvideo_args.py +++ b/fastvideo/v1/fastvideo_args.py @@ -10,7 +10,7 @@ from typing import List, Optional from fastvideo.v1.logger import init_logger from fastvideo.v1.utils import FlexibleArgumentParser -from fastvideo.v1.configs.models import VAEConfig +from fastvideo.v1.configs.models import VAEConfig, DiTConfig logger = init_logger(__name__) @@ -41,7 +41,8 @@ class FastVideoArgs: output_type: str = "pil" - # Model configuration + # DiT configuration + dit_config: DiTConfig = DiTConfig() precision: str = "bf16" # VAE configuration diff --git a/fastvideo/v1/models/dits/base.py b/fastvideo/v1/models/dits/base.py index 410b27ed0..2a3a3031d 100644 --- a/fastvideo/v1/models/dits/base.py +++ b/fastvideo/v1/models/dits/base.py @@ -6,12 +6,11 @@ import torch from torch import nn from fastvideo.v1.platforms import _Backend - +from fastvideo.v1.configs.models import DiTConfig # TODO class BaseDiT(nn.Module, ABC): _fsdp_shard_conditions: list = [] - attention_head_dim: int | None = None _param_names_mapping: dict hidden_size: int num_attention_heads: int @@ -31,8 +30,10 @@ class BaseDiT(nn.Module, ABC): f"Subclasses of BaseDiT must define '{attr}' class variable" ) - def __init__(self, *args, **kwargs) -> None: + def __init__(self, config: DiTConfig, **kwargs) -> None: super().__init__() + self.config = config + self.arch_config = config.arch_config if not self.supported_attention_backends: raise ValueError( f"Subclass {self.__class__.__name__} must define _supported_attention_backends" diff --git a/fastvideo/v1/models/dits/hunyuanvideo.py b/fastvideo/v1/models/dits/hunyuanvideo.py index fb7f99ed5..558ae45db 100644 --- a/fastvideo/v1/models/dits/hunyuanvideo.py +++ b/fastvideo/v1/models/dits/hunyuanvideo.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 -from typing import List, Optional, Tuple, Union +from typing import List, Optional, Tuple, Union, cast import torch import torch.nn as nn @@ -20,6 +20,7 @@ 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): @@ -431,240 +432,101 @@ class HunyuanVideoTransformer3DModel(BaseDiT): # PY: we make the input args the same as HF config # shard single stream, double stream blocks, and refiner_blocks - _fsdp_shard_conditions = [ - lambda n, m: "double" in n and str.isdigit(n.split(".")[-1]), - lambda n, m: "single" in n and str.isdigit(n.split(".")[-1]), - lambda n, m: "refiner" in n and str.isdigit(n.split(".")[-1]), - ] - _supported_attention_backends = (_Backend.SLIDING_TILE_ATTN, - _Backend.SAGE_ATTN, _Backend.FLASH_ATTN, - _Backend.TORCH_SDPA) - _param_names_mapping = { - # 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", - - # 3. x_embedder mapping: - r"^x_embedder\.proj\.(.*)$": - r"img_in.proj.\1", - - # 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", - - # 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", - } + _fsdp_shard_conditions = HunyuanVideoArchConfig()._fsdp_shard_conditions + _supported_attention_backends = HunyuanVideoArchConfig()._supported_attention_backends + _param_names_mapping = HunyuanVideoArchConfig()._param_names_mapping def __init__( self, - patch_size: int = 2, - patch_size_t: int = 1, - in_channels: int = 16, - out_channels: int = 16, - num_attention_heads: int = 24, - attention_head_dim: int = 128, - mlp_ratio: float = 4.0, - num_layers: int = 20, - num_single_layers: int = 40, - num_refiner_layers: int = 2, - rope_axes_dim: Tuple[int, int, int] = (16, 56, 56), - guidance_embeds: bool = False, - dtype: Optional[torch.dtype] = None, - text_embed_dim: int = 4096, - pooled_projection_dim: int = 768, - rope_theta: int = 256, - qk_norm: str = "rms_norm", #TODO(PY) - prefix="Hunyuan", + config: HunyuanVideoConfig ): - super().__init__() - 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 + 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] + 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 self.unpatchify_channels = self.out_channels - self.guidance_embeds = guidance_embeds - self.rope_dim_list = list(rope_axes_dim) - self.rope_theta = rope_theta - self.text_states_dim = text_embed_dim - self.text_states_dim_2 = pooled_projection_dim + self.guidance_embeds = arch_config.guidance_embeds + self.rope_dim_list = list(arch_config.rope_axes_dim) + self.rope_theta = arch_config.rope_theta + self.text_states_dim = arch_config.text_embed_dim + self.text_states_dim_2 = arch_config.pooled_projection_dim # TODO(will): hack? - self.dtype = dtype + self.dtype = arch_config.dtype - if hidden_size % num_attention_heads != 0: + pe_dim = arch_config.hidden_size // arch_config.num_attention_heads + if sum(arch_config.rope_axes_dim) != pe_dim: raise ValueError( - f"Hidden size {hidden_size} must be divisible by num_attention_heads {num_attention_heads}" - ) + f"Got {arch_config.rope_axes_dim} but expected positional dim {pe_dim}") - pe_dim = hidden_size // num_attention_heads - if sum(rope_axes_dim) != pe_dim: - raise ValueError( - f"Got {rope_axes_dim} but expected positional dim {pe_dim}") - - self.hidden_size = hidden_size - self.num_attention_heads = num_attention_heads + self.hidden_size = arch_config.hidden_size + self.num_attention_heads = arch_config.num_attention_heads + self.num_channels_latents = arch_config.num_channels_latents # Image projection self.img_in = PatchEmbed(self.patch_size, self.in_channels, self.hidden_size, - dtype=dtype, - prefix=f"{prefix}.img_in") + dtype=arch_config.dtype, + prefix=f"{config.prefix}.img_in") self.txt_in = SingleTokenRefiner(self.text_states_dim, - hidden_size, - num_attention_heads, - depth=num_refiner_layers, - dtype=dtype, - prefix=f"{prefix}.txt_in") + arch_config.hidden_size, + arch_config.num_attention_heads, + depth=arch_config.num_refiner_layers, + dtype=arch_config.dtype, + prefix=f"{config.prefix}.txt_in") # Time modulation self.time_in = TimestepEmbedder(self.hidden_size, act_layer="silu", - dtype=dtype, - prefix=f"{prefix}.time_in") + dtype=arch_config.dtype, + prefix=f"{config.prefix}.time_in") # Text modulation self.vector_in = MLP(self.text_states_dim_2, self.hidden_size, self.hidden_size, act_type="silu", - dtype=dtype, - prefix=f"{prefix}.vector_in") + dtype=arch_config.dtype, + prefix=f"{config.prefix}.vector_in") # Guidance modulation self.guidance_in = (TimestepEmbedder(self.hidden_size, act_layer="silu", - dtype=dtype, - prefix=f"{prefix}.guidance_in") + dtype=arch_config.dtype, + prefix=f"{config.prefix}.guidance_in") if self.guidance_embeds else None) # Double blocks self.double_blocks = nn.ModuleList([ MMDoubleStreamBlock( - hidden_size, - num_attention_heads, - mlp_ratio=mlp_ratio, - dtype=dtype, + arch_config.hidden_size, + arch_config.num_attention_heads, + mlp_ratio=arch_config.mlp_ratio, + dtype=arch_config.dtype, supported_attention_backends=self._supported_attention_backends, - prefix=f"{prefix}.double_blocks.{i}") for i in range(num_layers) + prefix=f"{config.prefix}.double_blocks.{i}") for i in range(arch_config.num_layers) ]) # Single blocks self.single_blocks = nn.ModuleList([ MMSingleStreamBlock( - hidden_size, - num_attention_heads, - mlp_ratio=mlp_ratio, - dtype=dtype, + arch_config.hidden_size, + arch_config.num_attention_heads, + mlp_ratio=arch_config.mlp_ratio, + dtype=arch_config.dtype, supported_attention_backends=self._supported_attention_backends, - prefix=f"{prefix}.single_blocks.{i+num_layers}") - for i in range(num_single_layers) + prefix=f"{config.prefix}.single_blocks.{i+arch_config.num_layers}") + for i in range(arch_config.num_single_layers) ]) - self.final_layer = FinalLayer(hidden_size, + self.final_layer = FinalLayer(arch_config.hidden_size, self.patch_size, self.out_channels, - dtype=dtype, - prefix=f"{prefix}.final_layer") + dtype=arch_config.dtype, + prefix=f"{config.prefix}.final_layer") self.__post_init__() diff --git a/fastvideo/v1/models/dits/wanvideo.py b/fastvideo/v1/models/dits/wanvideo.py index 2a2fcc515..13d1ec3e4 100644 --- a/fastvideo/v1/models/dits/wanvideo.py +++ b/fastvideo/v1/models/dits/wanvideo.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 import math -from typing import List, Optional, Tuple, Union +from typing import List, Optional, Tuple, Union, cast import torch import torch.nn as nn @@ -22,6 +22,7 @@ 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): @@ -350,116 +351,61 @@ class WanTransformerBlock(nn.Module): class WanTransformer3DModel(BaseDiT): - _fsdp_shard_conditions = [ - lambda n, m: "blocks" in n and str.isdigit(n.split(".")[-1]), - ] - _supported_attention_backends = (_Backend.SLIDING_TILE_ATTN, - _Backend.SAGE_ATTN, _Backend.FLASH_ATTN, - _Backend.TORCH_SDPA) - _param_names_mapping = { - 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", - } + _fsdp_shard_conditions = WanVideoArchConfig()._fsdp_shard_conditions + _supported_attention_backends = WanVideoArchConfig()._supported_attention_backends + _param_names_mapping = WanVideoArchConfig()._param_names_mapping def __init__(self, - patch_size: Tuple[int, int, int] = (1, 2, 2), - text_len=512, - num_attention_heads: int = 40, - attention_head_dim: int = 128, - in_channels: int = 16, - out_channels: int = 16, - text_dim: int = 4096, - freq_dim: int = 256, - ffn_dim: int = 13824, - num_layers: int = 40, - cross_attn_norm: bool = True, - qk_norm: str = "rms_norm_across_heads", - eps: float = 1e-6, - image_dim: Optional[int] = None, - added_kv_proj_dim: Optional[int] = None, - rope_max_seq_len: int = 1024, - prefix="Wan") -> None: - super().__init__() + config: WanVideoConfig) -> None: + super().__init__(config=config) - inner_dim = num_attention_heads * attention_head_dim - self.hidden_size = inner_dim - 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 + 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 + self.num_attention_heads = arch_config.num_attention_heads + self.in_channels = arch_config.in_channels + self.out_channels = arch_config.out_channels + self.num_channels_latents = arch_config.num_channels_latents + self.patch_size = arch_config.patch_size + self.text_len = arch_config.text_len # 1. Patch & position embedding - self.patch_embedding = PatchEmbed(in_chans=in_channels, + self.patch_embedding = PatchEmbed(in_chans=arch_config.in_channels, embed_dim=inner_dim, - patch_size=patch_size, + patch_size=arch_config.patch_size, flatten=False) # 2. Condition embeddings self.condition_embedder = WanTimeTextImageEmbedding( dim=inner_dim, - time_freq_dim=freq_dim, - text_embed_dim=text_dim, - image_embed_dim=image_dim, + time_freq_dim=arch_config.freq_dim, + text_embed_dim=arch_config.text_dim, + image_embed_dim=arch_config.image_dim, ) # 3. Transformer blocks self.blocks = nn.ModuleList([ WanTransformerBlock(inner_dim, - ffn_dim, - num_attention_heads, - qk_norm, - cross_attn_norm, - eps, - added_kv_proj_dim, + arch_config.ffn_dim, + arch_config.num_attention_heads, + arch_config.qk_norm, + arch_config.cross_attn_norm, + arch_config.eps, + arch_config.added_kv_proj_dim, self._supported_attention_backends, - prefix=f"{prefix}.blocks.{i}") - for i in range(num_layers) + prefix=f"{config.prefix}.blocks.{i}") + for i in range(arch_config.num_layers) ]) # 4. Output norm & projection self.norm_out = LayerNormScaleShift(inner_dim, norm_type="layer", - eps=eps, + eps=arch_config.eps, elementwise_affine=False, dtype=torch.float32) self.proj_out = nn.Linear(inner_dim, - out_channels * math.prod(patch_size)) + 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) diff --git a/fastvideo/v1/models/loader/component_loader.py b/fastvideo/v1/models/loader/component_loader.py index 83184db3b..0cf63d78c 100644 --- a/fastvideo/v1/models/loader/component_loader.py +++ b/fastvideo/v1/models/loader/component_loader.py @@ -339,17 +339,17 @@ class TransformerLoader(ComponentLoader): def load(self, model_path: str, architecture: str, fastvideo_args: FastVideoArgs): """Load the transformer based on the model path, architecture, and inference args.""" - model_config = get_diffusers_config(model=model_path) - cls_name = model_config.pop("_class_name") + config = get_diffusers_config(model=model_path) + cls_name = config.pop("_class_name") if cls_name is None: raise ValueError( "Model config does not contain a _class_name attribute. " "Only diffusers format is supported.") - model_config.pop("_diffusers_version") + config.pop("_diffusers_version") # Config from Diffusers supercedes fastvideo's model config - # dit_config = fastvideo_args.dit_config - # model_config.update(dit_config) + dit_config = fastvideo_args.dit_config + dit_config.update_model_arch(config) model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) @@ -368,7 +368,7 @@ class TransformerLoader(ComponentLoader): # Load the model using FSDP loader logger.info("Loading model from %s", cls_name) model = load_fsdp_model(model_cls=model_cls, - init_params=model_config, + init_params={"config": dit_config}, weight_dir_list=safetensors_list, device=fastvideo_args.device, cpu_offload=fastvideo_args.use_cpu_offload, diff --git a/fastvideo/v1/tests/transformers/test_hunyuanvideo_load.py b/fastvideo/v1/tests/transformers/test_hunyuanvideo_load.py index 1b4d9c7c0..2ec322dd2 100644 --- a/fastvideo/v1/tests/transformers/test_hunyuanvideo_load.py +++ b/fastvideo/v1/tests/transformers/test_hunyuanvideo_load.py @@ -10,11 +10,11 @@ from fastvideo.v1.distributed.parallel_state import ( get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size) from fastvideo.v1.logger import init_logger -from fastvideo.v1.models.dits.hunyuanvideo import ( - HunyuanVideoTransformer3DModel as HunyuanVideoDit) -from fastvideo.v1.models.loader.fsdp_load import load_fsdp_model +from fastvideo.v1.models.loader.component_loader import TransformerLoader +from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.utils import maybe_download_model from fastvideo.v1.forward_context import set_forward_context +from fastvideo.v1.configs.models.dits import HunyuanVideoConfig logger = init_logger(__name__) @@ -59,13 +59,15 @@ def test_hunyuanvideo_distributed(): config.pop("_class_name") config.pop("_diffusers_version") - weight_dir_list = glob.glob(os.path.join(TRANSFORMER_PATH, "*.safetensors")) - weight_dir_list = [str(path) for path in weight_dir_list] - model = load_fsdp_model(HunyuanVideoDit, - init_params=config, - weight_dir_list=weight_dir_list, - device=torch.device(f"cuda:{LOCAL_RANK}"), - cpu_offload=False) + precision_str = "bf16" + args = FastVideoArgs(model_path=TRANSFORMER_PATH, + use_cpu_offload=False, + precision=precision_str) + args.device = torch.device(f"cuda:{LOCAL_RANK}") + args.dit_config = HunyuanVideoConfig() + + loader = TransformerLoader() + model = loader.load(TRANSFORMER_PATH, "", args) model.eval() diff --git a/fastvideo/v1/tests/transformers/test_wanvideo.py b/fastvideo/v1/tests/transformers/test_wanvideo.py index f792ec00f..2eac2cd08 100644 --- a/fastvideo/v1/tests/transformers/test_wanvideo.py +++ b/fastvideo/v1/tests/transformers/test_wanvideo.py @@ -11,6 +11,7 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.logger import init_logger from fastvideo.v1.models.loader.component_loader import TransformerLoader from fastvideo.v1.utils import maybe_download_model +from fastvideo.v1.configs.models.dits import WanVideoConfig logger = init_logger(__name__) @@ -33,6 +34,7 @@ def test_wan_transformer(): use_cpu_offload=False, precision=precision_str) args.device = device + args.dit_config = WanVideoConfig() loader = TransformerLoader() model2 = loader.load(TRANSFORMER_PATH, "", args).to(device, dtype=precision) @@ -113,6 +115,8 @@ def test_wan_transformer(): # Check if outputs are similar (allowing for small numerical differences) max_diff = torch.max(torch.abs(output1 - output2)) mean_diff = torch.mean(torch.abs(output1 - output2)) + logger.info("Max Diff: %s", max_diff.item()) + logger.info("Mean Diff: %s", mean_diff.item()) assert max_diff < 1e-1, f"Maximum difference between outputs: {max_diff.item()}" # mean diff assert mean_diff < 1e-2, f"Mean difference between outputs: {mean_diff.item()}" \ No newline at end of file