Update app.py

app.py

@@ -1,252 +1,449 @@
-# -*- coding: utf-8 -*-
-import os
-import time
-import tempfile
-import gradio as gr
-from PIL import Image, ImageFilter, ImageEnhance
-from gradio_client import Client, handle_file
-from huggingface_hub import login
-
-SPACE = "yisol/IDM-VTON"
-API_NAME = "/tryon"
-
-HF_TOKEN = "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"
-
-print("HF_TOKEN set:", bool(HF_TOKEN), "len:", len(HF_TOKEN) if HF_TOKEN else 0)
-
-if HF_TOKEN and HF_TOKEN != "PASTE_YOUR_HF_TOKEN_HERE":
-    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 (token placeholder or empty)")
-
-_client = None
-
-
-def get_client():
-    global _client
-    if _client is None:
-        _client = Client(SPACE)
-    return _client
-
-
-def pil_save_temp(pil_img: Image.Image, suffix: str) -> str:
-    f = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
-    path = f.name
-    f.close()
-    pil_img.save(path)
-    return path
-
-
-def clamp_int(x, lo, hi):
-    try:
-        x = int(x)
-    except Exception:
-        x = lo
-    return max(lo, min(hi, x))
-
-
-def postprocess_text_boost(img: Image.Image, boost: int) -> Image.Image:
-    boost = clamp_int(boost, 0, 100)
-    if boost <= 0:
-        return img
-
-    t = boost / 100.0
-    out = img.convert("RGB")
-
-    out = out.filter(
-        ImageFilter.UnsharpMask(
-            radius=1.2 + 1.0 * t,
-            percent=int(120 + 80 * t),
-            threshold=2
-        )
-    )
-    out = ImageEnhance.Contrast(out).enhance(1.0 + 0.10 * t)
-    out = ImageEnhance.Sharpness(out).enhance(1.0 + 0.30 * t)
-
-    w, h = out.size
-    box = (int(w * 0.25), int(h * 0.22), int(w * 0.75), int(h * 0.72))
-
-    region = out.crop(box)
-    region = region.filter(
-        ImageFilter.UnsharpMask(
-            radius=1.8 + 1.6 * t,
-            percent=int(160 + 120 * t),
-            threshold=2
-        )
-    )
-    region = ImageEnhance.Contrast(region).enhance(1.0 + 0.15 * t)
-    region = ImageEnhance.Sharpness(region).enhance(1.0 + 0.45 * t)
-
-    blended = Image.blend(out.crop(box), region, alpha=min(0.55 + 0.35 * t, 0.9))
-    out.paste(blended, box)
-    return out
-
-
-def apply_mode_defaults(mode: str, denoise_steps: int, seed: int, crop_center: bool, garment_desc: str):
-    denoise_steps = clamp_int(denoise_steps, 10, 40)
-    seed = clamp_int(seed, 0, 999999)
-
-    if mode == "Balanced":
-        denoise_steps = max(denoise_steps, 28)
-    elif mode == "Quality":
-        denoise_steps = max(denoise_steps, 38)
-    elif mode == "Text/Logo":
-        denoise_steps = 40
-        crop_center = False
-        if seed == 0:
-            seed = 42
-        if garment_desc and "text" not in garment_desc.lower() and "logo" not in garment_desc.lower():
-            garment_desc = garment_desc.strip() + ", with clear text/logo print"
-
-    return denoise_steps, seed, crop_center, garment_desc
-
-
-def tryon_remote(
-    person_pil,
-    garment_pil,
-    garment_desc,
-    auto_mask,
-    crop_center,
-    denoise_steps,
-    seed,
-    mode,
-    text_boost,
-):
-    if person_pil is None:
-        return None, "❌ Загрузите фото человека"
-    if garment_pil is None:
-        return None, "❌ Загрузите одежду"
-
-    denoise_steps, seed, crop_center, garment_desc = apply_mode_defaults(
-        mode=mode,
-        denoise_steps=denoise_steps,
-        seed=seed,
-        crop_center=crop_center,
-        garment_desc=garment_desc or "a t-shirt"
-    )
-
-    p_path = pil_save_temp(person_pil.convert("RGB"), ".png")
-    g_path = pil_save_temp(garment_pil.convert("RGB"), ".png")
-
-    try:
-        client = get_client()
-
-        last_err = None
-        for attempt in range(1, 4):
-            try:
-                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")
-                out = postprocess_text_boost(out, text_boost)
-
-                return out, f"✅ Готово (mode={mode}, steps={denoise_steps}, seed={seed}, crop={crop_center})"
-
-            except Exception as e:
-                last_err = e
-                msg = str(e)
-
-                if "ZeroGPU" in msg or "Unlogged user" in msg or "quota" in msg.lower():
-                    return None, (
-                        "❌ ZeroGPU-квота/нагрузка Hugging Face.\n"
-                        "Если HF login = OK, то это лимит/очередь Space. Попробуйте позже или смените seed."
-                    )
-
-                time.sleep(1.2 * attempt)
-
-        return None, f"❌ Ошибка Space: {str(last_err)[:250]}"
-
-    finally:
-        for path in (p_path, g_path):
-            try:
-                os.remove(path)
-            except Exception:
-                pass
-
-
-def reset_ui():
-    return None, None, None, "Ожидание..."
-
-
-CUSTOM_CSS = """
-footer {display:none !important;}
-#api-info {display:none !important;}
-div[class*="footer"] {display:none !important;}
-button[aria-label="Settings"] {display:none !important;}
-"""
-
-with gr.Blocks(title="Virtual Try-On Rendez-vous", css=CUSTOM_CSS) as demo:
-    gr.Markdown(
-        "# Virtual Try-On Rendez-vous\n\n"
-        "**Для текста/логотипов:** режим *Text/Logo*, steps=40, попробуйте разные *seed*."
-    )
-
-    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="a t-shirt with clear text/logo print"
-            )
-
-            mode = gr.Radio(
-                ["Balanced", "Quality", "Text/Logo"],
-                value="Text/Logo",
-                label="Режим качества"
-            )
-
-            with gr.Accordion("Настройки", open=False):
-                auto_mask = gr.Checkbox(label="Auto-mask", value=True)
-                crop_center = gr.Checkbox(label="Crop по центру", value=True)
-                denoise_steps = gr.Slider(10, 40, value=28, step=1, label="Denoise steps")
-                seed = gr.Slider(0, 999999, value=42, step=1, label="Seed")
-                text_boost = gr.Slider(0, 100, value=55, step=1, label="Text boost (резкость/контраст принта)")
-
-            with gr.Row():
-                run = gr.Button("Примерить", variant="primary")
-                reset = gr.Button("Сбросить", variant="secondary")
-
-            status = gr.Textbox(value="Ожидание...", interactive=False)
-
-        with gr.Column():
-            out = gr.Image(label="Результат", type="pil", height=760)
-
-    run.click(
-        fn=tryon_remote,
-        inputs=[person, garment, garment_desc, auto_mask, crop_center, denoise_steps, seed, mode, text_boost],
-        outputs=[out, status]
-    )
-
-    reset.click(
-        fn=reset_ui,
-        inputs=[],
-        outputs=[person, garment, out, status]
-    )
-
-if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7863,
-        share=False
-    )
+# -*- coding: utf-8 -*-
+import os
+import time
+import tempfile
+
+import gradio as gr
+import numpy as np
+from PIL import Image, ImageFilter, ImageEnhance
+
+from gradio_client import Client, handle_file
+from huggingface_hub import login
+
+SPACE = "yisol/IDM-VTON"
+API_NAME = "/tryon"
+
+# ✅ IMPORTANT: do NOT hardcode tokens in a Space repo
+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)")
+
+_client = None
+
+
+def reset_client():
+    global _client
+    _client = None
+
+
+def get_client():
+    """
+    Keep it compatible with old gradio_client (no hf_token=...).
+    We rely on huggingface_hub login() above.
+    """
+    global _client
+    if _client is None:
+        _client = Client(SPACE)
+    return _client
+
+
+# ---------------- IO helpers ----------------
+
+def clamp_int(x, lo, hi):
+    try:
+        x = int(x)
+    except Exception:
+        x = lo
+    return max(lo, min(hi, x))
+
+
+def clamp_float(x, lo, hi):
+    try:
+        x = float(x)
+    except Exception:
+        x = lo
+    return max(lo, min(hi, x))
+
+
+def pil_save_temp(
+    pil_img: Image.Image,
+    suffix: str = ".jpg",
+    max_side: int = 1400,
+    quality: int = 92
+) -> str:
+    """
+    Reduce upload timeouts:
+    - convert to RGB
+    - downscale to max_side
+    - save JPEG optimized
+    """
+    img = pil_img.convert("RGB")
+    w, h = img.size
+    m = max(w, h)
+    if m > max_side:
+        scale = max_side / m
+        img = img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
+
+    f = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
+    path = f.name
+    f.close()
+    img.save(path, format="JPEG", quality=int(quality), optimize=True)
+    return path
+
+
+# ---------------- Quality / Mode ----------------
+
+def apply_mode_defaults(mode: str, denoise_steps: int, seed: int, crop_center: bool, garment_desc: str):
+    denoise_steps = clamp_int(denoise_steps, 10, 40)
+    seed = clamp_int(seed, 0, 999999)
+
+    if mode == "Balanced":
+        denoise_steps = max(denoise_steps, 28)
+    elif mode == "Quality":
+        denoise_steps = max(denoise_steps, 38)
+    elif mode == "Text/Logo":
+        denoise_steps = 40
+        crop_center = False
+        if seed == 0:
+            seed = 42
+        if garment_desc and "text" not in garment_desc.lower() and "logo" not in garment_desc.lower():
+            garment_desc = garment_desc.strip() + ", with clear text/logo print"
+
+    return denoise_steps, seed, crop_center, garment_desc
+
+
+# ---------------- Overlay (preserve print/texture) ----------------
+
+def _to_np(img: Image.Image) -> np.ndarray:
+    return np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
+
+
+def _to_pil(arr: np.ndarray) -> Image.Image:
+    arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8)
+    return Image.fromarray(arr, mode="RGB")
+
+
+def make_auto_mask_from_garment(garment_rgb: Image.Image) -> Image.Image:
+    """
+    Heuristic mask for print/text: edge density + center focus.
+    """
+    g = garment_rgb.convert("L")
+    edges = g.filter(ImageFilter.FIND_EDGES).filter(ImageFilter.GaussianBlur(radius=1.2))
+    arr = np.asarray(edges, dtype=np.uint8)
+
+    m = float(arr.mean())
+    s = float(arr.std())
+    thr = int(min(255, max(25, m + 0.7 * s)))
+
+    mask = (arr > thr).astype(np.uint8) * 255
+    mask_img = Image.fromarray(mask, mode="L").filter(ImageFilter.GaussianBlur(radius=2.0))
+
+    w, h = mask_img.size
+    focus = Image.new("L", (w, h), 0)
+    fx0, fy0, fx1, fy1 = int(w * 0.18), int(h * 0.12), int(w * 0.82), int(h * 0.88)
+    focus.paste(255, (fx0, fy0, fx1, fy1))
+    focus = focus.filter(ImageFilter.GaussianBlur(radius=12.0))
+
+    mask_arr = (np.asarray(mask_img, dtype=np.float32) / 255.0) * (np.asarray(focus, dtype=np.float32) / 255.0)
+    return Image.fromarray(np.clip(mask_arr * 255.0, 0, 255).astype(np.uint8), mode="L")
+
+
+def apply_overlay_preserve(
+    out_rgb: Image.Image,
+    garment_rgb: Image.Image,
+    strength: float,
+    detail_strength: float,
+    box_scale: float,
+    offset_x: int,
+    offset_y: int,
+    use_drawn_mask: bool,
+    garment_mask_draw: Image.Image | None,
+) -> Image.Image:
+    """
+    Blend garment colors + high-frequency detail into torso region.
+    If a drawn mask is provided, it wins (best for exact logos/text).
+    """
+    strength = clamp_float(strength, 0.0, 1.0)
+    detail_strength = clamp_float(detail_strength, 0.0, 2.0)
+    box_scale = clamp_float(box_scale, 0.6, 1.4)
+
+    if strength <= 0.0 and detail_strength <= 0.0:
+        return out_rgb
+
+    out = out_rgb.convert("RGB")
+    g = garment_rgb.convert("RGB")
+
+    ow, oh = out.size
+
+    # approximate torso box
+    x0 = int(ow * 0.22)
+    x1 = int(ow * 0.78)
+    y0 = int(oh * 0.20)
+    y1 = int(oh * 0.78)
+
+    bw, bh = (x1 - x0), (y1 - y0)
+
+    cx = (x0 + x1) // 2 + int(offset_x)
+    cy = (y0 + y1) // 2 + int(offset_y)
+    bw2 = int(bw * box_scale)
+    bh2 = int(bh * box_scale)
+
+    x0b = max(0, cx - bw2 // 2)
+    x1b = min(ow, cx + bw2 // 2)
+    y0b = max(0, cy - bh2 // 2)
+    y1b = min(oh, cy + bh2 // 2)
+
+    box_w, box_h = (x1b - x0b), (y1b - y0b)
+    if box_w <= 10 or box_h <= 10:
+        return out
+
+    g_fit = g.resize((box_w, box_h), Image.LANCZOS)
+
+    # mask
+    if use_drawn_mask and garment_mask_draw is not None:
+        gm = garment_mask_draw
+        if gm.mode == "RGBA":
+            mask = gm.split()[-1]
+        else:
+            mask = gm.convert("L")
+        mask = mask.resize((box_w, box_h), Image.LANCZOS).filter(ImageFilter.GaussianBlur(radius=2.0))
+    else:
+        mask = make_auto_mask_from_garment(g).resize((box_w, box_h), Image.LANCZOS)
+
+    m = np.asarray(mask, dtype=np.float32) / 255.0
+    m = np.clip(m, 0.0, 1.0)
+
+    out_arr = _to_np(out)
+    crop = out_arr[y0b:y1b, x0b:x1b, :]
+
+    g_arr = _to_np(g_fit)
+
+    # 1) color/print transfer
+    if strength > 0:
+        alpha = strength * (m[..., None])
+        crop = crop * (1.0 - alpha) + g_arr * alpha
+
+    # 2) texture/detail transfer
+    if detail_strength > 0:
+        g_blur = _to_np(g_fit.filter(ImageFilter.GaussianBlur(radius=3.0)))
+        detail = g_arr - g_blur
+        alpha2 = detail_strength * (m[..., None]) * 0.65
+        crop = np.clip(crop + detail * alpha2, 0.0, 1.0)
+
+    out_arr[y0b:y1b, x0b:x1b, :] = crop
+    return _to_pil(out_arr)
+
+
+# ---------------- Main inference ----------------
+
+def tryon_remote(
+    person_pil,
+    garment_pil,
+    garment_desc,
+    auto_mask,
+    crop_center,
+    denoise_steps,
+    seed,
+    mode_quality,
+    overlay_enable,
+    overlay_strength,
+    overlay_detail,
+    overlay_box_scale,
+    overlay_offset_x,
+    overlay_offset_y,
+    use_garment_mask_draw,
+    garment_mask_draw,
+):
+    if person_pil is None:
+        return None, "❌ Загрузите фото человека"
+    if garment_pil is None:
+        return None, "❌ Загрузите одежду"
+
+    denoise_steps, seed, crop_center, garment_desc = apply_mode_defaults(
+        mode=mode_quality,
+        denoise_steps=denoise_steps,
+        seed=seed,
+        crop_center=crop_center,
+        garment_desc=garment_desc or "a t-shirt"
+    )
+
+    # Save smaller for upload stability (still good enough for text with overlay)
+    p_path = pil_save_temp(person_pil, ".jpg", max_side=1400, quality=92)
+    g_path = pil_save_temp(garment_pil, ".jpg", max_side=1400, quality=95)
+
+    try:
+        last_err = None
+
+        # more attempts to survive ZeroGPU queue
+        for attempt in range(1, 7):
+            try:
+                client = get_client()
+
+                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")
+
+                if overlay_enable:
+                    out = apply_overlay_preserve(
+                        out_rgb=out,
+                        garment_rgb=garment_pil.convert("RGB"),
+                        strength=float(overlay_strength),
+                        detail_strength=float(overlay_detail),
+                        box_scale=float(overlay_box_scale),
+                        offset_x=int(overlay_offset_x),
+                        offset_y=int(overlay_offset_y),
+                        use_drawn_mask=bool(use_garment_mask_draw),
+                        garment_mask_draw=garment_mask_draw,
+                    )
+
+                return out, (
+                    f"✅ Готово | mode={mode_quality}, steps={denoise_steps}, seed={seed}, crop={crop_center} | "
+                    f"overlay={'on' if overlay_enable else 'off'}"
+                )
+
+            except Exception as e:
+                last_err = e
+                msg = str(e).lower()
+
+                is_timeout = ("write operation timed out" in msg) or ("timed out" in msg)
+                is_queue = ("queue" in msg) or ("too busy" in msg) or ("overloaded" in msg) or ("capacity" in msg)
+                is_zerogpu = ("zerogpu" in msg)
+
+                # Handle queue/timeouts: reset client and wait longer
+                if is_timeout or is_queue or is_zerogpu:
+                    reset_client()
+
+                    # escalate waiting
+                    wait_s = 6.0 * attempt  # 6s, 12s, 18s, ...
+                    # on later attempts, compress more to reduce upload chance of timeout
+                    if attempt >= 3:
+                        try:
+                            os.remove(p_path); os.remove(g_path)
+                        except Exception:
+                            pass
+                        p_path = pil_save_temp(person_pil, ".jpg", max_side=1100, quality=85)
+                        g_path = pil_save_temp(garment_pil, ".jpg", max_side=1100, quality=88)
+
+                    time.sleep(wait_s)
+                    continue
+
+                # generic backoff
+                time.sleep(1.2 * attempt)
+
+        return None, f"❌ Ошибка Space: {str(last_err)[:250]}"
+
+    finally:
+        for path in (p_path, g_path):
+            try:
+                os.remove(path)
+            except Exception:
+                pass
+
+
+def reset_ui():
+    return None, None, 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;}
+"""
+
+with gr.Blocks(title="Virtual Try-On Rendez-vous") as demo:
+    gr.Markdown(
+        "# Virtual Try-On Rendez-vous\n\n"
+        "**Максимум по тексту/фактуре:** выбери *Text/Logo* и включи *Overlay*.\n"
+        "Если авто-маска промахивается — включи *Use garment mask (draw)* и нарисуй маску по тексту/логотипу."
+    )
+
+    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="a t-shirt with clear text/logo print"
+            )
+
+            mode_quality = gr.Radio(
+                ["Balanced", "Quality", "Text/Logo"],
+                value="Text/Logo",
+                label="Режим качества"
+            )
+
+            with gr.Accordion("Настройки генерации", open=False):
+                auto_mask = gr.Checkbox(label="Auto-mask (Space)", value=True)
+                crop_center = gr.Checkbox(label="Crop по центру", value=True)
+                denoise_steps = gr.Slider(10, 40, value=28, step=1, label="Denoise steps")
+                seed = gr.Slider(0, 999999, value=42, step=1, label="Seed")
+
+            with gr.Accordion("Сохранение текста/текстуры (Overlay)", open=True):
+                overlay_enable = gr.Checkbox(label="Enable Overlay", value=True)
+                overlay_strength = gr.Slider(0, 1, value=0.55, step=0.01, label="Overlay strength (цвет/принт)")
+                overlay_detail = gr.Slider(0, 2, value=1.10, step=0.05, label="Overlay detail (фактура/буквы)")
+                overlay_box_scale = gr.Slider(0.6, 1.4, value=1.0, step=0.02, label="Overlay box scale")
+                overlay_offset_x = gr.Slider(-300, 300, value=0, step=1, label="Overlay offset X (px)")
+                overlay_offset_y = gr.Slider(-300, 300, value=0, step=1, label="Overlay offset Y (px)")
+
+                use_garment_mask_draw = gr.Checkbox(
+                    label="Use garment mask (draw) — нарисуй маску по тексту/принту",
+                    value=False
+                )
+                garment_mask_draw = gr.Image(
+                    label="Garment mask (опционально): нарисуй белым по тексту/принту",
+                    type="pil",
+                    height=260,
+                )
+
+            with gr.Row():
+                run = gr.Button("Примерить", variant="primary")
+                reset = gr.Button("Сбросить", variant="secondary")
+
+            status = gr.Textbox(value="Ожидание...", interactive=False)
+
+        with gr.Column():
+            out = gr.Image(label="Результат", type="pil", height=760)
+
+    run.click(
+        fn=tryon_remote,
+        inputs=[
+            person, garment, garment_desc,
+            auto_mask, crop_center, denoise_steps, seed, mode_quality,
+            overlay_enable, overlay_strength, overlay_detail, overlay_box_scale,
+            overlay_offset_x, overlay_offset_y,
+            use_garment_mask_draw, garment_mask_draw
+        ],
+        outputs=[out, status]
+    )
+
+    reset.click(
+        fn=reset_ui,
+        inputs=[],
+        outputs=[person, garment, out, status]
+    )
+
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,   # Spaces default
+        share=False,
+        debug=False,
+        css=CUSTOM_CSS,
+        ssr_mode=False
+    )