Update app.py

app.py

@@ -73,13 +73,14 @@ DEFAULTS = {
     "Turbo": {"steps": 8, "guidance": 0.0},
 }
 
-# Resolution presets. The prompt guide notes Turbo renders up to 2K and every
-# official sample is generated at 2K, so the presets default users toward it.
+# Resolution presets. The model renders up to 2K, but the compiled transformer
+# block can exceed this Space's GPU memory above 1024, so 1024 is the default
+# and larger sizes are opt-in (see the OOM guard in generate).
 RESOLUTIONS = {
     "Square · 1024": (1024, 1024),
+    "Portrait · 1024": (832, 1216),
+    "Landscape · 1024": (1216, 832),
     "Square · 2K": (2048, 2048),
-    "Portrait · 2K": (1536, 2048),
-    "Landscape · 2K": (2048, 1536),
 }
 
 PROMPT_TIPS = """\
@@ -88,7 +89,7 @@ Krea 2 is tuned for natural language. Describe the image the way you would descr
 - Write in full sentences or rich phrases. Longer, more specific prompts give the best results, but short prompts work too.
 - Name the things that matter: subject, setting, lighting, color, framing, medium, and mood.
 - To render text in the image, wrap the words in quotes, for example: a storefront window with a neon sign that reads "open late".
-- Turbo renders up to 2K. The example prompts below are all 2K.
+- 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.
 """
@@ -162,7 +163,10 @@ PLACEHOLDER = (
 
 
 def _duration(prompt, negative_prompt, model, steps, guidance, width, height, seed, randomize, progress=None):
-    return int(steps) * 2 + 25
+    # Scale the GPU reservation by step count and pixel area so larger renders
+    # are not killed before they finish.
+    megapixels = max(1.0, (int(width) * int(height)) / (1024 * 1024))
+    return int(int(steps) * 2 * megapixels + 25)
 
 
 @spaces.GPU(duration=_duration, size="xlarge")
@@ -185,15 +189,25 @@ def generate(
     seed = int(seed)
     generator = torch.Generator("cuda").manual_seed(seed)
     pipe = PIPES[model]
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=(negative_prompt or None) if guidance > 0 else None,
-        height=int(height),
-        width=int(width),
-        num_inference_steps=int(steps),
-        guidance_scale=float(guidance),
-        generator=generator,
-    ).images[0]
+    try:
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=(negative_prompt or None) if guidance > 0 else None,
+            height=int(height),
+            width=int(width),
+            num_inference_steps=int(steps),
+            guidance_scale=float(guidance),
+            generator=generator,
+        ).images[0]
+    except RuntimeError as exc:
+        # At high resolution the compiled transformer block can exhaust GPU
+        # memory, which surfaces as a CUDA allocation / AOTI runtime error.
+        # Recover the worker and tell the user how to fix it.
+        torch.cuda.empty_cache()
+        raise gr.Error(
+            f"Generation failed at {int(width)}x{int(height)}. This is usually the GPU running "
+            "out of memory at high resolution. Try 1024x1024 or a smaller size."
+        ) from exc
     return image, seed
 
 
@@ -279,7 +293,15 @@ CSS = """
   margin: 0;
   max-width: 60ch;
 }
-#krea-header .badges { margin-top: 14px; display: flex; gap: 8px; }
+#krea-header .meta {
+  margin-top: 18px;
+  display: flex;
+  justify-content: space-between;
+  align-items: center;
+  flex-wrap: wrap;
+  gap: 12px;
+}
+#krea-header .badges { display: flex; gap: 8px; }
 #krea-header .badge {
   font-family: 'JetBrains Mono', ui-monospace, monospace;
   font-size: 10px;
@@ -290,10 +312,17 @@ CSS = """
   border-radius: 999px;
   padding: 4px 10px;
 }
-
-/* Mono utility type for control labels keeps the technical, tool-like register. */
-.panel .block .gr-form label span,
-.panel span[data-testid="block-info"] { letter-spacing: 0.01em; }
+#krea-header .links { display: flex; gap: 16px; }
+#krea-header .links a {
+  font-family: 'JetBrains Mono', ui-monospace, monospace;
+  font-size: 11px;
+  letter-spacing: 0.08em;
+  text-transform: uppercase;
+  color: #737373;
+  text-decoration: none;
+  transition: color 0.15s ease;
+}
+#krea-header .links a:hover { color: #f5f5f5; }
 
 #generate-btn { font-weight: 600; letter-spacing: 0.01em; }
 
@@ -310,10 +339,16 @@ with gr.Blocks(theme=theme, css=CSS, title="Krea 2", fill_height=True) 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 up to 2K.</p>
-              <div class="badges">
-                <span class="badge">Raw · CFG</span>
-                <span class="badge">Turbo · few-step</span>
+              <p class="subtitle">Generate images from natural language. Pick Raw for CFG-guided control or Turbo for fast, few-step results.</p>
+              <div class="meta">
+                <div class="badges">
+                  <span class="badge">Raw · CFG</span>
+                  <span class="badge">Turbo · few-step</span>
+                </div>
+                <div class="links">
+                  <a href="https://www.krea.ai/blog/krea-2-technical-report" target="_blank" rel="noopener">Technical report ↗</a>
+                  <a href="https://github.com/krea-ai/krea-2" target="_blank" rel="noopener">GitHub ↗</a>
+                </div>
               </div>
             </header>
             """
@@ -336,7 +371,7 @@ with gr.Blocks(theme=theme, css=CSS, title="Krea 2", fill_height=True) as demo:
 
                 resolution = gr.Radio(
                     list(RESOLUTIONS.keys()),
-                    value="Square · 2K",
+                    value="Square · 1024",
                     label="Resolution",
                 )
 
@@ -350,8 +385,8 @@ with gr.Blocks(theme=theme, css=CSS, title="Krea 2", fill_height=True) as demo:
                     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")
                     with gr.Row():
-                        width = gr.Slider(512, 2048, value=2048, step=16, label="Width")
-                        height = gr.Slider(512, 2048, value=2048, step=16, label="Height")
+                        width = gr.Slider(512, 2048, value=1024, step=16, label="Width")
+                        height = gr.Slider(512, 2048, value=1024, step=16, label="Height")
                     with gr.Row():
                         seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                         randomize = gr.Checkbox(value=True, label="Randomize seed")