Merge pull request #129 from filipstrand/inpaint

FLUX.1 Tools | Fill

README.md

@@ -31,6 +31,10 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
   * [Available Styles](#available-styles)
   * [How It Works](#how-it-works)
   * [Tips for Best Results](#tips-for-best-results)
+- [🛠️ Flux Tools](#%EF%B8%8F-flux-tools)
+  * [🖌️ Fill](#%EF%B8%8F-fill)
+    + [Inpainting](#inpainting)
+    + [Outpainting](#outpainting)
 - [🎛️ Dreambooth fine-tuning](#-dreambooth-fine-tuning)
   * [Training configuration](#training-configuration)
   * [Training example](#training-example)
@@ -220,7 +224,7 @@ The `mflux-generate-controlnet` command supports most of the same arguments as `
 
 - **`--controlnet-save-canny`** (optional, bool, default: `False`): If set, saves the Canny edge detection reference image used by ControlNet.
 
-See the [ControlNet](#%EF%B8%8F-controlnet) section for more details on how to use this feature effectively.
+See the [Controlnet](#%EF%B8%8F-controlnet) section for more details on how to use this feature effectively.
 
 
 #### 📜 Batch Image Generation Arguments
@@ -769,6 +773,104 @@ This prompt clearly describes both the reference image (after `[IMAGE1]`) and th
 
 ---
 
+### 🛠️ Flux Tools
+
+MFLUX supports the official [Flux.1 Tools](https://blackforestlabs.ai/flux-1-tools/).
+
+#### 🖌️ Fill
+
+The Fill tool uses the [FLUX.1-Fill-dev](https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev) model to allow you to selectively edit parts of an image by providing a binary mask. This is useful for inpainting (replacing specific areas) and outpainting (expanding the canvas).
+
+#### Inpainting
+
+Inpainting allows you to selectively regenerate specific parts of an image while preserving the rest. This is perfect for removing unwanted objects, adding new elements, or changing specific areas of an image without affecting the surrounding content. The Fill tool understands the context of the entire image and creates seamless edits that blend naturally with the preserved areas.
+
+![inpaint example](src/mflux/assets/inpaint_example.jpg)
+*Original dog image credit: [Julio Bernal on Unsplash](https://unsplash.com/photos/brown-and-white-long-coated-dog-WLGx0fKFSeI)*
+
+##### Creating Masks
+
+Before using the Fill tool, you need an image and a corresponding mask. You can create a mask using the included tool:
+
+```bash
+python -m tools.fill_mask_tool /path/to/your/image.jpg
+```
+
+This will open an interactive interface where you can paint over the areas you want to regenerate.
+Pressing the `s` key will save the mask at the same location as the image.
+
+##### Example
+
+To regenerate specific parts of an image:
+
+```bash
+mflux-generate-fill \
+  --prompt "A professionally shot studio photograph of a dog wearing a red hat and ski goggles. The dog is centered against a uniformly bright yellow background, with well-balanced lighting and sharp details." \
+  --steps 20 \
+  --seed 42 \
+  --height 1280 \
+  --width 851 \
+  --guidance 30 \
+  -q 8 \
+  --image-path "dog.png" \
+  --masked-image-path "dog_mask.png" \
+```
+
+#### Outpainting
+
+![outpaint example](src/mflux/assets/outpaint_example.jpg)
+*Original room image credit: [Alexey Aladashvili on Unsplash](https://unsplash.com/photos/a-living-room-with-a-couch-and-a-table-JFzglhmwlck)*
+
+Outpainting extends your image beyond its original boundaries, allowing you to expand the canvas in any direction while maintaining visual consistency with the original content. This is useful for creating wider landscapes, revealing more of a scene, or transforming a portrait into a full-body image. The Fill tool intelligently generates new content that seamlessly connects with the existing image.
+
+You can expand the canvas of your image using the provided tool:
+
+```bash
+python -m tools.create_outpaint_image_canvas_and_mask \
+  /path/to/your/image.jpg \
+  --image-outpaint-padding "0,30%,20%,0"
+```
+
+As an example, here's how to add 25% padding to both the left and right sides of an image:
+
+```bash
+python -m tools.create_outpaint_image_canvas_and_mask \
+  room.png \
+  --image-outpaint-padding "0,25%,0,25%"
+```
+
+The padding format is "top,right,bottom,left" where each value can be in pixels or as a percentage. For example, "0,30%,20%,0" expands the canvas by 30% to the right and 20% to the bottom.
+
+After running this command, you'll get two files: an expanded canvas with your original image and a corresponding mask. You can then run the `mflux-generate-fill` command similar to the inpainting example, using these files as input.
+
+##### Example
+
+Once you've created the expanded canvas and mask files, run the Fill tool on them:
+
+```bash
+mflux-generate-fill \
+  --prompt "A detailed interior room photograph with natural lighting, extended in a way that perfectly complements the original space. The expanded areas continue the architectural style, color scheme, and overall aesthetic of the room seamlessly." \
+  --steps 25 \
+  --seed 43 \
+  --guidance 30 \
+  -q 8 \
+  --image-path "room.png" \
+  --masked-image-path "room_mask.png" \
+```
+
+##### Tips for Best Results
+
+- **Model**: The model is always `FLUX.1-Fill-dev` and should not be specified.
+- **Guidance**: Higher guidance values (around 30) typically yield better results. This is the default if not specified.
+- **Resolution**: Higher resolution images generally produce better results with the Fill tool.
+- **Masks**: Make sure your mask clearly defines the areas you want to regenerate. White areas in the mask will be regenerated, while black areas will be preserved.
+- **Prompting**: For best results, provide detailed prompts that describe both what you want in the newly generated areas and how it should relate to the preserved parts of the image.
+- **Steps**: Using 20-30 denoising steps generally produces higher quality results.
+
+⚠️ *Note: Using the Fill tool requires an additional 33.92 GB download from [black-forest-labs/FLUX.1-Fill-dev](https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev). The download happens automatically on first use.*
+
+---
+
 ### 🎛️ 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.
@@ -963,6 +1065,8 @@ with different prompts and LoRA adapters active.
 
 ![image](src/mflux/assets/controlnet2.jpg)
 
+---
+
 ### 🚧 Current limitations
 
 - Images are generated one by one.

pyproject.toml

@@ -57,6 +57,7 @@ homepage = "https://github.com/filipstrand/mflux"
 mflux-generate = "mflux.generate:main"
 mflux-generate-controlnet = "mflux.generate_controlnet:main"
 mflux-generate-in-context = "mflux.generate_in_context:main"
+mflux-generate-fill = "mflux.generate_fill:main"
 mflux-save = "mflux.save:main"
 mflux-train = "mflux.train:main"
 

src/mflux/__init__.py

@@ -3,11 +3,13 @@ from mflux.config.model_config import ModelConfig
 from mflux.controlnet.flux_controlnet import Flux1Controlnet
 from mflux.error.exceptions import StopImageGenerationException
 from mflux.flux.flux import Flux1
+from mflux.flux_tools.fill.flux_fill import Flux1Fill
 from mflux.post_processing.image_util import ImageUtil
 
 __all__ = [
     "Flux1",
     "Flux1Controlnet",
+    "Flux1Fill",
     "Config",
     "ModelConfig",
     "ImageUtil",

src/mflux/config/config.py

@@ -17,6 +17,7 @@ class Config:
         guidance: float = 4.0,
         image_path: Path | None = None,
         image_strength: float | None = None,
+        masked_image_path: Path | None = None,
         controlnet_strength: float | None = None,
     ):
         if width % 16 != 0 or height % 16 != 0:
@@ -27,4 +28,5 @@ class Config:
         self.guidance = guidance
         self.image_path = image_path
         self.image_strength = image_strength
+        self.masked_image_path = masked_image_path
         self.controlnet_strength = controlnet_strength

src/mflux/config/model_config.py

@@ -13,6 +13,8 @@ class ModelConfig:
         num_train_steps: int,
         max_sequence_length: int,
         supports_guidance: bool,
+        requires_sigma_shift: bool,
+        priority: int,
     ):
         self.alias = alias
         self.model_name = model_name
@@ -20,65 +22,57 @@ class ModelConfig:
         self.num_train_steps = num_train_steps
         self.max_sequence_length = max_sequence_length
         self.supports_guidance = supports_guidance
+        self.requires_sigma_shift = requires_sigma_shift
+        self.priority = priority
 
     @staticmethod
     @lru_cache
     def dev() -> "ModelConfig":
-        return ModelConfig(
-            alias="dev",
-            model_name="black-forest-labs/FLUX.1-dev",
-            base_model=None,
-            num_train_steps=1000,
-            max_sequence_length=512,
-            supports_guidance=True,
-        )
+        return AVAILABLE_MODELS["dev"]
+
+    @staticmethod
+    @lru_cache
+    def dev_fill() -> "ModelConfig":
+        return AVAILABLE_MODELS["dev-fill"]
 
     @staticmethod
     @lru_cache
     def schnell() -> "ModelConfig":
-        return ModelConfig(
-            alias="schnell",
-            model_name="black-forest-labs/FLUX.1-schnell",
-            base_model=None,
-            num_train_steps=1000,
-            max_sequence_length=256,
-            supports_guidance=False,
-        )
+        return AVAILABLE_MODELS["schnell"]
 
     @staticmethod
     def from_name(
         model_name: str,
-        base_model: Literal["dev", "schnell"] | None = None,
+        base_model: Literal["dev", "schnell", "dev-fill"] | None = None,
     ) -> "ModelConfig":
-        dev = ModelConfig.dev()
-        schnell = ModelConfig.schnell()
+        # 0. Get all base models (where base_model is None) sorted by priority
+        base_models = sorted(
+            [model for model in AVAILABLE_MODELS.values() if model.base_model is None], key=lambda x: x.priority
+        )
 
-        # 0. Validate explicit base_model
-        allowed_names = [dev.alias, dev.model_name, schnell.alias, schnell.model_name]
+        # 1. Check if model_name matches any base model's alias or full name
+        for base in base_models:
+            if model_name in (base.alias, base.model_name):
+                return base
+
+        # 2. Validate explicit base_model
+        allowed_names = []
+        for base in base_models:
+            allowed_names.extend([base.alias, base.model_name])
         if base_model and base_model not in allowed_names:
             raise InvalidBaseModel(f"Invalid base_model. Choose one of {allowed_names}")
 
-        # 1. If model_name is "dev" or "schnell" then simply return
-        if model_name == dev.model_name or model_name == dev.alias:
-            return dev
-        if model_name == schnell.model_name or model_name == schnell.alias:
-            return schnell
+        # 3. Determine the base model (explicit or inferred)
+        if base_model:
+            # Find by explicit base_model name
+            default_base = next((b for b in base_models if base_model in (b.alias, b.model_name)), None)
+        else:
+            # Infer from model_name substring (priority order via sorted base_models)
+            default_base = next((b for b in base_models if b.alias and b.alias in model_name), None)
+            if not default_base:
+                raise ModelConfigError(f"Cannot infer base_model from {model_name}")
 
-        # 1. Determine the appropriate base model
-        default_base = None
-        if not base_model:
-            if "dev" in model_name:
-                default_base = dev
-            elif "schnell" in model_name:
-                default_base = schnell
-            else:
-                raise ModelConfigError(f"Cannot infer base_model from {model_name}. Specify --base-model.")
-        elif base_model == dev.model_name or base_model == dev.alias:
-            default_base = dev
-        elif base_model == schnell.model_name or base_model == schnell.alias:
-            default_base = schnell
-
-        # 2. Construct the config based on the model name and base default
+        # 4. Construct the config
         return ModelConfig(
             alias=default_base.alias,
             model_name=model_name,
@@ -86,7 +80,40 @@ class ModelConfig:
             num_train_steps=default_base.num_train_steps,
             max_sequence_length=default_base.max_sequence_length,
             supports_guidance=default_base.supports_guidance,
+            requires_sigma_shift=default_base.requires_sigma_shift,
+            priority=default_base.priority,
         )
 
-    def is_dev(self) -> bool:
-        return self.alias == "dev"
+
+AVAILABLE_MODELS = {
+    "schnell": ModelConfig(
+        alias="schnell",
+        model_name="black-forest-labs/FLUX.1-schnell",
+        base_model=None,
+        num_train_steps=1000,
+        max_sequence_length=256,
+        supports_guidance=False,
+        requires_sigma_shift=False,
+        priority=2,
+    ),
+    "dev": ModelConfig(
+        alias="dev",
+        model_name="black-forest-labs/FLUX.1-dev",
+        base_model=None,
+        num_train_steps=1000,
+        max_sequence_length=512,
+        supports_guidance=True,
+        requires_sigma_shift=True,
+        priority=1,
+    ),
+    "dev-fill": ModelConfig(
+        alias="dev-fill",
+        model_name="black-forest-labs/FLUX.1-Fill-dev",
+        base_model=None,
+        num_train_steps=1000,
+        max_sequence_length=512,
+        supports_guidance=True,
+        requires_sigma_shift=True,
+        priority=0,
+    ),
+}

src/mflux/config/runtime_config.py

@@ -32,7 +32,7 @@ class RuntimeConfig:
         self.config.width = value
 
     @property
-    def guidance(self) -> float:
+    def guidance(self) -> float | None:
         return self.config.guidance
 
     @property
@@ -55,6 +55,10 @@ class RuntimeConfig:
     def image_strength(self) -> float | None:
         return self.config.image_strength
 
+    @property
+    def masked_image_path(self) -> str | None:
+        return self.config.masked_image_path
+
     @property
     def init_time_step(self) -> int:
         is_img2img = (
@@ -82,7 +86,7 @@ class RuntimeConfig:
     @staticmethod
     def _create_sigmas(config: Config, model_config: ModelConfig) -> mx.array:
         sigmas = RuntimeConfig._create_sigmas_values(config.num_inference_steps)
-        if model_config.is_dev():
+        if model_config.requires_sigma_shift:
             sigmas = RuntimeConfig._shift_sigmas(sigmas=sigmas, width=config.width, height=config.height)
         return sigmas
 

src/mflux/flux/flux.py

@@ -129,7 +129,7 @@ class Flux1(nn.Module):
             seed=seed,
             prompt=prompt,
             latents=latents,
-            config=config
+            config=config,
         )  # fmt: off
 
         # 7. Decode the latent array and return the image

src/mflux/flux_tools/fill/flux_fill.py

@@ -0,0 +1,160 @@
+import mlx.core as mx
+from mlx import nn
+from tqdm import tqdm
+
+from mflux.callbacks.callbacks import Callbacks
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.config.runtime_config import RuntimeConfig
+from mflux.error.exceptions import StopImageGenerationException
+from mflux.flux.flux_initializer import FluxInitializer
+from mflux.flux_tools.fill.mask_util import MaskUtil
+from mflux.latent_creator.latent_creator import LatentCreator
+from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
+from mflux.models.text_encoder.prompt_encoder import PromptEncoder
+from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
+from mflux.models.transformer.transformer import Transformer
+from mflux.models.vae.vae import VAE
+from mflux.post_processing.array_util import ArrayUtil
+from mflux.post_processing.generated_image import GeneratedImage
+from mflux.post_processing.image_util import ImageUtil
+from mflux.weights.model_saver import ModelSaver
+
+
+class Flux1Fill(nn.Module):
+    vae: VAE
+    transformer: Transformer
+    t5_text_encoder: T5Encoder
+    clip_text_encoder: CLIPEncoder
+
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        quantize: int | None = None,
+        local_path: str | None = None,
+        lora_paths: list[str] | None = None,
+        lora_scales: list[float] | None = None,
+    ):
+        super().__init__()
+        FluxInitializer.init(
+            flux_model=self,
+            model_config=model_config,
+            quantize=quantize,
+            local_path=local_path,
+            lora_paths=lora_paths,
+            lora_scales=lora_scales,
+        )
+
+    def generate_image(
+        self,
+        seed: int,
+        prompt: str,
+        config: Config,
+    ) -> GeneratedImage:
+        # 0. Create a new runtime config based on the model type and input parameters
+        config = RuntimeConfig(config, self.model_config)
+        time_steps = tqdm(range(config.init_time_step, config.num_inference_steps))
+
+        # 1. Create the initial latents
+        latents = LatentCreator.create(
+            seed=seed,
+            height=config.height,
+            width=config.width,
+        )
+
+        # 2. Encode the prompt
+        prompt_embeds, pooled_prompt_embeds = PromptEncoder.encode_prompt(
+            prompt=prompt,
+            prompt_cache=self.prompt_cache,
+            t5_tokenizer=self.t5_tokenizer,
+            clip_tokenizer=self.clip_tokenizer,
+            t5_text_encoder=self.t5_text_encoder,
+            clip_text_encoder=self.clip_text_encoder,
+        )
+
+        # 3. Create the static masked latents
+        static_masked_latents = MaskUtil.create_masked_latents(
+            vae=self.vae,
+            config=config,
+            latents=latents,
+            img_path=config.image_path,
+            mask_path=config.masked_image_path,
+        )
+
+        # (Optional) Call subscribers for beginning of loop
+        Callbacks.before_loop(
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=config,
+        )  # fmt: off
+
+        for t in time_steps:
+            try:
+                # 4.t Concatenate the updated latents with the static masked latents
+                hidden_states = mx.concatenate([latents, static_masked_latents], axis=-1)
+
+                # 5.t Predict the noise
+                noise = self.transformer(
+                    t=t,
+                    config=config,
+                    hidden_states=hidden_states,
+                    prompt_embeds=prompt_embeds,
+                    pooled_prompt_embeds=pooled_prompt_embeds,
+                )
+
+                # 6.t Take one denoise step
+                dt = config.sigmas[t + 1] - config.sigmas[t]
+                latents += noise * dt
+
+                # (Optional) Call subscribers in-loop
+                Callbacks.in_loop(
+                    t=t,
+                    seed=seed,
+                    prompt=prompt,
+                    latents=latents,
+                    config=config,
+                    time_steps=time_steps,
+                )  # fmt: off
+
+                # (Optional) Evaluate to enable progress tracking
+                mx.eval(latents)
+
+            except KeyboardInterrupt:  # noqa: PERF203
+                Callbacks.interruption(
+                    t=t,
+                    seed=seed,
+                    prompt=prompt,
+                    latents=latents,
+                    config=config,
+                    time_steps=time_steps,
+                )
+                raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
+
+        # (Optional) Call subscribers after 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)
+        return ImageUtil.to_image(
+            decoded_latents=decoded,
+            config=config,
+            seed=seed,
+            prompt=prompt,
+            quantization=self.bits,
+            lora_paths=self.lora_paths,
+            lora_scales=self.lora_scales,
+            image_path=config.image_path,
+            image_strength=config.image_strength,
+            masked_image_path=config.masked_image_path,
+            generation_time=time_steps.format_dict["elapsed"],
+        )
+
+    def save_model(self, base_path: str) -> None:
+        ModelSaver.save_model(self, self.bits, base_path)

src/mflux/flux_tools/fill/mask_util.py

@@ -0,0 +1,57 @@
+import mlx.core as mx
+
+from mflux.config.runtime_config import RuntimeConfig
+from mflux.models.vae.vae import VAE
+from mflux.post_processing.array_util import ArrayUtil
+from mflux.post_processing.image_util import ImageUtil
+
+
+class MaskUtil:
+    @staticmethod
+    def create_masked_latents(
+        vae: VAE,
+        config: RuntimeConfig,
+        latents: mx.array,
+        img_path: str,
+        mask_path: str | None
+    ) -> mx.array:  # fmt: off
+        if not img_path or not mask_path:
+            # Return empty latents if no image or mask is provided
+            return mx.zeros((1, 0, 0))
+
+        # 1. Get the reference image
+        scaled_image = ImageUtil.scale_to_dimensions(
+            image=ImageUtil.load_image(img_path).convert("RGB"),
+            target_width=config.width,
+            target_height=config.height,
+        )
+        image = ImageUtil.to_array(scaled_image)
+
+        # 2. Get the mask
+        scaled = ImageUtil.scale_to_dimensions(
+            image=ImageUtil.load_image(mask_path).convert("RGB"),
+            target_width=config.width,
+            target_height=config.height,
+        )
+        the_mask = ImageUtil.to_array(scaled, is_mask=True)
+
+        # 3. Create and pack the masked image
+        masked_image = image * (1 - the_mask)
+        masked_image = vae.encode(masked_image)
+        masked_image = ArrayUtil.pack_latents(latents=masked_image, height=config.height, width=config.width)
+
+        # 4. Resize mask and pack latents
+        mask = MaskUtil._reshape_mask(the_mask=the_mask, height=config.height, width=config.width)
+        mask = ArrayUtil.pack_latents(latents=mask, height=config.height, width=config.width, num_channels_latents=64)
+
+        # 5. Concat the masked_image and the mask
+        masked_image_latents = mx.concatenate([masked_image, mask], axis=-1)
+        return masked_image_latents
+
+    @staticmethod
+    def _reshape_mask(the_mask: mx.array, height: int, width: int) -> mx.array:
+        mask = the_mask[:, 0, :, :]
+        mask = mx.reshape(mask, (1, height // 8, 8, width // 8, 8))
+        mask = mx.transpose(mask, (0, 2, 4, 1, 3))
+        mask = mx.reshape(mask, (1, 64, height // 8, width // 8))
+        return mask

src/mflux/generate_fill.py

@@ -0,0 +1,73 @@
+from mflux import Config, ModelConfig, StopImageGenerationException
+from mflux.callbacks.callback_registry import CallbackRegistry
+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
+from mflux.ui.cli.parsers import CommandLineParser
+
+
+def main():
+    # 0. Parse command line arguments
+    parser = CommandLineParser(description="Generate an image using the fill tool to complete masked areas.")
+    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_fill_arguments()
+    parser.add_output_arguments()
+    args = parser.parse_args()
+
+    # 0. Default to a higher guidance value for fill related tasks.
+    if args.guidance is None:
+        args.guidance = 30
+
+    # 1. Load the model
+    flux = Flux1Fill(
+        model_config=ModelConfig.dev_fill(),
+        quantize=args.quantize,
+        local_path=args.path,
+        lora_paths=args.lora_paths,
+        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)
+        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
+            image = flux.generate_image(
+                seed=seed,
+                prompt=args.prompt,
+                config=Config(
+                    num_inference_steps=args.steps,
+                    height=args.height,
+                    width=args.width,
+                    guidance=args.guidance,
+                    image_path=args.image_path,
+                    masked_image_path=args.masked_image_path,
+                ),
+            )
+
+            # 4. Save the image
+            image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata)
+    except StopImageGenerationException as stop_exc:
+        print(stop_exc)
+    finally:
+        if memory_saver:
+            print(memory_saver.memory_stats())
+
+
+if __name__ == "__main__":
+    main()

src/mflux/post_processing/array_util.py

@@ -10,8 +10,8 @@ class ArrayUtil:
         return latents
 
     @staticmethod
-    def pack_latents(latents: mx.array, height: int, width: int) -> mx.array:
-        latents = mx.reshape(latents, (1, 16, height // 16, 2, width // 16, 2))
+    def pack_latents(latents: mx.array, height: int, width: int, num_channels_latents: int = 16) -> mx.array:
+        latents = mx.reshape(latents, (1, num_channels_latents, height // 16, 2, width // 16, 2))
         latents = mx.transpose(latents, (0, 2, 4, 1, 3, 5))
-        latents = mx.reshape(latents, (1, (width // 16) * (height // 16), 64))
+        latents = mx.reshape(latents, (1, (width // 16) * (height // 16), num_channels_latents * 4))
         return latents

src/mflux/post_processing/generated_image.py

@@ -27,6 +27,7 @@ class GeneratedImage:
         controlnet_strength: float | None = None,
         image_path: str | pathlib.Path | None = None,
         image_strength: float | None = None,
+        masked_image_path: str | pathlib.Path | None = None,
     ):
         self.image = image
         self.model_config = model_config
@@ -43,6 +44,7 @@ class GeneratedImage:
         self.controlnet_strength = controlnet_strength
         self.image_path = image_path
         self.image_strength = image_strength
+        self.masked_image_path = masked_image_path
 
     def get_right_half(self) -> "GeneratedImage":
         # Calculate the coordinates for the right half
@@ -66,6 +68,7 @@ class GeneratedImage:
             controlnet_strength=self.controlnet_strength,
             image_path=self.image_path,
             image_strength=self.image_strength,
+            masked_image_path=self.masked_image_path,
         )
 
     def save(
@@ -100,6 +103,7 @@ class GeneratedImage:
             "image_strength": self.image_strength if self.image_path else None,
             "controlnet_image_path": str(self.controlnet_image_path) if self.controlnet_image_path else None,
             "controlnet_strength": round(self.controlnet_strength, 2) if self.controlnet_strength else None,
+            "masked_image_path": str(self.masked_image_path) if self.masked_image_path else None,
             "prompt": self.prompt,
         }
 

src/mflux/post_processing/image_util.py

@@ -30,6 +30,7 @@ class ImageUtil:
         controlnet_image_path: str | None = None,
         image_path: str | None = None,
         image_strength: float | None = None,
+        masked_image_path: str | None = None,
     ) -> GeneratedImage:
         normalized = ImageUtil._denormalize(decoded_latents)
         normalized_numpy = ImageUtil._to_numpy(normalized)
@@ -50,6 +51,7 @@ class ImageUtil:
             image_strength=image_strength,
             controlnet_image_path=controlnet_image_path,
             controlnet_strength=config.controlnet_strength,
+            masked_image_path=masked_image_path,
         )
 
     @staticmethod
@@ -72,6 +74,10 @@ class ImageUtil:
     def _normalize(images: mx.array) -> mx.array:
         return 2.0 * images - 1.0
 
+    @staticmethod
+    def _binarize(image: mx.array) -> mx.array:
+        return mx.where(image < 0.5, mx.zeros_like(image), mx.ones_like(image))
+
     @staticmethod
     def _to_numpy(images: mx.array) -> np.ndarray:
         images = mx.transpose(images, (0, 2, 3, 1))
@@ -92,11 +98,14 @@ class ImageUtil:
         return images
 
     @staticmethod
-    def to_array(image: PIL.Image.Image) -> mx.array:
+    def to_array(image: PIL.Image.Image, is_mask: bool = False) -> mx.array:
         image = ImageUtil._pil_to_numpy(image)
         array = mx.array(image)
         array = mx.transpose(array, (0, 3, 1, 2))
-        array = ImageUtil._normalize(array)
+        if is_mask:
+            array = ImageUtil._binarize(array)
+        else:
+            array = ImageUtil._normalize(array)
         return array
 
     @staticmethod

src/mflux/post_processing/mask_util.py

@@ -0,0 +1,54 @@
+import mlx.core as mx
+
+from mflux.config.runtime_config import RuntimeConfig
+from mflux.models.vae.vae import VAE
+from mflux.post_processing.array_util import ArrayUtil
+
+
+class MaskUtil:
+    @staticmethod
+    def create_masked_latents(
+        vae: VAE,
+        config: RuntimeConfig,
+        latents: mx.array,
+        img_path: str,
+        mask_path: str | None
+    ) -> mx.array:  # fmt: off
+        from mflux import ImageUtil
+
+        # 1. Get the reference image
+        scaled_image = ImageUtil.scale_to_dimensions(
+            image=ImageUtil.load_image(config.image_path).convert("RGB"),
+            target_width=config.width,
+            target_height=config.height,
+        )
+        image = ImageUtil.to_array(scaled_image)
+
+        # 2. Get the mask
+        scaled = ImageUtil.scale_to_dimensions(
+            image=ImageUtil.load_image(mask_path).convert("RGB"),
+            target_width=config.width,
+            target_height=config.height,
+        )
+        the_mask = ImageUtil.to_array(scaled, is_mask=True)
+
+        # 3. Create and pack the masked image
+        masked_image = image * (1 - the_mask)
+        masked_image = vae.encode(masked_image)
+        masked_image = ArrayUtil.pack_latents(latents=masked_image, height=config.height, width=config.width)
+
+        # 4. Resize mask and pack latents
+        mask = MaskUtil._reshape_mask(the_mask=the_mask, height=config.height, width=config.width)
+        mask = ArrayUtil.pack_latents(latents=mask, height=config.height, width=config.width, num_channels_latents=64)
+
+        # 5. Concat the masked_image and the mask
+        masked_image_latents = mx.concatenate([masked_image, mask], axis=-1)
+        return masked_image_latents
+
+    @staticmethod
+    def _reshape_mask(the_mask: mx.array, height: int, width: int):
+        mask = the_mask[:, 0, :, :]
+        mask = mx.reshape(mask, (1, height // 8, 8, width // 8, 8))
+        mask = mx.transpose(mask, (0, 2, 4, 1, 3))
+        mask = mx.reshape(mask, (1, 64, height // 8, width // 8))
+        return mask

src/mflux/ui/cli/parsers.py

@@ -14,13 +14,14 @@ from mflux.ui import (
 
 class ModelSpecAction(argparse.Action):
     def __call__(self, parser, namespace, values, option_string=None):
-        if values in ["dev", "schnell"]:
+        if values in ui_defaults.MODEL_CHOICES:
             setattr(namespace, self.dest, values)
             return
 
         if values.count("/") != 1:
             raise argparse.ArgumentError(
-                self, f'Value must be either "dev", "schnell", or "in format "org/model". Got: {values}'
+                self,
+                (f'Value must be either {" ".join(ui_defaults.MODEL_CHOICES)} or in format "org/model". Got: {values}'),
             )
 
         # If we got here, values contains exactly one slash
@@ -38,12 +39,13 @@ class CommandLineParser(argparse.ArgumentParser):
         self.supports_image_to_image = False
         self.supports_image_outpaint = False
         self.supports_lora = False
+        self.require_model_arg = True
 
     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.require_model_arg = require_model_arg
         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).")
         if path_type == "load":
             self.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk")
@@ -87,6 +89,10 @@ class CommandLineParser(argparse.ArgumentParser):
         self.add_argument("--global-seed", type=int, default=argparse.SUPPRESS, help="Entropy Seed (used for all prompts in the batch)")
         self._add_image_generator_common_arguments()
 
+    def add_fill_arguments(self) -> None:
+        self.add_argument("--image-path", type=Path, required=True, help="Local path to the source image")
+        self.add_argument("--masked-image-path", type=Path, required=True, help="Local path to the mask image")
+
     def add_output_arguments(self) -> None:
         self.add_argument("--metadata", action="store_true", help="Export image metadata as a JSON file.")
         self.add_argument("--output", type=str, default="image.png", help="The filename for the output image. Default is \"image.png\".")
@@ -184,8 +190,8 @@ class CommandLineParser(argparse.ArgumentParser):
                 if namespace.image_outpaint_padding is None:
                     namespace.image_outpaint_padding = prior_gen_metadata.get("image_outpaint_padding", None)
 
-        # Only require model if we're not in training mode
-        if hasattr(namespace, "model") and namespace.model is None and not has_training_args:
+        # Only require model if we're not in training mode and require_model_arg is True
+        if hasattr(namespace, "model") and namespace.model is None and not has_training_args and self.require_model_arg:
             self.error("--model / -m must be provided, or 'model' must be specified in the config file.")
 
         if self.supports_image_generation and namespace.seed is None and namespace.auto_seeds > 0:

src/mflux/ui/defaults.py

@@ -2,9 +2,10 @@ CONTROLNET_STRENGTH = 0.4
 GUIDANCE_SCALE = 3.5
 HEIGHT, WIDTH = 1024, 1024
 IMAGE_STRENGTH = 0.4
-MODEL_CHOICES = ["dev", "schnell"]
+MODEL_CHOICES = ["dev", "dev-fill", "schnell"]
 MODEL_INFERENCE_STEPS = {
     "dev": 14,
+    "dev-fill": 14,
     "schnell": 4,
 }
 QUANTIZE_CHOICES = [3, 4, 6, 8]

tests/arg_parser/test_cli_argparser.py

@@ -5,14 +5,15 @@ from unittest.mock import patch
 
 import pytest
 
+from mflux.ui import defaults as ui_defaults
 from mflux.ui.box_values import BoxValues
 from mflux.ui.cli.parsers import CommandLineParser
 
 
-def _create_mflux_generate_parser(with_controlnet=False) -> CommandLineParser:
+def _create_mflux_generate_parser(with_controlnet=False, require_model_arg=False) -> CommandLineParser:
     parser = CommandLineParser(description="Generate an image based on a prompt.")
     parser.add_general_arguments()
-    parser.add_model_arguments(require_model_arg=False)
+    parser.add_model_arguments(require_model_arg=require_model_arg)
     parser.add_image_generator_arguments(supports_metadata_config=True)
     parser.add_lora_arguments()
     parser.add_image_to_image_arguments(required=False)
@@ -25,12 +26,12 @@ def _create_mflux_generate_parser(with_controlnet=False) -> CommandLineParser:
 
 @pytest.fixture
 def mflux_generate_parser() -> CommandLineParser:
-    return _create_mflux_generate_parser(with_controlnet=False)
+    return _create_mflux_generate_parser(with_controlnet=False, require_model_arg=False)
 
 
 @pytest.fixture
 def mflux_generate_controlnet_parser() -> CommandLineParser:
-    return _create_mflux_generate_parser(with_controlnet=True)
+    return _create_mflux_generate_parser(with_controlnet=True, require_model_arg=False)
 
 
 @pytest.fixture
@@ -85,6 +86,23 @@ def base_metadata_dict() -> dict:
     }
 
 
+@pytest.fixture
+def mflux_fill_parser() -> CommandLineParser:
+    parser = CommandLineParser(description="Generate an image using the fill tool to complete masked areas.")
+    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_fill_arguments()
+    parser.add_output_arguments()
+    return parser
+
+
+@pytest.fixture
+def mflux_fill_minimal_argv() -> list[str]:
+    return ["mflux-fill", "--prompt", "meaning of life", "--image-path", "image.png", "--masked-image-path", "mask.png"]
+
+
 def test_model_path_requires_model_arg(mflux_generate_parser):
     # when loading a model via --path, the model name still need to be specified
     with patch("sys.argv", "mflux-generate", "--path", "/some/saved/model"):
@@ -96,9 +114,14 @@ def test_model_arg_not_in_file(mflux_generate_parser, mflux_generate_minimal_arg
     with metadata_file.open("wt") as m:
         del base_metadata_dict["model"]
         json.dump(base_metadata_dict, m, indent=4)
-    # test model arg not provided in either flag or file
+
+    # Create a parser that requires the model argument
+    parser_with_required_model = _create_mflux_generate_parser(require_model_arg=True)
+
+    # test model arg not provided in either flag or file should raise SystemExit with required parser
     with patch('sys.argv', mflux_generate_minimal_argv + ['--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
-        pytest.raises(SystemExit, mflux_generate_parser.parse_args)
+        pytest.raises(SystemExit, parser_with_required_model.parse_args)
+
     # test value read from flag
     with patch('sys.argv', mflux_generate_minimal_argv + ['--model', 'dev', '--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
         args = mflux_generate_parser.parse_args()
@@ -411,3 +434,98 @@ def test_save_args(mflux_save_parser):
         # required --path not provided, exits to error
         args = mflux_save_parser.parse_args()
         assert args.path == "/some/model/folder"
+
+
+def test_fill_args(mflux_fill_parser, mflux_fill_minimal_argv):
+    # Test required arguments
+    with patch("sys.argv", mflux_fill_minimal_argv):
+        args = mflux_fill_parser.parse_args()
+        assert args.prompt == "meaning of life"
+        assert args.image_path == Path("image.png")
+        assert args.masked_image_path == Path("mask.png")
+        # Default guidance for fill should be None (will be set to 30 in generate_fill.py)
+        assert args.guidance == pytest.approx(3.5)  # default guidance value
+
+    # Test with missing required arguments
+    with patch("sys.argv", ["mflux-fill", "--prompt", "test"]):
+        # Missing image_path and masked_image_path should raise SystemExit
+        pytest.raises(SystemExit, mflux_fill_parser.parse_args)
+
+    with patch("sys.argv", ["mflux-fill", "--image-path", "image.png", "--masked-image-path", "mask.png"]):
+        # Missing prompt should raise SystemExit
+        pytest.raises(SystemExit, mflux_fill_parser.parse_args)
+
+    # Test with custom values
+    custom_argv = mflux_fill_minimal_argv + ["--guidance", "30", "--steps", "20", "--height", "512", "--width", "512"]
+    with patch("sys.argv", custom_argv):
+        args = mflux_fill_parser.parse_args()
+        assert args.guidance == pytest.approx(30.0)
+        assert args.steps == 20
+        assert args.height == 512
+        assert args.width == 512
+
+
+def test_fill_args_with_metadata(mflux_fill_parser, mflux_fill_minimal_argv, base_metadata_dict, temp_dir):
+    metadata_file = temp_dir / "fill_metadata.json"
+    # Set up metadata with fill-related values
+    with metadata_file.open("wt") as m:
+        # Add masked_image_path to the metadata dictionary
+        base_metadata_dict["masked_image_path"] = "metadata_mask.png"
+        base_metadata_dict["prompt"] = "from metadata file"
+        json.dump(base_metadata_dict, m, indent=4)
+
+    # Test with minimal args and metadata
+    # First modify the parser to support metadata config
+    mflux_fill_parser.supports_metadata_config = True
+    mflux_fill_parser.add_metadata_config()
+
+    # Create a modified version of minimal_argv that includes all required arguments
+    # that aren't in metadata
+    minimal_metadata_argv = [
+        "mflux-fill",
+        "--config-from-metadata",
+        metadata_file.as_posix(),
+        "--prompt",
+        "CLI prompt",
+        "--image-path",
+        "image.png",
+        "--masked-image-path",
+        "cli_mask.png",
+    ]
+
+    # Test command line arguments overriding metadata
+    with patch("sys.argv", minimal_metadata_argv):
+        args = mflux_fill_parser.parse_args()
+        assert args.prompt == "CLI prompt"  # From CLI
+        assert args.image_path == Path("image.png")  # From CLI
+        assert args.masked_image_path == Path("cli_mask.png")  # From CLI
+
+
+def test_fill_default_guidance():
+    # Create a parser just like in generate_fill.py
+    parser = CommandLineParser(description="Generate an image using the fill tool to complete masked areas.")
+    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_fill_arguments()
+    parser.add_output_arguments()
+
+    # Parse minimal arguments
+    with patch(
+        "sys.argv", ["mflux-fill", "--prompt", "test", "--image-path", "img.png", "--masked-image-path", "mask.png"]
+    ):
+        args = parser.parse_args()
+
+        # Verify initial guidance value is the UI default
+        assert args.guidance == ui_defaults.GUIDANCE_SCALE
+
+        # Simulate what happens in generate_fill.py
+        if args.guidance is None:
+            args.guidance = 30
+        else:
+            # In our test we'll just override it to simulate the behavior
+            args.guidance = 30
+
+        # Now check that guidance is correctly set to 30
+        assert args.guidance == 30

tests/image_generation/helpers/image_generation_fill_test_helper.py

@@ -0,0 +1,68 @@
+import os
+
+import numpy as np
+from PIL import Image
+
+from mflux import Config, ModelConfig
+from mflux.flux_tools.fill.flux_fill import Flux1Fill
+from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper
+
+
+class ImageGeneratorFillTestHelper:
+    @staticmethod
+    def assert_matches_reference_image(
+        reference_image_path: str,
+        output_image_path: str,
+        model_config: ModelConfig,
+        steps: int,
+        seed: int,
+        height: int,
+        width: int,
+        prompt: str,
+        image_path: str,
+        masked_image_path: str,
+        lora_paths: list[str] | None = None,
+        lora_scales: list[float] | None = None,
+    ):
+        # resolve paths
+        reference_image_path = ImageGeneratorTestHelper.resolve_path(reference_image_path)
+        output_image_path = ImageGeneratorTestHelper.resolve_path(output_image_path)
+        image_path = str(ImageGeneratorTestHelper.resolve_path(image_path))
+        masked_image_path = str(ImageGeneratorTestHelper.resolve_path(masked_image_path))
+        lora_paths = [str(ImageGeneratorTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None
+
+        try:
+            # given
+            flux = Flux1Fill(
+                model_config=model_config,
+                quantize=8,
+                lora_paths=lora_paths,
+                lora_scales=lora_scales,
+            )
+
+            # when
+            image = flux.generate_image(
+                seed=seed,
+                prompt=prompt,
+                config=Config(
+                    num_inference_steps=steps,
+                    height=height,
+                    width=width,
+                    image_path=image_path,
+                    masked_image_path=masked_image_path,
+                    guidance=30,
+                ),
+            )
+            image.save(path=output_image_path, overwrite=True)
+
+            # then
+            np.testing.assert_array_equal(
+                np.array(Image.open(output_image_path)),
+                np.array(Image.open(reference_image_path)),
+                err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}",
+            )
+
+        finally:
+            # cleanup
+            if os.path.exists(output_image_path):
+                os.remove(output_image_path)

tests/image_generation/test_generate_image_fill.py

@@ -0,0 +1,22 @@
+from mflux import ModelConfig
+from tests.image_generation.helpers.image_generation_fill_test_helper import ImageGeneratorFillTestHelper
+
+
+class TestImageGeneratorFill:
+    OUTPUT_IMAGE_FILENAME = "output.png"
+    SOURCE_IMAGE_FILENAME = "reference_dev_image_to_image_result.png"
+    MASK_IMAGE_FILENAME = "mask.png"
+
+    def test_inpaint(self):
+        ImageGeneratorFillTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_dev_image_to_image_result_inpaint.png",
+            output_image_path=TestImageGeneratorFill.OUTPUT_IMAGE_FILENAME,
+            model_config=ModelConfig.dev_fill(),
+            steps=15,
+            seed=42,
+            height=341,
+            width=768,
+            prompt="A burger and crispy golden french fries next to it",
+            image_path=TestImageGeneratorFill.SOURCE_IMAGE_FILENAME,
+            masked_image_path=TestImageGeneratorFill.MASK_IMAGE_FILENAME,
+        )

tests/image_generation/test_image_util.py

@@ -1,3 +1,5 @@
+import mlx.core as mx
+import numpy as np
 import PIL.Image
 import pytest
 
@@ -132,3 +134,63 @@ def test_create_outpaint_mask_image():
 
     # Left edge of center box
     assert mask.getpixel((left_padding, top_padding + 10)) == (0, 0, 0)
+
+
+def test_binarize():
+    # Create test data using numpy arrays and convert to mx arrays
+    # Create a gradient array from 0 to 1
+    gradient = np.linspace(0, 1, 10).reshape(1, 1, 1, 10).astype(np.float32)
+    gradient_mx = mx.array(gradient)
+
+    # Apply binarization
+    result = ImageUtil._binarize(gradient_mx)
+
+    # Expected: values < 0.5 should be 0, values >= 0.5 should be 1
+    expected = mx.where(gradient_mx < 0.5, mx.zeros_like(gradient_mx), mx.ones_like(gradient_mx))
+
+    # Convert results to numpy for easier comparison
+    result_np = np.array(result)
+    expected_np = np.array(expected)
+
+    # Check values
+    assert np.array_equal(result_np, expected_np)
+
+    # Specifically check that the first 5 values are 0 and the rest are 1
+    assert np.all(result_np[:, :, :, :5] == 0)
+    assert np.all(result_np[:, :, :, 5:] == 1)
+
+
+def test_to_array_with_mask():
+    # Create a test image with gradient colors
+    from PIL import Image, ImageDraw
+
+    # Create a simple mask image (white square on black background)
+    mask_img = Image.new("RGB", (100, 100), color="black")
+    draw = ImageDraw.Draw(mask_img)
+    draw.rectangle((25, 25, 75, 75), fill="white")
+
+    # Convert to array with is_mask=False (should normalize)
+    regular_array = ImageUtil.to_array(mask_img, is_mask=False)
+
+    # Convert to array with is_mask=True (should binarize)
+    mask_array = ImageUtil.to_array(mask_img, is_mask=True)
+
+    # Check shapes
+    assert regular_array.shape == mask_array.shape
+
+    # Regular array should have normalized values between -1 and 1
+    assert mx.min(regular_array) < 0
+    assert mx.max(regular_array) <= 1.0
+
+    # Mask array should only have binary values (0 or 1)
+    unique_values = set(mx.array.flatten(mask_array).tolist())
+    assert unique_values == {0.0, 1.0} or unique_values == {0.0} or unique_values == {1.0}
+
+    # The center of the mask should be 1 (white square)
+    assert mask_array[0, 0, 50, 50] == 1.0
+
+    # The corners should be 0 (black background)
+    assert mask_array[0, 0, 0, 0] == 0.0
+    assert mask_array[0, 0, 0, 99] == 0.0
+    assert mask_array[0, 0, 99, 0] == 0.0
+    assert mask_array[0, 0, 99, 99] == 0.0

tests/model_config/test_model_config.py

@@ -5,65 +5,134 @@ from mflux.error.error import InvalidBaseModel, ModelConfigError
 
 
 def test_bfl_dev():
-    model_attrs = ModelConfig.from_name("dev")
-    assert model_attrs.model_name.startswith("black-forest-labs/")
-    assert model_attrs.max_sequence_length == 512
-    assert model_attrs.supports_guidance is True
+    model = ModelConfig.from_name("dev")
+    assert model.alias == "dev"
+    assert model.model_name.startswith("black-forest-labs/")
+    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_full_name():
-    model_attrs = ModelConfig.from_name("black-forest-labs/FLUX.1-dev")
-    assert model_attrs.model_name.startswith("black-forest-labs/")
-    assert model_attrs.max_sequence_length == 512
-    assert model_attrs.supports_guidance is True
+    model = ModelConfig.from_name("black-forest-labs/FLUX.1-dev")
+    assert model.alias == "dev"
+    assert model.model_name.startswith("black-forest-labs/")
+    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_schnell():
-    model_attrs = ModelConfig.from_name("schnell")
-    assert model_attrs.model_name.startswith("black-forest-labs/")
-    assert model_attrs.max_sequence_length == 256
-    assert model_attrs.supports_guidance is False
+    model = ModelConfig.from_name("schnell")
+    assert model.alias == "schnell"
+    assert model.model_name.startswith("black-forest-labs/")
+    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
 
 
 def test_bfl_schnell_full_name():
-    model_attrs = ModelConfig.from_name("black-forest-labs/FLUX.1-schnell")
-    assert model_attrs.model_name.startswith("black-forest-labs/")
-    assert model_attrs.max_sequence_length == 256
-    assert model_attrs.supports_guidance is False
+    model = ModelConfig.from_name("black-forest-labs/FLUX.1-schnell")
+    assert model.alias == "schnell"
+    assert model.model_name.startswith("black-forest-labs/")
+    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
+
+
+def test_bfl_dev_fill():
+    model = ModelConfig.from_name("dev-fill")
+    assert model.alias == "dev-fill"
+    assert model.model_name == "black-forest-labs/FLUX.1-Fill-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_fill_full_name():
+    model = ModelConfig.from_name("black-forest-labs/FLUX.1-Fill-dev")
+    assert model.alias == "dev-fill"
+    assert model.model_name == "black-forest-labs/FLUX.1-Fill-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_community_dev_fill_implicit_base_model():
+    model = ModelConfig.from_name("acme-lab/some-dev-fill-model")
+    assert model.alias == "dev-fill"
+    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_community_dev_fill_explicit_base_model():
+    model = ModelConfig.from_name("acme-lab/some-model", base_model="dev-fill")
+    assert model.alias == "dev-fill"
+    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_implicit_base_model_prefers_dev_fill_over_dev():
+    model = ModelConfig.from_name("acme-lab/dev-fill-based-model")
+    assert model.alias == "dev-fill"
+    assert model.base_model == "black-forest-labs/FLUX.1-Fill-dev"
+    assert model.max_sequence_length == 512
+    assert model.num_train_steps == 1000
+    assert model.requires_sigma_shift is True
 
 
 def test_community_dev_implicit_base_model():
-    model_attrs = ModelConfig.from_name("acme-lab/some-awesome-dev-model")
-    assert model_attrs.max_sequence_length == 512
-    assert model_attrs.supports_guidance is True
+    model = ModelConfig.from_name("acme-lab/some-awesome-dev-model")
+    assert model.alias == "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_community_schnell_implicit_base_model():
-    model_attrs = ModelConfig.from_name("acme-lab/some-quick-schnell-model")
-    assert model_attrs.max_sequence_length == 256
-    assert model_attrs.supports_guidance is False
+    model = ModelConfig.from_name("acme-lab/some-quick-schnell-model")
+    assert model.alias == "schnell"
+    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
 
 
 def test_community_dev_explicit_base_model():
-    model_attrs = ModelConfig.from_name("acme-lab/some-awesome-model", base_model="dev")
-    assert model_attrs.max_sequence_length == 512
-    assert model_attrs.supports_guidance is True
+    model = ModelConfig.from_name("acme-lab/some-awesome-model", base_model="dev")
+    assert model.alias == "dev"
+    assert model.base_model == "black-forest-labs/FLUX.1-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_community_schnell_explicit_base_model():
-    model_attrs = ModelConfig.from_name("acme-lab/some-awesome-model", base_model="schnell")
-    assert model_attrs.max_sequence_length == 256
-    assert model_attrs.supports_guidance is False
+    model = ModelConfig.from_name("acme-lab/some-awesome-model", base_model="schnell")
+    assert model.base_model == "black-forest-labs/FLUX.1-schnell"
+    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
 
 
 def test_model_config_error():
-    assert pytest.raises(ModelConfigError, ModelConfig.from_name, "acme-lab/some-model-who-knows-what-its-based-on")
+    with pytest.raises(ModelConfigError):
+        ModelConfig.from_name("acme-lab/some-model-who-knows-what-its-based-on")
 
 
 def test_invalid_base_model_error():
-    assert pytest.raises(
-        InvalidBaseModel,
-        ModelConfig.from_name,
-        "acme-lab/some-model-who-knows-what-its-based-on",
-        base_model="something_unknown",
-    )
+    with pytest.raises(InvalidBaseModel):
+        ModelConfig.from_name("acme-lab/some-model-who-knows-what-its-based-on", base_model="something_unknown")

tools/inpaint_mask_tool.py

@@ -0,0 +1,154 @@
+from pathlib import Path
+
+import cv2
+import numpy as np
+
+
+class MaskCreator:
+    BRUSH_SIZES = {"1": 2, "2": 4, "3": 8, "4": 16, "5": 32, "6": 48, "7": 96, "8": 192, "9": 384}
+
+    def __init__(self, image_path: Path):
+        self.original_image = cv2.imread(image_path)
+        if self.original_image is None:
+            raise FileNotFoundError(f"Could not open or find the image: {image_path}")
+
+        self.mask_output_path = image_path.with_name(f"{image_path.stem}_mask").with_suffix(".png")
+
+        # Create a window and set up display image
+        self.window_name = "MFlux Inpaint Mask Creator - Draw with mouse or trackpad (hot keys: (s)ave, (r)eset, (q)uit"
+        self.display_image = self.original_image.copy()
+
+        # Create a blank mask the same size as the image
+        self.mask = np.zeros(self.original_image.shape[:2], dtype=np.uint8)
+
+        # Set up drawing parameters
+        self.drawing = False
+        self.brush_size = self.BRUSH_SIZES["5"]
+        self.last_point = None
+
+        self.overlay = np.zeros_like(self.original_image)
+
+        # Update display every N drawing events - lower is more responsive
+        self.update_frequency = 1
+        self.event_counter = 0
+
+        # Show the initial display
+        self.update_display()
+
+    def mouse_callback(self, event, x, y, flags, param):
+        # Start drawing
+        if event == cv2.EVENT_LBUTTONDOWN:
+            self.drawing = True
+            self.last_point = (x, y)
+            cv2.circle(self.mask, (x, y), self.brush_size, 255, -1)
+            self.event_counter += 1
+            if self.event_counter % self.update_frequency == 0:
+                self.update_display()
+
+        # Continue drawing
+        elif event == cv2.EVENT_MOUSEMOVE and self.drawing:
+            # Use thickness based on brush size for smoother lines
+            if self.last_point:  # Ensure we have a last point
+                # Draw a line between the last point and current point
+                cv2.line(self.mask, self.last_point, (x, y), 255, self.brush_size * 2)
+                # Also draw a circle at the current point to avoid gaps in fast movements
+                cv2.circle(self.mask, (x, y), self.brush_size, 255, -1)
+            self.last_point = (x, y)
+            self.event_counter += 1
+            if self.event_counter % self.update_frequency == 0:
+                self.update_display()
+
+        # Stop drawing
+        elif event == cv2.EVENT_LBUTTONUP:
+            self.drawing = False
+            self.update_display()  # Always update display when stopping
+
+    def update_display(self):
+        # Create a copy of the original image
+        self.display_image = self.original_image.copy()
+
+        # Create a colored overlay for the mask (semi-transparent red)
+        self.overlay[:] = 0  # Reset overlay
+        self.overlay[self.mask > 0] = [0, 0, 255]  # Red overlay
+
+        # Apply the overlay
+        alpha = 0.5  # Transparency level
+        cv2.addWeighted(self.overlay, alpha, self.display_image, 1 - alpha, 0, self.display_image)
+
+        # Draw brush size indicator in the corner
+        text = f"Brush Size: {self.brush_size} (Hotkeys 1-9: change brush size) "
+        cv2.putText(
+            self.display_image,
+            text,
+            (10, 30),
+            cv2.FONT_HERSHEY_DUPLEX,
+            0.5,
+            (255, 0, 0),
+            1,
+            cv2.LINE_AA,  # line type: anti-aliased
+        )
+
+        # Display the result with OpenCV's high GUI priority
+        cv2.imshow(self.window_name, self.display_image)
+        cv2.waitKey(1)  # Process events to force display update
+
+    def save_mask(self, output_path):
+        cv2.imwrite(output_path, self.mask)
+        print(f"Mask saved to {output_path}")
+
+    def reset_mask(self):
+        self.mask = np.zeros(self.original_image.shape[:2], dtype=np.uint8)
+        self.update_display()
+
+    def set_brush_size(self, size_key):
+        if size_key in self.BRUSH_SIZES:
+            self.brush_size = self.BRUSH_SIZES[size_key]
+            print(f"Brush size {size_key}: {self.brush_size}")
+            self.update_display()
+
+    def run(self):
+        cv2.namedWindow(self.window_name)
+        cv2.setMouseCallback(self.window_name, self.mouse_callback)
+
+        while True:
+            key = cv2.waitKey(1) & 0xFF
+            key_char = chr(key) if key < 128 else ""
+
+            # Check for brush size hotkeys (1-5)
+            if key_char in self.BRUSH_SIZES:
+                self.set_brush_size(key_char)
+
+            # Save mask (press 's')
+            elif key == ord("s"):
+                self.save_mask(self.mask_output_path)
+
+            # Reset mask (press 'r')
+            elif key == ord("r"):
+                self.reset_mask()
+                print("Mask reset")
+
+            # Quit (press 'q' or ESC)
+            elif key == ord("q") or key == 27:
+                break
+
+        cv2.destroyAllWindows()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Create binary mask image from source image to use as the complementary --masked-image-path arg."
+    )
+    parser.add_argument("image_path", type=Path, help="Path to the input image")
+    args = parser.parse_args()
+
+    try:
+        mask_creator = MaskCreator(args.image_path)
+        mask_creator.run()
+    except FileNotFoundError as e:
+        print(f"Error: {e}")
+    except Exception as e:  # noqa
+        print(f"An unexpected error occurred: {e}")
+    except KeyboardInterrupt:
+        pass