Phase 2

toolkit/models/v2/diffusion_models/wan.py

@@ -84,3 +84,54 @@ class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin):
         )
 
         return convert_wan_transformer_to_diffusers(dict(state_dict))
+
+    # inverse of diffusers' rename table, for the base/t2v/i2v variants (no
+    # vace/animate: their extra mappings collide in reverse). Order matters:
+    # longest/most-specific first, and the norm2/norm3 swap uses a placeholder.
+    _SAVE_RENAMES = [
+        ("condition_embedder.time_embedder.linear_1", "time_embedding.0"),
+        ("condition_embedder.time_embedder.linear_2", "time_embedding.2"),
+        ("condition_embedder.text_embedder.linear_1", "text_embedding.0"),
+        ("condition_embedder.text_embedder.linear_2", "text_embedding.2"),
+        ("condition_embedder.time_proj", "time_projection.1"),
+        ("condition_embedder.image_embedder.norm1", "img_emb.proj.0"),
+        ("condition_embedder.image_embedder.ff.net.0.proj", "img_emb.proj.1"),
+        ("condition_embedder.image_embedder.ff.net.2", "img_emb.proj.3"),
+        ("condition_embedder.image_embedder.norm2", "img_emb.proj.4"),
+        ("ffn.net.0.proj", "ffn.0"),
+        ("ffn.net.2", "ffn.2"),
+        (".norm_added_k.", ".norm_k_img."),
+        (".add_k_proj.", ".k_img."),
+        (".add_v_proj.", ".v_img."),
+        (".to_out.0.", ".o."),
+        (".to_q.", ".q."),
+        (".to_k.", ".k."),
+        (".to_v.", ".v."),
+        ("attn2", "cross_attn"),
+        ("attn1", "self_attn"),
+        # norm2 <-> norm3 swap back
+        ("norm3", "norm__placeholder"),
+        ("norm2", "norm3"),
+        ("norm__placeholder", "norm2"),
+        ("proj_out", "head.head"),
+    ]
+
+    @classmethod
+    def convert_state_dict_on_save(cls, state_dict):
+        """Diffusers layout back to the original/comfy wan key layout
+        (rename-only, so quantized weight/scale/marker keys ride along)."""
+        if any(".self_attn." in k or ".cross_attn." in k for k in state_dict):
+            return state_dict  # already original layout
+        new_sd = {}
+        for key, value in state_dict.items():
+            k = key
+            # scale_shift_table: blocks keep the name as `modulation`, the
+            # top-level one belongs to the output head
+            if k == "scale_shift_table":
+                k = "head.modulation"
+            elif k.endswith(".scale_shift_table"):
+                k = k[: -len("scale_shift_table")] + "modulation"
+            for src, dst in cls._SAVE_RENAMES:
+                k = k.replace(src, dst)
+            new_sd[k] = value
+        return new_sd