app.py

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

# Functions used by runtime.load()/run(); implemented in CatVTON/utils.py.
# The runtime sys.path injection makes `CatVTON/utils.py` importable as `utils`.
from CatVTON.utils import init_weight_dtype, resize_and_crop, resize_and_padding

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"))
DEFAULT_DEVICE = os.getenv("CATVTON_DEVICE", "cuda")
# If container has no CUDA driver (CPU-only), fall back to CPU to avoid crash.
DEVICE = DEFAULT_DEVICE
if DEVICE.startswith("cuda"):
    # CPU-only Spaces may still have torch built with CUDA, but no driver at runtime.
    # Robustly detect that case and fall back to CPU.
    if not torch.cuda.is_available():
        DEVICE = "cpu"
    else:
        try:
            _ = torch.cuda.current_device()
        except Exception:
            DEVICE = "cpu"

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

        # Help debug missing code/weights on HF Spaces.
        # (This will show up in Space logs during first request/build.)
        # NOTE: keep as lightweight prints.
        print("[CatVTON] load(): repo_dir=", self.repo_dir)
        print("[CatVTON] CATVTON_RESUME_PATH=", CATVTON_RESUME_PATH)
        print("[CatVTON] CATVTON_MODEL_DIR=", os.getenv("CATVTON_MODEL_DIR"))


        if not self.repo_dir.exists():
            raise RuntimeError(f"CatVTON repository not found at '{self.repo_dir}'.")

        # --- Resolve real python root that contains `model/` ---
        # Some HF environments mount the code differently; env/debug values can be wrong.
        # We therefore detect the repo root by searching for `model/cloth_masker.py`.

        # Force importability in HF Spaces: repo root is typically the working directory.
        # Ensure both repo root and repo_root/CatVTON are importable.
        repo_root = Path.cwd().resolve()
        # Ensure that imports like `import model.*` work in HF.
        # Different Space layouts may put the actual CatVTON code under:
        #   - ./CatVTON/model/...
        #   - ./model/...
        #   - /app/CatVTON/model/...
        candidate_code_roots = [
            repo_root,
            repo_root / "CatVTON",
            self.repo_dir,
            self.repo_dir / "CatVTON",
        ]
        for p in candidate_code_roots:
            ps = str(p)
            if ps not in sys.path and p.exists():
                sys.path.insert(0, ps)

        # Also add the directory that directly contains `model/` if present.
        direct_model_root = None
        for p in candidate_code_roots:
            if (p / "model" / "cloth_masker.py").exists():
                direct_model_root = p
                break
        if direct_model_root is not None:
            dm = str(direct_model_root)
            if dm not in sys.path:
                sys.path.insert(0, dm)


        found_model_parent: Path | None = None

        # Search for CatVTON's `model/` package starting from the working directory.
        # (Avoid scanning large absolute paths like /app that may not exist in the container.)
        for cloth_masker in repo_root.rglob("model/cloth_masker.py"):
            candidate_root = cloth_masker.parent.parent  # .../<repo_root>/model/cloth_masker.py
            if (candidate_root / "model" / "pipeline.py").exists():
                found_model_parent = candidate_root.resolve()
                break

        if found_model_parent is None:
            # Keep existing behavior as last resort.
            found_model_parent = self.repo_dir.resolve()

        repo_path = str(found_model_parent)
        if repo_path not in sys.path:
            sys.path.insert(0, repo_path)

        self.repo_dir = found_model_parent

        # If CatVTON code is under `<repo_root>/CatVTON/` then `import model.*` expects
        # `sys.path` to include that inner code root (so `model/` is importable).
        # Ensure this regardless of which candidate_root was selected.
        inner_code_root = repo_root / "CatVTON"
        if (inner_code_root / "model" / "cloth_masker.py").exists():
            sys.path.insert(0, str(inner_code_root.resolve()))
        else:
            # If the HF Space layout is different, fall back to adding `repo_root/model`.
            fallback_model_root = repo_root / "model"
            if (fallback_model_root / "cloth_masker.py").exists():
                sys.path.insert(0, str(fallback_model_root.resolve().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


if __name__ == "__main__":
    # HF debugging: confirm CatVTON code presence inside container
    _p1 = Path('/app/CatVTON/model/cloth_masker.py')
    _p2 = Path('/app/CatVTON/model/pipeline.py')
    _p3 = Path('./CatVTON/model/cloth_masker.py')
    print('[HF Debug] /app/CatVTON/model/cloth_masker.py exists:', _p1.exists())
    print('[HF Debug] /app/CatVTON/model/pipeline.py exists:', _p2.exists())
    print('[HF Debug] ./CatVTON/model/cloth_masker.py exists:', _p3.exists())

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)