wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-04-15 04:52:31 +08:00
Code Issues Actions 16 Packages Projects Releases Wiki Activity

Make ImageUpscaleWithModel node work with intermediate device and dtype.

Browse Source
This commit is contained in:
comfyanonymous 2026-04-10 21:29:11 -04:00
parent a134423890
commit d854a6c6ed
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
comfy_extras/nodes_upscale_model.py
View File

@ -6,6 +6,7 @@ import comfy.utils
import folder_paths
from typing_extensions import override
from comfy_api.latest import ComfyExtension, io
import comfy.model_management
try:
from spandrel_extra_arches import EXTRA_REGISTRY
@ -78,13 +79,15 @@ class ImageUpscaleWithModel(io.ComfyNode):
tile = 512
overlap = 32
output_device = comfy.model_management.intermediate_device()
oom = True
try:
while oom:
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)
pbar = comfy.utils.ProgressBar(steps)
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)
s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a.float()), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar, output_device=output_device)
oom = False
except Exception as e:
model_management.raise_non_oom(e)
@ -94,7 +97,7 @@ class ImageUpscaleWithModel(io.ComfyNode):
finally:
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).to(comfy.model_management.intermediate_dtype())
return io.NodeOutput(s)
upscale = execute # TODO: remove