From eff38124213934403ef9dcfbcc01adf8fa0f54df Mon Sep 17 00:00:00 2001
From: Filip Strand <strand.filip@gmail.com>
Date: Thu, 22 May 2025 13:02:18 +0200
Subject: [PATCH] Cleanup FluxFill mask args (#195)

---
 src/mflux/flux_tools/fill/flux_fill.py |  4 ++--
 src/mflux/flux_tools/fill/mask_util.py | 19 +++++++++----------
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/src/mflux/flux_tools/fill/flux_fill.py b/src/mflux/flux_tools/fill/flux_fill.py
index 7d56e4d..92336c3 100644
--- a/src/mflux/flux_tools/fill/flux_fill.py
+++ b/src/mflux/flux_tools/fill/flux_fill.py
@@ -73,8 +73,8 @@ class Flux1Fill(nn.Module):
         # 3. Create the static masked latents
         static_masked_latents = MaskUtil.create_masked_latents(
             vae=self.vae,
-            config=config,
-            latents=latents,
+            height=config.height,
+            width=config.width,
             img_path=config.image_path,
             mask_path=config.masked_image_path,
         )
diff --git a/src/mflux/flux_tools/fill/mask_util.py b/src/mflux/flux_tools/fill/mask_util.py
index b681dd1..835529c 100644
--- a/src/mflux/flux_tools/fill/mask_util.py
+++ b/src/mflux/flux_tools/fill/mask_util.py
@@ -2,7 +2,6 @@ from pathlib import Path
 
 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
@@ -12,8 +11,8 @@ class MaskUtil:
     @staticmethod
     def create_masked_latents(
         vae: VAE,
-        config: RuntimeConfig,
-        latents: mx.array,
+        height: int,
+        width: int,
         img_path: str | Path,
         mask_path: str | Path | None,
     ) -> mx.array:
@@ -24,27 +23,27 @@ class MaskUtil:
         # 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,
+            target_width=width,
+            target_height=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,
+            target_width=width,
+            target_height=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)
+        masked_image = ArrayUtil.pack_latents(latents=masked_image, height=height, width=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)
+        mask = MaskUtil._reshape_mask(the_mask=the_mask, height=height, width=width)
+        mask = ArrayUtil.pack_latents(latents=mask, height=height, width=width, num_channels_latents=64)
 
         # 5. Concat the masked_image and the mask
         masked_image_latents = mx.concatenate([masked_image, mask], axis=-1)