diff options
author | AUTOMATIC1111 <16777216c@gmail.com> | 2023-01-04 15:57:14 +0000 |
---|---|---|
committer | GitHub <noreply@github.com> | 2023-01-04 15:57:14 +0000 |
commit | 32547f2721c92794779e6ff9fb325243d5857cae (patch) | |
tree | d4d5f1a9705e59eef5029cc1be3bff57fbd389c2 /modules/bsrgan_model.py | |
parent | fe6e2362e8fa5d739de6997ab155a26686d20a49 (diff) | |
parent | 3dae545a03f5102ba5d9c3f27bb6241824c5a916 (diff) | |
download | stable-diffusion-webui-gfx803-32547f2721c92794779e6ff9fb325243d5857cae.tar.gz stable-diffusion-webui-gfx803-32547f2721c92794779e6ff9fb325243d5857cae.tar.bz2 stable-diffusion-webui-gfx803-32547f2721c92794779e6ff9fb325243d5857cae.zip |
Merge branch 'master' into xygrid_infotext_improvements
Diffstat (limited to 'modules/bsrgan_model.py')
-rw-r--r-- | modules/bsrgan_model.py | 78 |
1 files changed, 0 insertions, 78 deletions
diff --git a/modules/bsrgan_model.py b/modules/bsrgan_model.py deleted file mode 100644 index e62c6657..00000000 --- a/modules/bsrgan_model.py +++ /dev/null @@ -1,78 +0,0 @@ -import os.path -import sys -import traceback - -import PIL.Image -import numpy as np -import torch -from basicsr.utils.download_util import load_file_from_url - -import modules.upscaler -from modules import shared, modelloader -from modules.bsrgan_model_arch import RRDBNet -from modules.paths import models_path - - -class UpscalerBSRGAN(modules.upscaler.Upscaler): - def __init__(self, dirname): - self.name = "BSRGAN" - self.model_path = os.path.join(models_path, self.name) - self.model_name = "BSRGAN 4x" - self.model_url = "https://github.com/cszn/KAIR/releases/download/v1.0/BSRGAN.pth" - self.user_path = dirname - super().__init__() - model_paths = self.find_models(ext_filter=[".pt", ".pth"]) - scalers = [] - if len(model_paths) == 0: - scaler_data = modules.upscaler.UpscalerData(self.model_name, self.model_url, self, 4) - scalers.append(scaler_data) - for file in model_paths: - if "http" in file: - name = self.model_name - else: - name = modelloader.friendly_name(file) - try: - scaler_data = modules.upscaler.UpscalerData(name, file, self, 4) - scalers.append(scaler_data) - except Exception: - print(f"Error loading BSRGAN model: {file}", file=sys.stderr) - print(traceback.format_exc(), file=sys.stderr) - self.scalers = scalers - - def do_upscale(self, img: PIL.Image, selected_file): - torch.cuda.empty_cache() - model = self.load_model(selected_file) - if model is None: - return img - model.to(shared.device) - torch.cuda.empty_cache() - img = np.array(img) - img = img[:, :, ::-1] - img = np.moveaxis(img, 2, 0) / 255 - img = torch.from_numpy(img).float() - img = img.unsqueeze(0).to(shared.device) - with torch.no_grad(): - output = model(img) - output = output.squeeze().float().cpu().clamp_(0, 1).numpy() - output = 255. * np.moveaxis(output, 0, 2) - output = output.astype(np.uint8) - output = output[:, :, ::-1] - torch.cuda.empty_cache() - return PIL.Image.fromarray(output, 'RGB') - - def load_model(self, path: str): - if "http" in path: - filename = load_file_from_url(url=self.model_url, model_dir=self.model_path, file_name="%s.pth" % self.name, - progress=True) - else: - filename = path - if not os.path.exists(filename) or filename is None: - print(f"BSRGAN: Unable to load model from {filename}", file=sys.stderr) - return None - model = RRDBNet(in_nc=3, out_nc=3, nf=64, nb=23, gc=32, sf=4) # define network - model.load_state_dict(torch.load(filename), strict=True) - model.eval() - for k, v in model.named_parameters(): - v.requires_grad = False - return model - |