Merge pull request #6625 from PlasmaPower/textual-inversion-safetensors

Support loading textual inversion embeddings from safetensors files

1 files changed, 3 insertions, 0 deletions

diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index 5420903f..3866c154 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -9,6 +9,7 @@ import tqdm
 import html

 import datetime

 import csv

+import safetensors.torch

 

 from PIL import Image, PngImagePlugin

 

@@ -150,6 +151,8 @@ class EmbeddingDatabase:
                 name = data.get('name', name)

         elif ext in ['.BIN', '.PT']:

             data = torch.load(path, map_location="cpu")

+        elif ext in ['.SAFETENSORS']:

+            data = safetensors.torch.load_file(path, device="cpu")

         else:

             return