From 543a0dc1eb432d99a2481a85a2e9491ca31db0c4 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sat, 23 Aug 2025 10:06:54 -0500 Subject: [PATCH] patching logic from model_patcher `load` --- distorch_2.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/distorch_2.py b/distorch_2.py index 3eb1614..929f594 100644 --- a/distorch_2.py +++ b/distorch_2.py @@ -87,6 +87,25 @@ def register_patched_safetensor_modelpatcher(): if hasattr(m, "comfy_cast_weights"): logging.info(f"Unpatching weight {weight_key} for Distorch2") wipe_lowvram_weight(m) + + 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}") + m.prev_comfy_cast_weights = m.comfy_cast_weights + m.comfy_cast_weights = True + + if weight_key in self.weight_wrapper_patches: + logging.info(f"Patching weight wrapper {m} for Distorch2") + m.weight_function.extend(self.weight_wrapper_patches[weight_key]) + + if bias_key in self.weight_wrapper_patches: + logging.info(f"Patching bias wrapper {bias_key} for Distorch2") + m.bias_function.extend(self.weight_wrapper_patches[bias_key]) + + mem_counter += move_weight_functions(m, device_to)