| 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, | |
| seed: int = 42, | |
| ): | |
| generator = torch.Generator().manual_seed(seed) | |
| input_image = input_image.resize((start_size, start_size)) | |
| upscaled_image = upscale_pipe( | |
| prompt=prompt, | |
| image=input_image, | |
| num_inference_steps=upscale_steps, | |
| generator=generator, | |
| ).images[0] | |
| return upscaled_image |