diff --git a/__init__.py b/__init__.py index da7e686..c4dee25 100644 --- a/__init__.py +++ b/__init__.py @@ -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) diff --git a/nodes.py b/nodes.py index 822d059..9c8afbd 100644 --- a/nodes.py +++ b/nodes.py @@ -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 diff --git a/wrappers.py b/wrappers.py index 9abf9bf..42ab83c 100644 --- a/wrappers.py +++ b/wrappers.py @@ -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)"""