VTRY / main.py
Eaglebelt's picture
Upload main.py
d93a849 verified
raw
history blame
1.92 kB
import os
import io
import PIL.Image # Workaround for PIL/Gradio bug :contentReference[oaicite:13]{index=13}
import gradio as gr
from gradio_client import Client, handle_file
from numpy import array
# 1. Load your HF token from env
HF_TOKEN = os.getenv("HF_TOKEN") # export HF_TOKEN="hf_..."
# 1) Connect to the Leffa Gradio app’s predict endpoint
# Use the full "/call/predict" API path as shown on the View API page
client = Client("franciszzj/Leffa", hf_token=HF_TOKEN, ) # Gradio Python client
def virtual_tryon(person_path, garment_path):
# 2) Wrap file inputs so Gradio client uploads them correctly
person_file = handle_file(person_path) # handle_file uploads the image :contentReference[oaicite:6]{index=6}
garment_file = handle_file(garment_path)
# 3) Build inputs in the exact order shown on the “Use via API” page :contentReference[oaicite:7]{index=7}
# 4) Call the named endpoint with handle_file inputs
result = client.predict(
person_file, # Person Image
garment_file, # Garment Image
ref_acceleration=False,
step=30,
scale=2.5,
seed=42,
vt_model_type="viton_hd",
vt_garment_type="upper_body",
vt_repaint=False,
api_name="/leffa_predict_vt"
)
# result[0] is the generated image filepath on the server
return result[0] # Gradio will download & display this file
# 5) Gradio UI
with gr.Blocks() as demo:
gr.Markdown("## Leffa Virtual Try-On")
with gr.Row():
src = gr.Image(sources="upload", type="filepath", label="Person Image")
ref = gr.Image(sources="upload", type="filepath", label="Garment Image")
out = gr.Image(type="filepath", label="Result", )
btn = gr.Button("Generate")
btn.click(virtual_tryon, [src, ref], out)
demo.launch(share=True,
show_error=True,
pwa=True,)