Add upscaling support (#184)

README.md

@@ -37,6 +37,7 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
   * [🔍 Depth](#-depth)
   * [🔄 Redux](#-redux)
 - [🕹️ Controlnet](#%EF%B8%8F-controlnet)
+- [🔎 Upscale](#-upscale)
 - [🎛️ Dreambooth fine-tuning](#-dreambooth-fine-tuning)
   * [Training configuration](#training-configuration)
   * [Training example](#training-example)
@@ -207,6 +208,10 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
 
 - **`--stepwise-image-output-dir`** (optional, `str`, default: `None`): [EXPERIMENTAL] Output directory to write step-wise images and their final composite image to. This feature may change in future versions. When specified, MFLUX will save an image for each denoising step, allowing you to visualize the generation process from noise to final image.
 
+- **`--vae-tiling`** (optional, flag): Enable VAE tiling to reduce memory usage during the decoding phase. This splits the image into smaller chunks for processing, which can prevent out-of-memory errors when generating high-resolution images. Note that this optimization may occasionally produce a subtle seam in the middle of the image, but it's often worth the tradeoff for being able to generate images that would otherwise cause your system to run out of memory.
+
+- **`--vae-tiling-split`** (optional, `str`, default: `"horizontal"`): When VAE tiling is enabled, this parameter controls the direction to split the latents. Options are `"horizontal"` (splits into top/bottom) or `"vertical"` (splits into left/right). Use this option to control where potential seams might appear in the final image.
+
 #### 📜 In-Context LoRA Command-Line Arguments
 
 The `mflux-generate-in-context` command supports most of the same arguments as `mflux-generate`, with these additional parameters:
@@ -1082,6 +1087,41 @@ with different prompts and LoRA adapters active.
 
 ---
 
+### 🔎 Upscale
+
+The upscale tool allows you to increase the resolution of an existing image while maintaining or enhancing its quality and details. It uses a specialized ControlNet model that's trained to intelligently upscale images without introducing artifacts or losing image fidelity.
+
+Under the hood, this is simply the controlnet pipeline with [jasperai/Flux.1-dev-Controlnet-Upscaler](https://huggingface.co/jasperai/Flux.1-dev-Controlnet-Upscaler), with a small modification that we don't process the image with canny edge detection.
+
+![upscale example](src/mflux/assets/upscale_example.jpg)
+*Image credit: [Kevin Mueller on Unsplash](https://unsplash.com/photos/gray-owl-on-black-background-xvwZJNaiRNo)*
+
+#### How to use
+
+```sh
+mflux-upscale \
+  --prompt "A gray owl on black background" \
+  --steps 28 \
+  --seed 42 \
+  --height 1363 \
+  --width 908 \
+  -q 8 \
+  --controlnet-image-path "low_res_image.png" \
+  --controlnet-strength 0.6 
+```
+
+This will upscale your input image to the specified dimensions. The upscaler works best when increasing the resolution by a factor of 2-4x.
+
+⚠️ *Note: Depending on the capability of your machine, you might run out of memory when trying to export the image. Try the `--vae-tiling` flag to process the image in smaller chunks, which significantly reduces memory usage during the VAE decoding phase at a minimal cost to performance. This optimization may occasionally produce a subtle seam in the middle of the image, but it's often worth the tradeoff for being able to generate higher resolution images. You can also use `--vae-tiling-split vertical` to change the split direction from the default horizontal (top/bottom) to vertical (left/right).*
+
+#### Tips for Best Results
+
+- For optimal results, try to maintain the same aspect ratio as the original image
+- Prompting matters: Try to acculturate describe the image when upscaling
+- The recommended `--controlnet-strength` is in the range between 0.5 to 0.7
+
+---
+
 ### 🎛️ Dreambooth fine-tuning
 
 As of release [v.0.5.0](https://github.com/filipstrand/mflux/releases/tag/v.0.5.0), MFLUX has support for fine-tuning your own LoRA adapters using the [Dreambooth](https://dreambooth.github.io) technique.

pyproject.toml

@@ -64,6 +64,7 @@ mflux-generate-redux = "mflux.generate_redux:main"
 mflux-save = "mflux.save:main"
 mflux-save-depth = "mflux.save_depth:main"
 mflux-train = "mflux.train:main"
+mflux-upscale = "mflux.upscale:main"
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/mflux/config/model_config.py

@@ -10,6 +10,7 @@ class ModelConfig:
         alias: str | None,
         model_name: str,
         base_model: str | None,
+        controlnet_model: str | None,
         num_train_steps: int,
         max_sequence_length: int,
         supports_guidance: bool,
@@ -19,6 +20,7 @@ class ModelConfig:
         self.alias = alias
         self.model_name = model_name
         self.base_model = base_model
+        self.controlnet_model = controlnet_model
         self.num_train_steps = num_train_steps
         self.max_sequence_length = max_sequence_length
         self.supports_guidance = supports_guidance
@@ -45,6 +47,21 @@ class ModelConfig:
     def dev_redux() -> "ModelConfig":
         return AVAILABLE_MODELS["dev-redux"]
 
+    @staticmethod
+    @lru_cache
+    def dev_controlnet_canny() -> "ModelConfig":
+        return AVAILABLE_MODELS["dev-controlnet-canny"]
+
+    @staticmethod
+    @lru_cache
+    def schnell_controlnet_canny() -> "ModelConfig":
+        return AVAILABLE_MODELS["schnell-controlnet-canny"]
+
+    @staticmethod
+    @lru_cache
+    def dev_controlnet_upscaler() -> "ModelConfig":
+        return AVAILABLE_MODELS["dev-controlnet-upscaler"]
+
     @staticmethod
     @lru_cache
     def schnell() -> "ModelConfig":
@@ -58,6 +75,9 @@ class ModelConfig:
         else:
             return 64
 
+    def is_canny(self) -> bool:
+        return self.alias == "dev-controlnet-canny" or self.alias == "schnell-controlnet-canny"
+
     @staticmethod
     def from_name(
         model_name: str,
@@ -95,6 +115,7 @@ class ModelConfig:
             alias=default_base.alias,
             model_name=model_name,
             base_model=default_base.model_name,
+            controlnet_model=default_base.controlnet_model,
             num_train_steps=default_base.num_train_steps,
             max_sequence_length=default_base.max_sequence_length,
             supports_guidance=default_base.supports_guidance,
@@ -108,6 +129,7 @@ AVAILABLE_MODELS = {
         alias="schnell",
         model_name="black-forest-labs/FLUX.1-schnell",
         base_model=None,
+        controlnet_model=None,
         num_train_steps=1000,
         max_sequence_length=256,
         supports_guidance=False,
@@ -118,6 +140,7 @@ AVAILABLE_MODELS = {
         alias="dev",
         model_name="black-forest-labs/FLUX.1-dev",
         base_model=None,
+        controlnet_model=None,
         num_train_steps=1000,
         max_sequence_length=512,
         supports_guidance=True,
@@ -128,6 +151,7 @@ AVAILABLE_MODELS = {
         alias="dev-fill",
         model_name="black-forest-labs/FLUX.1-Fill-dev",
         base_model=None,
+        controlnet_model=None,
         num_train_steps=1000,
         max_sequence_length=512,
         supports_guidance=True,
@@ -138,6 +162,7 @@ AVAILABLE_MODELS = {
         alias="dev-depth",
         model_name="black-forest-labs/FLUX.1-Depth-dev",
         base_model=None,
+        controlnet_model=None,
         num_train_steps=1000,
         max_sequence_length=512,
         supports_guidance=True,
@@ -148,10 +173,44 @@ AVAILABLE_MODELS = {
         alias="dev-redux",
         model_name="black-forest-labs/FLUX.1-Redux-dev",
         base_model=None,
+        controlnet_model=None,
         num_train_steps=1000,
         max_sequence_length=512,
         supports_guidance=True,
         requires_sigma_shift=True,
         priority=3,
     ),
+    "dev-controlnet-canny": ModelConfig(
+        alias="dev-controlnet-canny",
+        model_name="black-forest-labs/FLUX.1-dev",
+        base_model=None,
+        controlnet_model="InstantX/FLUX.1-dev-Controlnet-Canny",
+        num_train_steps=1000,
+        max_sequence_length=512,
+        supports_guidance=True,
+        requires_sigma_shift=True,
+        priority=5,
+    ),
+    "schnell-controlnet-canny": ModelConfig(
+        alias="schnell-controlnet-canny",
+        model_name="black-forest-labs/FLUX.1-schnell",
+        base_model=None,
+        controlnet_model="InstantX/FLUX.1-dev-Controlnet-Canny",
+        num_train_steps=1000,
+        max_sequence_length=256,
+        supports_guidance=False,
+        requires_sigma_shift=False,
+        priority=6,
+    ),
+    "dev-controlnet-upscaler": ModelConfig(
+        alias="dev-controlnet-upscaler",
+        model_name="black-forest-labs/FLUX.1-dev",
+        base_model=None,
+        controlnet_model="jasperai/Flux.1-dev-Controlnet-Upscaler",
+        num_train_steps=1000,
+        max_sequence_length=512,
+        supports_guidance=False,
+        requires_sigma_shift=False,
+        priority=7,
+    ),
 }

src/mflux/controlnet/controlnet_util.py

@@ -18,15 +18,18 @@ class ControlnetUtil:
         height: int,
         width: int,
         controlnet_image_path: str,
+        is_canny: bool,
     ) -> tuple[mx.array, PIL.Image.Image]:
         from mflux import ImageUtil
 
         control_image = ImageUtil.load_image(controlnet_image_path)
         control_image = ControlnetUtil._scale_image(height=height, width=width, img=control_image)
-        control_image = ControlnetUtil._preprocess_canny(control_image)
+        if is_canny:
+            control_image = ControlnetUtil._preprocess_canny(control_image)
         controlnet_cond = ImageUtil.to_array(control_image)
         controlnet_cond = vae.encode(controlnet_cond)
-        controlnet_cond = (controlnet_cond / vae.scaling_factor) + vae.shift_factor
+        if is_canny:
+            controlnet_cond = (controlnet_cond / vae.scaling_factor) + vae.shift_factor
         controlnet_cond = ArrayUtil.pack_latents(latents=controlnet_cond, height=height, width=width)
         return controlnet_cond, control_image
 

src/mflux/controlnet/flux_controlnet.py

@@ -65,6 +65,7 @@ class Flux1Controlnet(nn.Module):
             height=config.height,
             width=config.width,
             controlnet_image_path=controlnet_image_path,
+            is_canny=self.model_config.is_canny(),
         )
 
         # 2. Create the initial latents

src/mflux/controlnet/weight_handler_controlnet.py

@@ -8,8 +8,6 @@ from mlx.utils import tree_unflatten
 from mflux.weights.weight_handler import MetaData
 from mflux.weights.weight_util import WeightUtil
 
-CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny"
-
 
 class WeightHandlerControlnet:
     def __init__(self, meta_data: MetaData, config: dict, controlnet_transformer: dict | None = None):
@@ -18,8 +16,8 @@ class WeightHandlerControlnet:
         self.config = config
 
     @staticmethod
-    def load_controlnet_transformer() -> "WeightHandlerControlnet":
-        controlnet_path = Path(snapshot_download(repo_id=CONTROLNET_ID, allow_patterns=["*.safetensors", "config.json"]))  # fmt:off
+    def load_controlnet_transformer(controlnet_model: str) -> "WeightHandlerControlnet":
+        controlnet_path = Path(snapshot_download(repo_id=controlnet_model, allow_patterns=["*.safetensors", "config.json"]))  # fmt:off
         file = next(controlnet_path.glob("diffusion_pytorch_model.safetensors"))
         quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level")
         weights = list(mx.load(str(file)).items())

src/mflux/flux/flux_initializer.py

@@ -175,7 +175,9 @@ class FluxInitializer:
         )
 
         # 2. Apply ControlNet-specific initialization
-        weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer()
+        weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer(
+            controlnet_model=model_config.controlnet_model
+        )
         flux_model.transformer_controlnet = TransformerControlnet(
             model_config=model_config,
             num_transformer_blocks=weights_controlnet.num_transformer_blocks(),

src/mflux/generate.py

@@ -62,6 +62,11 @@ def _register_callbacks(args: Namespace, flux: Flux1) -> MemorySaver | None:
     battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit)
     CallbackRegistry.register_before_loop(battery_saver)
 
+    # VAE Tiling
+    if args.vae_tiling:
+        flux.vae.decoder.enable_tiling = True
+        flux.vae.decoder.split_direction = args.vae_tiling_split
+
     # Stepwise Handler
     if args.stepwise_image_output_dir:
         handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)

src/mflux/generate_controlnet.py

@@ -24,7 +24,7 @@ def main():
 
     # 1. Load the model
     flux = Flux1Controlnet(
-        model_config=ModelConfig.from_name(model_name=args.model, base_model=args.base_model),
+        model_config=_get_controlnet_model_config(args.model),
         quantize=args.quantize,
         local_path=args.path,
         lora_paths=args.lora_paths,
@@ -59,11 +59,22 @@ def main():
             print(memory_saver.memory_stats())
 
 
+def _get_controlnet_model_config(model_name: str) -> ModelConfig:
+    if model_name == "schnell":
+        return ModelConfig.schnell_controlnet_canny()
+    return ModelConfig.dev_controlnet_canny()
+
+
 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)
 
+    # VAE Tiling
+    if args.vae_tiling:
+        flux.vae.decoder.enable_tiling = True
+        flux.vae.decoder.split_direction = args.vae_tiling_split
+
     # Canny Image Saver
     if args.controlnet_save_canny:
         canny_image_saver = CannyImageSaver(path=args.output)

src/mflux/generate_depth.py

@@ -68,6 +68,11 @@ def _register_callbacks(args: Namespace, flux: Flux1Depth) -> MemorySaver | None
     battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit)
     CallbackRegistry.register_before_loop(battery_saver)
 
+    # VAE Tiling
+    if args.vae_tiling:
+        flux.vae.decoder.enable_tiling = True
+        flux.vae.decoder.split_direction = args.vae_tiling_split
+
     # Depth Image Saver
     if args.save_depth_map:
         depth_image_saver = DepthImageSaver(path=args.output)

src/mflux/generate_fill.py

@@ -67,6 +67,11 @@ def _register_callbacks(args: Namespace, flux: Flux1Fill) -> MemorySaver | None:
     battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit)
     CallbackRegistry.register_before_loop(battery_saver)
 
+    # VAE Tiling
+    if args.vae_tiling:
+        flux.vae.decoder.enable_tiling = True
+        flux.vae.decoder.split_direction = args.vae_tiling_split
+
     # Stepwise Handler
     if args.stepwise_image_output_dir:
         handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)

src/mflux/generate_in_context.py

@@ -75,6 +75,11 @@ def _register_callbacks(args: Namespace, flux: Flux1InContextLoRA) -> MemorySave
     battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit)
     CallbackRegistry.register_before_loop(battery_saver)
 
+    # VAE Tiling
+    if args.vae_tiling:
+        flux.vae.decoder.enable_tiling = True
+        flux.vae.decoder.split_direction = "vertical"
+
     # Stepwise Handler
     if args.stepwise_image_output_dir:
         handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)

src/mflux/generate_redux.py

@@ -71,6 +71,11 @@ def _register_callbacks(args: Namespace, flux: Flux1Redux) -> MemorySaver | None
     battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit)
     CallbackRegistry.register_before_loop(battery_saver)
 
+    # VAE Tiling
+    if args.vae_tiling:
+        flux.vae.decoder.enable_tiling = True
+        flux.vae.decoder.split_direction = args.vae_tiling_split
+
     # Stepwise Handler
     if args.stepwise_image_output_dir:
         handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)

src/mflux/models/vae/decoder/decoder.py

@@ -11,7 +11,7 @@ from mflux.models.vae.decoder.up_block_4 import UpBlock4
 
 
 class Decoder(nn.Module):
-    def __init__(self):
+    def __init__(self, enable_tiling: bool = False, split_direction: str = "horizontal"):
         super().__init__()
         self.conv_in = ConvIn()
         self.mid_block = UnetMidBlock()
@@ -23,13 +23,78 @@ class Decoder(nn.Module):
         ]
         self.conv_norm_out = ConvNormOut()
         self.conv_out = ConvOut()
+        self.enable_tiling = enable_tiling
+        self.split_direction = split_direction
 
     def __call__(self, latents: mx.array) -> mx.array:
         latents = self.conv_in(latents)
         latents = self.mid_block(latents)
-        for up_block in self.up_blocks:
-            latents = up_block(latents)
+        for i, up_block in enumerate(self.up_blocks):
+            latents = self._apply_up_block(i, up_block, latents)
         latents = self.conv_norm_out(latents)
         latents = nn.silu(latents)
         latents = self.conv_out(latents)
         return latents
+
+    def _apply_up_block(
+        self,
+        block_id: int,
+        up_block: nn.Module,
+        latents: mx.array,
+    ) -> mx.array:
+        if self.enable_tiling:
+            return Decoder._apply_up_block_with_tiling(block_id, up_block, latents, self.split_direction)
+        else:
+            return up_block(latents)
+
+    @staticmethod
+    def _apply_up_block_with_tiling(
+        block_id: int,
+        up_block: nn.Module,
+        latents: mx.array,
+        split_direction: str = "horizontal",
+    ) -> mx.array:
+        # Turns out that the third block is the first to cause OOM issues for larger resolutions.
+        # This trick is a "lossy optimization" that splits the latents and processes them separately.
+        # The resulting image may have a slightly visible seam at the split point.
+        if block_id == 2:
+            latents = Decoder._process_block_3_in_tiles(latents, up_block, split_direction)
+        else:
+            latents = up_block(latents)
+        return latents
+
+    @staticmethod
+    def _process_block_3_in_tiles(
+        latents: mx.array, up_block: nn.Module, split_direction: str = "horizontal"
+    ) -> mx.array:
+        B, C, H, W = latents.shape
+        if split_direction == "horizontal":
+            return Decoder._process_horizontal(H, latents, up_block)
+        else:
+            return Decoder._process_vertical(W, latents, up_block)
+
+    @staticmethod
+    def _process_vertical(width: int, latents: mx.array, up_block: nn.Module) -> mx.array:
+        # 1. Tile left and right
+        left_tile_input = latents[:, :, :, : width // 2]
+        right_tile_input = latents[:, :, :, width // 2 :]
+
+        # 2. Process each tile individually
+        processed_left_tile = up_block(left_tile_input)
+        processed_right_tile = up_block(right_tile_input)
+
+        # 3. Concatenate along width (axis=3)
+        return mx.concatenate([processed_left_tile, processed_right_tile], axis=3)
+
+    @staticmethod
+    def _process_horizontal(height: int, latents: mx.array, up_block: nn.Module) -> mx.array:
+        # 1. Tile top and bottom
+        top_tile_input = latents[:, :, : height // 2, :]
+        bottom_tile_input = latents[:, :, height // 2 :, :]
+
+        # 2. Process each tile individually
+        processed_top_tile = up_block(top_tile_input)
+        processed_bottom_tile = up_block(bottom_tile_input)
+
+        # 3. Concatenate along height (axis=2)
+        return mx.concatenate([processed_top_tile, processed_bottom_tile], axis=2)

src/mflux/post_processing/image_util.py

@@ -115,7 +115,7 @@ class ImageUtil:
 
     @staticmethod
     def load_image(path: str | Path) -> PIL.Image.Image:
-        return PIL.Image.open(path)
+        return PIL.Image.open(path).convert("RGB")
 
     @staticmethod
     def expand_image(

src/mflux/ui/cli/parsers.py

@@ -44,6 +44,8 @@ class CommandLineParser(argparse.ArgumentParser):
     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("--vae-tiling", action="store_true", help="Enable VAE tiling to reduce memory usage at the cost of a potential seam")
+        self.add_argument("--vae-tiling-split", type=str, choices=["horizontal", "vertical"], default="horizontal", help="Direction to split the latents when using VAE tiling (horizontal = top/bottom, vertical = left/right). Default is horizontal.")
 
     def add_model_arguments(self, path_type: t.Literal["load", "save"] = "load", require_model_arg: bool = True) -> None:
         self.require_model_arg = require_model_arg

src/mflux/upscale.py

@@ -0,0 +1,88 @@
+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.memory_saver import MemorySaver
+from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
+from mflux.error.exceptions import PromptFileReadError
+from mflux.ui.cli.parsers import CommandLineParser
+from mflux.ui.prompt_utils import get_effective_prompt
+
+
+def main():
+    # 0. Parse command line arguments
+    parser = CommandLineParser(description="Upscale an image.")
+    parser.add_general_arguments()
+    parser.add_model_arguments(require_model_arg=False)
+    parser.add_lora_arguments()
+    parser.add_image_generator_arguments(supports_metadata_config=False)
+    parser.add_controlnet_arguments()
+    parser.add_output_arguments()
+    args = parser.parse_args()
+
+    # 1. Load the model
+    flux = Flux1Controlnet(
+        model_config=ModelConfig.dev_controlnet_upscaler(),
+        quantize=args.quantize,
+        local_path=args.path,
+        lora_paths=args.lora_paths,
+        lora_scales=args.lora_scales,
+    )
+
+    # 2. Register the optional callbacks
+    memory_saver = _register_callbacks(args=args, flux=flux)
+
+    try:
+        for seed in args.seed:
+            # 3. Generate an upscaled image for each seed value
+            image = flux.generate_image(
+                seed=seed,
+                prompt=get_effective_prompt(args),
+                controlnet_image_path=args.controlnet_image_path,
+                config=Config(
+                    num_inference_steps=args.steps,
+                    height=args.height,
+                    width=args.width,
+                    controlnet_strength=args.controlnet_strength,
+                ),
+            )
+
+            # 4. Save the image
+            image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata)
+    except (StopImageGenerationException, PromptFileReadError) as exc:
+        print(exc)
+    finally:
+        if memory_saver:
+            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)
+
+    # VAE Tiling
+    if args.vae_tiling:
+        flux.vae.decoder.enable_tiling = True
+        flux.vae.decoder.split_direction = args.vae_tiling_split
+
+    # 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()

tests/callbacks/test_battery_saver.py

@@ -6,24 +6,6 @@ from mflux.callbacks.instances.battery_saver import BatterySaver, get_battery_pe
 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:

tests/image_generation/helpers/__init__.py

@@ -1,9 +0,0 @@
-from tests.image_generation.helpers.image_generation_controlnet_test_helper import ImageGeneratorControlnetTestHelper
-from tests.image_generation.helpers.image_generation_in_context_test_helper import ImageGeneratorInContextTestHelper
-from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper
-
-__all__ = [
-    "ImageGeneratorTestHelper",
-    "ImageGeneratorControlnetTestHelper",
-    "ImageGeneratorInContextTestHelper",
-]

tests/image_generation/helpers/image_generation_controlnet_test_helper.py

@@ -17,6 +17,8 @@ class ImageGeneratorControlnetTestHelper:
         prompt: str,
         steps: int,
         seed: int,
+        height: int,
+        width: int,
         controlnet_strength: float,
         lora_paths: list[str] | None = None,
         lora_scales: list[float] | None = None,
@@ -43,8 +45,8 @@ class ImageGeneratorControlnetTestHelper:
                 controlnet_image_path=controlnet_image_path,
                 config=Config(
                     num_inference_steps=steps,
-                    height=768,
-                    width=493,
+                    height=height,
+                    width=width,
                     controlnet_strength=controlnet_strength,
                 ),
             )

tests/image_generation/test_generate_image_controlnet.py

@@ -10,9 +10,11 @@ class TestImageGeneratorControlnet:
             reference_image_path="reference_controlnet_schnell.png",
             output_image_path="output_controlnet_schnell.png",
             controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
-            model_config=ModelConfig.schnell(),
+            model_config=ModelConfig.schnell_controlnet_canny(),
             steps=2,
             seed=43,
+            height=768,
+            width=493,
             prompt="The joker with a hat and a cane",
             controlnet_strength=0.4,
         )
@@ -22,9 +24,11 @@ class TestImageGeneratorControlnet:
             reference_image_path="reference_controlnet_dev.png",
             output_image_path="output_controlnet_dev.png",
             controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
-            model_config=ModelConfig.dev(),
+            model_config=ModelConfig.dev_controlnet_canny(),
             steps=15,
             seed=42,
+            height=768,
+            width=493,
             prompt="The joker with a hat and a cane",
             controlnet_strength=0.4,
         )
@@ -34,11 +38,27 @@ class TestImageGeneratorControlnet:
             reference_image_path="reference_controlnet_dev_lora.png",
             output_image_path="output_controlnet_dev_lora.png",
             controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
-            model_config=ModelConfig.dev(),
+            model_config=ModelConfig.dev_controlnet_canny(),
             steps=15,
             seed=43,
+            height=768,
+            width=493,
             prompt="mkym this is made of wool, The joker with a hat and a cane",
             lora_paths=["FLUX-dev-lora-MiaoKa-Yarn-World.safetensors"],
             lora_scales=[1.0],
             controlnet_strength=0.4,
         )
+
+    def test_image_upscaling(self):
+        ImageGeneratorControlnetTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_upscaled.png",
+            output_image_path="output_upscaler.png",
+            controlnet_image_path="low_res.jpg",
+            model_config=ModelConfig.dev_controlnet_upscaler(),
+            steps=20,
+            seed=42,
+            height=int(192 * 2),
+            width=int(320 * 2),
+            prompt="A man holding up a hand",
+            controlnet_strength=0.6,
+        )

tests/model_config/test_model_config.py

@@ -64,6 +64,82 @@ def test_bfl_dev_fill_full_name():
     assert model.requires_sigma_shift is True
 
 
+def test_bfl_dev_depth():
+    model = ModelConfig.from_name("dev-depth")
+    assert model.alias == "dev-depth"
+    assert model.model_name == "black-forest-labs/FLUX.1-Depth-dev"
+    assert model.max_sequence_length == 512
+    assert model.num_train_steps == 1000
+    assert model.supports_guidance is True
+    assert model.requires_sigma_shift is True
+
+
+def test_bfl_dev_depth_full_name():
+    model = ModelConfig.from_name("black-forest-labs/FLUX.1-Depth-dev")
+    assert model.alias == "dev-depth"
+    assert model.model_name == "black-forest-labs/FLUX.1-Depth-dev"
+    assert model.max_sequence_length == 512
+    assert model.num_train_steps == 1000
+    assert model.supports_guidance is True
+    assert model.requires_sigma_shift is True
+
+
+def test_bfl_dev_redux():
+    model = ModelConfig.from_name("dev-redux")
+    assert model.alias == "dev-redux"
+    assert model.model_name == "black-forest-labs/FLUX.1-Redux-dev"
+    assert model.max_sequence_length == 512
+    assert model.num_train_steps == 1000
+    assert model.supports_guidance is True
+    assert model.requires_sigma_shift is True
+
+
+def test_bfl_dev_redux_full_name():
+    model = ModelConfig.from_name("black-forest-labs/FLUX.1-Redux-dev")
+    assert model.alias == "dev-redux"
+    assert model.model_name == "black-forest-labs/FLUX.1-Redux-dev"
+    assert model.max_sequence_length == 512
+    assert model.num_train_steps == 1000
+    assert model.supports_guidance is True
+    assert model.requires_sigma_shift is True
+
+
+def test_bfl_dev_controlnet_canny():
+    model = ModelConfig.from_name("dev-controlnet-canny")
+    assert model.alias == "dev-controlnet-canny"
+    assert model.model_name == "black-forest-labs/FLUX.1-dev"
+    assert model.controlnet_model == "InstantX/FLUX.1-dev-Controlnet-Canny"
+    assert model.max_sequence_length == 512
+    assert model.num_train_steps == 1000
+    assert model.supports_guidance is True
+    assert model.requires_sigma_shift is True
+    assert model.is_canny() is True
+
+
+def test_bfl_schnell_controlnet_canny():
+    model = ModelConfig.from_name("schnell-controlnet-canny")
+    assert model.alias == "schnell-controlnet-canny"
+    assert model.model_name == "black-forest-labs/FLUX.1-schnell"
+    assert model.controlnet_model == "InstantX/FLUX.1-dev-Controlnet-Canny"
+    assert model.max_sequence_length == 256
+    assert model.num_train_steps == 1000
+    assert model.supports_guidance is False
+    assert model.requires_sigma_shift is False
+    assert model.is_canny() is True
+
+
+def test_bfl_dev_controlnet_upscaler():
+    model = ModelConfig.from_name("dev-controlnet-upscaler")
+    assert model.alias == "dev-controlnet-upscaler"
+    assert model.model_name == "black-forest-labs/FLUX.1-dev"
+    assert model.controlnet_model == "jasperai/Flux.1-dev-Controlnet-Upscaler"
+    assert model.max_sequence_length == 512
+    assert model.num_train_steps == 1000
+    assert model.supports_guidance is False
+    assert model.requires_sigma_shift is False
+    assert model.is_canny() is False
+
+
 def test_community_dev_fill_implicit_base_model():
     model = ModelConfig.from_name("acme-lab/some-dev-fill-model")
     assert model.alias == "dev-fill"