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

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

View File

@ -78,6 +78,7 @@ class ImageUpscaleWithModel(io.ComfyNode):
overlap = 32 overlap = 32
oom = True oom = True
try:
while oom: while oom:
try: try:
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) 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)
@ -88,8 +89,9 @@ class ImageUpscaleWithModel(io.ComfyNode):
tile //= 2 tile //= 2
if tile < 128: if tile < 128:
raise e 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)