Use Redcraft checkpoint and public AoTI kernels

app.py

@@ -6,7 +6,7 @@ import sys
 import spaces
 import torch
 import gradio as gr
-from huggingface_hub import login
+from huggingface_hub import hf_hub_download, login
 
 if os.environ.get("HF_TOKEN"):
     login(token=os.environ["HF_TOKEN"])
@@ -14,44 +14,37 @@ if os.environ.get("HF_TOKEN"):
 from diffusers import Krea2Pipeline
 
 DTYPE = torch.bfloat16
-RAW_REPO = "krea/Krea-2-Raw"
-TURBO_REPO = "krea/Krea-2-Turbo"
+MODEL_REPO = "Daankular/redcraft-krea2-fp8"
+MODEL_FILE = "redcraftKREA2RedMix_krea2Edition.safetensors"
+AOTI_REPO = "multimodalart/krea2-aoti-kernels"
+AOTI_FILE = "Krea2TransformerBlock-lora-r64/package.pt2"
 MAX_SEED = 2**31 - 1
 
-# Both models are loaded at global scope. They share the architecture, so Turbo
-# reuses Raw's text encoder / tokenizer / VAE and only loads its own transformer.
-pipe_raw = Krea2Pipeline.from_pretrained(RAW_REPO, torch_dtype=DTYPE)
-pipe_turbo = Krea2Pipeline.from_pretrained(
-    TURBO_REPO,
-    text_encoder=pipe_raw.text_encoder,
-    tokenizer=pipe_raw.tokenizer,
-    vae=pipe_raw.vae,
-    torch_dtype=DTYPE,
+checkpoint_path = hf_hub_download(
+    repo_id=MODEL_REPO,
+    filename=MODEL_FILE,
+    token=os.environ.get("HF_TOKEN"),
 )
-pipe_raw.to("cuda")
-pipe_turbo.to("cuda")
+pipe_redcraft = Krea2Pipeline.from_single_file(checkpoint_path, torch_dtype=DTYPE)
+pipe_redcraft.to("cuda")
 
 
 
 def _load_aoti():
-    # One compiled Krea2TransformerBlock (kernels only, weights stay live) serves
-    # both pipelines since their transformer configs are identical. Mirrors
-    # spaces.aoti_blocks_load, but downloads from the private artifact dataset
-    # with an explicit write token instead of ambient model-repo auth.
-    from huggingface_hub import hf_hub_download
+    # The compiled Krea2TransformerBlock kernels are public and separate from
+    # the checkpoint weights. If they cannot be loaded, the Space falls back to
+    # eager mode.
     from spaces.zero.torch.aoti import LazyAOTIModel
 
     pt2 = hf_hub_download(
-        repo_id="multimodalart/Krea-2-aoti",
-        filename="Krea2TransformerBlock/package.pt2",
+        repo_id=AOTI_REPO,
+        filename=AOTI_FILE,
         repo_type="dataset",
-        token=os.environ.get("HF_WRITE_TOKEN"),
     )
     aoti_model = LazyAOTIModel(pt2)
-    for pipe in (pipe_raw, pipe_turbo):
-        for block in pipe.transformer.modules():
-            if block.__class__.__name__ == "Krea2TransformerBlock":
-                spaces.aoti_patch(block, aoti_model)
+    for block in pipe_redcraft.transformer.modules():
+        if block.__class__.__name__ == "Krea2TransformerBlock":
+            spaces.aoti_patch(block, aoti_model)
 
 
 try:
@@ -60,10 +53,9 @@ try:
 except Exception as e:
     print(f"AoTI load skipped ({e}); running eager.")
 
-PIPES = {"Raw": pipe_raw, "Turbo": pipe_turbo}
+PIPES = {"Redcraft": pipe_redcraft}
 DEFAULTS = {
-    "Raw": {"steps": 28, "guidance": 4.5},
-    "Turbo": {"steps": 8, "guidance": 0.0},
+    "Redcraft": {"steps": 8, "guidance": 0.0},
 }
 
 # Resolution presets. The model renders up to 2K, but the compiled transformer
@@ -84,7 +76,7 @@ Krea 2 is tuned for natural language. Describe the image the way you would descr
 - To render text in the image, wrap the words in quotes, for example: a storefront window with a neon sign that reads "open late".
 - The model can render up to 2K, but very high resolutions may run out of GPU memory on this Space. 1024 is the reliable default.
 
-Want help writing longer prompts? An `expansion.txt` system prompt is provided in the [model repo](https://huggingface.co/krea/Krea-2-Turbo) for use with any LLM.
+This Space uses the Redcraft checkpoint from [Daankular/redcraft-krea2-fp8](https://huggingface.co/Daankular/redcraft-krea2-fp8).
 """
 
 # Drawn from the official Krea 2 prompt guide. These demonstrate the
@@ -191,7 +183,7 @@ def _duration(prompt, negative_prompt, model, steps, guidance, width, height, se
 def generate(
     prompt,
     negative_prompt="",
-    model="Turbo",
+    model="Redcraft",
     steps=8,
     guidance=0,
     width=1024,
@@ -369,11 +361,11 @@ with gr.Blocks(title="Krea 2") as demo:
             <header id="krea-header">
               <div class="eyebrow">KREA · TEXT-TO-IMAGE</div>
               <h1>Krea 2</h1>
-              <p class="subtitle">Generate images from natural language. Pick Raw for CFG-guided control or Turbo for fast, few-step results.</p>
+              <p class="subtitle">Generate images from natural language with the Redcraft Krea 2 checkpoint.</p>
               <div class="meta">
                 <div class="badges">
-                  <span class="badge">Raw · CFG</span>
-                  <span class="badge">Turbo · few-step</span>
+                  <span class="badge">Redcraft</span>
+                  <span class="badge">Public AoTI kernels</span>
                 </div>
                 <div class="links">
                   <a href="https://www.krea.ai/blog/krea-2-technical-report" target="_blank" rel="noopener">Technical report ↗</a>
@@ -393,7 +385,7 @@ with gr.Blocks(title="Krea 2") as demo:
                     show_label=True,
                     autofocus=True,
                 )
-                model = gr.Radio(["Turbo", "Raw"], value="Turbo", label="Model")
+                model = gr.Radio(["Redcraft"], value="Redcraft", label="Model")
                 run = gr.Button("Generate", variant="primary", elem_id="generate-btn")
 
                 with gr.Accordion("Prompting tips", open=False):
@@ -410,7 +402,7 @@ with gr.Blocks(title="Krea 2") as demo:
                         label="Negative prompt",
                         lines=1,
                         interactive=False,
-                        info="Available with Raw, where guidance is above 0.",
+                        info="Enabled when guidance is above 0.",
                     )
                     steps = gr.Slider(1, 50, value=8, step=1, label="Steps")
                     guidance = gr.Slider(0.0, 10.0, value=0.0, step=0.1, label="Guidance scale")
@@ -460,4 +452,4 @@ with gr.Blocks(title="Krea 2") as demo:
     run.click(generate, inputs, [output, seed])
     prompt.submit(generate, inputs, [output, seed])
 
-demo.launch(theme=theme, css=CSS)
+demo.launch(theme=theme, css=CSS)