xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-08-14 11:50:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(upscale): Add parallel support for multiple GPU

Browse Source
This commit is contained in:
MagicBear 2025-07-04 00:31:15 +08:00
parent 2f386dcba1
commit 71b505f902
1 changed files with 24 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
comfy_extras/nodes_upscale_model.py
View File

@ -4,6 +4,7 @@ from comfy import model_management
import torch
import comfy.utils
import folder_paths
from torch.nn import DataParallel
try:
from spandrel_extra_arches import EXTRA_REGISTRY
@ -39,16 +40,29 @@ class UpscaleModelLoader:
class ImageUpscaleWithModel:
@classmethod
def INPUT_TYPES(s):
return {"required": { "upscale_model": ("UPSCALE_MODEL",),
inputs = {"required": {
"upscale_model": ("UPSCALE_MODEL",),
"image": ("IMAGE",),
}}
},
"optional": {}}
for i in range(torch.cuda.device_count()):
inputs["optional"]["cuda_%d" % i] = ("BOOLEAN", {"default": True, "tooltip": "Use device %s" % torch.cuda.get_device_name(i)})
return inputs
RETURN_TYPES = ("IMAGE",)
FUNCTION = "upscale"
CATEGORY = "image/upscaling"
def upscale(self, upscale_model, image):
def upscale(self, upscale_model, image, **kwargs):
device = model_management.get_torch_device()
device_ids = []
for k, v in kwargs.items():
if k.startswith("cuda_") and v:
device_ids.append(int(k[5:]))
if kwargs.get("cuda_0"):
device = "cuda:0"
elif len(device_ids) > 0:
device = "cuda:%d" % device_ids[0]
memory_required = model_management.module_size(upscale_model.model)
memory_required += (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0 #The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate
@ -58,6 +72,11 @@ class ImageUpscaleWithModel:
upscale_model.to(device)
in_img = image.movedim(-1,-3).to(device)
if len(device_ids) > 1:
parallel_model = DataParallel(upscale_model.model, device_ids=device_ids)
else:
parallel_model = upscale_model
tile = 512
overlap = 32
@ -66,7 +85,7 @@ class ImageUpscaleWithModel:
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, parallel_model, tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar, output_device=device)
oom = False
except model_management.OOM_EXCEPTION as e:
tile //= 2
@ -75,6 +94,7 @@ class ImageUpscaleWithModel:
upscale_model.to("cpu")
s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
s.to("cpu")
return (s,)
NODE_CLASS_MAPPINGS = {