hotfix: Corrects for corner case when DisTorch VirtualVRAM=0.0 GB (previous refactor shunted to standard loader. This replicates that required logic across all nodes using DisTorch2 for allocations. Next time I will wait for the final test work flow to finish VAE conversion (where is the only place I test this.)

checkpoint_multigpu.py

@@ -107,11 +107,13 @@ def patched_load_state_dict_guess_config(sd, output_vae=True, output_clip=True,
             multigpu_memory_log(f"unet:{config_hash[:8]}", "post-model")
 
             if distorch_config and 'unet_allocation' in distorch_config:
-                register_patched_safetensor_modelpatcher()
-                inner_model = model_patcher.model
-                inner_model._distorch_v2_meta = {"full_allocation": distorch_config['unet_allocation']}
-                logger.info(f"[CHECKPOINT_META] UNET inner_model id=0x{id(inner_model):x}")
-                model._distorch_high_precision_loras = distorch_config.get('high_precision_loras', True)
+                unet_alloc = distorch_config['unet_allocation']
+                if unet_alloc:
+                    register_patched_safetensor_modelpatcher()
+                    inner_model = model_patcher.model
+                    inner_model._distorch_v2_meta = {"full_allocation": unet_alloc}
+                    logger.info(f"[CHECKPOINT_META] UNET inner_model id=0x{id(inner_model):x}")
+                    model._distorch_high_precision_loras = distorch_config.get('high_precision_loras', True)
 
             model.load_model_weights(sd, diffusion_model_prefix)
             multigpu_memory_log(f"unet:{config_hash[:8]}", "post-weights")
@@ -141,10 +143,11 @@ def patched_load_state_dict_guess_config(sd, output_vae=True, output_clip=True,
                     clip = CLIP(clip_target, embedding_directory=embedding_directory, tokenizer_data=clip_sd, parameters=clip_params, model_options=te_model_options)
 
                     if distorch_config and 'clip_allocation' in distorch_config:
-                         if hasattr(clip, 'patcher'):
+                        clip_alloc = distorch_config['clip_allocation']
+                        if clip_alloc and hasattr(clip, 'patcher'):
                             register_patched_safetensor_modelpatcher()
                             inner_clip = clip.patcher.model
-                            inner_clip._distorch_v2_meta = {"full_allocation": distorch_config['clip_allocation']}
+                            inner_clip._distorch_v2_meta = {"full_allocation": clip_alloc}
                             logger.info(f"[CHECKPOINT_META] CLIP inner_model id=0x{id(inner_clip):x}")
                             clip.patcher.model._distorch_high_precision_loras = distorch_config.get('high_precision_loras', True)
 
@@ -257,10 +260,19 @@ class CheckpointLoaderAdvancedDisTorch2MultiGPU:
             'vae_device': vae_device
         }
 
-        unet_vram_str = f"{unet_compute_device};{unet_virtual_vram_gb};{unet_donor_device}"
-        unet_alloc = f"{unet_expert_mode_allocations}#{unet_vram_str}"
-        clip_vram_str = f"{clip_compute_device};{clip_virtual_vram_gb};{clip_donor_device}"
-        clip_alloc = f"{clip_expert_mode_allocations}#{clip_vram_str}"
+        unet_vram_str = ""
+        if unet_virtual_vram_gb > 0:
+            unet_vram_str = f"{unet_compute_device};{unet_virtual_vram_gb};{unet_donor_device}"
+        elif unet_expert_mode_allocations:
+            unet_vram_str = unet_compute_device
+        unet_alloc = f"{unet_expert_mode_allocations}#{unet_vram_str}" if unet_expert_mode_allocations or unet_vram_str else ""
+        
+        clip_vram_str = ""
+        if clip_virtual_vram_gb > 0:
+            clip_vram_str = f"{clip_compute_device};{clip_virtual_vram_gb};{clip_donor_device}"
+        elif clip_expert_mode_allocations:
+            clip_vram_str = clip_compute_device
+        clip_alloc = f"{clip_expert_mode_allocations}#{clip_vram_str}" if clip_expert_mode_allocations or clip_vram_str else ""
 
         checkpoint_distorch_config[config_hash] = {
             'unet_allocation': unet_alloc,

distorch_2.py

@@ -322,12 +322,16 @@ def analyze_safetensor_loading(model_patcher, allocations_string, is_clip=False)
     """
     DEVICE_RATIOS_DISTORCH = {}
     device_table = {}
-    distorch_alloc = allocations_string
+    distorch_alloc = ""
+    virtual_vram_str = ""
     virtual_vram_gb = 0.0
 
-    distorch_alloc, virtual_vram_str = allocations_string.split('#')
+    if '#' in allocations_string:
+        distorch_alloc, virtual_vram_str = allocations_string.split('#', 1)
+    else:
+        distorch_alloc = allocations_string
 
-    compute_device = virtual_vram_str.split(';')[0]
+    compute_device = virtual_vram_str.split(';')[0] if virtual_vram_str else "cuda:0"
     logger.debug(f"[MultiGPU DisTorch V2] Compute Device: {compute_device}")
 
     if not distorch_alloc:

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-multigpu"
 description = "Provides a suite of custom nodes to manage multiple GPUs for ComfyUI, including advanced model offloading for both GGUF and Safetensor formats with DisTorch, and bespoke MultiGPU support for WanVideoWrapper and other custom nodes."
-version = "2.5.6"
+version = "2.5.7"
 license = {file = "LICENSE"}
 
 [project.urls]

wrappers.py

@@ -112,7 +112,7 @@ def _create_distorch_safetensor_v2_override(cls, device_param_name, device_sette
             elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'):
                 model_to_check = out[0].patcher
 
-            if model_to_check:
+            if model_to_check and full_allocation:
                 inner_model = model_to_check.model
                 inner_model._distorch_v2_meta = {"full_allocation": full_allocation}
 
@@ -217,7 +217,7 @@ def override_class_with_distorch_gguf(cls):
             elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'):
                 model_to_check = out[0].patcher
             
-            if model_to_check:
+            if model_to_check and full_allocation:
                 inner_model = model_to_check.model
                 inner_model._distorch_v2_meta = {"full_allocation": full_allocation}
 
@@ -279,7 +279,7 @@ def override_class_with_distorch_gguf_v2(cls):
             elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'):
                 model_to_check = out[0].patcher
             
-            if model_to_check:
+            if model_to_check and full_allocation:
                 inner_model = model_to_check.model
                 inner_model._distorch_v2_meta = {"full_allocation": full_allocation}
 
@@ -343,7 +343,7 @@ def override_class_with_distorch_clip(cls):
             elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'):
                 model_to_check = out[0].patcher
             
-            if model_to_check:
+            if model_to_check and full_allocation:
                 inner_model = model_to_check.model
                 inner_model._distorch_v2_meta = {"full_allocation": full_allocation}
 
@@ -407,7 +407,7 @@ def override_class_with_distorch_clip_no_device(cls):
             elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'):
                 model_to_check = out[0].patcher
             
-            if model_to_check:
+            if model_to_check and full_allocation:
                 inner_model = model_to_check.model
                 inner_model._distorch_v2_meta = {"full_allocation": full_allocation}