update garment pil

app.py

@@ -1,499 +1,3 @@
-# import os
-# import sys
-
-# # ---------------------------------------------------------
-# # 0) Make sure local packages (diffusers3, preprocess, etc.) are importable on HF Spaces
-# # ---------------------------------------------------------
-# ROOT = os.path.dirname(os.path.abspath(__file__))
-# if ROOT not in sys.path:
-#     sys.path.insert(0, ROOT)
-
-# print("[BOOT] ROOT =", ROOT, flush=True)
-# print("[BOOT] sys.path[:5] =", sys.path[:5], flush=True)
-
-# import tempfile
-# from dataclasses import dataclass
-# from functools import lru_cache
-# from typing import Optional, Tuple
-
-# import gradio as gr
-# import torch
-# import numpy as np
-# import cv2
-# import imageio
-# from PIL import Image, ImageOps
-
-# from huggingface_hub import hf_hub_download
-
-# from diffusers import UniPCMultistepScheduler
-
-# # Show where diffusers3 is imported from (helps diagnose import collisions on Spaces)
-# import diffusers3
-# print("[BOOT] diffusers3 loaded from:", getattr(diffusers3, "__file__", "<?>"), flush=True)
-
-# from diffusers3.models.controlnet import ControlNetModel
-# from diffusers3.pipelines.controlnet.pipeline_controlnet_sd_xl_img2img_img import (
-#     StableDiffusionXLControlNetImg2ImgPipeline,
-# )
-# from ip_adapter import IPAdapterXL
-
-# # 너 코드에서 쓰는 extractor API 그대로( run(category, input_path, model_restore=...) )
-# from preprocess.simple_extractor import run as run_simple_extractor
-
-
-# # =========================
-# # HF Hub repo ids
-# # =========================
-# BASE_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
-# CONTROLNET_ID = "diffusers/controlnet-depth-sdxl-1.0"
-
-# # 네 assets dataset repo (가중치 저장소)
-# ASSETS_REPO = os.getenv("ASSETS_REPO", "soye/VISTA_assets")
-# ASSETS_REPO_TYPE = "dataset"  # dataset repo로 올렸으니
-
-
-# def asset_path(relpath: str) -> str:
-#     """
-#     relpath 예:
-#       - "image_encoder/model.safetensors"
-#       - "image_encoder/config.json"
-#       - "ip_adapter/ip-adapter_sdxl_vit-h.bin"
-#       - "preprocess_ckpts/exp-schp-201908301523-atr.pth"
-#     """
-#     return hf_hub_download(
-#         repo_id=ASSETS_REPO,
-#         repo_type=ASSETS_REPO_TYPE,
-#         filename=relpath,
-#     )
-
-
-# @lru_cache(maxsize=1)
-# def get_assets():
-#     """
-#     Lazily downloads required assets on first use.
-
-#     Returns:
-#         (image_encoder_dir, ip_ckpt_path, schp_ckpt_path)
-#     """
-#     print("[ASSETS] Downloading assets from:", ASSETS_REPO, flush=True)
-
-#     # Image encoder folder is needed by IPAdapterXL
-#     image_encoder_weight = asset_path("image_encoder/model.safetensors")
-#     _ = asset_path("image_encoder/config.json")  # ensure config exists locally
-#     image_encoder_dir = os.path.dirname(image_encoder_weight)
-
-#     ip_ckpt = asset_path("ip_adapter/ip-adapter_sdxl_vit-h.bin")
-#     schp_ckpt = asset_path("preprocess_ckpts/exp-schp-201908301523-atr.pth")
-
-#     print("[ASSETS] image_encoder_dir =", image_encoder_dir, flush=True)
-#     print("[ASSETS] ip_ckpt =", ip_ckpt, flush=True)
-#     print("[ASSETS] schp_ckpt =", schp_ckpt, flush=True)
-#     return image_encoder_dir, ip_ckpt, schp_ckpt
-
-
-# DEFAULT_STEPS = 40
-# DEBUG_SAVE = False
-
-# # 전역 resize params
-# H: Optional[int] = None  # 1024
-# W: Optional[int] = None  # aspect 유지 기반
-
-
-# @dataclass
-# class Paths:
-#     person_path: str
-#     depth_path: str
-#     style_path: str
-#     output_path: str
-
-
-# def _imread_or_raise(path: str, flag=cv2.IMREAD_COLOR):
-#     img = cv2.imread(path, flag)
-#     if img is None:
-#         raise FileNotFoundError(f"cv2.imread failed: {path} (exists={os.path.exists(path)})")
-#     return img
-
-
-# def compute_hw_from_person(person_path: str):
-#     """
-#     - height=1024 고정, aspect 유지로 W 계산
-#     - demo 안정성 위해 W를 1024로 cap.
-#     """
-#     img = _imread_or_raise(person_path)
-#     orig_h, orig_w = img.shape[:2]
-#     scale = 1024.0 / float(orig_h)
-#     new_h = 1024
-#     new_w = int(round(orig_w * scale))
-#     if new_w > 1024:
-#         new_w = 1024
-#     return new_h, new_w
-
-
-# def invert_sketch_area(sketch_pil: Image.Image) -> Image.Image:
-#     return ImageOps.invert(sketch_pil.convert("L")).convert("RGB")
-
-
-# def fill_sketch_from_image_path_to_pil(image_path: str) -> Image.Image:
-#     global H, W
-#     if H is None or W is None:
-#         raise RuntimeError("Global H/W not set.")
-#     img = _imread_or_raise(image_path, cv2.IMREAD_GRAYSCALE)
-#     img = cv2.resize(img, (W, H), interpolation=cv2.INTER_NEAREST)
-#     _, binary = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY_INV)
-#     contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-#     filled = np.zeros_like(binary)
-#     cv2.drawContours(filled, contours, -1, 255, thickness=cv2.FILLED)
-#     filled_rgb = cv2.cvtColor(filled, cv2.COLOR_GRAY2RGB)
-#     return Image.fromarray(filled_rgb)
-
-
-# def merge_white_regions_or(img1: Image.Image, img2: Image.Image) -> Image.Image:
-#     a = np.array(img1.convert("RGB"), dtype=np.uint8)
-#     b = np.array(img2.convert("RGB"), dtype=np.uint8)
-#     white_a = np.all(a == 255, axis=-1)
-#     white_b = np.all(b == 255, axis=-1)
-#     out = a.copy()
-#     out[white_a | white_b] = 255
-#     return Image.fromarray(out)
-
-
-# def preprocess_mask(mask_img: Image.Image) -> Image.Image:
-#     """
-#     (2번째 첨부 코드 규칙 반영)
-#     - 입력(mask_img)을 전역 (W,H)로 resize (NEAREST)
-#     - 이진화(threshold)
-#     - padding 목표 width는 항상 1024로 고정
-#       * width < 1024: 좌/우 padding (value=0)
-#       * width > 1024: 중앙 crop
-#     - dilation (kernel 17x17, iter=1)
-#     - return: RGB mask PIL
-#     """
-#     global H, W
-
-#     # 1) to grayscale np
-#     m = np.array(mask_img.convert("L"), dtype=np.uint8)
-
-#     # 2) resize to global (W, H)  (2번째 코드의 "resize는 전역 (W,H)" 컨벤션)
-#     if (H is not None) and (W is not None):
-#         # mask는 경계 보존 위해 NEAREST 권장
-#         m = cv2.resize(m, (W, H), interpolation=cv2.INTER_NEAREST)
-
-#     # 3) binarize
-#     _, m = cv2.threshold(m, 127, 255, cv2.THRESH_BINARY)
-
-#     # 4) pad/crop to target width=1024 (항상 고정)
-#     target_width = 1024
-#     h, w = m.shape[:2]
-
-#     if w < target_width:
-#         total_padding = target_width - w
-#         left_padding = total_padding // 2
-#         right_padding = total_padding - left_padding
-#         m = cv2.copyMakeBorder(
-#             m,
-#             top=0, bottom=0,
-#             left=left_padding, right=right_padding,
-#             borderType=cv2.BORDER_CONSTANT,
-#             value=0,
-#         )
-#     elif w > target_width:
-#         left = (w - target_width) // 2
-#         m = m[:, left:left + target_width]
-
-#     # 5) dilate (2번째 코드 동일)
-#     kernel = np.ones((17, 17), np.uint8)
-#     m = cv2.dilate(m, kernel, iterations=1)
-
-#     if DEBUG_SAVE:
-#         cv2.imwrite("mask_final_1024.png", m)
-
-#     # 6) return as RGB PIL
-#     return Image.fromarray(m, mode="L").convert("RGB")
-
-
-
-# def make_depth(depth_path: str) -> Image.Image:
-#     global H, W
-#     if H is None or W is None:
-#         raise RuntimeError("Global H/W not set.")
-#     depth_img = _imread_or_raise(depth_path, 0)
-#     inverted = cv2.bitwise_not(depth_img)
-#     contours, _ = cv2.findContours(inverted, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-#     filled = inverted.copy()
-#     cv2.drawContours(filled, contours, -1, (255), thickness=cv2.FILLED)
-#     filled = cv2.resize(filled, (W, H), interpolation=cv2.INTER_AREA)
-#     filled_rgb = cv2.cvtColor(filled, cv2.COLOR_GRAY2RGB)
-
-#     target_width = 1024
-#     cur_w = filled_rgb.shape[1]
-#     if cur_w < target_width:
-#         pad = (target_width - cur_w) // 2
-#         filled_rgb = cv2.copyMakeBorder(
-#             filled_rgb, 0, 0, pad, pad,
-#             borderType=cv2.BORDER_CONSTANT, value=[0, 0, 0]
-#         )
-#     elif cur_w > target_width:
-#         left = (cur_w - target_width) // 2
-#         filled_rgb = filled_rgb[:, left:left + target_width]
-#     return Image.fromarray(filled_rgb)
-
-
-# def center_crop_lr_to_768x1024(arr: np.ndarray) -> np.ndarray:
-#     target_h, target_w = 1024, 768
-#     h, w = arr.shape[:2]
-#     if h != target_h:
-#         arr = cv2.resize(arr, (w, target_h), interpolation=cv2.INTER_AREA)
-#         h, w = arr.shape[:2]
-#     if w < target_w:
-#         pad = (target_w - w) // 2
-#         arr = cv2.copyMakeBorder(arr, 0, 0, pad, pad, cv2.BORDER_CONSTANT, value=[255, 255, 255])
-#         w = arr.shape[1]
-#     left = (w - target_w) // 2
-#     return arr[:, left:left + target_w]
-
-
-# def save_cropped(imgs, out_path: str):
-#     np_imgs = [np.asarray(im) for im in imgs]
-#     cropped = [center_crop_lr_to_768x1024(x) for x in np_imgs]
-#     out = np.concatenate(cropped, axis=1)
-#     os.makedirs(os.path.dirname(out_path), exist_ok=True)
-#     imageio.imsave(out_path, out)
-
-# from diffusers import UniPCMultistepScheduler, AutoencoderKL, UNet2DConditionModel
-
-# @lru_cache(maxsize=1)
-# def get_pipe_and_device() -> Tuple[StableDiffusionXLControlNetImg2ImgPipeline, str, torch.dtype]:
-#     device = "cuda" if torch.cuda.is_available() else "cpu"
-#     dtype = torch.float32 if device == "cuda" else torch.float32
-
-#     print(f"[PIPE] device={device}, dtype={dtype}", flush=True)
-
-#     # ControlNet
-#     controlnet = ControlNetModel.from_pretrained(
-#         CONTROLNET_ID,
-#         torch_dtype=dtype,
-#         use_safetensors=True,
-#     ).to(device)
-
-#     # ✅ VAE: safetensors 강제 로드 후 주입
-#     vae = AutoencoderKL.from_pretrained(
-#         BASE_MODEL_ID,
-#         subfolder="vae",
-#         torch_dtype=dtype,
-#         use_safetensors=True,
-#     ).to(device)
-
-#     # ✅ UNet: safetensors 강제 로드 후 주입
-#     unet = UNet2DConditionModel.from_pretrained(
-#         BASE_MODEL_ID,
-#         subfolder="unet",
-#         torch_dtype=dtype,
-#         use_safetensors=True,
-#     ).to(device)
-
-#     # Pipeline
-#     pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-#         BASE_MODEL_ID,
-#         controlnet=controlnet,
-#         vae=vae,
-#         unet=unet,
-#         torch_dtype=dtype,
-#         use_safetensors=True,
-#         add_watermarker=False,
-#     ).to(device)
-
-#     # dtype mismatch 방지(vae)
-#     if device == "cuda":
-#         try:
-#             pipe.vae.to(dtype=dtype)
-#             if hasattr(pipe.vae, "config") and hasattr(pipe.vae.config, "force_upcast"):
-#                 pipe.vae.config.force_upcast = False
-#         except Exception as e:
-#             print("[PIPE] VAE dtype cast failed:", repr(e), flush=True)
-
-#     pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-#     pipe.enable_attention_slicing()
-#     try:
-#         pipe.enable_xformers_memory_efficient_attention()
-#     except Exception as e:
-#         print("[PIPE] xformers not enabled:", repr(e), flush=True)
-
-#     return pipe, device, dtype
-
-
-
-
-# def run_one(paths: Paths, prompt: str, steps: int = DEFAULT_STEPS):
-#     global H, W
-#     pipe, device, _dtype = get_pipe_and_device()
-
-#     # lazy assets download here (NOT at import time)
-#     image_encoder_dir, ip_ckpt, schp_ckpt = get_assets()
-
-#     H, W = compute_hw_from_person(paths.person_path)
-
-#     # parsing extractor (원본 호출 형태 유지)
-#     res = run_simple_extractor(
-#         category="Upper-clothes",
-#         input_path=os.path.abspath(paths.person_path),
-#         model_restore=schp_ckpt,
-#     )
-#     parsing_img = res["images"][0] if res.get("images") else None
-#     if parsing_img is None:
-#         raise RuntimeError("run_simple_extractor returned no parsing images.")
-
-#     sketch_area = fill_sketch_from_image_path_to_pil(paths.depth_path)
-#     sketch_area_inv = invert_sketch_area(sketch_area)
-#     merged_img = merge_white_regions_or(parsing_img, sketch_area_inv)
-#     mask_pil = preprocess_mask(merged_img)
-
-#     # person resize + width=1024 pad/crop
-#     person_bgr = _imread_or_raise(paths.person_path)
-#     person_bgr = cv2.resize(person_bgr, (W, H), interpolation=cv2.INTER_AREA)
-
-#     target_width = 1024
-#     cur_w = person_bgr.shape[1]
-#     if cur_w < target_width:
-#         pad = (target_width - cur_w) // 2
-#         padded_person = cv2.copyMakeBorder(
-#             person_bgr, 0, 0, pad, pad,
-#             borderType=cv2.BORDER_CONSTANT, value=[255, 255, 255]
-#         )
-#     elif cur_w > target_width:
-#         left = (cur_w - target_width) // 2
-#         padded_person = person_bgr[:, left:left + target_width]
-#     else:
-#         padded_person = person_bgr
-
-#     person_rgb = cv2.cvtColor(padded_person, cv2.COLOR_BGR2RGB)
-#     person_pil = Image.fromarray(person_rgb)
-
-#     depth_map = make_depth(paths.depth_path)
-
-#     # garment / garment_mask
-#     garment_pil = person_pil.copy()
-
-#     gm = np.array(parsing_img.convert("L"), dtype=np.uint8)
-#     gm = cv2.resize(gm, (W, H), interpolation=cv2.INTER_AREA)
-#     gm = cv2.cvtColor(gm, cv2.COLOR_GRAY2RGB)
-#     cur_w2 = gm.shape[1]
-#     if cur_w2 < target_width:
-#         pad2 = (target_width - cur_w2) // 2
-#         gm = cv2.copyMakeBorder(gm, 0, 0, pad2, pad2, cv2.BORDER_CONSTANT, value=[0, 0, 0])
-#     elif cur_w2 > target_width:
-#         left2 = (cur_w2 - target_width) // 2
-#         gm = gm[:, left2:left2 + target_width]
-#     garment_mask_pil = Image.fromarray(gm)
-
-#     ip_model = IPAdapterXL(
-#         pipe,
-#         image_encoder_dir,
-#         ip_ckpt,
-#         device,
-#         mask_pil,
-#         person_pil,
-#         content_scale=0.3,
-#         style_scale=0.5,
-#         garment_images=garment_pil,
-#         garment_mask=garment_mask_pil,
-#     )
-    
- 
-#     if device == "cuda":
-#         pipe.to(dtype=torch.float32)
-#         try:
-#             for _, proc in pipe.unet.attn_processors.items():
-#                 proc.to(dtype=torch.float32)
-#         except Exception:
-#             pass
-
-
-
-
-#     style_img = Image.open(paths.style_path).convert("RGB")
-    
-#     person_pil.save('./person_pil.png')
-#     person_pil.save('./mask_pil.png')
-#     depth_map.save('./depth_map.png')
-#     garment_pil.save('./garment_pil.png')
-#     garment_mask_pil.save('./garment_mask_pil.png')
-
-#     with torch.inference_mode():
-#         images = ip_model.generate(
-#             pil_image=style_img,
-#             image=person_pil,
-#             control_image=depth_map,
-#             strength=1.0,
-#             num_samples=1,
-#             num_inference_steps=int(steps),
-#             shape_prompt="",
-#             prompt=prompt or "",
-#             num=0,
-#             scale=None,
-#             controlnet_conditioning_scale=0.7,
-#             guidance_scale=7.5,
-#         )
-
-#     save_cropped(images, paths.output_path)
-
-
-# def set_seed(seed: int):
-#     if seed is None or seed < 0:
-#         return
-#     np.random.seed(seed)
-#     torch.manual_seed(seed)
-#     if torch.cuda.is_available():
-#         torch.cuda.manual_seed_all(seed)
-
-
-# def infer_web(person_fp, sketch_fp, style_fp, prompt, steps, seed):
-#     if person_fp is None or sketch_fp is None or style_fp is None:
-#         raise gr.Error("person / sketch / style 이미지를 모두 업로드해야 합니다.")
-
-#     set_seed(int(seed) if seed is not None else -1)
-
-#     tmp_dir = tempfile.mkdtemp(prefix="vista_demo_")
-#     out_path = os.path.join(tmp_dir, "result.png")
-
-#     paths = Paths(
-#         person_path=person_fp,
-#         depth_path=sketch_fp,
-#         style_path=style_fp,
-#         output_path=out_path,
-#     )
-#     run_one(paths, prompt=prompt, steps=int(steps))
-
-#     out_img = Image.open(out_path).convert("RGB")
-#     return out_img, out_path
-
-
-# with gr.Blocks(title="VISTA Demo (HF Spaces)") as demo:
-#     gr.Markdown("## VISTA Demo\nperson / sketch(guide) / style 입력으로 결과를 생성합니다.")
-
-#     with gr.Row():
-#         person_in = gr.Image(label="Person Image", type="filepath")
-#         sketch_in = gr.Image(label="Sketch / Guide (depth_path)", type="filepath")
-#         style_in = gr.Image(label="Style Image", type="filepath")
-
-#     with gr.Row():
-#         prompt_in = gr.Textbox(label="Prompt", value="upper garment", lines=2)
-#         steps_in = gr.Slider(1, 80, value=DEFAULT_STEPS, step=1, label="Steps")
-#         seed_in = gr.Number(label="Seed (-1=random)", value=-1, precision=0)
-
-#     run_btn = gr.Button("Run")
-#     out_img = gr.Image(label="Output", type="pil")
-#     out_file = gr.File(label="Download result.png")
-
-#     run_btn.click(
-#         fn=infer_web,
-#         inputs=[person_in, sketch_in, style_in, prompt_in, steps_in, seed_in],
-#         outputs=[out_img, out_file],
-#     )
-
-# demo.queue()
-# if __name__ == "__main__":
-#     demo.launch(server_name="0.0.0.0", server_port=7860)
-
 import os
 import sys
 
