app.py

import gradio as gr
import requests
import base64
import time
import os
import numpy as np
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 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_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,
        "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_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:
        return None, "Please select or upload a clothing image."
    if user is None:
        return None, "Please select or upload a user image."
    if not background:
        return None, "Please select or enter a background scene."

    try:
        result_image = generate_tryon(cloth, user, background)
        return result_image, "Try-on image generated successfully!"
    except Exception as e:
        return None, f"Error: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("# TryItOut.AI")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("## Available Clothing")
            cloth_gallery = gr.Gallery(
                [img["url"] for img in processed_cloth_images],
                label="Click to select clothing",
                columns=4,
                height=500
            )
            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 processed_user_images],
                label="Click to select user image",
                columns=3,
                height=500
            )
            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")
        custom_scene = gr.Textbox(label="Or Enter Custom Scene")

    generate_button = gr.Button("TryItOut!!")
    
    output_image = gr.Image(label="Try-On Result")
    output_text = gr.Textbox(label="Status")

    def update_selected(evt: gr.SelectData):
        return evt.index

    cloth_gallery.select(
        update_selected,
        outputs=[cloth_index]
    )
    
    user_gallery.select(
        update_selected,
        outputs=[user_index]
    )

    generate_button.click(
        tryon_interface,
        inputs=[cloth_index, cloth_upload, user_index, user_upload, scene_selection, custom_scene],
        outputs=[output_image, output_text]
    )

demo.launch()