wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-03-10 11:47:34 +08:00
Code Issues Actions 14 Packages Projects Releases Wiki Activity

Fix case where upscale model wouldn't be moved to cpu.

Browse Source
This commit is contained in:
comfyanonymous 2026-01-04 19:07:25 -05:00 committed by GitHub
parent acbf08cd60
commit 54b1ea7ad7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 13 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
comfy_extras/nodes_upscale_model.py
View File

@ -78,18 +78,20 @@ class ImageUpscaleWithModel(io.ComfyNode):
overlap = 32 overlap = 32
oom = True oom = True
while oom: try:
try: while oom:
steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap) try:
pbar = comfy.utils.ProgressBar(steps) steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap)
s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar) pbar = comfy.utils.ProgressBar(steps)
oom = False s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar)
except model_management.OOM_EXCEPTION as e: oom = False
tile //= 2 except model_management.OOM_EXCEPTION as e:
if tile < 128: tile //= 2
raise e if tile < 128:
raise e
finally:
upscale_model.to("cpu")
upscale_model.to("cpu")
s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0) s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
return io.NodeOutput(s) return io.NodeOutput(s)