Add model config for DiTs
This commit is contained in:
JerryZhou54
@@ -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"
|
||||
]
|
||||
@@ -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}")
|
||||
raise AttributeError(f"Invalid field: {key}")
|
||||
|
||||
if hasattr(self, "__post_init__"):
|
||||
self.__post_init__()
|
||||
@@ -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"
|
||||
]
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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"
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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__()
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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()}"
|
||||
Reference in New Issue
Block a user