Repository: filipstrand/mflux
Pull request #189: 🪫 Battery Saver - manage power over long batch operations
Author: login: anthonywu
Created: 2025-05-12T20:45:00Z | Merged: 2025-05-19T15:13:41Z

When using `--auto-seeds` or any batch operation that can deplete your MacBook {Air,Pro} battery over time, add some reasonable defaults and user flags to boundary the image generation to prevent the laptop from undesirable shut downs. This is a useful boundary if you rely on your Mac being remotely accessible by ssh or Screen Sharing.

# reference

This is the battery charging history for a M4 Max running mflux continuously while plugged into a 95W charger.

![image](https://github.com/user-attachments/assets/ba356592-b6b4-42aa-90a7-b049167e3f26)



Diff:
diff --git a/README.md b/README.md
index 1b7b17049..d399e000c 100644
--- a/README.md
+++ b/README.md
@@ -199,6 +199,8 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
 
 - **`--low-ram`** (optional): Reduces GPU memory usage by limiting the MLX cache size and releasing text encoders and transformer components after use (single image generation only). While this may slightly decrease performance, it helps prevent system memory swapping to disk, allowing image generation on systems with limited RAM.
 
+- **`--battery-percentage-stop-limit`** or **`-B`** (optional, `int`, default: `5`): On Mac laptops powered by battery, automatically stops image generation when battery percentage reaches this threshold. Prevents your Mac from shutting down and becoming unresponsive during long generation sessions.
+
 - **`--lora-name`** (optional, `str`, default: `None`): The name of the LoRA to download from Hugging Face.
 
 - **`--lora-repo-id`** (optional, `str`, default: `"ali-vilab/In-Context-LoRA"`): The Hugging Face repository ID for LoRAs.
@@ -1236,6 +1238,7 @@ See `uv run tools/rename_images.py --help` for full CLI usage help.
   - shortcut for dev model: `alias mflux-dev='mflux-generate --model dev'`
   - shortcut for schnell model *and* always save metadata: `alias mflux-schnell='mflux-generate --model schnell --metadata'`
 - For systems with limited memory, use the `--low-ram` flag to reduce memory usage by constraining the MLX cache size and releasing components after use
+- On battery-powered Macs, use `--battery-percentage-stop-limit` (or `-B`) to prevent your laptop from shutting down during long generation sessions
 - When generating multiple images with different seeds, use `--seed` with multiple values or `--auto-seeds` to automatically generate a series of random seeds
 - Use `--stepwise-image-output-dir` to save intermediate images at each denoising step, which can be useful for debugging or creating animations of the generation process
 
diff --git a/src/mflux/callbacks/instances/battery_saver.py b/src/mflux/callbacks/instances/battery_saver.py
new file mode 100644
index 000000000..f77333416
--- /dev/null
+++ b/src/mflux/callbacks/instances/battery_saver.py
@@ -0,0 +1,63 @@
+import json
+import logging
+import re
+import subprocess
+
+
+from mflux.callbacks.callback import BeforeLoopCallback
+from mflux.error.exceptions import StopImageGenerationException
+
+PMSET_AC_POWER_STATUS = "Now drawing from 'AC Power'"
+PMSET_BATT_STATUS_PATTERN = r"InternalBattery-.+?(\d+)%"
+
+logger = logging.getLogger(__name__)
+
+
+def _get_machine_model() -> str:
+    """Get the Mac machine model using system_profiler."""
+    try:
+        result = subprocess.run(
+            ["system_profiler", "-json", "SPHardwareDataType"], capture_output=True, text=True, check=True
+        )
+        data = json.loads(result.stdout)
+        return data["SPHardwareDataType"][0]["machine_model"]
+    except (subprocess.CalledProcessError, json.JSONDecodeError, IndexError, KeyError) as e:
+        logger.warning(f"Cannot determine machine model via 'system_profiler -json SPHardwareDataType': {e}")
+        return "Unknown"
+
+
+MACHINE_MODEL = _get_machine_model()
+# assumption: all Apple Silicon models powered by battery are "MacBook"s
+MACHINE_IS_BATTERY_POWERED = "MacBook" in MACHINE_MODEL
+
+
+def get_battery_percentage() -> int | None:
+    """Get the current battery percentage of a battery-powered Mac.
+    Returns None if Mac is not a battery-powered machine."""
+    if not MACHINE_IS_BATTERY_POWERED:
+        return None
+    percentage = None
+    try:
+        # running the subprocess would be expensive in a tight loop
+        # but in mflux use case, we would call this only once every
+        # few minutes due to N-minutes-long generation times
+        result = subprocess.run(["pmset", "-g", "batt"], capture_output=True, text=True, check=True)
+        if PMSET_AC_POWER_STATUS not in result.stdout:
+            if match := re.search(PMSET_BATT_STATUS_PATTERN, result.stdout):
+                percentage = int(match.group(1))
+    except (subprocess.CalledProcessError, TypeError) as e:
+        logger.warning(
+            f"Cannot read battery percentage via 'pmset -g batt': {e}. Battery saver functionality is disabled and the program will continue running."
+        )
+
+    return percentage
+
+
+class BatterySaver(BeforeLoopCallback):
+    def __init__(self, battery_percentage_stop_limit=10):
+        self.limit = battery_percentage_stop_limit
+
+    def call_before_loop(self, **kwargs) -> None:
+        current_pct: int | None = get_battery_percentage()
+        if current_pct is not None and current_pct <= self.limit:
+            raise StopImageGenerationException(f"Battery below {self.limit}% threshold: {current_pct}%")
diff --git a/src/mflux/generate.py b/src/mflux/generate.py
index dbd77657d..e9923d3f7 100644
--- a/src/mflux/generate.py
+++ b/src/mflux/generate.py
@@ -1,5 +1,8 @@
+from argparse import Namespace
+
 from mflux import Config, Flux1, ModelConfig, StopImageGenerationException
 from mflux.callbacks.callback_registry import CallbackRegistry
+from mflux.callbacks.instances.battery_saver import BatterySaver
 from mflux.callbacks.instances.memory_saver import MemorySaver
 from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
 from mflux.ui.cli.parsers import CommandLineParser
@@ -25,19 +28,8 @@ def main():
         lora_scales=args.lora_scales,
     )
 
