Don't block on tensor access in postprocessing (#245)

Author: jonatanklosko

lib/bumblebee/diffusion/stable_diffusion.ex

@@ -315,6 +315,9 @@ defmodule Bumblebee.Diffusion.StableDiffusion do
   end
 
   defp client_postprocessing({outputs, _metadata}, multi?, safety_checker?) do
+    # We use binary backend so we are not blocked by the serving computation
+    outputs = Nx.backend_transfer(outputs, Nx.BinaryBackend)
+
     for outputs <- Bumblebee.Utils.Nx.batch_to_list(outputs) do
       results =
         for outputs = %{image: image} <- Bumblebee.Utils.Nx.batch_to_list(outputs) do
@@ -336,7 +339,7 @@ defmodule Bumblebee.Diffusion.StableDiffusion do
 
   defp zeroed(tensor) do
     0
-    |> Nx.tensor(type: Nx.type(tensor))
+    |> Nx.tensor(type: Nx.type(tensor), backend: Nx.BinaryBackend)
     |> Nx.broadcast(Nx.shape(tensor))
   end
 

lib/bumblebee/text/question_answering.ex

@@ -102,6 +102,10 @@ defmodule Bumblebee.Text.QuestionAnswering do
       {batch, {all_inputs, raw_inputs, multi?}}
     end)
     |> Nx.Serving.client_postprocessing(fn {outputs, _metadata}, {inputs, raw_inputs, multi?} ->
+      # We use binary backend so we are not blocked by the serving computation
+      inputs = Nx.backend_transfer(inputs, Nx.BinaryBackend)
+      outputs = Nx.backend_transfer(outputs, Nx.BinaryBackend)
+
       Enum.zip_with(
         [raw_inputs, Utils.Nx.batch_to_list(inputs), Utils.Nx.batch_to_list(outputs)],
         fn [{_question_text, context_text}, inputs, outputs] ->

lib/bumblebee/text/text_embedding.ex

@@ -121,6 +121,9 @@ defmodule Bumblebee.Text.TextEmbedding do
       {batch, multi?}
     end)
     |> Nx.Serving.client_postprocessing(fn {embeddings, _metadata}, multi? ->
+      # We use binary backend so we are not blocked by the serving computation
+      embeddings = Nx.backend_transfer(embeddings, Nx.BinaryBackend)
+
       for embedding <- Bumblebee.Utils.Nx.batch_to_list(embeddings) do
         %{embedding: embedding}
       end

lib/bumblebee/text/token_classification.ex

@@ -87,6 +87,10 @@ defmodule Bumblebee.Text.TokenClassification do
       {batch, {all_inputs, multi?}}
     end)
     |> Nx.Serving.client_postprocessing(fn {scores, _metadata}, {inputs, multi?} ->
+      # We use binary backend so we are not blocked by the serving computation
+      scores = Nx.backend_transfer(scores, Nx.BinaryBackend)
+      inputs = Nx.backend_transfer(inputs, Nx.BinaryBackend)
+
       Enum.zip_with(
         Utils.Nx.batch_to_list(inputs),
         Utils.Nx.batch_to_list(scores),

lib/bumblebee/vision/image_embedding.ex

@@ -85,6 +85,9 @@ defmodule Bumblebee.Vision.ImageEmbedding do
       {Nx.Batch.concatenate([inputs]), multi?}
     end)
     |> Nx.Serving.client_postprocessing(fn {embeddings, _metadata}, multi? ->
+      # We use binary backend so we are not blocked by the serving computation
+      embeddings = Nx.backend_transfer(embeddings, Nx.BinaryBackend)
+
       for embedding <- Bumblebee.Utils.Nx.batch_to_list(embeddings) do
         %{embedding: embedding}
       end