How to release memory after generation? #34

Activity

1. atiorh commented on Sep 12, 2024

   

   atiorh

   on Sep 12, 2024

   Contributor

   More actions

   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.
2. camba1 commented on Sep 13, 2024

   

   camba1

   on Sep 13, 2024

   Author

   More actions

   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!