From 439c178c2cbb5cb07445314ab99384b86a3c3c5e Mon Sep 17 00:00:00 2001
From: Rattus <rattus128@gmail.com>
Date: Tue, 13 Jan 2026 15:16:41 +1000
Subject: [PATCH] pinned_memory: add python

Add a python for managing pinned memory of the weight/bias module level.
This allocates, pins and attached a tensor to a module for the pin for this
module. It does not set the weight, just allocates a singular ram buffer
for population and bulk DMA transfer.
---
 comfy/model_management.py |  2 +-
 comfy/pinned_memory.py    | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 comfy/pinned_memory.py

diff --git a/comfy/model_management.py b/comfy/model_management.py
index 790236ede..21761d971 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -1185,7 +1185,7 @@ if not args.disable_pinned_memory:
             MAX_PINNED_MEMORY = get_total_memory(torch.device("cpu")) * 0.95
         logging.info("Enabled pinned memory {}".format(MAX_PINNED_MEMORY // (1024 * 1024)))
 
-PINNING_ALLOWED_TYPES = set(["Parameter", "QuantizedTensor"])
+PINNING_ALLOWED_TYPES = set(["Tensor", "Parameter", "QuantizedTensor"])
 
 def discard_cuda_async_error():
     try:
diff --git a/comfy/pinned_memory.py b/comfy/pinned_memory.py
new file mode 100644
index 000000000..be303b4f1
--- /dev/null
+++ b/comfy/pinned_memory.py
@@ -0,0 +1,34 @@
+import torch
+import logging
+import comfy.model_management
+import comfy.memory_management
+
+from comfy.cli_args import args
+
+def get_pin(module):
+    return getattr(module, "_pin", None)
+
+def pin_memory(module):
+    if module.pin_failed or args.disable_pinned_memory or get_pin(module) is not None:
+        return
+    #FIXME: This is a RAM cache trigger event
+    params = [ module.weight, module.bias ]
+    size = comfy.memory_management.vram_aligned_size(params)
+    try:
+        pin = torch.empty((size,), dtype=torch.uint8)
+        if comfy.model_management.pin_memory(pin):
+            module._pin = pin
+        else:
+            module.pin_failed = True
+            return False
+    except:
+        module.pin_failed = True
+        return False
+    return True
+
+def unpin_memory(module):
+    if get_pin(module) is None:
+        return 0
+    size = module._pin.numel() * module._pin.element_size()
+    comfy.model_management.unpin_memory(module._pin)
+    return size