Merge pull request #14801 from wangshuai09/npu_support

Add NPU Support

5 files changed, 56 insertions, 3 deletions

diff --git a/modules/devices.py b/modules/devices.py
index 8f49f7a4..28c0c54d 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
@@ -57,6 +57,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"
 
 
@@ -84,6 +87,16 @@ def torch_gc():
     if has_xpu():
         xpu_specific.torch_xpu_gc()
 
+    if npu_specific.has_npu:
+        torch_npu_set_device()
+        npu_specific.torch_npu_gc()
+
+
+def torch_npu_set_device():
+    # Work around due to bug in torch_npu, revert me after fixed, @see https://gitee.com/ascend/pytorch/issues/I8KECW?from=project-issue
+    if npu_specific.has_npu:
+        torch.npu.set_device(0)
+
 
 def enable_tf32():
     if torch.cuda.is_available():
@@ -256,4 +269,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 7c1ac99e..f7313ff4 100644
--- a/modules/initialize.py
+++ b/modules/initialize.py
@@ -142,13 +142,14 @@ 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

+        devices.torch_npu_set_device()

 

         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/launch_utils.py b/modules/launch_utils.py
index 3ff4576a..107c72b0 100644
--- a/modules/launch_utils.py
+++ b/modules/launch_utils.py
@@ -338,6 +338,7 @@ def prepare_environment():
             torch_index_url = os.environ.get('TORCH_INDEX_URL', "https://pytorch-extension.intel.com/release-whl/stable/xpu/us/")

             torch_command = os.environ.get('TORCH_COMMAND', f"pip install torch==2.0.0a0 intel-extension-for-pytorch==2.0.110+gitba7f6c1 --extra-index-url {torch_index_url}")

     requirements_file = os.environ.get('REQS_FILE', "requirements_versions.txt")

+    requirements_file_for_npu = os.environ.get('REQS_FILE_FOR_NPU', "requirements_npu.txt")

 

     xformers_package = os.environ.get('XFORMERS_PACKAGE', 'xformers==0.0.23.post1')

     clip_package = os.environ.get('CLIP_PACKAGE', "https://github.com/openai/CLIP/archive/d50d76daa670286dd6cacf3bcd80b5e4823fc8e1.zip")

@@ -421,6 +422,13 @@ def prepare_environment():
         run_pip(f"install -r \"{requirements_file}\"", "requirements")

         startup_timer.record("install requirements")

 

+    if not os.path.isfile(requirements_file_for_npu):

+        requirements_file_for_npu = os.path.join(script_path, requirements_file_for_npu)

+

+    if "torch_npu" in torch_command and not requirements_met(requirements_file_for_npu):

+        run_pip(f"install -r \"{requirements_file_for_npu}\"", "requirements_for_npu")

+        startup_timer.record("install requirements_for_npu")

+

     if not args.skip_install:

         run_extensions_installers(settings_file=args.ui_settings_file)

 

diff --git a/modules/npu_specific.py b/modules/npu_specific.py
new file mode 100644
index 00000000..94100691
--- /dev/null
+++ b/modules/npu_specific.py
@@ -0,0 +1,31 @@
+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
+
+    try:
+        # Will raise a RuntimeError if no NPU is found
+        _ = torch_npu.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():
+    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 c6bcab15..6d815c0b 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -150,6 +150,7 @@ class EmbeddingDatabase:
         return embedding

 

     def get_expected_shape(self):

+        devices.torch_npu_set_device()

         vec = shared.sd_model.cond_stage_model.encode_embedding_init_text(",", 1)

         return vec.shape[1]