[review & refactor] Wan-Animate-2

src/diffusers/__init__.py

@@ -352,6 +352,7 @@
             "UVit2DModel",
             "VQModel",
             "WanAnimateTransformer3DModel",
+            "WanAnimate2Transformer3DModel",
             "WanTransformer3DModel",
             "WanVACETransformer3DModel",
             "ZImageControlNetModel",
@@ -844,6 +845,7 @@
             "VisualClozePipeline",
             "VQDiffusionPipeline",
             "WanAnimatePipeline",
+            "WanAnimate2Pipeline",
             "WanImageToVideoPipeline",
             "WanPipeline",
             "WanVACEPipeline",
@@ -1197,6 +1199,7 @@
             UNetSpatioTemporalConditionModel,
             UVit2DModel,
             VQModel,
+            WanAnimate2Transformer3DModel,
             WanAnimateTransformer3DModel,
             WanTransformer3DModel,
             WanVACETransformer3DModel,
@@ -1663,6 +1666,7 @@
             VisualClozeGenerationPipeline,
             VisualClozePipeline,
             VQDiffusionPipeline,
+            WanAnimate2Pipeline,
             WanAnimatePipeline,
             WanImageToVideoPipeline,
             WanPipeline,

src/diffusers/loaders/single_file_model.py

@@ -54,6 +54,7 @@
     convert_sana_transformer_to_diffusers,
     convert_sd3_transformer_checkpoint_to_diffusers,
     convert_stable_cascade_unet_single_file_to_diffusers,
+    convert_wan_animate_2_transformer_to_diffusers,
     convert_wan_transformer_to_diffusers,
     convert_wan_vae_to_diffusers,
     convert_z_image_controlnet_checkpoint_to_diffusers,
@@ -172,6 +173,10 @@
         "checkpoint_mapping_fn": convert_wan_transformer_to_diffusers,
         "default_subfolder": "transformer",
     },
+    "WanAnimate2Transformer3DModel": {
+        "checkpoint_mapping_fn": convert_wan_animate_2_transformer_to_diffusers,
+        "default_subfolder": "transformer",
+    },
     "AutoencoderKLWan": {
         "checkpoint_mapping_fn": convert_wan_vae_to_diffusers,
         "default_subfolder": "vae",

src/diffusers/loaders/single_file_utils.py

@@ -3289,6 +3289,38 @@ def reshape_bias_handler(key, state_dict):
     return converted_state_dict
 
 
+def convert_wan_animate_2_transformer_to_diffusers(checkpoint, **kwargs):
+    r"""
+    Converts the state dict of the Wan-Animate-2 transformer from the official checkpoint format to the diffusers
+    format.
+    """
+    attention_renames = {
+        ".q.": ".to_q.",
+        ".k.": ".to_k.",
+        ".v.": ".to_v.",
+        ".o.": ".to_out.0.",
+        ".k_img.": ".add_k_proj.",
+        ".v_img.": ".add_v_proj.",
+        ".norm_k_img.": ".norm_added_k.",
+    }
+
+    converted_state_dict = {}
+    for key in list(checkpoint.keys()):
+        new_key = key.replace("model.diffusion_model.", "")
+        # The official checkpoint wraps every transformer block in an in-context module the
+        # diffusers layout does not have: `blocks.N.block.X` -> `blocks.N.X`.
+        if new_key.startswith("blocks."):
+            new_key = new_key.replace(".block.", ".", 1)
+        if ".self_attn." in new_key or ".cross_attn." in new_key:
+            for old, new in attention_renames.items():
+                if old in new_key:
+                    new_key = new_key.replace(old, new)
+                    break
+        converted_state_dict[new_key] = checkpoint.pop(key)
+
+    return converted_state_dict
+
+
 def convert_wan_vae_to_diffusers(checkpoint, **kwargs):
     converted_state_dict = {}
 

src/diffusers/models/__init__.py

@@ -144,6 +144,7 @@
     _import_structure["transformers.transformer_temporal"] = ["TransformerTemporalModel"]
     _import_structure["transformers.transformer_wan"] = ["WanTransformer3DModel"]
     _import_structure["transformers.transformer_wan_animate"] = ["WanAnimateTransformer3DModel"]
+    _import_structure["transformers.transformer_wan_animate_2"] = ["WanAnimate2Transformer3DModel"]
     _import_structure["transformers.transformer_wan_vace"] = ["WanVACETransformer3DModel"]
     _import_structure["transformers.transformer_z_image"] = ["ZImageTransformer2DModel"]
     _import_structure["unets.unet_1d"] = ["UNet1DModel"]
@@ -284,6 +285,7 @@
             Transformer2DModel,
             TransformerTemporalModel,
             WanAnimateTransformer3DModel,
+            WanAnimate2Transformer3DModel,
             WanTransformer3DModel,
             WanVACETransformer3DModel,
             ZImageTransformer2DModel,

src/diffusers/models/transformers/__init__.py

@@ -64,5 +64,6 @@
     from .transformer_temporal import TransformerTemporalModel
     from .transformer_wan import WanTransformer3DModel
     from .transformer_wan_animate import WanAnimateTransformer3DModel
+    from .transformer_wan_animate_2 import WanAnimate2Transformer3DModel
     from .transformer_wan_vace import WanVACETransformer3DModel
     from .transformer_z_image import ZImageTransformer2DModel