Update app.py

app.py

@@ -1,30 +1,12 @@
 # -*- coding: utf-8 -*-
-"""
-Virtual Try-On Rendez-vous — production wrapper for IDM-VTON (SDXL)
-
-Что изменено по твоему запросу (убрано/исправлено):
-1) НЕТ “жёстко upper_body для всего” — маска выбирается АВТО по имени/папке одежды (dress/lower/upper),
-   либо можно отключить авто-маску полностью.
-2) НЕТ fixed strength=1.0 — strength настраиваемый (по умолчанию 0.9).
-3) НЕТ фиксированных промптов “a garment” — промпт генерируется из имени файла/папки одежды + эвристики,
-   можно переопределить вручную.
-4) НЕТ crop-center + paste обратно — используется letterbox (масштаб с сохранением пропорций + padding),
-   затем padding убирается, и результат возвращается в исходный размер.
-5) НЕТ принудительного 768×1024 “всегда” — размер выбирается ДИНАМИЧЕСКИ от входного фото (с ограничением max_side),
-   кратно 8.
-6) НЕТ низких/фиксированных CFG/steps/seed — все параметры управляемые в UI; seed может быть -1 (рандом).
-
-Остальное (датасет одежды, галерея, queue, patch gradio_client) оставлено как инфраструктура.
-"""
 import os
 import re
 import time
-import math
-from typing import List, Optional, Tuple, Dict, Any
+from typing import List, Optional, Tuple, Dict
 
 import spaces
 import gradio as gr
-from PIL import Image, ImageOps
+from PIL import Image
 
 # =========================
 # FIX: gradio 4.24 / gradio_client crashes on boolean JSON Schemas in /api_info
@@ -116,14 +98,11 @@ APP_AUTH = (DEMO_USER, DEMO_PASS) if (DEMO_USER and DEMO_PASS) else None
 # =========================
 GARMENT_DIR = "garments"
 ALLOWED_EXTS = (".png", ".jpg", ".jpeg", ".webp")
-GARMENTS_DATASET = os.getenv("GARMENTS_DATASET", "").strip()  # e.g. "ArmanRV/armanrv-garments"
+GARMENTS_DATASET = os.getenv("GARMENTS_DATASET", "").strip()
 HF_TOKEN = os.getenv("HF_TOKEN", "").strip()
 
 
 def ensure_garments_downloaded() -> None:
-    """
-    Downloads garments from HF Dataset into ./garments to avoid Space repo 1GB limit.
-    """
     os.makedirs(GARMENT_DIR, exist_ok=True)
 
     if HF_TOKEN:
@@ -151,9 +130,6 @@ def ensure_garments_downloaded() -> None:
 
 
 def list_garments() -> List[str]:
-    """
-    Recursively list images inside ./garments (handles dataset subfolders).
-    """
     files: List[str] = []
     if not os.path.isdir(GARMENT_DIR):
         return files
@@ -183,7 +159,6 @@ def load_garment_pil(filename: str) -> Optional[Image.Image]:
 
 
 def build_gallery_items(files: List[str]):
-    # Gallery items format: [(filepath, caption), ...]
     return [(garment_path(f), "") for f in files]
 
 
@@ -225,17 +200,15 @@ def round_to_multiple(x: int, m: int = 8) -> int:
 
 def pick_target_size_keep_aspect(w: int, h: int, max_side: int) -> Tuple[int, int]:
     """
-    Возвращает (tw, th) <= max_side по большей стороне, кратно 8.
+    (tw, th) <= max_side по большей стороне, кратно 8
     """
     if w <= 0 or h <= 0:
         return 768, 1024
     scale = min(max_side / float(max(w, h)), 1.0)
     tw = round_to_multiple(int(w * scale), 8)
     th = round_to_multiple(int(h * scale), 8)
-    # защитимся от слишком маленьких
     tw = max(512, tw)
     th = max(512, th)
