import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from diffusers import StableDiffusionPipeline
import gradio as gr

# Load the model and pipeline
model_id = "ares1123/virtual-dress-try-on"
pipeline = StableDiffusionPipeline.from_pretrained(model_id)
pipeline.to("cuda" if torch.cuda.is_available() else "cpu")

def virtual_try_on(image, clothing_image):
    # Process the images using the model
    try_on_image = pipeline(image, clothing_image).images[0]
    return try_on_image

# Set up a simple Gradio interface for testing
interface = gr.Interface(
    fn=virtual_try_on,
    inputs=[gr.inputs.Image(type="pil", label="User Image"), 
            gr.inputs.Image(type="pil", label="Clothing Image")],
    outputs="image",
    title="Virtual Dress Try-On",
    description="Upload an image of yourself and a clothing image to try it on virtually!"
)

# Launch the interface
interface.launch(share=True)