Upload app.py with huggingface_hub

app.py

@@ -11,18 +11,115 @@ TURBO_REPO = "krea/Krea-2-Turbo"
 HD_VAE_REPO = "wikeeyang/Krea2-Turbo-HD-V1"
 MAX_SEED = 2**31 - 1
 
+
+def _remap_vae_keys(state_dict):
+    """Remap ComfyUI-format VAE state dict keys to diffusers AutoencoderKLQwenImage keys.
+
+    The ComfyUI checkpoint uses a flat sequential naming convention (residual.0,
+    residual.2, etc.) while diffusers uses semantic names (norm1, conv1, norm2,
+    conv2). This function translates between the two.
+    """
+    # Mapping from ComfyUI decoder upsamples index to diffusers up_blocks path.
+    # Derived from the AutoencoderKLQwenImage architecture with dim_mult=[1,2,4,4]
+    # and num_res_blocks=2: each level has 2 resnets, with upsamplers between levels.
+    up_map = {
+        0: "up_blocks.0.resnets.0",
+        1: "up_blocks.0.resnets.1",
+        2: "up_blocks.0.resnets.2",
+        3: "up_blocks.0.upsamplers.0",
+        4: "up_blocks.1.resnets.0",
+        5: "up_blocks.1.resnets.1",
+        6: "up_blocks.1.resnets.2",
+        7: "up_blocks.1.upsamplers.0",
+        8: "up_blocks.2.resnets.0",
+        9: "up_blocks.2.resnets.1",
+        10: "up_blocks.2.resnets.2",
+        11: "up_blocks.2.upsamplers.0",
+        12: "up_blocks.3.resnets.0",
+        13: "up_blocks.3.resnets.1",
+        14: "up_blocks.3.resnets.2",
+    }
+
+    def _fix_resnet(rest):
+        """Map sequential residual indices to semantic names."""
+        rest = rest.replace("residual.0.", "norm1.")
+        rest = rest.replace("residual.2.", "conv1.")
+        rest = rest.replace("residual.3.", "norm2.")
+        rest = rest.replace("residual.6.", "conv2.")
+        rest = rest.replace("shortcut.", "conv_shortcut.")
+        return rest
+
+    def _fix_middle(key, side):
+        """Map encoder/decoder middle.X to mid_block structure."""
+        parts = key.split(".")
+        idx = int(parts[2])
+        rest = ".".join(parts[3:])
+        if idx == 0:  # first resnet
+            return f"{side}.mid_block.resnets.0.{_fix_resnet(rest)}"
+        elif idx == 1:  # attention
+            return f"{side}.mid_block.attentions.0.{rest}"
+        elif idx == 2:  # second resnet
+            return f"{side}.mid_block.resnets.1.{_fix_resnet(rest)}"
+
+    new_state = {}
+    for key, val in state_dict.items():
+        new_key = key
+
+        # Top-level convs: quant_conv and post_quant_conv
+        if key.startswith("conv1."):
+            new_key = "quant_conv." + key[len("conv1."):]
+        elif key.startswith("conv2."):
+            new_key = "post_quant_conv." + key[len("conv2."):]
+
+        # Encoder mapping
+        elif key.startswith("encoder.conv1."):
+            new_key = "encoder.conv_in." + key[len("encoder.conv1."):]
+        elif key.startswith("encoder.head.0."):
+            new_key = "encoder.norm_out." + key[len("encoder.head.0."):]
+        elif key.startswith("encoder.head.2."):
+            new_key = "encoder.conv_out." + key[len("encoder.head.2."):]
+        elif key.startswith("encoder.downsamples."):
+            parts = key.split(".")
+            idx = int(parts[2])
+            rest = ".".join(parts[3:])
+            rest = _fix_resnet(rest)
+            new_key = f"encoder.down_blocks.{idx}.{rest}"
+        elif key.startswith("encoder.middle."):
+            new_key = _fix_middle(key, "encoder")
+
+        # Decoder mapping
+        elif key.startswith("decoder.conv1."):
+            new_key = "decoder.conv_in." + key[len("decoder.conv1."):]
+        elif key.startswith("decoder.head.0."):
+            new_key = "decoder.norm_out." + key[len("decoder.head.0."):]
+        elif key.startswith("decoder.head.2."):
+            new_key = "decoder.conv_out." + key[len("decoder.head.2."):]
+        elif key.startswith("decoder.upsamples."):
+            parts = key.split(".")
+            idx = int(parts[2])
+            rest = ".".join(parts[3:])
+            rest = _fix_resnet(rest)
+            new_key = f"decoder.{up_map[idx]}.{rest}"
+        elif key.startswith("decoder.middle."):
+            new_key = _fix_middle(key, "decoder")
+
+        new_state[new_key] = val
+
+    return new_state
+
+
 # Load the Krea-2-Turbo pipeline (full diffusers format) then swap in the
 # HD-optimized VAE from wikeeyang/Krea2-Turbo-HD-V1 for enhanced detail,
-# clarity, and contrast. The HD VAE is a single safetensors checkpoint of the
-# same AutoencoderKLQwenImage architecture, so we load the base pipeline's VAE
-# (which has the proper config) and overwrite its weights.
+# clarity, and contrast. The HD VAE is a ComfyUI-format single safetensors
+# checkpoint; we remap its keys to the diffusers AutoencoderKLQwenImage layout.
 pipe = Krea2Pipeline.from_pretrained(TURBO_REPO, torch_dtype=torch.bfloat16)
 
-# Load the HD VAE weights and replace the pipeline's VAE state dict
+# Load the HD VAE weights, remap keys from ComfyUI format, and replace the VAE
 hd_vae_path = hf_hub_download(HD_VAE_REPO, "Krea2-HD-vae.safetensors")
 hd_vae_state = load_safetensors(hd_vae_path)
-pipe.vae.load_state_dict(hd_vae_state, strict=True)
-print("HD VAE weights loaded successfully.")
+remapped = _remap_vae_keys(hd_vae_state)
+pipe.vae.load_state_dict(remapped, strict=True)
+print("HD VAE weights loaded and remapped successfully.")
 
 pipe.to("cuda")