1 files changed, 3 insertions, 3 deletions

diff --git a/infer.py b/infer.py
index b29b136..507d0cf 100644
--- a/infer.py
+++ b/infer.py
@@ -28,7 +28,7 @@ from transformers import CLIPTextModel
 from models.clip.embeddings import patch_managed_embeddings
 from models.clip.tokenizer import MultiCLIPTokenizer
 from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
-from common import load_config, load_embeddings_from_dir
+from util import load_config, load_embeddings_from_dir
 
 
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -192,12 +192,12 @@ def save_args(basepath, args, extra={}):
 
 
 def load_embeddings(pipeline, embeddings_dir):
-    added_tokens = load_embeddings_from_dir(
+    added_tokens, added_ids = load_embeddings_from_dir(
         pipeline.tokenizer,
         pipeline.text_encoder.text_model.embeddings,
         Path(embeddings_dir)
     )
-    print(f"Added {len(added_tokens)} tokens from embeddings dir: {added_tokens}")
+    print(f"Added {len(added_tokens)} tokens from embeddings dir: {zip(added_tokens, added_ids)}")
 
 
 def create_pipeline(model, dtype):