diff options
Diffstat (limited to 'modules')
| -rw-r--r-- | modules/devices.py | 9 | ||||
| -rw-r--r-- | modules/initialize.py | 6 | ||||
| -rw-r--r-- | modules/npu_specific.py | 34 | ||||
| -rw-r--r-- | modules/textual_inversion/textual_inversion.py | 4 |
4 files changed, 50 insertions, 3 deletions
diff --git a/modules/devices.py b/modules/devices.py index ea1f712f..f1e56501 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -3,7 +3,7 @@ import contextlib from functools import lru_cache import torch -from modules import errors, shared +from modules import errors, shared, npu_specific if sys.platform == "darwin": from modules import mac_specific @@ -40,6 +40,9 @@ def get_optimal_device_name(): if has_xpu(): return xpu_specific.get_xpu_device_string() + if npu_specific.has_npu: + return npu_specific.get_npu_device_string() + return "cpu" @@ -67,6 +70,9 @@ def torch_gc(): if has_xpu(): xpu_specific.torch_xpu_gc() + if npu_specific.has_npu: + npu_specific.torch_npu_gc() + def enable_tf32(): if torch.cuda.is_available(): @@ -164,4 +170,3 @@ def first_time_calculation(): x = torch.zeros((1, 1, 3, 3)).to(device, dtype) conv2d = torch.nn.Conv2d(1, 1, (3, 3)).to(device, dtype) conv2d(x) - diff --git a/modules/initialize.py b/modules/initialize.py index ac95fc6f..3285cc3c 100644 --- a/modules/initialize.py +++ b/modules/initialize.py @@ -143,13 +143,17 @@ def initialize_rest(*, reload_script_modules=False): its optimization may be None because the list of optimizaers has neet been filled
by that time, so we apply optimization again.
"""
+ from modules import devices
+ # Work around due to bug in torch_npu, revert me after fixed, @see https://gitee.com/ascend/pytorch/issues/I8KECW?from=project-issue
+ if devices.npu_specific.has_npu:
+ import torch
+ torch.npu.set_device(0)
shared.sd_model # noqa: B018
if sd_hijack.current_optimizer is None:
sd_hijack.apply_optimizations()
- from modules import devices
devices.first_time_calculation()
if not shared.cmd_opts.skip_load_model_at_start:
Thread(target=load_model).start()
diff --git a/modules/npu_specific.py b/modules/npu_specific.py new file mode 100644 index 00000000..d8aebf9c --- /dev/null +++ b/modules/npu_specific.py @@ -0,0 +1,34 @@ +import importlib +import torch + +from modules import shared + + +def check_for_npu(): + if importlib.util.find_spec("torch_npu") is None: + return False + import torch_npu + torch_npu.npu.set_device(0) + + try: + # Will raise a RuntimeError if no NPU is found + _ = torch.npu.device_count() + return torch.npu.is_available() + except RuntimeError: + return False + + +def get_npu_device_string(): + if shared.cmd_opts.device_id is not None: + return f"npu:{shared.cmd_opts.device_id}" + return "npu:0" + + +def torch_npu_gc(): + # Work around due to bug in torch_npu, revert me after fixed, @see https://gitee.com/ascend/pytorch/issues/I8KECW?from=project-issue + torch.npu.set_device(0) + with torch.npu.device(get_npu_device_string()): + torch.npu.empty_cache() + + +has_npu = check_for_npu() diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index 04dda585..9c062503 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/modules/textual_inversion/textual_inversion.py @@ -151,6 +151,10 @@ class EmbeddingDatabase: return embedding
def get_expected_shape(self):
+ # workaround
+ if devices.npu_specific.has_npu:
+ import torch
+ torch.npu.set_device(0)
vec = shared.sd_model.cond_stage_model.encode_embedding_init_text(",", 1)
return vec.shape[1]
|