conversion script

Author: a-r-r-o-w

scripts/convert_wan_to_diffusers.py

@@ -438,7 +438,7 @@
             "VersatileDiffusionTextToImagePipeline",
             "VideoToVideoSDPipeline",
             "VQDiffusionPipeline",
-            "WanI2VPipeline",
+            "WanImageToVideoPipeline",
             "WanPipeline",
             "WuerstchenCombinedPipeline",
             "WuerstchenDecoderPipeline",
@@ -939,7 +939,7 @@
             VersatileDiffusionTextToImagePipeline,
             VideoToVideoSDPipeline,
             VQDiffusionPipeline,
-            WanI2VPipeline,
+            WanImageToVideoPipeline,
             WanPipeline,
             WuerstchenCombinedPipeline,
             WuerstchenDecoderPipeline,

src/diffusers/__init__.py

@@ -127,7 +127,7 @@ def __init__(
         time_freq_dim: int,
         time_proj_dim: int,
         text_embed_dim: int,
-        image_embedding_dim: Optional[int] = None,
+        image_embed_dim: Optional[int] = None,
     ):
         super().__init__()
 
@@ -138,8 +138,8 @@ def __init__(
         self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim, act_fn="gelu_tanh")
 
         self.image_embedder = None
-        if image_embedding_dim is not None:
-            self.image_embedder = WanImageEmbedding(image_embedding_dim, dim)
+        if image_embed_dim is not None:
+            self.image_embedder = WanImageEmbedding(image_embed_dim, dim)
 
     def forward(
         self,
@@ -348,7 +348,7 @@ def __init__(
         cross_attn_norm: bool = True,
         qk_norm: Optional[str] = "rms_norm_across_heads",
         eps: float = 1e-6,
-        image_embedding_dim: Optional[int] = None,
+        image_dim: Optional[int] = None,
         added_kv_proj_dim: Optional[int] = None,
         rope_max_seq_len: int = 1024,
     ) -> None:
@@ -368,7 +368,7 @@ def __init__(
             time_freq_dim=freq_dim,
             time_proj_dim=inner_dim * 6,
             text_embed_dim=text_dim,
-            image_embedding_dim=image_embedding_dim,
+            image_embed_dim=image_dim,
         )
 
         # 3. Transformer blocks

src/diffusers/models/transformers/transformer_wan.py

@@ -346,7 +346,7 @@
         "WuerstchenDecoderPipeline",
         "WuerstchenPriorPipeline",
     ]
-    _import_structure["wan"] = ["WanPipeline", "WanI2VPipeline"]
+    _import_structure["wan"] = ["WanPipeline", "WanImageToVideoPipeline"]
 try:
     if not is_onnx_available():
         raise OptionalDependencyNotAvailable()
@@ -689,7 +689,7 @@
             UniDiffuserPipeline,
             UniDiffuserTextDecoder,
         )
-        from .wan import WanI2VPipeline, WanPipeline
+        from .wan import WanImageToVideoPipeline, WanPipeline
         from .wuerstchen import (
             WuerstchenCombinedPipeline,
             WuerstchenDecoderPipeline,

src/diffusers/pipelines/__init__.py

@@ -23,7 +23,7 @@
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
 else:
     _import_structure["pipeline_wan"] = ["WanPipeline"]
-    _import_structure["pipeline_wan_i2v"] = ["WanI2VPipeline"]
+    _import_structure["pipeline_wan_i2v"] = ["WanImageToVideoPipeline"]
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     try:
@@ -34,7 +34,7 @@
         from ...utils.dummy_torch_and_transformers_objects import *
     else:
         from .pipeline_wan import WanPipeline
-        from .pipeline_wan_i2v import WanI2VPipeline
+        from .pipeline_wan_i2v import WanImageToVideoPipeline
 
 else:
     import sys

src/diffusers/pipelines/wan/__init__.py

@@ -46,7 +46,7 @@
     Examples:
         ```python
         >>> import torch
-        >>> from diffusers import WanI2VPipeline, WanTransformer3DModel
+        >>> from diffusers import WanImageToVideoPipeline, WanTransformer3DModel
         >>> from transformers import CLIPVisionModel, CLIPImageProcessor, UMT5EncoderModel
         >>> from diffusers.utils import load_image, export_to_video
 
@@ -56,7 +56,7 @@
         >>> text_encoder = UMT5EncoderModel.from_pretrained(model_id, subfolder="text_encoder")
         >>> transformer_i2v = WanTransformer3DModel.from_pretrained(model_id, subfolder="transformer_i2v_720p")
         >>> image_processor = CLIPImageProcessor.from_pretrained(model_id, subfolder="image_processor")
-        >>> pipe = WanI2VPipeline.from_pretrained(
+        >>> pipe = WanImageToVideoPipeline.from_pretrained(
         ...     model_id,
         ...     transformer=transformer_i2v,
         ...     text_encoder=text_encoder,
@@ -125,7 +125,7 @@ def retrieve_latents(
         raise AttributeError("Could not access latents of provided encoder_output")
 
 
-class WanI2VPipeline(DiffusionPipeline):
+class WanImageToVideoPipeline(DiffusionPipeline):
     r"""
     Pipeline for image-to-video generation using Wan.
 

src/diffusers/pipelines/wan/pipeline_wan_i2v.py

@@ -2552,7 +2552,7 @@ def from_pretrained(cls, *args, **kwargs):
         requires_backends(cls, ["torch", "transformers"])
 
 
-class WanI2VPipeline(metaclass=DummyObject):
+class WanImageToVideoPipeline(metaclass=DummyObject):
     _backends = ["torch", "transformers"]
 
     def __init__(self, *args, **kwargs):

src/diffusers/utils/dummy_torch_and_transformers_objects.py

@@ -19,7 +19,7 @@
 from PIL import Image
 from transformers import AutoTokenizer, T5EncoderModel, CLIPVisionConfig, CLIPVisionModel, CLIPImageProcessor
 
-from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, WanI2VPipeline, WanTransformer3DModel
+from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, WanImageToVideoPipeline, WanTransformer3DModel
 from diffusers.utils.testing_utils import enable_full_determinism
 
 from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
@@ -30,7 +30,7 @@
 
 
 class WanImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
-    pipeline_class = WanI2VPipeline
+    pipeline_class = WanImageToVideoPipeline
     params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs", "height", "width"}
     batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
     image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
@@ -78,7 +78,7 @@ def get_dummy_components(self):
             cross_attn_norm=True,
             qk_norm="rms_norm_across_heads",
             rope_max_seq_len=32,
-            image_embedding_dim=4,
+            image_dim=4,
         )
 
         torch.manual_seed(0)