From 389c3346317b0128b47f9ca1dfe8f2f68884f76a Mon Sep 17 00:00:00 2001
From: Rattus <rattus128@gmail.com>
Date: Tue, 13 Jan 2026 15:54:49 +1000
Subject: [PATCH] main: Go live with --fast dynamic_vram

Add the optional command line switch --fast dynamic_vram.

This is mutually exclusing --high-vram and --gpu-only which contradict
aimdos underlying feature.

Add appropriate installation warning and a startup message, match the
comfy debug level inconfiguring aimdo.

Add comfy-aimdo pip requirement. This will safely stub to a nop for
unsupported platforms.
---
 comfy/cli_args.py |  4 ++++
 cuda_malloc.py    |  7 ++++++-
 main.py           | 35 ++++++++++++++++++++++++++++++++++-
 requirements.txt  |  1 +
 4 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/comfy/cli_args.py b/comfy/cli_args.py
index 1716c3de7..63daca861 100644
--- a/comfy/cli_args.py
+++ b/comfy/cli_args.py
@@ -159,6 +159,7 @@ class PerformanceFeature(enum.Enum):
     Fp8MatrixMultiplication = "fp8_matrix_mult"
     CublasOps = "cublas_ops"
     AutoTune = "autotune"
+    DynamicVRAM = "dynamic_vram"
 
 parser.add_argument("--fast", nargs="*", type=PerformanceFeature, help="Enable some untested and potentially quality deteriorating optimizations. This is used to test new features so using it might crash your comfyui. --fast with no arguments enables everything. You can pass a list specific optimizations if you only want to enable specific ones. Current valid optimizations: {}".format(" ".join(map(lambda c: c.value, PerformanceFeature))))
 
@@ -257,3 +258,6 @@ elif args.fast == []:
 # '--fast' is provided with a list of performance features, use that list
 else:
     args.fast = set(args.fast)
+
+def enables_dynamic_vram():
+    return PerformanceFeature.DynamicVRAM in args.fast and not args.highvram and not args.gpu_only
diff --git a/cuda_malloc.py b/cuda_malloc.py
index 00ee7b633..3c7c8593e 100644
--- a/cuda_malloc.py
+++ b/cuda_malloc.py
@@ -1,8 +1,10 @@
 import os
 import importlib.util
-from comfy.cli_args import args, PerformanceFeature
+from comfy.cli_args import args, PerformanceFeature, enables_dynamic_vram
 import subprocess
 
+import comfy_aimdo.control
+
 #Can't use pytorch to get the GPU names because the cuda malloc has to be set before the first import.
 def get_gpu_names():
     if os.name == 'nt':
@@ -85,6 +87,9 @@ if not args.cuda_malloc:
     except:
         pass
 
+if enables_dynamic_vram() and comfy_aimdo.control.lib is not None:
+    args.cuda_malloc = False
+    os.environ['PYTORCH_CUDA_ALLOC_CONF'] = ""
 
 if args.disable_cuda_malloc:
     args.cuda_malloc = False
diff --git a/main.py b/main.py
index 37b06c1fa..52f11bfff 100644
--- a/main.py
+++ b/main.py
@@ -5,7 +5,7 @@ import os
 import importlib.util
 import folder_paths
 import time
-from comfy.cli_args import args
+from comfy.cli_args import args, enables_dynamic_vram
 from app.logger import setup_logger
 from app.assets.scanner import seed_assets
 import itertools
@@ -173,6 +173,30 @@ import gc
 if 'torch' in sys.modules:
     logging.warning("WARNING: Potential Error in code: Torch already imported, torch should never be imported before this point.")
 
+
+has_aimdo = False
+
+import comfy_aimdo.control
+
+if comfy_aimdo.control.lib is not None:
+    if args.verbose == 'DEBUG':
+        comfy_aimdo.control.set_log_debug()
+    elif args.verbose == 'CRITICAL':
+        comfy_aimdo.control.set_log_critical()
+    elif args.verbose == 'ERROR':
+        comfy_aimdo.control.set_log_error()
+    elif args.verbose == 'WARNING':
+        comfy_aimdo.control.set_log_warning()
+    else: #INFO
+        comfy_aimdo.control.set_log_info()
+
+    if enables_dynamic_vram():
+        logging.info("DynamicVRAM support detected and enabled")
+        has_aimdo = True
+else:
+    if enables_dynamic_vram():
+        logging.info("No native comfy-aimdo install detected. Falling back to legacy ModelPatcher. VRAM estimates may be unreliable especially on Windows")
+
 import comfy.utils
 
 import execution
@@ -184,6 +208,15 @@ import comfyui_version
 import app.logger
 import hook_breaker_ac10a0
 
+import comfy.memory_management
+import comfy.model_patcher
+
+if has_aimdo:
+    comfy.model_patcher.CoreModelPatcher = comfy.model_patcher.ModelPatcherDynamic
+    comfy_aimdo.control.init_vram_guard(comfy.model_management.get_torch_device().index)
+else:
+    comfy.memory_management.aimdo_allocator = None
+
 def cuda_malloc_warning():
     device = comfy.model_management.get_torch_device()
     device_name = comfy.model_management.get_torch_device_name(device)
diff --git a/requirements.txt b/requirements.txt
index 8650d28ec..950f51c15 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -22,6 +22,7 @@ alembic
 SQLAlchemy
 av>=14.2.0
 comfy-kitchen>=0.2.6
+comfy-aimdo>=0.1.0
 
 #non essential dependencies:
 kornia>=0.7.1