diff --git a/ldm_patched/modules/args_parser.py b/ldm_patched/modules/args_parser.py index e5b84dc..272deb8 100644 --- a/ldm_patched/modules/args_parser.py +++ b/ldm_patched/modules/args_parser.py @@ -100,7 +100,7 @@ vram_group.add_argument("--always-high-vram", action="store_true") vram_group.add_argument("--always-normal-vram", action="store_true") vram_group.add_argument("--always-low-vram", action="store_true") vram_group.add_argument("--always-no-vram", action="store_true") -vram_group.add_argument("--always-cpu", action="store_true") +vram_group.add_argument("--always-cpu", type=int, nargs="?", metavar="CPU_NUM_THREADS", const=-1) parser.add_argument("--always-offload-from-vram", action="store_true") diff --git a/ldm_patched/modules/model_management.py b/ldm_patched/modules/model_management.py index 6f88579..840d79a 100644 --- a/ldm_patched/modules/model_management.py +++ b/ldm_patched/modules/model_management.py @@ -60,6 +60,9 @@ except: pass if args.always_cpu: + if args.always_cpu > 0: + torch.set_num_threads(args.always_cpu) + print(f"Running on {torch.get_num_threads()} CPU threads") cpu_state = CPUState.CPU def is_intel_xpu(): diff --git a/readme.md b/readme.md index fa7e829..18b48f3 100644 --- a/readme.md +++ b/readme.md @@ -370,7 +370,7 @@ entry_with_update.py [-h] [--listen [IP]] [--port PORT] [--attention-split | --attention-quad | --attention-pytorch] [--disable-xformers] [--always-gpu | --always-high-vram | --always-normal-vram | - --always-low-vram | --always-no-vram | --always-cpu] + --always-low-vram | --always-no-vram | --always-cpu [CPU_NUM_THREADS]] [--always-offload-from-vram] [--disable-server-log] [--debug-mode] [--is-windows-embedded-python] [--disable-server-info] [--share] [--preset PRESET]