You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
"""Functions for image upscaling using RealESRGAN"""
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from imaginairy.model_manager import get_cached_url_path
|
|
from imaginairy.utils import get_device
|
|
from imaginairy.utils.model_cache import memory_managed_model
|
|
from imaginairy.vendored.basicsr.rrdbnet_arch import RRDBNet
|
|
from imaginairy.vendored.realesrgan import RealESRGANer
|
|
|
|
|
|
@memory_managed_model("realesrgan_upsampler", memory_usage_mb=70)
|
|
def realesrgan_upsampler():
|
|
model = RRDBNet(
|
|
num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4
|
|
)
|
|
url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
|
|
model_path = get_cached_url_path(url)
|
|
device = get_device()
|
|
|
|
upsampler = RealESRGANer(
|
|
scale=4, model_path=model_path, model=model, tile=512, device=device
|
|
)
|
|
|
|
upsampler.device = torch.device(device)
|
|
upsampler.model.to(device)
|
|
|
|
return upsampler
|
|
|
|
|
|
def upscale_image(img):
|
|
img = img.convert("RGB")
|
|
|
|
np_img = np.array(img, dtype=np.uint8)
|
|
upsampler_output, img_mode = realesrgan_upsampler().enhance(np_img[:, :, ::-1])
|
|
return Image.fromarray(upsampler_output[:, :, ::-1], mode=img_mode)
|