Bumblebee Documentation

Bumblebee is an Elixir counterpart to the Python Transformers library. It does not store the actual model weights itself, but rather implements the logic to fetch and pair trained parameters with a model architecture.

Key Requirements:

1. Implementation: The model architecture must be implemented in Bumblebee. You can verify this by checking if the class name under "architectures" in the repository's config.json exists in the Bumblebee codebase.
2. Format: Bumblebee supports pytorch_model.bin (PyTorch) and model.safetensors (Safetensors). It does not support Flax (flax_model.msgpack) or Tensorflow (tf_model.h5) formats.
3. Subdirectories: If a repository contains multiple models (e.g., stabilityai/stable-diffusion-2), use the subdir option: Bumblebee.load_model({:hf, "model-repo", subdir: "..."}).

Bumblebee requires a tokenizer.json file (the "fast tokenizer" format used by Rust). If a repository only provides a "slow tokenizer" (via tokenizer_config.json), you have two options:

1. Reuse a tokenizer: If the model is a fine-tuned version of a base model, load the tokenizer from the original base model repository.
2. Generate a tokenizer.json: Use a tokenizer generator tool to create the file. You can then load it using a specific Git revision: Bumblebee.load_tokenizer({:hf, "model-repo", revision: "..."}).

Integrating Bumblebee tasks into an application follows a two-step pattern:

1. Startup: On application startup, load the necessary models and preprocessors, build an instance of Nx.Serving, and start it under the application supervision tree.
2. Inference: To make predictions from user input, use Nx.Serving.batched_run/2. It is highly recommended to wrap this call in a Task to avoid blocking other user interactions.

Using Nx.Serving.batched_run/2 is critical because it automatically batches requests from concurrent users and distributes the results back to them transparently, improving throughput.

To use Bumblebee, add it and exla to your mix.exs dependencies. exla is highly recommended as it allows for just-in-time compilation and execution on CPU or GPU.

To use GPUs, you must set the XLA_TARGET environment variable according to the XLA documentation.

def deps do
  [
    {:bumblebee, "~> 0.6.0"},
    {:exla, ">= 0.0.0"}
  ]
end

In Livebook notebooks or standalone Elixir scripts, use Mix.install/2 to install both Bumblebee and EXLA, and to configure the Nx default backend in a single step.

Mix.install(
  [
    {:bumblebee, "~> 0.6.0"},
    {:exla, ">= 0.0.0"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)

To reduce server load and network bandwidth, move preprocessing work to the client (browser) before sending data to the Bumblebee server:

Images

Instead of uploading full-resolution PNG/JPEG files, use the browser's Canvas API to:

• Resize the image to the dimensions required by the Neural Network.
• Decode the image into raw pixel values immediately.

Audio

Decoding audio on the server typically requires ffmpeg. To avoid this, use the client to:

• Preprocess the audio.
• Send raw PCM data with a single channel directly to the server.

For real-time streaming requirements, consider using the Membrane Framework.

A featurizer is an abstraction used to convert raw data (such as text, images, or audio) into a format that a machine learning model can consume, typically a batched Nx.t or Nx.Container.t.

Every featurizer module in Bumblebee implements the Bumblebee.Featurizer behaviour and is expected to define a configuration struct. The featurization process typically follows these stages:

1. process_input/2: Converts a single raw input into a tensor.
2. process_batch/2 (Optional): A numerical function that performs batch processing on the tensors produced by process_input/2. This stage is highly efficient when used with Nx.Serving, as it can be merged with the model computation and compiled together.
3. batch_template/2 (Optional): Provides a template tensor of a specific batch size, used to prepare for batch processing.

You can adjust how the progress bar behaves during model downloads in your application configuration.

# Update every 10% instead of every 1%
config :bumblebee, :progress_bar_step, 10

# Disable progress bar entirely
config :bumblebee, :progress_bar_enabled, false

PNDM (Pseudo Numerical Methods for Diffusion Models) sampling relies on two numerical methods for solving ODEs: the Runge-Kutta (RK) method and the linear multi-step (LMS) method.

Because the transfer part (approximating the next sample based on the current sample and gradient) is non-linear, these are referred to as 'pseudo' numerical methods (PRK and PLMS).

Sampling Modes

1. Standard (PRK + PLMS): Uses the Runge-Kutta method for the first few steps (to establish a gradient history) and then switches to the linear multi-step method. This typically requires at least 4 steps.
2. Reduced Warmup (reduce_warmup: true): Uses lower-order linear multi-step for the initial samples instead of Runge-Kutta. This is more efficient as it requires fewer forward passes of the model. This requires at least 2 steps.

Bumblebee's generation logic uses iterative inference (autoregression). For encoder-decoder models, the encoder is run once, and its state is reused across all decoding iterations.

Streaming Hook Data

If you use the streaming capability, the hook receives a map for each token generated. Each attribute is a tensor with a leading batch dimension: