diff --git a/src/mflux/callbacks/callback.py b/src/mflux/callbacks/callback.py index 48281dd..7263b77 100644 --- a/src/mflux/callbacks/callback.py +++ b/src/mflux/callbacks/callback.py @@ -12,6 +12,8 @@ class BeforeLoopCallback(Protocol): self, seed: int, prompt: str, + latents: mx.array, + config: RuntimeConfig, canny_image: PIL.Image.Image | None = None, ) -> None: # fmt: off ... @@ -20,9 +22,9 @@ class BeforeLoopCallback(Protocol): class InLoopCallback(Protocol): def call_in_loop( self, + t: int, seed: int, prompt: str, - step: int, latents: mx.array, config: RuntimeConfig, time_steps: tqdm @@ -33,9 +35,9 @@ class InLoopCallback(Protocol): class InterruptCallback(Protocol): def call_interrupt( self, + t: int, seed: int, prompt: str, - step: int, latents: mx.array, config: RuntimeConfig, time_steps: tqdm diff --git a/src/mflux/callbacks/callbacks.py b/src/mflux/callbacks/callbacks.py index 220b7b9..10ad3b8 100644 --- a/src/mflux/callbacks/callbacks.py +++ b/src/mflux/callbacks/callbacks.py @@ -11,29 +11,33 @@ class Callbacks: def before_loop( seed: int, prompt: str, + latents: mx.array, + config: RuntimeConfig, canny_image: PIL.Image.Image | None = None, ): # fmt: off for subscriber in CallbackRegistry.before_loop_callbacks(): subscriber.call_before_loop( seed=seed, prompt=prompt, + latents=latents, + config=config, canny_image=canny_image, ) @staticmethod def in_loop( + t: int, 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( + t=t, seed=seed, prompt=prompt, - step=step, latents=latents, config=config, time_steps=time_steps, @@ -41,18 +45,18 @@ class Callbacks: @staticmethod def interruption( + t: int, 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( + t=t, seed=seed, prompt=prompt, - step=step, latents=latents, config=config, time_steps=time_steps, diff --git a/src/mflux/callbacks/instances/canny_saver.py b/src/mflux/callbacks/instances/canny_saver.py index 441a0b1..548b291 100644 --- a/src/mflux/callbacks/instances/canny_saver.py +++ b/src/mflux/callbacks/instances/canny_saver.py @@ -1,10 +1,12 @@ import os from pathlib import Path +import mlx.core as mx import PIL.Image from mflux import ImageUtil from mflux.callbacks.callback import BeforeLoopCallback +from mflux.config.runtime_config import RuntimeConfig class CannyImageSaver(BeforeLoopCallback): @@ -15,8 +17,10 @@ class CannyImageSaver(BeforeLoopCallback): self, seed: int, prompt: str, + latents: mx.array, + config: RuntimeConfig, canny_image: PIL.Image.Image | None = None, - ) -> None: # fmt: off + ) -> None: base, ext = os.path.splitext(self.path) ImageUtil.save_image( image=canny_image, diff --git a/src/mflux/callbacks/instances/stepwise_handler.py b/src/mflux/callbacks/instances/stepwise_handler.py index 1f425e7..ea74c60 100644 --- a/src/mflux/callbacks/instances/stepwise_handler.py +++ b/src/mflux/callbacks/instances/stepwise_handler.py @@ -1,16 +1,17 @@ from pathlib import Path import mlx.core as mx +import PIL.Image import tqdm from mflux import StopImageGenerationException -from mflux.callbacks.callback import InLoopCallback, InterruptCallback +from mflux.callbacks.callback import BeforeLoopCallback, 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): +class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback): def __init__( self, flux, @@ -23,17 +24,65 @@ class StepwiseHandler(InLoopCallback, InterruptCallback): if self.output_dir: self.output_dir.mkdir(parents=True, exist_ok=True) - def call_in_loop( + def call_before_loop( self, seed: int, prompt: str, + latents: mx.array, + config: RuntimeConfig, + canny_image: PIL.Image.Image | None = None, + ) -> None: # fmt: off + self._save_image( + step=config.init_time_step, + seed=seed, + prompt=prompt, + latents=latents, + config=config, + time_steps=None, + ) + + def call_in_loop( + self, + t: int, + seed: int, + prompt: str, + latents: mx.array, + config: RuntimeConfig, + time_steps: tqdm + ) -> None: # fmt: off + self._save_image( + step=t + 1, + seed=seed, + prompt=prompt, + latents=latents, + config=config, + time_steps=time_steps, + ) + + def call_interrupt( + self, + t: int, + seed: int, + prompt: str, + latents: mx.array, + config: RuntimeConfig, + time_steps: tqdm + ) -> None: # fmt: off + self._save_composite(seed=seed) + raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}") + + def _save_image( + self, step: int, + seed: int, + prompt: str, 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) + generation_time = time_steps.format_dict["elapsed"] if time_steps is not None else 0 stepwise_img = ImageUtil.to_image( decoded_latents=stepwise_decoded, config=config, @@ -42,28 +91,15 @@ class StepwiseHandler(InLoopCallback, InterruptCallback): quantization=self.flux.bits, lora_paths=self.flux.lora_paths, lora_scales=self.flux.lora_scales, - generation_time=time_steps.format_dict["elapsed"], + generation_time=generation_time, ) - self.step_wise_images.append(stepwise_img) - stepwise_img.save( - path=self.output_dir / f"seed_{seed}_step{step}of{len(time_steps)}.png", + path=self.output_dir / f"seed_{seed}_step{step}of{config.num_inference_steps}.png", export_json_metadata=False, ) + self.step_wise_images.append(stepwise_img) 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) diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index 78a71e9..73479e3 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -87,10 +87,12 @@ class Flux1Controlnet(nn.Module): Callbacks.before_loop( seed=seed, prompt=prompt, + latents=latents, + config=config, canny_image=canny_image ) # fmt: off - for gen_step, t in enumerate(time_steps, 1): + for t in time_steps: try: # 4.t Compute controlnet samples controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet( @@ -119,9 +121,9 @@ class Flux1Controlnet(nn.Module): # (Optional) Call subscribes at end of loop Callbacks.in_loop( + t=t, seed=seed, prompt=prompt, - step=gen_step, latents=latents, config=config, time_steps=time_steps, @@ -132,9 +134,9 @@ class Flux1Controlnet(nn.Module): except KeyboardInterrupt: # noqa: PERF203 Callbacks.interruption( + t=t, seed=seed, prompt=prompt, - step=gen_step, latents=latents, config=config, time_steps=time_steps, diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index 31c4f6b..ad996b2 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -79,10 +79,12 @@ class Flux1(nn.Module): # (Optional) Call subscribers for beginning of loop Callbacks.before_loop( seed=seed, - prompt=prompt + prompt=prompt, + latents=latents, + config=config, ) # fmt: off - for gen_step, t in enumerate(time_steps, 1): + for t in time_steps: try: # 3.t Predict the noise noise = self.transformer( @@ -99,9 +101,9 @@ class Flux1(nn.Module): # (Optional) Call subscribes at end of loop Callbacks.in_loop( + t=t, seed=seed, prompt=prompt, - step=gen_step, latents=latents, config=config, time_steps=time_steps, @@ -112,9 +114,9 @@ class Flux1(nn.Module): except KeyboardInterrupt: # noqa: PERF203 Callbacks.interruption( + t=t, seed=seed, prompt=prompt, - step=gen_step, latents=latents, config=config, time_steps=time_steps, diff --git a/src/mflux/generate.py b/src/mflux/generate.py index 5666c80..d3a1cfe 100644 --- a/src/mflux/generate.py +++ b/src/mflux/generate.py @@ -26,6 +26,7 @@ def main(): # 2. Register the optional callbacks if args.stepwise_image_output_dir: handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir) + CallbackRegistry.register_before_loop(handler) CallbackRegistry.register_in_loop(handler) CallbackRegistry.register_interrupt(handler) diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index bd5deab..bf92231 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/generate_controlnet.py @@ -29,6 +29,7 @@ def main(): 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_before_loop(handler) CallbackRegistry.register_in_loop(handler) CallbackRegistry.register_interrupt(handler)