Update app.py

app.py

@@ -2,7 +2,7 @@
 import os
 import time
 import tempfile
-from typing import Optional, Tuple
+from typing import Optional, Tuple, List
 
 import gradio as gr
 import spaces  # ZeroGPU runtime requires at least one @spaces.GPU-decorated function
@@ -17,7 +17,7 @@ SPACE = "yisol/IDM-VTON"
 API_NAME = "/tryon"
 
 # ----------------------------
-# Auth for company demo (no HF accounts needed)
+# Demo auth (no HF accounts needed)
 # Set these in HF Space Secrets:
 #   DEMO_USER=companydemo
 #   DEMO_PASS=your-strong-password
@@ -26,6 +26,23 @@ DEMO_USER = os.getenv("DEMO_USER", "").strip()
 DEMO_PASS = os.getenv("DEMO_PASS", "").strip()
 APP_AUTH = (DEMO_USER, DEMO_PASS) if (DEMO_USER and DEMO_PASS) else None
 
+# ----------------------------
+# Garment catalog folder in repo
+# ----------------------------
+GARMENT_DIR = "garments"
+ALLOWED_EXTS = (".png", ".jpg", ".jpeg", ".webp")
+
+def list_garments() -> List[str]:
+    try:
+        files = []
+        for f in os.listdir(GARMENT_DIR):
+            if f.lower().endswith(ALLOWED_EXTS) and not f.startswith("."):
+                files.append(f)
+        files.sort()
+        return files
+    except Exception:
+        return []
+
 # ----------------------------
 # HF token (optional)
 # ----------------------------
@@ -47,12 +64,10 @@ else:
 # ----------------------------
 _client: Optional[Client] = None
 
-
 def reset_client():
     global _client
     _client = None
 
-
 def get_client() -> Client:
     """
     gradio_client differs by version. Newer versions support hf_token=...
@@ -69,7 +84,6 @@ def get_client() -> Client:
             _client = Client(SPACE)
     return _client
 
-
 # ----------------------------
 # Helpers
 # ----------------------------
@@ -80,7 +94,6 @@ def clamp_int(x, lo, hi):
         x = lo
     return max(lo, min(hi, x))
 
-
 def save_pil_temp(pil_img: Image.Image, suffix: str = ".png") -> str:
     f = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
     path = f.name
@@ -88,6 +101,16 @@ def save_pil_temp(pil_img: Image.Image, suffix: str = ".png") -> str:
     pil_img.save(path, format="PNG")  # no resize/compress
     return path
 
+def load_garment_pil(filename: str) -> Optional[Image.Image]:
+    if not filename:
+        return None
+    path = os.path.join(GARMENT_DIR, filename)
+    if not os.path.exists(path):
+        return None
+    try:
+        return Image.open(path).convert("RGB")
+    except Exception:
+        return None
 
 # ----------------------------
 # Simple global rate limit (anti spam)
@@ -96,7 +119,6 @@ def save_pil_temp(pil_img: Image.Image, suffix: str = ".png") -> str:
 # ----------------------------
 _last_call_ts = 0.0
 
-
 def allow_call(min_interval_sec: float = 3.0) -> Tuple[bool, str]:
     global _last_call_ts
     now = time.time()
@@ -106,22 +128,25 @@ def allow_call(min_interval_sec: float = 3.0) -> Tuple[bool, str]:
     _last_call_ts = now
     return True, ""
 
-
 # ----------------------------
 # Core inference (remote call)
 # ZeroGPU: keep duration LOW to avoid burning quota
 # ----------------------------
 @spaces.GPU(duration=20)
-def tryon_remote(person_pil, garment_pil, garment_desc, auto_mask, crop_center, denoise_steps, seed):
-    # anti spam
+def tryon_remote(person_pil, garment_filename, garment_desc, auto_mask, crop_center, denoise_steps, seed):
     ok, msg = allow_call(3.0)
     if not ok:
         return None, msg
 
     if person_pil is None:
         return None, "❌ Загрузите фото человека"
+
+    if not garment_filename:
+        return None, "❌ Выберите одежду из списка"
+
+    garment_pil = load_garment_pil(garment_filename)
     if garment_pil is None:
-        return None, "❌ Загрузите одежду"
+        return None, "❌ Не удалось загрузить выбранную одежду (проверьте файл в папке garments/)"
 
     denoise_steps = clamp_int(denoise_steps, 10, 40)
     seed = clamp_int(seed, 0, 999999)
@@ -137,9 +162,6 @@ def tryon_remote(person_pil, garment_pil, garment_desc, auto_mask, crop_center,
         last_err = None
 
         for attempt in range(1, 7):
-            # Status per attempt (helps user confidence)
-            # We'll update status by returning it only on success/final fail,
-            # but we can still encode attempt info in the final error message.
             try:
                 client = get_client()
 
@@ -158,12 +180,11 @@ def tryon_remote(person_pil, garment_pil, garment_desc, auto_mask, crop_center,
                     result = result[0]
 
                 out = Image.open(result).convert("RGB")
-                return out, f"✅ Готово (steps={denoise_steps}, seed={seed}, crop={crop_center})"
+                return out, f"✅ Готово (одежда: {garment_filename}, steps={denoise_steps}, seed={seed})"
 
             except Exception as e:
                 last_err = e
-                msg = str(e)
-                msg_l = msg.lower()
+                msg_l = str(e).lower()
 
                 is_timeout = (
                     "write operation timed out" in msg_l
@@ -182,26 +203,13 @@ def tryon_remote(person_pil, garment_pil, garment_desc, auto_mask, crop_center,
 
                 is_expired = "expired zerogpu proxy token" in msg_l or "zerogpu proxy token" in msg_l
 
-                # Make retry reason human-friendly
-                if is_expired:
-                    reason = "истёк токен ZeroGPU"
-                elif is_busy:
-                    reason = "очередь/перегрузка"
-                elif is_timeout:
-                    reason = "таймаут сети"
-                else:
-                    reason = "ошибка"
-
-                # Retry on known transient errors
                 if is_timeout or is_busy or is_expired:
                     reset_client()
                     time.sleep(4.0 * attempt)
                     continue
 
-                # Unknown error: short backoff and continue retrying a few times anyway
                 time.sleep(1.2 * attempt)
 
-        # Final fail
         tail = str(last_err)[:240] if last_err else "unknown error"
         return None, f"❌ Ошибка Space после 6 попыток: {tail}"
 
@@ -212,11 +220,28 @@ def tryon_remote(person_pil, garment_pil, garment_desc, auto_mask, crop_center,
             except Exception:
                 pass
 
+# ----------------------------
+# UI helpers
+# ----------------------------
+def refresh_garments():
+    files = list_garments()
+    value = files[0] if files else None
+    preview = load_garment_pil(value) if value else None
+    status = "✅ Каталог обновлён" if files else "⚠️ В папке garments/ пока нет изображений"
+    return gr.Dropdown(choices=files, value=value), preview, status
 
-def reset_ui():
-    return None, None, "", "Ожидание...", None
+def on_select_garment(filename: str):
+    return load_garment_pil(filename)
 
+def reset_ui():
+    files = list_garments()
+    value = files[0] if files else None
+    preview = load_garment_pil(value) if value else None
+    return None, value, preview, "", "Ожидание...", None
 
+# ----------------------------
+# UI
+# ----------------------------
 CUSTOM_CSS = """
 footer {display:none !important;}
 #api-info {display:none !important;}
