Update app.py

app.py

@@ -2,7 +2,7 @@
 import os
 import time
 import tempfile
-from typing import Optional, Tuple, List
+from typing import Optional, Tuple, List, Dict
 
 import gradio as gr
 from PIL import Image
@@ -17,7 +17,7 @@ API_NAME = "/tryon"
 
 # ----------------------------
 # Auth for company demo (no HF accounts needed)
-# Set these in HF Space Secrets:
+# Secrets:
 #   DEMO_USER=RVtest
 #   DEMO_PASS=rv2026
 # ----------------------------
@@ -34,10 +34,10 @@ 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 = [
+            f for f in os.listdir(GARMENT_DIR)
+            if f.lower().endswith(ALLOWED_EXTS) and not f.startswith(".")
+        ]
         files.sort()
         return files
     except Exception:
@@ -61,7 +61,6 @@ def load_garment_pil(filename: str) -> Optional[Image.Image]:
 
 
 def build_gallery_items(files: List[str]):
-    # (image, caption). Caption empty = cleaner UI
     return [(garment_path(f), "") for f in files]
 
 
@@ -70,7 +69,6 @@ def build_gallery_items(files: List[str]):
 # ----------------------------
 HF_TOKEN = os.getenv("HF_TOKEN", "")
 print("HF_TOKEN set:", bool(HF_TOKEN), "len:", len(HF_TOKEN) if HF_TOKEN else 0)
-
 if HF_TOKEN:
     try:
         login(token=HF_TOKEN, add_to_git_credential=False)
@@ -81,50 +79,14 @@ else:
     print("HF login: skipped (no token in env)")
 
 
-# ----------------------------
-# Client caching
-# ----------------------------
-_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=...
-    Older versions don't. We fallback gracefully.
-    """
-    global _client
-    if _client is None:
-        try:
-            if HF_TOKEN:
-                _client = Client(SPACE, hf_token=HF_TOKEN)  # may raise TypeError on older versions
-            else:
-                _client = Client(SPACE)
-        except TypeError:
-            _client = Client(SPACE)
-    return _client
-
-
 # ----------------------------
 # Helpers
 # ----------------------------
-def clamp_int(x, lo, hi):
-    try:
-        x = int(x)
-    except Exception:
-        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
     f.close()
-    pil_img.save(path, format="PNG")  # no resize/compress
+    pil_img.save(path, format="PNG")
     return path
 
 
@@ -145,10 +107,33 @@ def allow_call(min_interval_sec: float = 3.0) -> Tuple[bool, str]:
     return True, ""
 
 
+def make_client_from_request(request: gr.Request) -> Client:
+    """
+    IMPORTANT for ZeroGPU Spaces:
+    Forward X-IP-Token so the downstream ZeroGPU Space applies rate limits/quota
+    per user correctly (instead of treating calls as unauthenticated).
+    """
+    headers: Dict[str, str] = {}
+    try:
+        # Gradio normalizes headers to lowercase keys
+        x_ip_token = request.headers.get("x-ip-token")
+        if x_ip_token:
+            headers["x-ip-token"] = x_ip_token
+    except Exception:
+        pass
+
+    # Some gradio_client versions accept headers=..., some may not. Fallback safely.
+    try:
+        return Client(SPACE, headers=headers) if headers else Client(SPACE)
+    except TypeError:
+        # older client: no headers kwarg
+        return Client(SPACE)
+
+
 # ----------------------------
 # Core inference (remote call)
 # ----------------------------
-def tryon_remote(person_pil, garment_filename):
+def tryon_remote(person_pil, garment_filename, request: gr.Request):
     ok, msg = allow_call(3.0)
     if not ok:
         return None, msg
@@ -178,7 +163,7 @@ def tryon_remote(person_pil, garment_filename):
 
         for attempt in range(1, 7):
             try:
-                client = get_client()
+                client = make_client_from_request(request)
 
                 result = client.predict(
                     dict={"background": handle_file(p_path), "layers": [], "composite": None},
@@ -206,7 +191,6 @@ def tryon_remote(person_pil, garment_filename):
                     or "read operation timed out" in msg_l
                     or "timed out" in msg_l
                 )
-
                 is_busy = (
                     "too many requests" in msg_l
                     or "queue" in msg_l
@@ -214,14 +198,19 @@ def tryon_remote(person_pil, garment_filename):
                     or "overloaded" in msg_l
                     or "capacity" in msg_l
                 )
+                is_quota = "quota" in msg_l and "zerogpu" in msg_l
 
-                is_expired = "expired zerogpu proxy token" in msg_l or "zerogpu proxy token" in msg_l
-
-                if is_timeout or is_busy or is_expired:
-                    reset_client()
+                # Retry on transient issues; quota will likely not improve immediately
+                if is_timeout or is_busy:
                     time.sleep(4.0 * attempt)
                     continue
 
+                if is_quota:
+                    return None, (
+                        "⚠️ Лимит ZeroGPU на стороне модели исчерпан для текущего пользователя.\n"
+                        "Попробуйте позже или используйте меньше попыток подряд."
+                    )
+
                 time.sleep(1.2 * attempt)
 
         tail = str(last_err)[:240] if last_err else "unknown error"
@@ -297,21 +286,18 @@ 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)
 
-    # Update selection on click
     garment_gallery.select(
         fn=on_gallery_select,
         inputs=[garment_files_state],
         outputs=[selected_garment_state, selected_label],
     )
 
-    # Refresh catalog after uploading new garments
     refresh_btn.click(
         fn=refresh_catalog,
         inputs=[],
         outputs=[garment_gallery, garment_files_state, selected_garment_state, status],
     )
 
-    # Run try-on
     run.click(
         fn=tryon_remote,
         inputs=[person, selected_garment_state],
@@ -325,5 +311,5 @@ if __name__ == "__main__":
         share=False,
         debug=False,
         ssr_mode=False,
-        auth=APP_AUTH,  # ✅ login/password gate
+        auth=APP_AUTH,
     )