Allow bf16 computations on CPUs with BF16 support

Modern CPUs have native AVX512 BF16 instructions, which significantly improves matmul and conv2d operations.

With Bfloat16 instructions UNET steps are 40-50% faster on both AMD and Intel CPUs. There are minor visible changes with bf16, but no avalanche effects, so this feature is enabled by default with new `--use-cpu-bf16=auto` option. It can be disabled with `--use-cpu-bf16=no`.

Signed-off-by: Sv. Lockal <lockalsash@gmail.com>

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.")

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
+        if is_device_cpu(device):
-            return False
+            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

requirements.txt

@@ -18,6 +18,7 @@ Pillow
 scipy
 tqdm
 psutil
+py-cpuinfo
 alembic
 SQLAlchemy