test

.history/app_20260617193632.py

@@ -0,0 +1,332 @@
+from __future__ import annotations
+
+import os
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+import gradio as gr
+import numpy as np
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from huggingface_hub import snapshot_download
+from PIL import Image, ImageOps
+
+APP_TITLE = "ChitraTech Virtual Try-On"
+APP_DESCRIPTION = (
+    "Upload a shopper photo and clothing image to run on-demand CatVTON virtual try-on inference "
+    "using the Zheng-Chong CatVTON implementation."
+)
+CATVTON_REPO_DIR_ENV = os.getenv("CATVTON_REPO_DIR")
+CATVTON_REPO_DIR = Path(CATVTON_REPO_DIR_ENV) if CATVTON_REPO_DIR_ENV else Path("./CatVTON")
+CATVTON_RESUME_PATH = os.getenv("CATVTON_RESUME_PATH", "zhengchong/CatVTON")
+
+
+def resolve_catvton_repo_dir(start_dir: Path) -> Path:
+    """Find the CatVTON repo root that contains `model/cloth_masker.py`.
+
+    HF Spaces sometimes mount code in unexpected places; relying on fixed paths like
+    `/app/CatVTON` can be wrong. We therefore:
+    1) try a few common candidates
+    2) then scan under `/app` (and `/workspace` if present) for `model/cloth_masker.py`
+    """
+
+    def looks_like_repo_dir(p: Path) -> bool:
+        return (p / "model" / "cloth_masker.py").exists() and (p / "model" / "pipeline.py").exists()
+
+    candidates: list[Path] = []
+
+    if start_dir is not None:
+        candidates.append(start_dir)
+
+    if CATVTON_REPO_DIR_ENV:
+        candidates.append(Path(CATVTON_REPO_DIR_ENV))
+
+    candidates.extend([
+        Path("/app/CatVTON"),
+        Path("/app"),
+        Path("./CatVTON"),
+        Path("./"),
+        Path("/workspace"),
+    ])
+
+    for c in candidates:
+        if c is not None and looks_like_repo_dir(c):
+            return c.resolve()
+
+    # Broad scan for the actual code root.
+    scan_roots = [Path("/app"), Path("/workspace")]
+    for root in scan_roots:
+        if not root.exists():
+            continue
+        for cloth_masker in root.rglob("model/cloth_masker.py"):
+            repo_root = cloth_masker.parent.parent  # .../<repo_root>/model/cloth_masker.py
+            if looks_like_repo_dir(repo_root):
+                return repo_root.resolve()
+
+    # Fallback: return the provided start_dir (so error message includes candidates).
+    return start_dir.resolve()
+
+
+CATVTON_BASE_MODEL = os.getenv("CATVTON_BASE_MODEL", "booksforcharlie/stable-diffusion-inpainting")
+CATVTON_OUTPUT_DIR = Path(os.getenv("CATVTON_OUTPUT_DIR", "./outputs"))
+DEVICE = os.getenv("CATVTON_DEVICE", "cuda")
+DEFAULT_WIDTH = int(os.getenv("CATVTON_WIDTH", "768"))
+DEFAULT_HEIGHT = int(os.getenv("CATVTON_HEIGHT", "1024"))
+DEFAULT_STEPS = int(os.getenv("CATVTON_STEPS", "50"))
+DEFAULT_GUIDANCE_SCALE = float(os.getenv("CATVTON_GUIDANCE_SCALE", "2.5"))
+DEFAULT_MIXED_PRECISION = os.getenv("CATVTON_MIXED_PRECISION", "bf16")
+DEFAULT_SEED = int(os.getenv("CATVTON_SEED", "42"))
+
+
+@dataclass
+class CatVTONRuntime:
+    repo_dir: Path
+    device: str
+    pipeline: object | None = field(default=None, init=False, repr=False)
+    automasker: object | None = field(default=None, init=False, repr=False)
+    mask_processor: object | None = field(default=None, init=False, repr=False)
+    resize_and_crop: object | None = field(default=None, init=False, repr=False)
+    resize_and_padding: object | None = field(default=None, init=False, repr=False)
+    vis_mask: object | None = field(default=None, init=False, repr=False)
+    ready: bool = False
+    status: str = "not loaded"
+
+    def load(self) -> None:
+        if self.ready:
+            return
+
+        if not self.repo_dir.exists():
+            raise RuntimeError(f"CatVTON repository not found at '{self.repo_dir}'.")
+
+        repo_path = str(self.repo_dir.resolve())
+        if repo_path not in sys.path:
+            sys.path.insert(0, repo_path)
+
+        # Ensure CatVTON's `model/` package is importable as `import model.*`.
+        # In many deployments the code is flattened under `/app`, so `CATVTON_REPO_DIR`
+        # might be relative and/or point to the wrong nesting level.
+        model_dir = (self.repo_dir / "model").resolve()
+        if model_dir.exists():
+            # Add the parent so `model` becomes a top-level module
+            model_parent = str(model_dir.parent)
+            if model_parent not in sys.path:
+                sys.path.insert(0, model_parent)
+
+        # Extra fallback: if repo_dir/model doesn't exist, try sibling 'model' under repo_path
+        if not (self.repo_dir / "model").exists():
+            guessed = (Path(repo_path) / "model").resolve()
+            if guessed.exists():
+                guessed_parent = str(guessed.parent)
+                if guessed_parent not in sys.path:
+                    sys.path.insert(0, guessed_parent)
+
+
+        # If this still fails inside HF, add debugging info.
+        try:
+            from model.cloth_masker import AutoMasker, vis_mask
+            from model.pipeline import CatVTONPipeline
+        except Exception as import_exc:
+            # Helpful diagnostics for HF Spaces.
+            repo_model_exists = (self.repo_dir / "model").exists()
+            candidate_roots = [
+                self.repo_dir,
+                self.repo_dir / "model",
+                (self.repo_dir / "model").parent,
+            ]
+            candidate_roots_str = ", ".join(str(p) for p in candidate_roots)
+
+            raise RuntimeError(
+                "CatVTON import failed. "
+                f"repo_dir={self.repo_dir} "
+                f"repo_dir/model_exists={repo_model_exists} "
+                f"repo_model_candidate_roots={candidate_roots_str} "
+                f"sys.path[0:10]={sys.path[:10]} "
+                f"import_error={import_exc}"
+            ) from import_exc
+
+
+        repo_weights_dir = Path(snapshot_download(repo_id=CATVTON_RESUME_PATH))
+        self.pipeline = CatVTONPipeline(
+            base_ckpt=CATVTON_BASE_MODEL,
+            attn_ckpt=str(repo_weights_dir),
+            attn_ckpt_version="mix",
+            weight_dtype=init_weight_dtype(DEFAULT_MIXED_PRECISION),
+            use_tf32=True,
+            device=self.device,
+        )
+        self.mask_processor = VaeImageProcessor(
+            vae_scale_factor=8,
+            do_normalize=False,
+            do_binarize=True,
+            do_convert_grayscale=True,
+        )
+        self.automasker = AutoMasker(
+            densepose_ckpt=os.path.join(repo_weights_dir, "DensePose"),
+            schp_ckpt=os.path.join(repo_weights_dir, "SCHP"),
+            device=self.device,
+        )
+        self.resize_and_crop = resize_and_crop
+        self.resize_and_padding = resize_and_padding
+        self.vis_mask = vis_mask
+        CATVTON_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+        self.ready = True
+        self.status = "loaded"
+
+    def run(
+        self,
+        person_image: Image.Image,
+        garment_image: Image.Image,
+        cloth_type: str,
+        num_inference_steps: int,
+        guidance_scale: float,
+        seed: int,
+        show_type: str,
+    ) -> Image.Image:
+        self.load()
+        assert self.pipeline is not None
+        assert self.automasker is not None
+        assert self.mask_processor is not None
+        assert self.resize_and_crop is not None
+        assert self.resize_and_padding is not None
+        assert self.vis_mask is not None
+
+        person_image = self.resize_and_crop(person_image.convert("RGB"), (DEFAULT_WIDTH, DEFAULT_HEIGHT))
+        garment_image = self.resize_and_padding(garment_image.convert("RGB"), (DEFAULT_WIDTH, DEFAULT_HEIGHT))
+
+        generated_mask = self.automasker(person_image, cloth_type)["mask"]
+        generated_mask = self.mask_processor.blur(generated_mask, blur_factor=9)
+
+        generator = None
+        if seed != -1:
+            generator = torch.Generator(device=self.device).manual_seed(seed)
+
+        result_image = self.pipeline(
+            image=person_image,
+            condition_image=garment_image,
+            mask=generated_mask,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        )[0]
+
+        if show_type == "result only":
+            return result_image.convert("RGB")
+
+        masked_person = self.vis_mask(person_image, generated_mask)
+        return compose_preview(person_image, garment_image, masked_person, result_image, show_type)
+
+
+runtime = CatVTONRuntime(repo_dir=resolve_catvton_repo_dir(CATVTON_REPO_DIR), device=DEVICE)
+
+
+
+def prepare_image(image: Image.Image) -> Image.Image:
+    return ImageOps.exif_transpose(image).convert("RGB")
+
+
+def image_grid(images: list[Image.Image], rows: int, cols: int) -> Image.Image:
+    if len(images) != rows * cols:
+        raise ValueError("The number of images does not match the grid shape.")
+    width, height = images[0].size
+    grid = Image.new("RGB", size=(cols * width, rows * height))
+    for index, image in enumerate(images):
+        grid.paste(image, box=(index % cols * width, index // cols * height))
+    return grid
+
+
+def compose_preview(
+    person_image: Image.Image,
+    garment_image: Image.Image,
+    masked_person: Image.Image,
+    result_image: Image.Image,
+    show_type: str,
+) -> Image.Image:
+    width, height = person_image.size
+    if show_type == "input & result":
+        side_panel = image_grid([person_image, garment_image], 2, 1).resize((width // 2, height), Image.NEAREST)
+    else:
+        side_panel = image_grid([person_image, masked_person, garment_image], 3, 1).resize((width // 3, height), Image.NEAREST)
+
+    preview = Image.new("RGB", (side_panel.width + 5 + width, height), color=(255, 255, 255))
+    preview.paste(side_panel, (0, 0))
+    preview.paste(result_image.convert("RGB"), (side_panel.width + 5, 0))
+    return preview
+
+
+def try_on(
+    person_image: Optional[Image.Image],
+    garment_image: Optional[Image.Image],
+    cloth_type: str,
+    num_inference_steps: int,
+    guidance_scale: float,
+    seed: int,
+    show_type: str,
+) -> Image.Image:
+    if person_image is None or garment_image is None:
+        raise gr.Error("Please upload both a shopper photo and a clothing image.")
+
+    prepared_person = prepare_image(person_image)
+    prepared_garment = prepare_image(garment_image)
+
+    try:
+        return runtime.run(
+            person_image=prepared_person,
+            garment_image=prepared_garment,
+            cloth_type=cloth_type,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            seed=seed,
+            show_type=show_type,
+        )
+    except Exception as exc:
+        raise gr.Error(f"CatVTON inference failed: {exc}") from exc
+
+
+with gr.Blocks(theme=gr.themes.Soft(), title=APP_TITLE) as demo:
+    gr.Markdown(f"# {APP_TITLE}")
+    gr.Markdown(APP_DESCRIPTION)
+    gr.Markdown(
+        f"**Runtime:** repo=`{CATVTON_REPO_DIR}` | weights=`{CATVTON_RESUME_PATH}` | device=`{DEVICE}`"
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            person_input = gr.Image(type="pil", label="Shopper photo")
+            garment_input = gr.Image(type="pil", label="Clothing image")
+            cloth_type_input = gr.Radio(
+                label="Garment type",
+                choices=["upper", "lower", "overall"],
+                value="upper",
+            )
+            submit_button = gr.Button("Try On", variant="primary")
+            with gr.Accordion("Advanced options", open=False):
+                step_input = gr.Slider(label="Inference steps", minimum=10, maximum=100, step=5, value=DEFAULT_STEPS)
+                guidance_input = gr.Slider(label="Guidance scale", minimum=0.0, maximum=7.5, step=0.5, value=DEFAULT_GUIDANCE_SCALE)
+                seed_input = gr.Slider(label="Seed", minimum=-1, maximum=10000, step=1, value=DEFAULT_SEED)
+                show_type_input = gr.Radio(
+                    label="Preview mode",
+                    choices=["result only", "input & result", "input & mask & result"],
+                    value="result only",
+                )
+        with gr.Column(scale=1):
+            result_output = gr.Image(type="pil", label="Try-on result")
+
+    gr.Markdown(
+        """
+        ### Notes
+        - This app is just for testing `CatVTON/` codebase.
+        - Model weights are downloaded on demand from Hugging Face using `zhengchong/CatVTON` by default.
+        - Just for testing purposes only.
+        """
+    )
+
+    submit_button.click(
+        fn=try_on,
+        inputs=[person_input, garment_input, cloth_type_input, step_input, guidance_input, seed_input, show_type_input],
+        outputs=result_output,
+    )
+
+
+demo.queue().launch(show_error=True)

app.py

@@ -24,18 +24,25 @@ CATVTON_RESUME_PATH = os.getenv("CATVTON_RESUME_PATH", "zhengchong/CatVTON")
 
 
 def resolve_catvton_repo_dir(start_dir: Path) -> Path:
-    """Try to find the CatVTON repo root that contains `model/`.
+    """Find the CatVTON repo root that contains `model/cloth_masker.py`.
 
-    HF spaces often mount code under `/app`, and env vars can be relative,
-    so we probe a few likely locations.
+    HF Spaces sometimes mount code in unexpected places; relying on fixed paths like
+    `/app/CatVTON` can be wrong. We therefore:
+    1) try a few common candidates
+    2) then scan under `/app` (and `/workspace` if present) for `model/cloth_masker.py`
     """
 
+    def looks_like_repo_dir(p: Path) -> bool:
+        return (p / "model" / "cloth_masker.py").exists() and (p / "model" / "pipeline.py").exists()
+
     candidates: list[Path] = []
 
     if start_dir is not None:
         candidates.append(start_dir)
 
-    # Common HF layout
+    if CATVTON_REPO_DIR_ENV:
+        candidates.append(Path(CATVTON_REPO_DIR_ENV))
+
     candidates.extend([
         Path("/app/CatVTON"),
         Path("/app"),
@@ -44,20 +51,24 @@ def resolve_catvton_repo_dir(start_dir: Path) -> Path:
         Path("/workspace"),
     ])
 
-    # Also allow using the env var directly if it was set
-    if CATVTON_REPO_DIR_ENV:
-        candidates.append(Path(CATVTON_REPO_DIR_ENV))
-
-    # Probe for the *real* python root: folder containing `model/cloth_masker.py`
     for c in candidates:
-        model_root = (c / "model" / "cloth_masker.py")
-        pipeline_root = (c / "model" / "pipeline.py")
-        if model_root.exists() and pipeline_root.exists():
+        if c is not None and looks_like_repo_dir(c):
             return c.resolve()
 
-    # Fallback: return the start directory (and let the existing error message show details)
+    # Broad scan for the actual code root.
+    scan_roots = [Path("/app"), Path("/workspace")]
+    for root in scan_roots:
+        if not root.exists():
+            continue
+        for cloth_masker in root.rglob("model/cloth_masker.py"):
+            repo_root = cloth_masker.parent.parent  # .../<repo_root>/model/cloth_masker.py
+            if looks_like_repo_dir(repo_root):
+                return repo_root.resolve()
+
+    # Fallback: return the provided start_dir (so error message includes candidates).
     return start_dir.resolve()
 
+
 CATVTON_BASE_MODEL = os.getenv("CATVTON_BASE_MODEL", "booksforcharlie/stable-diffusion-inpainting")
 CATVTON_OUTPUT_DIR = Path(os.getenv("CATVTON_OUTPUT_DIR", "./outputs"))
 DEVICE = os.getenv("CATVTON_DEVICE", "cuda")