diff --git a/CHANGELOG.md b/CHANGELOG.md index a8bc842a4..5e77b46eb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -204,11 +204,14 @@ And other goodies like multiple *XYZ grid* improvements, additional *Flux Contro - **extensions** - add mechanism to lock-down extension to specific working commit - added `sd-webui-controlnet` and `adetailer` last-known working commits +- **upscaling** + - interruptible operations - **refactor** - modularize main process loop - massive log cleanup - full lint pass - improve inference mode handling + - unify quant lib loading ## Update for 2024-09-13 diff --git a/modules/postprocess/esrgan_model.py b/modules/postprocess/esrgan_model.py index b827f334e..8ff75abab 100644 --- a/modules/postprocess/esrgan_model.py +++ b/modules/postprocess/esrgan_model.py @@ -3,9 +3,8 @@ import torch from PIL import Image from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn import modules.postprocess.esrgan_model_arch as arch -from modules import images, devices +from modules import images, devices, shared from modules.upscaler import Upscaler, UpscalerData, compile_upscaler -from modules.shared import opts, log, console def mod2normal(state_dict): @@ -130,9 +129,9 @@ class UpscalerESRGAN(Upscaler): return img model.to(devices.device) img = esrgan_upscale(model, img) - if opts.upscaler_unload and selected_model in self.models: + if shared.opts.upscaler_unload and selected_model in self.models: del self.models[selected_model] - log.debug(f"Upscaler unloaded: type={self.name} model={selected_model}") + shared.log.debug(f"Upscaler unloaded: type={self.name} model={selected_model}") devices.torch_gc(force=True) return img @@ -141,10 +140,10 @@ class UpscalerESRGAN(Upscaler): if info is None: return if self.models.get(info.local_data_path, None) is not None: - log.debug(f"Upscaler cached: type={self.name} model={info.local_data_path}") + shared.log.debug(f"Upscaler cached: type={self.name} model={info.local_data_path}") return self.models[info.local_data_path] state_dict = torch.load(info.local_data_path, map_location='cpu' if devices.device.type == 'mps' else None) - log.info(f"Upscaler loaded: type={self.name} model={info.local_data_path}") + shared.log.info(f"Upscaler loaded: type={self.name} model={info.local_data_path}") if "params_ema" in state_dict: state_dict = state_dict["params_ema"] @@ -190,21 +189,25 @@ def upscale_without_tiling(model, img): def esrgan_upscale(model, img): - if opts.upscaler_tile_size == 0: + if shared.opts.upscaler_tile_size == 0: return upscale_without_tiling(model, img) - grid = images.split_grid(img, opts.upscaler_tile_size, opts.upscaler_tile_size, opts.upscaler_tile_overlap) + grid = images.split_grid(img, shared.opts.upscaler_tile_size, shared.opts.upscaler_tile_size, shared.opts.upscaler_tile_overlap) newtiles = [] scale_factor = 1 - with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=console) as progress: + with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=shared.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: + if shared.state.interrupted: + break newrow = [] for tiledata in row: + if shared.state.interrupted: + break x, w, tile = tiledata output = upscale_without_tiling(model, tile) scale_factor = output.width // tile.width diff --git a/modules/postprocess/realesrgan_model_arch.py b/modules/postprocess/realesrgan_model_arch.py index 30b8e65ac..7947b0c7d 100644 --- a/modules/postprocess/realesrgan_model_arch.py +++ b/modules/postprocess/realesrgan_model_arch.py @@ -8,8 +8,7 @@ import torch from torch import nn from torch.nn import functional as F from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn -from modules import devices -from modules.shared import log, console +from modules import devices, shared from modules.upscaler import compile_upscaler ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -67,7 +66,7 @@ class RealESRGANer(): from modules.modelloader import load_file_from_url model_path = load_file_from_url(url=model_path, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None) loadnet = torch.load(model_path, map_location=torch.device('cpu')) - log.info(f"Upscaler loaded: type={self.name} model={model_path}") + shared.log.info(f"Upscaler loaded: type={self.name} model={model_path}") # prefer to use params_ema if 'params_ema' in loadnet: @@ -139,11 +138,15 @@ class RealESRGANer(): tiles_y = math.ceil(height / self.tile_size) # loop over all tiles - with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=console) as progress: + with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=shared.console) as progress: task = progress.add_task(description="Upscaling", total=tiles_y * tiles_x) with torch.no_grad(): for y in range(tiles_y): + if shared.state.interrupted: + break for x in range(tiles_x): + if shared.state.interrupted: + break # extract tile from input image ofs_x = x * self.tile_size ofs_y = y * self.tile_size @@ -169,7 +172,7 @@ class RealESRGANer(): try: output_tile = self.model(input_tile) except Exception as e: - log.error(f'Upscale error: type=R-ESRGAN {e}') + shared.log.error(f'Upscale error: type=R-ESRGAN {e}') # output tile area on total image output_start_x = input_start_x * self.scale diff --git a/modules/postprocess/swinir_model.py b/modules/postprocess/swinir_model.py index 4d4d1b476..86cc2e77f 100644 --- a/modules/postprocess/swinir_model.py +++ b/modules/postprocess/swinir_model.py @@ -125,6 +125,8 @@ def inference(img, model, tile, tile_overlap, window_size, scale): 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 shared.state.interrupted: + break for w_idx in w_idx_list: if shared.state.interrupted or shared.state.skipped: break