diff --git a/README.md b/README.md index 19f783e..8fdc598 100644 --- a/README.md +++ b/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. +
parameters supported by config files diff --git a/src/mflux/callbacks/callback.py b/src/mflux/callbacks/callback.py index 7263b77..fe428ca 100644 --- a/src/mflux/callbacks/callback.py +++ b/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, diff --git a/src/mflux/callbacks/callback_registry.py b/src/mflux/callbacks/callback_registry.py index 8fd3974..c2caadd 100644 --- a/src/mflux/callbacks/callback_registry.py +++ b/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 diff --git a/src/mflux/callbacks/callbacks.py b/src/mflux/callbacks/callbacks.py index 10ad3b8..ba30e76 100644 --- a/src/mflux/callbacks/callbacks.py +++ b/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, diff --git a/src/mflux/callbacks/instances/memory_saver.py b/src/mflux/callbacks/instances/memory_saver.py new file mode 100644 index 0000000..e8c1a95 --- /dev/null +++ b/src/mflux/callbacks/instances/memory_saver.py @@ -0,0 +1,72 @@ +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 + gc.collect() + mx.metal.clear_cache() + + def memory_stats(self) -> str: + return f"Peak MLX memory: {self.peak_memory / 10**9:.2f} GB" diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index ad996b2..b1426c0 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -122,6 +122,14 @@ class Flux1(nn.Module): time_steps=time_steps, ) + # (Optional) Call subscribes 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) diff --git a/src/mflux/generate.py b/src/mflux/generate.py index d3a1cfe..c1c0589 100644 --- a/src/mflux/generate.py +++ b/src/mflux/generate.py @@ -2,11 +2,13 @@ from mflux import Config, Flux1, ModelConfig, StopImageGenerationException from mflux.callbacks.callback_registry import CallbackRegistry from mflux.callbacks.instances.stepwise_handler import StepwiseHandler from mflux.ui.cli.parsers import CommandLineParser +from mflux.callbacks.instances.memory_saver import MemorySaver 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__": diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index 8135ac2..e6ec546 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/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