diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py index bac53f7..d2b0584 100644 --- a/src/mflux/config/config.py +++ b/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 diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py index 8b65bca..25723b7 100644 --- a/src/mflux/config/runtime_config.py +++ b/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 diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index d0ab954..1fc62b8 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/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, diff --git a/src/mflux/controlnet/transformer_controlnet.py b/src/mflux/controlnet/transformer_controlnet.py index 232fcfe..50c2db3 100644 --- a/src/mflux/controlnet/transformer_controlnet.py +++ b/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) diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index 0105c21..0410e17 100644 --- a/src/mflux/flux/flux.py +++ b/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, diff --git a/src/mflux/latent_creator/latent_creator.py b/src/mflux/latent_creator/latent_creator.py index 90b4ee8..84e816f 100644 --- a/src/mflux/latent_creator/latent_creator.py +++ b/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 diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index 832d55f..c232cc4 100644 --- a/src/mflux/post_processing/image_util.py +++ b/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: diff --git a/src/mflux/post_processing/stepwise_handler.py b/src/mflux/post_processing/stepwise_handler.py index ee604f6..60332bb 100644 --- a/src/mflux/post_processing/stepwise_handler.py +++ b/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, diff --git a/tests/helpers/image_generation_test_helper.py b/tests/helpers/image_generation_test_helper.py index 7bcdc8a..d1d8929 100644 --- a/tests/helpers/image_generation_test_helper.py +++ b/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 diff --git a/tests/resources/reference_dev_image_to_image_init.png b/tests/resources/reference_dev_image_to_image_init.png deleted file mode 100644 index 21c7781..0000000 Binary files a/tests/resources/reference_dev_image_to_image_init.png and /dev/null differ diff --git a/tests/resources/reference_dev_image_to_image_result.png b/tests/resources/reference_dev_image_to_image_result.png index e52531b..6fd8baa 100644 Binary files a/tests/resources/reference_dev_image_to_image_result.png 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 34e5f8d..d234361 100644 --- a/tests/test_generate_image.py +++ b/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", )