From a24f0a6e87cbb561ac355eac5b9122f1693f100d Mon Sep 17 00:00:00 2001 From: John Pollock Date: Mon, 13 Oct 2025 22:08:05 -0500 Subject: [PATCH] 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 | 34 +++++++++++++++++++++++----------- distorch_2.py | 10 +++++++--- pyproject.toml | 2 +- wrappers.py | 10 +++++----- 4 files changed, 36 insertions(+), 20 deletions(-) diff --git a/checkpoint_multigpu.py b/checkpoint_multigpu.py index f4d068a..6b5f7ef 100644 --- a/checkpoint_multigpu.py +++ b/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, diff --git a/distorch_2.py b/distorch_2.py index b796ba6..f1bf931 100644 --- a/distorch_2.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index cd9d304..c8df6f2 100644 --- a/pyproject.toml +++ b/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] diff --git a/wrappers.py b/wrappers.py index 3663780..30e85c6 100644 --- a/wrappers.py +++ b/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}