refactor devices

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

installer.py

@@ -664,8 +664,8 @@ def check_torch():
     allow_ipex = not (args.use_cuda or args.use_rocm or args.use_directml or args.use_openvino)
     allow_directml = not (args.use_cuda or args.use_rocm or args.use_ipex or args.use_openvino)
     allow_openvino = not (args.use_cuda or args.use_rocm or args.use_ipex or args.use_directml)
-    log.debug(f'Torch overrides: cuda={args.use_cuda} rocm={args.use_rocm} ipex={args.use_ipex} diml={args.use_directml} openvino={args.use_openvino}')
-    log.debug(f'Torch allowed: cuda={allow_cuda} rocm={allow_rocm} ipex={allow_ipex} diml={allow_directml} openvino={allow_openvino}')
+    log.debug(f'Torch overrides: cuda={args.use_cuda} rocm={args.use_rocm} ipex={args.use_ipex} diml={args.use_directml} openvino={args.use_openvino} zluda={args.use_zluda}')
+    # log.debug(f'Torch allowed: cuda={allow_cuda} rocm={allow_rocm} ipex={allow_ipex} diml={allow_directml} openvino={allow_openvino}')
     torch_command = os.environ.get('TORCH_COMMAND', '')
 
     if torch_command != '':

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

modules/interrogate.py

@@ -67,7 +67,7 @@ class InterrogateModels:
         self.loaded_categories = None
         self.skip_categories = []
         self.content_dir = content_dir
-        self.running_on_cpu = devices.device_interrogate == torch.device("cpu")
+        self.running_on_cpu = False
 
     def categories(self):
         if not os.path.exists(self.content_dir):
@@ -123,7 +123,7 @@ class InterrogateModels:
         else:
             model, preprocess = clip.load(clip_model_name, download_root=shared.opts.clip_models_path)
         model.eval()
-        model = model.to(devices.device_interrogate)
+        model = model.to(devices.device)
         return model, preprocess
 
     def load(self):
@@ -131,12 +131,12 @@ class InterrogateModels:
             self.blip_model = self.load_blip_model()
             if not shared.opts.no_half and not self.running_on_cpu:
                 self.blip_model = self.blip_model.half()
-        self.blip_model = self.blip_model.to(devices.device_interrogate)
+        self.blip_model = self.blip_model.to(devices.device)
         if self.clip_model is None:
             self.clip_model, self.clip_preprocess = self.load_clip_model()
             if not shared.opts.no_half and not self.running_on_cpu:
                 self.clip_model = self.clip_model.half()
-        self.clip_model = self.clip_model.to(devices.device_interrogate)
+        self.clip_model = self.clip_model.to(devices.device)
         self.dtype = next(self.clip_model.parameters()).dtype
 
     def send_clip_to_ram(self):
@@ -160,10 +160,10 @@ class InterrogateModels:
         if shared.opts.interrogate_clip_dict_limit != 0:
             text_array = text_array[0:int(shared.opts.interrogate_clip_dict_limit)]
         top_count = min(top_count, len(text_array))
-        text_tokens = clip.tokenize(list(text_array), truncate=True).to(devices.device_interrogate)
+        text_tokens = clip.tokenize(list(text_array), truncate=True).to(devices.device)
         text_features = self.clip_model.encode_text(text_tokens).type(self.dtype)
         text_features /= text_features.norm(dim=-1, keepdim=True)
-        similarity = torch.zeros((1, len(text_array))).to(devices.device_interrogate)
+        similarity = torch.zeros((1, len(text_array))).to(devices.device)
         for i in range(image_features.shape[0]):
             similarity += (100.0 * image_features[i].unsqueeze(0) @ text_features.T).softmax(dim=-1)
         similarity /= image_features.shape[0]
@@ -175,7 +175,7 @@ class InterrogateModels:
             transforms.Resize((blip_image_eval_size, blip_image_eval_size), interpolation=InterpolationMode.BICUBIC),
             transforms.ToTensor(),
             transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
