How to release memory after generation?

[camba1]

Great repo, thanks for putting it together!
Question: I incorporated DiffusionKit to a FastApi API and it works great. However, I noticed that after generating an image, the memory used is not released. Given that Flux uses 26GB, it would be great if it would release all that memory. Is there a way to accomplish this? I tried to gc.collect, but that does not seem to do anything. You can see the code in my sample project.

Thanks!

Juan

[atiorh]

Hi @camba1, could you try calling https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.metal.clear_cache.html after the image generation is complete? I believe the peak memory consumption should be much less than 26GB for low_memory=True.

[camba1]

Hi @atiorh
Thanks for the quick reply. Using clear_cache did the trick. The memory gets released, which I great. In terms of memory usage, the terminal debug logging messages (when the pipeline verbose value is set to true) prints peak usage as 16 GB, but Activity monitor shows the python process using 26 GB.

I tried a basic image generation using the sample code from the readme (to rule out my other code from causing the high memory usage) and it still shows 26 GB in Activity monitor. See code at the bottom.

By the way, which version of Torch are you using? I downgraded to 2.3.0 since it kept complaining that "Torch version 2.4.1 has not been tested with coremltools. You may run into unexpected errors. Torch 2.3.0 is the most recent version that has been tested." Did not make a difference though, but I figured I'd ask just in case.

from diffusionkit.mlx import FluxPipeline

pipeline = FluxPipeline(
  shift=1.0,
  model_version="argmaxinc/mlx-FLUX.1-schnell", # model_version="argmaxinc/mlx-FLUX.1-dev" for FLUX.1-dev
  low_memory_mode=True,
  a16=True,
  w16=True,
)

HEIGHT = 512
WIDTH = 512
NUM_STEPS = 4  #  4 for FLUX.1-schnell, 50 for SD3 and FLUX.1-dev
CFG_WEIGHT = 0. # for FLUX.1-schnell, 5. for SD3

image, _ = pipeline.generate_image(
  "a photo of a cat",
  cfg_weight=CFG_WEIGHT,
  num_steps=NUM_STEPS,
  latent_size=(HEIGHT // 8, WIDTH // 8),
)

image.save("image.png")

Thanks again!