app.py

import numpy as np
from PIL import Image
from huggingface_hub import snapshot_download
from leffa.transform import LeffaTransform
from leffa.model import LeffaModel
from leffa.inference import LeffaInference
from utils.garment_agnostic_mask_predictor import AutoMasker
from utils.densepose_predictor import DensePosePredictor
from utils.utils import resize_and_center
import spaces
import requests
from io import BytesIO

import gradio as gr

print("Imports done, downloading the model...")

# Download checkpoints
snapshot_download(repo_id="franciszzj/Leffa", local_dir="./ckpts")

mask_predictor = AutoMasker(
    densepose_path="./ckpts/densepose",
    schp_path="./ckpts/schp",
)

densepose_predictor = DensePosePredictor(
    config_path="./ckpts/densepose/densepose_rcnn_R_50_FPN_s1x.yaml",
    weights_path="./ckpts/densepose/model_final_162be9.pkl",
)

vt_model = LeffaModel(
    pretrained_model_name_or_path="./ckpts/stable-diffusion-inpainting",
    pretrained_model="./ckpts/virtual_tryon.pth",
)
vt_inference = LeffaInference(model=vt_model)

pt_model = LeffaModel(
    pretrained_model_name_or_path="./ckpts/stable-diffusion-xl-1.0-inpainting-0.1",
    pretrained_model="./ckpts/pose_transfer.pth",
)
pt_inference = LeffaInference(model=pt_model)

print("Model downloaded, ready to serve!")

@spaces.GPU
def leffa_predict(src_image_path, ref_image_path, control_type):
    assert control_type in [
        "virtual_tryon", "pose_transfer"], "Invalid control type: {}".format(control_type)
    
    if isinstance(src_image_path, str):
        src_image = Image.open(src_image_path)
    else:
        src_image = src_image_path
    
    if isinstance(ref_image_path, str):
        ref_image = Image.open(ref_image_path)
    else:
        ref_image = ref_image_path
    
    src_image = resize_and_center(src_image, 768, 1024)
    ref_image = resize_and_center(ref_image, 768, 1024)

    src_image_array = np.array(src_image)
    ref_image_array = np.array(ref_image)

    # Mask
    if control_type == "virtual_tryon":
        src_image = src_image.convert("RGB")
        mask = mask_predictor(src_image, "upper")["mask"]
    elif control_type == "pose_transfer":
        mask = Image.fromarray(np.ones_like(src_image_array) * 255)

    # DensePose
    src_image_iuv_array = densepose_predictor.predict_iuv(src_image_array)
    src_image_seg_array = densepose_predictor.predict_seg(src_image_array)
    src_image_iuv = Image.fromarray(src_image_iuv_array)
    src_image_seg = Image.fromarray(src_image_seg_array)
    if control_type == "virtual_tryon":
        densepose = src_image_seg
    elif control_type == "pose_transfer":
        densepose = src_image_iuv

    # Leffa
    transform = LeffaTransform()

    data = {
        "src_image": [src_image],
        "ref_image": [ref_image],
        "mask": [mask],
        "densepose": [densepose],
    }
    data = transform(data)
    if control_type == "virtual_tryon":
        inference = vt_inference
    elif control_type == "pose_transfer":
        inference = pt_inference
    output = inference(data)
    gen_image = output["generated_image"][0]
    # gen_image.save("gen_image.png")
    return np.array(gen_image)

@spaces.GPU
def leffa_predict_vt(src_image_path, ref_image_path):
    return leffa_predict(src_image_path, ref_image_path, "virtual_tryon")

@spaces.GPU
def leffa_predict_vt_image_url(person_url, garment_url):
    if not person_url or not garment_url:
        return None
    
    src_image = fetch_image_from_url(person_url)
    if not src_image:
        return None
    print("fetched person image")
    
    ref_image = fetch_image_from_url(garment_url)
    if not ref_image:
        return None
    print("fetched garment image")

    return leffa_predict(src_image, ref_image, "virtual_tryon")

@spaces.GPU
def leffa_predict_pt(src_image_path, ref_image_path):
    return leffa_predict(src_image_path, ref_image_path, "pose_transfer")

def fetch_image_from_url(url):
    try:
        response = requests.get(url)
        img = Image.open(BytesIO(response.content))
        return img
    except Exception as e:
        print(e)
        return None

def handle_image_input(image_input):
    if image_input.startswith('http'):
        return fetch_image_from_url(image_input)
    else:
        return Image.open(image_input)

# if __name__ == "__main__":
#     # import sys

#     # src_image_path = sys.argv[1]
#     # ref_image_path = sys.argv[2]
#     # control_type = sys.argv[3]
#     # leffa_predict(src_image_path, ref_image_path, control_type)

#     title = "## Leffa: Learning Flow Fields in Attention for Controllable Person Image Generation"
#     link = "[📚 Paper](https://arxiv.org/abs/2412.08486) - [🔥 Demo](https://huggingface.co/spaces/franciszzj/Leffa) - [🤗 Model](https://huggingface.co/franciszzj/Leffa)"
#     description = "Leffa is a unified framework for controllable person image generation that enables precise manipulation of both appearance (i.e., virtual try-on) and pose (i.e., pose transfer)."
#     note = "Note: The models used in the demo are trained solely on academic datasets. Virtual try-on uses VITON-HD, and pose transfer uses DeepFashion."

#     with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.pink, secondary_hue=gr.themes.colors.red)).queue() as demo:
#         gr.Markdown(title)
#         gr.Markdown(link)
#         gr.Markdown(description)

