From e908c1cede4fda91e18ad0dfbc4664d90aa5a3a1 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Wed, 9 Oct 2024 16:48:59 +0200 Subject: [PATCH] Add stepwise handler --- src/mflux/__init__.py | 2 +- src/mflux/controlnet/flux_controlnet.py | 71 ++++++------------- src/mflux/error/__init__.py | 0 src/mflux/{ => error}/exceptions.py | 0 src/mflux/flux/flux.py | 70 ++++++++---------- src/mflux/generate.py | 2 +- src/mflux/generate_controlnet.py | 2 +- src/mflux/post_processing/array_util.py | 17 +++++ src/mflux/post_processing/stepwise_handler.py | 55 ++++++++++++++ 9 files changed, 123 insertions(+), 96 deletions(-) create mode 100644 src/mflux/error/__init__.py rename src/mflux/{ => error}/exceptions.py (100%) create mode 100644 src/mflux/post_processing/array_util.py create mode 100644 src/mflux/post_processing/stepwise_handler.py diff --git a/src/mflux/__init__.py b/src/mflux/__init__.py index e7c48dc..e53c33f 100644 --- a/src/mflux/__init__.py +++ b/src/mflux/__init__.py @@ -3,7 +3,7 @@ from mflux.config.config import ConfigControlnet from mflux.config.model_config import ModelConfig from mflux.controlnet.flux_controlnet import Flux1Controlnet from mflux.flux.flux import Flux1 -from mflux.exceptions import StopImageGenerationException +from mflux.error.exceptions import StopImageGenerationException from mflux.post_processing.image_util import ImageUtil __all__ = [ diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index 3241c30..3cca900 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -10,12 +10,14 @@ from mflux.config.runtime_config import RuntimeConfig from mflux.controlnet.controlnet_util import ControlnetUtil from mflux.controlnet.transformer_controlnet import TransformerControlnet from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet -from mflux.exceptions import StopImageGenerationException +from mflux.error.exceptions import StopImageGenerationException 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 from mflux.models.vae.vae import VAE +from mflux.post_processing.array_util import ArrayUtil from mflux.post_processing.image_util import ImageUtil +from mflux.post_processing.stepwise_handler import StepwiseHandler from mflux.tokenizer.clip_tokenizer import TokenizerCLIP from mflux.tokenizer.t5_tokenizer import TokenizerT5 from mflux.tokenizer.tokenizer_handler import TokenizerHandler @@ -112,11 +114,18 @@ class Flux1Controlnet: controlnet_save_canny: bool = False, config: ConfigControlnet = ConfigControlnet(), stepwise_output_dir: Path = None, - stepwise_composite_only=False ) -> "GeneratedImage": # fmt: off # 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, + ) # Embedd the controlnet reference image control_image = ImageUtil.load_image(controlnet_image_path) @@ -127,7 +136,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 = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width) + controlnet_cond = ArrayUtil.pack_latents(controlnet_cond, config.height, config.width) # 1. Create the initial latents latents = mx.random.normal( @@ -135,15 +144,12 @@ class Flux1Controlnet: key=mx.random.key(seed) ) # fmt: off - # 2. Embedd the prompt + # 2. Embed the prompt t5_tokens = self.t5_tokenizer.tokenize(prompt) clip_tokens = self.clip_tokenizer.tokenize(prompt) prompt_embeds = self.t5_text_encoder.forward(t5_tokens) pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens) - step_wise_images = [] - if stepwise_output_dir: - stepwise_output_dir.mkdir(parents=True, exist_ok=True) for t in time_steps: try: # Compute controlnet samples @@ -171,41 +177,18 @@ class Flux1Controlnet: dt = config.sigmas[t + 1] - config.sigmas[t] latents += noise * dt + # Handle stepwise output if enabled + stepwise_handler.process_step(t, latents) + # Evaluate to enable progress tracking mx.eval(latents) - if stepwise_output_dir: - stepwise_decoded = self.vae.decode( - Flux1Controlnet._unpack_latents(latents, config.height, config.width) - ) - # performance todo: Pillow mostly uses CPU, - # can try to improve image generation performance by offloading - # stepwise image processing to the CPU via threading - stepwise_img = ImageUtil.to_image( - decoded_latents=stepwise_decoded, - seed=seed, - prompt=prompt, - quantization=self.bits, - generation_time=time_steps.format_dict["elapsed"], - lora_paths=self.lora_paths, - lora_scales=self.lora_scales, - config=config, - ) - step_wise_images.append(stepwise_img) - if not stepwise_composite_only: - stepwise_img.save( - path=stepwise_output_dir / f"seed_{seed}_step{t+1}of{len(time_steps)}.png", - export_json_metadata=False, - ) - except KeyboardInterrupt: # noqa: PERF203 - raise StopImageGenerationException(f"Stopping image generation at step {t+1}/{len(time_steps)}") - finally: - if step_wise_images and stepwise_output_dir: - composite_img = ImageUtil.to_composite_image(step_wise_images) - composite_img.save(stepwise_output_dir / f"seed_{seed}_composite.png") + except KeyboardInterrupt: + stepwise_handler.handle_interruption() + raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}") # 5. Decode the latent array and return the image - latents = Flux1Controlnet._unpack_latents(latents, config.height, config.width) + latents = ArrayUtil.unpack_latents(latents, config.height, config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, @@ -219,20 +202,6 @@ class Flux1Controlnet: controlnet_image_path=controlnet_image_path, ) - @staticmethod - def _unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: - latents = mx.reshape(latents, (1, height // 16, width // 16, 16, 2, 2)) - latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5)) - latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2)) - return latents - - @staticmethod - def _pack_latents(latents: mx.array, height: int, width: int) -> mx.array: - latents = mx.reshape(latents, (1, 16, height // 16, 2, width // 16, 2)) - latents = mx.transpose(latents, (0, 2, 4, 1, 3, 5)) - latents = mx.reshape(latents, (1, (width // 16) * (height // 16), 64)) - return latents - def _set_model_weights(self, weights): self.vae.update(weights.vae) self.transformer.update(weights.transformer) diff --git a/src/mflux/error/__init__.py b/src/mflux/error/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/exceptions.py b/src/mflux/error/exceptions.py similarity index 100% rename from src/mflux/exceptions.py rename to src/mflux/error/exceptions.py diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index 41de838..634bf32 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -1,16 +1,19 @@ -import mlx.core as mx from pathlib import Path + +import mlx.core as mx from mlx import nn from tqdm import tqdm from mflux.config.config import Config from mflux.config.model_config import ModelConfig from mflux.config.runtime_config import RuntimeConfig -from mflux.exceptions import StopImageGenerationException +from mflux.error.exceptions import StopImageGenerationException +from mflux.post_processing.stepwise_handler import StepwiseHandler 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 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.tokenizer.clip_tokenizer import TokenizerCLIP @@ -71,10 +74,24 @@ class Flux1: if weights.quantization_level is not None: self._set_model_weights(weights) - def generate_image(self, seed: int, prompt: str, config: Config = Config(), stepwise_output_dir: Path = None, stepwise_composite_only=False) -> GeneratedImage: # fmt: off + def generate_image( + self, + seed: int, + prompt: str, + config: Config = Config(), + stepwise_output_dir: Path = None, + ) -> GeneratedImage: # 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, + ) # 1. Create the initial latents latents = mx.random.normal( @@ -82,15 +99,12 @@ class Flux1: key=mx.random.key(seed) ) # fmt: off - # 2. Embedd the prompt + # 2. Embed the prompt t5_tokens = self.t5_tokenizer.tokenize(prompt) clip_tokens = self.clip_tokenizer.tokenize(prompt) prompt_embeds = self.t5_text_encoder.forward(t5_tokens) pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens) - step_wise_images = [] - if stepwise_output_dir: - stepwise_output_dir.mkdir(parents=True, exist_ok=True) for t in time_steps: try: # 3.t Predict the noise @@ -106,39 +120,18 @@ class Flux1: dt = config.sigmas[t + 1] - config.sigmas[t] latents += noise * dt + # Handle stepwise output if enabled + stepwise_handler.process_step(t, latents) + # Evaluate to enable progress tracking mx.eval(latents) - if stepwise_output_dir: - stepwise_decoded = self.vae.decode(Flux1._unpack_latents(latents, config.height, config.width)) - # performance todo: Pillow mostly uses CPU, - # can try to improve image generation performance by offloading - # stepwise image processing to the CPU via threading - stepwise_img = ImageUtil.to_image( - decoded_latents=stepwise_decoded, - seed=seed, - prompt=prompt, - quantization=self.bits, - generation_time=time_steps.format_dict["elapsed"], - lora_paths=self.lora_paths, - lora_scales=self.lora_scales, - config=config, - ) - step_wise_images.append(stepwise_img) - if not stepwise_composite_only: - stepwise_img.save( - path=stepwise_output_dir / f"seed_{seed}_step{t+1}of{len(time_steps)}.png", - export_json_metadata=False, - ) - except KeyboardInterrupt: # noqa: PERF203 - raise StopImageGenerationException(f"Stopping image generation at step {t+1}/{len(time_steps)}") - finally: - if step_wise_images and stepwise_output_dir: - composite_img = ImageUtil.to_composite_image(step_wise_images) - composite_img.save(stepwise_output_dir / f"seed_{seed}_composite.png") + except KeyboardInterrupt: + stepwise_handler.handle_interruption() + raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}") # 5. Decode the latent array and return the image - latents = Flux1._unpack_latents(latents, config.height, config.width) + latents = ArrayUtil.unpack_latents(latents, config.height, config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, @@ -151,13 +144,6 @@ class Flux1: config=config, ) - @staticmethod - def _unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: - latents = mx.reshape(latents, (1, height // 16, width // 16, 16, 2, 2)) - latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5)) - latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2)) - return latents - @staticmethod def from_alias(alias: str, quantize: int | None = None) -> "Flux1": return Flux1( diff --git a/src/mflux/generate.py b/src/mflux/generate.py index 36d1cf5..16dbf28 100644 --- a/src/mflux/generate.py +++ b/src/mflux/generate.py @@ -46,13 +46,13 @@ def main(): image = flux.generate_image( seed=int(time.time()) if args.seed is None else args.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, width=args.width, guidance=args.guidance, ), - stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None, ) # Save the image diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index c8813d8..470236d 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/generate_controlnet.py @@ -52,6 +52,7 @@ def main(): 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=ConfigControlnet( num_inference_steps=args.steps, height=args.height, @@ -59,7 +60,6 @@ def main(): guidance=args.guidance, controlnet_strength=args.controlnet_strength, ), - stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None, ) # Save the image diff --git a/src/mflux/post_processing/array_util.py b/src/mflux/post_processing/array_util.py new file mode 100644 index 0000000..9fe82b8 --- /dev/null +++ b/src/mflux/post_processing/array_util.py @@ -0,0 +1,17 @@ +import mlx.core as mx + + +class ArrayUtil: + @staticmethod + def unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: + latents = mx.reshape(latents, (1, height // 16, width // 16, 16, 2, 2)) + latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5)) + latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2)) + return latents + + @staticmethod + def pack_latents(latents: mx.array, height: int, width: int) -> mx.array: + latents = mx.reshape(latents, (1, 16, height // 16, 2, width // 16, 2)) + latents = mx.transpose(latents, (0, 2, 4, 1, 3, 5)) + latents = mx.reshape(latents, (1, (width // 16) * (height // 16), 64)) + return latents diff --git a/src/mflux/post_processing/stepwise_handler.py b/src/mflux/post_processing/stepwise_handler.py new file mode 100644 index 0000000..d3e9cd6 --- /dev/null +++ b/src/mflux/post_processing/stepwise_handler.py @@ -0,0 +1,55 @@ +from pathlib import Path + +import mlx.core as mx + +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, + 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 process_step(self, step: int, latents: mx.array): + if self.output_dir: + unpack_latents = ArrayUtil.unpack_latents(latents, self.config.height, self.config.width) + stepwise_decoded = self.flux.vae.decode(unpack_latents) + stepwise_img = ImageUtil.to_image( + decoded_latents=stepwise_decoded, + seed=self.seed, + prompt=self.prompt, + quantization=self.flux.bits, + generation_time=self.time_steps.format_dict["elapsed"], + lora_paths=self.flux.lora_paths, + lora_scales=self.flux.lora_scales, + config=self.config, + ) + self.step_wise_images.append(stepwise_img) + + stepwise_img.save( + path=self.output_dir / f"seed_{self.seed}_step{step + 1}of{len(self.time_steps)}.png", + export_json_metadata=False, + ) + + def handle_interruption(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")