Update app.py

app.py

@@ -36,25 +36,33 @@ scene_options = [
     "Shopping Mall", "Concert", "Sports Stadium"
 ]
 
-def get_base64_from_url(url):
-    response = requests.get(url)
-    return base64.b64encode(response.content).decode('utf-8')
+def fetch_and_process_images(image_list):
+    processed_images = []
+    for img in image_list:
+        try:
+            response = requests.get(img['url'])
+            image = Image.open(io.BytesIO(response.content))
+            processed_images.append({
+                "image": image,
+                "label": img['label'],
+                "url": img['url']
+            })
+        except Exception as e:
+            print(f"Error processing image {img['url']}: {str(e)}")
+    return processed_images
+
+# Fetch and process images at startup
+processed_cloth_images = fetch_and_process_images(cloth_images)
+processed_user_images = fetch_and_process_images(user_images)
 
 def get_base64_from_image(image):
     buffered = io.BytesIO()
     image.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode('utf-8')
 
-def generate_tryon(cloth_input, user_input, background):
-    if isinstance(cloth_input, str):  # URL selected
-        cloth_base64 = get_base64_from_url(cloth_input)
-    else:  # Image uploaded
-        cloth_base64 = get_base64_from_image(cloth_input)
-    
-    if isinstance(user_input, str):  # URL selected
-        user_base64 = get_base64_from_url(user_input)
-    else:  # Image uploaded
-        user_base64 = get_base64_from_image(user_input)
+def generate_tryon(cloth_image, user_image, background):
+    cloth_base64 = get_base64_from_image(cloth_image)
+    user_base64 = get_base64_from_image(user_image)
 
     input_data = {
         "user_image": user_base64,
@@ -95,14 +103,10 @@ def generate_tryon(cloth_input, user_input, background):
             raise Exception(f"Job processing failed: {status_data}")
         
         time.sleep(2)
-def tryon_interface(cloth_selected, cloth_upload, user_selected, user_upload, scene_selection, custom_scene):
-    # Use selected cloth if available, otherwise use uploaded cloth
-    cloth = cloth_selected if cloth_selected is not None else cloth_upload
-    
-    # Use selected user image if available, otherwise use uploaded user image
-    user = user_selected if user_selected is not None else user_upload
-    
-    # Use custom scene if provided, otherwise use selected scene
+
+def tryon_interface(cloth_index, cloth_upload, user_index, user_upload, scene_selection, custom_scene):
+    cloth = processed_cloth_images[cloth_index]['image'] if cloth_index is not None else cloth_upload
+    user = processed_user_images[user_index]['image'] if user_index is not None else user_upload
     background = custom_scene if custom_scene else scene_selection
 
     if cloth is None:
@@ -118,7 +122,6 @@ def tryon_interface(cloth_selected, cloth_upload, user_selected, user_upload, sc
     except Exception as e:
         return None, f"Error: {str(e)}"
 
-
 with gr.Blocks() as demo:
     gr.Markdown("# TryItOut.AI")
     
@@ -126,25 +129,24 @@ with gr.Blocks() as demo:
         with gr.Column():
             gr.Markdown("## Available Clothing")
             cloth_gallery = gr.Gallery(
-                [img["url"] for img in cloth_images],
+                [img["url"] for img in processed_cloth_images],
                 label="Click to select clothing",
                 columns=4,
                 height=500
             )
-            cloth_selected = gr.State(value=None)
-            print(cloth_selected)
-            cloth_upload = gr.Image(label="Or Upload Custom Clothing", type="numpy", visible=False)
+            cloth_index = gr.State(value=None)
+            cloth_upload = gr.Image(label="Or Upload Custom Clothing", type="pil")
 
         with gr.Column():
             gr.Markdown("## Available User Images")
             user_gallery = gr.Gallery(
-                [img["url"] for img in user_images],
+                [img["url"] for img in processed_user_images],
                 label="Click to select user image",
                 columns=3,
                 height=500
             )
-            user_selected = gr.State(value=None)
-            user_upload = gr.Image(label="Or Upload Custom User Image", type="numpy", visible=False)
+            user_index = gr.State(value=None)
+            user_upload = gr.Image(label="Or Upload Custom User Image", type="pil")
 
     with gr.Row():
         scene_selection = gr.Dropdown(choices=scene_options, label="Select Scene")
@@ -155,38 +157,23 @@ with gr.Blocks() as demo:
     output_image = gr.Image(label="Try-On Result")
     output_text = gr.Textbox(label="Status")
 
-    def update_selected(evt: gr.SelectData, gallery_images):
-        return gallery_images[evt.index]["url"]
-
-    def toggle_upload(gallery_selected, upload_component):
-        return gr.update(visible=gallery_selected is None)
+    def update_selected(evt: gr.SelectData):
+        return evt.index
 
     cloth_gallery.select(
         update_selected,
-        inputs=[cloth_gallery],
-        outputs=[cloth_selected]
-    ).then(
-        toggle_upload,
-        inputs=[cloth_selected],
-        outputs=[cloth_upload]
+        outputs=[cloth_index]
     )
     
     user_gallery.select(
         update_selected,
-        inputs=[user_gallery],
-        outputs=[user_selected]
-    ).then(
-        toggle_upload,
-        inputs=[user_selected],
-        outputs=[user_upload]
+        outputs=[user_index]
     )
 
-
     generate_button.click(
         tryon_interface,
-        inputs=[cloth_selected, cloth_upload, user_selected, user_upload, scene_selection, custom_scene],
+        inputs=[cloth_index, cloth_upload, user_index, user_upload, scene_selection, custom_scene],
         outputs=[output_image, output_text]
     )
 
-
 demo.launch()