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.)

Author: pollockjj

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 @@ def load_checkpoint(self, ckpt_name, unet_compute_device, unet_virtual_vram_gb,
             '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 override(self, *args, virtual_vram_gb=4.0, donor_device="cpu",
             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(self, *args, device=None, expert_mode_allocations="", use_other_vra
             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(self, *args, compute_device=None, virtual_vram_gb=4.0, donor_device
             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(self, *args, device=None, expert_mode_allocations="", use_other_vra
             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(self, *args, device=None, expert_mode_allocations="", use_other_vra
             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}