Phase 2

toolkit/models/v2/diffusion_models/qwen_image.py

@@ -29,7 +29,9 @@ class QwenImageTransformer2DModel(
         return ["transformer_blocks"]
 
     @classmethod
-    def _load_single_file(cls, file_path, dtype, config_path=None, subfolder=None):
+    def _load_single_file(
+        cls, file_path, dtype, config_path=None, config=None, subfolder=None
+    ):
         from safetensors import safe_open
 
         with safe_open(file_path, framework="pt") as f:
@@ -38,7 +40,11 @@ class QwenImageTransformer2DModel(
             # comfy prequantized checkpoint (diffusers key layout): the mixin
             # path attaches the quantized layers
             return super()._load_single_file(
-                file_path, dtype, config_path=config_path, subfolder=subfolder
+                file_path,
+                dtype,
+                config_path=config_path,
+                config=config,
+                subfolder=subfolder,
             )
         # other single-file checkpoints carry diffusers or original key
         # layouts; diffusers' single-file machinery owns that conversion

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