1 files changed, 4 insertions, 12 deletions

diff --git a/infer.py b/infer.py
index 01010eb..ac05955 100644
--- a/infer.py
+++ b/infer.py
@@ -12,7 +12,7 @@ from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler, DDIMSc
 from transformers import CLIPTextModel, CLIPTokenizer
 from slugify import slugify
 
-from schedulers.scheduling_euler_a import EulerAScheduler
+from schedulers.scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler
 from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 
 
@@ -175,16 +175,8 @@ def load_embeddings_ti(tokenizer, text_encoder, embeddings_dir):
     embeddings_dir = Path(embeddings_dir)
     embeddings_dir.mkdir(parents=True, exist_ok=True)
 
-    for file in embeddings_dir.iterdir():
-        if file.is_file():
-            placeholder_token = file.stem
-
-            num_added_tokens = tokenizer.add_tokens(placeholder_token)
-            if num_added_tokens == 0:
-                raise ValueError(
-                    f"The tokenizer already contains the token {placeholder_token}. Please pass a different"
-                    " `placeholder_token` that is not already in the tokenizer."
-                )
+    placeholder_tokens = [file.stem for file in embeddings_dir.iterdir() if file.is_file()]
+    tokenizer.add_tokens(placeholder_tokens)
 
     text_encoder.resize_token_embeddings(len(tokenizer))
 
@@ -231,7 +223,7 @@ def create_pipeline(model, scheduler, ti_embeddings_dir, dtype):
             beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False
         )
     else:
-        scheduler = EulerAScheduler(
+        scheduler = EulerAncestralDiscreteScheduler(
             beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
         )