From aa0a025df9ccc0d514a6b0ee28b450bc4cc29052 Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Mon, 10 Nov 2025 14:59:07 -0500
Subject: [PATCH] Try to disable pinned memory if wrong behavior detected.

Not sure if this is going to improve anything
---
 comfy/model_management.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/comfy/model_management.py b/comfy/model_management.py
index a4410f2ec..a51737db3 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -1128,6 +1128,16 @@ def pin_memory(tensor):
 
     return False
 
+def disable_pinned():
+    global MAX_PINNED_MEMORY
+    global TOTAL_PINNED_MEMORY
+    MAX_PINNED_MEMORY = -1
+    logging.warning("Bad behavior detected, disabling pinned memory")
+    for ptr in PINNED_MEMORY:  # TODO: check if this makes things worse
+        torch.cuda.cudart().cudaHostUnregister(ptr)
+    PINNED_MEMORY.clear()
+    TOTAL_PINNED_MEMORY = 0
+
 def unpin_memory(tensor):
     global TOTAL_PINNED_MEMORY
     if MAX_PINNED_MEMORY <= 0:
@@ -1142,6 +1152,7 @@ def unpin_memory(tensor):
     size_stored = PINNED_MEMORY.get(ptr, None)
     if size_stored is None:
         logging.warning("Tried to unpin tensor not pinned by ComfyUI")
+        disable_pinned()
         return False
 
     if size != size_stored: