Bumblebee Examples

Mix.install([
  {:bumblebee, "~> 0.1"},
  {:axon, "~> 0.5"},
  {:exla, "~> 0.5"},
  {:nx, "~> 0.5"},
  {:kino_bumblebee, "~> 0.1"}
])

Labelling a sentence

Set default backend for Nx

Nx.default_backend(EXLA.Backend)

Load model from Hugging Face

{:ok, model} = Bumblebee.load_model({:hf, "facebook/bart-large-mnli"})
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "facebook/bart-large-mnli"})

labels = ["New booking", "Update booking", "Cancel booking", "Refund"]

zero_shot_serving = Bumblebee.Text.zero_shot_classification(model, tokenizer, labels)

Lebel a sentence.

input = "I need a new seat"
Nx.Serving.run(zero_shot_serving, input)

Text Generation

{:ok, gpt2} = Bumblebee.load_model({:hf, "openai-community/gpt2-medium"})
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "openai-community/gpt2-medium"})
{:ok, generation_config} = Bumblebee.load_generation_config({:hf, "openai-community/gpt2"})

chatgpt_serving = Bumblebee.Text.generation(gpt2, tokenizer, generation_config)

text_input = Kino.Input.text("Text Suggestion", default: "I like to travel around the world")

text = Kino.Input.read(text_input)
Nx.Serving.run(chatgpt_serving, text)

Classifying Images

Load model from Hugging Face.

{:ok, model_info} = Bumblebee.load_model({:hf, "google/vit-base-patch16-224"})
{:ok, featurizer} = Bumblebee.load_featurizer({:hf, "google/vit-base-patch16-224"})

serving_for_classifying =
  Bumblebee.Vision.image_classification(model_info, featurizer,
    top_k: 1,
    compile: [batch_size: 1],
    defn_options: [compiler: EXLA]
  )

Setup input & result:

image_input = Kino.Input.image("Image", size: {224, 224})
form = Kino.Control.form([image: image_input], submit: "Run")
frame = Kino.Frame.new()

form
|> Kino.Control.stream()
|> Stream.filter(& &1.data.image)
|> Kino.listen(fn %{data: %{image: image}} ->
  Kino.Frame.render(frame, Kino.Markdown.new("Running..."))
  image_path = Kino.Input.file_path(image.file_ref)

  content = File.read!(image_path)

  image =
    content
    |> Nx.from_binary(:u8)
    |> Nx.reshape({image.height, image.width, 3})

  output = Nx.Serving.run(serving_for_classifying, image)

  output.predictions
  |> Enum.map(&{&1.label, &1.score})
  |> Kino.Bumblebee.ScoredList.new()
  |> then(&Kino.Frame.render(frame, &1))
end)

Kino.Layout.grid([form, frame], boxed: true, gap: 16)

Text to image

Setup Stable Diffusion

repo_id = "CompVis/stable-diffusion-v1-4"
opts = [params_variant: "fp16", type: :bf16, backend: EXLA.Backend]

{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "openai/clip-vit-large-patch14"})
{:ok, clip} = Bumblebee.load_model({:hf, repo_id, subdir: "text_encoder"}, opts)
{:ok, unet} = Bumblebee.load_model({:hf, repo_id, subdir: "unet"}, opts)
{:ok, vae} = Bumblebee.load_model({:hf, repo_id, subdir: "vae"}, [architecture: :decoder] ++ opts)
{:ok, scheduler} = Bumblebee.load_scheduler({:hf, repo_id, subdir: "scheduler"})
{:ok, featurizer} = Bumblebee.load_featurizer({:hf, repo_id, subdir: "feature_extractor"})
{:ok, safety_checker} = Bumblebee.load_model({:hf, repo_id, subdir: "safety_checker"}, opts)

Start serving

serving =
  Bumblebee.Diffusion.StableDiffusion.text_to_image(clip, unet, vae, tokenizer, scheduler,
    num_steps: 30,
    num_images_per_prompt: 3,
    safety_checker: safety_checker,
    safety_checker_featurizer: featurizer,
    compile: [batch_size: 1, sequence_length: 70],
    # Option 1
    defn_options: [compiler: EXLA]
    # Option 2 (reduces GPU usage, but runs noticeably slower)
    # Also remove `backend: EXLA.Backend` from the loading options above
    # defn_options: [compiler: EXLA, lazy_transfers: :always]
  )

Kino.start_child({Nx.Serving, name: StableDiffusion, serving: serving})

Setup text input

prompt_input =
  Kino.Input.text("Prompt", default: "pig, medium quality")

negative_prompt_input = Kino.Input.text("Negative Prompt", default: "")

Kino.Layout.grid([prompt_input, negative_prompt_input])

Generate image

prompt = Kino.Input.read(prompt_input)
negative_prompt = Kino.Input.read(negative_prompt_input)

output =
  Nx.Serving.batched_run(StableDiffusion, %{prompt: prompt, negative_prompt: negative_prompt})

for result <- output.results do
  Kino.Image.new(result.image)
end
|> Kino.Layout.grid(columns: 2)