fix(api): handle empty prompts when parsing

ssube · Nov 15, 2023 · 14ddfb4 · 14ddfb4
1 parent a3519f5
commit 14ddfb4
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/api/onnx_web/diffusers/utils.py b/api/onnx_web/diffusers/utils.py
@@ -515,6 +515,9 @@ def parse_reseed(prompt: str) -> Tuple[str, List[Reseed]]:
     return get_tokens_from_prompt(prompt, RESEED_TOKEN, parser=parse_reseed_group)
 
 
-def remove_tokens(prompt: str) -> str:
+def remove_tokens(prompt: Optional[str]) -> Optional[str]:
+    if prompt is None:
+        return prompt
+
     remainder, tokens = get_tokens_from_prompt(prompt, ANY_TOKEN)
     return remainder