Models v2 - phase 0

toolkit/models/v2/diffusion_models/z_image.py

@@ -0,0 +1,101 @@
+import torch
+from diffusers.models.transformers import (
+    ZImageTransformer2DModel as DiffusersZImageTransformer2DModel,
+)
+
+from .._mixin import OstrisModelMixin
+
+
+class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMixin):
+    aitk_subfolder = "transformer"
+    # repo to pull the config from when loading a single-file checkpoint
+    aitk_config_repo = "Tongyi-MAI/Z-Image-Turbo"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["layers"]
+
+    @classmethod
+    def get_quantization_exclude_modules(cls):
+        # sensitive modules kept in full precision (fnmatch patterns on module
+        # names within ZImageTransformer2DModel):
+        #   t_embedder*      - timestep embedder; feeds every block's
+        #                      adaLN_modulation and the final layers
+        #   cap_embedder*    - caption feature -> model width projection
+        #   all_x_embedder*  - patchified latent input projections
+        #   all_final_layer* - final adaLN-modulated output projections
+        #   siglip_embedder* - siglip feature projection (edit models only)
+        return [
+            "t_embedder*",
+            "cap_embedder*",
+            "all_x_embedder*",
+            "all_final_layer*",
+            "siglip_embedder*",
+        ]
+
+    @classmethod
+    def convert_state_dict_on_load(cls, state_dict):
+        """Convert a single-file Z-Image checkpoint to diffusers transformer keys."""
+        new_sd = {}
+        for key, value in state_dict.items():
+            k = key
+            if k.endswith(".attention.qkv.weight"):
+                # the single file fuses q,k,v into one tensor (in that order); diffusers keeps them split
+                prefix = k[: -len(".attention.qkv.weight")]
+                q, k_proj, v = torch.chunk(value, 3, dim=0)
+                new_sd[prefix + ".attention.to_q.weight"] = q
+                new_sd[prefix + ".attention.to_k.weight"] = k_proj
+                new_sd[prefix + ".attention.to_v.weight"] = v
+                continue
+            k = k.replace(".attention.out.weight", ".attention.to_out.0.weight")
+            k = k.replace(".attention.q_norm.weight", ".attention.norm_q.weight")
+            k = k.replace(".attention.k_norm.weight", ".attention.norm_k.weight")
+            if k.startswith("x_embedder."):
+                k = "all_x_embedder.2-1." + k[len("x_embedder.") :]
+            elif k.startswith("final_layer."):
+                k = "all_final_layer.2-1." + k[len("final_layer.") :]
+            new_sd[k] = value
+        return new_sd
+
+    @classmethod
+    def convert_state_dict_on_save(cls, state_dict):
+        """Convert a diffusers transformer state dict back to the single-file layout."""
+        new_sd = {}
+        qkv_cache = {}
+        for key, value in state_dict.items():
+            k = key
+            matched = False
+            for suffix in (
+                ".attention.to_q.weight",
+                ".attention.to_k.weight",
+                ".attention.to_v.weight",
+            ):
+                if k.endswith(suffix):
+                    prefix = k[: -len(suffix)]
+                    cache = qkv_cache.setdefault(prefix, {})
+                    cache[suffix] = value
+                    if len(cache) == 3:
+                        # the single file expects q,k,v fused in that order
+                        qkv = torch.cat(
+                            [
+                                cache[".attention.to_q.weight"],
+                                cache[".attention.to_k.weight"],
+                                cache[".attention.to_v.weight"],
+                            ],
+                            dim=0,
+                        )
+                        new_sd[prefix + ".attention.qkv.weight"] = qkv
+                        del qkv_cache[prefix]
+                    matched = True
+                    break
+            if matched:
+                continue
+            k = k.replace(".attention.to_out.0.weight", ".attention.out.weight")
+            k = k.replace(".attention.norm_q.weight", ".attention.q_norm.weight")
+            k = k.replace(".attention.norm_k.weight", ".attention.k_norm.weight")
+            if k.startswith("all_x_embedder.2-1."):
+                k = "x_embedder." + k[len("all_x_embedder.2-1.") :]
+            elif k.startswith("all_final_layer.2-1."):
+                k = "final_layer." + k[len("all_final_layer.2-1.") :]
+            new_sd[k] = value
+        return new_sd