app.py

# -*- coding: utf-8 -*-
import os
import time
import tempfile
from typing import Optional, Tuple, List, Dict

import gradio as gr
from PIL import Image
from gradio_client import Client, handle_file
from huggingface_hub import login

# ----------------------------
# Remote Space (IDM-VTON)
# ----------------------------
SPACE = "yisol/IDM-VTON"
API_NAME = "/tryon"

# ----------------------------
# Auth for company demo (no HF accounts needed)
# Secrets:
#   DEMO_USER=RVtest
#   DEMO_PASS=rv2026
# ----------------------------
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 = [
            f for f in os.listdir(GARMENT_DIR)
            if f.lower().endswith(ALLOWED_EXTS) and not f.startswith(".")
        ]
        files.sort()
        return files
    except Exception:
        return []


def garment_path(filename: str) -> str:
    return os.path.join(GARMENT_DIR, filename)


def load_garment_pil(filename: str) -> Optional[Image.Image]:
    if not filename:
        return None
    path = garment_path(filename)
    if not os.path.exists(path):
        return None
    try:
        return Image.open(path).convert("RGB")
    except Exception:
        return None


def build_gallery_items(files: List[str]):
    return [(garment_path(f), "") for f in files]


# ----------------------------
# HF token (optional)
# ----------------------------
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)
        print("HF login: OK")
    except Exception as e:
        print("HF login: FAILED:", str(e)[:200])
else:
    print("HF login: skipped (no token in env)")


# ----------------------------
# Helpers
# ----------------------------
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")
    return path


# ----------------------------
# Simple global rate limit (anti spam)
# NOTE: global across all users. Good enough for internal demo.
# ----------------------------
_last_call_ts = 0.0


def allow_call(min_interval_sec: float = 3.0) -> Tuple[bool, str]:
    global _last_call_ts
    now = time.time()
    if now - _last_call_ts < min_interval_sec:
        wait = max(0.0, min_interval_sec - (now - _last_call_ts))
        return False, f"⏳ Слишком часто. Подождите {wait:.1f} сек."
    _last_call_ts = now
    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, request: gr.Request):
    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, "❌ Не удалось загрузить выбранную одежду (проверьте garments/)"

    # Fixed params for simple demo
    garment_desc = "a photo of a garment"
    auto_mask = True
    crop_center = True
    denoise_steps = 25
    seed = 42

    p_path = save_pil_temp(person_pil)
    g_path = save_pil_temp(garment_pil)

    try:
        last_err = None

        for attempt in range(1, 7):
            try:
                client = make_client_from_request(request)

                result = client.predict(
                    dict={"background": handle_file(p_path), "layers": [], "composite": None},
                    garm_img=handle_file(g_path),
                    garment_des=garment_desc,
                    is_checked=bool(auto_mask),
                    is_checked_crop=bool(crop_center),
                    denoise_steps=int(denoise_steps),
                    seed=int(seed),
                    api_name=API_NAME,
                )

                if isinstance(result, (list, tuple)):
                    result = result[0]

                out = Image.open(result).convert("RGB")
                return out, "✅ Готово"

            except Exception as e:
                last_err = e
                msg_l = str(e).lower()

                is_timeout = (
                    "write operation timed out" in msg_l
                    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
                    or "too busy" in msg_l
                    or "overloaded" in msg_l
                    or "capacity" in msg_l
                )
                is_quota = "quota" in msg_l and "zerogpu" in msg_l

                # 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"
        return None, f"❌ Ошибка Space после 6 попыток: {tail}"

    finally:
        for path in (p_path, g_path):
            try:
                os.remove(path)
            except Exception:
                pass


# ----------------------------
# UI helpers
# ----------------------------
def refresh_catalog():
    files = list_garments()
    items = build_gallery_items(files)
    status = "✅ Каталог обновлён" if files else "⚠️ В папке garments/ пока нет изображений"
    return items, files, None, status


def on_gallery_select(files: List[str], evt: gr.SelectData):
    if not files:
        return None, "⚠️ Каталог пуст"
    try:
        idx = int(evt.index) if evt.index is not None else 0
        idx = max(0, min(idx, len(files) - 1))
        return files[idx], f"👕 Выбрано: {files[idx]}"
    except Exception:
        return None, "⚠️ Не удалось выбрать одежду"


# ----------------------------
# UI
# ----------------------------
CUSTOM_CSS = """
footer {display:none !important;}
#api-info {display:none !important;}
div[class*="footer"] {display:none !important;}
button[aria-label="Settings"] {display:none !important;}
"""

initial_files = list_garments()
initial_items = build_gallery_items(initial_files)

with gr.Blocks(title="Virtual Try-On Rendez-vous", css=CUSTOM_CSS) as demo:
    gr.Markdown("# Virtual Try-On Rendez-vous")

    garment_files_state = gr.State(initial_files)
    selected_garment_state = gr.State(None)

    with gr.Row():
        with gr.Column():
            person = gr.Image(label="Фото человека", type="pil", height=420)

            with gr.Row():
                refresh_btn = gr.Button("🔄 Обновить каталог", variant="secondary")
                selected_label = gr.Markdown("👕 Выберите одежду, кликнув по превью ниже")

            garment_gallery = gr.Gallery(
                label="Каталог одежды (кликните на превью)",
                value=initial_items,
                columns=4,
                height=340,
                allow_preview=True,
            )

            run = gr.Button("Примерить", variant="primary")
            status = gr.Textbox(value="Ожидание...", interactive=False)

        with gr.Column():
            out = gr.Image(label="Результат", type="pil", height=760)

    garment_gallery.select(
        fn=on_gallery_select,
        inputs=[garment_files_state],
        outputs=[selected_garment_state, selected_label],
    )

    refresh_btn.click(
        fn=refresh_catalog,
        inputs=[],
        outputs=[garment_gallery, garment_files_state, selected_garment_state, status],
    )

    run.click(
        fn=tryon_remote,
        inputs=[person, selected_garment_state],
        outputs=[out, status],
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=False,
        ssr_mode=False,
        auth=APP_AUTH,
    )