support-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,27 @@ 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.
+    """
+    converted_state_dict = {}
+
+    # Strip model.diffusion_model prefix if present
+    keys = list(checkpoint.keys())
+    for k in keys:
+        if "model.diffusion_model." in k:
+            checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k)
+
+    # The official checkpoint already uses the same key format as the diffusers model
+    # (blocks.N.block.*), so no remapping is needed.
+    for key in list(checkpoint.keys()):
+        converted_state_dict[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