Phase 2

toolkit/models/v2/diffusion_models/qwen_image.py

@@ -11,13 +11,36 @@ class QwenImageTransformer2DModel(
     aitk_subfolder = "transformer"
     aitk_config_repo = "Qwen/Qwen-Image"
 
+    aitk_comfy_repo = "Comfy-Org/Qwen-Image_ComfyUI"
+    # the comfy files use the diffusers key layout directly (no conversion);
+    # fp8mixed carries float8_e4m3fn comfy_quant markers that attach straight
+    # onto this class's modules
+    aitk_comfy_weight_names = {
+        "Qwen/Qwen-Image": [
+            "split_files/diffusion_models/qwen_image_fp8mixed.safetensors",
+            # raw fp8 cast (no markers, diffusers keys) — loads via from_single_file
+            "split_files/diffusion_models/qwen_image_fp8_e4m3fn.safetensors",
+            "split_files/diffusion_models/qwen_image_bf16.safetensors",
+        ],
+    }
+
     @classmethod
     def get_transformer_block_names(cls):
         return ["transformer_blocks"]
 
     @classmethod
     def _load_single_file(cls, file_path, dtype, config_path=None, subfolder=None):
-        # single-file checkpoints in the wild carry diffusers or original key
+        from safetensors import safe_open
+
+        with safe_open(file_path, framework="pt") as f:
+            has_markers = any(k.endswith(".comfy_quant") for k in f.keys())
+        if has_markers:
+            # 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
+            )
+        # other single-file checkpoints carry diffusers or original key
         # layouts; diffusers' single-file machinery owns that conversion
         model = cls.from_single_file(
             file_path,

toolkit/models/v2/diffusion_models/wan.py

@@ -8,6 +8,79 @@ class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin):
 
     aitk_subfolder = "transformer"
 
+    aitk_comfy_repo = "Comfy-Org/Wan_2.2_ComfyUI_Repackaged"
+    # comfy wan files use the original key layout; convert_state_dict_on_load
+    # renames them (pure substring renames, so the legacy scaled_fp8 weight /
+    # scale keys ride along with their modules). wan2.2 A14B checkpoints hold
+    # two DiTs, keyed by (repo, subfolder): transformer = high noise,
+    # transformer_2 = low noise. wan2.1 entries override the comfy repo.
+    aitk_comfy_weight_names = {
+        "Wan-AI/Wan2.2-TI2V-5B-Diffusers": [
+            "split_files/diffusion_models/wan2.2_ti2v_5B_fp16.safetensors",
+        ],
+        ("Wan-AI/Wan2.2-T2V-A14B-Diffusers", "transformer"): [
+            "split_files/diffusion_models/wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors",
+        ],
+        ("Wan-AI/Wan2.2-T2V-A14B-Diffusers", "transformer_2"): [
+            "split_files/diffusion_models/wan2.2_t2v_low_noise_14B_fp8_scaled.safetensors",
+        ],
+        ("Wan-AI/Wan2.2-I2V-A14B-Diffusers", "transformer"): [
+            "split_files/diffusion_models/wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors",
+        ],
+        ("Wan-AI/Wan2.2-I2V-A14B-Diffusers", "transformer_2"): [
+            "split_files/diffusion_models/wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors",
+        ],
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": {
+            "repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
+            "files": [
+                "split_files/diffusion_models/wan2.1_t2v_1.3B_bf16.safetensors",
+                "split_files/diffusion_models/wan2.1_t2v_1.3B_fp16.safetensors",
+            ],
+        },
+        "Wan-AI/Wan2.1-T2V-14B-Diffusers": {
+            "repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
+            "files": [
+                "split_files/diffusion_models/wan2.1_t2v_14B_fp8_scaled.safetensors",
+                "split_files/diffusion_models/wan2.1_t2v_14B_bf16.safetensors",
+                "split_files/diffusion_models/wan2.1_t2v_14B_fp16.safetensors",
+            ],
+        },
+        "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": {
+            "repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
+            "files": [
+                "split_files/diffusion_models/wan2.1_i2v_480p_14B_fp8_scaled.safetensors",
+                "split_files/diffusion_models/wan2.1_i2v_480p_14B_bf16.safetensors",
+                "split_files/diffusion_models/wan2.1_i2v_480p_14B_fp16.safetensors",
+            ],
+        },
+        "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": {
+            "repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
+            "files": [
+                "split_files/diffusion_models/wan2.1_i2v_720p_14B_fp8_scaled.safetensors",
+                "split_files/diffusion_models/wan2.1_i2v_720p_14B_bf16.safetensors",
+                "split_files/diffusion_models/wan2.1_i2v_720p_14B_fp16.safetensors",
+            ],
+        },
+    }
+
     @classmethod
     def get_transformer_block_names(cls):
         return ["blocks"]
