From e8ce4290ffdfc8ec792a71bd23cf42784624790c Mon Sep 17 00:00:00 2001 From: Peiyuan Zhang Date: Wed, 11 Dec 2024 20:27:18 +0000 Subject: [PATCH] update --- fastvideo/distill.py | 6 +- fastvideo/fsdp_util.py | 8 +-- fastvideo/models/hunyuan/modules/models.py | 5 +- fastvideo/utils/load_model.py | 70 +++++++++++++++++----- 4 files changed, 66 insertions(+), 23 deletions(-) diff --git a/fastvideo/distill.py b/fastvideo/distill.py index 19e9464da..b73a37ba8 100644 --- a/fastvideo/distill.py +++ b/fastvideo/distill.py @@ -30,6 +30,7 @@ from fastvideo.fsdp_util import get_dit_fsdp_kwargs, apply_fsdp_checkpointing from diffusers import ( FlowMatchEulerDiscreteScheduler, ) +from fastvideo.utils.load_model import get_no_split_modules, load_transformer from fastvideo.distill.solver import EulerSolver, extract_into_tensor from copy import deepcopy from diffusers.optimization import get_scheduler @@ -315,7 +316,7 @@ def main(args): main_print(f"--> loading model from {args.pretrained_model_name_or_path}") - transformer = load_transformer(args) + transformer = load_transformer(args.model_type,args.dit_model_name_or_path, args.pretrained_model_name_or_path) # keep the master weight to float32 @@ -342,6 +343,7 @@ def main(args): f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}" ) fsdp_kwargs = get_dit_fsdp_kwargs( + transformer, args.fsdp_sharding_startegy, args.use_lora, args.use_cpu_offload, @@ -352,7 +354,7 @@ def main(args): transformer.config.lora_rank = args.lora_rank transformer.config.lora_alpha = args.lora_alpha transformer.config.lora_target_modules = ["to_k", "to_q", "to_v", "to_out.0"] - transformer._no_split_modules = get_no_split_modules(args) + transformer._no_split_modules = get_no_split_modules(transformer) fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](transformer) transformer = FSDP( diff --git a/fastvideo/fsdp_util.py b/fastvideo/fsdp_util.py index 96a459767..c1662d85f 100644 --- a/fastvideo/fsdp_util.py +++ b/fastvideo/fsdp_util.py @@ -17,6 +17,7 @@ from torch.distributed.fsdp import ( # ShardedStateDictConfig, # un-flattened param but shards, usable by other parallel schemes. ) +from fastvideo.utils.load_model import get_no_split_modules from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformerBlock from functools import partial @@ -80,16 +81,15 @@ def get_mixed_precision(master_weight_type="fp32"): def get_dit_fsdp_kwargs( - sharding_strategy, use_lora=False, cpu_offload=False, master_weight_type="fp32" + transformer, sharding_strategy, use_lora=False, cpu_offload=False, master_weight_type="fp32" ): + no_split_modules = get_no_split_modules(transformer) if use_lora: auto_wrap_policy = fsdp_auto_wrap_policy else: auto_wrap_policy = functools.partial( transformer_auto_wrap_policy, - transformer_layer_cls={ - MochiTransformerBlock, - }, + transformer_layer_cls=no_split_modules, ) # we use float32 for fsdp but autocast during training diff --git a/fastvideo/models/hunyuan/modules/models.py b/fastvideo/models/hunyuan/modules/models.py index 9c3ef7dbe..6a2eb9c70 100644 --- a/fastvideo/models/hunyuan/modules/models.py +++ b/fastvideo/models/hunyuan/modules/models.py @@ -433,7 +433,6 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): @register_to_config def __init__( self, - args: Any, patch_size: list = [1, 2, 2], in_channels: int = 4, # Should be VAE.config.latent_channels. out_channels: int = None, @@ -452,6 +451,8 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): use_attention_mask: bool = True, dtype: Optional[torch.dtype] = None, device: Optional[torch.device] = None, + text_states_dim: int = 4096, + text_states_dim_2: int = 768, ): factory_kwargs = {"device": device, "dtype": dtype} super().__init__() @@ -468,8 +469,6 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): self.use_attention_mask = use_attention_mask self.text_projection = text_projection - self.text_states_dim = args.text_states_dim - self.text_states_dim_2 = args.text_states_dim_2 if hidden_size % heads_num != 0: raise ValueError( diff --git a/fastvideo/utils/load_model.py b/fastvideo/utils/load_model.py index 4873892ab..751601f98 100644 --- a/fastvideo/utils/load_model.py +++ b/fastvideo/utils/load_model.py @@ -1,28 +1,70 @@ import torch -from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel +from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel, MochiTransformerBlock +from fastvideo.models.hunyuan.modules.models import HYVideoDiffusionTransformer, MMDoubleStreamBlock, MMSingleStreamBlock -def load_transformer(args): - if args.model_type == "mochi": - if args.dit_model_name_or_path: +hunyuan_config = { + "mm_double_blocks_depth": 20, + "mm_single_blocks_depth": 40, + "rope_dim_list": [16, 56, 56], + "hidden_size": 3072, + "heads_num": 24, + "mlp_width_ratio": 4, + "guidance_embed": True, +} + + +def load_hunyuan_state_dict(model, dit_model_name_or_path): + load_key = "module" + model_path = dit_model_name_or_path + bare_model = "unknown" + + state_dict = torch.load(model_path, map_location=lambda storage, loc: storage) + + if bare_model == "unknown" and ("ema" in state_dict or "module" in state_dict): + bare_model = False + if bare_model is False: + if load_key in state_dict: + state_dict = state_dict[load_key] + else: + raise KeyError( + f"Missing key: `{load_key}` in the checkpoint: {model_path}. The keys in the checkpoint " + f"are: {list(state_dict.keys())}." + ) + model.load_state_dict(state_dict, strict=True) + return model + +def load_transformer(model_type,dit_model_name_or_path, pretrained_model_name_or_path ): + if model_type == "mochi": + if dit_model_name_or_path: transformer = MochiTransformer3DModel.from_pretrained( - args.dit_model_name_or_path, + dit_model_name_or_path, torch_dtype=torch.float32, # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32, ) else: transformer = MochiTransformer3DModel.from_pretrained( - args.pretrained_model_name_or_path, + pretrained_model_name_or_path, subfolder="transformer", torch_dtype=torch.float32, # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32, ) - elif args.model_type == "hunyuan" - pass - + elif model_type == "hunyuan": + transformer = HYVideoDiffusionTransformer( + in_channels=16, + out_channels=16, + **hunyuan_config, + dtype=torch.bfloat16, + ) + transformer = load_hunyuan_state_dict(transformer, dit_model_name_or_path) + else: + raise ValueError(f"Unsupported model type: {model_type}") return transformer -def get_no_split_modules(args): - if args.model_type == "mochi": - return ["MochiTransformerBlock"] - elif args.model_type == "hunyuan": - pass \ No newline at end of file +def get_no_split_modules(transformer): + # if of type MochiTransformer3DModel + if isinstance(transformer, MochiTransformer3DModel): + return [MochiTransformerBlock] + elif isinstance(transformer, HYVideoDiffusionTransformer): + return [MMDoubleStreamBlock, MMSingleStreamBlock] + else: + raise ValueError(f"Unsupported transformer type: {type(transformer)}") \ No newline at end of file