Update app.py

app.py

@@ -1,3 +1,100 @@
+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'
+
+cloth_images = [
+    {"url": "https://i.postimg.cc/7ZzLZtbc/hmgoepprod-4-1.jpg", "label": "Jacket 1"},
+    {"url": "https://i.postimg.cc/7Yd6DrP0/hmgoepprod-6.jpg", "label": "Jacket 2"},
+    {"url": "https://i.postimg.cc/8z40MNFr/hnm.png", "label": "Jacket 3"},
+    {"url": "https://i.postimg.cc/mkqbb74B/hmgoepprod-5.jpg", "label": "Jacket 4"},
+    {"url": "https://i.postimg.cc/vBwySGzH/hmgoepprod.jpg", "label": "Jacket 5"},
+    {"url": "https://i.postimg.cc/6qSYYprM/Whats-App-Image-2024-08-18-at-13-45-37-2.jpg", "label": "Jacket 6"},
+    {"url": "https://i.postimg.cc/tJHhTnzQ/Screenshot-2024-08-20-at-10-04-29-PM.png", "label": "Jacket 7"},
+    {"url": "https://i.postimg.cc/yxQYGwxQ/Screenshot-2024-08-20-at-10-06-28-PM.png", "label": "Jacket 8"},
+]
+
+user_images = [
+    {"url": "https://iili.io/dEk7qtp.md.jpg", "label": "User 1 (Male)"},
+    {"url": "https://iili.io/d1RpKBa.md.jpg", "label": "User 2 (Female)"},
+    {"url": "https://i.postimg.cc/QtzQrnFt/Whats-App-Image-2024-08-18-at-13-42-10.jpg", "label": "User 3 (Female)"},
+    {"url": "https://i.postimg.cc/rsKLDzyj/Whats-App-Image-2024-08-19-at-12-13-41.jpg", "label": "User 4 (Female)"},
+    {"url": "https://i.postimg.cc/cHkpGZ7n/Whats-App-Image-2024-08-19-at-00-27-16-1.jpg", "label": "User 5 (Female)"},
+    {"url": "https://i.postimg.cc/g2bfWghN/Whats-App-Image-2024-08-19-at-12-13-41-1.jpg", "label": "User 6 (Female)"},
+]
+
+scene_options = [
+    "Gym", "City Street", "Party", "Beach", "Office", "Park", "Cafe", 
+    "Shopping Mall", "Concert", "Sports Stadium"
+]
+
+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
+    }
+
+    response = requests.post(
+        f"{API_URL}/run",
+        headers={
+            "Authorization": f"Bearer {API_KEY}",
+            "Content-Type": "application/json"
+        },
+        data=json.dumps({"input": 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
+        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_selected if cloth_selected else cloth_upload
     user = user_selected if user_selected else user_upload