Add cache-limit option (#351)

Co-authored-by: Cursor Assistant <assistant@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 3 people

src/mflux/callbacks/callback_manager.py

@@ -1,5 +1,7 @@
 from argparse import Namespace
 
+import mlx.core as mx
+
 from mflux.callbacks.instances.battery_saver import BatterySaver
 from mflux.callbacks.instances.canny_saver import CannyImageSaver
 from mflux.callbacks.instances.depth_saver import DepthImageSaver
@@ -62,12 +64,28 @@ def _register_stepwise_handler(args: Namespace, model, latent_creator) -> None:
     @staticmethod
     def _register_memory_saver(args: Namespace, model) -> MemorySaver | None:
         memory_saver = None
+        cache_limit_bytes = CallbackManager._resolve_cache_limit_bytes(getattr(args, "mlx_cache_limit_gb", None))
         if args.low_ram:
             seeds = getattr(args, "seed", []) or []
             images = getattr(args, "image_path", [])
             if not isinstance(images, list):
                 images = [images] if images is not None else []
             keep_transformer = len(seeds) > 1 or len(images) > 1
-            memory_saver = MemorySaver(model=model, keep_transformer=keep_transformer, args=args)
+            memory_saver = MemorySaver(
+                model=model,
+                keep_transformer=keep_transformer,
+                cache_limit_bytes=cache_limit_bytes or 1000**3,
+                args=args,
+            )
             model.callbacks.register(memory_saver)
+        elif cache_limit_bytes is not None:
+            mx.set_cache_limit(cache_limit_bytes)
+            mx.clear_cache()
+            mx.reset_peak_memory()
         return memory_saver
+
+    @staticmethod
+    def _resolve_cache_limit_bytes(mlx_cache_limit_gb: float | None) -> int | None:
+        if mlx_cache_limit_gb is None:
+            return None
+        return int(mlx_cache_limit_gb * (1000**3))

src/mflux/cli/parser/parsers.py

@@ -35,6 +35,16 @@ def int_or_special_value(value) -> int | scale_factor.ScaleFactor:
         )
 
 
+def positive_float(value: str) -> float:
+    try:
+        parsed = float(value)
+    except ValueError:
+        raise argparse.ArgumentTypeError(f"'{value}' is not a valid number")
+    if parsed <= 0:
+        raise argparse.ArgumentTypeError(f"'{value}' must be > 0")
+    return parsed
+
+
 # fmt: off
 class CommandLineParser(argparse.ArgumentParser):
 
@@ -52,6 +62,7 @@ def __init__(self, *args, **kwargs):
     def add_general_arguments(self) -> None:
         self.add_argument("--battery-percentage-stop-limit", "-B", type=lambda v: max(min(int(v), 99), 1), default=ui_defaults.BATTERY_PERCENTAGE_STOP_LIMIT, help=f"On Macs powered by battery, stop image generation when battery reaches this percentage. Default: {ui_defaults.BATTERY_PERCENTAGE_STOP_LIMIT}")
         self.add_argument("--low-ram", action="store_true", help="Enable low-RAM mode to reduce memory usage (may impact performance).")
+        self.add_argument("--mlx-cache-limit-gb", type=positive_float, default=None, help="Limit MLX cache size in GB without enabling full low-RAM mode (e.g. 8 or 16).")
 
     def add_seedvr2_upscale_arguments(self) -> None:
         self.supports_image_generation = True

src/mflux/models/common/README.md