#         with gr.Tab("Control Appearance (Virtual Try-on)"):
#             with gr.Row():
#                 with gr.Column():
#                     gr.Markdown("#### Person Image")
#                     vt_src_image = gr.Image(
#                         sources=["upload", "url"],
#                         type="filepath",
#                         label="Person Image",
#                         width=512,
#                         height=512,
#                     )

#                     gr.Examples(
#                         inputs=vt_src_image,
#                         examples_per_page=5,
#                         examples=["./ckpts/examples/person1/01350_00.jpg",
#                                   "./ckpts/examples/person1/01376_00.jpg",
#                                   "./ckpts/examples/person1/01416_00.jpg",
#                                   "./ckpts/examples/person1/05976_00.jpg",
#                                   "./ckpts/examples/person1/06094_00.jpg",],
#                     )

#                 with gr.Column():
#                     gr.Markdown("#### Garment Image")
#                     vt_ref_image = gr.Image(
#                         sources=["upload", "url"],
#                         type="filepath",
#                         label="Garment Image",
#                         width=512,
#                         height=512,
#                     )

#                     gr.Examples(
#                         inputs=vt_ref_image,
#                         examples_per_page=5,
#                         examples=["./ckpts/examples/garment/01449_00.jpg",
#                                   "./ckpts/examples/garment/01486_00.jpg",
#                                   "./ckpts/examples/garment/01853_00.jpg",
#                                   "./ckpts/examples/garment/02070_00.jpg",
#                                   "./ckpts/examples/garment/03553_00.jpg",],
#                     )

#                 with gr.Column():
#                     gr.Markdown("#### Generated Image")
#                     vt_gen_image = gr.Image(
#                         label="Generated Image",
#                         width=512,
#                         height=512,
#                     )

#                     with gr.Row():
#                         vt_gen_button = gr.Button("Generate")

#                 vt_gen_button.click(fn=leffa_predict_vt, inputs=[
#                     vt_src_image, vt_ref_image], outputs=[vt_gen_image])

#         with gr.Tab("Control Pose (Pose Transfer)"):
#             with gr.Row():
#                 with gr.Column():
#                     gr.Markdown("#### Person Image")
#                     pt_ref_image = gr.Image(
#                         sources=["upload"],
#                         type="filepath",
#                         label="Person Image",
#                         width=512,
#                         height=512,
#                     )

#                     gr.Examples(
#                         inputs=pt_ref_image,
#                         examples_per_page=5,
#                         examples=["./ckpts/examples/person1/01350_00.jpg",
#                                   "./ckpts/examples/person1/01376_00.jpg",
#                                   "./ckpts/examples/person1/01416_00.jpg",
#                                   "./ckpts/examples/person1/05976_00.jpg",
#                                   "./ckpts/examples/person1/06094_00.jpg",],
#                     )

#                 with gr.Column():
#                     gr.Markdown("#### Target Pose Person Image")
#                     pt_src_image = gr.Image(
#                         sources=["upload"],
#                         type="filepath",
#                         label="Target Pose Person Image",
#                         width=512,
#                         height=512,
#                     )

#                     gr.Examples(
#                         inputs=pt_src_image,
#                         examples_per_page=5,
#                         examples=["./ckpts/examples/person2/01850_00.jpg",
#                                   "./ckpts/examples/person2/01875_00.jpg",
#                                   "./ckpts/examples/person2/02532_00.jpg",
#                                   "./ckpts/examples/person2/02902_00.jpg",
#                                   "./ckpts/examples/person2/05346_00.jpg",],
#                     )

#                 with gr.Column():
#                     gr.Markdown("#### Generated Image")
#                     pt_gen_image = gr.Image(
#                         label="Generated Image",
#                         width=512,
#                         height=512,
#                     )

#                     with gr.Row():
#                         pose_transfer_gen_button = gr.Button("Generate")

#                 pose_transfer_gen_button.click(fn=leffa_predict_pt, inputs=[
#                     pt_src_image, pt_ref_image], outputs=[pt_gen_image])

#         gr.Markdown(note)

#         demo.launch(share=True, server_port=7860)

def create_demo():
    title = "## Virtual Try-on with URLs"
    description = "Enter URLs for both the person image and the garment image to generate a virtual try-on result."
    
    with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.pink)) as demo:
        gr.Markdown(title)
        gr.Markdown(description)
        
        with gr.Row():
            with gr.Column():
                person_url = gr.Textbox(
                    label="Person Image URL",
                    placeholder="Enter URL of the person image..."
                )
                garment_url = gr.Textbox(
                    label="Garment Image URL",
                    placeholder="Enter URL of the garment image..."
                )
                
                # Example URLs
                gr.Examples(
                    inputs=[person_url, garment_url],
                    examples=[
                        ["https://example.com/person1.jpg", "https://example.com/garment1.jpg"],
                        ["https://example.com/person2.jpg", "https://example.com/garment2.jpg"],
                    ],
                    label="Example URLs"
                )
                
                generate_btn = gr.Button("Generate Try-on")
            
            with gr.Column():
                output_image = gr.Image(
                    label="Generated Result",
                    width=512,
                    height=512
                )
        
        generate_btn.click(
            fn=leffa_predict_vt_image_url,
            inputs=[person_url, garment_url],
            outputs=output_image
        )
        
        gr.Markdown("Note: This model is trained solely on academic datasets (VITON-HD).")
    
    return demo

if __name__ == "__main__":
    demo = create_demo()
    demo.launch(share=True, server_port=7860)