All tools in one Gradio Space (ZeroGPU, no Daggr)

app.py

@@ -1,43 +1,251 @@
 """
-Imageat Workflow Agent — ZeroGPU Gradio Space.
-Uses the exact pattern from https://huggingface.co/docs/hub/spaces-zerogpu:
-  - import spaces
-  - @spaces.GPU on GPU-dependent code
-  - gr.Interface(..., fn=...).launch()
+Imageat Workflow Agent — All tools in one Gradio Space (ZeroGPU).
+Single Space, single quota; no Daggr, no iframe. Pattern: spaces + @spaces.GPU + demo.launch().
 """
+import base64
+import os
+import tempfile
+import urllib.request
 import gradio as gr
 import spaces
+from gradio_client import Client, handle_file
 
-# Full Daggr workflow runs on the sibling Space (same code as imageat-workflow).
-# Iframe embedding can break session/token — use "Open in new tab" so nodes work.
-WORKFLOW_SPACE_URL = "https://incendies-imageat-workflow.hf.space"
 
+def _hf_token():
+    """HF token for ZeroGPU quota (Space's HF_TOKEN or huggingface_hub)."""
+    token = os.environ.get("HF_TOKEN")
+    if token:
+        return token
+    try:
+        from huggingface_hub import get_token
+        return get_token()
+    except Exception:
+        return None
 
+
+def _url_to_path(url):
+    """Download image URL to temp file for Gradio display."""
+    if not url or not isinstance(url, str) or not url.startswith("http"):
+        return url
+    try:
+        ext = "png"
+        if ".jpg" in url or ".jpeg" in url:
+            ext = "jpg"
+        elif ".webp" in url:
+            ext = "webp"
+        f = tempfile.NamedTemporaryFile(suffix=f".{ext}", delete=False)
+        req = urllib.request.Request(url, headers={"User-Agent": "Gradio-Imageat/1.0"})
+        with urllib.request.urlopen(req, timeout=60) as r:
+            f.write(r.read())
+        f.close()
+        return f.name
+    except Exception:
+        return url
+
+
+def _image_to_path(image):
+    """Gradio Image → single filepath string."""
+    if image is None:
+        return None
+    if isinstance(image, str):
+        return image
+    if isinstance(image, dict) and image.get("path"):
+        return image["path"]
+    if isinstance(image, (list, tuple)) and len(image) > 0:
+        first = image[0]
+        if isinstance(first, str):
+            return first
+        if isinstance(first, dict) and first.get("path"):
+            return first["path"]
+    return None
+
+
+def _path_for_api(path):
+    """Data URL → temp file path; else return path."""
+    if not path or not isinstance(path, str):
+        return None
+    if path.startswith("data:"):
+        try:
+            header, b64 = path.split(",", 1)
+            ext = "png"
+            if "jpeg" in header or "jpg" in header:
+                ext = "jpg"
+            elif "webp" in header:
+                ext = "webp"
+            data = base64.b64decode(b64)
+            f = tempfile.NamedTemporaryFile(suffix=f".{ext}", delete=False)
+            f.write(data)
+            f.close()
+            return f.name
+        except Exception as e:
+            raise RuntimeError(f"Could not decode data URL: {e}") from e
+    return path
+
+
+# ---------- 1. Background Removal ----------
+def run_bg_removal(image):
+    if image is None:
+        return None
+    path = _image_to_path(image)
+    if not path:
+        return None
+    path = _path_for_api(path)
+    if not path:
+        return None
+    try:
+        client = Client("hf-applications/background-removal", token=_hf_token())
+        result = client.predict(handle_file(path), api_name="/image")
+    except Exception as e:
+        raise RuntimeError(f"Background removal error: {e}") from e
+    # API may return (original, result); we need the result.
+    out = result[-1] if isinstance(result, (list, tuple)) and result else result
+    if out and isinstance(out, str) and out.startswith("http"):
+        return _url_to_path(out)
+    return out
+
+
+# ---------- 2. Upscaler ----------
 @spaces.GPU