-    # еще раз не превышать max_side
     if max(tw, th) > max_side:
         scale2 = max_side / float(max(tw, th))
         tw = round_to_multiple(int(tw * scale2), 8)
@@ -243,10 +216,10 @@ def pick_target_size_keep_aspect(w: int, h: int, max_side: int) -> Tuple[int, in
     return tw, th
 
 
-def letterbox(img: Image.Image, target_w: int, target_h: int, fill=(0, 0, 0)) -> Tuple[Image.Image, Dict[str, int]]:
+def letterbox(img: Image.Image, target_w: int, target_h: int, fill=(127, 127, 127)) -> Tuple[Image.Image, Dict[str, int]]:
     """
-    Масштабирует с сохранением пропорций + padding до target_w/target_h.
-    Возвращает (img_lb, meta) где meta содержит offset/size для обратного unletterbox.
+    Resize with aspect + padding to (target_w,target_h).
+    meta: x,y,w,h for core region inside padded canvas
     """
     src_w, src_h = img.size
     if src_w <= 0 or src_h <= 0:
@@ -262,45 +235,44 @@ def letterbox(img: Image.Image, target_w: int, target_h: int, fill=(0, 0, 0)) ->
     x = (target_w - new_w) // 2
     y = (target_h - new_h) // 2
     canvas.paste(img_rs, (x, y))
-    meta = {"x": x, "y": y, "w": new_w, "h": new_h, "src_w": src_w, "src_h": src_h}
-    return canvas, meta
+    return canvas, {"x": x, "y": y, "w": new_w, "h": new_h, "src_w": src_w, "src_h": src_h}
 
 
 def unletterbox(img_lb: Image.Image, meta: Dict[str, int]) -> Image.Image:
+    x, y, w, h = meta["x"], meta["y"], meta["w"], meta["h"]
+    return img_lb.crop((x, y, x + w, y + h))
+
+
+def paste_into_canvas(canvas_mode: str, canvas_size: Tuple[int, int], core_img: Image.Image, meta: Dict[str, int], fill):
     """
-    Вырезает область без padding и возвращает как есть (потом можно resize к исходнику).
+    Вклеивает core_img в канвас (target_w,target_h) по meta x,y.
     """
     x, y, w, h = meta["x"], meta["y"], meta["w"], meta["h"]
-    return img_lb.crop((x, y, x + w, y + h))
+    canvas = Image.new(canvas_mode, canvas_size, fill)
+    if core_img.size != (w, h):
+        core_img = core_img.resize((w, h), Image.BILINEAR)
+    canvas.paste(core_img, (x, y))
+    return canvas
 
 
 def infer_garment_class_from_path(relpath: str) -> str:
     """
-    Возвращает тип для get_mask_location: 'upper_body' | 'lower_body' | 'dresses'
-    Это НЕ “жестко upper_body” — эвристика по папке/имени.
+    'upper_body' | 'lower_body' | 'dresses'
     """
-    s = (relpath or "").lower()
-    # папки/имена под платья
-    if any(k in s for k in ["dress", "dresses", "suk", "plate", "плать", "sarafan"]):
+    s = (relpath or "").lower().replace("\\", "/")
+    if any(k in s for k in ["dress", "dresses", "sarafan", "plate", "плать", "сараф"]):
         return "dresses"
-    # низ
     if any(k in s for k in ["pants", "trouser", "jeans", "skirt", "short", "брюк", "джин", "юбк", "шорт"]):
         return "lower_body"
-    # верх по умолчанию
     return "upper_body"
 
 
 def guess_garment_description(relpath: str) -> str:
-    """
-    Генерирует более полезное текстовое описание одежды из имени файла/папки.
-    (Это замена твоего фиксированного 'a garment'.)
-    """
-    s = (relpath or "").replace("\\", "/").lower()
-    # словарь эвристик
+    s = (relpath or "").lower().replace("\\", "/")
     mapping = [
-        (["shearling", "dub", "дублен", "sheepskin"], "a shearling jacket"),
+        (["shearling", "дублен", "sheepskin"], "a shearling jacket"),
         (["coat", "пальт", "overcoat"], "a coat"),
-        (["jacket", "куртк", "bomber", "парка", "parka"], "a jacket"),
+        (["jacket", "куртк", "парка", "parka", "bomber"], "a jacket"),
         (["blazer", "пидж", "suit"], "a blazer"),
         (["hoodie", "худи"], "a hoodie"),
         (["sweater", "свит", "jumper"], "a sweater"),
@@ -314,18 +286,56 @@ def guess_garment_description(relpath: str) -> str:
         if any(k in s for k in keys):
             return desc
 
-    # иначе — попытка вытащить “человеческое” имя
     base = os.path.splitext(os.path.basename(s))[0]
     base = re.sub(r"[_\-]+", " ", base)
     base = re.sub(r"\d+", " ", base)
     base = re.sub(r"\s+", " ", base).strip()
     if len(base) >= 3:
-        # ограничим длину
-        words = base.split()[:4]
-        return "a " + " ".join(words)
+        return "a " + " ".join(base.split()[:4])
     return "a piece of clothing"
 
 
+def apply_safety_clamp(mask_full: Image.Image, meta: Dict[str, int], garment_class: str, clamp_strength: float) -> Image.Image:
+    """
+    Универсальная страховка от “уехало вниз/вверх”:
+    - upper_body: оставляем маску выше линии бёдер (чем больше clamp_strength, тем “выше” граница)
+    - lower_body: оставляем маску ниже линии талии/бёдер (чем больше clamp_strength, тем “ниже” граница)
+    - dresses: не трогаем
+
+    clamp_strength: 0..1 (0 = почти не влияет, 1 = сильнее)
+    """
+    if garment_class == "dresses":
+        return mask_full
+
+    tw, th = mask_full.size
+    x, y, w, h = meta["x"], meta["y"], meta["w"], meta["h"]
+
+    # базовые линии (проценты по core высоте) — эмпирика для full-body
+    # upper_body: граница где-то около 0.60..0.72 от высоты core
+    # lower_body: граница около 0.34..0.48 от высоты core
+    clamp_strength = clamp_float(clamp_strength, 0.0, 1.0)
+
+    if garment_class == "upper_body":
+        lo, hi = 0.60, 0.72
+        frac = lo + (hi - lo) * (1.0 - clamp_strength)  # clamp_strength↑ => граница ближе к lo (выше)
+        cut_y = y + int(frac * h)
+        keep = mask_full.crop((0, 0, tw, max(0, min(th, cut_y))))
+        out = Image.new("L", (tw, th), 0)
+        out.paste(keep, (0, 0))
+        return out
+
+    if garment_class == "lower_body":
+        lo, hi = 0.34, 0.48
+        frac = lo + (hi - lo) * (clamp_strength)  # clamp_strength↑ => граница ближе к hi (ниже)
+        cut_y = y + int(frac * h)
+        keep = mask_full.crop((0, max(0, min(th, cut_y)), tw, th))
+        out = Image.new("L", (tw, th), 0)
+        out.paste(keep, (0, max(0, min(th, cut_y))))
+        return out
+
+    return mask_full
+
+
 # =========================
 # Model init (local IDM-VTON)
 # =========================
@@ -337,7 +347,6 @@ print("DEVICE:", DEVICE, "DTYPE:", DTYPE, flush=True)
 
 tensor_transfrom = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
 
-# Components
 unet = UNet2DConditionModel.from_pretrained(base_path, subfolder="unet", torch_dtype=DTYPE)
 unet.requires_grad_(False)
 
@@ -355,7 +364,6 @@ vae = AutoencoderKL.from_pretrained(base_path, subfolder="vae", torch_dtype=DTYP
 UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(base_path, subfolder="unet_encoder", torch_dtype=DTYPE)
 UNet_Encoder.requires_grad_(False)
 
-# Preprocessors
 parsing_model = Parsing(0)
 openpose_model = OpenPose(0)
 
@@ -379,58 +387,74 @@ pipe.unet_encoder = UNet_Encoder
 
 
 # =========================
-# Inference (returns ONLY final image)
+# Inference
 # =========================
 @spaces.GPU
 def start_tryon(
     human_pil: Image.Image,
     garm_img: Image.Image,
     garm_relpath: str = "",
+    garment_type_override: str = "auto",  # auto | upper_body | lower_body | dresses
     auto_mask: bool = True,
-    denoise_steps: int = 30,
-    guidance_scale: float = 3.5,
+    safety_clamp: bool = True,
+    clamp_strength: float = 0.55,  # 0..1
+    denoise_steps: int = 34,
+    guidance_scale: float = 3.8,
     strength: float = 0.90,
     seed: int = -1,
     max_side: int = 1024,
     prompt_override: str = "",
     negative_prompt: str = "monochrome, lowres, bad anatomy, worst quality, low quality",
 ) -> Image.Image:
-    # pick device/dtype
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.float16 if device == "cuda" else torch.float32
 
-    # Move models
     if device == "cuda":
         openpose_model.preprocessor.body_estimation.model.to(device)
+
     pipe.to(device)
     pipe.unet_encoder.to(device)
 
-    # --- sizes (dynamic, no forced 768x1024) ---
     human_img_orig = human_pil.convert("RGB")
     src_w, src_h = human_img_orig.size
+
     target_w, target_h = pick_target_size_keep_aspect(src_w, src_h, max_side=max_side)
 
-    # letterbox to target size (no crop-center, no paste-back)
-    human_lb, lb_meta = letterbox(human_img_orig, target_w, target_h, fill=(0, 0, 0))
+    # letterbox for model canvas (important: gray padding)
+    human_lb, lb_meta = letterbox(human_img_orig, target_w, target_h, fill=(127, 127, 127))
     garm_img = garm_img.convert("RGB")
-    garm_lb, _ = letterbox(garm_img, target_w, target_h, fill=(0, 0, 0))
+    garm_lb, _ = letterbox(garm_img, target_w, target_h, fill=(127, 127, 127))
 
-    # --- Mask (not fixed upper_body) ---
-    if auto_mask:
-        # preprocess runs on 384x512; use letterbox to avoid distortion
-        human_384, _m = letterbox(human_lb, 384, 512, fill=(0, 0, 0))
-        keypoints = openpose_model(human_384)
-        model_parse, _ = parsing_model(human_384)
+    # Core region (no padding) — IMPORTANT for preprocessors
+    human_core = unletterbox(human_lb, lb_meta)
+    x, y, w, h = lb_meta["x"], lb_meta["y"], lb_meta["w"], lb_meta["h"]
 
+    # garment class
+    if garment_type_override and garment_type_override != "auto":
+        cloth_class = garment_type_override
+    else:
         cloth_class = infer_garment_class_from_path(garm_relpath)
-        mask, _ = get_mask_location("hd", cloth_class, model_parse, keypoints)
-        # upscale mask back to target size
-        mask = mask.resize((target_w, target_h), Image.BILINEAR)
+
+    # ---- MASK (compute on core -> paste to full) ----
+    if auto_mask:
+        human_core_384 = human_core.resize((384, 512), Image.BILINEAR)
+        keypoints = openpose_model(human_core_384)
+        model_parse, _ = parsing_model(human_core_384)
+
+        mask_core_384, _ = get_mask_location("hd", cloth_class, model_parse, keypoints)
+        mask_core = mask_core_384.resize((w, h), Image.BILINEAR)
+
+        mask_full = Image.new("L", (target_w, target_h), 0)
+        mask_full.paste(mask_core, (x, y))
+
+        if safety_clamp:
+            mask_full = apply_safety_clamp(mask_full, lb_meta, cloth_class, clamp_strength)
+        mask = mask_full
     else:
         mask = Image.new("L", (target_w, target_h), 0)
 
-    # --- DensePose ---
-    human_dp = _apply_exif_orientation(human_lb.resize((384, 512)))
+    # ---- DensePose (compute on core -> paste to full) ----
+    human_dp = _apply_exif_orientation(human_core.resize((384, 512), Image.BILINEAR))
     human_dp = convert_PIL_to_numpy(human_dp, format="BGR")
 
     args = apply_net.create_argument_parser().parse_args(
@@ -445,11 +469,12 @@ def start_tryon(
             "cuda" if device == "cuda" else "cpu",
         )
     )
-    pose_img = args.func(args, human_dp)
-    pose_img = pose_img[:, :, ::-1]
-    pose_img = Image.fromarray(pose_img).resize((target_w, target_h), Image.BILINEAR)
+    pose_core = args.func(args, human_dp)
+    pose_core = pose_core[:, :, ::-1]
+    pose_core = Image.fromarray(pose_core).resize((w, h), Image.BILINEAR)
+    pose_img = paste_into_canvas("RGB", (target_w, target_h), pose_core, lb_meta, (127, 127, 127))
 
-    # --- prompts (not fixed “a garment”) ---
+    # ---- prompts (not fixed) ----
     garment_desc = guess_garment_description(garm_relpath)
     if prompt_override and prompt_override.strip():
         garment_desc = prompt_override.strip()
@@ -457,15 +482,14 @@ def start_tryon(
     prompt_main = f"model is wearing {garment_desc}"
     prompt_cloth = f"a photo of {garment_desc}"
 
-    # --- params (no fixed low steps/cfg/seed) ---
+    # ---- params ----
     denoise_steps = clamp_int(denoise_steps, 15, 60)
     guidance_scale = clamp_float(guidance_scale, 0.0, 12.0)
     strength = clamp_float(strength, 0.50, 1.00)
-    if seed is None:
-        seed = -1
-    seed = int(seed)
+    max_side = clamp_int(max_side, 640, 2048)
+
+    seed = int(seed) if seed is not None else -1
     if seed < 0:
-        # random but reproducible per call if needed
         seed = int.from_bytes(os.urandom(2), "big") + int(time.time() * 1000) % 1000000
 
     with torch.no_grad():
@@ -504,7 +528,6 @@ def start_tryon(
 
             pose_t = tensor_transfrom(pose_img).unsqueeze(0).to(device=device, dtype=dtype)
             garm_t = tensor_transfrom(garm_lb).unsqueeze(0).to(device=device, dtype=dtype)
-
             generator = torch.Generator(device).manual_seed(seed)
 
             images = pipe(
@@ -514,7 +537,7 @@ def start_tryon(
                 negative_pooled_prompt_embeds=negative_pooled_prompt_embeds.to(device=device, dtype=dtype),
                 num_inference_steps=denoise_steps,
                 generator=generator,
-                strength=strength,  # <-- not fixed 1.0
+                strength=strength,
                 pose_img=pose_t,
                 text_embeds_cloth=prompt_embeds_c.to(device=device, dtype=dtype),
                 cloth=garm_t,
@@ -522,13 +545,13 @@ def start_tryon(
                 image=human_lb,
                 height=target_h,
                 width=target_w,
-                ip_adapter_image=garm_lb,  # keep conditioning, but not hard-resized 768x1024
-                guidance_scale=guidance_scale,  # <-- not fixed low value
+                ip_adapter_image=garm_lb,
+                guidance_scale=guidance_scale,
             )[0]
 
     out_img_lb = images[0].convert("RGB")
 
-    # remove letterbox padding and resize back to original size (no crop-center paste)
+    # remove padding and return to original resolution
     out_core = unletterbox(out_img_lb, lb_meta)
     out_final = out_core.resize((src_w, src_h), Image.LANCZOS)
     return out_final
@@ -563,7 +586,10 @@ def on_gallery_select(files_list: List[str], evt: gr.SelectData):
 def tryon_ui(
     person_pil,
     selected_filename,
+    garment_type_override,
     auto_mask,
+    safety_clamp,
+    clamp_strength,
     steps,
     cfg,
     strength,
@@ -595,7 +621,10 @@ def tryon_ui(
             human_pil=person_pil,
             garm_img=garm,
             garm_relpath=selected_filename,
+            garment_type_override=str(garment_type_override),
             auto_mask=bool(auto_mask),
+            safety_clamp=bool(safety_clamp),
+            clamp_strength=float(clamp_strength),
             denoise_steps=int(steps),
             guidance_scale=float(cfg),
             strength=float(strength),
@@ -608,7 +637,7 @@ def tryon_ui(
         yield None, f"❌ Ошибка: {type(e).__name__}: {str(e)[:220]}"
 
 
-# Preload garments
+# preload garments
 ensure_garments_downloaded()
 _initial_files = list_garments()
 _initial_items = build_gallery_items(_initial_files)
@@ -635,17 +664,34 @@ with gr.Blocks(title="Virtual Try-On Rendez-vous", css=CUSTOM_CSS) as demo:
                 allow_preview=True,
             )
 
-            with gr.Accordion("⚙️ Настройки качества", open=False):
-                auto_mask = gr.Checkbox(value=True, label="Auto mask (парсинг + поза)")
-                steps = gr.Slider(15, 60, value=30, step=1, label="Шаги (num_inference_steps)")
-                cfg = gr.Slider(0.0, 12.0, value=3.5, step=0.1, label="Guidance scale (CFG)")
-                strength = gr.Slider(0.50, 1.00, value=0.90, step=0.01, label="Strength (насколько сильно перерисовывать)")
+            with gr.Accordion("⚙️ Настройки", open=False):
+                garment_type_override = gr.Dropdown(
+                    choices=["auto", "upper_body", "lower_body", "dresses"],
+                    value="auto",
+                    label="Тип одежды (override)",
+                )
+                auto_mask = gr.Checkbox(value=True, label="Auto mask (parsing + openpose)")
+
+                safety_clamp = gr.Checkbox(
+                    value=True,
+                    label="Safety clamp (защита от съезда зоны редактирования)",
+                )
+                clamp_strength = gr.Slider(
+                    0.0, 1.0, value=0.55, step=0.01,
+                    label="Clamp strength (0 = мягко, 1 = сильнее)",
+                )
+
+                steps = gr.Slider(15, 60, value=34, step=1, label="Шаги (num_inference_steps)")
+                cfg = gr.Slider(0.0, 12.0, value=3.8, step=0.1, label="Guidance scale (CFG)")
+                strength = gr.Slider(0.50, 1.00, value=0.90, step=0.01, label="Strength")
+
                 seed = gr.Number(value=-1, precision=0, label="Seed (-1 = случайный)")
-                max_side = gr.Slider(768, 1408, value=1024, step=64, label="Максимальный размер стороны (динамический)")
+                max_side = gr.Slider(768, 1536, value=1024, step=64, label="Макс. сторона (динамический размер)")
+
                 prompt_override = gr.Textbox(
                     value="",
                     label="Описание одежды (опц.)",
-                    placeholder="Напр.: a black leather jacket / a blazer / a coat ... (если пусто — авто по имени файла)",
+                    placeholder="Напр.: a blazer / a dress / a t-shirt ... (если пусто — авто по имени файла)",
                 )
 
             run = gr.Button("Примерить", variant="primary")
@@ -668,7 +714,20 @@ with gr.Blocks(title="Virtual Try-On Rendez-vous", css=CUSTOM_CSS) as demo:
 
     run.click(
         fn=tryon_ui,
-        inputs=[person, selected_garment_state, auto_mask, steps, cfg, strength, seed, max_side, prompt_override],
+        inputs=[
+            person,
+            selected_garment_state,
+            garment_type_override,
+            auto_mask,
+            safety_clamp,
+            clamp_strength,
+            steps,
+            cfg,
+            strength,
+            seed,
+            max_side,
+            prompt_override,
+        ],
         outputs=[out, status],
         concurrency_limit=1,
     )