From 60f01567d1c6a3075ddae9cf986be1c6631442e0 Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Tue, 29 Apr 2025 03:24:04 +0000 Subject: [PATCH] Pass pre-commit tests --- fastvideo/v1/configs/models/__init__.py | 10 +- fastvideo/v1/configs/models/base.py | 36 ++- fastvideo/v1/configs/models/dits/__init__.py | 12 +- fastvideo/v1/configs/models/dits/base.py | 8 +- .../v1/configs/models/dits/hunyuanvideo.py | 243 +++++++++--------- fastvideo/v1/configs/models/dits/wanvideo.py | 93 +++---- .../v1/configs/models/encoders/__init__.py | 22 +- fastvideo/v1/configs/models/encoders/base.py | 18 +- fastvideo/v1/configs/models/encoders/clip.py | 19 +- fastvideo/v1/configs/models/encoders/llama.py | 9 +- fastvideo/v1/configs/models/encoders/t5.py | 7 +- fastvideo/v1/configs/models/vaes/__init__.py | 12 +- fastvideo/v1/configs/models/vaes/base.py | 2 + .../v1/configs/models/vaes/hunyuanvae.py | 9 +- fastvideo/v1/configs/models/vaes/wanvae.py | 13 +- fastvideo/v1/configs/pipelines/__init__.py | 12 +- fastvideo/v1/configs/pipelines/base.py | 36 +-- fastvideo/v1/configs/pipelines/hunyuan.py | 7 +- fastvideo/v1/configs/pipelines/registry.py | 7 +- fastvideo/v1/configs/pipelines/wan.py | 10 +- fastvideo/v1/configs/quantization/__init__.py | 4 +- fastvideo/v1/configs/quantization/base.py | 3 +- fastvideo/v1/configs/sample/__init__.py | 4 +- fastvideo/v1/configs/sample/base.py | 30 ++- fastvideo/v1/configs/sample/hunyuan.py | 4 +- fastvideo/v1/configs/sample/registry.py | 16 +- fastvideo/v1/configs/sample/wan.py | 4 +- fastvideo/v1/entrypoints/video_generator.py | 26 +- fastvideo/v1/fastvideo_args.py | 19 +- fastvideo/v1/inference_engine.py | 2 + fastvideo/v1/models/dits/base.py | 7 +- fastvideo/v1/models/dits/hunyuanvideo.py | 38 +-- fastvideo/v1/models/dits/wanvideo.py | 16 +- fastvideo/v1/models/encoders/base.py | 7 +- fastvideo/v1/models/encoders/clip.py | 24 +- fastvideo/v1/models/encoders/llama.py | 39 +-- fastvideo/v1/models/encoders/t5.py | 6 +- .../v1/models/loader/component_loader.py | 19 +- fastvideo/v1/models/vaes/common.py | 6 +- fastvideo/v1/models/vaes/hunyuanvae.py | 13 +- fastvideo/v1/models/vaes/wanvae.py | 36 ++- .../v1/pipelines/composed_pipeline_base.py | 2 +- fastvideo/v1/pipelines/pipeline_batch_info.py | 2 +- fastvideo/v1/pipelines/stages/decoding.py | 2 +- fastvideo/v1/pipelines/stages/encoding.py | 11 +- .../v1/pipelines/stages/input_validation.py | 1 + .../v1/pipelines/stages/latent_preparation.py | 6 +- fastvideo/v1/sample/v1_fastvideo_inference.py | 2 + fastvideo/v1/utils.py | 6 +- 49 files changed, 509 insertions(+), 431 deletions(-) diff --git a/fastvideo/v1/configs/models/__init__.py b/fastvideo/v1/configs/models/__init__.py index dcc139a9d..d20f41c61 100644 --- a/fastvideo/v1/configs/models/__init__.py +++ b/fastvideo/v1/configs/models/__init__.py @@ -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" -] \ No newline at end of file + "ArchConfig", "ModelConfig", "VAEArchConfig", "VAEConfig", "DiTArchConfig", + "DiTConfig", "EncoderConfig" +] diff --git a/fastvideo/v1/configs/models/base.py b/fastvideo/v1/configs/models/base.py index c2652300b..ddd19e5dd 100644 --- a/fastvideo/v1/configs/models/base.py +++ b/fastvideo/v1/configs/models/base.py @@ -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__() \ No newline at end of file + self.__post_init__() diff --git a/fastvideo/v1/configs/models/dits/__init__.py b/fastvideo/v1/configs/models/dits/__init__.py index d7b4cd7cc..a64bc4102 100644 --- a/fastvideo/v1/configs/models/dits/__init__.py +++ b/fastvideo/v1/configs/models/dits/__init__.py @@ -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" -] \ No newline at end of file + "HunyuanVideoArchConfig", "HunyuanVideoConfig", "WanVideoArchConfig", + "WanVideoConfig" +] diff --git a/fastvideo/v1/configs/models/dits/base.py b/fastvideo/v1/configs/models/dits/base.py index d8139cdb2..7bff63c6f 100644 --- a/fastvideo/v1/configs/models/dits/base.py +++ b/fastvideo/v1/configs/models/dits/base.py @@ -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() diff --git a/fastvideo/v1/configs/models/dits/hunyuanvideo.py b/fastvideo/v1/configs/models/dits/hunyuanvideo.py index fcc70ab1e..77288848c 100644 --- a/fastvideo/v1/configs/models/dits/hunyuanvideo.py +++ b/fastvideo/v1/configs/models/dits/hunyuanvideo.py @@ -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" \ No newline at end of file + prefix: str = "Hunyuan" diff --git a/fastvideo/v1/configs/models/dits/wanvideo.py b/fastvideo/v1/configs/models/dits/wanvideo.py index 853c1d4ba..e73110928 100644 --- a/fastvideo/v1/configs/models/dits/wanvideo.py +++ b/fastvideo/v1/configs/models/dits/wanvideo.py @@ -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" \ No newline at end of file + prefix: str = "Wan" diff --git a/fastvideo/v1/configs/models/encoders/__init__.py b/fastvideo/v1/configs/models/encoders/__init__.py index 41d9aac28..ee1c5b0fe 100644 --- a/fastvideo/v1/configs/models/encoders/__init__.py +++ b/fastvideo/v1/configs/models/encoders/__init__.py @@ -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" -] \ No newline at end of file + "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", + "CLIPTextConfig", "CLIPVisionConfig", "LlamaConfig", "T5Config" +] diff --git a/fastvideo/v1/configs/models/encoders/base.py b/fastvideo/v1/configs/models/encoders/base.py index f3b8d6ae5..fb98334ea 100644 --- a/fastvideo/v1/configs/models/encoders/base.py +++ b/fastvideo/v1/configs/models/encoders/base.py @@ -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() \ No newline at end of file + arch_config: ArchConfig = ImageEncoderArchConfig() diff --git a/fastvideo/v1/configs/models/encoders/clip.py b/fastvideo/v1/configs/models/encoders/clip.py index fb9c22b8f..42dd1ded9 100644 --- a/fastvideo/v1/configs/models/encoders/clip.py +++ b/fastvideo/v1/configs/models/encoders/clip.py @@ -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" \ No newline at end of file + prefix: str = "clip" diff --git a/fastvideo/v1/configs/models/encoders/llama.py b/fastvideo/v1/configs/models/encoders/llama.py index cc3958597..73057350a 100644 --- a/fastvideo/v1/configs/models/encoders/llama.py +++ b/fastvideo/v1/configs/models/encoders/llama.py @@ -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" \ No newline at end of file + prefix: str = "llama" diff --git a/fastvideo/v1/configs/models/encoders/t5.py b/fastvideo/v1/configs/models/encoders/t5.py index 99aaba085..c8f469490 100644 --- a/fastvideo/v1/configs/models/encoders/t5.py +++ b/fastvideo/v1/configs/models/encoders/t5.py @@ -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" \ No newline at end of file + prefix: str = "t5" diff --git a/fastvideo/v1/configs/models/vaes/__init__.py b/fastvideo/v1/configs/models/vaes/__init__.py index 2fe9d5725..b20debf0d 100644 --- a/fastvideo/v1/configs/models/vaes/__init__.py +++ b/fastvideo/v1/configs/models/vaes/__init__.py @@ -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" -] \ No newline at end of file + "HunyuanVAEConfig", "HunyuanVAEArchConfig", "WanVAEConfig", + "WanVAEArchConfig" +] diff --git a/fastvideo/v1/configs/models/vaes/base.py b/fastvideo/v1/configs/models/vaes/base.py index 57f814004..9e92b3afa 100644 --- a/fastvideo/v1/configs/models/vaes/base.py +++ b/fastvideo/v1/configs/models/vaes/base.py @@ -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() diff --git a/fastvideo/v1/configs/models/vaes/hunyuanvae.py b/fastvideo/v1/configs/models/vaes/hunyuanvae.py index 3bd368aff..7a4a216bf 100644 --- a/fastvideo/v1/configs/models/vaes/hunyuanvae.py +++ b/fastvideo/v1/configs/models/vaes/hunyuanvae.py @@ -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() \ No newline at end of file + arch_config: VAEArchConfig = HunyuanVAEArchConfig() diff --git a/fastvideo/v1/configs/models/vaes/wanvae.py b/fastvideo/v1/configs/models/vaes/wanvae.py index 46ddee66f..683ee752e 100644 --- a/fastvideo/v1/configs/models/vaes/wanvae.py +++ b/fastvideo/v1/configs/models/vaes/wanvae.py @@ -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 \ No newline at end of file + self.blend_num_frames = (self.tile_sample_min_num_frames - + self.tile_sample_stride_num_frames) * 2 diff --git a/fastvideo/v1/configs/pipelines/__init__.py b/fastvideo/v1/configs/pipelines/__init__.py index 81a34dcde..f3e945b92 100644 --- a/fastvideo/v1/configs/pipelines/__init__.py +++ b/fastvideo/v1/configs/pipelines/__init__.py @@ -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", diff --git a/fastvideo/v1/configs/pipelines/base.py b/fastvideo/v1/configs/pipelines/base.py index 2c475b78a..1a4463f9d 100644 --- a/fastvideo/v1/configs/pipelines/base.py +++ b/fastvideo/v1/configs/pipelines/base.py @@ -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 \ No newline at end of file + use_overlap_optimization: bool = True diff --git a/fastvideo/v1/configs/pipelines/hunyuan.py b/fastvideo/v1/configs/pipelines/hunyuan.py index 1ef682568..d8e0951b0 100644 --- a/fastvideo/v1/configs/pipelines/hunyuan.py +++ b/fastvideo/v1/configs/pipelines/hunyuan.py @@ -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): diff --git a/fastvideo/v1/configs/pipelines/registry.py b/fastvideo/v1/configs/pipelines/registry.py index 3b0252a04..e9e253cb8 100644 --- a/fastvideo/v1/configs/pipelines/registry.py +++ b/fastvideo/v1/configs/pipelines/registry.py @@ -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) diff --git a/fastvideo/v1/configs/pipelines/wan.py b/fastvideo/v1/configs/pipelines/wan.py index 3ec31a426..931c9a9a2 100644 --- a/fastvideo/v1/configs/pipelines/wan.py +++ b/fastvideo/v1/configs/pipelines/wan.py @@ -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.""" diff --git a/fastvideo/v1/configs/quantization/__init__.py b/fastvideo/v1/configs/quantization/__init__.py index f594fd8bc..0273f8a29 100644 --- a/fastvideo/v1/configs/quantization/__init__.py +++ b/fastvideo/v1/configs/quantization/__init__.py @@ -1,5 +1,3 @@ from fastvideo.v1.configs.quantization.base import QuantizationConfig -__all__ = [ - "QuantizationConfig" -] \ No newline at end of file +__all__ = ["QuantizationConfig"] diff --git a/fastvideo/v1/configs/quantization/base.py b/fastvideo/v1/configs/quantization/base.py index b3d7c3c9d..6a8ec17c7 100644 --- a/fastvideo/v1/configs/quantization/base.py +++ b/fastvideo/v1/configs/quantization/base.py @@ -1,5 +1,6 @@ from dataclasses import dataclass + @dataclass class QuantizationConfig: - pass \ No newline at end of file + pass diff --git a/fastvideo/v1/configs/sample/__init__.py b/fastvideo/v1/configs/sample/__init__.py index 15e4b1455..241d00d87 100644 --- a/fastvideo/v1/configs/sample/__init__.py +++ b/fastvideo/v1/configs/sample/__init__.py @@ -1,5 +1,3 @@ from fastvideo.v1.configs.sample.base import SamplingParam -__all__ = [ - "SamplingParam" -] \ No newline at end of file +__all__ = ["SamplingParam"] diff --git a/fastvideo/v1/configs/sample/base.py b/fastvideo/v1/configs/sample/base.py index 14dcb1b56..906abde99 100644 --- a/fastvideo/v1/configs/sample/base.py +++ b/fastvideo/v1/configs/sample/base.py @@ -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 \ No newline at end of file + return sampling_param diff --git a/fastvideo/v1/configs/sample/hunyuan.py b/fastvideo/v1/configs/sample/hunyuan.py index 01bc53cc6..7953b9957 100644 --- a/fastvideo/v1/configs/sample/hunyuan.py +++ b/fastvideo/v1/configs/sample/hunyuan.py @@ -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 \ No newline at end of file + num_inference_steps: int = 6 diff --git a/fastvideo/v1/configs/sample/registry.py b/fastvideo/v1/configs/sample/registry.py index 217e41c60..f9003846e 100644 --- a/fastvideo/v1/configs/sample/registry.py +++ b/fastvideo/v1/configs/sample/registry.py @@ -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 \ No newline at end of file + logger.warning( + "No match found for pipeline %s, using fallback sampling param %s.", + pipeline_name_or_path, fallback_config) + return fallback_config diff --git a/fastvideo/v1/configs/sample/wan.py b/fastvideo/v1/configs/sample/wan.py index 8ae8b2861..2579dfe80 100644 --- a/fastvideo/v1/configs/sample/wan.py +++ b/fastvideo/v1/configs/sample/wan.py @@ -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 \ No newline at end of file + num_inference_steps: int = 40 diff --git a/fastvideo/v1/entrypoints/video_generator.py b/fastvideo/v1/entrypoints/video_generator.py index c35cfd79c..6a2e9663f 100644 --- a/fastvideo/v1/entrypoints/video_generator.py +++ b/fastvideo/v1/entrypoints/video_generator.py @@ -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 } diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py index 3bf88cf4c..514cf72e9 100644 --- a/fastvideo/v1/fastvideo_args.py +++ b/fastvideo/v1/fastvideo_args.py @@ -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( diff --git a/fastvideo/v1/inference_engine.py b/fastvideo/v1/inference_engine.py index 1ab02a2c9..87e7d82a1 100644 --- a/fastvideo/v1/inference_engine.py +++ b/fastvideo/v1/inference_engine.py @@ -1,3 +1,5 @@ +# type: ignore + # SPDX-License-Identifier: Apache-2.0 """ Inference module for diffusion models. diff --git a/fastvideo/v1/models/dits/base.py b/fastvideo/v1/models/dits/base.py index 2a3a3031d..df532f97f 100644 --- a/fastvideo/v1/models/dits/base.py +++ b/fastvideo/v1/models/dits/base.py @@ -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( diff --git a/fastvideo/v1/models/dits/hunyuanvideo.py b/fastvideo/v1/models/dits/hunyuanvideo.py index 558ae45db..0fffa0547 100644 --- a/fastvideo/v1/models/dits/hunyuanvideo.py +++ b/fastvideo/v1/models/dits/hunyuanvideo.py @@ -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) ]) diff --git a/fastvideo/v1/models/dits/wanvideo.py b/fastvideo/v1/models/dits/wanvideo.py index 13d1ec3e4..f5414f937 100644 --- a/fastvideo/v1/models/dits/wanvideo.py +++ b/fastvideo/v1/models/dits/wanvideo.py @@ -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) diff --git a/fastvideo/v1/models/encoders/base.py b/fastvideo/v1/models/encoders/base.py index b97dff5bf..a34ee91f8 100644 --- a/fastvideo/v1/models/encoders/base.py +++ b/fastvideo/v1/models/encoders/base.py @@ -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__() diff --git a/fastvideo/v1/models/encoders/clip.py b/fastvideo/v1/models/encoders/clip.py index 8c10bc597..f1e618153 100644 --- a/fastvideo/v1/models/encoders/clip.py +++ b/fastvideo/v1/models/encoders/clip.py @@ -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, diff --git a/fastvideo/v1/models/encoders/llama.py b/fastvideo/v1/models/encoders/llama.py index d04618a61..6b16449b3 100644 --- a/fastvideo/v1/models/encoders/llama.py +++ b/fastvideo/v1/models/encoders/llama.py @@ -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( diff --git a/fastvideo/v1/models/encoders/t5.py b/fastvideo/v1/models/encoders/t5.py index a27b55516..2767df211 100644 --- a/fastvideo/v1/models/encoders/t5.py +++ b/fastvideo/v1/models/encoders/t5.py @@ -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: diff --git a/fastvideo/v1/models/loader/component_loader.py b/fastvideo/v1/models/loader/component_loader.py index 151e90392..bd2b69093 100644 --- a/fastvideo/v1/models/loader/component_loader.py +++ b/fastvideo/v1/models/loader/component_loader.py @@ -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) diff --git a/fastvideo/v1/models/vaes/common.py b/fastvideo/v1/models/vaes/common.py index e480c97ce..369b9e470 100644 --- a/fastvideo/v1/models/vaes/common.py +++ b/fastvideo/v1/models/vaes/common.py @@ -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 diff --git a/fastvideo/v1/models/vaes/hunyuanvae.py b/fastvideo/v1/models/vaes/hunyuanvae.py index 8a066bf07..e49c3687e 100644 --- a/fastvideo/v1/models/vaes/hunyuanvae.py +++ b/fastvideo/v1/models/vaes/hunyuanvae.py @@ -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, diff --git a/fastvideo/v1/models/vaes/wanvae.py b/fastvideo/v1/models/vaes/wanvae.py index d50e47981..c3e5b9721 100644 --- a/fastvideo/v1/models/vaes/wanvae.py +++ b/fastvideo/v1/models/vaes/wanvae.py @@ -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 diff --git a/fastvideo/v1/pipelines/composed_pipeline_base.py b/fastvideo/v1/pipelines/composed_pipeline_base.py index beab21fae..a8b22b0a3 100644 --- a/fastvideo/v1/pipelines/composed_pipeline_base.py +++ b/fastvideo/v1/pipelines/composed_pipeline_base.py @@ -118,7 +118,7 @@ class ComposedPipelineBase(ABC): """ Initialize the pipeline. """ - pass + return def load_modules(self, fastvideo_args: FastVideoArgs) -> Dict[str, Any]: """ diff --git a/fastvideo/v1/pipelines/pipeline_batch_info.py b/fastvideo/v1/pipelines/pipeline_batch_info.py index 7c5e98e6c..5d34df5cc 100644 --- a/fastvideo/v1/pipelines/pipeline_batch_info.py +++ b/fastvideo/v1/pipelines/pipeline_batch_info.py @@ -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 diff --git a/fastvideo/v1/pipelines/stages/decoding.py b/fastvideo/v1/pipelines/stages/decoding.py index af700887b..663fd5df9 100644 --- a/fastvideo/v1/pipelines/stages/decoding.py +++ b/fastvideo/v1/pipelines/stages/decoding.py @@ -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__) diff --git a/fastvideo/v1/pipelines/stages/encoding.py b/fastvideo/v1/pipelines/stages/encoding.py index 61fdd9dd6..88718d316 100644 --- a/fastvideo/v1/pipelines/stages/encoding.py +++ b/fastvideo/v1/pipelines/stages/encoding.py @@ -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( diff --git a/fastvideo/v1/pipelines/stages/input_validation.py b/fastvideo/v1/pipelines/stages/input_validation.py index f4afad434..036c7a0c6 100644 --- a/fastvideo/v1/pipelines/stages/input_validation.py +++ b/fastvideo/v1/pipelines/stages/input_validation.py @@ -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... diff --git a/fastvideo/v1/pipelines/stages/latent_preparation.py b/fastvideo/v1/pipelines/stages/latent_preparation.py index b6ea3d713..9b7b8181b 100644 --- a/fastvideo/v1/pipelines/stages/latent_preparation.py +++ b/fastvideo/v1/pipelines/stages/latent_preparation.py @@ -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 diff --git a/fastvideo/v1/sample/v1_fastvideo_inference.py b/fastvideo/v1/sample/v1_fastvideo_inference.py index 508a0cc5c..dd33a5b91 100644 --- a/fastvideo/v1/sample/v1_fastvideo_inference.py +++ b/fastvideo/v1/sample/v1_fastvideo_inference.py @@ -1,3 +1,5 @@ +# type: ignore + # SPDX-License-Identifier: Apache-2.0 import os diff --git a/fastvideo/v1/utils.py b/fastvideo/v1/utils.py index 4ae2fd8d2..f229754f7 100644 --- a/fastvideo/v1/utils.py +++ b/fastvideo/v1/utils.py @@ -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