Merge pull request #82 from anthonywu/fixes-for-image-to-image

tests and fixes for image to image

src/mflux/config/config.py

@@ -1,5 +1,4 @@
 import logging
-import time
 from pathlib import Path
 
 import mlx.core as mx
@@ -16,9 +15,8 @@ class Config:
         width: int = 1024,
         height: int = 1024,
         guidance: float = 4.0,
-        init_image: Path | None = None,
+        init_image_path: Path | None = None,
         init_image_strength: float | None = None,
-        seed: float | None = None,
     ):
         if width % 16 != 0 or height % 16 != 0:
             log.warning("Width and height should be multiples of 16. Rounding down.")
@@ -26,9 +24,8 @@ class Config:
         self.height = 16 * (height // 16)
         self.num_inference_steps = num_inference_steps
         self.guidance = guidance
-        self.init_image = init_image
+        self.init_image_path = init_image_path
         self.init_image_strength = init_image_strength
-        self.seed = seed or int(time.time())
 
 
 class ConfigControlnet(Config):

src/mflux/config/runtime_config.py

@@ -6,8 +6,6 @@ import numpy as np
 
 from mflux.config.config import Config, ConfigControlnet
 from mflux.config.model_config import ModelConfig
-from mflux.post_processing.array_util import ArrayUtil
-from mflux.post_processing.image_util import ImageUtil
 
 logger = logging.getLogger(__name__)
 
@@ -15,11 +13,10 @@ VAE: t.TypeAlias = "mflux.models.vae.vae.VAE"  # noqa: F821
 
 
 class RuntimeConfig:
-    def __init__(self, config: Config | ConfigControlnet, model_config: ModelConfig, vae: VAE):
+    def __init__(self, config: Config | ConfigControlnet, model_config):
         self.config = config
         self.model_config = model_config
         self.sigmas = self._create_sigmas(config, model_config)
-        self.vae = vae
 
     @property
     def height(self) -> int:
@@ -29,10 +26,6 @@ class RuntimeConfig:
     def width(self) -> int:
         return self.config.width
 
-    @property
-    def seed(self) -> int:
-        return self.config.seed
-
     @property
     def guidance(self) -> float:
         return self.config.guidance
@@ -51,7 +44,7 @@ class RuntimeConfig:
 
     @property
     def init_time_step(self) -> int:
-        if self.config.init_image is None:
+        if self.config.init_image_path is None:
             # text to image, always begin at time step 0
             return 0
         else:
@@ -63,26 +56,6 @@ class RuntimeConfig:
             t = max(1, int(self.num_inference_steps * strength))
             return t
 
-    @property
-    def init_latents(self) -> mx.array:
-        noise = mx.random.normal(
-            shape=[1, (self.height // 16) * (self.width // 16), 64],
-            key=mx.random.key(self.seed)
-        )  # fmt: off
-        if self.config.init_image is not None:
-            user_image = ImageUtil.load_image(self.config.init_image).convert("RGB")
-            latents = ArrayUtil.pack_latents(
-                self.vae.encode(ImageUtil.to_array(ImageUtil.scale_to_dimensions(user_image, self.width, self.height))),
-                self.width,
-                self.height,
-            )
-            sigmas_for_init_image_strength = self.sigmas[self.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
-        else:
-            return noise
-
     @property
     def controlnet_strength(self) -> float:
         if isinstance(self.config, ConfigControlnet):

src/mflux/controlnet/controlnet_util.py

@@ -17,6 +17,13 @@ class ControlnetUtil:
         image_to_canny = np.concatenate([image_to_canny, image_to_canny, image_to_canny], axis=2)
         return PIL.Image.fromarray(image_to_canny)
 
+    @staticmethod
+    def scale_image(height: int, width: int, img: PIL.Image) -> PIL.Image:
+        if height != img.height or width != img.width:
+            log.warning(f"Control image has different dimensions than the model. Resizing to {width}x{height}")
+            img = img.resize((width, height), PIL.Image.LANCZOS)
+        return img
+
     @staticmethod
     def save_canny_image(control_image: PIL.Image, path: str):
         from mflux import ImageUtil

src/mflux/controlnet/flux_controlnet.py

@@ -131,7 +131,7 @@ class Flux1Controlnet:
 
         # Embed the controlnet reference image
         control_image = ImageUtil.load_image(controlnet_image_path)
-        control_image = ImageUtil.scale_to_dimensions(control_image, config.height, config.width)
+        control_image = ControlnetUtil.scale_image(config.height, config.width, control_image)
         control_image = ControlnetUtil.preprocess_canny(control_image)
         if controlnet_save_canny:
             ControlnetUtil.save_canny_image(control_image, output)

src/mflux/flux/flux.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+from typing import Union
 
 import mlx.core as mx
 from mlx import nn
@@ -74,6 +75,27 @@ class Flux1:
         if weights.quantization_level is not None:
             self._set_model_weights(weights)
 
+    def create_initial_latents(self, seed: int, config: Config, init_image_path: Union[str | Path | None] = None):
+        noise = mx.random.normal(
+            shape=[1, (config.height // 16) * (config.width // 16), 64],
+            key=mx.random.key(seed)
+        )  # fmt: off
+        if init_image_path is not None:
+            user_image = ImageUtil.load_image(init_image_path).convert("RGB")
+            latents = ArrayUtil.pack_latents(
+                self.vae.encode(
+                    ImageUtil.to_array(ImageUtil.scale_to_dimensions(user_image, config.width, config.height))
+                ),
+                config.width,
+                config.height,
+            )
+            sigmas_for_init_image_strength = config.sigmas[config.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
+        else:
+            return noise
+
     def generate_image(
         self,
         seed: int,
@@ -82,10 +104,12 @@ class Flux1:
         stepwise_output_dir: Path = None,
     ) -> GeneratedImage:
         # Create a new runtime config based on the model type and input parameters
-        config = RuntimeConfig(config, self.model_config, self.vae)
+        config = RuntimeConfig(config, self.model_config)
 
         # 1. Create the initial latents: text-to-image and image-to-image
-        latents = config.init_latents
+        # todo: the nested config.config and confusing Config relations
+        #       should be refactored to bring more clarity
+        latents = self.create_initial_latents(seed, config, config.config.init_image_path)
         time_steps = tqdm(range(config.init_time_step, config.num_inference_steps))
 
         stepwise_handler = StepwiseHandler(

src/mflux/generate.py

@@ -35,9 +35,8 @@ def main():
                 height=args.height,
                 width=args.width,
                 guidance=args.guidance,
-                init_image=args.init_image,
-                init_image_strength=args.init_image_strength,
-                seed=args.seed
+                init_image_path=args.init_image_path,
+                init_image_strength=args.init_image_strength
             ),
         )
 

src/mflux/ui/cli/parsers.py

@@ -28,7 +28,7 @@ class CommandLineParser(argparse.ArgumentParser):
         self._add_image_generator_common_arguments()
 
     def add_image_to_image_arguments(self, required=False) -> None:
-        self.add_argument("--init-image", type=Path, required=required, help="Local path to init image")
+        self.add_argument("--init-image-path", type=Path, required=required, help="Local path to init image")
         self.add_argument("--init-image-strength", type=float, required=False, default=ui_defaults.INIT_IMAGE_STRENGTH, help=f"Controls how strongly the init image influences the output image. A value of 0.0 means no influence. (Default is {ui_defaults.INIT_IMAGE_STRENGTH})")
 
     def add_batch_image_generator_arguments(self) -> None:

tests/helpers/image_generation_test_helper.py

@@ -16,6 +16,10 @@ class ImageGeneratorTestHelper:
         prompt: str,
         steps: int,
         seed: int,
+        height: int = None,
+        width: int = None,
+        init_image_path: str | None = None,
+        init_image_strength: float | None = None,
         lora_paths: list[str] | None = None,
         lora_scales: list[float] | None = None,
     ):
@@ -32,15 +36,16 @@ class ImageGeneratorTestHelper:
                 lora_paths=lora_paths,
                 lora_scales=lora_scales
             )  # fmt: off
-
             # when
             image = flux.generate_image(
                 seed=seed,
                 prompt=prompt,
                 config=Config(
                     num_inference_steps=steps,
-                    height=341,
-                    width=768,
+                    init_image_path=ImageGeneratorTestHelper.resolve_path(init_image_path),
+                    init_image_strength=init_image_strength,
+                    height=height,
+                    width=width,
                 ),
             )
             image.save(path=output_image_path)
@@ -49,7 +54,7 @@ class ImageGeneratorTestHelper:
             np.testing.assert_array_equal(
                 np.array(Image.open(output_image_path)),
                 np.array(Image.open(reference_image_path)),
-                err_msg="Generated image doesn't match reference image",
+                err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}",
             )
 
         finally:
@@ -59,4 +64,6 @@ class ImageGeneratorTestHelper:
 
     @staticmethod
     def resolve_path(path) -> Path:
+        if path is None:
+            return None
         return Path(__file__).parent.parent / "resources" / path

tests/test_generate_image.py

@@ -12,6 +12,8 @@ class TestImageGenerator:
             model_config=ModelConfig.FLUX1_SCHNELL,
             steps=2,
             seed=42,
+            height=341,
+            width=768,
             prompt="Luxury food photograph",
         )
 
@@ -22,6 +24,8 @@ class TestImageGenerator:
             model_config=ModelConfig.FLUX1_DEV,
             steps=15,
             seed=42,
+            height=341,
+            width=768,
             prompt="Luxury food photograph",
         )
 
@@ -32,7 +36,23 @@ class TestImageGenerator:
             model_config=ModelConfig.FLUX1_DEV,
             steps=15,
             seed=42,
+            height=341,
+            width=768,
             prompt="mkym this is made of wool, burger",
             lora_paths=["FLUX-dev-lora-MiaoKa-Yarn-World.safetensors"],
             lora_scales=[1.0],
         )
+
+    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
+            output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
+            model_config=ModelConfig.FLUX1_DEV,
+            steps=8,
+            seed=42,
+            height=768,
+            width=768,
+            prompt="astronauts in a jungle",
+        )