update

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(

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

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(

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
+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)}")