@@ -224,15 +249,28 @@ div[class*="footer"] {display:none !important;}
 button[aria-label="Settings"] {display:none !important;}
 """
 
+initial_files = list_garments()
+initial_value = initial_files[0] if initial_files else None
+initial_preview = load_garment_pil(initial_value) if initial_value else None
+
 with gr.Blocks(title="Virtual Try-On Rendez-vous", css=CUSTOM_CSS) as demo:
     gr.Markdown("# Virtual Try-On Rendez-vous")
 
     with gr.Row():
         with gr.Column():
             person = gr.Image(label="Фото человека", type="pil", height=420)
-            garment = gr.Image(label="Одежда", type="pil", height=320)
 
-            garment_desc = gr.Textbox(label="Описание одежды", value="")
+            with gr.Row():
+                garment_selector = gr.Dropdown(
+                    choices=initial_files,
+                    value=initial_value,
+                    label="Выберите одежду из каталога (garments/)"
+                )
+                refresh_btn = gr.Button("🔄 Обновить каталог", variant="secondary")
+
+            garment_preview = gr.Image(label="Предпросмотр одежды", type="pil", height=320)
+
+            garment_desc = gr.Textbox(label="Описание одежды (опционально)", value="")
 
             with gr.Accordion("Настройки", open=False):
                 auto_mask = gr.Checkbox(label="Auto-mask (Space)", value=True)
@@ -249,17 +287,32 @@ with gr.Blocks(title="Virtual Try-On Rendez-vous", css=CUSTOM_CSS) as demo:
         with gr.Column():
             out = gr.Image(label="Результат", type="pil", height=760)
 
-    # Better user feedback during processing (shows spinner and disables button)
+    # Preview updates
+    garment_selector.change(
+        fn=on_select_garment,
+        inputs=[garment_selector],
+        outputs=[garment_preview]
+    )
+
+    # Refresh catalog (after uploading new files)
+    refresh_btn.click(
+        fn=refresh_garments,
+        inputs=[],
+        outputs=[garment_selector, garment_preview, status]
+    )
+
+    # Run try-on
     run.click(
         fn=tryon_remote,
-        inputs=[person, garment, garment_desc, auto_mask, crop_center, denoise_steps, seed],
+        inputs=[person, garment_selector, garment_desc, auto_mask, crop_center, denoise_steps, seed],
         outputs=[out, status],
     )
 
+    # Reset UI
     reset.click(
         fn=reset_ui,
         inputs=[],
-        outputs=[person, garment, garment_desc, status, out],
+        outputs=[person, garment_selector, garment_preview, garment_desc, status, out],
     )
 
 if __name__ == "__main__":