Update app.py

app.py

@@ -1,51 +1,160 @@
-def virtual_tryon(job):
-    print('Started processing')
+import gradio as gr
+import requests
+import base64
+import time
+import os
+from PIL import Image
+import io
+import json
+
+API_KEY = 'N80HWHVG3DV8URRNYZY382UPSHP1N8G1SNPYG0E9'
+API_URL = 'https://api.runpod.ai/v2/31jyh9kh7nwyga'
+
+# ... [previous code for cloth_images, user_images, and scene_options remains unchanged] ...
+
+def get_base64_from_url(url):
+    response = requests.get(url)
+    return base64.b64encode(response.content).decode('utf-8')
+
+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)
+
+    input_data = {
+        "user_image": user_base64,
+        "product_image": cloth_base64,
+        "background": background
+    }
+
+    # Prepare log message
+    log_data = {
+        "user_image": user_base64[:20] + '...' if user_base64 else 'undefined',
+        "product_image": cloth_base64[:20] + '...' if cloth_base64 else 'undefined',
+        "background": background,
+    }
+    log_message = f"Sending data:\n{json.dumps(log_data, indent=2)}"
+
+    response = requests.post(
+        f"{API_URL}/run",
+        headers={
+            "Authorization": f"Bearer {API_KEY}",
+            "Content-Type": "application/json"
+        },
+        json=input_data
+    )
+
+    if not response.ok:
+        error_text = response.text
+        raise Exception(f"Failed to upload image: {response.status_code} {response.reason} - {error_text}")
+
+    job_id = response.json()['id']
+    
+    while True:
+        status_response = requests.get(
+            f"{API_URL}/status/{job_id}",
+            headers={"Authorization": f"Bearer {API_KEY}"}
+        )
+        
+        if not status_response.ok:
+            raise Exception(f"Status check failed: {status_response.status_code} {status_response.reason} - {status_response.text}")
+        
+        status_data = status_response.json()
+        if status_data['status'] == 'COMPLETED':
+            output_base64 = status_data['output']['output']
+            output_image = Image.open(io.BytesIO(base64.b64decode(output_base64)))
+            return output_image, log_message
+        elif status_data['status'] == 'FAILED':
+            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):
+    cloth = cloth_upload if cloth_upload is not None else cloth_selected
+    user = user_upload if user_upload is not None else user_selected
+    background = custom_scene if custom_scene else scene_selection
+
+    if not cloth:
+        return None, "Please select or upload a clothing image.", ""
+    if not user:
+        return None, "Please select or upload a user image.", ""
+    if not background:
+        return None, "Please select or enter a background scene.", ""
+
     try:
-        # Parse the JSON input
-        job_input = job["input"]
-        product_image_base64 = job_input.get('product_image')
-        user_image_base64 = job_input.get('user_image')
-        background = job_input.get('background')
-        print('Input parsed')
-
-        # Convert base64 images to PIL Image
-        user_image = base64_to_pil_image(user_image_base64)
-        user_image.convert('RGB').save('/tmp/user.jpg')
-        print('User image saved')
-
-        product_image = base64_to_pil_image(product_image_base64)
-        product_image.convert('RGB').save('/tmp/product.jpg')
-        print('Product image saved')
-
-        # Create SDXL pipeline
-        pipeline = create_sdxl_pipeline_with_dinov2(MODEL_NAME)
-        print('Pipeline generated')
-        segment_cloth = "dress cloth sweater shirt tshirt"
-
-        # Describe the garment
-        garment_description, gender_description, garment_description_full = describe_garment('/tmp/product.jpg', '/tmp/user.jpg')
-
-        if 'hood' in garment_description:
-            garment_description = 'hoodie'
-        prompt = f"photo at dusk, a {gender_description} wearing {garment_description}, background is {background}, 35 mm film camera"
-        insight = Insight_Face2(MODEL_NAME, LORA_CKPT, IMAGE_ENCODER_CKPT, FACE_MODEL_PATH)
-        single_image = insight.generate(prompt, NEGATIVE_PROMPT, 1, user_image)
-        single_image.save('/tmp/Single_IP_Result.jpg')
-
-        # Virtual Try-on
-        ip_weight = 0.5
-        virt_try = VirtualTryon(MODEL_NAME, pipeline, ip_weight, IMAGE_ENCODER_CKPT, segment_cloth)
-        final_im = virt_try.generate(garment_description_full, NEGATIVE_PROMPT, product_image, single_image)
-        final_im.save('/tmp/Try_on.jpg')
-
-        final_im = np.asarray(final_im)
-        final_im = cv2.cvtColor(final_im, cv2.COLOR_BGR2RGB)
-
-        _, buffer = cv2.imencode('.png', final_im)
-        image_base64 = base64.b64encode(buffer).decode('utf-8')
-
-        return {"output": image_base64}
+        result_image, log_message = generate_tryon(cloth, user, background)
+        return result_image, "Try-on image generated successfully!", log_message
     except Exception as e:
-        print(f"An error occurred: {str(e)}")
-        return {"error": str(e)}
+        return None, f"Error: {str(e)}", ""
+
+def select_image(evt: gr.SelectData, image_list):
+    if evt.index < len(image_list):
+        return image_list[evt.index]["url"]
+    return None
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Virtual Try-On Application")
+    
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("## Available Clothing")
+            cloth_gallery = gr.Gallery(
+                [img["url"] for img in cloth_images],
+                label="Click to select clothing",
+                columns=4,
+                height=500
+            )
+            cloth_selected = gr.Textbox(label="Selected Clothing URL", visible=False)
+            cloth_upload = gr.Image(label="Or Upload Custom Clothing")
+
+        with gr.Column():
+            gr.Markdown("## Available User Images")
+            user_gallery = gr.Gallery(
+                [img["url"] for img in user_images],
+                label="Click to select user image",
+                columns=3,
+                height=500
+            )
+            user_selected = gr.Textbox(label="Selected User Image URL", visible=False)
+            user_upload = gr.Image(label="Or Upload Custom User Image")
+
+    with gr.Row():
+        scene_selection = gr.Dropdown(choices=scene_options, label="Select Scene")
+        custom_scene = gr.Textbox(label="Or Enter Custom Scene")
+
+    generate_button = gr.Button("Generate Try-On")
+    
+    output_image = gr.Image(label="Try-On Result")
+    output_text = gr.Textbox(label="Status")
+    log_output = gr.Textbox(label="API Request Details", lines=10)
+
+    cloth_gallery.select(
+        fn=lambda evt: select_image(evt, cloth_images),
+        inputs=[cloth_gallery],
+        outputs=[cloth_selected]
+    )
+    
+    user_gallery.select(
+        fn=lambda evt: select_image(evt, user_images),
+        inputs=[user_gallery],
+        outputs=[user_selected]
+    )
+
+    generate_button.click(
+        tryon_interface,
+        inputs=[cloth_selected, cloth_upload, user_selected, user_upload, scene_selection, custom_scene],
+        outputs=[output_image, output_text, log_output]
+    )
+
+demo.launch()