Update app.py

app.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 import os
 import time
-import tempfile
 from typing import List, Optional, Tuple
 
 import spaces
@@ -11,7 +10,6 @@ from PIL import Image
 import torch
 import numpy as np
 from torchvision import transforms
-from torchvision.transforms.functional import to_pil_image
 
 from huggingface_hub import login, snapshot_download
 
@@ -51,6 +49,7 @@ ALLOWED_EXTS = (".png", ".jpg", ".jpeg", ".webp")
 GARMENTS_DATASET = os.getenv("GARMENTS_DATASET", "").strip()  # e.g. "ArmanRV/armanrv-garments"
 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.
@@ -69,7 +68,6 @@ def ensure_garments_downloaded() -> None:
         return
 
     try:
-        # Download snapshot to local garments/ (no symlinks for HF container)
         snapshot_download(
             repo_id=GARMENTS_DATASET,
             repo_type="dataset",
@@ -81,20 +79,27 @@ def ensure_garments_downloaded() -> None:
     except Exception as e:
         print("Garments download FAILED:", str(e)[:300])
 
+
 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.sort()
+    """
+    Recursively list images inside ./garments (handles dataset subfolders).
+    """
+    files: List[str] = []
+    if not os.path.isdir(GARMENT_DIR):
         return files
-    except Exception:
-        return []
+    for root, _, fnames in os.walk(GARMENT_DIR):
+        for f in fnames:
+            if f.lower().endswith(ALLOWED_EXTS) and not f.startswith("."):
+                rel = os.path.relpath(os.path.join(root, f), GARMENT_DIR)
+                files.append(rel)
+    files.sort()
+    return files
+
 
 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
@@ -106,8 +111,8 @@ def load_garment_pil(filename: str) -> Optional[Image.Image]:
     except Exception:
         return None
 
+
 def build_gallery_items(files: List[str]):
-    # (image_path, caption) — caption empty for clean UI
     return [(garment_path(f), "") for f in files]
 
 
@@ -121,14 +126,7 @@ def clamp_int(x, lo, hi):
         x = lo
     return max(lo, min(hi, x))
 
-def pil_to_binary_mask(pil_image, threshold=0):
-    np_image = np.array(pil_image)
-    grayscale_image = Image.fromarray(np_image).convert("L")
-    binary_mask = np.array(grayscale_image) > threshold
-    mask = (binary_mask.astype(np.uint8) * 255)
-    return Image.fromarray(mask)
 
-# global simple rate limit (helps avoid spam during internal demo)
 _last_call_ts = 0.0
 def allow_call(min_interval_sec: float = 2.5) -> Tuple[bool, str]:
     global _last_call_ts
@@ -145,13 +143,15 @@ def allow_call(min_interval_sec: float = 2.5) -> Tuple[bool, str]:
 # =========================
 base_path = "yisol/IDM-VTON"
 
-# device policy
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
-
 print("DEVICE:", DEVICE, "DTYPE:", DTYPE)
 
-# Load components
+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)
 
