refactor devices

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/devices.py

@@ -1,26 +1,76 @@
 import os
-import gc
 import sys
 import time
 import contextlib
-from functools import wraps
 import torch
-from modules.errors import log
-from modules import cmd_args, shared, memstats, errors, timer
-
-if sys.platform == "darwin":
-    from modules import mac_specific # pylint: disable=ungrouped-imports
+from modules.errors import log, display, install
 
 
-previous_oom = 0
 debug = os.environ.get('SD_DEVICE_DEBUG', None) is not None
+install() # traceback handler
+opts = None # initialized in get_backend to avoid circular import
+args = None # initialized in get_backend to avoid circular import
+cuda_ok = torch.cuda.is_available()
+inference_context = torch.no_grad
+cpu = torch.device("cpu")
+
+fp16_ok = None # set once by test_fp16
+bf16_ok = None # set once by test_bf16
+
+backend = None # set by get_backend
+device = None # set by get_optimal_device
+dtype = None # set by set_dtype
+dtype_vae = None
+dtype_unet = None
+unet_needs_upcast = False # compatibility item
+onnx = None
+previous_oom = 0 # oom counter
+if debug:
+    log.info(f'Torch build config: {torch.__config__.show()}')
+# set_cuda_sync_mode('block') # none/auto/spin/yield/block
 
 
 def has_mps() -> bool:
     if sys.platform != "darwin":
         return False
     else:
-        return mac_specific.has_mps # pylint: disable=used-before-assignment
+        from modules import devices_mac # pylint: disable=ungrouped-imports
+        return devices_mac.has_mps # pylint: disable=used-before-assignment
+
+
+def get_backend(shared_cmd_opts, shared_opts):
+    global opts, args # pylint: disable=global-statement
+    opts = shared_opts
+    args = shared_cmd_opts
+    if args.use_openvino:
+        from modules.intel import openvino # pylint: disable=unused-import
+        name = 'openvino'
+        if hasattr(torch, 'xpu') and torch.xpu.is_available():
+            torch.xpu.is_available = lambda *args, **kwargs: False
+        torch.cuda.is_available = lambda *args, **kwargs: False
+    elif args.use_ipex or (hasattr(torch, 'xpu') and torch.xpu.is_available()):
+        name = 'ipex'
+        from modules.intel.ipex import ipex_init
+        ok, e = ipex_init()
+        if not ok:
+            log.error(f'IPEX initialization failed: {e}')
+            name = 'cpu'
+    elif args.use_directml:
+        name = 'directml'
+        from modules.dml import directml_init
+        ok, e = directml_init()
+        if not ok:
+            log.error(f'DirectML initialization failed: {e}')
+            name = 'cpu'
+    elif torch.cuda.is_available() and torch.version.cuda:
+        name = 'cuda'
+    elif torch.cuda.is_available() and torch.version.hip:
+        name = 'rocm'
+    elif sys.platform == 'darwin':
+        name = 'mps'
+    else:
+        name = 'cpu'
+    return name
 
 
 def get_gpu_info():
@@ -44,12 +94,13 @@ def get_gpu_info():
 
     if not torch.cuda.is_available():
         try:
-            if shared.cmd_opts.use_openvino:
+            if backend == 'openvino':
+                from modules.intel.openvino import get_openvino_device
                 return {
                     'device': get_openvino_device(), # pylint: disable=used-before-assignment
                     'openvino': get_package_version("openvino"),
                 }
