Merge pull request #131 from ssakar/memory_saver

Memory usage optimizations

README.md

@@ -191,6 +191,8 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
 
 - **`--config-from-metadata`** or **`-C`** (optional, `str`): [EXPERIMENTAL] Path to a prior file saved via `--metadata`, or a compatible handcrafted config file adhering to the expected args schema.
 
+- **`--low-ram`** (optional): Reduces GPU memory usage by constraining the MLX cache size and releasing text encoders and transformer components after use. This option is only compatible with single image generation. While it may slightly decrease performance, it helps prevent system memory swapping to disk, allowing generation on systems with limited RAM.
+
 <details>
 <summary>parameters supported by config files</summary>
 

src/mflux/callbacks/callback.py

@@ -32,6 +32,17 @@ class InLoopCallback(Protocol):
         ...
 
 
+class AfterLoopCallback(Protocol):
+    def call_after_loop(
+        self,
+        seed: int,
+        prompt: str,
+        latents: mx.array,
+        config: RuntimeConfig
+    ) -> None:  # fmt: off
+        ...
+
+
 class InterruptCallback(Protocol):
     def call_interrupt(
         self,

src/mflux/callbacks/callback_registry.py

@@ -1,10 +1,11 @@
-from mflux.callbacks.callback import BeforeLoopCallback, InLoopCallback, InterruptCallback
+from mflux.callbacks.callback import BeforeLoopCallback, InLoopCallback, InterruptCallback, AfterLoopCallback
 
 
 class CallbackRegistry:
     in_loop = []
     before_loop = []
     interrupt = []
+    after_loop = []
 
     @staticmethod
     def register_in_loop(callback: InLoopCallback) -> None:
@@ -15,17 +16,25 @@ class CallbackRegistry:
         CallbackRegistry.before_loop.append(callback)
 
     @staticmethod
-    def register_interrupt(callback: InterruptCallback) -> None:
-        CallbackRegistry.interrupt.append(callback)
+    def register_after_loop(callback: AfterLoopCallback) -> None:
+        CallbackRegistry.after_loop.append(callback)
 
     @staticmethod
-    def in_loop_callbacks() -> list[InLoopCallback]:
-        return CallbackRegistry.in_loop
+    def register_interrupt(callback: InterruptCallback) -> None:
+        CallbackRegistry.interrupt.append(callback)
 
     @staticmethod
     def before_loop_callbacks() -> list[BeforeLoopCallback]:
         return CallbackRegistry.before_loop
 
+    @staticmethod
+    def in_loop_callbacks() -> list[InLoopCallback]:
+        return CallbackRegistry.in_loop
+
+    @staticmethod
+    def after_loop_callbacks() -> list[AfterLoopCallback]:
+        return CallbackRegistry.after_loop
+
     @staticmethod
     def interrupt_callbacks() -> list[InterruptCallback]:
         return CallbackRegistry.interrupt

src/mflux/callbacks/callbacks.py

@@ -43,6 +43,21 @@ class Callbacks:
                 time_steps=time_steps,
             )
 
