|
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." |
|
|
|
|
|
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, |
|
seed=42, |
|
vt_model_type="viton_hd", |
|
vt_garment_type=garment_type, |
|
vt_repaint=False, |
|
api_name="/leffa_predict_vt", |
|
) |
|
|
|
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, |
|
) |