diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py index 2285b3a..bac53f7 100644 --- a/src/mflux/config/config.py +++ b/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): diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py index ece21f2..a4f6ecd 100644 --- a/src/mflux/config/runtime_config.py +++ b/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): diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index e2a8a52..fd59a7f 100644 --- a/src/mflux/flux/flux.py +++ b/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( diff --git a/src/mflux/generate.py b/src/mflux/generate.py index 1a0457c..b7c5c17 100644 --- a/src/mflux/generate.py +++ b/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 ), ) diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index 0814c8e..2521a57 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/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: diff --git a/tests/helpers/image_generation_test_helper.py b/tests/helpers/image_generation_test_helper.py index 8f35a11..7bcdc8a 100644 --- a/tests/helpers/image_generation_test_helper.py +++ b/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 diff --git a/tests/resources/reference_dev_image_to_image_init.png b/tests/resources/reference_dev_image_to_image_init.png new file mode 100644 index 0000000..21c7781 Binary files /dev/null and b/tests/resources/reference_dev_image_to_image_init.png differ diff --git a/tests/resources/reference_dev_image_to_image_result.png b/tests/resources/reference_dev_image_to_image_result.png new file mode 100644 index 0000000..e52531b Binary files /dev/null and b/tests/resources/reference_dev_image_to_image_result.png differ diff --git a/tests/test_generate_image.py b/tests/test_generate_image.py index c52d5a8..34e5f8d 100644 --- a/tests/test_generate_image.py +++ b/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", + )