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

---
 modules/textual_inversion/textual_inversion.py | 4 ++++
 1 file changed, 4 insertions(+)

(limited to 'modules/textual_inversion/textual_inversion.py')

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]
 
-- 
cgit v1.2.3