File size: 4,068 Bytes
9dce458 |
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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
import os
import subprocess
import tempfile
from sys import platform
from typing import List
from PIL import Image
import shutil
from .common import OfflineUpscaler
if platform == 'win32':
waifu2x_base_folder = 'waifu2x-win'
waifu2x_executable_path = os.path.join(waifu2x_base_folder, 'waifu2x-ncnn-vulkan.exe')
model_mapping = {
'waifu2x-win': {
'url': 'https://github.com/nihui/waifu2x-ncnn-vulkan/releases/download/20220728/waifu2x-ncnn-vulkan-20220728-windows.zip',
'hash': '3f60ba0b26763c602cb75178c2051bf0c46f3cc9d13975a052a902773988a34b',
'archive': {
'waifu2x-ncnn-vulkan-20220728-windows': waifu2x_base_folder,
},
},
}
elif platform == 'darwin':
waifu2x_base_folder = 'waifu2x-macos'
waifu2x_executable_path = os.path.join(waifu2x_base_folder, 'waifu2x-ncnn-vulkan')
model_mapping = {
'waifu2x-macos': {
'url': 'https://github.com/nihui/waifu2x-ncnn-vulkan/releases/download/20220728/waifu2x-ncnn-vulkan-20220728-macos.zip',
'hash': '9801839aa1a73a3a22e86d05f09a9b7412c289d1a03215e1b2713cc969690ba4',
'archive': {
'waifu2x-ncnn-vulkan-20220728-macos': waifu2x_base_folder,
},
},
}
else:
waifu2x_base_folder = 'waifu2x-linux'
waifu2x_executable_path = os.path.join(waifu2x_base_folder, 'waifu2x-ncnn-vulkan')
model_mapping = {
'waifu2x-linux': {
'url': 'https://github.com/nihui/waifu2x-ncnn-vulkan/releases/download/20220728/waifu2x-ncnn-vulkan-20220728-ubuntu.zip',
'hash': 'f2244412aeaf474d58e262f636737abca24ee24cd632d86eb8f0a4c4f9649aaa',
'archive': {
'waifu2x-ncnn-vulkan-20220728-ubuntu': waifu2x_base_folder,
},
'executables': [
waifu2x_executable_path
],
},
}
# https://github.com/nihui/waifu2x-ncnn-vulkan
class Waifu2xUpscaler(OfflineUpscaler): # ~2GB of vram
_MODEL_MAPPING = model_mapping
_VALID_UPSCALE_RATIOS = [2, 4, 8, 16, 32]
def __init__(self, *args, **kwargs):
os.makedirs(self.model_dir, exist_ok=True)
if os.path.exists(os.path.join('models', waifu2x_base_folder)):
shutil.move(os.path.join('models', waifu2x_base_folder), self._get_file_path(waifu2x_base_folder))
super().__init__(*args, **kwargs)
async def _load(self, device: str):
pass
async def _unload(self):
pass
async def _infer(self, image_batch: List[Image.Image], upscale_ratio: float) -> List[Image.Image]:
# Has to cache images because chosen upscaler doesn't support piping
in_dir = tempfile.mkdtemp()
out_dir = tempfile.mkdtemp()
for i, image in enumerate(image_batch):
image.save(os.path.join(in_dir, f'{i}.png'))
try:
self._run_waifu2x_executable(in_dir, out_dir, upscale_ratio, 0)
except Exception:
# Maybe throw exception instead
self.logger.warn(f'Process returned non-zero exit status. Skipping upscaling.')
return image_batch
output_batch = []
for i, image in enumerate(image_batch):
img_path = os.path.join(out_dir, f'{i}.png')
if os.path.exists(img_path):
img = Image.open(img_path)
img.load()
output_batch.append(img)
else:
output_batch.append(image)
shutil.rmtree(in_dir)
shutil.rmtree(out_dir)
return output_batch
def _run_waifu2x_executable(self, image_directory: str, output_directory: str, upscale_ratio: float, denoise_level: int):
cmds = [
self._get_file_path(waifu2x_executable_path),
'-i', image_directory,
'-o', output_directory,
'-m', self._get_file_path(os.path.join(waifu2x_base_folder, 'models-cunet')),
'-s', str(upscale_ratio),
'-n', str(denoise_level),
]
subprocess.check_call(cmds)
|