Make SwinIR interruptible (f9935258) · Commits · github_fork / Stable Diffusion Webui

extensions-builtin/SwinIR/scripts/swinir_model.py

+7 −1

Original line number	Diff line number	Diff line
		@@ -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:
		break

		for w_idx in w_idx_list:
		if state.interrupted:
		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)