-    # 2. Register the optional callbacks
-    if args.stepwise_image_output_dir:
-        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
-        CallbackRegistry.register_before_loop(handler)
-        CallbackRegistry.register_in_loop(handler)
-        CallbackRegistry.register_interrupt(handler)
-
-    memory_saver = None
-    if args.low_ram:
-        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
-        CallbackRegistry.register_before_loop(memory_saver)
-        CallbackRegistry.register_in_loop(memory_saver)
-        CallbackRegistry.register_after_loop(memory_saver)
+    # 2. Register callbacks
+    memory_saver = _register_callbacks(args=args, flux=flux)
 
     try:
         for seed in args.seed:
@@ -63,5 +55,27 @@ def main():
             print(memory_saver.memory_stats())
 
 
+def _register_callbacks(args: Namespace, flux: Flux1) -> MemorySaver | None:
+    # Battery saver
+    battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit)
+    CallbackRegistry.register_before_loop(battery_saver)
+
+    # Stepwise Handler
+    if args.stepwise_image_output_dir:
+        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
+        CallbackRegistry.register_before_loop(handler)
+        CallbackRegistry.register_in_loop(handler)
+        CallbackRegistry.register_interrupt(handler)
+
+    # Memory Saver
+    memory_saver = None
+    if args.low_ram:
+        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
+        CallbackRegistry.register_before_loop(memory_saver)
+        CallbackRegistry.register_in_loop(memory_saver)
+        CallbackRegistry.register_after_loop(memory_saver)
+    return memory_saver
+
+
 if __name__ == "__main__":
     main()
diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py
index 41dfd79b7..6bb3536c1 100644
--- a/src/mflux/generate_controlnet.py
+++ b/src/mflux/generate_controlnet.py
@@ -1,5 +1,8 @@
+from argparse import Namespace
+
 from mflux import Config, Flux1Controlnet, ModelConfig, StopImageGenerationException
 from mflux.callbacks.callback_registry import CallbackRegistry
+from mflux.callbacks.instances.battery_saver import BatterySaver
 from mflux.callbacks.instances.canny_saver import CannyImageSaver
 from mflux.callbacks.instances.memory_saver import MemorySaver
 from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
@@ -26,21 +29,8 @@ def main():
         lora_scales=args.lora_scales,
     )
 
-    # 2. Register the optional callbacks
-    if args.controlnet_save_canny:
-        CallbackRegistry.register_before_loop(CannyImageSaver(path=args.output))
-    if args.stepwise_image_output_dir:
-        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
-        CallbackRegistry.register_before_loop(handler)
-        CallbackRegistry.register_in_loop(handler)
-        CallbackRegistry.register_interrupt(handler)
-
-    memory_saver = None
-    if args.low_ram:
-        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
-        CallbackRegistry.register_before_loop(memory_saver)
-        CallbackRegistry.register_in_loop(memory_saver)
-        CallbackRegistry.register_after_loop(memory_saver)
+    # 2. Register callbacks
+    memory_saver = _register_callbacks(args=args, flux=flux)
 
     try:
         for seed in args.seed:
@@ -67,5 +57,32 @@ def main():
             print(memory_saver.memory_stats())
 
 
+def _register_callbacks(args: Namespace, flux: Flux1Controlnet) -> MemorySaver | None:
+    # Battery saver
+    battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit)
+    CallbackRegistry.register_before_loop(battery_saver)
+
+    # Canny Image Saver
+    if args.controlnet_save_canny:
+        canny_image_saver = CannyImageSaver(path=args.output)
+        CallbackRegistry.register_before_loop(canny_image_saver)
+
+    # Stepwise Handler
+    if args.stepwise_image_output_dir:
+        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
+        CallbackRegistry.register_before_loop(handler)
+        CallbackRegistry.register_in_loop(handler)
+        CallbackRegistry.register_interrupt(handler)
+
+    # Memory Saver
+    memory_saver = None
+    if args.low_ram:
+        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
+        CallbackRegistry.register_before_loop(memory_saver)
+        CallbackRegistry.register_in_loop(memory_saver)
+        CallbackRegistry.register_after_loop(memory_saver)
+    return memory_saver
+
+
 if __name__ == "__main__":
     main()
