298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-04-20 03:13:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update nodes_upscale_model.py

Browse Source 
Fix GPU utilization in upscale model node by keeping tensors on GPU. Added output_device parameter to tiled_scale function to prevent unnecessary CPU transfers, resulting in 2x faster processing. Commented out model CPU offloading to maintain GPU acceleration throughout the pipeline.
This commit is contained in:
yash-711 2025-03-11 18:16:11 +05:30 committed by GitHub
parent bc219a6487
commit 5cd7530637
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
comfy_extras/nodes_upscale_model.py
View File

@ -66,14 +66,15 @@ class ImageUpscaleWithModel:
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)
pbar = comfy.utils.ProgressBar(steps) 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) # KEY CHANGE: Pass device as output_device instead of default "cpu"
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, output_device=device, pbar=pbar)
oom = False oom = False
except model_management.OOM_EXCEPTION as e: except model_management.OOM_EXCEPTION as e:
tile //= 2 tile //= 2
if tile < 128: if tile < 128:
raise e raise e
upscale_model.to("cpu") # upscale_model.to("cpu") # Commented out to keep model on GPU because when processing batch images then model unnecessarily moves 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 (s,) return (s,)