Add Turbo/Raw model switch (default Turbo)

README.md

@@ -1,5 +1,5 @@
 ---
-title: Krea 2 Turbo Demo
+title: Krea 2 Demo
 emoji: 🎨
 colorFrom: indigo
 colorTo: purple
@@ -9,9 +9,10 @@ app_file: app.py
 python_version: "3.12"
 startup_duration_timeout: 1h
 pinned: true
-short_description: Krea 2 Turbo text-to-image demo with the Z-Image UI layout
+short_description: Krea 2 demo with a Turbo/Raw model switch
 models:
   - krea/Krea-2-Turbo
+  - krea/Krea-2-Raw
 ---
 
-Krea 2 Turbo text-to-image demo using the same compact control layout as the Z-Image demo.
+Krea 2 text-to-image demo with a Turbo/Raw model switch, using the same compact control layout as the Z-Image demo.

app.py

@@ -4,31 +4,61 @@ import torch
 from diffusers import Krea2Pipeline
 
 
-MODEL_ID = "krea/Krea-2-Turbo"
+MODELS = {
+    "Turbo": {
+        "id": "krea/Krea-2-Turbo",
+        "steps": 8,
+        "cfg": 0.0,
+    },
+    "Raw": {
+        "id": "krea/Krea-2-Raw",
+        "steps": 28,
+        "cfg": 4.5,
+    },
+}
+DEFAULT_MODE = "Turbo"
 DTYPE = torch.bfloat16
 MAX_SEED = 2**31 - 1
 
-pipe = None
+pipes = {}
 
 
-def get_pipe():
-    global pipe
-    if pipe is None:
-        if not torch.cuda.is_available():
-            raise RuntimeError("CUDA is not available. Set this Space hardware to ZeroGPU before generating.")
-        print(f"Loading {MODEL_ID} pipeline...")
-        pipe = Krea2Pipeline.from_pretrained(MODEL_ID, torch_dtype=DTYPE).to("cuda")
-        print("Pipeline loaded!")
-    return pipe
+def get_pipe(mode):
+    if mode not in MODELS:
+        raise gr.Error(f"Unknown model mode: {mode}")
+    if mode in pipes:
+        return pipes[mode]
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is not available. Set this Space hardware to ZeroGPU before generating.")
+
+    model_id = MODELS[mode]["id"]
+
+    def _load():
+        print(f"Loading {model_id} pipeline...")
+        loaded = Krea2Pipeline.from_pretrained(model_id, torch_dtype=DTYPE).to("cuda")
+        print(f"{mode} pipeline loaded!")
+        return loaded
+
+    try:
+        pipes[mode] = _load()
+    except (torch.cuda.OutOfMemoryError, RuntimeError) as exc:
+        # Not enough VRAM to hold both checkpoints at once: drop the others and retry.
+        print(f"Load failed ({exc}); freeing other cached pipelines and retrying.")
+        for other in list(pipes):
+            del pipes[other]
+        torch.cuda.empty_cache()
+        pipes[mode] = _load()
+    return pipes[mode]
 
 
 if torch.cuda.is_available():
