From 2331710c501381bb86a4a7e0573d53722df41d02 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sat, 23 Aug 2025 14:31:32 -0500 Subject: [PATCH] Enhance partially_load with fallback and reduced logging - Add force_patch_weights parameter to new_partially_load signature for better control - Implement check for _distorch_high_precision_loras with fallback to original loading behavior - Include cleanup for _distorch_block_assignments attribute - Comment out debug logging statements to minimize noise during execution --- distorch_2.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/distorch_2.py b/distorch_2.py index 929f594..ff80015 100644 --- a/distorch_2.py +++ b/distorch_2.py @@ -58,9 +58,19 @@ def register_patched_safetensor_modelpatcher(): if not hasattr(comfy.model_patcher.ModelPatcher, '_distorch_patched'): original_partially_load = comfy.model_patcher.ModelPatcher.partially_load - def new_partially_load(self, device_to, extra_memory=0, full_load=False, **kwargs): + def new_partially_load(self, device_to, extra_memory=0, full_load=False, force_patch_weights=False, **kwargs): """Override to use our static device assignments""" global safetensor_allocation_store + + if not hasattr(self.model, '_distorch_high_precision_loras'): + logger.info(f"[DEBUG_NEW_LOAD] high_precision_loras flag not retrieved from model. DisTorchV2 Loader not used. Reverting to normal loading behavior") + result = original_partially_load(self, device_to, extra_memory, force_patch_weights) + + # Clean up + if hasattr(self, '_distorch_block_assignments'): + del self._distorch_block_assignments + + return result # Check if we have a device allocation for this model debug_hash = create_safetensor_model_hash(self, "partial_load") @@ -85,15 +95,15 @@ def register_patched_safetensor_modelpatcher(): cast_weight = self.force_cast_weights if hasattr(m, "comfy_cast_weights"): - logging.info(f"Unpatching weight {weight_key} for Distorch2") + #logging.info(f"Unpatching weight {weight_key} for Distorch2") wipe_lowvram_weight(m) - logging.info(f"Adding {n} to 'load_completely' list") + #logging.info(f"Adding {n} to 'load_completely' list") mem_counter += module_mem load_completely.append((module_mem, n, m, params)) if cast_weight and hasattr(m, "comfy_cast_weights"): - logging.info(f"Setting cast weights for {weight_key}") + #logging.info(f"Setting cast weights for {weight_key}") m.prev_comfy_cast_weights = m.comfy_cast_weights m.comfy_cast_weights = True