Bumblebee Example

bumblebee_example.exs

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

repository_id = "nitrosocke/Ghibli-Diffusion"
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "openai/clip-vit-large-patch14"})

{:ok, clip} =
  Bumblebee.load_model({:hf, repository_id, subdir: "text_encoder"},
    backend: {EXLA.Backend, client: :host}n
  )

{:ok, unet} =
  Bumblebee.load_model({:hf, repository_id, subdir: "unet"},
    backend: {EXLA.Backend, client: :host}
  )

{:ok, vae} =
  Bumblebee.load_model({:hf, repository_id, subdir: "vae"},
    architecture: :decoder,
    backend: {EXLA.Backend, client: :host}
  )

{:ok, scheduler} = Bumblebee.load_scheduler({:hf, repository_id, subdir: "scheduler"})

{:ok, featurizer} =
  Bumblebee.load_featurizer({:hf, repository_id, subdir: "feature_extractor"})

{:ok, safety_checker} =
  Bumblebee.load_model({:hf, repository_id, subdir: "safety_checker"},
    backend: {EXLA.Backend, client: :host}
  )

serving =
  Bumblebee.Diffusion.StableDiffusion.text_to_image(clip, unet, vae, tokenizer, scheduler,
    num_steps: 20,
    num_images_per_prompt: 1,
    safety_checker: safety_checker,
    safety_checker_featurizer: featurizer,
    compile: [batch_size: 1, sequence_length: 50],
    defn_options: [compiler: EXLA, lazy_transfers: :always, preallocate_params: true]
  )
  
  text_input = Kino.Input.textarea("Text", default: "ghibli style numbat in forest")
seed_input = Kino.Input.number("Seed")
form = Kino.Control.form([text: text_input, seed: seed_input], submit: "Run")
frame = Kino.Frame.new()

Kino.listen(form, fn %{data: %{text: text, seed: seed}} ->
  Kino.Frame.render(frame, Kino.Text.new("Running..."))
  output = Nx.Serving.run(serving, %{prompt: text, seed: seed})

  for result <- output.results do
    Kino.Image.new(result.image)
  end
  |> Kino.Layout.grid(columns: 2)
  |> then(&Kino.Frame.render(frame, &1))
end)

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