From b35a84f505545aee221d49ba74c92379b239b5ed Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 18 Sep 2023 14:48:41 -0400 Subject: [PATCH] add upscaler logging --- .../SwinIR/scripts/swinir_model.py | 12 +++++--- modules/esrgan_model.py | 30 +++++++++++-------- modules/gfpgan_model.py | 3 +- modules/images.py | 2 +- modules/realesrgan_model.py | 21 ++++++------- modules/ui_common.py | 2 +- 6 files changed, 39 insertions(+), 31 deletions(-) diff --git a/extensions-builtin/SwinIR/scripts/swinir_model.py b/extensions-builtin/SwinIR/scripts/swinir_model.py index e470db874..7df4a50b8 100644 --- a/extensions-builtin/SwinIR/scripts/swinir_model.py +++ b/extensions-builtin/SwinIR/scripts/swinir_model.py @@ -2,8 +2,8 @@ import os import numpy as np import torch from PIL import Image +from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn from basicsr.utils.download_util import load_file_from_url -from tqdm.rich import tqdm from swinir_model_arch import SwinIR as net from swinir_model_arch_v2 import Swin2SR as net2 from modules import modelloader, devices, script_callbacks, shared @@ -45,11 +45,12 @@ class UpscalerSwinIR(Upscaler): def load_model(self, path, scale=4): if "http" in path: - dl_name = "%s%s" % (self.model_name.replace(" ", "_"), ".pth") + dl_name = "%s%s" % (self.model_name.replace(" ", "_"), ".pth") # pylint: disable=consider-using-f-string filename = load_file_from_url(url=path, model_dir=self.model_download_path, file_name=dl_name, progress=True) else: filename = path if filename is None or not os.path.exists(filename): + shared.log.error(f"Model failed loading: type=SwinIR model={filename}") return None model_v2 = net2( upscale=scale, @@ -78,6 +79,8 @@ class UpscalerSwinIR(Upscaler): resi_connection="3conv", ) pretrained_model = torch.load(filename) + shared.log.info(f"Model loaded: type=SwinIR model={filename}") + for model in [model_v1, model_v2]: for param in ["params_ema", "params", None]: try: @@ -140,7 +143,8 @@ def inference(img, model, tile, tile_overlap, window_size, scale): E = torch.zeros(b, c, h * sf, w * sf, dtype=devices.dtype, device=device_swinir).type_as(img) W = torch.zeros_like(E, dtype=devices.dtype, device=device_swinir) - with tqdm(total=len(h_idx_list) * len(w_idx_list), desc="Upscaling SwinIR") as pbar: + with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=shared.console) as progress: + task = progress.add_task(description="Upscaling Initializing", total=len(h_idx_list) * len(w_idx_list)) for h_idx in h_idx_list: if state.interrupted or state.skipped: break @@ -159,7 +163,7 @@ def inference(img, model, tile, tile_overlap, window_size, scale): W[ ..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf ].add_(out_patch_mask) - pbar.update(1) + progress.update(task, advance=1, description="Upscaling") output = E.div_(W) return output diff --git a/modules/esrgan_model.py b/modules/esrgan_model.py index 7ed5c33ec..3a4bf599b 100644 --- a/modules/esrgan_model.py +++ b/modules/esrgan_model.py @@ -4,11 +4,12 @@ import numpy as np import torch from PIL import Image from basicsr.utils.download_util import load_file_from_url +from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn import modules.esrgan_model_arch as arch from modules import modelloader, images, devices from modules.upscaler import Upscaler, UpscalerData -from modules.shared import opts +from modules.shared import opts, log, console @@ -161,10 +162,11 @@ class UpscalerESRGAN(Upscaler): else: filename = path if not os.path.exists(filename) or filename is None: - print(f"Unable to load {self.model_path} from {filename}") + log.error(f"Model failed loading: type=ESRGAN model={filename}") return None state_dict = torch.load(filename, map_location='cpu' if devices.device_esrgan.type == 'mps' else None) + log.info(f"Model loaded: type=ESRGAN model={filename}") if "params_ema" in state_dict: state_dict = state_dict["params_ema"] @@ -216,16 +218,20 @@ def esrgan_upscale(model, img): newtiles = [] scale_factor = 1 - for y, h, row in grid.tiles: - newrow = [] - for tiledata in row: - x, w, tile = tiledata - - output = upscale_without_tiling(model, tile) - scale_factor = output.width // tile.width - - newrow.append([x * scale_factor, w * scale_factor, output]) - newtiles.append([y * scale_factor, h * scale_factor, newrow]) + with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=console) as progress: + total = 0 + for y, h, row in grid.tiles: + total += len(row) + task = progress.add_task(description="Upscaling", total=total) + for y, h, row in grid.tiles: + newrow = [] + for tiledata in row: + x, w, tile = tiledata + output = upscale_without_tiling(model, tile) + scale_factor = output.width // tile.width + newrow.append([x * scale_factor, w * scale_factor, output]) + progress.update(task, advance=1, description="Upscaling") + newtiles.append([y * scale_factor, h * scale_factor, newrow]) newgrid = images.Grid(newtiles, grid.tile_w * scale_factor, grid.tile_h * scale_factor, grid.image_w * scale_factor, grid.image_h * scale_factor, grid.overlap * scale_factor) output = images.combine_grid(newgrid) diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py index 728df70bf..2988cd407 100644 --- a/modules/gfpgan_model.py +++ b/modules/gfpgan_model.py @@ -29,12 +29,13 @@ def gfpgann(): latest_file = max(models, key=os.path.getctime) model_file = latest_file else: - print("Unable to load gfpgan model!") + shared.log.error(f"Model failed loading: type=GFPGAN model={model_file}") return None if hasattr(facexlib.detection.retinaface, 'device'): facexlib.detection.retinaface.device = devices.device_gfpgan model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None, device=devices.device_gfpgan) loaded_gfpgan_model = model + shared.log.info(f"Model loaded: type=GFPGAN model={model_file}") return model diff --git a/modules/images.py b/modules/images.py index 39d9e7246..3a2e8ded0 100644 --- a/modules/images.py +++ b/modules/images.py @@ -320,7 +320,7 @@ class FilenameGenerator: } default_time_format = '%Y%m%d%H%M%S' - def __init__(self, p, seed, prompt, image, index): + def __init__(self, p, seed, prompt, image, index = 0): self.p = p self.seed = seed self.prompt = prompt diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py index 5f5b132e5..318fe21fc 100644 --- a/modules/realesrgan_model.py +++ b/modules/realesrgan_model.py @@ -1,12 +1,10 @@ import os -import sys import numpy as np from PIL import Image from basicsr.utils.download_util import load_file_from_url from modules.upscaler import Upscaler, UpscalerData -from modules.shared import opts, device +from modules.shared import opts, device, log from modules import modelloader -import modules.errors as errors class UpscalerRealESRGAN(Upscaler): @@ -30,9 +28,8 @@ class UpscalerRealESRGAN(Upscaler): scaler.local_data_path = local_model_candidates[0] if scaler.name in opts.realesrgan_enabled_models: self.scalers.append(scaler) - except Exception as e: - errors.display(e, 'real-esrgan') + log.error(f"Error loading Real-ESRGAN: model={path} {e}") self.enable = False self.scalers = [] @@ -43,12 +40,11 @@ class UpscalerRealESRGAN(Upscaler): try: from realesrgan import RealESRGANer except Exception: - print("Error importing Real-ESRGAN:", file=sys.stderr) + log.error("Error importing Real-ESRGAN:") return img info = self.load_model(selected_model) - if not os.path.exists(info.local_data_path): - print(f"Unable to load RealESRGAN model: {info.name}") + if info is None or not os.path.exists(info.local_data_path): return img upsampler = RealESRGANer( @@ -70,13 +66,14 @@ class UpscalerRealESRGAN(Upscaler): try: info = next(iter([scaler for scaler in self.scalers if scaler.data_path == path]), None) if info is None: - print(f"Unable to find model info: {path}") + log.error(f"Model failed loading: type=R-ESRGAN model={info.name}") return None if info.local_data_path.startswith("http"): info.local_data_path = load_file_from_url(url=info.data_path, model_dir=self.model_download_path, progress=True) + log.info(f"Model loaded: type=R-ESRGAN model={info.name}") return info except Exception as e: - errors.display(e, 'real-esrgan model list') + log.error(f"Model failed loading: type=R-ESRGAN model={info.name} {e}") return None def load_models(self, _): @@ -132,6 +129,6 @@ def get_realesrgan_models(scaler): ), ] return models - except Exception: - print("Error creating Real-ESRGAN models list", file=sys.stderr) + except Exception as e: + log.error(f'Error creating Real-ESRGAN models list: {e}') return [] diff --git a/modules/ui_common.py b/modules/ui_common.py index b4f31734c..cac1363a9 100644 --- a/modules/ui_common.py +++ b/modules/ui_common.py @@ -110,7 +110,7 @@ def save_files(js_data, images, html_info, index): fullfns.append(fullfn) destination = shared.opts.outdir_save if shared.opts.use_save_to_dirs_for_ui: - namegen = modules.images.FilenameGenerator(p, seed=p.all_seeds[i], prompt=p.all_prompts[i], image=None) # pylint: disable=no-member + namegen = modules.images.FilenameGenerator(p, seed=p.all_seeds[i], prompt=p.all_prompts[i], image=None, index=image_index) # pylint: disable=no-member dirname = namegen.apply(shared.opts.directories_filename_pattern or "[prompt_words]").lstrip(' ').rstrip('\\ /') destination = os.path.join(destination, dirname) os.makedirs(destination, exist_ok = True)