interruptible upscale ops

Signed-off-by: Vladimir Mandic <mandic00@live.com>

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

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

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

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