import gradio as gr
from diffusers import DiffusionPipeline
from PIL import Image

# Load your diffusion model
model_id = "models/gouthaml/raos-virtual-try-on-model"  # Update with the correct model ID
pipe = DiffusionPipeline.from_pretrained(model_id)

# Function to handle the virtual try-on process
def virtual_tryon(person_image, dress_image):
    # Call the model to generate the virtual try-on result
    output_image = pipe(person_image, dress_image).images[0]  # Update based on actual pipeline usage

    return output_image

# Define Gradio interface
iface = gr.Interface(
    fn=virtual_tryon,
    inputs=[gr.inputs.Image(type="pil"), gr.inputs.Image(type="pil")],
    outputs=gr.outputs.Image(type="pil"),
    title="Virtual Try-On",
    description="Upload a picture of a person and a dress, and see the result of the virtual try-on."
)

# Launch the interface
iface.launch()