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

Author: pollockjj

__init__.py

@@ -23,7 +23,7 @@
 )
 
 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 @@ def text_encoder_device_patched():
 
 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 @@ def register_and_count(module_names, node_map):
 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

@@ -176,10 +176,12 @@ def _load_vae(self, weights, config=None):
         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 @@ def INPUT_TYPES(s):
             },
             "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 @@ def INPUT_TYPES(s):
     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 @@ def INPUT_TYPES(s):
                     '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 @@ def INPUT_TYPES(s):
             },
             "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 @@ def INPUT_TYPES(s):
     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

@@ -484,6 +484,36 @@ def override(self, *args, device=None, **kwargs):
 
     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)"""