Merge pull request #87 from filipstrand/fix-img2img-for-non-square-resolution

Fix img2img for non-square resolutions

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

@@ -12,7 +12,7 @@ class LatentCreator:
         seed: int,
         height: int,
         width: int,
-    ):
+    ) -> mx.array:
         return mx.random.normal(
             shape=[1, (height // 16) * (width // 16), 64],
             key=mx.random.key(seed)
@@ -23,8 +23,8 @@ class LatentCreator:
         seed: int,
         runtime_conf: RuntimeConfig,
         vae: nn.Module,
-    ):
-        noise = LatentCreator.create(
+    ) -> mx.array:
+        pure_noise = LatentCreator.create(
             seed=seed,
             height=runtime_conf.height,
             width=runtime_conf.width,
@@ -32,14 +32,24 @@ class LatentCreator:
 
         if runtime_conf.config.init_image_path is None:
             # Text2Image
-            return noise
+            return pure_noise
         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.width, runtime_conf.height)
-            sigmas_for_init_image_strength = runtime_conf.sigmas[runtime_conf.init_time_step]
-            latents_adjusted = latents * (1.0 - sigmas_for_init_image_strength)
-            noise_adjusted = noise * sigmas_for_init_image_strength
-            return latents_adjusted + noise_adjusted
+            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,
+                noise=pure_noise,
+                sigma=sigma
+            )  # fmt: off
+
+    @staticmethod
+    def add_noise_by_interpolation(clean: mx.array, noise: mx.array, sigma: float) -> mx.array:
+        return (1 - sigma) * clean + sigma * noise

src/mflux/post_processing/image_util.py

@@ -108,9 +108,7 @@ class ImageUtil:
         image: PIL.Image.Image,
         target_width: int,
         target_height: int,
-    ):
-        # for now, naively resize the user image to the output image size
-        # todo: consider alt strategies: cropping / resize+padding
+    ) -> PIL.Image.Image:
         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,

tests/helpers/image_generation_test_helper.py

@@ -63,7 +63,7 @@ class ImageGeneratorTestHelper:
                 os.remove(output_image_path)
 
     @staticmethod
-    def resolve_path(path) -> Path:
+    def resolve_path(path) -> Path | None:
         if path is None:
             return None
         return Path(__file__).parent.parent / "resources" / path

tests/test_generate_image.py

@@ -46,13 +46,13 @@ class TestImageGenerator:
     def test_image_generation_dev_image_to_image(self):
         ImageGeneratorTestHelper.assert_matches_reference_image(
             reference_image_path="reference_dev_image_to_image_result.png",
-            init_image_path="reference_dev_image_to_image_init.png",
-            init_image_strength=0.126,  # fill 1/8 steps
+            init_image_path="reference_dev_lora.png",
+            init_image_strength=0.4,
             output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
             model_config=ModelConfig.FLUX1_DEV,
             steps=8,
-            seed=42,
-            height=768,
+            seed=44,
+            height=341,
             width=768,
-            prompt="astronauts in a jungle",
+            prompt="Luxury food photograph of a burger",
         )