+    @staticmethod
+    def after_loop(
+        seed: int,
+        prompt: str,
+        latents: mx.array,
+        config: RuntimeConfig
+    ):  # fmt: off
+        for subscriber in CallbackRegistry.after_loop_callbacks():
+            subscriber.call_after_loop(
+                seed=seed,
+                prompt=prompt,
+                latents=latents,
+                config=config
+            )
+
     @staticmethod
     def interruption(
         t: int,

src/mflux/callbacks/instances/memory_saver.py

@@ -0,0 +1,74 @@
+import gc
+
+import mlx.core as mx
+import PIL.Image
+from tqdm import tqdm
+
+from mflux.callbacks.callback import (
+    AfterLoopCallback,
+    BeforeLoopCallback,
+    InLoopCallback
+)
+from mflux.config.runtime_config import RuntimeConfig
+
+
+class MemorySaver(BeforeLoopCallback, InLoopCallback, AfterLoopCallback):
+    """
+    Optimizes memory usage by clearing caches and removing unused model 
+    components at strategic points in the execution cycle.
+    """
+    
+    def __init__(self, flux, cache_limit_bytes: int = 1000**3):
+        self.flux = flux
+        self.peak_memory: int = 0
+        mx.metal.set_cache_limit(cache_limit_bytes)
+        mx.metal.clear_cache()
+        mx.metal.reset_peak_memory()
+    
+    def call_before_loop(
+        self,
+        seed: int,
+        prompt: str,
+        latents: mx.array,
+        config: RuntimeConfig,
+        canny_image: PIL.Image.Image | None = None,
+    ) -> None:
+        self.peak_memory = mx.metal.get_peak_memory()
+        self._delete_encoders()
+    
+    def call_in_loop(
+        self,
+        t: int,
+        seed: int,
+        prompt: str,
+        latents: mx.array,
+        config: RuntimeConfig,
+        time_steps: tqdm,
+    ) -> None:
+        self.peak_memory = mx.metal.get_peak_memory()
+    
+    def call_after_loop(
+        self,
+        seed: int,
+        prompt: str,
+        latents: mx.array,
+        config: RuntimeConfig,
+    ) -> None:
+        self.peak_memory = mx.metal.get_peak_memory()
+        self._delete_transformer()
+        
+    def _delete_encoders(self) -> None:
+        self.flux.clip_text_encoder = None
+        self.flux.t5_text_encoder = None
+        gc.collect()
+        mx.metal.clear_cache()
+    
+    def _delete_transformer(self) -> None:
+        self.flux.transformer = None
+        if hasattr(self.flux, 'transformer_controlnet'):
+            self.flux.transformer_controlnet = None
+        gc.collect()
+        mx.metal.clear_cache()
+
+    def memory_stats(self) -> str:
+        return f"Peak MLX memory: {self.peak_memory / 10**9:.2f} GB"

src/mflux/controlnet/flux_controlnet.py

@@ -119,7 +119,7 @@ class Flux1Controlnet(nn.Module):
                 dt = config.sigmas[t + 1] - config.sigmas[t]
                 latents += noise * dt
 
-                # (Optional) Call subscribes at end of loop
+                # (Optional) Call subscribers at end of loop
                 Callbacks.in_loop(
                     t=t,
                     seed=seed,
@@ -142,6 +142,14 @@ class Flux1Controlnet(nn.Module):
                     time_steps=time_steps,
                 )
 
+        # (Optional) Call subscribers at end of loop
+        Callbacks.after_loop(
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=config,
+        )  # fmt: off
+
         # 7. Decode the latent array and return the image
         latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
         decoded = self.vae.decode(latents)

src/mflux/flux/flux.py

@@ -99,7 +99,7 @@ class Flux1(nn.Module):
                 dt = config.sigmas[t + 1] - config.sigmas[t]
                 latents += noise * dt
 
-                # (Optional) Call subscribes at end of loop
+                # (Optional) Call subscribers at end of loop
                 Callbacks.in_loop(
                     t=t,
                     seed=seed,
@@ -122,6 +122,14 @@ class Flux1(nn.Module):
                     time_steps=time_steps,
                 )
 
+        # (Optional) Call subscribers at end of loop
+        Callbacks.after_loop(
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=config
+        )  # fmt: off
+
         # 7. Decode the latent array and return the image
         latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
         decoded = self.vae.decode(latents)

src/mflux/generate.py

@@ -1,12 +1,14 @@
 from mflux import Config, Flux1, ModelConfig, StopImageGenerationException
 from mflux.callbacks.callback_registry import CallbackRegistry
 from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
+from mflux.callbacks.instances.memory_saver import MemorySaver
 from mflux.ui.cli.parsers import CommandLineParser
 
 
 def main():
     # 0. Parse command line arguments
     parser = CommandLineParser(description="Generate an image based on a prompt.")
+    parser.add_general_arguments()
     parser.add_model_arguments(require_model_arg=False)
     parser.add_lora_arguments()
     parser.add_image_generator_arguments(supports_metadata_config=True)
@@ -30,6 +32,13 @@ def main():
         CallbackRegistry.register_in_loop(handler)
         CallbackRegistry.register_interrupt(handler)
 
+    memory_saver = None
+    if args.low_ram:
+        memory_saver = MemorySaver(flux)
+        CallbackRegistry.register_before_loop(memory_saver)
+        CallbackRegistry.register_in_loop(memory_saver)
+        CallbackRegistry.register_after_loop(memory_saver)
+
     try:
         for seed in args.seed:
             # 3. Generate an image for each seed value
@@ -49,6 +58,9 @@ def main():
             image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata)
     except StopImageGenerationException as stop_exc:
         print(stop_exc)
