Upload app.py with huggingface_hub

app.py

@@ -4,22 +4,25 @@ import spaces
 import torch
 import gradio as gr
 from huggingface_hub import hf_hub_download
-from diffusers import Krea2Pipeline, AutoencoderKLQwenImage
+from safetensors.torch import load_file as load_safetensors
+from diffusers import Krea2Pipeline
 
 TURBO_REPO = "krea/Krea-2-Turbo"
 HD_VAE_REPO = "wikeeyang/Krea2-Turbo-HD-V1"
 MAX_SEED = 2**31 - 1
 
-# Load the Krea-2-Turbo pipeline (full diffusers format) and swap in the HD VAE
-# from wikeeyang/Krea2-Turbo-HD-V1 for enhanced detail, clarity, and contrast.
+# 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.
 pipe = Krea2Pipeline.from_pretrained(TURBO_REPO, torch_dtype=torch.bfloat16)
 
-# Load and replace the VAE with the HD-optimized version
+# Load the HD VAE weights and replace the pipeline's VAE state dict
 hd_vae_path = hf_hub_download(HD_VAE_REPO, "Krea2-HD-vae.safetensors")
-hd_vae = AutoencoderKLQwenImage.from_single_file(
-    hd_vae_path, torch_dtype=torch.bfloat16
-)
-pipe.vae = hd_vae
+hd_vae_state = load_safetensors(hd_vae_path)
+pipe.vae.load_state_dict(hd_vae_state, strict=True)
+print("HD VAE weights loaded successfully.")
 
 pipe.to("cuda")
 
@@ -79,51 +82,6 @@ CSS = """
 .dark .gradio-container { color: var(--body-text-color); }
 """
 
-with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(
-            "# Krea 2 Turbo HD\n"
-            "Text-to-image generation with the HD-optimized VAE from "
-            "[wikeeyang/Krea2-Turbo-HD-V1](https://huggingface.co/wikeeyang/Krea2-Turbo-HD-V1), "
-            "built on [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo). "
-            "The fine-tuned VAE enhances detail rendering, clarity, and contrast."
-        )
-
-        with gr.Row():
-            prompt = gr.Textbox(
-                show_label=False,
-                placeholder=PLACEHOLDER,
-                container=False,
-                scale=4,
-                lines=3,
-                autofocus=True,
-            )
-            run = gr.Button("Generate", variant="primary", scale=1)
-
-        output = gr.Image(label="Result", format="png")
-
-        with gr.Accordion("Advanced settings", open=False):
-            with gr.Row():
-                width = gr.Slider(512, 2048, value=1024, step=16, label="Width")
-                height = gr.Slider(512, 2048, value=1024, step=16, label="Height")
-            steps = gr.Slider(1, 50, value=8, step=1, label="Steps")
-            with gr.Row():
-                seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
-                randomize = gr.Checkbox(value=True, label="Randomize seed")
-
-            with gr.Accordion("Prompting tips", open=False):
-                gr.Markdown(PROMPT_TIPS)
-
-        gr.Examples(
-            examples=EXAMPLE_PROMPTS,
-            inputs=[prompt],
-            outputs=output,
-            fn=lambda p: generate(p),
-            cache_examples=True,
-            cache_mode="lazy",
-            label="Example prompts",
-        )
-
 
 def _duration(prompt, width, height, steps, seed, randomize):
     """Estimate GPU time based on steps and pixel area."""
@@ -174,16 +132,61 @@ def generate(
     return image, seed
 
 
-run.click(
-    generate,
-    inputs=[prompt, width, height, steps, seed, randomize],
-    outputs=[output, seed],
-    api_name="generate",
-)
-prompt.submit(
-    generate,
-    inputs=[prompt, width, height, steps, seed, randomize],
-    outputs=[output, seed],
-)
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            "# Krea 2 Turbo HD\n"
+            "Text-to-image generation with the HD-optimized VAE from "
+            "[wikeeyang/Krea2-Turbo-HD-V1](https://huggingface.co/wikeeyang/Krea2-Turbo-HD-V1), "
+            "built on [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo). "
+            "The fine-tuned VAE enhances detail rendering, clarity, and contrast."
+        )
+
+        with gr.Row():
+            prompt = gr.Textbox(
+                show_label=False,
+                placeholder=PLACEHOLDER,
+                container=False,
+                scale=4,
+                lines=3,
+                autofocus=True,
+            )
+            run = gr.Button("Generate", variant="primary", scale=1)
+
+        output = gr.Image(label="Result", format="png")
+
+        with gr.Accordion("Advanced settings", open=False):
+            with gr.Row():
+                width = gr.Slider(512, 2048, value=1024, step=16, label="Width")
+                height = gr.Slider(512, 2048, value=1024, step=16, label="Height")
+            steps = gr.Slider(1, 50, value=8, step=1, label="Steps")
+            with gr.Row():
+                seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
+                randomize = gr.Checkbox(value=True, label="Randomize seed")
+
+            with gr.Accordion("Prompting tips", open=False):
+                gr.Markdown(PROMPT_TIPS)
+
+        gr.Examples(
+            examples=EXAMPLE_PROMPTS,
+            inputs=[prompt],
+            outputs=[output, seed],
+            fn=generate,
+            cache_examples=True,
+            cache_mode="lazy",
+            label="Example prompts",
+        )
+
+    run.click(
+        generate,
+        inputs=[prompt, width, height, steps, seed, randomize],
+        outputs=[output, seed],
+        api_name="generate",
+    )
+    prompt.submit(
+        generate,
+        inputs=[prompt, width, height, steps, seed, randomize],
+        outputs=[output, seed],
+    )
 
 demo.launch(mcp_server=True)