fix(api): do not parse tokens when removing them

ssube · Nov 15, 2023 · 1bdfb4a · 1bdfb4a
1 parent 14ddfb4
commit 1bdfb4a
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/api/onnx_web/diffusers/utils.py b/api/onnx_web/diffusers/utils.py
@@ -3,7 +3,7 @@
 from logging import getLogger
 from math import ceil
 from re import Pattern, compile
-from typing import Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -515,9 +515,13 @@ def parse_reseed(prompt: str) -> Tuple[str, List[Reseed]]:
     return get_tokens_from_prompt(prompt, RESEED_TOKEN, parser=parse_reseed_group)
 
 
+def skip_group(group) -> Any:
+    return group
+
+
 def remove_tokens(prompt: Optional[str]) -> Optional[str]:
     if prompt is None:
         return prompt
 
-    remainder, tokens = get_tokens_from_prompt(prompt, ANY_TOKEN)
+    remainder, tokens = get_tokens_from_prompt(prompt, ANY_TOKEN, parser=skip_group)
     return remainder