1 files changed, 2 insertions, 2 deletions

diff --git a/infer.py b/infer.py
index cf59bba..ed86ab1 100644
--- a/infer.py
+++ b/infer.py
@@ -28,7 +28,7 @@ from diffusers import (
 )
 from transformers import CLIPTextModel
 
-from data.keywords import prompt_to_keywords, keywords_to_prompt
+from data.keywords import str_to_keywords, keywords_to_str
 from models.clip.embeddings import patch_managed_embeddings
 from models.clip.tokenizer import MultiCLIPTokenizer
 from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
@@ -296,7 +296,7 @@ def create_pipeline(model, dtype):
 
 
 def shuffle_prompts(prompts: list[str]) -> list[str]:
-    return [keywords_to_prompt(prompt_to_keywords(prompt), shuffle=True) for prompt in prompts]
+    return [keywords_to_str(str_to_keywords(prompt), shuffle=True) for prompt in prompts]
 
 
 @torch.inference_mode()