+
+    @classmethod
+    def convert_state_dict_on_load(cls, state_dict):
+        # original/comfy wan keys -> diffusers layout via diffusers' own
+        # converter (rename-only for the base/i2v/vace variants)
+        is_original = any(
+            ".self_attn." in k
+            or ".cross_attn." in k
+            or k.startswith("model.diffusion_model.")
+            for k in state_dict
+        )
+        if not is_original:
+            return state_dict
+        from diffusers.loaders.single_file_utils import (
+            convert_wan_transformer_to_diffusers,
+        )
+
+        return convert_wan_transformer_to_diffusers(dict(state_dict))

toolkit/models/v2/diffusion_models/z_image.py

@@ -11,6 +11,17 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix
     # repo to pull the config from when loading a single-file checkpoint
     aitk_config_repo = "Tongyi-MAI/Z-Image-Turbo"
 
+    aitk_comfy_repo = "Comfy-Org/z_image_turbo"
+    # the int8_convrot file marks the FUSED attention.qkv modules; the load
+    # converter splits those entries exactly into to_q/to_k/to_v (row-sliced
+    # qdata + scales), so it attaches onto this split layout directly
+    aitk_comfy_weight_names = {
+        "Tongyi-MAI/Z-Image-Turbo": [
+            "split_files/diffusion_models/z_image_turbo_int8_convrot.safetensors",
+            "split_files/diffusion_models/z_image_turbo_bf16.safetensors",
+        ],
+    }
+
     @classmethod
     def get_transformer_block_names(cls):
         return ["layers"]
@@ -36,6 +47,22 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix
     @classmethod
     def convert_state_dict_on_load(cls, state_dict):
         """Convert a single-file Z-Image checkpoint to diffusers transformer keys."""
+        from toolkit.util.comfy_quant_import import split_fused_quantized_keys
+
+        state_dict = dict(state_dict)
+        # quantized fused qkv entries (comfy convrot/fp8/nvfp4 files) split
+        # exactly into the three projections: row-consecutive weights/scales
+        for marker_key in [
+            k for k in list(state_dict) if k.endswith(".attention.qkv.comfy_quant")
+        ]:
+            prefix = marker_key[: -len(".comfy_quant")]
+            base = prefix[: -len(".qkv")]
+            split_fused_quantized_keys(
+                state_dict,
+                prefix,
+                [f"{base}.to_q", f"{base}.to_k", f"{base}.to_v"],
+            )
+
         new_sd = {}
         for key, value in state_dict.items():
             k = key
@@ -47,9 +74,10 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix
                 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")
+            # module-prefix renames so weight/bias/scale/marker keys all map
+            k = k.replace(".attention.out.", ".attention.to_out.0.")
+            k = k.replace(".attention.q_norm.", ".attention.norm_q.")
+            k = k.replace(".attention.k_norm.", ".attention.norm_k.")
             if k.startswith("x_embedder."):
                 k = "all_x_embedder.2-1." + k[len("x_embedder.") :]
             elif k.startswith("final_layer."):
@@ -60,6 +88,20 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix
     @classmethod
     def convert_state_dict_on_save(cls, state_dict):
         """Convert a diffusers transformer state dict back to the single-file layout."""
+        from toolkit.util.comfy_quant_import import fuse_split_quantized_keys
+
+        state_dict = dict(state_dict)
+        # quantized split projections fuse back into the single-file qkv entry
+        for marker_key in [
+            k for k in list(state_dict) if k.endswith(".attention.to_q.comfy_quant")
+        ]:
+            base = marker_key[: -len(".to_q.comfy_quant")]
+            fuse_split_quantized_keys(
+                state_dict,
+                [f"{base}.to_q", f"{base}.to_k", f"{base}.to_v"],
+                f"{base}.qkv",
+            )
+
         new_sd = {}
         qkv_cache = {}
         for key, value in state_dict.items():
@@ -90,9 +132,9 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix
                     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")
+            k = k.replace(".attention.to_out.0.", ".attention.out.")
+            k = k.replace(".attention.norm_q.", ".attention.q_norm.")
+            k = k.replace(".attention.norm_k.", ".attention.k_norm.")
             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."):