@@ -594,6 +98,18 @@ def _imread_or_raise(path: str, flag=cv2.IMREAD_COLOR):
         raise FileNotFoundError(f"cv2.imread failed: {path} (exists={os.path.exists(path)})")
     return img
 
+def apply_parsing_white_mask_to_person_cv2(
+    person_pil: Image.Image,
+    parsing_img: Image.Image
+) -> np.ndarray:
+    person_rgb = np.array(person_pil.convert("RGB"), dtype=np.uint8)
+    mask = np.array(parsing_img.convert("L"), dtype=np.uint8)
+    white_mask = mask == 255
+    result_rgb = np.full_like(person_rgb, 255, dtype=np.uint8)
+    result_rgb[white_mask] = person_rgb[white_mask]
+    result_bgr = cv2.cvtColor(result_rgb, cv2.COLOR_RGB2BGR)
+    return result_bgr
+
 
 def compute_hw_from_person(person_path: str):
     img = _imread_or_raise(person_path)
@@ -824,6 +340,25 @@ def run_one(paths: Paths, prompt: str, steps: int = DEFAULT_STEPS):
     depth_map = make_depth(paths.depth_path)
 
     garment_pil = person_pil.copy()
+    
+    garment_ = apply_parsing_white_mask_to_person_cv2(
+        garment_pil,
+        parsing_img
+    )
+
+    garment_rgb = cv2.cvtColor(garment_, cv2.COLOR_BGR2RGB)
+
+    # ✅ (중요) garment_는 원본 person 크기일 수 있으니 전역 (W,H)로 맞춘 뒤 padding
+    garment_rgb = cv2.resize(garment_rgb, (W, H), interpolation=cv2.INTER_AREA)
+
+    garment_rgb = cv2.copyMakeBorder(
+        garment_rgb,
+        top=0, bottom=0,
+        left=padding, right=padding,
+        borderType=cv2.BORDER_CONSTANT,
+        value=[255, 255, 255],
+    )
+    garment_pil = Image.fromarray(garment_rgb)
 
     gm = np.array(parsing_img.convert("L"), dtype=np.uint8)
     gm = cv2.resize(gm, (W, H), interpolation=cv2.INTER_AREA)