Create app.py

app.py

@@ -0,0 +1,202 @@
+import os
+import io
+import base64
+import time
+import logging
+from pathlib import Path
+
+import gradio as gr
+from gradio_client import Client
+from PIL import Image
+
+# ─────────  Logging ─────────
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# ─────────  Constants ─────────
+PREDICT_TIMEOUT = 600               # ⏱  increase backend timeout (seconds)
+GPU_WARM_WINDOW = 900               # 🕒  keep GPU alive 15 min (must match backend)
+MAX_RETRIES = 1                     # 🔄 avoid multiple backend calls per click
+
+# ─────────  Backend connection ─────────
+# HF_TOKEN = os.getenv("HF_TOKEN")
+# if not HF_TOKEN:
+#     raise ValueError("HF_TOKEN environment variable is required")
+#-------
+
+backend_status = {
+    "client": None,
+    "connected": False,
+    "last_check": None,
+    "error_message": ""
+}
+
+
+def check_backend_connection():
+    """Ping the HF Space and cache the client object."""
+    try:
+        test_client = Client("SnapwearAI/SAKS-backend", hf_token=HF_TOKEN)
+        backend_status.update(
+            {
+                "client": test_client,
+                "connected": True,
+                "error_message": "",
+                "last_check": time.time(),
+            }
+        )
+        logger.info("✅ Backend connection established")
+        return True, "🟢 Model is ready"
+    except Exception as e:
+        backend_status.update(
+            {
+                "client": None,
+                "connected": False,
+                "last_check": time.time(),
+                "error_message": str(e),
+            }
+        )
+        err = str(e).lower()
+        if "timeout" in err or "read operation timed out" in err:
+            return False, "🟡 Model is starting up. Please wait 3‑4 min."
+        return False, f"🔴 Backend error: {e}"
+
+
+# initial probe
+# check_backend_connection()
+
+# ─────────  Helpers ─────────
+
+
+def image_to_base64(image: Image.Image) -> str:
+    if image is None:
+        return ""
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    buf = io.BytesIO()
+    image.save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode()
+
+
+def base64_to_image(b64: str) -> Image.Image | None:
+    if not b64:
+        return None
+    try:
+        return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
+    except Exception as e:
+        logger.error(f"Failed to decode base64 → image: {e}")
+        return None
+
+
+# ─────────  UI ↔ Backend bridge ─────────
+
+
+def call_backend_with_retry(input_image: Image.Image, category: str, gender: str, *, max_retries: int = MAX_RETRIES):
+    """Single‑shot call (no more than `max_retries` times)."""
+
+    if input_image is None:
+        return None, None, "❌ Please upload an image."
+
+    # lazy reconnect
+    if not backend_status["connected"]:
+        ok, msg = check_backend_connection()
+        if not ok:
+            return None, None, msg
+
+    client: Client = backend_status["client"]
+    img_b64 = image_to_base64(input_image)
+
+    for attempt in range(max_retries):
+        try:
+            logger.info(f"Backend call #{attempt+1}")
+            start = time.time()
+            result = client.predict(
+                img_b64,
+                category,
+                gender,
+                api_name="/predict",
+            )
+            dt = time.time() - start
+
+            if not result or len(result) < 4:
+                raise ValueError("Invalid response structure from backend")
+
+            _, overlay_b64, bg_b64, status = result
+            overlay_img = base64_to_image(overlay_b64)
+            bg_img = base64_to_image(bg_b64)
+
+            if overlay_img is None or bg_img is None:
+                raise ValueError("Failed to decode backend images")
+
+            if not status.startswith("✅"):
+                status = "✅ " + status
+            status += f" (⏱ {dt:.1f}s)"
+            return overlay_img, bg_img, status
+
+        except Exception as e:
+            logger.error(f"Attempt {attempt+1} failed: {e}")
+            if attempt == max_retries - 1:
+                return None, None, f"❌ {e}"
+            time.sleep(1)
+
+    return None, None, "❌ Unknown error"
+
+
+# ─────────  CSS ─────────
+CSS_PATH = Path(__file__).with_name("style.css")
+CUSTOM_CSS = CSS_PATH.read_text() if CSS_PATH.exists() else ""  # external CSS if you prefer
+
+
+# ─────────  Gradio Blocks ─────────
+with gr.Blocks(css=CUSTOM_CSS, title="SAKS Product Photography Demo") as demo:
+    # Hero
+    gr.HTML(
+        """
+        <div class="hero-section">
+            <h1 style="font-size:48px;margin:0;background:linear-gradient(45deg,#fff,#f0f8ff);-webkit-background-clip:text;-webkit-text-fill-color:transparent;">💍 SAKS Product Photography</h1>
+            <p style="font-size:18px;margin:10px 0;opacity:0.8;">Detect jewellery → segment → generate pro model photo</p>
+        </div>
+        """
+    )
+
+    # Status banner
+    status_html = gr.HTML()
+
+    def _update_status():
+        ok, msg = check_backend_connection()
+        cls = "status-ready" if ok else ("status-starting" if "🟡" in msg else "status-error")
+        return f'<div class="status-banner {cls}">{msg}</div>'
+
+    status_html.value = _update_status()
+    gr.Button("🔄 Check Status").click(fn=_update_status, outputs=status_html)
+
+    # 2‑column layout
+    with gr.Row():
+        # left inputs
+        with gr.Column():
+            input_img = gr.Image(label="Upload image", type="pil", height=400)
+            category = gr.Dropdown(label="Jewellery category", choices=["Rings", "Bracelets", "Watches", "Earrings"], value="Bracelets")
+            gender = gr.Dropdown(label="Model gender", choices=["male", "female"], value="female")
+            run_btn = gr.Button("🎯 Generate", elem_id="button")
+
+        # right outputs
+        with gr.Column():
+            with gr.Tabs():
+                with gr.TabItem("Detection overlay"):
+                    out_overlay = gr.Image(height=400)
+                with gr.TabItem("Final result"):
+                    out_bg = gr.Image(height=400)
+            out_status = gr.Text(label="Status", interactive=False)
+
+    # wire button → backend
+    run_btn.click(
+        fn=call_backend_with_retry,
+        inputs=[input_img, category, gender],
+        outputs=[out_overlay, out_bg, out_status],
+        concurrency_limit=1,
+        show_progress=True,
+    )
+
+
+# ─────────  Launch ─────────
+if __name__ == "__main__":
+    demo.queue(max_size=20, default_concurrency_limit=1).launch(share=False)