diff --git a/src/mflux/generate_depth.py b/src/mflux/generate_depth.py
index 4860bc073..ca62a2981 100644
--- a/src/mflux/generate_depth.py
+++ b/src/mflux/generate_depth.py
@@ -1,5 +1,8 @@
+from argparse import Namespace
+
 from mflux import Config, StopImageGenerationException
 from mflux.callbacks.callback_registry import CallbackRegistry
+from mflux.callbacks.instances.battery_saver import BatterySaver
 from mflux.callbacks.instances.depth_saver import DepthImageSaver
 from mflux.callbacks.instances.memory_saver import MemorySaver
 from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
@@ -30,21 +33,8 @@ def main():
         lora_scales=args.lora_scales,
     )
 
-    # 2. Register the optional callbacks
-    if args.save_depth_map:
-        CallbackRegistry.register_before_loop(DepthImageSaver(path=args.output))
-    if args.stepwise_image_output_dir:
-        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
-        CallbackRegistry.register_before_loop(handler)
-        CallbackRegistry.register_in_loop(handler)
-        CallbackRegistry.register_interrupt(handler)
-
-    memory_saver = None
-    if args.low_ram:
-        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
-        CallbackRegistry.register_before_loop(memory_saver)
-        CallbackRegistry.register_in_loop(memory_saver)
-        CallbackRegistry.register_after_loop(memory_saver)
+    # 2. Register callbacks
+    memory_saver = _register_callbacks(args=args, flux=flux)
 
     try:
         for seed in args.seed:
@@ -71,5 +61,32 @@ def main():
             print(memory_saver.memory_stats())
 
 
+def _register_callbacks(args: Namespace, flux: Flux1Depth) -> MemorySaver | None:
+    # Battery saver
+    battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit)
+    CallbackRegistry.register_before_loop(battery_saver)
+
+    # Depth Image Saver
+    if args.save_depth_map:
+        depth_image_saver = DepthImageSaver(path=args.output)
+        CallbackRegistry.register_before_loop(depth_image_saver)
+
+    # Stepwise Handler
+    if args.stepwise_image_output_dir:
+        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
+        CallbackRegistry.register_before_loop(handler)
+        CallbackRegistry.register_in_loop(handler)
+        CallbackRegistry.register_interrupt(handler)
+
+    # Memory Saver
+    memory_saver = None
+    if args.low_ram:
+        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
+        CallbackRegistry.register_before_loop(memory_saver)
+        CallbackRegistry.register_in_loop(memory_saver)
+        CallbackRegistry.register_after_loop(memory_saver)
+    return memory_saver
+
+
 if __name__ == "__main__":
     main()
diff --git a/src/mflux/generate_fill.py b/src/mflux/generate_fill.py
index ae9e1313e..46ee3ce95 100644
--- a/src/mflux/generate_fill.py
+++ b/src/mflux/generate_fill.py
@@ -1,5 +1,8 @@
+from argparse import Namespace
+
 from mflux import Config, StopImageGenerationException
 from mflux.callbacks.callback_registry import CallbackRegistry
+from mflux.callbacks.instances.battery_saver import BatterySaver
 from mflux.callbacks.instances.memory_saver import MemorySaver
 from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
 from mflux.flux_tools.fill.flux_fill import Flux1Fill
@@ -29,19 +32,8 @@ def main():
         lora_scales=args.lora_scales,
     )
 
-    # 2. Register the optional callbacks
-    if args.stepwise_image_output_dir:
-        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
-        CallbackRegistry.register_before_loop(handler)
-        CallbackRegistry.register_in_loop(handler)
-        CallbackRegistry.register_interrupt(handler)
-
-    memory_saver = None
-    if args.low_ram:
-        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
-        CallbackRegistry.register_before_loop(memory_saver)
-        CallbackRegistry.register_in_loop(memory_saver)
-        CallbackRegistry.register_after_loop(memory_saver)
+    # 2. Register callbacks
+    memory_saver = _register_callbacks(args=args, flux=flux)
 
     try:
         for seed in args.seed:
@@ -68,5 +60,27 @@ def main():
             print(memory_saver.memory_stats())
 
 
+def _register_callbacks(args: Namespace, flux: Flux1Fill) -> MemorySaver | None:
+    # Battery saver
+    battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit)
+    CallbackRegistry.register_before_loop(battery_saver)
+
+    # Stepwise Handler
+    if args.stepwise_image_output_dir:
+        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
+        CallbackRegistry.register_before_loop(handler)
+        CallbackRegistry.register_in_loop(handler)
+        CallbackRegistry.register_interrupt(handler)
+
+    # Memory Saver
+    memory_saver = None
+    if args.low_ram:
+        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
+        CallbackRegistry.register_before_loop(memory_saver)
+        CallbackRegistry.register_in_loop(memory_saver)
+        CallbackRegistry.register_after_loop(memory_saver)
+    return memory_saver
+
+
 if __name__ == "__main__":
     main()