-def _zero_gpu_ready():
-    """Ensures HF detects @spaces.GPU at startup and enables ZeroGPU."""
-    return "ZeroGPU ready. Use **Open workflow in new tab** below so the workflow runs correctly."
+def run_upscaler(image, model_selection="4xBHI_dat2_real"):
+    path = _image_to_path(image)
+    if not path:
+        return None
+    path = _path_for_api(path)
+    if not path:
+        return None
+    try:
+        image_arg = handle_file(path)
+    except Exception as e:
+        raise RuntimeError(f"Upscaler: could not load image: {e}") from e
+    try:
+        client = Client("Phips/Upscaler", token=_hf_token())
+        result = client.predict(image_arg, model_selection, api_name="/upscale_image")
+    except Exception as e:
+        raise RuntimeError(f"Upscaler API error: {e}") from e
+    out = None
+    if result and len(result) >= 2 and result[1]:
+        out = result[1]
+    elif result and len(result) >= 1 and result[0]:
+        r0 = result[0]
+        if isinstance(r0, (list, tuple)) and len(r0) >= 2 and r0[1]:
+            out = r0[1]
+        elif isinstance(r0, str):
+            out = r0
+    if out and isinstance(out, str) and out.startswith("http"):
+        return _url_to_path(out)
+    return out
 
 
-with gr.Blocks(title="Imageat Workflow Agent") as demo:
+# ---------- 3. Z-Image Turbo ----------
+@spaces.GPU
+def run_z_image_turbo(prompt, height=1024, width=1024, seed=42):
+    if not prompt or not str(prompt).strip():
+        return None
+    try:
+        client = Client("hf-applications/Z-Image-Turbo", token=_hf_token())
+        result = client.predict(
+            prompt=str(prompt).strip(),
+            height=float(height),
+            width=float(width),
+            seed=int(seed),
+            api_name="/generate_image",
+        )
+    except Exception as e:
+        raise RuntimeError(f"Z-Image-Turbo API error: {e}") from e
+    if result and len(result) >= 1 and result[0]:
+        img = result[0]
+        out = img.get("url") or img.get("path") if isinstance(img, dict) else (img if isinstance(img, str) else None)
+        if out and isinstance(out, str) and out.startswith("http"):
+            return _url_to_path(out)
+        return out
+    return None
+
+
+# ---------- 4. FLUX.2 Klein 9B ----------
+@spaces.GPU
+def run_flux_klein(prompt, mode_choice="Distilled (4 steps)", seed=0, randomize_seed=True, width=1024, height=1024):
+    if not prompt or not str(prompt).strip():
+        return None
+    try:
+        client = Client("black-forest-labs/FLUX.2-klein-9B", token=_hf_token())
+        result = client.predict(
+            prompt=str(prompt).strip(),
+            input_images=[],
+            mode_choice=mode_choice,
+            seed=float(seed),
+            randomize_seed=bool(randomize_seed),
+            width=float(width),
+            height=float(height),
+            num_inference_steps=4.0 if "Distilled" in mode_choice else 50.0,
+            guidance_scale=1.0 if "Distilled" in mode_choice else 3.5,
+            prompt_upsampling=False,
+            api_name="/generate",
+        )
+    except Exception as e:
+        raise RuntimeError(f"FLUX.2-klein-9B API error: {e}") from e
+    if result and len(result) >= 1 and result[0]:
+        img = result[0]
+        out = img.get("url") or img.get("path") if isinstance(img, dict) else (img if isinstance(img, str) else None)
+        if out and isinstance(out, str) and out.startswith("http"):
+            return _url_to_path(out)
+        return out
+    return None
+
+
+# ---------- Gradio UI: one app, all tools in tabs ----------
+with gr.Blocks(title="Imageat Workflow Agent", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# Imageat Workflow — All tools on one Space (ZeroGPU)")
     with gr.Tabs():
-        with gr.TabItem("Workflow"):
-            gr.Markdown(
-                "If the workflow below does not run (e.g. nodes show \"No image\"), "
-                "**open it in a new tab** so your session and API calls work:"
+        with gr.TabItem("Background Removal"):
+            gr.Interface(
+                fn=run_bg_removal,
+                inputs=gr.Image(label="Input Image", type="filepath"),
+                outputs=gr.Image(label="Background Removed"),
+                title="Background Removal",
             )
-            gr.HTML(
-                f'<p style="margin:0.5em 0"><a href="{WORKFLOW_SPACE_URL}" target="_blank" rel="noopener" style="font-weight:600">▶ Open workflow in new tab</a></p>'
+        with gr.TabItem("Upscaler"):
+            gr.Interface(
+                fn=run_upscaler,
+                inputs=[
+                    gr.Image(label="Input Image", type="filepath"),
+                    gr.Dropdown(
+                        label="Model",
+                        choices=["4xBHI_dat2_real", "4xNomos8kDAT", "4xHFA2k", "2xEvangelion_dat2"],
+                        value="4xBHI_dat2_real",
+                    ),
+                ],
+                outputs=gr.Image(label="Upscaled Image"),
+                title="Upscaler",
             )
-            gr.HTML(
-                f'<iframe src="{WORKFLOW_SPACE_URL}" style="width:100%;height:75vh;border:none" title="Daggr workflow"></iframe>'
+        with gr.TabItem("Z-Image Turbo"):
+            gr.Interface(
+                fn=run_z_image_turbo,
+                inputs=[
+                    gr.Textbox(label="Prompt", lines=3),
+                    gr.Slider(512, 1024, value=1024, step=64, label="Height"),
+                    gr.Slider(512, 1024, value=1024, step=64, label="Width"),
+                    gr.Number(value=42, label="Seed", precision=0),
+                ],
+                outputs=gr.Image(label="Generated Image"),
+                title="Z-Image Turbo",
             )
-        with gr.TabItem("GPU check"):
+        with gr.TabItem("FLUX.2 Klein 9B"):
             gr.Interface(
-                fn=_zero_gpu_ready,
-                inputs=None,
-                outputs=gr.Textbox(label="Status"),
-                title="ZeroGPU",
+                fn=run_flux_klein,
+                inputs=[
+                    gr.Textbox(label="Prompt", lines=3),
+                    gr.Radio(
+                        choices=["Distilled (4 steps)", "Base (50 steps)"],
+                        value="Distilled (4 steps)",
+                        label="Mode",
+                    ),
+                    gr.Number(value=0, label="Seed", precision=0),
+                    gr.Checkbox(value=True, label="Randomize seed"),
+                    gr.Slider(512, 1024, value=1024, step=64, label="Width"),
+                    gr.Slider(512, 1024, value=1024, step=64, label="Height"),
+                ],
+                outputs=gr.Image(label="Generated Image"),
+                title="FLUX.2 Klein 9B",
             )
 
 demo.launch()

requirements.txt

@@ -1,2 +1,3 @@
 gradio
 spaces
+gradio_client