import torch
from diffusers import StableDiffusionInstructPix2PixPipeline
import gradio as gr
from PIL import Image

# Ensure we use GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"

def load_model():
    """Load InstructPix2Pix model."""
    pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
        "timbrooks/instruct-pix2pix",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
    ).to(device)
    
    return pipe

pipe = load_model()

def edit_image(image, prompt):
    """Perform image editing using InstructPix2Pix model."""
    if image is None:
        raise ValueError("Please upload an image.")
    
    image = Image.fromarray(image).convert("RGB")  # Ensure image is in RGB format
    
    result = pipe(image=image, prompt=prompt, num_inference_steps=10, image_guidance_scale=1.5).images[0]
    return result

# Gradio UI
demo = gr.Interface(
    fn=edit_image,
    inputs=[
        gr.Image(type="numpy", label="Upload Image"),
        gr.Textbox(label="Editing Instruction"),
    ],
    outputs=gr.Image(type="pil", label="Edited Image"),
    title="InstructPix2Pix Image Editing",
    description="Upload an image and provide a text instruction to modify it.",
)

demo.launch()