diff --git a/src/mflux/generate_in_context.py b/src/mflux/generate_in_context.py
index 8ad2cb948..9da70a8b2 100644
--- a/src/mflux/generate_in_context.py
+++ b/src/mflux/generate_in_context.py
@@ -1,7 +1,9 @@
+from argparse import Namespace
 from pathlib import Path
 
 from mflux import Config, StopImageGenerationException
 from mflux.callbacks.callback_registry import CallbackRegistry
+from mflux.callbacks.instances.battery_saver import BatterySaver
 from mflux.callbacks.instances.memory_saver import MemorySaver
 from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
 from mflux.community.in_context_lora.flux_in_context_lora import Flux1InContextLoRA
@@ -37,19 +39,8 @@ def main():
         lora_scales=args.lora_scales,
     )
 
-    # 2. Register the optional callbacks
-    if args.stepwise_image_output_dir:
-        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
-        CallbackRegistry.register_before_loop(handler)
-        CallbackRegistry.register_in_loop(handler)
-        CallbackRegistry.register_interrupt(handler)
-
-    memory_saver = None
-    if args.low_ram:
-        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
-        CallbackRegistry.register_before_loop(memory_saver)
-        CallbackRegistry.register_in_loop(memory_saver)
-        CallbackRegistry.register_after_loop(memory_saver)
+    # 2. Register callbacks
+    memory_saver = _register_callbacks(args=args, flux=flux)
 
     try:
         for seed in args.seed:
@@ -77,5 +68,27 @@ def main():
             print(memory_saver.memory_stats())
 
 
+def _register_callbacks(args: Namespace, flux: Flux1InContextLoRA) -> MemorySaver | None:
+    # Battery saver
+    battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit)
+    CallbackRegistry.register_before_loop(battery_saver)
+
+    # Stepwise Handler
+    if args.stepwise_image_output_dir:
+        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
+        CallbackRegistry.register_before_loop(handler)
+        CallbackRegistry.register_in_loop(handler)
+        CallbackRegistry.register_interrupt(handler)
+
+    # Memory Saver
+    memory_saver = None
+    if args.low_ram:
+        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
+        CallbackRegistry.register_before_loop(memory_saver)
+        CallbackRegistry.register_in_loop(memory_saver)
+        CallbackRegistry.register_after_loop(memory_saver)
+    return memory_saver
+
+
 if __name__ == "__main__":
     main()
diff --git a/src/mflux/generate_redux.py b/src/mflux/generate_redux.py
index b52ed3f22..af399a247 100644
--- a/src/mflux/generate_redux.py
+++ b/src/mflux/generate_redux.py
@@ -1,7 +1,10 @@
+from argparse import Namespace
+
 from pathlib import Path
 
 from mflux import Config, ModelConfig, StopImageGenerationException
 from mflux.callbacks.callback_registry import CallbackRegistry
+from mflux.callbacks.instances.battery_saver import BatterySaver
 from mflux.callbacks.instances.memory_saver import MemorySaver
 from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
 from mflux.flux_tools.redux.flux_redux import Flux1Redux
@@ -34,19 +37,8 @@ def main():
         lora_scales=args.lora_scales,
     )
 
-    # 2. Register the optional callbacks
-    if args.stepwise_image_output_dir:
-        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
-        CallbackRegistry.register_before_loop(handler)
-        CallbackRegistry.register_in_loop(handler)
-        CallbackRegistry.register_interrupt(handler)
-
-    memory_saver = None
-    if args.low_ram:
-        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
-        CallbackRegistry.register_before_loop(memory_saver)
-        CallbackRegistry.register_in_loop(memory_saver)
-        CallbackRegistry.register_after_loop(memory_saver)
+    # 2. Register callbacks
+    memory_saver = _register_callbacks(args=args, flux=flux)
 
     try:
         for seed in args.seed:
@@ -73,6 +65,28 @@ def main():
             print(memory_saver.memory_stats())
 
 
+def _register_callbacks(args: Namespace, flux: Flux1Redux) -> MemorySaver | None:
+    # Battery saver
+    battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit)
+    CallbackRegistry.register_before_loop(battery_saver)
+
+    # Stepwise Handler
+    if args.stepwise_image_output_dir:
+        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
+        CallbackRegistry.register_before_loop(handler)
+        CallbackRegistry.register_in_loop(handler)
+        CallbackRegistry.register_interrupt(handler)
+
+    # Memory Saver
+    memory_saver = None
+    if args.low_ram:
+        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
+        CallbackRegistry.register_before_loop(memory_saver)
+        CallbackRegistry.register_in_loop(memory_saver)
+        CallbackRegistry.register_after_loop(memory_saver)
+    return memory_saver
+
+
 def _validate_redux_image_strengths(
     redux_image_paths: list[Path],
     redux_image_strengths: list[float] | None,
diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py
index 343c00c09..87109e712 100644
--- a/src/mflux/ui/cli/parsers.py
+++ b/src/mflux/ui/cli/parsers.py
@@ -42,6 +42,7 @@ def __init__(self, *args, **kwargs):
         self.require_model_arg = True
 
     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).")
 
     def add_model_arguments(self, path_type: t.Literal["load", "save"] = "load", require_model_arg: bool = True) -> None:
