Create app.py

app.py

@@ -0,0 +1,167 @@
+import gradio as gr
+import requests
+import base64
+import time
+import os
+from PIL import Image
+import io
+
+API_KEY = os.environ.get('RUNPOD_API_KEY')
+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, scene):
+    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": scene
+    }
+
+    response = requests.post(
+        f"{API_URL}/run",
+        headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"},
+        json={"input": input_data}
+    )
+
+    if response.status_code != 200:
+        raise Exception(f"API request failed with status {response.status_code}: {response.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 status_response.status_code != 200:
+            raise Exception(f"Status check failed with status {status_response.status_code}: {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_selection, cloth_upload, user_selection, user_upload, scene_selection, custom_scene):
+    cloth_input = cloth_upload if cloth_upload is not None else cloth_selection
+    user_input = user_upload if user_upload is not None else user_selection
+    scene = custom_scene if custom_scene else scene_selection
+
+    if not cloth_input or not user_input or not scene:
+        return None, "Please select or upload both clothing and user images, and specify a scene."
+
+    try:
+        result_image = generate_tryon(cloth_input, user_input, scene)
+        return result_image, "Try-on image generated successfully!"
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
+def create_image_grid(images, labels, columns=4):
+    rows = -(-len(images) // columns)  # Ceiling division
+    cell_width = 200
+    cell_height = 250
+    grid = Image.new('RGB', (columns * cell_width, rows * cell_height), color='white')
+    
+    for i, (url, label) in enumerate(zip(images, labels)):
+        response = requests.get(url)
+        img = Image.open(io.BytesIO(response.content))
+        img.thumbnail((cell_width, cell_height - 30))
+        
+        x = (i % columns) * cell_width
+        y = (i // columns) * cell_height
+        
+        grid.paste(img, (x + (cell_width - img.width) // 2, y))
+        
+        # Add label
+        from PIL import ImageDraw, ImageFont
+        draw = ImageDraw.Draw(grid)
+        font = ImageFont.load_default()
+        label_width = draw.textlength(label, font=font)
+        draw.text((x + (cell_width - label_width) // 2, y + cell_height - 20), label, fill="black", font=font)
+    
+    return grid
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Virtual Try-On Application")
+    
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("## Available Clothing")
+            cloth_grid = gr.Image(create_image_grid([img["url"] for img in cloth_images], 
+                                                    [img["label"] for img in cloth_images]))
+            cloth_selection = gr.Dropdown(choices=[img["url"] for img in cloth_images], 
+                                          label="Select Clothing")
+            cloth_upload = gr.Image(label="Or Upload Custom Clothing")
+
+        with gr.Column():
+            gr.Markdown("## Available User Images")
+            user_grid = gr.Image(create_image_grid([img["url"] for img in user_images], 
+                                                   [img["label"] for img in user_images]))
+            user_selection = gr.Dropdown(choices=[img["url"] for img in user_images], 
+                                         label="Select User Image")
+            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")
+
+    generate_button.click(
+        tryon_interface,
+        inputs=[cloth_selection, cloth_upload, user_selection, user_upload, scene_selection, custom_scene],
+        outputs=[output_image, output_text]
+    )
+
+demo.launch()