import torch

from PIL import Image
from diffusers import StableDiffusionUpscalePipeline


device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "stabilityai/stable-diffusion-x4-upscaler"
upscale_pipe = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
upscale_pipe = upscale_pipe.to(device)

def upscale_image(
    input_image: Image,
    prompt: str,
    start_size: int = 128,
    upscale_steps: int = 30,
):
    input_image = input_image.resize((start_size, start_size))
    upscaled_image = upscale_pipe(
        prompt=prompt, 
        image=input_image,
        num_inference_steps=upscale_steps,
    ).images[0]
    
    return upscaled_image