Add named args for height and width for clarity

src/mflux/config/config.py

@@ -37,5 +37,5 @@ class ConfigControlnet(Config):
         guidance: float = 4.0,
         controlnet_strength: float = 1.0,
     ):
-        super().__init__(num_inference_steps, width, height, guidance)
+        super().__init__(num_inference_steps=num_inference_steps, width=width, height=height, guidance=guidance)
         self.controlnet_strength = controlnet_strength

src/mflux/config/runtime_config.py

@@ -72,7 +72,7 @@ class RuntimeConfig:
     def _create_sigmas(config, model) -> mx.array:
         sigmas = RuntimeConfig._create_sigmas_values(config.num_inference_steps)
         if model == ModelConfig.FLUX1_DEV:
-            sigmas = RuntimeConfig._shift_sigmas(sigmas, config.width, config.height)
+            sigmas = RuntimeConfig._shift_sigmas(sigmas=sigmas, width=config.width, height=config.height)
         return sigmas
 
     @staticmethod

src/mflux/controlnet/flux_controlnet.py

@@ -139,7 +139,7 @@ class Flux1Controlnet:
         controlnet_cond = ImageUtil.to_array(control_image)
         controlnet_cond = self.vae.encode(controlnet_cond)
         controlnet_cond = (controlnet_cond / self.vae.scaling_factor) + self.vae.shift_factor
-        controlnet_cond = ArrayUtil.pack_latents(controlnet_cond, config.height, config.width)
+        controlnet_cond = ArrayUtil.pack_latents(latents=controlnet_cond, height=config.height, width=config.width)
 
         # 1. Create the initial latents
         latents = LatentCreator.create(seed=seed, height=config.height, width=config.width)
@@ -188,7 +188,7 @@ class Flux1Controlnet:
                 raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
 
         # 5. Decode the latent array and return the image
-        latents = ArrayUtil.unpack_latents(latents, config.height, config.width)
+        latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
         decoded = self.vae.decode(latents)
         return ImageUtil.to_image(
             decoded_latents=decoded,

src/mflux/controlnet/transformer_controlnet.py

@@ -54,7 +54,7 @@ class TransformerControlnet(nn.Module):
         text_embeddings = self.time_text_embed.forward(time_step, pooled_prompt_embeds, guidance)
         encoder_hidden_states = self.context_embedder(prompt_embeds)
         txt_ids = Transformer.prepare_text_ids(seq_len=prompt_embeds.shape[1])
-        img_ids = Transformer.prepare_latent_image_ids(config.height, config.width)
+        img_ids = Transformer.prepare_latent_image_ids(height=config.height, width=config.width)
         ids = mx.concatenate((txt_ids, img_ids), axis=1)
         image_rotary_emb = self.pos_embed.forward(ids)
 

src/mflux/flux/flux.py

@@ -129,7 +129,7 @@ class Flux1:
                 raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
 
         # 5. Decode the latent array and return the image
-        latents = ArrayUtil.unpack_latents(latents, config.height, config.width)
+        latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
         decoded = self.vae.decode(latents)
         return ImageUtil.to_image(
             decoded_latents=decoded,

src/mflux/latent_creator/latent_creator.py

@@ -36,9 +36,13 @@ class LatentCreator:
         else:
             # Image2Image
             user_image = ImageUtil.load_image(runtime_conf.config.init_image_path).convert("RGB")
-            scaled_user_image = ImageUtil.scale_to_dimensions(user_image, runtime_conf.width, runtime_conf.height)
+            scaled_user_image = ImageUtil.scale_to_dimensions(
+                image=user_image,
+                target_width=runtime_conf.width,
+                target_height=runtime_conf.height,
+            )
             encoded = vae.encode(ImageUtil.to_array(scaled_user_image))
-            latents = ArrayUtil.pack_latents(encoded, runtime_conf.height, runtime_conf.width)
+            latents = ArrayUtil.pack_latents(latents=encoded, height=runtime_conf.height, width=runtime_conf.width)
             sigma = runtime_conf.sigmas[runtime_conf.init_time_step]
             return LatentCreator.add_noise_by_interpolation(
                 clean=latents,

src/mflux/post_processing/image_util.py

@@ -108,9 +108,7 @@ class ImageUtil:
         image: PIL.Image.Image,
         target_width: int,
         target_height: int,
-    ):
+    ) -> PIL.Image.Image:
-        # for now, naively resize the user image to the output image size
-        # todo: consider alt strategies: cropping / resize+padding
         if (image.width, image.height) != (target_width, target_height):
             return image.resize((target_width, target_height), PIL.Image.LANCZOS)
         else:

src/mflux/post_processing/stepwise_handler.py

@@ -36,7 +36,7 @@ class StepwiseHandler:
 
     def process_step(self, gen_step: int, latents: mx.array):
         if self.output_dir:
-            unpack_latents = ArrayUtil.unpack_latents(latents, self.config.height, self.config.width)
+            unpack_latents = ArrayUtil.unpack_latents(latents=latents, height=self.config.height, width=self.config.width)  # fmt: off
             stepwise_decoded = self.flux.vae.decode(unpack_latents)
             stepwise_img = ImageUtil.to_image(
                 decoded_latents=stepwise_decoded,