diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 7234a7ba0..930205be1 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -131,6 +131,13 @@ parser.add_argument("--reserve-vram", type=float, default=None, help="Set the am parser.add_argument("--async-offload", action="store_true", help="Use async weight offloading.") +class CpuBf16Mode(enum.Enum): + Auto = "auto" + Yes = "yes" + No = "no" + +parser.add_argument("--use-cpu-bf16", type=CpuBf16Mode, default=CpuBf16Mode.Auto, help="When CPU mode is enabled use bf16 instructions to improve performance.", action=EnumAction) + parser.add_argument("--default-hashing-function", type=str, choices=['md5', 'sha1', 'sha256', 'sha512'], default='sha256', help="Allows you to choose the hash function to use for duplicate filename / contents comparison. Default is sha256.") 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.") diff --git a/comfy/model_management.py b/comfy/model_management.py index 816caf18f..76b34b883 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -19,7 +19,7 @@ import psutil import logging from enum import Enum -from comfy.cli_args import args, PerformanceFeature +from comfy.cli_args import args, PerformanceFeature, CpuBf16Mode import torch import sys import platform @@ -151,6 +151,21 @@ def is_mlu(): return True return False +def use_cpu_bf16(): + if args.use_cpu_bf16 == CpuBf16Mode.No or cpu_state != CPUState.CPU: + return False + if args.use_cpu_bf16 == CpuBf16Mode.Yes: + return True + + try: + from cpuinfo import get_cpu_info + except: + logging.warning('py-cpuinfo is not installed, rerun "pip install -r requirements.txt"') + return False + + cpu_info = get_cpu_info() + return 'avx512_bf16' in cpu_info['flags'] + def get_torch_device(): global directml_enabled global cpu_state @@ -1218,8 +1233,8 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma def should_use_bf16(device=None, model_params=0, prioritize_performance=True, manual_cast=False): if device is not None: - if is_device_cpu(device): #TODO ? bf16 works on CPU but is extremely slow - return False + if is_device_cpu(device): + return use_cpu_bf16() if FORCE_FP32: return False @@ -1233,7 +1248,7 @@ def should_use_bf16(device=None, model_params=0, prioritize_performance=True, ma return True if cpu_mode(): - return False + return use_cpu_bf16() if is_intel_xpu(): return True diff --git a/requirements.txt b/requirements.txt index 479a29eec..c84b79742 100644 --- a/requirements.txt +++ b/requirements.txt @@ -18,6 +18,7 @@ Pillow scipy tqdm psutil +py-cpuinfo alembic SQLAlchemy