import gradio as gr
from diffusers import StableDiffusionUpscalePipeline
from diffusers.utils import load_image
import torch
from PIL import Image
import base64
from io import BytesIO

# Load model and scheduler
model_id = "stabilityai/stable-diffusion-x4-upscaler"
pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id)
pipeline = pipeline.to("cpu")  # Use CPU instead of GPU

def upscale_image(image, prompt):
    image = image.resize((128, 128))  # Resize to the expected input size
    upscaled_image = pipeline(prompt=prompt, image=image).images[0]
    return upscaled_image

def image_to_base64(image):
    buffered = BytesIO()
    image.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode()

def base64_to_image(base64_str):
    image_data = base64.b64decode(base64_str)
    return Image.open(BytesIO(image_data))

def handle_upload(image, prompt):
    upscaled_image = upscale_image(image, prompt)
    base64_str = image_to_base64(upscaled_image)
    return base64_str, upscaled_image

def main():
    with gr.Blocks() as demo:
        gr.Markdown("# Stable Diffusion Upscaler")

        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(type="pil", label="Low-Resolution Image")
                prompt_input = gr.Textbox(label="Prompt", value="a white cat")

                upload_btn = gr.Button("Upload and Upscale")
                base64_output = gr.Textbox(label="Base64 Encoded Image")
                upscaled_image_output = gr.Image(type="pil", label="Upscaled Image")

                upload_btn.click(fn=handle_upload, inputs=[image_input, prompt_input], outputs=[base64_output, upscaled_image_output])
            
    demo.launch()

if __name__ == "__main__":
    main()