From c8923a4cd7b7ddea104ea1939b4d715bf1849ccb Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Tue, 3 Sep 2024 16:37:35 +0200 Subject: [PATCH] mmap unload fix for lowvram --- nodes.py | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/nodes.py b/nodes.py index 061e2b9..ef505cf 100644 --- a/nodes.py +++ b/nodes.py @@ -180,9 +180,32 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): device_to = device_to if comfy.model_management.DISABLE_SMART_MEMORY else None return super().unpatch_model(device_to=device_to, unpatch_weights=unpatch_weights) + mmap_released = False def load(self, *args, force_patch_weights=False, **kwargs): # always call `patch_weight_to_device` even for lowvram - return super().load(*args, force_patch_weights=True, **kwargs) + super().load(*args, force_patch_weights=True, **kwargs) + + # make sure nothing stays linked to mmap after first load + if not self.mmap_released: + linked = [] + if kwargs.get("lowvram_model_memory", 0) > 0: + for n, m in self.model.named_modules(): + if hasattr(m, "weight"): + device = getattr(m.weight, "device", None) + if device == self.offload_device: + linked.append((n, m)) + continue + if hasattr(m, "bias"): + device = getattr(m.bias, "device", None) + if device == self.offload_device: + linked.append((n, m)) + continue + if linked: + print(f"Attempting to release mmap ({len(linked)})") + for n, m in linked: + # TODO: possible to OOM, find better way to detach + m.to(self.load_device).to(self.offload_device) + self.mmap_released = True def clone(self, *args, **kwargs): n = GGUFModelPatcher(self.model, self.load_device, self.offload_device, self.size, weight_inplace_update=self.weight_inplace_update)