浏览代码

interrupt upscale

Andray 1 年之前
父节点
当前提交
c321680b3d
共有 2 个文件被更改,包括 5 次插入0 次删除
  1. 3 0
      modules/upscaler.py
  2. 2 0
      modules/upscaler_utils.py

+ 3 - 0
modules/upscaler.py

@@ -60,6 +60,9 @@ class Upscaler:
             if img.width >= dest_w and img.height >= dest_h:
             if img.width >= dest_w and img.height >= dest_h:
                 break
                 break
 
 
+            if shared.state.interrupted:
+                break
+
             shape = (img.width, img.height)
             shape = (img.width, img.height)
 
 
             img = self.do_upscale(img, selected_model)
             img = self.do_upscale(img, selected_model)

+ 2 - 0
modules/upscaler_utils.py

@@ -69,6 +69,8 @@ def upscale_with_model(
         for y, h, row in grid.tiles:
         for y, h, row in grid.tiles:
             newrow = []
             newrow = []
             for x, w, tile in row:
             for x, w, tile in row:
+                if shared.state.interrupted:
+                    return img
                 output = upscale_pil_patch(model, tile)
                 output = upscale_pil_patch(model, tile)
                 scale_factor = output.width // tile.width
                 scale_factor = output.width // tile.width
                 newrow.append([x * scale_factor, w * scale_factor, output])
                 newrow.append([x * scale_factor, w * scale_factor, output])