+    finally:
+        if memory_saver:
+            print(memory_saver.memory_stats())
 
 
 if __name__ == "__main__":

src/mflux/generate_controlnet.py

@@ -2,12 +2,14 @@ from mflux import Config, Flux1Controlnet, ModelConfig, StopImageGenerationExcep
 from mflux.callbacks.callback_registry import CallbackRegistry
 from mflux.callbacks.instances.canny_saver import CannyImageSaver
 from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
+from mflux.callbacks.instances.memory_saver import MemorySaver
 from mflux.ui.cli.parsers import CommandLineParser
 
 
 def main():
     # 0. Parse command line arguments
     parser = CommandLineParser(description="Generate an image based on a prompt and a controlnet reference image.")  # fmt: off
+    parser.add_general_arguments()
     parser.add_model_arguments(require_model_arg=True)
     parser.add_lora_arguments()
     parser.add_image_generator_arguments(supports_metadata_config=False)
@@ -33,6 +35,13 @@ def main():
         CallbackRegistry.register_in_loop(handler)
         CallbackRegistry.register_interrupt(handler)
 
+    memory_saver = None
+    if args.low_ram:
+        memory_saver = MemorySaver(flux)
+        CallbackRegistry.register_before_loop(memory_saver)
+        CallbackRegistry.register_in_loop(memory_saver)
+        CallbackRegistry.register_after_loop(memory_saver)
+
     try:
         for seed in args.seed:
             # 3. Generate an image for each seed value
@@ -53,7 +62,9 @@ def main():
             image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata)
     except StopImageGenerationException as stop_exc:
         print(stop_exc)
-
+    finally:
+        if memory_saver:
+            print(memory_saver.memory_stats())
 
 if __name__ == "__main__":
     main()

src/mflux/ui/cli/parsers.py

@@ -34,6 +34,9 @@ class CommandLineParser(argparse.ArgumentParser):
         self.supports_image_to_image = False
         self.supports_lora = False
 
+    def add_general_arguments(self) -> None:
+        self.add_argument("--low-ram", action="store_true", help="Enable low-RAM mode to reduce memory usage (may impact performance).")
+
     def add_model_arguments(self, path_type: t.Literal["load", "save"] = "load", require_model_arg: bool = True) -> None:
 
         self.add_argument("--model", "-m", type=str, required=require_model_arg, action=ModelSpecAction, help=f"The model to use ({' or '.join(ui_defaults.MODEL_CHOICES)} or a compatible huggingface repo_id org/model).")
@@ -187,4 +190,7 @@ class CommandLineParser(argparse.ArgumentParser):
         if self.supports_image_generation and namespace.steps is None:
             namespace.steps = ui_defaults.MODEL_INFERENCE_STEPS.get(namespace.model, None)
 
+        if namespace.low_ram and len(namespace.seed) > 1:
+            self.error("--low-ram cannot be used with multiple seeds")
+
         return namespace