diff --git a/src/mflux/ui/defaults.py b/src/mflux/ui/defaults.py
index 2f4a33cce..455a477e9 100644
--- a/src/mflux/ui/defaults.py
+++ b/src/mflux/ui/defaults.py
@@ -1,3 +1,4 @@
+BATTERY_PERCENTAGE_STOP_LIMIT = 5
 CONTROLNET_STRENGTH = 0.4
 GUIDANCE_SCALE = 3.5
 HEIGHT, WIDTH = 1024, 1024
diff --git a/tests/callbacks/__init__.py b/tests/callbacks/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/callbacks/test_battery_saver.py b/tests/callbacks/test_battery_saver.py
new file mode 100644
index 000000000..94acb4f9f
--- /dev/null
+++ b/tests/callbacks/test_battery_saver.py
@@ -0,0 +1,100 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from mflux.callbacks.instances.battery_saver import BatterySaver, get_battery_percentage
+from mflux.error.exceptions import StopImageGenerationException
+
+
+def test_get_battery_percentage_success():
+    """Test that the battery percentage is correctly extracted when subprocess returns a valid output."""
+    with patch("subprocess.run") as mock_run:
+        # Set up mock to return a sample battery status output
+        mock_result = MagicMock()
+        mock_result.stdout = "Now drawing from 'Battery Power'\n -InternalBattery-0 (id=1234567)	42%;\n"
+        mock_run.return_value = mock_result
+
+        # Call the function
+        percentage = get_battery_percentage()
+
+        # Verify subprocess.run was called with the correct arguments
+        mock_run.assert_called_once_with(["pmset", "-g", "batt"], capture_output=True, text=True, check=True)
+
+        # Assert the function properly extracted the battery percentage
+        assert percentage == 42
+
+
+def test_get_battery_percentage_while_charging():
+    """Test that the function returns None when the output doesn't match the expected pattern."""
+    with patch("subprocess.run") as mock_run:
+        # Set up mock to return an output that doesn't match the expected pattern
+        mock_result = MagicMock()
+        mock_result.stdout = "Now drawing from 'AC Power'"
+        mock_run.return_value = mock_result
+
+        # Call the function
+        percentage = get_battery_percentage()
+
+        # Assert the function returns None when no match is found
+        assert percentage is None
+
+
+def test_battery_saver_below_limit():
+    """Test that BatterySaver raises an exception when the battery is below the limit."""
+    with patch("mflux.callbacks.instances.battery_saver.get_battery_percentage") as mock_get:
+        # Configure mock to return a battery percentage below the limit
+        mock_get.return_value = 5
+
+        # Create a BatterySaver instance with a limit of 10%
+        battery_saver = BatterySaver(battery_percentage_stop_limit=10)
+
+        # Assert that calling before_loop raises StopImageGenerationException
+        with pytest.raises(StopImageGenerationException) as excinfo:
+            battery_saver.call_before_loop()
+
+        # Check the exception message contains the correct limit and percentage
+        assert "10%" in str(excinfo.value)
+        assert "5%" in str(excinfo.value)
+
+
+def test_battery_saver_above_limit():
+    """Test that BatterySaver does not raise an exception when the battery is above the limit."""
+    with patch("mflux.callbacks.instances.battery_saver.get_battery_percentage") as mock_get:
+        # Configure mock to return a battery percentage above the limit
+        mock_get.return_value = 20
+
+        # Create a BatterySaver instance with a limit of 10%
+        battery_saver = BatterySaver(battery_percentage_stop_limit=10)
+
+        # Assert that calling before_loop does not raise an exception
+        battery_saver.call_before_loop()
+
+
+def test_battery_saver_none_percentage():
+    """Test that BatterySaver does not raise an exception when percentage is None."""
+    with patch("mflux.callbacks.instances.battery_saver.get_battery_percentage") as mock_get:
+        # Configure mock to return None (e.g., on non-battery systems)
+        mock_get.return_value = None
+
+        # Create a BatterySaver instance
+        battery_saver = BatterySaver()
+
+        # Assert that calling before_loop does not raise an exception when percentage is None
+        battery_saver.call_before_loop()
+
+
+def test_battery_saver_equal_to_limit():
+    """Test that BatterySaver raises an exception when the battery equals the limit."""
+    with patch("mflux.callbacks.instances.battery_saver.get_battery_percentage") as mock_get:
+        # Configure mock to return a battery percentage equal to the limit
+        mock_get.return_value = 10
+
+        # Create a BatterySaver instance with a limit of 10%
+        battery_saver = BatterySaver(battery_percentage_stop_limit=10)
+
+        # Assert that calling before_loop raises StopImageGenerationException
+        with pytest.raises(StopImageGenerationException) as excinfo:
+            battery_saver.call_before_loop()
+
+        # Check the exception message contains the correct limit and percentage
+        assert "10%" in str(excinfo.value)


Reviews:
user: login: anthonywu
body: My observations on my M1 Max and M4 Max laptops, having run batches > 24 and > 48 hours.

