import gradio as gr
from diffusers import StableDiffusionInstructPix2PixPipeline
import torch

# Load the uploaded model from Hugging Face
model = StableDiffusionInstructPix2PixPipeline.from_pretrained("poojan1202/image_modification")
model.to("cuda")  # If you have access to GPU

# Define the image-to-image function
def instruct_pix2pix(image, instruction):
    edited_image = model(image=image, prompt=instruction).images[0]
    return edited_image

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# InstructPix2Pix: Image-to-Image Editing")
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(source="upload", type="pil")
            instruction_input = gr.Textbox(label="Instruction", placeholder="Describe the edit you want...")
            submit_button = gr.Button("Generate")
        with gr.Column():
            image_output = gr.Image()

    submit_button.click(instruct_pix2pix, inputs=[image_input, instruction_input], outputs=image_output)

# Launch the app
demo.launch()