1 files changed, 3 insertions, 2 deletions

diff --git a/infer.py b/infer.py
index 1fd11e2..efeb24d 100644
--- a/infer.py
+++ b/infer.py
@@ -181,7 +181,7 @@ def save_args(basepath, args, extra={}):
         json.dump(info, f, indent=4)
 
 
-def create_pipeline(model, ti_embeddings_dir, dtype):
+def create_pipeline(model, embeddings_dir, dtype):
     print("Loading Stable Diffusion pipeline...")
 
     tokenizer = CLIPTokenizer.from_pretrained(model, subfolder='tokenizer', torch_dtype=dtype)
@@ -190,7 +190,8 @@ def create_pipeline(model, ti_embeddings_dir, dtype):
     unet = UNet2DConditionModel.from_pretrained(model, subfolder='unet', torch_dtype=dtype)
     scheduler = DDIMScheduler.from_pretrained(model, subfolder='scheduler', torch_dtype=dtype)
 
-    load_text_embeddings(tokenizer, text_encoder, Path(ti_embeddings_dir))
+    added_tokens = load_text_embeddings(tokenizer, text_encoder, embeddings_dir)
+    print(f"Added {len(added_tokens)} tokens from embeddings dir: {added_tokens}")
 
     pipeline = VlpnStableDiffusion(
         text_encoder=text_encoder,