Update app.py

app.py

@@ -1,3 +1,81 @@
+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):
     
@@ -52,19 +130,240 @@ def call_backend_with_retry(input_image: Image.Image, category: str, gender: str
 def disable_button():
     return gr.update(interactive=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,
-)
+# ─────────  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)