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 # ...//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 # ...//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 `/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)