Răsfoiți Sursa

Merge pull request #14981 from wangshuai09/gpu_info_for_ascend

Add training support and change lspci for Ascend NPU
AUTOMATIC1111 1 an în urmă
părinte
comite
eee46a5094
3 a modificat fișierele cu 4 adăugiri și 3 ștergeri
  1. 1 0
      modules/hypernetworks/hypernetwork.py
  2. 1 1
      modules/sd_hijack_clip.py
  3. 2 2
      webui.sh

+ 1 - 0
modules/hypernetworks/hypernetwork.py

@@ -95,6 +95,7 @@ class HypernetworkModule(torch.nn.Module):
                         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):

+ 1 - 1
modules/sd_hijack_clip.py

@@ -230,7 +230,7 @@ class FrozenCLIPEmbedderWithCustomWordsBase(torch.nn.Module):
             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)
 

+ 2 - 2
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