Update app.py

app.py

@@ -1,11 +1,30 @@
 # -*- 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
-from typing import List, Optional, Tuple
+import math
+from typing import List, Optional, Tuple, Dict, Any
 
 import spaces
 import gradio as gr
-from PIL import Image
+from PIL import Image, ImageOps
 
 # =========================
 # FIX: gradio 4.24 / gradio_client crashes on boolean JSON Schemas in /api_info
@@ -179,10 +198,18 @@ def clamp_int(x, lo, hi):
     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))
+
+
 _last_call_ts = 0.0
 
 
-def allow_call(min_interval_sec: float = 2.5) -> Tuple[bool, str]:
+def allow_call(min_interval_sec: float = 2.0) -> Tuple[bool, str]:
     global _last_call_ts
     now = time.time()
     if now - _last_call_ts < min_interval_sec:
@@ -192,6 +219,113 @@ def allow_call(min_interval_sec: float = 2.5) -> Tuple[bool, str]:
     return True, ""
 
 
+def round_to_multiple(x: int, m: int = 8) -> int:
+    return max(m, int(round(x / m) * m))
+
+
+def pick_target_size_keep_aspect(w: int, h: int, max_side: int) -> Tuple[int, int]:
+    """
+    Возвращает (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)
+        th = round_to_multiple(int(th * scale2), 8)
+    return tw, th
+
+
+def letterbox(img: Image.Image, target_w: int, target_h: int, fill=(0, 0, 0)) -> Tuple[Image.Image, Dict[str, int]]:
+    """
+    Масштабирует с сохранением пропорций + padding до target_w/target_h.
+    Возвращает (img_lb, meta) где meta содержит offset/size для обратного unletterbox.
+    """
+    src_w, src_h = img.size
+    if src_w <= 0 or src_h <= 0:
+        out = img.resize((target_w, target_h))
+        return out, {"x": 0, "y": 0, "w": target_w, "h": target_h, "src_w": src_w, "src_h": src_h}
+
+    scale = min(target_w / src_w, target_h / src_h)
+    new_w = max(1, int(src_w * scale))
+    new_h = max(1, int(src_h * scale))
+
+    img_rs = img.resize((new_w, new_h), Image.LANCZOS)
+    canvas = Image.new("RGB", (target_w, target_h), fill)
+    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
+
+
+def unletterbox(img_lb: Image.Image, meta: Dict[str, int]) -> Image.Image:
+    """
+    Вырезает область без padding и возвращает как есть (потом можно resize к исходнику).
+    """
+    x, y, w, h = meta["x"], meta["y"], meta["w"], meta["h"]
+    return img_lb.crop((x, y, x + w, y + h))
+
+
+def infer_garment_class_from_path(relpath: str) -> str:
+    """
+    Возвращает тип для get_mask_location: 'upper_body' | 'lower_body' | 'dresses'
+    Это НЕ “жестко upper_body” — эвристика по папке/имени.
+    """
+    s = (relpath or "").lower()
+    # папки/имена под платья
+    if any(k in s for k in ["dress", "dresses", "suk", "plate", "плать", "sarafan"]):
+        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()
+    # словарь эвристик
+    mapping = [
+        (["shearling", "dub", "дублен", "sheepskin"], "a shearling jacket"),
+        (["coat", "пальт", "overcoat"], "a coat"),
+        (["jacket", "куртк", "bomber", "парка", "parka"], "a jacket"),
+        (["blazer", "пидж", "suit"], "a blazer"),
+        (["hoodie", "худи"], "a hoodie"),
+        (["sweater", "свит", "jumper"], "a sweater"),
+        (["shirt", "рубаш"], "a shirt"),
+        (["tshirt", "tee", "футбол"], "a t-shirt"),
+        (["dress", "плать", "sarafan"], "a dress"),
+        (["pants", "jeans", "брюк", "джин"], "pants"),
+        (["skirt", "юбк"], "a skirt"),
+    ]
+    for keys, desc in mapping:
+        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 piece of clothing"
+
+
 # =========================
 # Model init (local IDM-VTON)
 # =========================
@@ -201,9 +335,7 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
 print("DEVICE:", DEVICE, "DTYPE:", DTYPE, flush=True)
 
-tensor_transfrom = transforms.Compose(
-    [transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]
-)
+tensor_transfrom = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
 
 # Components
 unet = UNet2DConditionModel.from_pretrained(base_path, subfolder="unet", torch_dtype=DTYPE)
@@ -253,11 +385,17 @@ pipe.unet_encoder = UNet_Encoder
 def start_tryon(
     human_pil: Image.Image,
     garm_img: Image.Image,
+    garm_relpath: str = "",
     auto_mask: bool = True,
-    crop_center: bool = True,
-    denoise_steps: int = 25,
-    seed: int = 42,
+    denoise_steps: int = 30,
+    guidance_scale: float = 3.5,
+    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
 
@@ -267,38 +405,33 @@ def start_tryon(
     pipe.to(device)
     pipe.unet_encoder.to(device)
 
-    garm_img = garm_img.convert("RGB").resize((768, 1024))
+    # --- 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)
 
-    # Crop
-    if crop_center:
-        width, height = human_img_orig.size
-        target_width = int(min(width, height * (3 / 4)))
-        target_height = int(min(height, width * (4 / 3)))
-        left = (width - target_width) / 2
-        top = (height - target_height) / 2
-        right = (width + target_width) / 2
-        bottom = (height + target_height) / 2
-        cropped_img = human_img_orig.crop((left, top, right, bottom))
-        crop_size = cropped_img.size
-        human_img = cropped_img.resize((768, 1024))
-    else:
-        human_img = human_img_orig.resize((768, 1024))
-        crop_size = None
-        left = top = 0
+    # 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))
+    garm_img = garm_img.convert("RGB")
+    garm_lb, _ = letterbox(garm_img, target_w, target_h, fill=(0, 0, 0))
 
-    # Mask (как раньше: upper_body всегда)
+    # --- Mask (not fixed upper_body) ---
     if auto_mask:
-        keypoints = openpose_model(human_img.resize((384, 512)))
-        model_parse, _ = parsing_model(human_img.resize((384, 512)))
-        mask, _ = get_mask_location("hd", "upper_body", model_parse, keypoints)
-        mask = mask.resize((768, 1024))
+        # 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)
+
+        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)
     else:
-        mask = Image.new("L", (768, 1024), 0)
+        mask = Image.new("L", (target_w, target_h), 0)
 
-    # DensePose
-    human_img_arg = _apply_exif_orientation(human_img.resize((384, 512)))
-    human_img_arg = convert_PIL_to_numpy(human_img_arg, format="BGR")
+    # --- DensePose ---
+    human_dp = _apply_exif_orientation(human_lb.resize((384, 512)))
+    human_dp = convert_PIL_to_numpy(human_dp, format="BGR")
 
     args = apply_net.create_argument_parser().parse_args(
         (
@@ -312,18 +445,28 @@ def start_tryon(
             "cuda" if device == "cuda" else "cpu",
         )
     )
-    pose_img = args.func(args, human_img_arg)
+    pose_img = args.func(args, human_dp)
     pose_img = pose_img[:, :, ::-1]
-    pose_img = Image.fromarray(pose_img).resize((768, 1024))
-
-    # Fixed prompts (как раньше)
-    garment_des = "a garment"
-    prompt_main = "model is wearing " + garment_des
-    prompt_cloth = "a photo of " + garment_des
-    negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"
-
-    denoise_steps = clamp_int(denoise_steps, 20, 40)
-    seed = clamp_int(seed, 0, 999999)
+    pose_img = Image.fromarray(pose_img).resize((target_w, target_h), Image.BILINEAR)
+
+    # --- prompts (not fixed “a garment”) ---
+    garment_desc = guess_garment_description(garm_relpath)
+    if prompt_override and prompt_override.strip():
+        garment_desc = prompt_override.strip()
+
+    prompt_main = f"model is wearing {garment_desc}"
+    prompt_cloth = f"a photo of {garment_desc}"
+
+    # --- params (no fixed low steps/cfg/seed) ---
+    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)
+    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():
         if device == "cuda":
@@ -352,12 +495,7 @@ def start_tryon(
                 negative_prompt=negative_prompt,
             )
 
-            (
-                prompt_embeds_c,
-                _,
-                _,
-                _,
-            ) = pipe.encode_prompt(
+            (prompt_embeds_c, _, _, _) = pipe.encode_prompt(
                 [prompt_cloth],
                 num_images_per_prompt=1,
                 do_classifier_free_guidance=False,
@@ -365,7 +503,7 @@ def start_tryon(
             )
 
             pose_t = tensor_transfrom(pose_img).unsqueeze(0).to(device=device, dtype=dtype)
-            garm_t = tensor_transfrom(garm_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)
 
@@ -376,24 +514,24 @@ def start_tryon(
                 negative_pooled_prompt_embeds=negative_pooled_prompt_embeds.to(device=device, dtype=dtype),
                 num_inference_steps=denoise_steps,
                 generator=generator,
-                strength=1.0,
+                strength=strength,  # <-- not fixed 1.0
                 pose_img=pose_t,
                 text_embeds_cloth=prompt_embeds_c.to(device=device, dtype=dtype),
                 cloth=garm_t,
                 mask_image=mask,
-                image=human_img,
-                height=1024,
-                width=768,
-                ip_adapter_image=garm_img.resize((768, 1024)),
-                guidance_scale=2.0,
+                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
             )[0]
 
-    out_img = images[0]
-    if crop_center and crop_size is not None:
-        out_img_rs = out_img.resize(crop_size)
-        human_img_orig.paste(out_img_rs, (int(left), int(top)))
-        return human_img_orig
-    return out_img
+    out_img_lb = images[0].convert("RGB")
+
+    # remove letterbox padding and resize back to original size (no crop-center paste)
+    out_core = unletterbox(out_img_lb, lb_meta)
+    out_final = out_core.resize((src_w, src_h), Image.LANCZOS)
+    return out_final
 
 
 # =========================
@@ -406,17 +544,6 @@ div[class*="footer"] {display:none !important;}
 button[aria-label="Settings"] {display:none !important;}
 """
 