-            elif shared.cmd_opts.use_directml:
+            elif backend == 'directml':
                 return {
                     'device': f'{torch.cuda.get_device_name(torch.cuda.current_device())} n={torch.cuda.device_count()}',
                     'directml': get_package_version("torch-directml"),
@@ -60,19 +111,19 @@ def get_gpu_info():
             return {}
     else:
         try:
-            if hasattr(torch, "xpu") and torch.xpu.is_available():
+            if backend == 'ipex':
                 return {
                     'device': f'{torch.xpu.get_device_name(torch.xpu.current_device())} n={torch.xpu.device_count()}',
                     'ipex': get_package_version('intel-extension-for-pytorch'),
                 }
-            elif torch.version.cuda:
+            elif backend == 'cuda':
                 return {
                     'device': f'{torch.cuda.get_device_name(torch.cuda.current_device())} n={torch.cuda.device_count()} arch={torch.cuda.get_arch_list()[-1]} capability={torch.cuda.get_device_capability(device)}',
                     'cuda': torch.version.cuda,
                     'cudnn': torch.backends.cudnn.version(),
                     'driver': get_driver(),
                 }
-            elif torch.version.hip:
+            elif backend == 'rocm':
                 return {
                     'device': f'{torch.cuda.get_device_name(torch.cuda.current_device())} n={torch.cuda.device_count()}',
                     'hip': torch.version.hip,
@@ -83,7 +134,7 @@ def get_gpu_info():
                 }
         except Exception as ex:
             if debug:
-                errors.display(ex, 'Device exception')
+                display(ex, 'Device exception')
             return { 'error': ex }
 
 
@@ -95,17 +146,18 @@ def extract_device_id(args, name): # pylint: disable=redefined-outer-name
 
 
 def get_cuda_device_string():
+    from modules.shared import cmd_opts
     if backend == 'ipex':
-        if shared.cmd_opts.device_id is not None:
-            return f"xpu:{shared.cmd_opts.device_id}"
+        if cmd_opts.device_id is not None:
+            return f"xpu:{cmd_opts.device_id}"
         return "xpu"
     elif backend == 'directml' and torch.dml.is_available():
-        if shared.cmd_opts.device_id is not None:
-            return f"privateuseone:{shared.cmd_opts.device_id}"
+        if cmd_opts.device_id is not None:
+            return f"privateuseone:{cmd_opts.device_id}"
         return torch.dml.get_device_string(torch.dml.default_device().index)
     else:
-        if shared.cmd_opts.device_id is not None:
-            return f"cuda:{shared.cmd_opts.device_id}"
+        if cmd_opts.device_id is not None:
+            return f"cuda:{cmd_opts.device_id}"
         return "cuda"
 
 
@@ -121,14 +173,17 @@ def get_optimal_device():
     return torch.device(get_optimal_device_name())
 
 
-def get_device_for(task):
-    if task in shared.cmd_opts.use_cpu:
-        log.debug(f'Forcing CPU for task: {task}')
-        return cpu
+def get_device_for(task): # pylint: disable=unused-argument
+    # if task in cmd_opts.use_cpu:
+    #    log.debug(f'Forcing CPU for task: {task}')
+    #    return cpu
     return get_optimal_device()
 
 
 def torch_gc(force=False, fast=False):
+    import gc
+    from modules import timer, memstats
+    from modules.shared import cmd_opts
     t0 = time.time()
     mem = memstats.memory_stats()
     gpu = mem.get('gpu', {})
@@ -140,7 +195,7 @@ def torch_gc(force=False, fast=False):
         used_gpu = round(100 * gpu.get('used', 0) / gpu.get('total', 1)) if gpu.get('total', 1) > 1 else 0
     used_ram = round(100 * ram.get('used', 0) / ram.get('total', 1)) if ram.get('total', 1) > 1 else 0
     global previous_oom # pylint: disable=global-statement
-    threshold = 0 if (shared.cmd_opts.lowvram and not shared.cmd_opts.use_zluda) else shared.opts.torch_gc_threshold
+    threshold = 0 if (cmd_opts.lowvram and not cmd_opts.use_zluda) else opts.torch_gc_threshold
     if force or threshold == 0 or used_gpu >= threshold or used_ram >= threshold:
         force = True
     if oom > previous_oom:
@@ -192,6 +247,12 @@ def set_cuda_sync_mode(mode):
 
 
 def test_fp16():
+    global fp16_ok # pylint: disable=global-statement
+    if fp16_ok is not None:
+        return fp16_ok
+    if sys.platform == "darwin" or backend == 'openvino': # override
+        fp16_ok = False
+        return fp16_ok
     try:
         x = torch.tensor([[1.5,.0,.0,.0]]).to(device=device, dtype=torch.float16)
         layerNorm = torch.nn.LayerNorm(4, eps=0.00001, elementwise_affine=True, dtype=torch.float16, device=device)
@@ -200,18 +261,20 @@ def test_fp16():
             raise RuntimeError('Torch FP16 test: dtype mismatch')
         if torch.all(torch.isnan(out)).item():
             raise RuntimeError('Torch FP16 test: NaN')
-        if debug:
-            log.debug('Torch FP16 test: pass')
-        return True
+        fp16_ok = True
     except Exception as ex:
         log.warning(f'Torch FP16 test fail: {ex}')
-        if shared.cmd_opts.experimental:
-            log.debug('Torch FP16 test fail: override experimental')
-            return True
-        return False
+        fp16_ok = False
+    return fp16_ok
 
 
 def test_bf16():
+    global bf16_ok # pylint: disable=global-statement
+    if bf16_ok is not None:
+        return bf16_ok
+    if sys.platform == "darwin" or backend == 'openvino': # override
+        bf16_ok = False
+        return bf16_ok
     try:
         import torch.nn.functional as F
         image = torch.randn(1, 4, 32, 32).to(device=device, dtype=torch.bfloat16)
@@ -220,17 +283,11 @@ def test_bf16():
             raise RuntimeError('Torch BF16 test: dtype mismatch')
         if torch.all(torch.isnan(out)).item():
             raise RuntimeError('Torch BF16 test: NaN')
-        if debug:
-            log.debug('Torch BF16 test: pass')
-        # if torch.cuda.is_available() and not torch.cuda.is_bf16_supported():
-        #    log.warning('Torch BF16 test: partial pass')
-        return True
+        bf16_ok = True
     except Exception as ex:
         log.warning(f'Torch BF16 test fail: {ex}')
-        if shared.cmd_opts.experimental:
-            log.debug('Torch FP16 test fail: override experimental')
-            return True
-        return False
+        bf16_ok = False
+    return bf16_ok
 
 
 def set_cudnn_params():
@@ -243,12 +300,12 @@ def set_cudnn_params():
             pass
         if torch.backends.cudnn.is_available():
             try:
-                torch.backends.cudnn.deterministic = shared.opts.cudnn_deterministic
-                torch.use_deterministic_algorithms(shared.opts.cudnn_deterministic)
-                if shared.opts.cudnn_deterministic:
+                torch.backends.cudnn.deterministic = opts.cudnn_deterministic
+                torch.use_deterministic_algorithms(opts.cudnn_deterministic)
+                if opts.cudnn_deterministic:
                     os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ':4096:8')
                 torch.backends.cudnn.benchmark = True
