HAL1993's picture
Update app.py
d785d43 verified
raw
history blame
2.93 kB
import os
from PIL import Image
import gradio as gr
from gradio_client import Client, handle_file
HF_TOKEN = os.getenv("HF_TOKEN")
client = Client(
"franciszzj/Leffa",
hf_token=HF_TOKEN,
)
def rename_to_image(input_path, output_dir="Uploads"):
os.makedirs(output_dir, exist_ok=True)
random_id = os.urandom(4).hex()
output_filename = f"Image_{random_id}.png"
output_path = os.path.join(output_dir, output_filename)
try:
with Image.open(input_path) as img:
img.convert('RGB').save(output_path, 'PNG', quality=100)
return output_path
except Exception as e:
raise ValueError(f"Error renaming: {str(e)}")
def virtual_tryon(person_path, garment_path, garment_type):
if not person_path or not garment_path:
return None, "Ju lutem ngarkoni të dyja imazhet."
# Rename uploads to Image_<random>.png
person_png = rename_to_image(person_path)
garment_png = rename_to_image(garment_path)
person_file = handle_file(person_png)
garment_file = handle_file(garment_png)
try:
result = client.predict(
person_file,
garment_file,
ref_acceleration=False,
step=40,
scale=2.5, # Corrected to 2.5 (was 42 in last version)
seed=42,
vt_model_type="viton_hd",
vt_garment_type=garment_type,
vt_repaint=False,
api_name="/leffa_predict_vt",
)
# Rename result to uniqueproc_<random>.png
input_image_path = result[0]
random_id = os.urandom(4).hex()
output_image_path = os.path.join("Outputs", f"uniqueproc_{random_id}.png")
os.makedirs("Outputs", exist_ok=True)
with Image.open(input_image_path) as img:
img.convert('RGB').save(output_image_path, 'PNG', quality=95)
return output_image_path, ""
except Exception as e:
return None, f"Error processing: {str(e)}"
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown("")
error_msg = gr.Markdown("", visible=False)
with gr.Row():
with gr.Column():
src = gr.Image(sources="upload", type="filepath", label="Foto e personit")
with gr.Column():
ref = gr.Image(sources="upload", type="filepath", label="Foto e veshjes")
with gr.Column():
garment_type = gr.Radio(
choices=[("Siper", "upper_body"),
("Posht", "lower_body"),
("Komplet", "dresses")],
value="upper_body",
label="Lloji i veshjes",
)
with gr.Column():
btn = gr.Button("VISHU")
out = gr.Image(type="filepath", label="Rezultati")
btn.click(
virtual_tryon,
inputs=[src, ref, garment_type],
outputs=[out, error_msg]
)
demo.launch(
share=True,
show_error=True,
pwa=True,
)