-        ])(pil_image).unsqueeze(0).type(self.dtype).to(devices.device_interrogate)
+        ])(pil_image).unsqueeze(0).type(self.dtype).to(devices.device)
         with devices.inference_context():
             caption = self.blip_model.generate(gpu_image, sample=False, num_beams=shared.opts.interrogate_clip_num_beams, min_length=shared.opts.interrogate_clip_min_length, max_length=shared.opts.interrogate_clip_max_length)
         return caption[0]
@@ -199,7 +199,7 @@ class InterrogateModels:
             self.send_blip_to_ram()
             devices.torch_gc()
             res = caption
-            clip_image = self.clip_preprocess(pil_image).unsqueeze(0).type(self.dtype).to(devices.device_interrogate)
+            clip_image = self.clip_preprocess(pil_image).unsqueeze(0).type(self.dtype).to(devices.device)
             with devices.inference_context(), devices.autocast():
                 image_features = self.clip_model.encode_image(clip_image).type(self.dtype)
                 image_features /= image_features.norm(dim=-1, keepdim=True)

modules/postprocess/codeformer_model.py

@@ -41,7 +41,7 @@ def setup_model(dirname):
                 from facelib.utils.face_restoration_helper import FaceRestoreHelper
                 from facelib.detection.retinaface import retinaface
                 if self.net is not None and self.face_helper is not None:
-                    self.net.to(devices.device_codeformer)
+                    self.net.to(devices.device)
                     return self.net, self.face_helper
                 model_paths = modelloader.load_models(model_path, model_url, self.cmd_dir, download_name='codeformer-v0.1.0.pth', ext_filter=['.pth'])
                 if len(model_paths) != 0:
@@ -49,14 +49,14 @@ def setup_model(dirname):
                 else:
                     shared.log.error(f"Model failed loading: type=CodeFormer model={model_path}")
                     return None, None