-                if shared.opts.cudnn_benchmark:
+                if opts.cudnn_benchmark:
                     log.debug('Torch cuDNN: enable benchmark')
                     torch.backends.cudnn.benchmark_limit = 0
                 torch.backends.cudnn.allow_tf32 = True
@@ -266,15 +323,16 @@ def override_ipex_math():
 
 def set_sdpa_params():
     try:
-        if shared.opts.cross_attention_optimization == "Scaled-Dot-Product":
-            torch.backends.cuda.enable_flash_sdp('Flash attention' in shared.opts.sdp_options)
-            torch.backends.cuda.enable_mem_efficient_sdp('Memory attention' in shared.opts.sdp_options)
-            torch.backends.cuda.enable_math_sdp('Math attention' in shared.opts.sdp_options)
+        if opts.cross_attention_optimization == "Scaled-Dot-Product":
+            torch.backends.cuda.enable_flash_sdp('Flash attention' in opts.sdp_options)
+            torch.backends.cuda.enable_mem_efficient_sdp('Memory attention' in opts.sdp_options)
+            torch.backends.cuda.enable_math_sdp('Math attention' in opts.sdp_options)
             if backend == "rocm":
-                if 'Flash attention' in shared.opts.sdp_options:
+                if 'Flash attention' in opts.sdp_options:
                     try:
                         # https://github.com/huggingface/diffusers/discussions/7172
                         from flash_attn import flash_attn_func