@@ -18,6 +18,7 @@ This README covers stable, shared patterns. For model-specific usage, see each m
 - [Metadata reuse](#metadata-reuse)
 - [Metadata inspection](#metadata-inspection)
 - [Resource and inspection options](#resource-and-inspection-options)
+- [MLX cache limit](#mlx-cache-limit)
 - [Cache locations](#cache-locations)
 - [Capabilities by task](#capabilities-by-task)
 - [Command reference](#command-reference)
@@ -522,6 +523,24 @@ image.save("image.png")
 
 ---
 
+## MLX cache limit
+
+Use `--mlx-cache-limit-gb` to cap MLX cache usage.
+
+```sh
+mflux-generate-z-image-turbo \
+  --model z-image-turbo \
+  --steps 9 \
+  --prompt "a portrait" \
+  --mlx-cache-limit-gb 2.5
+```
+
+- Value must be positive (`> 0`).
+- The value is converted internally using decimal gigabytes (`GB * 1000^3`).
+- Works in both normal mode and `--low-ram` mode.
+
+---
+
 ## Cache locations
 
 - **MFLUX cache**: set `MFLUX_CACHE_DIR` to override the default (`~/Library/Caches/mflux` on macOS, `~/.cache/mflux` on Linux).

tests/arg_parser/test_mlx_cache_limit_argparser.py

@@ -0,0 +1,26 @@
+from unittest.mock import patch
+
+import pytest
+
+from mflux.cli.parser.parsers import CommandLineParser
+
+
+@pytest.fixture
+def parser() -> CommandLineParser:
+    parser = CommandLineParser(description="Parser for MLX cache limit flag tests.")
+    parser.add_general_arguments()
+    return parser
+
+
+@pytest.mark.fast
+def test_mlx_cache_limit_gb_parses(parser: CommandLineParser):
+    with patch("sys.argv", ["mflux-generate", "--mlx-cache-limit-gb", "8"]):
+        args = parser.parse_args()
+        assert args.mlx_cache_limit_gb == 8.0
+
+
+@pytest.mark.fast
+def test_mlx_cache_limit_gb_rejects_non_positive_value(parser: CommandLineParser):
+    with patch("sys.argv", ["mflux-generate", "--mlx-cache-limit-gb", "0"]):
+        with pytest.raises(SystemExit):
+            parser.parse_args()

tests/callbacks/test_callback_manager_cache_limit.py

@@ -0,0 +1,41 @@
+from argparse import Namespace
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from mflux.callbacks.callback_manager import CallbackManager
+
+
+@pytest.mark.fast
+def test_register_memory_saver_sets_mlx_cache_limit_without_low_ram():
+    args = Namespace(low_ram=False, mlx_cache_limit_gb=2.5)
+    model = SimpleNamespace(callbacks=MagicMock())
+
+    with (
+        patch("mflux.callbacks.callback_manager.mx.set_cache_limit") as mock_set_cache_limit,
+        patch("mflux.callbacks.callback_manager.mx.clear_cache") as mock_clear_cache,
+        patch("mflux.callbacks.callback_manager.mx.reset_peak_memory") as mock_reset_peak_memory,
+    ):
+        memory_saver = CallbackManager._register_memory_saver(args=args, model=model)
+
+    assert memory_saver is None
+    mock_set_cache_limit.assert_called_once_with(int(2.5 * (1000**3)))
+    mock_clear_cache.assert_called_once()
+    mock_reset_peak_memory.assert_called_once()
+    model.callbacks.register.assert_not_called()
+
+
+@pytest.mark.fast
+def test_register_memory_saver_uses_mlx_cache_limit_for_low_ram_mode():
+    args = Namespace(low_ram=True, mlx_cache_limit_gb=3.0, seed=[42], image_path=None)
+    model = SimpleNamespace(callbacks=MagicMock())
+    mocked_memory_saver = object()
+
+    with patch("mflux.callbacks.callback_manager.MemorySaver", return_value=mocked_memory_saver) as mock_memory_saver:
+        memory_saver = CallbackManager._register_memory_saver(args=args, model=model)
+
+    assert memory_saver is mocked_memory_saver
+    _, kwargs = mock_memory_saver.call_args
+    assert kwargs["cache_limit_bytes"] == int(3.0 * (1000**3))
+    model.callbacks.register.assert_called_once_with(mocked_memory_saver)