Spaces:
Runtime error
Runtime error
import os | |
from abc import abstractmethod | |
import PIL | |
import numpy as np | |
import torch | |
from PIL import Image | |
import modules.shared | |
from modules import modelloader, shared | |
LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS) | |
from modules.paths import models_path | |
class Upscaler: | |
name = None | |
model_path = None | |
model_name = None | |
model_url = None | |
enable = True | |
filter = None | |
model = None | |
user_path = None | |
scalers: [] | |
tile = True | |
def __init__(self, create_dirs=False): | |
self.mod_pad_h = None | |
self.tile_size = modules.shared.opts.ESRGAN_tile | |
self.tile_pad = modules.shared.opts.ESRGAN_tile_overlap | |
self.device = modules.shared.device | |
self.img = None | |
self.output = None | |
self.scale = 1 | |
self.half = not modules.shared.cmd_opts.no_half | |
self.pre_pad = 0 | |
self.mod_scale = None | |
if self.name is not None and create_dirs: | |
self.model_path = os.path.join(models_path, self.name) | |
if not os.path.exists(self.model_path): | |
os.makedirs(self.model_path) | |
try: | |
import cv2 | |
self.can_tile = True | |
except: | |
pass | |
def do_upscale(self, img: PIL.Image, selected_model: str): | |
return img | |
def upscale(self, img: PIL.Image, scale: int, selected_model: str = None): | |
self.scale = scale | |
dest_w = img.width * scale | |
dest_h = img.height * scale | |
for i in range(3): | |
if img.width >= dest_w and img.height >= dest_h: | |
break | |
img = self.do_upscale(img, selected_model) | |
if img.width != dest_w or img.height != dest_h: | |
img = img.resize((int(dest_w), int(dest_h)), resample=LANCZOS) | |
return img | |
def load_model(self, path: str): | |
pass | |
def find_models(self, ext_filter=None) -> list: | |
return modelloader.load_models(model_path=self.model_path, model_url=self.model_url, command_path=self.user_path) | |
def update_status(self, prompt): | |
print(f"\nextras: {prompt}", file=shared.progress_print_out) | |
class UpscalerData: | |
name = None | |
data_path = None | |
scale: int = 4 | |
scaler: Upscaler = None | |
model: None | |
def __init__(self, name: str, path: str, upscaler: Upscaler = None, scale: int = 4, model=None): | |
self.name = name | |
self.data_path = path | |
self.scaler = upscaler | |
self.scale = scale | |
self.model = model | |
class UpscalerNone(Upscaler): | |
name = "None" | |
scalers = [] | |
def load_model(self, path): | |
pass | |
def do_upscale(self, img, selected_model=None): | |
return img | |
def __init__(self, dirname=None): | |
super().__init__(False) | |
self.scalers = [UpscalerData("None", None, self)] | |
class UpscalerLanczos(Upscaler): | |
scalers = [] | |
def do_upscale(self, img, selected_model=None): | |
return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=LANCZOS) | |
def load_model(self, _): | |
pass | |
def __init__(self, dirname=None): | |
super().__init__(False) | |
self.name = "Lanczos" | |
self.scalers = [UpscalerData("Lanczos", None, self)] | |