Update

fastvideo/v1/configs/models/dits/wanvideo.py

@@ -13,6 +13,73 @@ def is_blocks(n: str, m) -> bool:
 class WanVideoArchConfig(DiTArchConfig):
     _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
 
+    _videox_param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^patch_embedding\.(.*)$":
+            r"patch_embedding.proj.\1",
+            r"^text_embedding\.0\.(.*)$":
+            r"condition_embedder.text_embedder.fc_in.\1",
+            r"^text_embedding\.2\.(.*)$":
+            r"condition_embedder.text_embedder.fc_out.\1",
+            r"^time_embedding\.0\.(.*)$":
+            r"condition_embedder.time_embedder.mlp.fc_in.\1",
+            r"^time_embedding\.2\.(.*)$":
+            r"condition_embedder.time_embedder.mlp.fc_out.\1",
+            r"^time_projection\.1\.(.*)$":
+            r"condition_embedder.time_modulation.linear.\1",
+            r"^img_emb\.proj\.0\.(.*)$":
+            r"condition_embedder.image_embedder.norm1.\1",
+            r"^img_emb\.proj\.1\.(.*)$":
+            r"condition_embedder.image_embedder.ff.fc_in.\1",
+            r"^img_emb\.proj\.3\.(.*)$":
+            r"condition_embedder.image_embedder.ff.fc_out.\1",
+            r"^img_emb\.proj\.4\.(.*)$":
+            r"condition_embedder.image_embedder.norm2.\1",
+            r"^head\.modulation":
+            r"scale_shift_table",
+            r"^head\.head\.(.*)$":
+            r"proj_out.\1",
+            r"^blocks\.(\d+)\.self_attn\.q\.(.*)$":
+            r"blocks.\1.to_q.\2",
+            r"^blocks\.(\d+)\.self_attn\.k\.(.*)$":
+            r"blocks.\1.to_k.\2",
+            r"^blocks\.(\d+)\.self_attn\.v\.(.*)$":
+            r"blocks.\1.to_v.\2",
+            r"^blocks\.(\d+)\.self_attn\.o\.(.*)$":
+            r"blocks.\1.to_out.\2",
+            r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$":
+            r"blocks.\1.norm_q.\2",
+            r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$":
+            r"blocks.\1.norm_k.\2",
+            r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$":
+            r"blocks.\1.attn2.to_q.\2",
+            r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$":
+            r"blocks.\1.attn2.to_k.\2",
+            r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$":
+            r"blocks.\1.attn2.add_k_proj.\2",
+            r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$":
+            r"blocks.\1.attn2.to_v.\2",
+            r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$":
+            r"blocks.\1.attn2.add_v_proj.\2",
+            r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$":
+            r"blocks.\1.attn2.to_out.\2",
+            r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$":
+            r"blocks.\1.attn2.norm_q.\2",
+            r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$":
+            r"blocks.\1.attn2.norm_k.\2",
+            r"^blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$":
+            r"blocks.\1.attn2.norm_added_k.\2",
+            r"^blocks\.(\d+)\.ffn\.0\.(.*)$":
+            r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.2\.(.*)$":
+            r"blocks.\1.ffn.fc_out.\2",
+            r"^blocks\.(\d+)\.modulation":
+            r"blocks.\1.scale_shift_table",
+            r"^blocks\.(\d+)\.norm3\.(.*)$":
+            r"blocks.\1.self_attn_residual_norm.norm.\2",
+        }
+    )
+
     _param_names_mapping: dict = field(
         default_factory=lambda: {
             r"^patch_embedding\.(.*)$":

fastvideo/v1/models/dits/wanvideo.py

@@ -518,6 +518,7 @@ class WanTransformer3DModel(CachableDiT):
     _supported_attention_backends = WanVideoConfig(
     )._supported_attention_backends
     _param_names_mapping = WanVideoConfig()._param_names_mapping
+    _videox_param_names_mapping = WanVideoConfig()._videox_param_names_mapping
     _reverse_param_names_mapping = WanVideoConfig()._reverse_param_names_mapping
     _lora_param_names_mapping = WanVideoConfig()._lora_param_names_mapping
 

fastvideo/v1/models/loader/component_loader.py

@@ -379,10 +379,19 @@ class TransformerLoader(ComponentLoader):
         # Config from Diffusers supersedes fastvideo's model config
         dit_config = fastvideo_args.pipeline_config.dit_config
         dit_config.update_model_arch(config)
+        # dit_config.update_model_arch({
+        #     "num_attention_heads": 12,
+        #     "ffn_dim": 8960,
+        #     "in_channels": 36,
+        #     "num_layers": 30,
+        #     "image_dim": 1280,
+        #     "added_kv_proj_dim": 0
+        # })
 
         model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
 
         # Find all safetensors files
+        # model_path="/mnt/weka/home/hao.zhang/wei/FastVideo"
         safetensors_list = glob.glob(
             os.path.join(str(model_path), "*.safetensors"))
         if not safetensors_list:
@@ -458,6 +467,8 @@ class SchedulerLoader(ComponentLoader):
         assert class_name is not None, "Model config does not contain a _class_name attribute. Only diffusers format is supported."
 
         scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+        from fastvideo.v1.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchDiscreteScheduler
+        scheduler_cls = FlowMatchDiscreteScheduler
 
         scheduler = scheduler_cls(**config)
         if fastvideo_args.pipeline_config.flow_shift is not None:

scripts/huggingface/convert_to_hf.py

@@ -0,0 +1,88 @@
+from huggingface_hub import save_torch_state_dict, load_state_dict_from_file
+from safetensors.torch import save_file
+import torch
+import re
+from collections import OrderedDict
+
+_param_names_mapping: dict = {
+    r"^text_embedding\.0\.(.*)$":
+    r"condition_embedder.text_embedder.linear_1.\1",
+    r"^text_embedding\.2\.(.*)$":
+    r"condition_embedder.text_embedder.linear_2.\1",
+    r"^time_embedding\.0\.(.*)$":
+    r"condition_embedder.time_embedder.linear_1.\1",
+    r"^time_embedding\.2\.(.*)$":
+    r"condition_embedder.time_embedder.linear_2.\1",
+    r"^time_projection\.1\.(.*)$":
+    r"condition_embedder.time_proj.\1",
+    r"^img_emb\.proj\.0\.(.*)$":
+    r"condition_embedder.image_embedder.norm1.\1",
+    r"^img_emb\.proj\.1\.(.*)$":
+    r"condition_embedder.image_embedder.ff.net.0.proj.\1",
+    r"^img_emb\.proj\.3\.(.*)$":
+    r"condition_embedder.image_embedder.ff.net.2.\1",
+    r"^img_emb\.proj\.4\.(.*)$":
+    r"condition_embedder.image_embedder.norm2.\1",
+    r"^head\.modulation":
+    r"scale_shift_table",
+    r"^head\.head\.(.*)$":
+    r"proj_out.\1",
+    r"^blocks\.(\d+)\.self_attn\.q\.(.*)$":
+    r"blocks.\1.attn1.to_q.\2",
+    r"^blocks\.(\d+)\.self_attn\.k\.(.*)$":
+    r"blocks.\1.attn1.to_k.\2",
+    r"^blocks\.(\d+)\.self_attn\.v\.(.*)$":
+    r"blocks.\1.attn1.to_v.\2",
+    r"^blocks\.(\d+)\.self_attn\.o\.(.*)$":
+    r"blocks.\1.attn1.to_out.0.\2",
+    r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$":
+    r"blocks.\1.attn1.norm_q.\2",
+    r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$":
+    r"blocks.\1.attn1.norm_k.\2",
+    r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$":
+    r"blocks.\1.attn2.to_q.\2",
+    r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$":
+    r"blocks.\1.attn2.to_k.\2",
+    r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$":
+    r"blocks.\1.attn2.add_k_proj.\2",
+    r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$":
+    r"blocks.\1.attn2.to_v.\2",
+    r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$":
+    r"blocks.\1.attn2.add_v_proj.\2",
+    r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$":
+    r"blocks.\1.attn2.to_out.0.\2",
+    r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$":
+    r"blocks.\1.attn2.norm_q.\2",
+    r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$":
+    r"blocks.\1.attn2.norm_k.\2",
+    r"^blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$":
+    r"blocks.\1.attn2.norm_added_k.\2",
+    r"^blocks\.(\d+)\.ffn\.0\.(.*)$":
+    r"blocks.\1.ffn.net.0.proj.\2",
+    r"^blocks\.(\d+)\.ffn\.2\.(.*)$":
+    r"blocks.\1.ffn.net.2.\2",
+    r"^blocks\.(\d+)\.modulation":
+    r"blocks.\1.scale_shift_table",
+    r"^blocks\.(\d+)\.norm3\.(.*)$":
+    r"blocks.\1.norm2.\2",
+}
+
+state_dict = load_state_dict_from_file("/mnt/weka/home/hao.zhang/wei/FastVideo/diffusion_pytorch_model.safetensors")
+new_state_dict = OrderedDict()
+for k, v in state_dict.items():
+    new_key = k
+    for pattern, replacement in _param_names_mapping.items():
+        if re.match(pattern, k):
+            new_key = re.sub(pattern, replacement, k)
+            break  # Stop at the first match
+    new_state_dict[new_key] = v
+    if "norm_added_k" in new_key:
+        dummy_key = new_key.replace("norm_added_k", "norm_added_q")
+        dummy_value = torch.zeros_like(v)
+        new_state_dict[dummy_key] = dummy_value
+del state_dict
+
+save_torch_state_dict(
+    new_state_dict,
+    "/mnt/weka/home/hao.zhang/wei/FastVideo/data/Wan2.1-Fun-1.3B-InP-Diffusers/transformer"
+)

scripts/huggingface/upload_hf.py

@@ -3,7 +3,7 @@ from huggingface_hub import HfApi
 api = HfApi()
 
 api.upload_folder(
-    folder_path="data/Black-Myth-Taylor-Src",
-    repo_id="FastVideo/Image-Vid-Finetune-Src",
-    repo_type="dataset",
+    folder_path="/mnt/weka/home/hao.zhang/wei/FastVideo/data/Wan2.1-Fun-1.3B-InP-Diffusers",
+    repo_id="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers",
+    repo_type="model",
 )

scripts/inference/v1_inference_wan_i2v.sh

@@ -1,8 +1,8 @@
 #!/bin/bash
-
+export HOME=/mnt/weka/home/hao.zhang/wei
 num_gpus=2
 export FASTVIDEO_ATTENTION_BACKEND=
-export MODEL_BASE=Wan-AI/Wan2.1-I2V-14B-480P-Diffusers
+export MODEL_BASE=$HOME/FastVideo/data/Wan2.1-Fun-1.3B-InP-Diffusers
 # export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
 # Note that the tp_size and sp_size should be the same and equal to the number
 # of GPUs. They are used for different parallel groups. sp_size is used for
@@ -14,13 +14,13 @@ fastvideo generate \
     --num-gpus $num_gpus \
     --height 480 \
     --width 832 \
-    --num-frames 77 \
+    --num-frames 29 \
     --num-inference-steps 40 \
     --fps 16 \
     --flow-shift 3.0 \
     --guidance-scale 5.0 \
-    --image-path "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" \
-    --prompt "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot." \
+    --image-path "assets/dog.png" \
+    --prompt "A brown dog is shaking its head while sitting on a light-colored sofa in a cozy room. Behind the dog, there is a framed painting on the shelf, surrounded by pink flowers. The soft, warm lighting in the room creates a comfortable atmosphere." \
     --negative-prompt "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" \
     --seed 1024 \
-    --output-path outputs_i2v/
+    --output-path outputs_videox_i2v/