app.py

"""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())