From 3aee33b54e53e032755edf253ba39005bae6e20e Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 17 Aug 2023 03:12:37 -0400 Subject: [PATCH] Add --disable-smart-memory for those that want the old behaviour. --- comfy/cli_args.py | 3 +++ comfy/model_management.py | 9 ++++++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/comfy/cli_args.py b/comfy/cli_args.py index ec7d34a..374dd2f 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -82,6 +82,9 @@ vram_group.add_argument("--novram", action="store_true", help="When lowvram isn' vram_group.add_argument("--cpu", action="store_true", help="To use the CPU for everything (slow).") +parser.add_argument("--disable-smart-memory", action="store_true", help="Force ComfyUI to agressively offload to regular ram instead of keeping models in vram when it can.") + + parser.add_argument("--dont-print-server", action="store_true", help="Don't print server output.") parser.add_argument("--quick-test-for-ci", action="store_true", help="Quick test for CI.") parser.add_argument("--windows-standalone-build", action="store_true", help="Windows standalone build: Enable convenient things that most people using the standalone windows build will probably enjoy (like auto opening the page on startup).") diff --git a/comfy/model_management.py b/comfy/model_management.py index 4c2441f..5c5d5ab 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -202,6 +202,10 @@ if cpu_state == CPUState.MPS: print(f"Set vram state to: {vram_state.name}") +DISABLE_SMART_MEMORY = args.disable_smart_memory + +if DISABLE_SMART_MEMORY: + print("Disabling smart memory management") def get_torch_device_name(device): if hasattr(device, 'type'): @@ -289,7 +293,10 @@ def unload_model_clones(model): def free_memory(memory_required, device, keep_loaded=[]): unloaded_model = False for i in range(len(current_loaded_models) -1, -1, -1): - current_free_mem = get_free_memory(device) + if DISABLE_SMART_MEMORY: + current_free_mem = 0 + else: + current_free_mem = get_free_memory(device) if current_free_mem > memory_required: break shift_model = current_loaded_models[i]