import gradio as gr
from diffusers import StableDiffusionInstructPix2PixPipeline
import torch

# Load the model
model_id = "timbrooks/instruct-pix2pix"
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
).to("cuda")

# Define the function
def edit_image(image, prompt):
    image = pipe(prompt, image=image, num_inference_steps=20).images[0]
    return image

# Create the UI
demo = gr.Interface(
    fn=edit_image,
    inputs=[gr.Image(label="Upload Image"), gr.Textbox(label="Instruction")],
    outputs=gr.Image(label="Result"),
    examples=[["input_example.jpg", "Make the sky stormy"]]
)

demo.launch()