add upscaler logging

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

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)

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
 

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

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 []

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)