import gradio as gr
from PIL import Image
from diffusers import StableDiffusionUpscalePipeline
import torch
import numpy as np

# Load model and scheduler
model_id = "stabilityai/stable-diffusion-x4-upscaler"
pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipeline = pipeline.to("cuda")

def upscale_image(image, prompt):
    # Convert uploaded image to PIL
    low_res_img = Image.fromarray(image).convert("RGB")

    # Upscale the image
    upscaled_image = pipeline(prompt=prompt, image=low_res_img).images[0]

    # Convert upscaled PIL image back to numpy array for Gradio
    upscaled_image_np = np.array(upscaled_image)
    
    return upscaled_image_np

# Create the Gradio interface
interface = gr.Interface(
    fn=upscale_image,
    inputs=[
        gr.Image(type="numpy", label="Upload Low-Resolution Image"),
        gr.Textbox(label="Upscaling Prompt", placeholder="Enter a prompt, e.g., 'a red box with glasses'")
    ],
    outputs=gr.Image(type="numpy", label="Upscaled Image"),
    title="Image Upscaler",
    description="Upload a low-resolution image and provide a prompt to upscale it using Stable Diffusion."
)

# Launch the Gradio app
interface.launch(share=True)