From 3601baf75ed07845c166eebf4a94064bd14e9ea7 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sun, 9 Feb 2025 11:29:50 +0100 Subject: [PATCH] Implement general callback mechanism --- src/mflux/callbacks/__init__.py | 0 src/mflux/callbacks/callback.py | 43 +++++++++++ src/mflux/callbacks/callback_registry.py | 31 ++++++++ src/mflux/callbacks/callbacks.py | 59 +++++++++++++++ src/mflux/callbacks/instances/__init__.py | 0 src/mflux/callbacks/instances/canny_saver.py | 24 +++++++ .../callbacks/instances/stepwise_handler.py | 70 ++++++++++++++++++ src/mflux/controlnet/controlnet_util.py | 23 ++---- src/mflux/controlnet/flux_controlnet.py | 69 +++++++++--------- src/mflux/flux/flux.py | 71 ++++++++++--------- src/mflux/generate.py | 23 +++--- src/mflux/generate_controlnet.py | 27 ++++--- src/mflux/latent_creator/latent_creator.py | 69 ++++++++++++------ src/mflux/post_processing/stepwise_handler.py | 60 ---------------- ...image_generation_controlnet_test_helper.py | 2 - 15 files changed, 385 insertions(+), 186 deletions(-) create mode 100644 src/mflux/callbacks/__init__.py create mode 100644 src/mflux/callbacks/callback.py create mode 100644 src/mflux/callbacks/callback_registry.py create mode 100644 src/mflux/callbacks/callbacks.py create mode 100644 src/mflux/callbacks/instances/__init__.py create mode 100644 src/mflux/callbacks/instances/canny_saver.py create mode 100644 src/mflux/callbacks/instances/stepwise_handler.py delete mode 100644 src/mflux/post_processing/stepwise_handler.py diff --git a/src/mflux/callbacks/__init__.py b/src/mflux/callbacks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/callbacks/callback.py b/src/mflux/callbacks/callback.py new file mode 100644 index 0000000..48281dd --- /dev/null +++ b/src/mflux/callbacks/callback.py @@ -0,0 +1,43 @@ +from typing import Protocol + +import mlx.core as mx +import PIL.Image +import tqdm + +from mflux.config.runtime_config import RuntimeConfig + + +class BeforeLoopCallback(Protocol): + def call_before_loop( + self, + seed: int, + prompt: str, + canny_image: PIL.Image.Image | None = None, + ) -> None: # fmt: off + ... + + +class InLoopCallback(Protocol): + def call_in_loop( + self, + seed: int, + prompt: str, + step: int, + latents: mx.array, + config: RuntimeConfig, + time_steps: tqdm + ) -> None: # fmt: off + ... + + +class InterruptCallback(Protocol): + def call_interrupt( + self, + seed: int, + prompt: str, + step: int, + latents: mx.array, + config: RuntimeConfig, + time_steps: tqdm + ) -> None: # fmt: off + ... diff --git a/src/mflux/callbacks/callback_registry.py b/src/mflux/callbacks/callback_registry.py new file mode 100644 index 0000000..8fd3974 --- /dev/null +++ b/src/mflux/callbacks/callback_registry.py @@ -0,0 +1,31 @@ +from mflux.callbacks.callback import BeforeLoopCallback, InLoopCallback, InterruptCallback + + +class CallbackRegistry: + in_loop = [] + before_loop = [] + interrupt = [] + + @staticmethod + def register_in_loop(callback: InLoopCallback) -> None: + CallbackRegistry.in_loop.append(callback) + + @staticmethod + def register_before_loop(callback: BeforeLoopCallback) -> None: + CallbackRegistry.before_loop.append(callback) + + @staticmethod + def register_interrupt(callback: InterruptCallback) -> None: + CallbackRegistry.interrupt.append(callback) + + @staticmethod + def in_loop_callbacks() -> list[InLoopCallback]: + return CallbackRegistry.in_loop + + @staticmethod + def before_loop_callbacks() -> list[BeforeLoopCallback]: + return CallbackRegistry.before_loop + + @staticmethod + def interrupt_callbacks() -> list[InterruptCallback]: + return CallbackRegistry.interrupt diff --git a/src/mflux/callbacks/callbacks.py b/src/mflux/callbacks/callbacks.py new file mode 100644 index 0000000..f3e2499 --- /dev/null +++ b/src/mflux/callbacks/callbacks.py @@ -0,0 +1,59 @@ +import mlx.core as mx +import PIL.Image +import tqdm + +from mflux.callbacks.callback_registry import CallbackRegistry +from mflux.config.runtime_config import RuntimeConfig + + +class Callbacks: + @staticmethod + def before_loop( + seed: int, + prompt: str, + canny_image: PIL.Image.Image | None = None, + ): # fmt: off + for subscriber in CallbackRegistry.before_loop_callbacks(): + subscriber.call_before_loop( + seed=seed, + prompt=prompt, + canny_image=canny_image, + ) + + @staticmethod + def in_loop( + seed: int, + prompt: str, + step: int, + latents: mx.array, + config: RuntimeConfig, + time_steps: tqdm + ): # fmt: off + for subscriber in CallbackRegistry.in_loop_callbacks(): + subscriber.call_in_loop( + seed=seed, + prompt=prompt, + step=step, + latents=latents, + config=config, + time_steps=time_steps, + ) + + @staticmethod + def interruption( + seed: int, + prompt: str, + step: int, + latents: mx.array, + config: RuntimeConfig, + time_steps: tqdm + ): # fmt: off + for subscriber in CallbackRegistry.interrupt_callbacks(): + subscriber.call_interrupt( + seed=seed, + prompt=prompt, + step=step, + latents=latents, + config=config, + time_steps=time_steps, + ) diff --git a/src/mflux/callbacks/instances/__init__.py b/src/mflux/callbacks/instances/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/callbacks/instances/canny_saver.py b/src/mflux/callbacks/instances/canny_saver.py new file mode 100644 index 0000000..441a0b1 --- /dev/null +++ b/src/mflux/callbacks/instances/canny_saver.py @@ -0,0 +1,24 @@ +import os +from pathlib import Path + +import PIL.Image + +from mflux import ImageUtil +from mflux.callbacks.callback import BeforeLoopCallback + + +class CannyImageSaver(BeforeLoopCallback): + def __init__(self, path: str): + self.path = Path(path) + + def call_before_loop( + self, + seed: int, + prompt: str, + canny_image: PIL.Image.Image | None = None, + ) -> None: # fmt: off + base, ext = os.path.splitext(self.path) + ImageUtil.save_image( + image=canny_image, + path=f"{base}_controlnet_canny{ext}" + ) # fmt: off diff --git a/src/mflux/callbacks/instances/stepwise_handler.py b/src/mflux/callbacks/instances/stepwise_handler.py new file mode 100644 index 0000000..1f425e7 --- /dev/null +++ b/src/mflux/callbacks/instances/stepwise_handler.py @@ -0,0 +1,70 @@ +from pathlib import Path + +import mlx.core as mx +import tqdm + +from mflux import StopImageGenerationException +from mflux.callbacks.callback import InLoopCallback, InterruptCallback +from mflux.config.runtime_config import RuntimeConfig +from mflux.post_processing.array_util import ArrayUtil +from mflux.post_processing.image_util import ImageUtil + + +class StepwiseHandler(InLoopCallback, InterruptCallback): + def __init__( + self, + flux, + output_dir: str, + ): + self.flux = flux + self.output_dir = Path(output_dir) + self.step_wise_images = [] + + if self.output_dir: + self.output_dir.mkdir(parents=True, exist_ok=True) + + def call_in_loop( + self, + seed: int, + prompt: str, + step: int, + latents: mx.array, + config: RuntimeConfig, + time_steps: tqdm + ) -> None: # fmt: off + unpack_latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) + stepwise_decoded = self.flux.vae.decode(unpack_latents) + stepwise_img = ImageUtil.to_image( + decoded_latents=stepwise_decoded, + config=config, + seed=seed, + prompt=prompt, + quantization=self.flux.bits, + lora_paths=self.flux.lora_paths, + lora_scales=self.flux.lora_scales, + generation_time=time_steps.format_dict["elapsed"], + ) + self.step_wise_images.append(stepwise_img) + + stepwise_img.save( + path=self.output_dir / f"seed_{seed}_step{step}of{len(time_steps)}.png", + export_json_metadata=False, + ) + self._save_composite(seed=seed) + + def call_interrupt( + self, + seed: int, + prompt: str, + step: int, + latents: mx.array, + config: RuntimeConfig, + time_steps: tqdm + ) -> None: # fmt: off + self._save_composite(seed=seed) + raise StopImageGenerationException(f"Stopping image generation at step {step + 1}/{len(time_steps)}") + + def _save_composite(self, seed: int) -> None: + if self.step_wise_images: + composite_img = ImageUtil.to_composite_image(self.step_wise_images) + composite_img.save(self.output_dir / f"seed_{seed}_composite.png") diff --git a/src/mflux/controlnet/controlnet_util.py b/src/mflux/controlnet/controlnet_util.py index b852163..a66f26f 100644 --- a/src/mflux/controlnet/controlnet_util.py +++ b/src/mflux/controlnet/controlnet_util.py @@ -1,12 +1,10 @@ import logging -import os import cv2 import mlx.core as mx import numpy as np import PIL.Image -from mflux.config.runtime_config import RuntimeConfig from mflux.models.vae.vae import VAE from mflux.post_processing.array_util import ArrayUtil @@ -17,29 +15,20 @@ class ControlnetUtil: @staticmethod def encode_image( vae: VAE, - config: RuntimeConfig, + height: int, + width: int, controlnet_image_path: str, - controlnet_save_canny: bool, - output: str, - ) -> mx.array: + ) -> (mx.array, PIL.Image): from mflux import ImageUtil control_image = ImageUtil.load_image(controlnet_image_path) - control_image = ControlnetUtil._scale_image(config.height, config.width, control_image) + control_image = ControlnetUtil._scale_image(height=height, width=width, img=control_image) control_image = ControlnetUtil._preprocess_canny(control_image) - - if controlnet_save_canny: - base, ext = os.path.splitext(output) - ImageUtil.save_image( - image=control_image, - path=f"{base}_controlnet_canny{ext}" - ) # fmt: off - controlnet_cond = ImageUtil.to_array(control_image) controlnet_cond = vae.encode(controlnet_cond) controlnet_cond = (controlnet_cond / vae.scaling_factor) + vae.shift_factor - controlnet_cond = ArrayUtil.pack_latents(latents=controlnet_cond, height=config.height, width=config.width) - return controlnet_cond + controlnet_cond = ArrayUtil.pack_latents(latents=controlnet_cond, height=height, width=width) + return controlnet_cond, control_image @staticmethod def _preprocess_canny(img: PIL.Image) -> PIL.Image: diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index 41c9b03..5d38a11 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -1,15 +1,13 @@ -from pathlib import Path - import mlx.core as mx from mlx import nn from tqdm import tqdm +from mflux.callbacks.callbacks import Callbacks from mflux.config.config import Config from mflux.config.model_config import ModelConfig from mflux.config.runtime_config import RuntimeConfig from mflux.controlnet.controlnet_util import ControlnetUtil from mflux.controlnet.transformer_controlnet import TransformerControlnet -from mflux.error.exceptions import StopImageGenerationException from mflux.flux.flux_initializer import FluxInitializer from mflux.latent_creator.latent_creator import LatentCreator from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder @@ -19,7 +17,6 @@ from mflux.models.vae.vae import VAE from mflux.post_processing.array_util import ArrayUtil from mflux.post_processing.generated_image import GeneratedImage from mflux.post_processing.image_util import ImageUtil -from mflux.post_processing.stepwise_handler import StepwiseHandler from mflux.weights.model_saver import ModelSaver @@ -53,49 +50,44 @@ class Flux1Controlnet(nn.Module): self, seed: int, prompt: str, - output: str, controlnet_image_path: str, - controlnet_save_canny: bool = False, config: Config = Config(), - stepwise_output_dir: Path = None, ) -> GeneratedImage: - # Convert the user config to a runtime config with derived parameters. + # 0. Create a new runtime config based on the model type and input parameters config = RuntimeConfig(config, self.model_config) time_steps = tqdm(range(config.num_inference_steps)) - stepwise_handler = StepwiseHandler( - flux=self, - config=config, - seed=seed, - prompt=prompt, - time_steps=time_steps, - output_dir=stepwise_output_dir, - ) - # 0. Encode the controlnet reference image - controlnet_condition = ControlnetUtil.encode_image( + # 1. Encode the controlnet reference image + controlnet_condition, canny_image = ControlnetUtil.encode_image( vae=self.vae, - config=config, + height=config.height, + width=config.width, controlnet_image_path=controlnet_image_path, - controlnet_save_canny=controlnet_save_canny, - output=output, ) - # 1. Create the initial latents + # 2. Create the initial latents latents = LatentCreator.create( seed=seed, height=config.height, width=config.width ) # fmt: off - # 2. Embed the prompt + # 3. Encode the prompt t5_tokens = self.t5_tokenizer.tokenize(prompt) clip_tokens = self.clip_tokenizer.tokenize(prompt) prompt_embeds = self.t5_text_encoder(t5_tokens) pooled_prompt_embeds = self.clip_text_encoder(clip_tokens) + # (Optional) Call subscribers for beginning of loop + Callbacks.before_loop( + seed=seed, + prompt=prompt, + canny_image=canny_image + ) # fmt: off + for gen_step, t in enumerate(time_steps, 1): try: - # 3.t Compute controlnet samples + # 4.t Compute controlnet samples controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet( t=t, config=config, @@ -105,7 +97,7 @@ class Flux1Controlnet(nn.Module): controlnet_condition=controlnet_condition, ) - # 4.t Predict the noise + # 5.t Predict the noise noise = self.transformer( t=t, config=config, @@ -116,21 +108,34 @@ class Flux1Controlnet(nn.Module): controlnet_single_block_samples=controlnet_single_block_samples, ) - # 5.t Take one denoise step + # 6.t Take one denoise step dt = config.sigmas[t + 1] - config.sigmas[t] latents += noise * dt - # Handle stepwise output if enabled - stepwise_handler.process_step(gen_step, latents) + # (Optional) Call subscribes at end of loop + Callbacks.in_loop( + seed=seed, + prompt=prompt, + step=gen_step, + latents=latents, + config=config, + time_steps=time_steps, + ) # fmt: off - # Evaluate to enable progress tracking + # (Optional) Evaluate to enable progress tracking mx.eval(latents) except KeyboardInterrupt: # noqa: PERF203 - stepwise_handler.handle_interruption() - raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}") + Callbacks.interruption( + seed=seed, + prompt=prompt, + step=gen_step, + latents=latents, + config=config, + time_steps=time_steps, + ) - # 5. Decode the latent array and return the image + # 7. Decode the latent array and return the image latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index 15ed094..178b68e 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -1,16 +1,13 @@ -import warnings -from pathlib import Path - import mlx.core as mx from mlx import nn from tqdm import tqdm +from mflux.callbacks.callbacks import Callbacks from mflux.config.config import Config from mflux.config.model_config import ModelConfig, ModelLookup from mflux.config.runtime_config import RuntimeConfig -from mflux.error.exceptions import StopImageGenerationException from mflux.flux.flux_initializer import FluxInitializer -from mflux.latent_creator.latent_creator import LatentCreator +from mflux.latent_creator.latent_creator import Img2Img, LatentCreator from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder from mflux.models.transformer.transformer import Transformer @@ -18,7 +15,6 @@ from mflux.models.vae.vae import VAE from mflux.post_processing.array_util import ArrayUtil from mflux.post_processing.generated_image import GeneratedImage from mflux.post_processing.image_util import ImageUtil -from mflux.post_processing.stepwise_handler import StepwiseHandler from mflux.weights.model_saver import ModelSaver @@ -51,33 +47,36 @@ class Flux1(nn.Module): seed: int, prompt: str, config: Config = Config(), - stepwise_output_dir: Path = None, ) -> GeneratedImage: - # Convert the user config to a runtime config with derived parameters. + # 0. Create a new runtime config based on the model type and input parameters config = RuntimeConfig(config, self.model_config) time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) - stepwise_handler = StepwiseHandler( - flux=self, - config=config, - seed=seed, - prompt=prompt, - time_steps=time_steps, - output_dir=stepwise_output_dir, - ) # 1. Create the initial latents latents = LatentCreator.create_for_txt2img_or_img2img( seed=seed, - vae=self.vae, - runtime_conf=config, + height=config.height, + width=config.width, + img2img=Img2Img( + vae=self.vae, + sigmas=config.sigmas, + init_time_step=config.init_time_step, + init_image_path=config.init_image_path, + ), ) - # 2. Embed the prompt + # 2. Encode the prompt t5_tokens = self.t5_tokenizer.tokenize(prompt) clip_tokens = self.clip_tokenizer.tokenize(prompt) prompt_embeds = self.t5_text_encoder(t5_tokens) pooled_prompt_embeds = self.clip_text_encoder(clip_tokens) + # (Optional) Call subscribers for beginning of loop + Callbacks.before_loop( + seed=seed, + prompt=prompt + ) # fmt: off + for gen_step, t in enumerate(time_steps, 1): try: # 3.t Predict the noise @@ -93,17 +92,30 @@ class Flux1(nn.Module): dt = config.sigmas[t + 1] - config.sigmas[t] latents += noise * dt - # Handle stepwise output if enabled - stepwise_handler.process_step(gen_step, latents) + # (Optional) Call subscribes at end of loop + Callbacks.in_loop( + seed=seed, + prompt=prompt, + step=gen_step, + latents=latents, + config=config, + time_steps=time_steps, + ) # fmt: off - # Evaluate to enable progress tracking + # (Optional) Evaluate to enable progress tracking mx.eval(latents) except KeyboardInterrupt: # noqa: PERF203 - stepwise_handler.handle_interruption() - raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}") + Callbacks.interruption( + seed=seed, + prompt=prompt, + step=gen_step, + latents=latents, + config=config, + time_steps=time_steps, + ) - # 5. Decode the latent array and return the image + # 7. Decode the latent array and return the image latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( @@ -119,15 +131,6 @@ class Flux1(nn.Module): generation_time=time_steps.format_dict["elapsed"], ) - @staticmethod - def from_alias(alias: str, quantize: int | None = None) -> "Flux1": - warnings.warn( - "from_alias is deprecated and will be removed in a future release. Please use from_name instead.", - DeprecationWarning, - stacklevel=2, - ) - return Flux1.from_name(model_name=alias, quantize=quantize) - @staticmethod def from_name(model_name: str, quantize: int | None = None) -> "Flux1": return Flux1( diff --git a/src/mflux/generate.py b/src/mflux/generate.py index 72f5a49..ea95a02 100644 --- a/src/mflux/generate.py +++ b/src/mflux/generate.py @@ -1,11 +1,11 @@ -from pathlib import Path - from mflux import Config, Flux1, ModelLookup, StopImageGenerationException +from mflux.callbacks.callback_registry import CallbackRegistry +from mflux.callbacks.instances.stepwise_handler import StepwiseHandler from mflux.ui.cli.parsers import CommandLineParser def main(): - # fmt: off + # 0. Parse command line arguments parser = CommandLineParser(description="Generate an image based on a prompt.") parser.add_model_arguments(require_model_arg=False) parser.add_lora_arguments() @@ -23,13 +23,18 @@ def main(): lora_scales=args.lora_scales, ) + # 2. Register the optional callbacks + if args.stepwise_image_output_dir: + handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir) + CallbackRegistry.register_in_loop(handler) + CallbackRegistry.register_interrupt(handler) + try: - for seed_value in args.seed: - # 2. Generate an image for each seed value + for seed in args.seed: + # 3. Generate an image for each seed value image = flux.generate_image( - seed=seed_value, + seed=seed, prompt=args.prompt, - stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None, config=Config( num_inference_steps=args.steps, height=args.height, @@ -39,8 +44,8 @@ def main(): init_image_strength=args.init_image_strength, ), ) - # 3. Save the image - image.save(path=args.output.format(seed=seed_value), export_json_metadata=args.metadata) + # 4. Save the image + image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) except StopImageGenerationException as stop_exc: print(stop_exc) diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index ccbf394..d467b4a 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/generate_controlnet.py @@ -1,10 +1,12 @@ -from pathlib import Path - from mflux import Config, Flux1Controlnet, ModelLookup, StopImageGenerationException +from mflux.callbacks.callback_registry import CallbackRegistry +from mflux.callbacks.instances.canny_saver import CannyImageSaver +from mflux.callbacks.instances.stepwise_handler import StepwiseHandler from mflux.ui.cli.parsers import CommandLineParser def main(): + # 0. Parse command line arguments parser = CommandLineParser(description="Generate an image based on a prompt and a controlnet reference image.") # fmt: off parser.add_model_arguments(require_model_arg=True) parser.add_lora_arguments() @@ -22,16 +24,21 @@ def main(): lora_scales=args.lora_scales, ) + # 2. Register the optional callbacks + if args.controlnet_save_canny: + CallbackRegistry.register_before_loop(CannyImageSaver(path=args.output)) + if args.stepwise_image_output_dir: + handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir) + CallbackRegistry.register_in_loop(handler) + CallbackRegistry.register_interrupt(handler) + try: - for seed_value in args.seed: - # 2. Generate an image for each seed value + for seed in args.seed: + # 3. Generate an image for each seed value image = flux.generate_image( - seed=seed_value, + seed=seed, prompt=args.prompt, - output=args.output, controlnet_image_path=args.controlnet_image_path, - controlnet_save_canny=args.controlnet_save_canny, - stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None, config=Config( num_inference_steps=args.steps, height=args.height, @@ -41,8 +48,8 @@ def main(): ), ) - # 3. Save the image - image.save(path=args.output.format(seed=seed_value), export_json_metadata=args.metadata) + # 4. Save the image + image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) except StopImageGenerationException as stop_exc: print(stop_exc) diff --git a/src/mflux/latent_creator/latent_creator.py b/src/mflux/latent_creator/latent_creator.py index 84e816f..addb970 100644 --- a/src/mflux/latent_creator/latent_creator.py +++ b/src/mflux/latent_creator/latent_creator.py @@ -1,11 +1,24 @@ import mlx.core as mx -from mlx import nn -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.vae.vae import VAE from mflux.post_processing.array_util import ArrayUtil from mflux.post_processing.image_util import ImageUtil +class Img2Img: + def __init__( + self, + vae: VAE, + sigmas: mx.array, + init_time_step: int, + init_image_path: int, + ): + self.vae = vae + self.sigmas = sigmas + self.init_time_step = init_time_step + self.init_image_path = init_image_path + + class LatentCreator: @staticmethod def create( @@ -21,29 +34,41 @@ class LatentCreator: @staticmethod def create_for_txt2img_or_img2img( seed: int, - runtime_conf: RuntimeConfig, - vae: nn.Module, + height: int, + width: int, + img2img: Img2Img, ) -> mx.array: - pure_noise = LatentCreator.create( - seed=seed, - height=runtime_conf.height, - width=runtime_conf.width, - ) + # 0. Determine type of image generation + is_text2img = img2img.init_image_path is None - if runtime_conf.config.init_image_path is None: - # Text2Image - 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( - image=user_image, - target_width=runtime_conf.width, - target_height=runtime_conf.height, + if is_text2img: + # 1. Create the pure noise + return LatentCreator.create( + seed=seed, + height=height, + width=width, ) - encoded = vae.encode(ImageUtil.to_array(scaled_user_image)) - latents = ArrayUtil.pack_latents(latents=encoded, height=runtime_conf.height, width=runtime_conf.width) - sigma = runtime_conf.sigmas[runtime_conf.init_time_step] + else: + # 1. Create the pure noise + pure_noise = LatentCreator.create( + seed=seed, + height=height, + width=width, + ) + + # 2. Encode the image + scaled_user_image = ImageUtil.scale_to_dimensions( + image=ImageUtil.load_image(img2img.init_image_path).convert("RGB"), + target_width=width, + target_height=height, + ) + encoded = img2img.vae.encode(ImageUtil.to_array(scaled_user_image)) + latents = ArrayUtil.pack_latents(latents=encoded, height=height, width=width) + + # 3. Find the appropriate sigma value + sigma = img2img.sigmas[img2img.init_time_step] + + # 4. Blend the appropriate amount of noise based on linear interpolation return LatentCreator.add_noise_by_interpolation( clean=latents, noise=pure_noise, diff --git a/src/mflux/post_processing/stepwise_handler.py b/src/mflux/post_processing/stepwise_handler.py deleted file mode 100644 index 59dcc25..0000000 --- a/src/mflux/post_processing/stepwise_handler.py +++ /dev/null @@ -1,60 +0,0 @@ -from pathlib import Path - -import mlx.core as mx -import tqdm - -from mflux.config.runtime_config import RuntimeConfig -from mflux.post_processing.array_util import ArrayUtil -from mflux.post_processing.image_util import ImageUtil - - -class StepwiseHandler: - def __init__( - self, - flux, - config: RuntimeConfig, - seed: int, - prompt: str, - time_steps: tqdm.std.tqdm, - output_dir: Path | None = None, - ): - self.flux = flux - self.config = config - self.seed = seed - self.prompt = prompt - self.output_dir = output_dir - self.time_steps = time_steps - self.step_wise_images = [] - - if self.output_dir: - self.output_dir.mkdir(parents=True, exist_ok=True) - - def save_composite(self): - if self.step_wise_images: - composite_img = ImageUtil.to_composite_image(self.step_wise_images) - composite_img.save(self.output_dir / f"seed_{self.seed}_composite.png") - - def process_step(self, gen_step: int, latents: mx.array): - if self.output_dir: - 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, - config=self.config, - seed=self.seed, - prompt=self.prompt, - quantization=self.flux.bits, - lora_paths=self.flux.lora_paths, - lora_scales=self.flux.lora_scales, - generation_time=self.time_steps.format_dict["elapsed"], - ) - self.step_wise_images.append(stepwise_img) - - stepwise_img.save( - path=self.output_dir / f"seed_{self.seed}_step{gen_step}of{len(self.time_steps)}.png", - export_json_metadata=False, - ) - self.save_composite() - - def handle_interruption(self): - self.save_composite() diff --git a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py index e2794a4..f89eb1f 100644 --- a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py +++ b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py @@ -40,9 +40,7 @@ class ImageGeneratorControlnetTestHelper: image = flux.generate_image( seed=seed, prompt=prompt, - output=str(output_image_path), controlnet_image_path=controlnet_image_path, - controlnet_save_canny=False, config=Config( num_inference_steps=steps, height=768,