aboutsummaryrefslogtreecommitdiff
path: root/extensions-builtin/SwinIR/scripts/swinir_model.py
diff options
context:
space:
mode:
authorunknown <mcgpapu@gmail.com>2023-01-28 03:40:51 -0600
committerunknown <mcgpapu@gmail.com>2023-01-28 03:40:51 -0600
commite79b7db4b47a33889551b9266ee3277879d4f560 (patch)
tree1c1944204e58e254bfea22ae44edccdbb54e6b3c /extensions-builtin/SwinIR/scripts/swinir_model.py
parentb921a52071cf2a5e551c31a6073af6eaebbf7847 (diff)
parente8a41df49fadd2cf9f23b1f02d75a4947bec5646 (diff)
Merge branch 'master' of github.com:AUTOMATIC1111/stable-diffusion-webui into gamepad
Diffstat (limited to 'extensions-builtin/SwinIR/scripts/swinir_model.py')
-rw-r--r--extensions-builtin/SwinIR/scripts/swinir_model.py8
1 files changed, 7 insertions, 1 deletions
diff --git a/extensions-builtin/SwinIR/scripts/swinir_model.py b/extensions-builtin/SwinIR/scripts/swinir_model.py
index 9a74b253..e8783bca 100644
--- a/extensions-builtin/SwinIR/scripts/swinir_model.py
+++ b/extensions-builtin/SwinIR/scripts/swinir_model.py
@@ -8,7 +8,7 @@ from basicsr.utils.download_util import load_file_from_url
from tqdm import tqdm
from modules import modelloader, devices, script_callbacks, shared
-from modules.shared import cmd_opts, opts
+from modules.shared import cmd_opts, opts, state
from swinir_model_arch import SwinIR as net
from swinir_model_arch_v2 import Swin2SR as net2
from modules.upscaler import Upscaler, UpscalerData
@@ -145,7 +145,13 @@ def inference(img, model, tile, tile_overlap, window_size, scale):
with tqdm(total=len(h_idx_list) * len(w_idx_list), desc="SwinIR tiles") as pbar:
for h_idx in h_idx_list:
+ if state.interrupted or state.skipped:
+ break
+
for w_idx in w_idx_list:
+ if state.interrupted or state.skipped:
+ break
+
in_patch = img[..., h_idx: h_idx + tile, w_idx: w_idx + tile]
out_patch = model(in_patch)
out_patch_mask = torch.ones_like(out_patch)