-PHOTO_TIPS_MD = """
-### Какое фото подойдёт
-✅ В полный рост или по пояс  
-✅ Руки �� предметы не закрывают тело  
-✅ Одежда по фигуре  
-✅ Вы стоите прямо и смотрите в камеру  
-✅ Хорошее освещение  
-✅ В кадре нет других людей
-"""
-
-
 def refresh_catalog():
     ensure_garments_downloaded()
     files = list_garments()
@@ -433,10 +560,20 @@ def on_gallery_select(files_list: List[str], evt: gr.SelectData):
     return files_list[idx], f"👕 Выбрано: {files_list[idx]}"
 
 
-def tryon_ui(person_pil, selected_filename):
+def tryon_ui(
+    person_pil,
+    selected_filename,
+    auto_mask,
+    steps,
+    cfg,
+    strength,
+    seed,
+    max_side,
+    prompt_override,
+):
     yield None, "⏳ Обработка... (первый запуск может быть дольше)"
 
-    ok, msg = allow_call(2.5)
+    ok, msg = allow_call(2.0)
     if not ok:
         yield None, msg
         return
@@ -457,10 +594,14 @@ def tryon_ui(person_pil, selected_filename):
         out_img = start_tryon(
             human_pil=person_pil,
             garm_img=garm,
-            auto_mask=True,
-            crop_center=True,
-            denoise_steps=25,
-            seed=42,
+            garm_relpath=selected_filename,
+            auto_mask=bool(auto_mask),
+            denoise_steps=int(steps),
+            guidance_scale=float(cfg),
+            strength=float(strength),
+            seed=int(seed),
+            max_side=int(max_side),
+            prompt_override=str(prompt_override or "").strip(),
         )
         yield out_img, "✅ Готово"
     except Exception as e:
@@ -482,9 +623,6 @@ with gr.Blocks(title="Virtual Try-On Rendez-vous", css=CUSTOM_CSS) as demo:
         with gr.Column():
             person = gr.Image(label="Фото человека", type="pil", height=420)
 
-            # Подсказка под загрузкой фото
-            gr.Markdown(PHOTO_TIPS_MD)
-
             with gr.Row():
                 refresh_btn = gr.Button("🔄 Обновить каталог одежды", variant="secondary")
                 selected_label = gr.Markdown("👕 Выберите одежду ниже")
@@ -497,6 +635,19 @@ 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 (насколько сильно перерисовывать)")
+                seed = gr.Number(value=-1, precision=0, label="Seed (-1 = случайный)")
+                max_side = gr.Slider(768, 1408, value=1024, step=64, label="Максимальный размер стороны (динамический)")
+                prompt_override = gr.Textbox(
+                    value="",
+                    label="Описание одежды (опц.)",
+                    placeholder="Напр.: a black leather jacket / a blazer / a coat ... (если пусто — авто по имени файла)",
+                )
+
             run = gr.Button("Примерить", variant="primary")
             status = gr.Textbox(value="Ожидание...", interactive=False)
 
@@ -517,7 +668,7 @@ with gr.Blocks(title="Virtual Try-On Rendez-vous", css=CUSTOM_CSS) as demo:
 
     run.click(
         fn=tryon_ui,
-        inputs=[person, selected_garment_state],
+        inputs=[person, selected_garment_state, auto_mask, steps, cfg, strength, seed, max_side, prompt_override],
         outputs=[out, status],
         concurrency_limit=1,
     )
@@ -532,5 +683,5 @@ if __name__ == "__main__":
         auth=APP_AUTH,
         max_threads=4,
         show_error=True,
-        show_api=False,  # важно: не показываем API, но /api_info могут дергать — патч это чинит
-    )
+        show_api=False,
+    )