import gradio as gr
import torch
from torchvision.transforms import ToPILImage, ToTensor
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 input and output components for the Gradio interface
input_image = gr.inputs.Image(type="filepath")
output_image = gr.outputs.Image(type="filepath")

# Define the function that will be called when the interface is used
def upscale_image(model, image):
    # Convert the image to a PyTorch tensor
    image_tensor = ToTensor()(image)
    
    # Upscale the image using the selected model
    if model == "SD 2.0 2x Latent Upscaler":
        upscaled_tensor = sd_2_0_2x(image_tensor)
    else:
        upscaled_tensor = sd_2_1_4x(image_tensor)
    
    # Convert the upscaled tensor back to a PIL image
    upscaled_image = ToPILImage()(upscaled_tensor)
    
    # 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"]), input_image],
    outputs=output_image,
    title="Image Upscaler",
    description="Upscale an image using either the SD 2.0 2x Latent Upscaler or the SD 2.1 4x Upscaler."
)

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