Update serving input specs to also allow nil for optional attributes

elixir-nx · Feb 15, 2024 · 8290545 · 8290545
1 parent ae3deac
commit 8290545
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 5 deletions.
diff --git a/lib/bumblebee/audio.ex b/lib/bumblebee/audio.ex
@@ -43,7 +43,7 @@ defmodule Bumblebee.Audio do
   @type audio :: Nx.t() | {:file, String.t()}
 
   @type speech_to_text_whisper_input ::
-          audio() | %{:audio => audio(), optional(:seed) => integer()}
+          audio() | %{:audio => audio(), optional(:seed) => integer() | nil}
   @type speech_to_text_whisper_output :: %{
           chunks: list(speech_to_text_whisper_chunk())
         }

diff --git a/lib/bumblebee/diffusion/stable_diffusion.ex b/lib/bumblebee/diffusion/stable_diffusion.ex
@@ -12,8 +12,8 @@ defmodule Bumblebee.Diffusion.StableDiffusion do
           String.t()
           | %{
               :prompt => String.t(),
-              optional(:negative_prompt) => String.t(),
-              optional(:seed) => integer()
+              optional(:negative_prompt) => String.t() | nil,
+              optional(:seed) => integer() | nil
             }
   @type text_to_image_output :: %{results: list(text_to_image_result())}
   @type text_to_image_result :: %{:image => Nx.Tensor.t(), optional(:is_safe) => boolean()}

diff --git a/lib/bumblebee/text.ex b/lib/bumblebee/text.ex
@@ -127,7 +127,7 @@ defmodule Bumblebee.Text do
     to: Bumblebee.Text.TokenClassification
 
   @type generation_input ::
-          String.t() | %{:text => String.t(), optional(:seed) => integer()}
+          String.t() | %{:text => String.t(), optional(:seed) => integer() | nil}
   @type generation_output :: %{results: list(generation_result())}
   @type generation_result :: %{text: String.t(), token_summary: token_summary()}
   @type token_summary :: %{

diff --git a/lib/bumblebee/vision.ex b/lib/bumblebee/vision.ex
@@ -86,7 +86,7 @@ defmodule Bumblebee.Vision do
   defdelegate image_classification(model_info, featurizer, opts \\ []),
     to: Bumblebee.Vision.ImageClassification
 
-  @type image_to_text_input :: image() | %{:image => image(), optional(:seed) => integer()}
+  @type image_to_text_input :: image() | %{:image => image(), optional(:seed) => integer() | nil}
   @type image_to_text_output :: %{results: list(image_to_text_result())}
   @type image_to_text_result :: %{text: String.t()}