Fix stable diffusion (#115)

elixir-nx · Dec 15, 2022 · 52632b1 · 52632b1
1 parent e867a7c
commit 52632b1
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/lib/bumblebee/diffusion/stable_diffusion.ex b/lib/bumblebee/diffusion/stable_diffusion.ex
@@ -402,6 +402,8 @@ defmodule Bumblebee.Diffusion.StableDiffusion do
 
   defp get_tokens(prompt) when is_binary(prompt), do: [{prompt, ""}]
 
+  defp get_tokens(prompts) when is_list(prompts), do: Enum.flat_map(prompts, &get_tokens/1)
+
   defp validate_input(input) when is_binary(input), do: true
 
   defp validate_input(%{prompt: input, negative: negative})