From ec124607f47371a6cfd61a795f86a7f1cbd44651 Mon Sep 17 00:00:00 2001
From: wangshuai09 <391746016@qq.com>
Date: Sat, 27 Jan 2024 17:21:32 +0800
Subject: [PATCH] Add NPU Support

---
 modules/devices.py                            |  9 +++--
 modules/initialize.py                         |  6 +++-
 modules/npu_specific.py                       | 34 +++++++++++++++++++
 .../textual_inversion/textual_inversion.py    |  4 +++
 requirements.txt                              |  4 +++
 requirements_versions.txt                     |  4 +++
 webui.sh                                      |  4 +++
 7 files changed, 62 insertions(+), 3 deletions(-)
 create mode 100644 modules/npu_specific.py

diff --git a/modules/devices.py b/modules/devices.py
index ea1f712f9..f1e565015 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 ac95fc6f0..3285cc3c1 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 000000000..d8aebf9c2
--- /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 04dda585c..9c062503e 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]
 
diff --git a/requirements.txt b/requirements.txt
index 80b438455..4537402bb 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,6 +5,8 @@ accelerate
 basicsr
 blendmodes
 clean-fid
+cloudpickle
+decorator
 einops
 fastapi>=0.90.1
 gfpgan
@@ -26,9 +28,11 @@ resize-right
 
 safetensors
 scikit-image>=0.19
+synr==0.5.0
 timm
 tomesd
 torch
 torchdiffeq
 torchsde
+tornado
 transformers==4.30.2
diff --git a/requirements_versions.txt b/requirements_versions.txt
index cb7403a9d..95515b55a 100644
--- a/requirements_versions.txt
+++ b/requirements_versions.txt
@@ -4,6 +4,8 @@ accelerate==0.21.0
 basicsr==1.4.2
 blendmodes==2022
 clean-fid==0.1.35
+cloudpickle==3.0.0
+decorator==5.1.1
 einops==0.4.1
 fastapi==0.94.0
 gfpgan==1.3.8
@@ -23,10 +25,12 @@ realesrgan==0.3.0
 resize-right==0.0.2
 safetensors==0.3.1
 scikit-image==0.21.0
+synr==0.5.0
 timm==0.9.2
 tomesd==0.1.3
 torch
 torchdiffeq==0.2.3
 torchsde==0.2.6
+tornado==6.4
 transformers==4.30.2
 httpx==0.24.1
diff --git a/webui.sh b/webui.sh
index cff433272..3f6e87fd8 100755
--- a/webui.sh
+++ b/webui.sh
@@ -159,6 +159,10 @@ then
     if echo "$gpu_info" | grep -q "AMD" && [[ -z "${TORCH_COMMAND}" ]]
     then
         export TORCH_COMMAND="pip install torch==2.0.1+rocm5.4.2 torchvision==0.15.2+rocm5.4.2 --index-url https://download.pytorch.org/whl/rocm5.4.2"
+    elif echo "$gpu_info" | grep -q "Huawei" && [[ -z "${TORCH_COMMAND}" ]]
+    then
+        export TORCH_COMMAND="pip install torch==2.1.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu; pip install torch_npu"
+    
     fi
 fi