nodes/pollockjj-ComfyUI-MultiGPU
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add UNet offload device support, and enhance Florence2 model loading with safetensors conversion option.

Browse Source
This commit is contained in:
John Pollock
2025-10-12 04:35:43 -05:00
parent 0d6dfc7524
commit 161174266c
3 changed files with 68 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+26 -4
View File
@@ -23,7 +23,7 @@ from .model_management_mgpu import (
)
WEB_DIRECTORY = "./web"
MGPU_MM_LOG = True
MGPU_MM_LOG = False
DEBUG_LOG = False
logger = logging.getLogger("MultiGPU")
@@ -148,6 +148,7 @@ def check_module_exists(module_path):
current_device = mm.get_torch_device()
current_text_encoder_device = mm.text_encoder_device()
current_unet_offload_device = mm.unet_offload_device()
def set_current_device(device):
"""Set the current device context for MultiGPU operations."""
@@ -161,6 +162,12 @@ def set_current_text_encoder_device(device):
current_text_encoder_device = device
logger.debug(f"[MultiGPU Initialization] current_text_encoder_device set to: {device}")
def set_current_unet_offload_device(device):
"""Set the current UNet offload device context."""
global current_unet_offload_device
current_unet_offload_device = device
logger.debug(f"[MultiGPU Initialization] current_unet_offload_device set to: {device}")
def get_torch_device_patched():
"""Return MultiGPU-aware device selection for patched mm.get_torch_device."""
device = None
@@ -183,11 +190,25 @@ def text_encoder_device_patched():
logger.info(f"[MultiGPU Core Patching] text_encoder_device_patched returning device: {device} (current_text_encoder_device={current_text_encoder_device})")
return device
logger.info(f"[MultiGPU Core Patching] Patching mm.get_torch_device and mm.text_encoder_device")
def unet_offload_device_patched():
"""Return MultiGPU-aware UNet offload device for patched mm.unet_offload_device."""
device = None
if (not is_accelerator_available() or mm.cpu_state == mm.CPUState.CPU or "cpu" in str(current_unet_offload_device).lower()):
device = torch.device("cpu")
else:
devs = set(get_device_list())
device = torch.device(current_unet_offload_device) if str(current_unet_offload_device) in devs else torch.device("cpu")
logger.debug(f"[MultiGPU Core Patching] unet_offload_device_patched returning device: {device} (current_unet_offload_device={current_unet_offload_device})")
return device
logger.info(f"[MultiGPU Core Patching] Patching mm.get_torch_device, mm.text_encoder_device, mm.unet_offload_device")
logger.info(f"[MultiGPU DEBUG] Initial current_device: {current_device}")
logger.info(f"[MultiGPU DEBUG] Initial current_text_encoder_device: {current_text_encoder_device}")
logger.info(f"[MultiGPU DEBUG] Initial current_unet_offload_device: {current_unet_offload_device}")
mm.get_torch_device = get_torch_device_patched
mm.text_encoder_device = text_encoder_device_patched
mm.unet_offload_device = unet_offload_device_patched
from .nodes import (
UnetLoaderGGUF,
@@ -235,6 +256,7 @@ from .wanvideo import (
from .wrappers import (
override_class,
override_class_offload,
override_class_clip,
override_class_clip_no_device,
override_class_with_distorch_gguf,
@@ -319,8 +341,8 @@ ltx_nodes = {"LTXVLoaderMultiGPU": override_class(LTXVLoader)}
register_and_count(["ComfyUI-LTXVideo", "comfyui-ltxvideo"], ltx_nodes)
florence_nodes = {
"Florence2ModelLoaderMultiGPU": override_class(Florence2ModelLoader),
"DownloadAndLoadFlorence2ModelMultiGPU": override_class(DownloadAndLoadFlorence2Model)
"Florence2ModelLoaderMultiGPU": override_class_offload(Florence2ModelLoader),
"DownloadAndLoadFlorence2ModelMultiGPU": override_class_offload(DownloadAndLoadFlorence2Model)
}
register_and_count(["ComfyUI-Florence2", "comfyui-florence2"], florence_nodes)
nodes.py
+12 -7
View File
@@ -176,10 +176,12 @@ class LTXVLoader:
return original_loader._load_vae(weights, config=None)
class Florence2ModelLoader:
@classmethod
def INPUT_TYPES(s):
all_llm_paths = folder_paths.get_folder_paths("LLM")
s.model_paths = create_path_dict(all_llm_paths, lambda x: x.is_dir())
return {"required": {
"model": ([item.name for item in Path(folder_paths.models_dir, "LLM").iterdir() if item.is_dir()], {"tooltip": "models are expected to be in Comfyui/models/LLM folder"}),
"model": ([*s.model_paths], {"tooltip": "models are expected to be in Comfyui/models/LLM folder"}),
"precision": (['fp16','bf16','fp32'],),
"attention": (
[ 'flash_attention_2', 'sdpa', 'eager'],
@@ -189,6 +191,7 @@ class Florence2ModelLoader:
},
"optional": {
"lora": ("PEFTLORA",),
"convert_to_safetensors": ("BOOLEAN", {"default": False, "tooltip": "Some of the older model weights are not saved in .safetensors format, which seem to cause longer loading times, this option converts the .bin weights to .safetensors"}),
}
}
@@ -197,10 +200,10 @@ class Florence2ModelLoader:
FUNCTION = "loadmodel"
CATEGORY = "Florence2"
def loadmodel(self, model, precision, attention, lora=None):
def loadmodel(self, model, precision, attention, lora=None, convert_to_safetensors=False):
"""Load Florence2 vision model with specified precision and attention mode."""
original_loader = NODE_CLASS_MAPPINGS["Florence2ModelLoader"]()
return original_loader.loadmodel(model, precision, attention, lora)
return original_loader.loadmodel(model, precision, attention, lora, convert_to_safetensors)
class DownloadAndLoadFlorence2Model:
@classmethod
@@ -220,7 +223,8 @@ class DownloadAndLoadFlorence2Model:
'MiaoshouAI/Florence-2-base-PromptGen-v1.5',
'MiaoshouAI/Florence-2-large-PromptGen-v1.5',
'MiaoshouAI/Florence-2-base-PromptGen-v2.0',
'MiaoshouAI/Florence-2-large-PromptGen-v2.0'
'MiaoshouAI/Florence-2-large-PromptGen-v2.0',
'PJMixers-Images/Florence-2-base-Castollux-v0.5'
],
{
"default": 'microsoft/Florence-2-base'
@@ -237,6 +241,7 @@ class DownloadAndLoadFlorence2Model:
},
"optional": {
"lora": ("PEFTLORA",),
"convert_to_safetensors": ("BOOLEAN", {"default": False, "tooltip": "Some of the older model weights are not saved in .safetensors format, which seem to cause longer loading times, this option converts the .bin weights to .safetensors"}),
}
}
@@ -245,10 +250,10 @@ class DownloadAndLoadFlorence2Model:
FUNCTION = "loadmodel"
CATEGORY = "Florence2"
def loadmodel(self, model, precision, attention, lora=None):
def loadmodel(self, model, precision, attention, lora=None, convert_to_safetensors=False):
"""Download and load Florence2 model from HuggingFace."""
original_loader = NODE_CLASS_MAPPINGS["DownloadAndLoadFlorence2Model"]()
return original_loader.loadmodel(model, precision, attention, lora)
return original_loader.loadmodel(model, precision, attention, lora, convert_to_safetensors)
class CheckpointLoaderNF4:
@classmethod
wrappers.py
+30
View File
@@ -484,6 +484,36 @@ def override_class(cls):
return NodeOverride
def override_class_offload(cls):
"""Standard MultiGPU device override for UNet/VAE models"""
from . import set_current_device, set_current_unet_offload_device
class NodeOverride(cls):
@classmethod
def INPUT_TYPES(s):
inputs = copy.deepcopy(cls.INPUT_TYPES())
devices = get_device_list()
default_device = devices[1] if len(devices) > 1 else devices[0]
inputs["optional"] = inputs.get("optional", {})
inputs["optional"]["device"] = (devices, {"default": default_device})
inputs["optional"]["offload_device"] = (devices, {"default": "cpu"})
return inputs
CATEGORY = "multigpu"
FUNCTION = "override"
def override(self, *args, device=None, offload_device=None, **kwargs):
if device is not None:
set_current_device(device)
if offload_device is not None:
set_current_unet_offload_device(offload_device)
fn = getattr(super(), cls.FUNCTION)
out = fn(*args, **kwargs)
return out
return NodeOverride
def override_class_clip(cls):
"""Standard MultiGPU device override for CLIP models (with device kwarg workaround)"""