Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 from PIL import Image
-from gradio_client import Client, handle_file
 import gradio as gr
+from gradio_client import Client, handle_file
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 
@@ -13,7 +13,7 @@ client = Client(
 # Variable to hold result path
 result_path = ""
 
-def convert_to_png(input_path, output_dir="uploads"):
+def convert_to_png(input_path, output_dir="Uploads"):
     # Create output directory if it doesn't exist
     os.makedirs(output_dir, exist_ok=True)
     # Generate unique output path
@@ -24,13 +24,27 @@ def convert_to_png(input_path, output_dir="uploads"):
         img.convert('RGB').save(output_path, 'PNG', quality=100)
     return output_path
 
+def handle_person_upload(person_path):
+    if person_path:
+        # Convert uploaded person image to PNG
+        person_png = convert_to_png(person_path)
+        return person_png
+    return None
+
+def handle_garment_upload(garment_path):
+    if garment_path:
+        # Convert uploaded garment image to PNG
+        garment_png = convert_to_png(garment_path)
+        return garment_png
+    return None
+
 def virtual_tryon(person_path, garment_path, garment_type):
     global result_path
-    # Convert uploaded images to PNG
-    person_png = convert_to_png(person_path)
-    garment_png = convert_to_png(garment_path)
-    person_file = handle_file(person_png)
-    garment_file = handle_file(garment_png)
+    if not person_path or not garment_path:
+        return None, gr.update(visible=False), gr.update(visible=True), "Ju lutem ngarkoni të dyja imazhet."
+    
+    person_file = handle_file(person_path)
+    garment_file = handle_file(garment_path)
 
     result = client.predict(
         person_file,
@@ -50,13 +64,14 @@ def virtual_tryon(person_path, garment_path, garment_type):
     with Image.open(input_image_path) as img:
         img.convert('RGB').save(output_image_path, 'JPEG', quality=95)
     result_path = output_image_path
-    return result_path, gr.update(visible=True), gr.update(visible=False)
+    return result_path, gr.update(visible=True), gr.update(visible=False), ""
 
 def trigger_download():
     return gr.update(value=result_path, visible=True, interactive=True)
 
 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")
@@ -76,7 +91,12 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             downloadable = gr.File(visible=False)
             out = gr.Image(type="filepath", label="Rezultati")
 
-    btn.click(virtual_tryon, [src, ref, garment_type], [out, shkarko_btn, downloadable])
+    # Convert uploads to PNG immediately
+    src.change(handle_person_upload, inputs=src, outputs=src)
+    ref.change(handle_garment_upload, inputs=ref, outputs=ref)
+    # Process and show result
+    btn.click(virtual_tryon, [src, ref, garment_type], [out, shkarko_btn, downloadable, error_msg])
+    # Trigger download
     shkarko_btn.click(trigger_download, outputs=downloadable)
 
 demo.launch(