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

# Define the path to the SafeTensor model
model_path = "/content/uberRealisticPornMerge_urpmv12.instruct-pix2pix.safetensors"

# Load the SafeTensor model
safe_pipe = torch.load(model_path, map_location=torch.device('cuda'))

def generate_edited_image(input_image):
    # Convert the Gradio Image object to a PIL Image
    input_image_pil = Image.fromarray(input_image.astype('uint8'), 'RGB')
    
    # Generate the edited image using the SafeTensor model
    edited_image = safe_pipe(instruction="", image=input_image_pil, num_inference_steps=50).images[0]
    
    # Convert the edited image back to Gradio Image format
    edited_image_gradio = edited_image.cpu().numpy().astype('uint8')
    
    return edited_image_gradio

# Define the input and output components for the Gradio app
input_image = gr.inputs.Image(label="Upload an Input Image")
output_image = gr.outputs.Image(label="Edited Image")

# Create the Gradio interface
gr.Interface(
    fn=generate_edited_image,
    inputs=input_image,
    outputs=output_image,
    title="SafeTensor Image Editing",
    description="Upload an image and generate an edited image using a SafeTensor model.",
    capture_session=True  # This ensures that we use the same session for model inference
).launch()