Always grey images when using Stable Diffusion with Bumblebee

[ffloyd]

I have a MacBook with M3 Pro and I'm trying to utilize GPU for Neural Network tasks with LiveBook & Bumblebee.

DistillBERT (question answering) works fine, but when I try to use "text-to-image" I always get results like this.

But if use 3090 RTX via EXLA and the same model with the same settings I get normal images.

I'm using LiveBook. My setup:

Mix.install(
  [
    {:kino_bumblebee, "~> 0.5.0"},
    {:emlx, github: "elixir-nx/emlx"}
  ],
  config: [nx: [default_backend: {EMLX.Backend, device: :cpu}]],
  system_env: []
)

And Bumble-related codeblocks:

Nx.Defn.default_options(compiler: EMLX)

repository_id = "CompVis/stable-diffusion-v1-4"
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "openai/clip-vit-large-patch14"})
{:ok, clip} = Bumblebee.load_model({:hf, repository_id, subdir: "text_encoder"})
{:ok, unet} = Bumblebee.load_model({:hf, repository_id, subdir: "unet"})

{:ok, vae} =
  Bumblebee.load_model({:hf, repository_id, subdir: "vae"}, architecture: :decoder)

{: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"})

serving =
  Bumblebee.Diffusion.StableDiffusion.text_to_image(clip, unet, vae, tokenizer, scheduler,
    num_steps: 40,
    num_images_per_prompt: 1,
    safety_checker: safety_checker,
    safety_checker_featurizer: featurizer,
    compile: [batch_size: 1, sequence_length: 50]
  )

And UI:

text_input =
  Kino.Input.textarea("Text",
    default: "numbat, forest, high quality, detailed, digital art"
  )

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)

I've tried:

• enabling/disabling LIBMLX_ENABLE_JIT
• installing XCode instead of command line tools
• enabling LIBMLX_BUILD and build from sources
• updating MacOS from 14 to 15
• changing number of steps. Either black, grey or very-very blurred results.

My main concern:

I'm newbie to NN world. So I want reliable environment for learning & experimenting. And I want to avoid Python when I can and use LiveBook. I have set it up on my Linux PC with Nvidia GPU and it works well. With MacOS I was able to run ollama with big models, use BERT in LiveBook, but image generation doesn't work. I spent 10+ hours trying to fix it and this issue is manifestation of my despair =)