Merge pull request AUTOMATIC1111#14981 from wangshuai09/gpu_info_for_…

…ascend Add training support and change lspci for Ascend NPU
anapnoe · Mar 4, 2024 · eee46a5 · eee46a5
2 parents 09b5ce6 + 3fb1c2e
commit eee46a5
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/modules/hypernetworks/hypernetwork.py b/modules/hypernetworks/hypernetwork.py
@@ -95,6 +95,7 @@ def __init__(self, dim, state_dict=None, layer_structure=None, activation_func=N
                         zeros_(b)
                     else:
                         raise KeyError(f"Key {weight_init} is not defined as initialization!")
+        devices.torch_npu_set_device()
         self.to(devices.device)
 
     def fix_old_state_dict(self, state_dict):

diff --git a/modules/sd_hijack_clip.py b/modules/sd_hijack_clip.py
@@ -230,7 +230,7 @@ def forward(self, texts):
             for fixes in self.hijack.fixes:
                 for _position, embedding in fixes:
                     used_embeddings[embedding.name] = embedding
-
+            devices.torch_npu_set_device()
             z = self.process_tokens(tokens, multipliers)
             zs.append(z)
 

diff --git a/webui.sh b/webui.sh
@@ -158,9 +158,9 @@ 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}" ]]
+    elif npu-smi info 2>/dev/null
     then
-        export TORCH_COMMAND="pip install torch==2.1.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu; pip install torch_npu"
+        export TORCH_COMMAND="pip install torch==2.1.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu; pip install torch_npu==2.1.0"
 
     fi
 fi