1. on 100% GPU utilization and screen usage over a long batch run, a 85W or 96W battery charger does not keep up with energy consumption and I see gradual depletion of the battery while plugged in. (may have to do with power or heat management settings)
2. the biggest charger shipped with MBPs is 96W
3. most monitors provide power to the laptop around 90W, I've used or owned multiple Thunderbolt monitors (Apple Studio Display, LG, Dell Thunderbolt displays) popular with Mac users)

In any case, when you are running on battery-only, this is definitely an issue!
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T20:54:05Z
user: login: azrahello
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T06:18:12Z
user: login: anthonywu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-13T06:20:57Z
user: login: azrahello
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T06:32:38Z
user: login: anthonywu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-13T06:34:16Z
user: login: anthonywu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-13T18:49:52Z
user: login: filipstrand
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-19T14:24:47Z
user: login: filipstrand
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-19T14:26:07Z
user: login: filipstrand
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-19T14:50:29Z
user: login: filipstrand
body: Looks good!
state: APPROVED
author_association: OWNER
submitted_at: 2025-05-19T14:57:39Z

Review comments:
diff_hunk: @@ -0,0 +1,100 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from mflux.callbacks.instances.battery_saver import BatterySaver, get_battery_percentage
+from mflux.error.exceptions import StopImageGenerationException
+
+
+def test_get_battery_percentage_success():
+    """Test that the battery percentage is correctly extracted when subprocess returns a valid output."""
+    with patch("subprocess.run") as mock_run:
+        # Set up mock to return a sample battery status output
+        mock_result = MagicMock()
+        mock_result.stdout = "Now drawing from 'Battery Power'\n -InternalBattery-0 (id=1234567)	42%;\n"
+        mock_run.return_value = mock_result
+
+        # Call the function
+        percentage = get_battery_percentage()
+
+        # Verify subprocess.run was called with the correct arguments
+        mock_run.assert_called_once_with(["pmset", "-g", "batt"], capture_output=True, text=True, check=True)
+
+        # Assert the function properly extracted the battery percentage
+        assert percentage == 42
+
+
+def test_get_battery_percentage_no_match():
+    """Test that the function returns None when the output doesn't match the expected pattern."""
+    with patch("subprocess.run") as mock_run:
+        # Set up mock to return an output that doesn't match the expected pattern
+        mock_result = MagicMock()
+        mock_result.stdout = "Now drawing from 'AC Power'\n-InternalBattery-0 (id=1234567)    No battery installed"
path: tests/callbacks/test_battery_saver.py
user: login: anthonywu
body: may need someone with a Mac mini / studio to confirm - I am pretty sure "No battery installed" is in the string but the rest of the return status string is my guess.
created_at: 2025-05-12T20:47:35Z
author_association: COLLABORATOR
diff_hunk: @@ -1,3 +1,4 @@
+BATTERY_PERCENTAGE_STOP_LIMIT = 5
path: src/mflux/ui/defaults.py
user: login: anthonywu
body: I am open to suggestions to lower this to like 1 or 2, given we are adding the callback by default.
created_at: 2025-05-12T20:50:58Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,100 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from mflux.callbacks.instances.battery_saver import BatterySaver, get_battery_percentage
+from mflux.error.exceptions import StopImageGenerationException
+
+
+def test_get_battery_percentage_success():
+    """Test that the battery percentage is correctly extracted when subprocess returns a valid output."""
+    with patch("subprocess.run") as mock_run:
+        # Set up mock to return a sample battery status output
+        mock_result = MagicMock()
+        mock_result.stdout = "Now drawing from 'Battery Power'\n -InternalBattery-0 (id=1234567)	42%;\n"
+        mock_run.return_value = mock_result
+
+        # Call the function
+        percentage = get_battery_percentage()
+
+        # Verify subprocess.run was called with the correct arguments
+        mock_run.assert_called_once_with(["pmset", "-g", "batt"], capture_output=True, text=True, check=True)
+
+        # Assert the function properly extracted the battery percentage
+        assert percentage == 42
+
+
+def test_get_battery_percentage_no_match():
+    """Test that the function returns None when the output doesn't match the expected pattern."""
+    with patch("subprocess.run") as mock_run:
+        # Set up mock to return an output that doesn't match the expected pattern
+        mock_result = MagicMock()
+        mock_result.stdout = "Now drawing from 'AC Power'\n-InternalBattery-0 (id=1234567)    No battery installed"
path: tests/callbacks/test_battery_saver.py
user: login: azrahello
body: i'have a Mac studio, (and a Macbook pro m1) ,  how can help? 
created_at: 2025-05-13T06:18:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,100 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from mflux.callbacks.instances.battery_saver import BatterySaver, get_battery_percentage
+from mflux.error.exceptions import StopImageGenerationException
+
+
+def test_get_battery_percentage_success():
+    """Test that the battery percentage is correctly extracted when subprocess returns a valid output."""
+    with patch("subprocess.run") as mock_run:
+        # Set up mock to return a sample battery status output
+        mock_result = MagicMock()
+        mock_result.stdout = "Now drawing from 'Battery Power'\n -InternalBattery-0 (id=1234567)	42%;\n"
+        mock_run.return_value = mock_result
+
+        # Call the function
+        percentage = get_battery_percentage()
+
+        # Verify subprocess.run was called with the correct arguments
+        mock_run.assert_called_once_with(["pmset", "-g", "batt"], capture_output=True, text=True, check=True)
+
+        # Assert the function properly extracted the battery percentage
+        assert percentage == 42
+
+
+def test_get_battery_percentage_no_match():
+    """Test that the function returns None when the output doesn't match the expected pattern."""
+    with patch("subprocess.run") as mock_run:
+        # Set up mock to return an output that doesn't match the expected pattern
+        mock_result = MagicMock()
+        mock_result.stdout = "Now drawing from 'AC Power'\n-InternalBattery-0 (id=1234567)    No battery installed"
path: tests/callbacks/test_battery_saver.py
user: login: anthonywu
body: Show me the output of terminal command "pmset -g batt" and paste me the exact strings. I'll add to unit test. In your case this feature should never activate for you.
created_at: 2025-05-13T06:20:57Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,100 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from mflux.callbacks.instances.battery_saver import BatterySaver, get_battery_percentage
+from mflux.error.exceptions import StopImageGenerationException
+
+
+def test_get_battery_percentage_success():
+    """Test that the battery percentage is correctly extracted when subprocess returns a valid output."""
+    with patch("subprocess.run") as mock_run:
+        # Set up mock to return a sample battery status output
+        mock_result = MagicMock()
+        mock_result.stdout = "Now drawing from 'Battery Power'\n -InternalBattery-0 (id=1234567)	42%;\n"
+        mock_run.return_value = mock_result
+
+        # Call the function
+        percentage = get_battery_percentage()
+
+        # Verify subprocess.run was called with the correct arguments
+        mock_run.assert_called_once_with(["pmset", "-g", "batt"], capture_output=True, text=True, check=True)
+
+        # Assert the function properly extracted the battery percentage
+        assert percentage == 42
+
+
+def test_get_battery_percentage_no_match():
+    """Test that the function returns None when the output doesn't match the expected pattern."""
+    with patch("subprocess.run") as mock_run:
+        # Set up mock to return an output that doesn't match the expected pattern
+        mock_result = MagicMock()
+        mock_result.stdout = "Now drawing from 'AC Power'\n-InternalBattery-0 (id=1234567)    No battery installed"
path: tests/callbacks/test_battery_saver.py
user: login: azrahello
body: > Show me the output of terminal command "pmset -g batt" and paste me the exact strings. I'll add to unit test. In your case this feature should never activate for you.