-                net = CodeFormer(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256']).to(devices.device_codeformer)
+                net = CodeFormer(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256']).to(devices.device)
                 checkpoint = torch.load(ckpt_path)['params_ema']
                 net.load_state_dict(checkpoint)
                 net.eval()
                 shared.log.info(f"Model loaded: type=CodeFormer model={ckpt_path}")
                 if hasattr(retinaface, 'device'):
-                    retinaface.device = devices.device_codeformer
-                face_helper = FaceRestoreHelper(1, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png', use_parse=True, device=devices.device_codeformer)
+                    retinaface.device = devices.device
+                face_helper = FaceRestoreHelper(1, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png', use_parse=True, device=devices.device)
                 self.net = net
                 self.face_helper = face_helper
                 return net, face_helper
@@ -74,7 +74,7 @@ def setup_model(dirname):
                 self.create_models()
                 if self.net is None or self.face_helper is None:
                     return np_image
-                self.send_model_to(devices.device_codeformer)
+                self.send_model_to(devices.device)
                 self.face_helper.clean_all()
                 self.face_helper.read_image(np_image)
                 self.face_helper.get_face_landmarks_5(only_center_face=False, resize=640, eye_dist_threshold=5)
@@ -82,7 +82,7 @@ def setup_model(dirname):
                 for cropped_face in self.face_helper.cropped_faces:
                     cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
                     normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
-                    cropped_face_t = cropped_face_t.unsqueeze(0).to(devices.device_codeformer)
+                    cropped_face_t = cropped_face_t.unsqueeze(0).to(devices.device)
                     try:
                         with devices.inference_context():
                             output = self.net(cropped_face_t, w=w if w is not None else shared.opts.code_former_weight, adain=True)[0] # pylint: disable=not-callable

modules/postprocess/esrgan_model.py

@@ -128,7 +128,7 @@ class UpscalerESRGAN(Upscaler):
         model = self.load_model(selected_model)
         if model is None:
             return img
-        model.to(devices.device_esrgan)
+        model.to(devices.device)
         img = esrgan_upscale(model, img)
         if opts.upscaler_unload and selected_model in self.models:
             del self.models[selected_model]
@@ -143,7 +143,7 @@ class UpscalerESRGAN(Upscaler):
         if self.models.get(info.local_data_path, None) is not None:
             log.debug(f"Upscaler cached: type={self.name} model={info.local_data_path}")
             return self.models[info.local_data_path]
-        state_dict = torch.load(info.local_data_path, map_location='cpu' if devices.device_esrgan.type == 'mps' else None)
+        state_dict = torch.load(info.local_data_path, map_location='cpu' if devices.device.type == 'mps' else None)
         log.info(f"Upscaler loaded: type={self.name} model={info.local_data_path}")
 
         if "params_ema" in state_dict:
@@ -179,7 +179,7 @@ def upscale_without_tiling(model, img):
     img = img[:, :, ::-1]
     img = np.ascontiguousarray(np.transpose(img, (2, 0, 1))) / 255
     img = torch.from_numpy(img).float()
-    img = img.unsqueeze(0).to(devices.device_esrgan)
+    img = img.unsqueeze(0).to(devices.device)
     with devices.inference_context():
         output = model(img)
     output = output.squeeze().float().cpu().clamp_(0, 1).detach().numpy()

modules/postprocess/gfpgan_model.py

@@ -16,7 +16,7 @@ def gfpgann():
     import gfpgan # pylint: disable=unused-import
     global loaded_gfpgan_model # pylint: disable=global-statement
     if loaded_gfpgan_model is not None:
-        loaded_gfpgan_model.gfpgan.to(devices.device_gfpgan)
+        loaded_gfpgan_model.gfpgan.to(devices.device)
         return loaded_gfpgan_model
     if gfpgan_constructor is None:
         return None
@@ -30,8 +30,8 @@ def gfpgann():
         shared.log.error(f"Model failed loading: type=GFPGAN model={model_file}")
         return None
     if hasattr(facexlib.detection.retinaface, 'device'):
-        facexlib.detection.retinaface.device = devices.device_gfpgan
-    model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None, device=devices.device_gfpgan)
+        facexlib.detection.retinaface.device = devices.device
+    model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None, device=devices.device)
     loaded_gfpgan_model = model
     shared.log.info(f"Model loaded: type=GFPGAN model={model_file}")
     return model
@@ -48,7 +48,7 @@ def gfpgan_fix_faces(np_image):
     if model is None:
         return np_image
 
-    send_model_to(model, devices.device_gfpgan)
+    send_model_to(model, devices.device)
 
     np_image_bgr = np_image[:, :, ::-1]
     _cropped_faces, _restored_faces, gfpgan_output_bgr = model.enhance(np_image_bgr, has_aligned=False, only_center_face=False, paste_back=True)

modules/postprocess/realesrgan_model_arch.py

@@ -55,7 +55,7 @@ class RealESRGANer():
             self.device = torch.device(
                 f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu') if device is None else device
         else:
-            self.device = devices.device_esrgan if device is None else device
+            self.device = devices.device if device is None else device
 
         if isinstance(model_path, list):
             # dni

modules/sd_models.py

@@ -1737,6 +1737,7 @@ def reload_text_encoder(initial=False):
 
 
 def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model', force=False):
+    devices.set_cuda_params()
     load_dict = shared.opts.sd_model_dict != model_data.sd_dict
     from modules import lowvram, sd_hijack
     checkpoint_info = info or select_checkpoint(op=op) # are we selecting model or dictionary

modules/shared.py

@@ -15,7 +15,7 @@ import fasteners
 import orjson
 import diffusers
 from rich.console import Console
-from modules import errors, shared_items, shared_state, cmd_args, theme
+from modules import errors, devices, shared_items, shared_state, cmd_args, theme
 from modules.paths import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611
 from modules.dml import memory_providers, default_memory_provider, directml_do_hijack
 from modules.onnx_impl import initialize_onnx, execution_providers
@@ -24,7 +24,6 @@ from modules.memstats import memory_stats
 import modules.interrogate
 import modules.memmon
 import modules.styles
-import modules.devices as devices # pylint: disable=R0402
 import modules.paths as paths
 from installer import print_dict
 from installer import log as central_logger # pylint: disable=E0611
@@ -1074,15 +1073,12 @@ if not native:
     log.warning('Backend=original is in maintainance-only mode')
     opts.data['diffusers_offload_mode'] = 'none'
 
-try:
-    log.info(f'Device: {print_dict(devices.get_gpu_info())}')
-except Exception as ex:
-    log.error(f'Device: {ex}')
-
 prompt_styles = modules.styles.StyleDatabase(opts)
 reference_models = readfile(os.path.join('html', 'reference.json'))
 cmd_opts.disable_extension_access = (cmd_opts.share or cmd_opts.listen or (cmd_opts.server_name or False)) and not cmd_opts.insecure
-devices.device, devices.device_interrogate, devices.device_gfpgan, devices.device_esrgan, devices.device_codeformer = (devices.cpu if any(y in cmd_opts.use_cpu for y in [x, 'all']) else devices.get_optimal_device() for x in ['sd', 'interrogate', 'gfpgan', 'esrgan', 'codeformer'])
+
+devices.backend = devices.get_backend(cmd_opts, opts)
+devices.device = devices.get_optimal_device()
 devices.onnx = [opts.onnx_execution_provider]
 if opts.onnx_cpu_fallback and 'CPUExecutionProvider' not in devices.onnx:
     devices.onnx.append('CPUExecutionProvider')
@@ -1096,6 +1092,10 @@ if devices.backend == "directml":
 elif devices.backend == "cuda":
     initialize_zluda()
 initialize_onnx()
+try:
+    log.info(f'Device: {print_dict(devices.get_gpu_info())}')
+except Exception as ex:
+    log.error(f'Device: {ex}')
 
 
 class TotalTQDM: # compatibility with previous global-tqdm

modules/styles.py

@@ -228,7 +228,7 @@ class StyleDatabase:
             import concurrent
             future_items = {}
             candidates = list(files_cache.list_files(folder, ext_filter=['.json'], recursive=files_cache.not_hidden))
-            with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
+            with concurrent.futures.ThreadPoolExecutor(max_workers=shared.max_workers) as executor:
                 for fn in candidates:
                     if os.path.isfile(fn) and fn.lower().endswith(".json"):
                         future_items[executor.submit(self.load_style, fn, None)] = fn

modules/ui_extra_networks.py

@@ -622,7 +622,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
             refresh_time = time.time()
         if not skip_indexing:
             import concurrent
-            with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
+            with concurrent.futures.ThreadPoolExecutor(max_workers=shared.max_workers) as executor:
                 for page in get_pages():
                     executor.submit(page.create_items, ui.tabname)
         for page in get_pages():

modules/zluda.py

@@ -59,8 +59,6 @@ def initialize_zluda():
         if shared.opts.onnx_execution_provider == ExecutionProvider.CUDA:
             shared.opts.onnx_execution_provider = ExecutionProvider.CPU
 
-        devices.device_codeformer = devices.cpu
-
         result = test(device)
         if result is not None:
             shared.log.warning(f'ZLUDA device failed to pass basic operation test: index={device.index}, device_name={torch.cuda.get_device_name(device)}')
@@ -68,4 +66,4 @@ def initialize_zluda():
             torch.cuda.is_available = lambda: False
             devices.cuda_ok = False
             devices.backend = 'cpu'
-            devices.device = devices.device_esrgan = devices.device_gfpgan = devices.device_interrogate = devices.cpu
+            devices.device = devices.cpu

webui.py

@@ -153,7 +153,6 @@ def initialize():
 
 
 def load_model():
-    modules.devices.set_cuda_params()
     if not shared.opts.sd_checkpoint_autoload or (shared.cmd_opts.ckpt is not None and shared.cmd_opts.ckpt.lower() != 'none'):
         log.debug('Model auto load disabled')
     else: