app.py

import os
from pathlib import Path
import gradio as gr
import numpy as np
from rembg import remove
from PIL import Image, ImageFilter
import shutil


print(f"[DCI] ENABLE_DCI={os.getenv('ENABLE_DCI','0')=='1'}")


ROOT = "/tmp/dci_onepair"
IMG_DIR = f"{ROOT}/test_img"
CLOTH_DIR = f"{ROOT}/test_color"
PAIRS = f"{ROOT}/test_pairs.txt"

def prepare_onepair_dataset(user_img_path: str, cloth_img_path: str):
    # Reset temp dataset
    if os.path.exists(ROOT):
        shutil.rmtree(ROOT)

    # Expected DCI-VTON structure
    img_dir   = os.path.join(ROOT, "test", "image")
    cloth_dir = os.path.join(ROOT, "test", "cloth")
    cm_dir    = os.path.join(ROOT, "test", "cloth-mask")
    os.makedirs(img_dir, exist_ok=True)
    os.makedirs(cloth_dir, exist_ok=True)
    os.makedirs(cm_dir, exist_ok=True)

    user_name  = "person.jpg"  # person image filename
    cloth_name = "cloth.jpg"   # garment image filename

    # Write pairs: CLOTH FIRST, then PERSON (as CPDataset expects)
    with open(os.path.join(ROOT, "test_pairs.txt"), "w") as f:
        f.write(f"{cloth_name} {user_name}\n")

    # Save images in expected places
    shutil.copy(user_img_path,  os.path.join(img_dir,   user_name))
    shutil.copy(cloth_img_path, os.path.join(cloth_dir, cloth_name))

    # --- Compatibility shim ---
    # Some codepaths also try to read image/cloth.jpg; make a duplicate.
    shutil.copy(user_img_path, os.path.join(img_dir, cloth_name))

    # Reuse the garment as a coarse cloth-mask (you can replace with a real mask later)
    shutil.copy(cloth_img_path, os.path.join(cm_dir, cloth_name))

    # Debug print
    print("PAIR LINE:", open(os.path.join(ROOT, "test_pairs.txt")).read().strip())
    for p in ["test/image", "test/cloth", "test/cloth-mask"]:
        full = os.path.join(ROOT, p)
        print(p, "=>", os.listdir(full) if os.path.exists(full) else "MISSING")

    return ROOT



def _lazy_imports():
    from dci_vton_infer import DciVtonPredictor
    from preprocess_onepair import build_temp_dataset
    return DciVtonPredictor, build_temp_dataset

try:
    from spaces import GPU
except Exception:
    def GPU(*a,**kw):
        def deco(fn): return fn
        return deco

TITLE="💙 venbab DCI-VTON Try-On (Preview)"
DESC=("Upload a person image and a garment image. • CPU = preview blend "
      "• GPU = build one-pair dataset & run real DCI-VTON via test.py.")

def _pil_to_rgba(im): return im if im.mode=="RGBA" else im.convert("RGBA")

def auto_mask_torso(human,top,bot,feather):
    w,h=human.size; y1=int(h*top); y2=int(h*bot)
    mask=Image.new("L",(w,h),0); band=Image.new("L",(w,max(1,y2-y1)),255); mask.paste(band,(0,y1))
    if feather>0: mask=mask.filter(ImageFilter.GaussianBlur(radius=feather))
    return mask

def blend_preview(human,garment,use_mask,feather,top,bot,fit,blend):
    if human is None or garment is None: return None
    h=_pil_to_rgba(human); g=_pil_to_rgba(garment)
    hw,hh=h.size; gw,gh=g.size
    fr={"Slim (75%)":0.75,"Relaxed (85%)":0.85,"Wide (95%)":0.95}.get(fit,0.85)
    tw=int(hw*fr); sc=tw/max(1,gw); th=int(gh*sc)
    g=g.resize((tw,th),Image.BICUBIC)
    y1=int(hh*top); y2=int(hh*bot); thh=max(1,y2-y1)
    x=(hw-tw)//2; y=y1+(thh-th)//2
    ov=Image.new("RGBA",(hw,hh),(0,0,0,0)); ov.paste(g,(x,y),g)
    mask=auto_mask_torso(h,top,bot,feather) if use_mask else Image.new("L",(hw,hh),255)
    return Image.blend(h,Image.composite(ov,h,mask),max(0,min(1,blend))).convert("RGB")

_PREDICTOR=None

@GPU
def tryon_gpu(person_img, garment_img, use_mask, feather, top, bot, fit, blend):
    global _PREDICTOR
    DciVtonPredictor, _ = _lazy_imports()  # we won't use build_temp_dataset anymore

    if _PREDICTOR is None:
        _PREDICTOR = DciVtonPredictor(device="cuda")

    # 1) save incoming PIL images to temp files
    tmp_user  = "/tmp/dci_user.jpg"
    tmp_cloth = "/tmp/dci_cloth.jpg"
    person_img.convert("RGB").save(tmp_user, quality=95)
    garment_img.convert("RGB").save(tmp_cloth, quality=95)

    # 2) build the one-pair dataset (creates test_img/, test_color/, test_pairs.txt)
    dataroot = prepare_onepair_dataset(tmp_user, tmp_cloth)

    mask = auto_mask_torso(person_img, top, bot, feather) if use_mask else None
    return _PREDICTOR.predict(
        person_img.convert("RGB"),
        garment_img.convert("RGB"),
        mask_img=mask,
        cfg=dict(dataroot=dataroot, fit=fit, blend=blend, torso=(top, bot))
    )

def tryon_dispatch(person,garment,use_mask,feather,top,bot,fit,blend):
    if person is None or garment is None: return None
    if os.getenv("ENABLE_DCI","0")=="1":
        return tryon_gpu(person,garment,use_mask,feather,top,bot,fit,blend)
    return blend_preview(person,garment,use_mask,feather,top,bot,fit,blend)

CSS="""
.gradio-container{max-width:1200px!important;margin:auto;}
.gr-button.primary{background:linear-gradient(90deg,#6366f1,#8b5cf6);}
"""

with gr.Blocks(title=TITLE,theme=gr.themes.Soft(),css=CSS) as demo:
    gr.Markdown(f"# {TITLE}"); gr.Markdown(DESC)
    with gr.Row():
        with gr.Column(scale=1):
            human=gr.Image(label="Human",type="pil",height=480)
            use=gr.Checkbox(value=True,label="Auto-mask torso")
            feather=gr.Slider(0,30,value=10,step=1,label="Feather (px)")
            ttop=gr.Slider(0.05,0.45,value=0.30,step=0.01,label="Torso top")
            tbot=gr.Slider(0.50,0.90,value=0.68,step=0.01,label="Torso bottom")
        with gr.Column(scale=1):
            garment=gr.Image(label="Garment",type="pil",height=480)
            fit=gr.Radio(["Slim (75%)","Relaxed (85%)","Wide (95%)"],value="Relaxed (85%)",label="Fit width")
            blend=gr.Slider(0.2,1.5,value=0.9,step=0.05,label="Blend strength")
        with gr.Column(scale=1):
            output=gr.Image(label="Result",type="pil",height=480)
    btn=gr.Button("Try-On",variant="primary")
    btn.click(fn=tryon_dispatch,
              inputs=[human,garment,use,feather,ttop,tbot,fit,blend],
              outputs=output)

if __name__=="__main__": demo.launch()