model selector, better offline UX

.gitignore

@@ -0,0 +1,44 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Distribution / packaging
+build/
+dist/
+*.egg-info/
+.eggs/
+
+# Virtual environments
+.venv/
+venv/
+env/
+ENV/
+
+# Environment / secrets
+.env
+.env.*
+
+# Caches
+.cache/
+.ruff_cache/
+.mypy_cache/
+.pytest_cache/
+
+# Hugging Face / model caches
+.huggingface/
+hf_cache/
+
+# Jupyter
+.ipynb_checkpoints/
+
+# OS / editor
+.DS_Store
+Thumbs.db
+.idea/
+.vscode/
+*.swp
+
+# Gradio
+.gradio/
+flagged/

app.py

@@ -111,10 +111,6 @@ def _load(model_id):
 
 def get_model(precision):
     """Return (model, processor); load the selected model, swapping out any previously-loaded one."""
-    if IS_ZERO_GPU:
-        precision = (
-            DEFAULT_PRECISION  # ZeroGPU serves bf16 only, whatever the caller passes
-        )
     model_id = MODELS[precision]
     if _current["id"] != model_id:
         # Free the current model before loading the new one (one resident at a time).
@@ -128,8 +124,30 @@ def get_model(precision):
     return _current["model"], _current["processor"]
 
 
-# Preload the default so the first caption is fast (and gives ZeroGPU its module-scope load).
-get_model(DEFAULT_PRECISION)
+# Two paths:
+#   ZeroGPU  -> the platform hands out the GPU only for the duration of a @spaces.GPU call and takes
+#              it right back, so there's nothing to "hog". We just preload bf16 at module scope (the
+#              canonical ZeroGPU pattern) and skip the whole load/unload + quantization machinery.
+#   Local    -> we hold a real GPU, so this support app must not squat on VRAM while the user runs the
+#              Ideogram generation app. Nothing loads until the Load button; Unload frees it again.
+if IS_ZERO_GPU:
+    get_model(DEFAULT_PRECISION)
+
+
+@spaces.GPU(duration=120)
+def load_model(precision):
+    """Load the selected model into VRAM (local only; nothing is resident until this runs)."""
+    get_model(precision)
+    return f"✅ Loaded: {precision}"
+
+
+def unload_model():
+    """Evict the resident model from VRAM and RAM so it stops hogging the GPU (local only)."""
+    _current.update(id=None, model=None, processor=None)
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    return "⚪ No model loaded"
 
 
 def _extract_json(text: str) -> str:
@@ -156,6 +174,8 @@ def caption(
     if image is None:
         raise gr.Error("An image is required.")
 
+    if IS_ZERO_GPU:
+        precision = DEFAULT_PRECISION  # ZeroGPU serves bf16 only
     model, processor = get_model(precision)
     messages = [
         {
@@ -188,27 +208,63 @@ def caption(
     return _extract_json(text)
 
 
-with gr.Blocks(title="Ideogram4 Caption (Gemma4)") as demo:
-    gr.Markdown(
-        "# Ideogram4 Caption — Gemma4\n"
-        "Upload an image → get Ideogram4's rich JSON caption (with pixel-space `bbox`es). Edit the JSON "
-        "(add / remove / move / resize elements) and feed it to the Ideogram4 generation Space."
+CSS = """
+#caption-btn {
+    background: linear-gradient(180deg, #2563eb 0%, #1e3a8a 100%) !important;
+    color: #fff !important; border: none !important;
+    box-shadow: inset 0 1px 0 rgba(255,255,255,0.15), 0 1px 2px rgba(0,0,0,0.25) !important;
+}
+#load-btn {
+    background: linear-gradient(180deg, #15803d 0%, #14532d 100%) !important;
+    color: #fff !important; border: none !important;
+    box-shadow: inset 0 1px 0 rgba(255,255,255,0.15), 0 1px 2px rgba(0,0,0,0.25) !important;
+}
+#unload-btn {
+    background: linear-gradient(180deg, #b91c1c 0%, #7f1d1d 100%) !important;
+    color: #fff !important; border: none !important;
+    box-shadow: inset 0 1px 0 rgba(255,255,255,0.15), 0 1px 2px rgba(0,0,0,0.25) !important;
+}
+#caption-btn, #load-btn, #unload-btn { min-height: 56px !important; height: 56px !important; }
+"""
+
+intro = (
+    "# Ideogram4 Caption — Gemma4\n"
+    "Upload an image → get Ideogram4's rich JSON caption (with pixel-space `bbox`es). Edit the JSON "
+    "(add / remove / move / resize elements) and feed it to the Ideogram4 generation Space."
+)
+if not IS_ZERO_GPU:
+    intro += (
+        "\n\nOne model at a time — switching + Load reloads. bf16 needs a >=24GB GPU; "
+        "SDNQ 4/8-bit fit 16GB."
     )
+
+with gr.Blocks(title="Ideogram4 Caption (Gemma4)", css=CSS) as demo:
+    gr.Markdown(intro)
     with gr.Row():
         with gr.Column():
+            run = gr.Button("Caption", variant="primary", elem_id="caption-btn")
+
+            if IS_ZERO_GPU:
+                # ZeroGPU: bf16 is preloaded and the GPU is per-request — no selector / load / unload.
+                precision = gr.State(value=DEFAULT_PRECISION)
+            else:
+                # Local: pick a model, load it on demand, and free it so we don't squat on VRAM.
+                with gr.Row(equal_height=True):
+                    precision = gr.Dropdown(
+                        choices=list(MODELS),
+                        value=DEFAULT_PRECISION,
+                        show_label=False,
+                        scale=2,
+                    )
+                    load_btn = gr.Button("Load model", elem_id="load-btn")
+                    unload_btn = gr.Button("Unload model", elem_id="unload-btn")
+                model_status = gr.Textbox(
+                    label="Model status", value="⚪ No model loaded", interactive=False
+                )
+
             image = gr.Image(label="Source image", type="pil")
-            # Hidden on ZeroGPU (bf16-only); locally the user picks one model at a time (switching reloads).
-            precision = gr.Dropdown(
-                choices=list(MODELS),
-                value=DEFAULT_PRECISION,
-                label="Model / precision",
-                info="One model at a time — switching reloads. bf16 (full) needs a >=24GB GPU; "
-                "use SDNQ 4-bit/8-bit on a 16GB card.",
-                visible=not IS_ZERO_GPU,
-            )
-            instruction = gr.Textbox(
-                label="Instruction", value=DEFAULT_INSTRUCTION, lines=10
-            )
+            # The Ideogram4 schema prompt is fixed and hidden; kept as a state so it's still passed to caption().
+            instruction = gr.State(value=DEFAULT_INSTRUCTION)
             with gr.Row():
                 max_new_tokens = gr.Slider(
                     256, 4096, value=2048, step=64, label="Max new tokens"
@@ -216,10 +272,15 @@ with gr.Blocks(title="Ideogram4 Caption (Gemma4)") as demo:
                 temperature = gr.Slider(
                     0.0, 2.0, value=0.0, step=0.1, label="Temperature (0 = greedy)"
                 )
-            run = gr.Button("Caption", variant="primary")
         with gr.Column():
             result = gr.Code(label="Ideogram4 JSON caption", language="json")
 
+    if not IS_ZERO_GPU:
+        load_btn.click(load_model, inputs=precision, outputs=model_status, api_name="load")
+        unload_btn.click(
+            unload_model, inputs=None, outputs=model_status, api_name="unload"
+        )
+
     run.click(
         caption,
         inputs=[image, instruction, precision, max_new_tokens, temperature],