Spaces:
Running
on
Zero
Running
on
Zero
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 | |
import gradio as gr | |
# Download checkpoints | |
snapshot_download(repo_id="franciszzj/Leffa", local_dir="./") | |
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) | |
src_image = Image.open(src_image_path) | |
ref_image = Image.open(ref_image_path) | |
src_image_array = np.array(src_image) | |
ref_image_array = np.array(ref_image) | |
# Mask | |
if control_type == "virtual_tryon": | |
automasker = AutoMasker( | |
densepose_path="./ckpts/densepose", | |
schp_path="./ckpts/schp", | |
) | |
src_image = src_image.convert("RGB") | |
mask = automasker(src_image, "upper")["mask"] | |
elif control_type == "pose_transfer": | |
mask = Image.fromarray(np.ones_like(src_image_array) * 255) | |
# DensePose | |
densepose_predictor = DensePosePredictor( | |
config_path="./ckpts/densepose/densepose_rcnn_R_50_FPN_s1x.yaml", | |
weights_path="./ckpts/densepose/model_final_162be9.pkl", | |
) | |
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() | |
if control_type == "virtual_tryon": | |
pretrained_model_name_or_path = "./ckpts/stable-diffusion-inpainting" | |
pretrained_model = "./ckpts/virtual_tryon.pth" | |
elif control_type == "pose_transfer": | |
pretrained_model_name_or_path = "./ckpts/stable-diffusion-xl-1.0-inpainting-0.1" | |
pretrained_model = "./ckpts/pose_transfer.pth" | |
model = LeffaModel( | |
pretrained_model_name_or_path=pretrained_model_name_or_path, | |
pretrained_model=pretrained_model, | |
) | |
inference = LeffaInference(model=model) | |
data = { | |
"src_image": [src_image], | |
"ref_image": [ref_image], | |
"mask": [mask], | |
"densepose": [densepose], | |
} | |
data = transform(data) | |
output = inference(data) | |
gen_image = output["generated_image"][0] | |
# gen_image.save("gen_image.png") | |
return np.array(gen_image) | |
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) | |
with gr.Blocks().queue() as demo: | |
gr.Markdown( | |
"## Leffa: Learning Flow Fields in Attention for Controllable Person Image Generation") | |
gr.Markdown("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).") | |
with gr.Row(): | |
with gr.Column(): | |
src_image = gr.Image( | |
sources=["upload"], | |
type="filepath", | |
label="Source Person Image", | |
width=384, | |
height=512, | |
) | |
with gr.Row(): | |
control_type = gr.Dropdown( | |
["virtual_tryon", "pose_transfer"], label="Control Type") | |
example = gr.Examples( | |
inputs=src_image, | |
examples_per_page=10, | |
examples=["./examples/14684_00_person.jpg", | |
"./examples/14092_00_person.jpg"], | |
) | |
with gr.Column(): | |
ref_image = gr.Image( | |
sources=["upload"], | |
type="filepath", | |
label="Reference Image", | |
width=384, | |
height=512, | |
) | |
with gr.Row(): | |
gen_button = gr.Button("Generate") | |
example = gr.Examples( | |
inputs=ref_image, | |
examples_per_page=10, | |
examples=["./examples/04181_00_garment.jpg", | |
"./examples/14684_00_person.jpg"], | |
) | |
with gr.Column(): | |
gen_image = gr.Image( | |
label="Generated Person Image", | |
width=384, | |
height=512, | |
) | |
gen_button.click(fn=leffa_predict, inputs=[ | |
src_image, ref_image, control_type], outputs=[gen_image]) | |
demo.launch(share=True, server_port=7860) | |