(base) @Alessandros-Mac-Studio ~ % pmset -g batt 
Now drawing from 'AC Power'
(base)  @Alessandros-Mac-Studio ~ % 

created_at: 2025-05-13T06:32:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,100 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from mflux.callbacks.instances.battery_saver import BatterySaver, get_battery_percentage
+from mflux.error.exceptions import StopImageGenerationException
+
+
+def test_get_battery_percentage_success():
+    """Test that the battery percentage is correctly extracted when subprocess returns a valid output."""
+    with patch("subprocess.run") as mock_run:
+        # Set up mock to return a sample battery status output
+        mock_result = MagicMock()
+        mock_result.stdout = "Now drawing from 'Battery Power'\n -InternalBattery-0 (id=1234567)	42%;\n"
+        mock_run.return_value = mock_result
+
+        # Call the function
+        percentage = get_battery_percentage()
+
+        # Verify subprocess.run was called with the correct arguments
+        mock_run.assert_called_once_with(["pmset", "-g", "batt"], capture_output=True, text=True, check=True)
+
+        # Assert the function properly extracted the battery percentage
+        assert percentage == 42
+
+
+def test_get_battery_percentage_no_match():
+    """Test that the function returns None when the output doesn't match the expected pattern."""
+    with patch("subprocess.run") as mock_run:
+        # Set up mock to return an output that doesn't match the expected pattern
+        mock_result = MagicMock()
+        mock_result.stdout = "Now drawing from 'AC Power'\n-InternalBattery-0 (id=1234567)    No battery installed"
path: tests/callbacks/test_battery_saver.py
user: login: anthonywu
body: Interesting. No mention of batteries at all. Will add to test case in next commit. Thanks!
created_at: 2025-05-13T06:34:15Z
author_association: COLLABORATOR
diff_hunk: @@ -1180,6 +1182,7 @@ See `uv run tools/rename_images.py --help` for full CLI usage help.
   - shortcut for dev model: `alias mflux-dev='mflux-generate --model dev'`
