EMLX compiler hangs when calling larger expressions via nif_call

[polvalente]

Pretty much the title. Running the script in the body of #87 with Nx.Defn.global_default_options(compiler: EMLX) hangs indefinitely.

[polvalente]

Mix.install([
   {:bumblebee, github: "elixir-nx/bumblebee", override: true},
  {:emlx, path: __DIR__},
  {:stb_image, "~> 0.1"},
  {:kino_bumblebee, "~> 0.5.1"}
])

Nx.global_default_backend({EMLX.Backend, device: :gpu})

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]
  )

IO.puts("running...")
%{results: [%{image: tensor}]} = Nx.Serving.run(serving, %{prompt: "numbat, forest, high quality, detailed, digital art", seed: 1})

tensor |> StbImage.from_nx() |> StbImage.write_file("/tmp/output.png")

[polvalente]

Decided to just defer compiler callbacks to Nx.Defn.Evaluator. MLX caches most things internally.