From 74173f543e8e4b57293b7468b413725bf2afb926 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Fri, 14 Feb 2025 05:21:18 +0100 Subject: [PATCH 1/2] Let before-loop callback accept latents and let StepwiseHandler include initial latent --- src/mflux/callbacks/callback.py | 6 +- src/mflux/callbacks/callbacks.py | 12 ++- src/mflux/callbacks/instances/canny_saver.py | 6 +- .../callbacks/instances/stepwise_handler.py | 74 ++++++++++++++----- src/mflux/controlnet/flux_controlnet.py | 8 +- src/mflux/flux/flux.py | 10 ++- src/mflux/generate.py | 1 + src/mflux/generate_controlnet.py | 1 + 8 files changed, 85 insertions(+), 33 deletions(-) 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) From 821238284ac93136ac6d091b71f0873ad7d21d0e Mon Sep 17 00:00:00 2001 From: filipstrand Date: Fri, 14 Feb 2025 05:21:35 +0100 Subject: [PATCH 2/2] Fix: Handle edge case for img2img --- src/mflux/config/runtime_config.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py index 2ae1a8c..5c7d9d2 100644 --- a/src/mflux/config/runtime_config.py +++ b/src/mflux/config/runtime_config.py @@ -53,14 +53,16 @@ class RuntimeConfig: @property def init_time_step(self) -> int: - if self.config.init_image_path is None: - # For text-to-image, always begin at time step 0. + is_txt2img = self.config.init_image_path is None or self.config.init_image_strength == 0.0 + + if is_txt2img: return 0 else: - # For image-to-image: higher strength means we skip more steps. + # 1. Clamp strength to [0, 1] strength = max(0.0, min(1.0, self.config.init_image_strength)) - t = max(1, int(self.num_inference_steps * strength)) - return t + + # 2. Return start time in [1, floor(num_steps * strength)] + return max(1, int(self.num_inference_steps * strength)) @property def controlnet_strength(self) -> float | None: