import gradio as gr
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import make_image_grid
from PIL import Image

# Load the pipeline
pipeline = AutoPipelineForImage2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0", 
    torch_dtype=torch.float16, 
    variant="fp16", 
    use_safetensors=True
)

# Offload model to reduce memory usage
pipeline.enable_model_cpu_offload()

# Gradio function for image generation
def generate_image(prompt, init_image, strength):
    result_image = pipeline(prompt, image=init_image, strength=strength).images[0]
    
    # Display both the initial and result images side by side
    grid_image = make_image_grid([init_image, result_image], rows=1, cols=2)
    return grid_image

# Gradio interface
gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(lines=1, label="Prompt", placeholder="Enter the image description prompt"),
        gr.Image(type="pil", label="Upload Initial Image"),  # Image input instead of URL
        gr.Slider(0.0, 1.0, value=0.5, label="Strength"),
    ],
    outputs=gr.Image(label="Image Comparison"),
    title="Stable Diffusion XL Refiner - Image to Image",
    description="Upload an initial image and provide a text prompt to generate a new image using the Stable Diffusion XL Refiner model.",
).launch()