"""SIFT-VTON — Geometric Correspondence Supervision on Cross-Attention for Virtual Try-On. Faithful port of the authors' `inference_hf.py` (github.com/takesukeDS/SIFT-VTON) to a Gradio / ZeroGPU Space, with the VITON-HD preprocessing chain (human parsing -> agnostic mask, OpenPose, DensePose) run on-the-fly so arbitrary photos can be used as input. """ import os import sys import time import spaces # noqa: F401 (must precede any torch/CUDA import) import cv2 import gradio as gr import numpy as np import torch from huggingface_hub import hf_hub_download from omegaconf import OmegaConf from PIL import Image sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) REPO_ID = "takesuke/SIFT-VTON" PREPROC_REPO = "yisol/IDM-VTON" # host of the DensePose / parsing / OpenPose checkpoints IMG_H, IMG_W = 512, 384 LATENT_SIZES = [(16, 12), (32, 24), (64, 48)] # -------------------------------------------------------------------------------------- # checkpoints # -------------------------------------------------------------------------------------- print("Fetching checkpoints ...", flush=True) CONFIG_PATH = hf_hub_download(REPO_ID, "config.yaml") WEIGHTS_PATH = hf_hub_download(REPO_ID, "model.ckpt") DENSEPOSE_CKPT = hf_hub_download(PREPROC_REPO, "densepose/model_final_162be9.pkl") ATR_ONNX = hf_hub_download(PREPROC_REPO, "humanparsing/parsing_atr.onnx") LIP_ONNX = hf_hub_download(PREPROC_REPO, "humanparsing/parsing_lip.onnx") BODY_POSE = hf_hub_download(PREPROC_REPO, "openpose/ckpts/body_pose_model.pth") # the vendored OpenposeDetector resolves its checkpoint relative to the app root _op_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ckpt", "openpose", "ckpts") os.makedirs(_op_dir, exist_ok=True) _op_dst = os.path.join(_op_dir, "body_pose_model.pth") if not os.path.exists(_op_dst): os.symlink(BODY_POSE, _op_dst) # -------------------------------------------------------------------------------------- # SIFT-VTON model # -------------------------------------------------------------------------------------- from cldm.model import create_model # noqa: E402 from cldm.plms_hacked import PLMSSampler # noqa: E402 from utils import tensor2img # noqa: E402 print("Building SIFT-VTON ...", flush=True) config = OmegaConf.load(CONFIG_PATH) config.model.params.img_H = IMG_H config.model.params.img_W = IMG_W config.model.params.unet_config.params.use_sift_loss = False config.model.params.unet_config.params.use_checkpoint = False # no-op at inference, faster model = create_model(config_path=None, config=config) _ck = torch.load(WEIGHTS_PATH, map_location="cpu", mmap=True, weights_only=False) _sd = _ck["state_dict"] if "state_dict" in _ck else _ck # The checkpoint was saved with an old `transformers`, whose CLIPVisionModel nested its # weights under `.vision_model.`; recent releases flattened that away. Re-key so the frozen # CLIP image tower still receives the checkpoint's own weights, and drop buffers (e.g. # `position_ids`) that no longer exist as state. _tgt = model.state_dict() _aligned, _dropped, _renamed = {}, 0, 0 for _k, _v in _sd.items(): if _k not in _tgt: _k2 = _k.replace("cond_stage_model.transformer.vision_model.", "cond_stage_model.transformer.") if _k2 in _tgt: _k = _k2 _renamed += 1 else: _dropped += 1 continue _aligned[_k] = _v _missing, _unexpected = model.load_state_dict(_aligned, strict=False) # the frozen CLIP tower is already correctly initialised from `openai/clip-vit-large-patch14` _missing = [k for k in _missing if not k.startswith("cond_stage_model.transformer.")] print(f"state_dict: renamed={_renamed} dropped={_dropped} missing={len(_missing)}", flush=True) if _missing: raise RuntimeError(f"missing {len(_missing)} weights, e.g. {_missing[:8]}") del _ck, _sd, _aligned, _tgt model = model.cuda().eval() sampler = PLMSSampler(model) print("SIFT-VTON ready", flush=True) # -------------------------------------------------------------------------------------- # preprocessing models (DensePose / parsing / OpenPose) # -------------------------------------------------------------------------------------- from detectron2.data.detection_utils import _apply_exif_orientation, convert_PIL_to_numpy # noqa: E402 from detectron2.engine.defaults import DefaultPredictor # noqa: E402 import apply_net # noqa: E402 from preprocess.humanparsing.run_parsing import Parsing # noqa: E402 from preprocess.openpose.run_openpose import OpenPose # noqa: E402 from utils_mask import get_mask_location # noqa: E402 print("Building preprocessors ...", flush=True) _dp_args = apply_net.create_argument_parser().parse_args( ( "show", "./configs/densepose_rcnn_R_50_FPN_s1x.yaml", DENSEPOSE_CKPT, "dp_segm", "-v", "--opts", "MODEL.DEVICE", "cuda", ) ) _dp_cfg = apply_net.ShowAction.setup_config(_dp_args.cfg, _dp_args.model, _dp_args, []) densepose_predictor = DefaultPredictor(_dp_cfg) densepose_context = apply_net.ShowAction.create_context(_dp_args, _dp_cfg) parsing_model = Parsing(ATR_ONNX, LIP_ONNX) openpose_model = OpenPose(0) openpose_model.preprocessor.body_estimation.model = ( openpose_model.preprocessor.body_estimation.model.cuda() ) print("Preprocessors ready", flush=True) # -------------------------------------------------------------------------------------- # helpers # -------------------------------------------------------------------------------------- def _center_crop_3_4(img: Image.Image) -> Image.Image: w, h = img.size tw = int(min(w, h * (3 / 4))) th = int(min(h, w * (4 / 3))) left, top = (w - tw) / 2, (h - th) / 2 return img.crop((left, top, left + tw, top + th)) def _norm(img_uint8: np.ndarray) -> np.ndarray: """uint8 RGB HWC -> float32 [-1, 1].""" return img_uint8.astype(np.float32) / 127.5 - 1.0 def build_densepose(person_rgb: Image.Image) -> np.ndarray: """VITON-HD style `image-densepose` (fine-segmentation visualisation), uint8 RGB.""" arg_img = _apply_exif_orientation(person_rgb) arg_img = convert_PIL_to_numpy(arg_img, format="BGR") with torch.no_grad(): outputs = densepose_predictor(arg_img)["instances"] vis = apply_net.ShowAction.execute_on_outputs( densepose_context, {"image": arg_img}, outputs ) return np.ascontiguousarray(vis[:, :, ::-1]) # BGR -> RGB def build_batch(person_rgb: Image.Image, garment_rgb: Image.Image, mask_pil: Image.Image, densepose_rgb: np.ndarray): person = np.array(person_rgb) # (512, 384, 3) uint8 garment = np.array(garment_rgb) inpaint = (np.array(mask_pil.convert("L")) >= 128).astype(np.float32)[:, :, None] keep = 1.0 - inpaint # dataset's `agn_mask` convention agn = _norm(person) * keep # masked region -> mid gray (0.0) batch = { "image": _norm(person), "agn": agn, "agn_mask": keep, "agn_mask_orig": keep, "image_densepose": _norm(densepose_rgb), "cloth": _norm(garment), "cloth_mask": np.ones((IMG_H, IMG_W, 1), dtype=np.float32), "gt_cloth_warped_mask": np.zeros((IMG_H, IMG_W, 1), dtype=np.float32), } batch = {k: torch.from_numpy(v).float().unsqueeze(0).cuda() for k, v in batch.items()} # SIFT correspondence histograms are a training-time signal only (zeros at inference) for (h, w), key in zip(LATENT_SIZES, ["hist16", "hist32", "hist64"]): batch[key] = torch.zeros(1, h, w, h, w, dtype=torch.float32).cuda() batch[key + "_mask"] = torch.zeros(1, h, w, dtype=torch.float32).cuda() batch["txt"] = [""] batch["img_fn"] = ["person.jpg"] batch["cloth_fn"] = ["cloth.jpg"] return batch # -------------------------------------------------------------------------------------- # inference # -------------------------------------------------------------------------------------- def _estimate_duration(person_image=None, garment_image=None, denoise_steps=50, *args, **kwargs): # measured on ZeroGPU: ~5 s preprocessing + ~0.16 s / PLMS step, plus 40 % headroom return int(min(90, (6 + 0.16 * float(denoise_steps)) * 1.4)) @spaces.GPU(duration=_estimate_duration) def try_on( person_image, garment_image, denoise_steps: int = 50, cfg_scale: float = 1.5, seed: int = 1235, auto_crop: bool = True, repaint: bool = True, progress=gr.Progress(track_tqdm=True), ): if person_image is None or garment_image is None: raise gr.Error("Please provide both a person image and a garment image.") _t0 = time.perf_counter() with torch.no_grad(): out = _try_on(person_image, garment_image, denoise_steps, cfg_scale, seed, auto_crop, repaint) print(f"[try_on] {denoise_steps} steps in {time.perf_counter() - _t0:.2f}s", flush=True) return out def _try_on(person_image, garment_image, denoise_steps, cfg_scale, seed, auto_crop, repaint): person = person_image.convert("RGB") if auto_crop: person = _center_crop_3_4(person) person = person.resize((IMG_W, IMG_H), Image.LANCZOS) garment = garment_image.convert("RGB").resize((IMG_W, IMG_H), Image.LANCZOS) # --- VITON-HD preprocessing --------------------------------------------------- keypoints = openpose_model(person) model_parse, _ = parsing_model(person) mask_pil, _ = get_mask_location("hd", "upper_body", model_parse, keypoints, width=IMG_W, height=IMG_H) densepose_rgb = build_densepose(person) # --- sampling (mirrors inference_hf.py) ---------------------------------------- torch.manual_seed(int(seed)) np.random.seed(int(seed) % (2**32)) batch = build_batch(person, garment, mask_pil, densepose_rgb) z, c = model.get_input(batch, config.model.params.first_stage_key) bs = z.shape[0] c_crossattn = c["c_crossattn"][0][:bs] if c_crossattn.ndim == 4: c["c_crossattn"] = [model.get_learned_conditioning(c_crossattn)] uc_full = { "c_concat": None, "c_crossattn": [model.learnable_vector.repeat(bs, 1, 1)], "first_stage_cond": c["first_stage_cond"], } sampler.model.batch = batch ts = torch.full((1,), 999, device=z.device, dtype=torch.long) start_code = model.q_sample(c["first_stage_cond"][:, :4], ts) # --start_from_noised_agn samples, _, _ = sampler.sample( int(denoise_steps), bs, (4, IMG_H // 8, IMG_W // 8), c, x_T=start_code, verbose=False, eta=0.0, unconditional_guidance_scale=float(cfg_scale), unconditional_conditioning=uc_full, ) x_samples = model.decode_first_stage(samples) out = tensor2img(x_samples[0], round=True) if repaint: orig = np.uint8((batch["image"][0].cpu().numpy() + 1) / 2 * 255 + 0.5) keep = batch["agn_mask_orig"][0].cpu().numpy() out = np.uint8(orig * keep + out * (1 - keep) + 0.5) masked_preview = Image.fromarray( np.uint8((batch["agn"][0].cpu().numpy() + 1) / 2 * 255 + 0.5) ) return Image.fromarray(out), masked_preview, Image.fromarray(densepose_rgb) # -------------------------------------------------------------------------------------- # UI # -------------------------------------------------------------------------------------- PERSONS = "examples/person" GARMENTS = "examples/garment" EXAMPLES = [ [f"{PERSONS}/00034_00.jpg", f"{GARMENTS}/04469_00.jpg"], [f"{PERSONS}/00035_00.jpg", f"{GARMENTS}/09133_00.jpg"], [f"{PERSONS}/00055_00.jpg", f"{GARMENTS}/09266_00.jpg"], [f"{PERSONS}/01992_00.jpg", f"{GARMENTS}/14673_00.jpg"], [f"{PERSONS}/00121_00.jpg", f"{GARMENTS}/09163_00.jpg"], ] DESCRIPTION = """ # SIFT-VTON — Virtual Try-On with Geometric Correspondence Supervision Upload a **person** photo and an **upper-body garment**, and SIFT-VTON dresses the person in it. [Paper](https://huggingface.co/papers/2605.01296) · [Model](https://huggingface.co/takesuke/SIFT-VTON) · [Code](https://github.com/takesukeDS/SIFT-VTON) SIFT-VTON (ICPR 2026) supervises the cross-attention maps of a StableVITON-style diffusion try-on model with SIFT correspondences between garment and person, which sharpens the geometric alignment of the transferred garment. Trained on VITON-HD at 512×384 — front-facing, full-torso shots on a plain background work best. """ with gr.Blocks(title="SIFT-VTON") as demo: gr.Markdown(DESCRIPTION) with gr.Row(): with gr.Column(): person_image = gr.Image(label="Person", type="pil", height=420) with gr.Column(): garment_image = gr.Image(label="Garment (upper body)", type="pil", height=420) with gr.Column(): result = gr.Image(label="Try-on result", type="pil", height=420) run_button = gr.Button("Try it on", variant="primary") with gr.Accordion("Advanced options", open=False): with gr.Row(): denoise_steps = gr.Slider(10, 100, value=50, step=1, label="PLMS denoising steps") cfg_scale = gr.Slider(1.0, 5.0, value=1.5, step=0.1, label="Classifier-free guidance scale") with gr.Row(): seed = gr.Slider(0, 2**31 - 1, value=1235, step=1, label="Seed") auto_crop = gr.Checkbox(value=True, label="Auto-crop person to 3:4") repaint = gr.Checkbox(value=True, label="Repaint (keep unmasked pixels)") with gr.Row(): agnostic_out = gr.Image(label="Agnostic person (masked input)", type="pil") densepose_out = gr.Image(label="DensePose", type="pil") gr.Examples( examples=EXAMPLES, inputs=[person_image, garment_image], outputs=[result, agnostic_out, densepose_out], fn=try_on, cache_examples=True, cache_mode="lazy", ) gr.Markdown( "Example person / garment images are VITON-HD test-set items redistributed from the " "[IDM-VTON Space](https://huggingface.co/spaces/yisol/IDM-VTON) under CC BY-NC-SA 4.0. " "The DensePose, human-parsing and OpenPose preprocessing checkpoints are loaded from " "[yisol/IDM-VTON](https://huggingface.co/yisol/IDM-VTON)." ) run_button.click( fn=try_on, inputs=[person_image, garment_image, denoise_steps, cfg_scale, seed, auto_crop, repaint], outputs=[result, agnostic_out, densepose_out], ) if __name__ == "__main__": demo.queue().launch(theme=gr.themes.Citrus())