From fcc9b21d19ecc240ff24e3fd68c63916671728f6 Mon Sep 17 00:00:00 2001 From: Anthony Wu <462072+anthonywu@users.noreply.github.com> Date: Mon, 23 Sep 2024 15:20:25 -0700 Subject: [PATCH] option to save stepwise and composite images, add interruptibility --- .gitignore | 5 +- pyproject.toml | 2 +- src/mflux/__init__.py | 2 + src/mflux/controlnet/flux_controlnet.py | 92 +++++++++++++++++-------- src/mflux/exceptions.py | 18 +++++ src/mflux/flux/flux.py | 62 +++++++++++++---- src/mflux/generate.py | 34 +++++---- src/mflux/generate_controlnet.py | 42 ++++++----- src/mflux/post_processing/image_util.py | 17 ++++- tests/test_readme_example.sh | 36 +++++++++- 10 files changed, 229 insertions(+), 81 deletions(-) create mode 100644 src/mflux/exceptions.py diff --git a/.gitignore b/.gitignore index db83963..c518e73 100644 --- a/.gitignore +++ b/.gitignore @@ -14,4 +14,7 @@ *.safetensors *.json -*.egg-info \ No newline at end of file +*.egg-info + +# build/ generated from 'pip install -e .' in developer mode +build/ diff --git a/pyproject.toml b/pyproject.toml index 7e4fd83..f441bba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,7 +57,7 @@ respect-gitignore = true # Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`) codes by default. # Unlike Flake8, Ruff doesn't enable pycodestyle warnings (`W`) or # McCabe complexity (`C901`) by default. -select = ["E4", "E7", "E9", "F"] +select = ["BLE", "E4", "E7", "E9", "F", "ICN", "LOG", "PERF", "W"] ignore = [] # Allow fix for all enabled rules (when `--fix`) is provided. diff --git a/src/mflux/__init__.py b/src/mflux/__init__.py index 3893264..e7c48dc 100644 --- a/src/mflux/__init__.py +++ b/src/mflux/__init__.py @@ -3,6 +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.post_processing.image_util import ImageUtil __all__ = [ @@ -12,4 +13,5 @@ __all__ = [ "ConfigControlnet", "ModelConfig", "ImageUtil", + "StopImageGenerationException", ] diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index fb68368..3241c30 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -1,16 +1,16 @@ import logging -from typing import TYPE_CHECKING - import mlx.core as mx from mlx import nn +from pathlib import Path from tqdm import tqdm - +from typing import TYPE_CHECKING from mflux.config.config import ConfigControlnet 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.controlnet.weight_handler_controlnet import WeightHandlerControlnet +from mflux.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 @@ -110,7 +110,9 @@ class Flux1Controlnet: output: str, controlnet_image_path: str, controlnet_save_canny: bool = False, - config: ConfigControlnet = ConfigControlnet() + 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) @@ -139,34 +141,68 @@ class Flux1Controlnet: 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: - # Compute controlnet samples - controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet.forward( - t=t, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - hidden_states=latents, - controlnet_cond=controlnet_cond, - config=config, - ) + try: + # Compute controlnet samples + controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet.forward( + t=t, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + hidden_states=latents, + controlnet_cond=controlnet_cond, + config=config, + ) - # 3.t Predict the noise - noise = self.transformer.predict( - t=t, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - hidden_states=latents, - config=config, - controlnet_block_samples=controlnet_block_samples, - controlnet_single_block_samples=controlnet_single_block_samples, - ) + # 3.t Predict the noise + noise = self.transformer.predict( + t=t, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + hidden_states=latents, + config=config, + controlnet_block_samples=controlnet_block_samples, + controlnet_single_block_samples=controlnet_single_block_samples, + ) - # 4.t Take one denoise step - dt = config.sigmas[t + 1] - config.sigmas[t] - latents += noise * dt + # 4.t Take one denoise step + dt = config.sigmas[t + 1] - config.sigmas[t] + latents += noise * dt - # Evaluate to enable progress tracking - mx.eval(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") # 5. Decode the latent array and return the image latents = Flux1Controlnet._unpack_latents(latents, config.height, config.width) diff --git a/src/mflux/exceptions.py b/src/mflux/exceptions.py new file mode 100644 index 0000000..a92b499 --- /dev/null +++ b/src/mflux/exceptions.py @@ -0,0 +1,18 @@ +class MFluxException(Exception): + """base class for all custom exceptions in mflux package.""" + + +class ImageSavingException(MFluxException): + """error ocurred while attempting to save image to storage.""" + + +class MetadataEmbedException(MFluxException): + """error ocurred while attempting to embed metadata in image""" + + +class MFluxUserException(MFluxException): + """an exception raised by user behavior or intention.""" + + +class StopImageGenerationException(MFluxUserException): + """user has requested to stop a image generation in progress.""" diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index 2b05e26..41de838 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -1,10 +1,12 @@ import mlx.core as mx +from pathlib import Path 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.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 @@ -69,7 +71,7 @@ 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()) -> GeneratedImage: + def generate_image(self, seed: int, prompt: str, config: Config = Config(), 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)) @@ -86,22 +88,54 @@ class Flux1: 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: - # 3.t Predict the noise - noise = self.transformer.predict( - t=t, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - hidden_states=latents, - config=config, - ) + try: + # 3.t Predict the noise + noise = self.transformer.predict( + t=t, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + hidden_states=latents, + config=config, + ) - # 4.t Take one denoise step - dt = config.sigmas[t + 1] - config.sigmas[t] - latents += noise * dt + # 4.t Take one denoise step + dt = config.sigmas[t + 1] - config.sigmas[t] + latents += noise * dt - # Evaluate to enable progress tracking - mx.eval(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") # 5. Decode the latent array and return the image latents = Flux1._unpack_latents(latents, config.height, config.width) diff --git a/src/mflux/generate.py b/src/mflux/generate.py index ded7371..36d1cf5 100644 --- a/src/mflux/generate.py +++ b/src/mflux/generate.py @@ -1,7 +1,8 @@ import argparse import time +from pathlib import Path -from mflux import Flux1, Config, ModelConfig +from mflux import Flux1, Config, ModelConfig, StopImageGenerationException def main(): @@ -14,6 +15,7 @@ def main(): parser.add_argument("--height", type=int, default=1024, help="Image height (Default is 1024)") parser.add_argument("--width", type=int, default=1024, help="Image width (Default is 1024)") parser.add_argument("--steps", type=int, default=None, help="Inference Steps") + parser.add_argument('--stepwise-image-output-dir', type=str, default=None, help='Output dir to write step-wise images and their final composite image to.') parser.add_argument("--guidance", type=float, default=3.5, help="Guidance Scale (Default is 3.5)") parser.add_argument("--quantize", "-q", type=int, choices=[4, 8], default=None, help="Quantize the model (4 or 8, Default is None)") parser.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk") @@ -39,20 +41,24 @@ def main(): lora_scales=args.lora_scales, ) - # Generate an image - image = flux.generate_image( - seed=int(time.time()) if args.seed is None else args.seed, - prompt=args.prompt, - config=Config( - num_inference_steps=args.steps, - height=args.height, - width=args.width, - guidance=args.guidance, - ), - ) + try: + # Generate an image + image = flux.generate_image( + seed=int(time.time()) if args.seed is None else args.seed, + prompt=args.prompt, + 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 - image.save(path=args.output, export_json_metadata=args.metadata) + # Save the image + image.save(path=args.output, export_json_metadata=args.metadata) + except StopImageGenerationException as stop_exc: + print(stop_exc) if __name__ == "__main__": diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index 5a2d861..c8813d8 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/generate_controlnet.py @@ -1,7 +1,8 @@ import argparse import time +from pathlib import Path -from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig +from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig, StopImageGenerationException def main(): @@ -17,6 +18,7 @@ def main(): parser.add_argument("--height", type=int, default=1024, help="Image height (Default is 1024)") parser.add_argument("--width", type=int, default=1024, help="Image width (Default is 1024)") parser.add_argument("--steps", type=int, default=None, help="Inference Steps") + parser.add_argument('--stepwise-image-output-dir', type=str, default=None, help='Output dir to write step-wise images and their final composite image to.') parser.add_argument("--guidance", type=float, default=3.5, help="Guidance Scale (Default is 3.5)") parser.add_argument("--quantize", "-q", type=int, choices=[4, 8], default=None, help="Quantize the model (4 or 8, Default is None)") parser.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk") @@ -42,24 +44,28 @@ def main(): lora_scales=args.lora_scales, ) - # Generate an image - image = flux.generate_image( - seed=int(time.time()) if args.seed is None else args.seed, - prompt=args.prompt, - output=args.output, - controlnet_image_path=args.controlnet_image_path, - controlnet_save_canny=args.controlnet_save_canny, - config=ConfigControlnet( - num_inference_steps=args.steps, - height=args.height, - width=args.width, - guidance=args.guidance, - controlnet_strength=args.controlnet_strength, - ), - ) + try: + # Generate an image + image = flux.generate_image( + seed=int(time.time()) if args.seed is None else args.seed, + prompt=args.prompt, + output=args.output, + controlnet_image_path=args.controlnet_image_path, + controlnet_save_canny=args.controlnet_save_canny, + config=ConfigControlnet( + num_inference_steps=args.steps, + height=args.height, + width=args.width, + 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 - image.save(path=args.output, export_json_metadata=args.metadata) + # Save the image + image.save(path=args.output, export_json_metadata=args.metadata) + except StopImageGenerationException as stop_exc: + print(stop_exc) if __name__ == "__main__": diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index 56b4145..726df64 100644 --- a/src/mflux/post_processing/image_util.py +++ b/src/mflux/post_processing/image_util.py @@ -1,7 +1,7 @@ +import typing as t import json import logging from pathlib import Path - import PIL import PIL.Image import mlx.core as mx @@ -47,6 +47,17 @@ class ImageUtil: controlnet_strength=config.controlnet_strength if isinstance(config.config, ConfigControlnet) else None, ) + def to_composite_image(generated_images: t.List[GeneratedImage]) -> Image: + # stitch horizontally + total_width = sum(gen_img.image.width for gen_img in generated_images) + max_height = max(gen_img.image.height for gen_img in generated_images) + composite_img = Image.new("RGB", (total_width, max_height)) + current_x = 0 + for index, gen_img in enumerate(generated_images): + composite_img.paste(gen_img.image, (current_x, 0)) + current_x += gen_img.image.width + return composite_img + @staticmethod def _denormalize(images: mx.array) -> mx.array: return mx.clip((images / 2 + 0.5), 0, 1) @@ -119,7 +130,7 @@ class ImageUtil: if metadata is not None: ImageUtil._embed_metadata(metadata, file_path) log.info(f"Metadata embedded successfully at: {file_path}") - except Exception as e: + except Exception as e: # noqa: BLE001 log.error(f"Error saving image: {e}") @staticmethod @@ -145,5 +156,5 @@ class ImageUtil: # Save the image with metadata image.save(path, exif=exif_bytes) - except Exception as e: + except Exception as e: # noqa: BLE001 log.error(f"Error embedding metadata: {e}") diff --git a/tests/test_readme_example.sh b/tests/test_readme_example.sh index 3fc6827..6346e3b 100755 --- a/tests/test_readme_example.sh +++ b/tests/test_readme_example.sh @@ -1,7 +1,39 @@ +#!/bin/zsh -e +# ^ safe to assume Mac devs have zsh installed +# default since Catalina in 2019 + +mkdir -p /tmp/mflux-test + mflux-generate \ --prompt "Luxury food photograph" \ --model schnell \ --steps 2 \ --seed 2 \ - --height 1024 \ - --width 1024 + --height 512 \ + --width 512 \ + --output /tmp/mflux-test/luxury_food.png + +# generate an image of a blue bird, then use it as input for the following test +mflux-generate \ + --prompt "blue bird, morning, spring" \ + --model schnell \ + --steps 2 \ + --seed 24 \ + --height 512 \ + --width 512 \ + --stepwise-image-output-dir /tmp/mflux-test \ + --output /tmp/mflux-test/sf_blue_bird.png + +# use the image from the prior test, generate an image with similar visual structure +mflux-generate-controlnet \ + --prompt "yellow bird, afternoon, snowy mountain" \ + --model schnell \ + --controlnet-image-path /tmp/mflux-test/sf_blue_bird.png \ + --controlnet-strength 0.7 \ + --controlnet-save-canny \ + --steps 2 \ + --seed 42 \ + --height 512 \ + --width 512 \ + --output /tmp/mflux-test/controlnet_sf_yellow_bird.png \ + --stepwise-image-output-dir /tmp/mflux-test