import gradio as gr
import torch
import modin.pandas as pd
from PIL import Image
from io import BytesIO
from diffusers import StableDiffusionLatentUpscalePipeline, StableDiffusionUpscalePipeline

device = "cuda" if torch.cuda.is_available() else "cpu"

# Define the models
model_2x = "stabilityai/sd-x2-latent-upscaler"
model_4x = "stabilityai/stable-diffusion-x4-upscaler"
# Load the models
sd_2_0_2x = StableDiffusionLatentUpscalePipelin.from_pretrained(model_2x, torch_dtype=torch.float16, revision="fp16") if torch.cuda.is_available() else StableDiffusionLatentUpscalePipeline.from_pretrained(model_2x)
sd_2_1_4x = StableDiffusionUpscalePipeline.from_pretrained(model_4x, torch_dtype=torch.float16, revision="fp16") if torch.cuda.is_available() else StableDiffusionUpscalePipeline.from_pretrained(model_4x)

# Define the function that will be called when the interface is used

def upscale_image(model, input_image):
    # Convert the image 
    generator = torch.manual_seed(999999)
    input_image = Image.open(input_image).convert("RGB")
    # Upscale the image using the selected model
    if model == "SD 2.0 2x Latent Upscaler":
        upscaled_image = sd_2_0_2x(prompt="", image=input_image, num_inference_steps=5).images[0]
    else:
        low_res_img = input_image.resize((128, 128))
        upscaled_image = sd_2_1_4x(prompt="", image=low_res_img, num_inference_steps=5).images[0]
    # Return the upscaled image
    return upscaled_image

# Define the Gradio interface
iface = gr.Interface(
    fn=upscale_image,
    inputs=[gr.Radio(["SD 2.0 2x Latent Upscaler", "SD 2.1 4x Upscaler"], label="Models:"), gr.Image(type="filepath", label = "Raw Image")],
    outputs=gr.Image(type="filepath", label = "Upscaled Image"),
    title="SD Image Upscaler",
    description="Upscale an image using either the SD 2.0 2x Latent Upscaler or the SD 2.1 4x Upscaler. Use the 4x Upscaler for images lower than 512x512. Use the 2x Upscaler for images 512x512 to 768x768"
)

# Launch the interface
iface.launch(debug=True)