File size: 697 Bytes
b38f27e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
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 |