+                        from functools import wraps
                         backup_sdpa = torch.nn.functional.scaled_dot_product_attention
                         @wraps(torch.nn.functional.scaled_dot_product_attention)
                         def sdpa_hijack(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
@@ -283,10 +341,10 @@ def set_sdpa_params():
                             else:
                                 return backup_sdpa(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale)
                         torch.nn.functional.scaled_dot_product_attention = sdpa_hijack
-                        shared.log.debug('ROCm Flash Attention Hijacked')
+                        log.debug('ROCm Flash Attention Hijacked')
                     except Exception as err:
                         log.error(f'ROCm Flash Attention failed: {err}')
-            if 'Dynamic attention' in shared.opts.sdp_options:
+            if 'Dynamic attention' in opts.sdp_options:
                 from modules.sd_hijack_dynamic_atten import sliced_scaled_dot_product_attention
                 torch.nn.functional.scaled_dot_product_attention = sliced_scaled_dot_product_attention
     except Exception:
@@ -294,15 +352,10 @@ def set_sdpa_params():
 
 
 def set_dtype():
-    global dtype, dtype_vae, dtype_unet, unet_needs_upcast, inference_context, fp16_ok, bf16_ok # pylint: disable=global-statement
-    if shared.opts.cuda_dtype == 'Auto': # detect
-        if sys.platform == "darwin" or shared.cmd_opts.use_openvino: # override
-            fp16_ok = False
-            bf16_ok = False
-        else:
-            fp16_ok = test_fp16() if fp16_ok is None else fp16_ok
-            bf16_ok = test_bf16() if bf16_ok is None else bf16_ok
-
+    global dtype, dtype_vae, dtype_unet, unet_needs_upcast, inference_context # pylint: disable=global-statement
+    test_fp16()
+    test_bf16()
+    if opts.cuda_dtype == 'Auto': # detect
         if bf16_ok:
             dtype = torch.bfloat16
             dtype_vae = torch.bfloat16
@@ -315,102 +368,51 @@ def set_dtype():
             dtype = torch.float32
             dtype_vae = torch.float32
             dtype_unet = torch.float32
-    elif shared.opts.cuda_dtype == 'FP32':
+    elif opts.cuda_dtype == 'FP32':
         dtype = torch.float32
         dtype_vae = torch.float32
         dtype_unet = torch.float32
-        fp16_ok = None
-        bf16_ok = None
-    elif shared.opts.cuda_dtype == 'BF16':
-        fp16_ok = test_fp16() if fp16_ok is None else fp16_ok
-        bf16_ok = test_bf16() if bf16_ok is None else bf16_ok
-        dtype = torch.bfloat16 if bf16_ok else torch.float16
-        dtype_vae = torch.bfloat16 if bf16_ok else torch.float16
-        dtype_unet = torch.bfloat16 if bf16_ok else torch.float16
-    elif shared.opts.cuda_dtype == 'FP16':
-        fp16_ok = test_fp16() if fp16_ok is None else fp16_ok
-        bf16_ok = None
-        dtype = torch.float16 if fp16_ok else torch.float32
-        dtype_vae = torch.float16 if fp16_ok else torch.float32
-        dtype_unet = torch.float16 if fp16_ok else torch.float32
+    elif opts.cuda_dtype == 'BF16':
+        if not bf16_ok:
+            log.warning(f'Torch device capability failed: device={device} dtype={torch.bfloat16}')
+        dtype = torch.bfloat16
+        dtype_vae = torch.bfloat16
+        dtype_unet = torch.bfloat16
+    elif opts.cuda_dtype == 'FP16':
+        if not fp16_ok:
+            log.warning(f'Torch device capability failed: device={device} dtype={torch.float16}')
+        dtype = torch.float16
+        dtype_vae = torch.float16
+        dtype_unet = torch.float16
 
-    if shared.opts.no_half:
+    if opts.no_half:
         log.info('Torch override dtype: no-half set')
         dtype = torch.float32
         dtype_vae = torch.float32
         dtype_unet = torch.float32
-    if shared.opts.no_half_vae:
+    if opts.no_half_vae:
         log.info('Torch override VAE dtype: no-half set')
         dtype_vae = torch.float32
-    unet_needs_upcast = shared.opts.upcast_sampling
-    if shared.opts.inference_mode == 'inference-mode':
+    unet_needs_upcast = opts.upcast_sampling
+    if opts.inference_mode == 'inference-mode':
         inference_context = torch.inference_mode
-    elif shared.opts.inference_mode == 'none':
+    elif opts.inference_mode == 'none':
         inference_context = contextlib.nullcontext
     else:
         inference_context = torch.no_grad
 
 
 def set_cuda_params():
-    if debug:
-        log.debug(f'Verifying Torch settings: cuda={cuda_ok}')
     override_ipex_math()
     set_cudnn_params()
     set_sdpa_params()
     set_dtype()
-    if shared.cmd_opts.profile:
-        shared.log.debug(f'Torch info: {torch.__config__.show()}')
-    device_name = get_raw_openvino_device() if shared.cmd_opts.use_openvino else torch.device(get_optimal_device_name()) # pylint: disable=used-before-assignment
-    log.info(f'Torch parameters: device={device_name} config={shared.opts.cuda_dtype} dtype={dtype} vae={dtype_vae} unet={dtype_unet} context={inference_context.__name__} nohalf={shared.opts.no_half} nohalfvae={shared.opts.no_half_vae} upscast={shared.opts.upcast_sampling} deterministic={shared.opts.cudnn_deterministic} test-fp16={fp16_ok} test-bf16={bf16_ok} optimization="{shared.opts.cross_attention_optimization}"')
-
-
-args = cmd_args.parser.parse_args()
-backend = 'not set'
-if args.use_openvino:
-    from modules.intel.openvino import get_openvino_device
-    from modules.intel.openvino import get_device as get_raw_openvino_device
-    backend = 'openvino'
-    if hasattr(torch, 'xpu') and torch.xpu.is_available():
-        torch.xpu.is_available = lambda *args, **kwargs: False
-    torch.cuda.is_available = lambda *args, **kwargs: False
-elif args.use_ipex or (hasattr(torch, 'xpu') and torch.xpu.is_available()):
-    backend = 'ipex'
-    from modules.intel.ipex import ipex_init
-    ok, e = ipex_init()
-    if not ok:
-        log.error(f'IPEX initialization failed: {e}')
-        backend = 'cpu'
-elif args.use_directml:
-    backend = 'directml'
-    from modules.dml import directml_init
-    ok, e = directml_init()
-    if not ok:
-        log.error(f'DirectML initialization failed: {e}')
-        backend = 'cpu'
-elif torch.cuda.is_available() and torch.version.cuda:
-    backend = 'cuda'
-elif torch.cuda.is_available() and torch.version.hip:
-    backend = 'rocm'
-elif sys.platform == 'darwin':
-    backend = 'mps'
-else:
-    backend = 'cpu'
-
-
-inference_context = torch.no_grad
-cuda_ok = torch.cuda.is_available()
-cpu = torch.device("cpu")
-device = device_interrogate = device_gfpgan = device_esrgan = device_codeformer = None
-dtype = None
-dtype_vae = None
-dtype_unet = None
-fp16_ok = None
-bf16_ok = None
-unet_needs_upcast = False
-onnx = None
-if args.profile:
-    log.info(f'Torch build config: {torch.__config__.show()}')
-# set_cuda_sync_mode('block') # none/auto/spin/yield/block
+    if backend == 'openvino':
+        from modules.intel.openvino import get_device as get_raw_openvino_device
+        device_name = get_raw_openvino_device()
+    else:
+        device_name = torch.device(get_optimal_device_name())
+    log.info(f'Torch parameters: backend={backend} device={device_name} config={opts.cuda_dtype} dtype={dtype} vae={dtype_vae} unet={dtype_unet} context={inference_context.__name__} nohalf={opts.no_half} nohalfvae={opts.no_half_vae} upscast={opts.upcast_sampling} deterministic={opts.cudnn_deterministic} test-fp16={fp16_ok} test-bf16={bf16_ok} optimization="{opts.cross_attention_optimization}"')
 
 
 def cond_cast_unet(tensor):
@@ -429,7 +431,7 @@ def randn(seed, shape=None):
         return None
     if device.type == 'mps':
         return torch.randn(shape, device=cpu).to(device)
-    elif shared.opts.diffusers_generator_device == "CPU":
+    elif opts.diffusers_generator_device == "CPU":
         return torch.randn(shape, device=cpu)
     else:
         return torch.randn(shape, device=device)
@@ -441,9 +443,9 @@ def randn_without_seed(shape):
     return torch.randn(shape, device=device)
 
 def autocast(disable=False):
-    if disable or dtype == torch.float32 or shared.cmd_opts.precision == "Full":
+    if disable or dtype == torch.float32:
         return contextlib.nullcontext()
-    if shared.cmd_opts.use_directml:
+    if backend == 'directml':
         return torch.dml.amp.autocast(dtype)
     if cuda_ok:
         return torch.autocast("cuda")
@@ -454,7 +456,7 @@ def autocast(disable=False):
 def without_autocast(disable=False):
     if disable:
         return contextlib.nullcontext()
-    if shared.cmd_opts.use_directml:
+    if backend == 'directml':
         return torch.dml.amp.autocast(enabled=False) if torch.is_autocast_enabled() else contextlib.nullcontext() # pylint: disable=unexpected-keyword-arg
     if cuda_ok:
         return torch.autocast("cuda", enabled=False) if torch.is_autocast_enabled() else contextlib.nullcontext()
@@ -467,17 +469,17 @@ class NansException(Exception):
 
 
 def test_for_nans(x, where):
-    if shared.opts.disable_nan_check:
+    if opts.disable_nan_check:
         return
     if not torch.all(torch.isnan(x)).item():
         return
     if where == "unet":
         message = "A tensor with all NaNs was produced in Unet."
-        if not shared.opts.no_half:
+        if not opts.no_half:
             message += " This could be either because there's not enough precision to represent the picture, or because your video card does not support half type. Try setting the \"Upcast cross attention layer to float32\" option in Settings > Stable Diffusion or using the --no-half commandline argument to fix this."
     elif where == "vae":
         message = "A tensor with all NaNs was produced in VAE."
-        if not shared.opts.no_half and not shared.opts.no_half_vae:
+        if not opts.no_half and not opts.no_half_vae:
             message += " This could be because there's not enough precision to represent the picture. Try adding --no-half-vae commandline argument to fix this."
     else:
         message = "A tensor with all NaNs was produced."