[WIP] [Alibaba Wan Team] Add Wan 2.1 support

src/diffusers/__init__.py

@@ -90,6 +90,7 @@
             "AuraFlowTransformer2DModel",
             "AutoencoderDC",
             "AutoencoderKL",
+            "AutoencoderKLWan",
             "AutoencoderKLAllegro",
             "AutoencoderKLCogVideoX",
             "AutoencoderKLHunyuanVideo",
@@ -148,6 +149,7 @@
             "UNetSpatioTemporalConditionModel",
             "UVit2DModel",
             "VQModel",
+            "WanTransformer3DModel",
         ]
     )
     _import_structure["optimization"] = [
@@ -439,6 +441,8 @@
             "WuerstchenCombinedPipeline",
             "WuerstchenDecoderPipeline",
             "WuerstchenPriorPipeline",
+            "WanPipeline",
+            "WanI2VPipeline",
         ]
     )
 
@@ -610,6 +614,7 @@
             AuraFlowTransformer2DModel,
             AutoencoderDC,
             AutoencoderKL,
+            AutoencoderKLWan,
             AutoencoderKLAllegro,
             AutoencoderKLCogVideoX,
             AutoencoderKLHunyuanVideo,
@@ -667,6 +672,7 @@
             UNetSpatioTemporalConditionModel,
             UVit2DModel,
             VQModel,
+            WanTransformer3DModel,
         )
         from .optimization import (
             get_constant_schedule,
@@ -936,6 +942,8 @@
             WuerstchenCombinedPipeline,
             WuerstchenDecoderPipeline,
             WuerstchenPriorPipeline,
+            WanPipeline,
+            WanI2VPipeline
         )
 
     try:

src/diffusers/models/__init__.py

@@ -29,6 +29,7 @@
     _import_structure["autoencoders.autoencoder_asym_kl"] = ["AsymmetricAutoencoderKL"]
     _import_structure["autoencoders.autoencoder_dc"] = ["AutoencoderDC"]
     _import_structure["autoencoders.autoencoder_kl"] = ["AutoencoderKL"]
+    _import_structure["autoencoders.autoencoder_kl_wan"] = ["AutoencoderKLWan"]
     _import_structure["autoencoders.autoencoder_kl_allegro"] = ["AutoencoderKLAllegro"]
     _import_structure["autoencoders.autoencoder_kl_cogvideox"] = ["AutoencoderKLCogVideoX"]
     _import_structure["autoencoders.autoencoder_kl_hunyuan_video"] = ["AutoencoderKLHunyuanVideo"]
@@ -89,6 +90,7 @@
     _import_structure["unets.unet_spatio_temporal_condition"] = ["UNetSpatioTemporalConditionModel"]
     _import_structure["unets.unet_stable_cascade"] = ["StableCascadeUNet"]
     _import_structure["unets.uvit_2d"] = ["UVit2DModel"]
+    _import_structure["transformers.transformer_wan"] = ["WanTransformer3DModel"]
 
 if is_flax_available():
     _import_structure["controlnets.controlnet_flax"] = ["FlaxControlNetModel"]
@@ -158,6 +160,7 @@
             T5FilmDecoder,
             Transformer2DModel,
             TransformerTemporalModel,
+            WanTransformer3DModel,
         )
         from .unets import (
             I2VGenXLUNet,

src/diffusers/models/attention_processor.py

@@ -278,6 +278,10 @@ def __init__(
             elif qk_norm == "rms_norm":
                 self.norm_added_q = RMSNorm(dim_head, eps=eps)
                 self.norm_added_k = RMSNorm(dim_head, eps=eps)
+            elif qk_norm == "rms_norm_across_heads":
+                # Wanx applies qk norm across all heads
+                self.norm_added_q = RMSNorm(dim_head * heads, eps=eps)
+                self.norm_added_k = RMSNorm(dim_head * kv_heads, eps=eps)
             else:
                 raise ValueError(
                     f"unknown qk_norm: {qk_norm}. Should be one of `None,'layer_norm','fp32_layer_norm','rms_norm'`"

src/diffusers/models/autoencoders/__init__.py

@@ -1,6 +1,7 @@
 from .autoencoder_asym_kl import AsymmetricAutoencoderKL
 from .autoencoder_dc import AutoencoderDC
 from .autoencoder_kl import AutoencoderKL
+from .autoencoder_kl_wan import AutoencoderKLWan
 from .autoencoder_kl_allegro import AutoencoderKLAllegro
 from .autoencoder_kl_cogvideox import AutoencoderKLCogVideoX
 from .autoencoder_kl_hunyuan_video import AutoencoderKLHunyuanVideo