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 296fa22..84e816f 100644 --- a/src/mflux/latent_creator/latent_creator.py +++ b/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, 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,