diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 792148f0a..ee9e1ce9f 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -172,6 +172,7 @@ vram_group.add_argument("--cpu", action="store_true", help="To use the CPU for e parser.add_argument("--reserve-vram", type=float, default=None, help="Set the amount of vram in GB you want to reserve for use by your OS/other software. By default some amount is reserved depending on your OS.") parser.add_argument("--vram-headroom", type=float, default=0, help="Set the amount of vram in GB for DynamicVRAM to maintain as extra headroom above default. ComfyUI will try and keep this much VRAM completely free and unused, even counting VRAM from other apps.") +parser.add_argument("--disable-nvml-pressure", action="store_true", help="Use CUDA instead of NVML for DynamicVRAM memory pressure.") parser.add_argument("--async-offload", nargs='?', const=2, type=int, default=None, metavar="NUM_STREAMS", help="Use async weight offloading. An optional argument controls the amount of offload streams. Default is 2. Enabled by default on Nvidia.") parser.add_argument("--disable-async-offload", action="store_true", help="Disable async weight offloading.") diff --git a/main.py b/main.py index 9c318fafe..361b1fc89 100644 --- a/main.py +++ b/main.py @@ -58,11 +58,16 @@ if __name__ == "__main__" and args.debug_hang: import comfy_aimdo.control if enables_dynamic_vram(): + simple_vram_headroom = None if args.reserve_vram is None else int(args.reserve_vram * 1024 ** 3) try: - comfy_aimdo.control.init(simple_vram_headroom=None if args.reserve_vram is None else int(args.reserve_vram * 1024 ** 3)) + comfy_aimdo.control.init(simple_vram_headroom=simple_vram_headroom, nvml_pressure=not args.disable_nvml_pressure) except TypeError: - # comfy-aimdo 0.4.9 protocol. - comfy_aimdo.control.init() + # comfy-aimdo 0.4.10 protocol. + try: + comfy_aimdo.control.init(simple_vram_headroom=simple_vram_headroom) + except TypeError: + # comfy-aimdo 0.4.9 protocol. + comfy_aimdo.control.init() if os.name == "nt": os.environ['MIMALLOC_PURGE_DELAY'] = '0' diff --git a/requirements.txt b/requirements.txt index aa3fbba8e..c1eeb9752 100644 --- a/requirements.txt +++ b/requirements.txt @@ -23,7 +23,7 @@ SQLAlchemy>=2.0.0 filelock av>=16.0.0 comfy-kitchen==0.2.26 -comfy-aimdo==0.4.10 +comfy-aimdo==0.4.11 requests simpleeval>=1.0.0 blake3