@@ -164,25 +164,18 @@ text_encoder_one = CLIPTextModel.from_pretrained(base_path, subfolder="text_enco
 text_encoder_two = CLIPTextModelWithProjection.from_pretrained(base_path, subfolder="text_encoder_2", torch_dtype=DTYPE)
 
 image_encoder = CLIPVisionModelWithProjection.from_pretrained(base_path, subfolder="image_encoder", torch_dtype=DTYPE)
-
 vae = AutoencoderKL.from_pretrained(base_path, subfolder="vae", torch_dtype=DTYPE)
 
 UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(base_path, subfolder="unet_encoder", torch_dtype=DTYPE)
 UNet_Encoder.requires_grad_(False)
 
-# Parsing/OpenPose init
-# These are heavy; GPU intended. On CPU it may be very slow.
+# Preprocessors
 parsing_model = Parsing(0)
 openpose_model = OpenPose(0)
 
-# Freeze
 for m in [UNet_Encoder, image_encoder, vae, unet, text_encoder_one, text_encoder_two]:
     m.requires_grad_(False)
 
-tensor_transfrom = transforms.Compose(
-    [transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]
-)
-
 pipe = TryonPipeline.from_pretrained(
     base_path,
     unet=unet,
@@ -200,9 +193,9 @@ pipe.unet_encoder = UNet_Encoder
 
 
 # =========================
-# Inference
+# Inference (returns ONLY final image)
 # =========================
-@spaces.GPU  # ok on dedicated GPU too
+@spaces.GPU
 def start_tryon(
     human_pil: Image.Image,
     garm_img: Image.Image,
@@ -210,25 +203,21 @@ def start_tryon(
     crop_center: bool = True,
     denoise_steps: int = 25,
     seed: int = 42,
-):
-    """
-    Simplified local try-on.
-    Returns: (output_image, masked_preview)
-    """
+) -> Image.Image:
+
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.float16 if device == "cuda" else torch.float32
 
-    # move heavy models
+    # Move models
     if device == "cuda":
         openpose_model.preprocessor.body_estimation.model.to(device)
     pipe.to(device)
     pipe.unet_encoder.to(device)
 
-    # resize inputs to expected
     garm_img = garm_img.convert("RGB").resize((768, 1024))
     human_img_orig = human_pil.convert("RGB")
 
-    # optional center crop
+    # Crop
     if crop_center:
         width, height = human_img_orig.size
         target_width = int(min(width, height * (3 / 4)))
@@ -243,20 +232,16 @@ def start_tryon(
     else:
         human_img = human_img_orig.resize((768, 1024))
 
-    # mask
+    # Mask
     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))
     else:
-        # if someday you add manual mask, you can pass it here
         mask = Image.new("L", (768, 1024), 0)
 
-    mask_gray = (1 - transforms.ToTensor()(mask)) * tensor_transfrom(human_img)
-    mask_gray = to_pil_image((mask_gray + 1.0) / 2.0)
-
-    # densepose
+    # DensePose
     human_img_arg = _apply_exif_orientation(human_img.resize((384, 512)))
     human_img_arg = convert_PIL_to_numpy(human_img_arg, format="BGR")
 
@@ -274,7 +259,7 @@ def start_tryon(
     pose_img = pose_img[:, :, ::-1]
     pose_img = Image.fromarray(pose_img).resize((768, 1024))
 
-    # prompts (fixed, like your API demo)
+    # Fixed prompts (like your API demo)
     garment_des = "a garment"
     prompt_main = "model is wearing " + garment_des
     prompt_cloth = "a photo of " + garment_des
@@ -283,19 +268,16 @@ def start_tryon(
     denoise_steps = clamp_int(denoise_steps, 20, 40)
     seed = clamp_int(seed, 0, 999999)
 
-    # inference
     with torch.no_grad():
         if device == "cuda":
             autocast_ctx = torch.cuda.amp.autocast()
         else:
-            # no autocast on cpu
             class _NoCtx:
                 def __enter__(self): return None
                 def __exit__(self, *args): return False
             autocast_ctx = _NoCtx()
 
         with autocast_ctx:
-            # encode prompts
             (
                 prompt_embeds,
                 negative_prompt_embeds,
@@ -348,8 +330,8 @@ def start_tryon(
     if crop_center:
         out_img_rs = out_img.resize(crop_size)
         human_img_orig.paste(out_img_rs, (int(left), int(top)))
-        return human_img_orig, mask_gray
-    return out_img, mask_gray
+        return human_img_orig
+    return out_img
 
 
 # =========================
@@ -366,7 +348,7 @@ def refresh_catalog():
     ensure_garments_downloaded()
     files = list_garments()
     items = build_gallery_items(files)
-    status = "✅ Каталог обновлён" if files else "⚠️ Каталог пуст (dataset не скачался или нет файлов)"
+    status = "✅ Каталог обновлён" if files else "⚠️ Каталог пуст (проверь dataset/токен)"
     return items, files, None, status
 
 def on_gallery_select(files_list: List[str], evt: gr.SelectData):
@@ -377,31 +359,41 @@ 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):
+    # Сразу показать что кнопка сработала
+    yield None, "⏳ Обработка... (первый запуск может быть дольше)"
+
     ok, msg = allow_call(2.5)
     if not ok:
-        return None, None, msg
+        yield None, msg
+        return
 
     if person_pil is None:
-        return None, None, "❌ Загрузите фото человека"
+        yield None, "❌ Загрузите фото человека"
+        return
     if not selected_filename:
-        return None, None, "❌ Выберите одежду из каталога"
+        yield None, "❌ Выберите одежду (клик по превью)"
+        return
 
     garm = load_garment_pil(selected_filename)
     if garm is None:
-        return None, None, "❌ Не удалось загрузить выбранную одежду"
-
-    out, masked = start_tryon(
-        human_pil=person_pil,
-        garm_img=garm,
-        auto_mask=True,
-        crop_center=True,
-        denoise_steps=25,
-        seed=42,
-    )
-    return out, masked, "✅ Готово"
+        yield None, "❌ Не удалось загрузить выбранную одежду"
+        return
 
+    try:
+        out = start_tryon(
+            human_pil=person_pil,
+            garm_img=garm,
+            auto_mask=True,
+            crop_center=True,
+            denoise_steps=25,
+            seed=42,
+        )
+        yield out, "✅ Готово"
+    except Exception as e:
+        yield None, f"❌ Ошибка: {type(e).__name__}: {str(e)[:220]}"
 
-# ensure garments present at startup (best effort)
+
+# Preload garments (best-effort)
 ensure_garments_downloaded()
 _initial_files = list_garments()
 _initial_items = build_gallery_items(_initial_files)
@@ -432,8 +424,7 @@ with gr.Blocks(title="Virtual Try-On Rendez-vous", css=CUSTOM_CSS) as demo:
             status = gr.Textbox(value="Ожидание...", interactive=False)
 
         with gr.Column():
-            out = gr.Image(label="Результат", type="pil", height=520)
-            masked = gr.Image(label="Маска/предпросмотр (служебное)", type="pil", height=320)
+            out = gr.Image(label="Результат", type="pil", height=760)
 
     garment_gallery.select(
         fn=on_gallery_select,
@@ -450,9 +441,12 @@ with gr.Blocks(title="Virtual Try-On Rendez-vous", css=CUSTOM_CSS) as demo:
     run.click(
         fn=tryon_ui,
         inputs=[person, selected_garment_state],
-        outputs=[out, masked, status],
+        outputs=[out, status],
     )
 
+# IMPORTANT: queue helps stability on GPU
+demo.queue(concurrency_count=1, max_size=20)
+
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",