path: README.md
user: login: anthonywu
body: I don't know why this diff is showing up so ugly in GH PR review. But it's adding a list item.
created_at: 2025-05-13T18:30:09Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,40 @@
+import logging
+import re
+import subprocess
+
+from mflux.callbacks.callback import BeforeLoopCallback
+from mflux.error.exceptions import StopImageGenerationException
+
+PMSET_AC_POWER_STATUS = "Now drawing from 'AC Power'"
+PMSET_BATT_STATUS_PATTERN = r"InternalBattery-.+?(\d+)%"
+
+logger = logging.getLogger(__name__)
+
+
+def get_battery_percentage() -> int | None:
+    """Get the current battery percentage of a Mac notebook."""
+    percentage = None
+    try:
+        # running the subprocess would be expensive in a tight loop
+        # but in mflux use case, we would call this only once every
+        # few minutes due to N-minutes-long generation times
+        result = subprocess.run(["pmset", "-g", "batt"], capture_output=True, text=True, check=True)
path: src/mflux/callbacks/instances/battery_saver.py
user: login: anthonywu
body: I had considered using the https://pypi.org/project/psutil/ lib, but the surface area is too small to merit adding the dependency, so going with built-in libs is the way to go for now.
created_at: 2025-05-13T18:31:00Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,63 @@
+import json
+import logging
+import re
+import subprocess
+
+
+from mflux.callbacks.callback import BeforeLoopCallback
+from mflux.error.exceptions import StopImageGenerationException
+
+PMSET_AC_POWER_STATUS = "Now drawing from 'AC Power'"
+PMSET_BATT_STATUS_PATTERN = r"InternalBattery-.+?(\d+)%"
+
+logger = logging.getLogger(__name__)
+
+
+def _get_machine_model() -> str:
+    """Get the Mac machine model using system_profiler."""
+    try:
+        result = subprocess.run(
+            ["system_profiler", "-json", "SPHardwareDataType"], capture_output=True, text=True, check=True
path: src/mflux/callbacks/instances/battery_saver.py
user: login: anthonywu
body: profiler data may be interesting going forward for creating a benchmarking script so that we can all share run stats:

Here's mine with relevant fields:

```json
{
  "SPHardwareDataType" : [
    {
      "chip_type" : "Apple M1 Max",
      "machine_model" : "MacBookPro18,4",
      "machine_name" : "MacBook Pro",
      "number_processors" : "proc 10:8:2",
      "physical_memory" : "64 GB",
    }
  ]
}
```
created_at: 2025-05-13T18:49:00Z
author_association: COLLABORATOR
diff_hunk: @@ -25,7 +26,15 @@ def main():
         lora_scales=args.lora_scales,
     )
 
-    # 2. Register the optional callbacks
+    # 2. Register callbacks
+
+    # Add BatterySaver callback unconditionally
path: src/mflux/generate.py
user: login: anthonywu
body: "unconditionally" but should be no-cost to desktops, see latest commit on excluding AC powered machines from the check.
created_at: 2025-05-13T18:49:41Z
author_association: COLLABORATOR
diff_hunk: @@ -1,3 +1,4 @@
+BATTERY_PERCENTAGE_STOP_LIMIT = 5
path: src/mflux/ui/defaults.py
user: login: filipstrand
body: I think this works fine! On older machines (e.g my older pre apple-silicon Macs) I noticed that as they get older they might report a 7% battery status, then suddenly die within the next 2 minutes or so if under any kind of load, so it might actually be good to have some extra margin here. 
created_at: 2025-05-19T14:24:47Z
author_association: OWNER
diff_hunk: @@ -25,7 +26,15 @@ def main():
         lora_scales=args.lora_scales,
     )
 
-    # 2. Register the optional callbacks
+    # 2. Register callbacks
+
+    # Add BatterySaver callback unconditionally
path: src/mflux/generate.py
user: login: filipstrand
body: Nice!
created_at: 2025-05-19T14:26:07Z
author_association: OWNER
diff_hunk: @@ -0,0 +1,63 @@
+import json
+import logging
+import re
+import subprocess
+
+
+from mflux.callbacks.callback import BeforeLoopCallback
+from mflux.error.exceptions import StopImageGenerationException
+
+PMSET_AC_POWER_STATUS = "Now drawing from 'AC Power'"
+PMSET_BATT_STATUS_PATTERN = r"InternalBattery-.+?(\d+)%"
+
+logger = logging.getLogger(__name__)
+
+
+def _get_machine_model() -> str:
+    """Get the Mac machine model using system_profiler."""
+    try:
+        result = subprocess.run(
+            ["system_profiler", "-json", "SPHardwareDataType"], capture_output=True, text=True, check=True
path: src/mflux/callbacks/instances/battery_saver.py
user: login: filipstrand
body: Good idea! Right now I think we only mention the command `system_profiler SPHardwareDataType SPDisplaysDataType` in the readme for the benchmark section, but it would be convenient to have a separate script that does all of this at once (including the generation itself) 
created_at: 2025-05-19T14:50:29Z
author_association: OWNER

PR comments:
user: login: filipstrand
created_at: 2025-05-19T13:57:21Z
body: Interesting! Never thought of adding this as a feature but it makes much sense to have (I am running MacBooks exclusivley and running out of juice has definitely happened hehe). 
author_association: OWNER
user: login: filipstrand
created_at: 2025-05-19T14:28:37Z
body: Now that we have an increasing number of callbacks being registered at startup, I moved these to a separate method (still duplicated across the different entry-point files for now).
author_association: OWNER
user: login: filipstrand
created_at: 2025-05-19T14:56:26Z
body: I don't think I will have the time to test this feature for real, but I assume it is working as we expect and the PR looks good to me! I'm happy to merge this to main anytime. @anthonywu Let me know if there is anything more to test/add or if it is ready to be merged.
author_association: OWNER