-    get_pipe()
+    get_pipe(DEFAULT_MODE)
 else:
     print("CUDA is not available at startup. The UI will load, but generation requires ZeroGPU hardware.")
 
 
 def gpu_duration(
+    mode,
     prompt,
     negative_prompt,
     height,
@@ -45,6 +75,7 @@ def gpu_duration(
 
 @spaces.GPU(duration=gpu_duration, size="xlarge")
 def generate_image(
+    mode,
     prompt,
     negative_prompt,
     height,
@@ -55,13 +86,13 @@ def generate_image(
     randomize_seed,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """Generate 4 Krea 2 Turbo images with seeds: seed, 2x, 3x, 4x."""
+    """Generate 4 Krea 2 images (Turbo or Raw) with seeds: seed, 2x, 3x, 4x."""
     if not prompt or not str(prompt).strip():
         raise gr.Error("Enter a prompt to generate images.")
     if not torch.cuda.is_available():
         raise gr.Error("CUDA is not available. Set this Space hardware to ZeroGPU before generating.")
 
-    pipe = get_pipe()
+    pipe = get_pipe(mode)
     if randomize_seed:
         seed = torch.randint(0, MAX_SEED, (1,)).item()
 
@@ -96,6 +127,12 @@ def generate_image(
     return images, ", ".join(str(s) for s in seeds)
 
 
+def apply_mode(mode):
+    """Reset Steps and CFG to the selected checkpoint's recommended defaults."""
+    cfg = MODELS.get(mode, MODELS[DEFAULT_MODE])
+    return gr.update(value=cfg["steps"]), gr.update(value=cfg["cfg"])
+
+
 examples = [
     [
         "A russet harvest mouse clinging to a branch, macro photograph, shallow depth of field, "
@@ -115,18 +152,27 @@ examples = [
 ]
 
 
-with gr.Blocks(title="Krea 2 Turbo Demo") as demo:
+with gr.Blocks(title="Krea 2 Demo") as demo:
     gr.Markdown(
         """
-        # 🎨 Krea 2 Turbo Demo
+        # 🎨 Krea 2 Demo
 
-        Generate images with [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo).
-        Turbo is the fast distilled Krea 2 checkpoint; the recommended default is 8 steps with CFG 0.0.
+        Generate images with Krea 2. Use the **Model** switch to choose
+        [Turbo](https://huggingface.co/krea/Krea-2-Turbo) (fast distilled — 8 steps, CFG 0.0) or
+        [Raw](https://huggingface.co/krea/Krea-2-Raw) (full quality — ~28-52 steps, CFG 4.5).
+        Switching the model resets Steps and CFG to that checkpoint's recommended defaults.
         """
     )
 
     with gr.Row():
         with gr.Column(scale=1):
+            mode = gr.Radio(
+                choices=["Turbo", "Raw"],
+                value=DEFAULT_MODE,
+                label="Model",
+                info="Turbo = fast preview. Raw = full quality, slower.",
+            )
+
             prompt = gr.Textbox(
                 label="Prompt",
                 placeholder='Describe the image in natural language. Wrap rendered text in quotes, e.g. a sign that reads "open late".',
@@ -158,20 +204,20 @@ with gr.Blocks(title="Krea 2 Turbo Demo") as demo:
             with gr.Row():
                 num_inference_steps = gr.Slider(
                     minimum=1,
-                    maximum=28,
-                    value=8,
+                    maximum=60,
+                    value=MODELS[DEFAULT_MODE]["steps"],
                     step=1,
                     label="Inference Steps",
-                    info="Krea 2 Turbo is designed for 8 steps.",
+                    info="Turbo ~8 steps; Raw ~28-52 steps. Resets with the Model switch.",
                 )
 
             guidance_scale = gr.Slider(
                 minimum=0.0,
-                maximum=5.0,
-                value=0.0,
+                maximum=10.0,
+                value=MODELS[DEFAULT_MODE]["cfg"],
                 step=0.1,
                 label="CFG Guidance Scale",
-                info="Turbo default is 0.0. Negative prompt is ignored when CFG is 0.",
+                info="Turbo uses 0.0; Raw uses ~4.5. Negative prompt is ignored when CFG is 0.",
             )
 
             with gr.Row():
@@ -207,13 +253,19 @@ with gr.Blocks(title="Krea 2 Turbo Demo") as demo:
     )
 
     gr.Markdown(
-        "Model by [Krea](https://huggingface.co/krea). "
-        "This Space follows the Krea 2 Community License and uses the Turbo checkpoint for demo inference."
+        "Models by [Krea](https://huggingface.co/krea). "
+        "This Space follows the Krea 2 Community License and uses the Turbo and Raw checkpoints for demo inference."
     )
 
-    inputs = [prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, seed, randomize_seed]
+    inputs = [mode, prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, seed, randomize_seed]
     outputs = [output_images, used_seeds]
 
+    mode.change(
+        fn=apply_mode,
+        inputs=[mode],
+        outputs=[num_inference_steps, guidance_scale],
+    )
+
     generate_btn.click(
         fn=generate_image,
         inputs=inputs,