Update app.py

app.py

@@ -1,100 +1,22 @@
-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
-GPU_WARM_WINDOW = 900
-MAX_RETRIES = 1
-
-# ─────────  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("milliyin/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."
+        return None, None, "❌ Please upload an image.", gr.update(interactive=True)
 
     if not backend_status["connected"]:
         ok, msg = check_backend_connection()
         if not ok:
-            return None, None, msg
+            return None, None, msg, gr.update(interactive=True)
 
     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}")
+            logger.info(f"Backend call #{attempt+1}")
             start = time.time()
             result = client.predict(
                 img_b64,
@@ -116,254 +38,33 @@ def call_backend_with_retry(input_image: Image.Image, category: str, gender: str
 
             if not status.startswith("✅"):
                 status = "✅ " + status
-            status += f" (⏱ {dt:.1f}s)"
-            return overlay_img, bg_img, status
+            status += f" (⏱ {dt:.1f}s)"
+            return overlay_img, bg_img, status, gr.update(interactive=True)
 
         except Exception as e:
-            logger.error(f"Attempt {attempt+1} failed: {e}")
+            logger.error(f"Attempt {attempt+1} failed: {e}")
             if attempt == max_retries - 1:
-                return None, None, f"❌ {e}"
+                return None, None, f"❌ {e}", gr.update(interactive=True)
             time.sleep(1)
     
-    return None, None, "❌ Unknown error",gr.update(interactive=True)
+    return None, None, "❌ Unknown error", gr.update(interactive=True)
 
 def disable_button():
     return gr.update(interactive=False)
 
-# ─────────  CSS ─────────
-custom_css = """
-.gradio-container {
-    background: linear-gradient(135deg, #3b4371 0%, #2d1b69 25%, #673ab7 50%, #8e24aa 75%, #6a1b9a 100%);
-    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
-    min-height: 100vh;
-}
-.contain {
-    background: rgba(255, 255, 255, 0.95);
-    border-radius: 15px;
-    padding: 25px;
-    margin: 15px;
-    box-shadow: 0 10px 30px rgba(0, 0, 0, 0.2);
-    backdrop-filter: blur(10px);
-}
-.title-container {
-    text-align: center;
-    margin-bottom: 25px;
-    padding: 20px;
-    background: linear-gradient(135deg, #673ab7, #8e24aa);
-    border-radius: 12px;
-    box-shadow: 0 5px 20px rgba(103, 58, 183, 0.4);
-}
-.title-container h1 {
-    color: white;
-    font-size: 2.2em;
-    font-weight: bold;
-    margin: 0;
-    text-shadow: 1px 1px 3px rgba(0, 0, 0, 0.3);
-}
-.info-bar {
-    background: linear-gradient(135deg, #7c4dff, #6a1b9a);
-    padding: 12px;
-    border-radius: 8px;
-    margin-bottom: 20px;
-    color: white;
-    text-align: center;
-    font-weight: 500;
-    box-shadow: 0 3px 12px rgba(124, 77, 255, 0.3);
-}
-.section-header {
-    background: linear-gradient(135deg, #e1bee7, #d1c4e9);
-    padding: 12px;
-    border-radius: 8px;
-    margin-bottom: 15px;
-    border-left: 4px solid #673ab7;
-}
-.section-header h3 {
-    margin: 0;
-    color: #333;
-    font-weight: 600;
-}
-.input-group {
-    background: rgba(255, 255, 255, 0.85);
-    padding: 18px;
-    border-radius: 12px;
-    margin-bottom: 15px;
-    border: 1px solid rgba(103, 58, 183, 0.2);
-    box-shadow: 0 3px 12px rgba(103, 58, 183, 0.1);
-}
-.result-section {
-    background: rgba(255, 255, 255, 0.9);
-    padding: 18px;
-    border-radius: 12px;
-    border: 1px solid rgba(103, 58, 183, 0.2);
-    box-shadow: 0 3px 12px rgba(103, 58, 183, 0.1);
-}
-.tip-box {
-    background: linear-gradient(135deg, #f3e5f5, #e8eaf6);
-    padding: 10px;
-    border-radius: 6px;
-    margin: 8px 0;
-    border-left: 3px solid #673ab7;
-    color: #4a148c;
-    font-weight: 500;
-}
-button.primary {
-    background: linear-gradient(135deg, #673ab7, #8e24aa) !important;
-    border: none !important;
-    border-radius: 20px !important;
-    padding: 12px 25px !important;
-    color: white !important;
-    font-weight: bold !important;
-    font-size: 15px !important;
-    box-shadow: 0 5px 15px rgba(103, 58, 183, 0.4) !important;
-}
-button.primary:hover {
-    box-shadow: 0 8px 25px rgba(103, 58, 183, 0.6) !important;
-    opacity: 0.9 !important;
-    transform: translateY(-2px) !important;
-}
-label {
-    color: #4a148c !important;
-    font-weight: 600 !important;
-}
-input, textarea, select {
-    border: 1px solid rgba(103, 58, 183, 0.3) !important;
-    border-radius: 6px !important;
-}
-input:focus, textarea:focus, select:focus {
-    border-color: #673ab7 !important;
-    box-shadow: 0 0 0 2px rgba(103, 58, 183, 0.2) !important;
-}
-.gr-slider input[type="range"] {
-    accent-color: #673ab7 !important;
-}
-input[type="checkbox"] {
-    accent-color: #673ab7 !important;
-}
-.preserve-aspect-ratio img {
-    object-fit: contain !important;
-    width: auto !important;
-    max-height: 512px !important;
-}
-.social-links {
-    text-align: center;
-    margin: 20px 0;
-}
-.social-links a {
-    margin: 0 10px;
-    padding: 8px 16px;
-    background: #667eea;
-    color: white;
-    text-decoration: none;
-    border-radius: 8px;
-    transition: all 0.3s ease;
-}
-.social-links a:hover {
-    background: #764ba2;
-    transform: translateY(-2px);
-}
-.feature-box {
-    background: #f8fafc;
-    border: 1px solid #e2e8f0;
-    padding: 20px;
-    border-radius: 12px;
-    margin: 10px 0;
-  
-"""
-
-# ─────────  Gradio Blocks ─────────
-with gr.Blocks(css=custom_css, title="Jewellery Photography Preview") as demo:
-    # Hero
-    gr.HTML("""
-        <div style="text-align: center; margin-bottom: 20px;">
-            <h1 style="font-size: 2.5em;">🎨 Raresence: AI-Powered Jewellery Photo Preview</h1>
-            <p style="color: #666;">Upload a jewellery image, select model, and get professional photos instantly</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)
-
-    with gr.Column():
-            with gr.Row():
-                
-                with gr.Column(scale=0.4):
-                    gr.HTML("""
-                    <div class="feature-box"">
-                        <h3>🖼️ Upload Jewellery Image</h3>
-                        <p style="color: #666; font-size: 14px;">Select a clear jewellery image for best results</p>
-                    </div>
-                    """)
-                    gr.Markdown("‎")
-                    gr.Markdown("‎")
-                    input_img = gr.Image(label="Upload image", type="pil", height=400)
-                
-            # with gr.Column():
-                
-                with gr.Column():
-                    gr.HTML("""
-                    <div class="feature-box">
-                        <h3>🎨 AI Generated Results</h3>
-                        <p style="color: #666; font-size: 14px;">Preview overlay detection and final professional background</p>
-                    </div>
-                """)
-                    
-                    with gr.Tabs():
-                        with gr.TabItem("Final result"):
-                            info2 = gr.Markdown(value="### Final result")
-                            out_bg = gr.Image(height=400)
-                        with gr.TabItem("Detection overlay"):
-                            info1 = gr.Markdown(value="### Detection overlay")
-                            out_overlay = gr.Image(height=400)
-                    run_btn = gr.Button("🎯 Generate", elem_id="button", variant="primary")
-                
-            with gr.Row():
-                    with gr.Column(scale=0.4):
-                        gr.Markdown(value="Setting")
-                        category = gr.Dropdown(label="Jewellery category", choices=["Rings", "Bracelets", "Watches", "Earrings"], value="Bracelets")
-                        gender = gr.Dropdown(label="Model gender", choices=["male", "female"], value="female")
-            
-            
-            out_status = gr.Text(label="Status", interactive=False)
-        # ──────── Footer ────────
-    gr.HTML("""
-        <div style="text-align:center;padding:40px 20px;background:#f8fafc;border:1px solid #e2e8f0;border-radius:16px;margin:30px 0;">
-            <h3 style="color:#333;">🚀 Powered by Snapwear AI</h3>
-            <p style="color:#666;">
-                Experience the future of virtual fashion and garment visualization.
-            </p>
-            <div class="social-links">
-                <a href="https://snapwear.io" target="_blank">🌐 Website</a>
-                <a href="https://www.instagram.com/snapwearai/" target="_blank">📸 Instagram</a>
-                <a href="https://huggingface.co/spaces/SnapwearAI/Snapwear-Texture-Transfer" target="_blank">🎨 Pattern Transfer</a>
-            </div>
-            <p style="font-size:12px;color:#999;margin-top:20px;">
-                © 2024 Snapwear AI. Professional AI tools for fashion and design.
-            </p>
-        </div>
-    """)
-
-    # wire button → backend
-    run_btn.click(
-        fn=disable_button,
-        inputs=None,
-        outputs=run_btn
-    ).then(
-        fn=call_backend_with_retry,
-        inputs=[input_img, category, gender],
-        outputs=[out_overlay, out_bg, out_status,run_btn],
-        concurrency_limit=1,
-        show_progress=True,
-    )
-    
-
-# ─────────  Launch ─────────
-if __name__ == "__main__":
-    demo.queue(max_size=20, default_concurrency_limit=1).launch(share=False)
+def enable_button():
+    return gr.update(interactive=True)
+
+# In your button click event, update the outputs list:
+# wire button → backend
+run_btn.click(
+    fn=disable_button,
+    inputs=None,
+    outputs=run_btn
+).then(
+    fn=call_backend_with_retry,
+    inputs=[input_img, category, gender],
+    outputs=[out_overlay, out_bg, out_status, run_btn],  # Make sure run_btn is included here
+    concurrency_limit=1,
+    show_progress=True,
+)