From 22f98fc5a46b90175040cd6f20c89c649e2cbdce Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 8 Feb 2025 05:09:33 +0100 Subject: [PATCH 1/3] Refactor: Add separate Flux initializer --- src/mflux/config/runtime_config.py | 6 +- src/mflux/controlnet/controlnet_util.py | 44 +++++++--- src/mflux/controlnet/flux_controlnet.py | 102 +++++++--------------- src/mflux/flux/flux.py | 55 ++++-------- src/mflux/flux/flux_initializer.py | 111 ++++++++++++++++++++++++ 5 files changed, 199 insertions(+), 119 deletions(-) create mode 100644 src/mflux/flux/flux_initializer.py diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py index 25723b7..3c0cbc1 100644 --- a/src/mflux/config/runtime_config.py +++ b/src/mflux/config/runtime_config.py @@ -10,7 +10,11 @@ logger = logging.getLogger(__name__) class RuntimeConfig: - def __init__(self, config: Config | ConfigControlnet, model_config): + def __init__( + self, + config: Config | ConfigControlnet, + model_config: ModelConfig, + ): self.config = config self.model_config = model_config self.sigmas = self._create_sigmas(config, model_config) diff --git a/src/mflux/controlnet/controlnet_util.py b/src/mflux/controlnet/controlnet_util.py index 64cb27c..b852163 100644 --- a/src/mflux/controlnet/controlnet_util.py +++ b/src/mflux/controlnet/controlnet_util.py @@ -2,15 +2,47 @@ 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 + log = logging.getLogger(__name__) class ControlnetUtil: @staticmethod - def preprocess_canny(img: PIL.Image) -> PIL.Image: + def encode_image( + vae: VAE, + config: RuntimeConfig, + controlnet_image_path: str, + controlnet_save_canny: bool, + output: str, + ) -> mx.array: + 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._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 + + @staticmethod + def _preprocess_canny(img: PIL.Image) -> PIL.Image: image_to_canny = np.array(img) image_to_canny = cv2.Canny(image_to_canny, 100, 200) image_to_canny = np.array(image_to_canny[:, :, None]) @@ -18,16 +50,8 @@ class ControlnetUtil: return PIL.Image.fromarray(image_to_canny) @staticmethod - def scale_image(height: int, width: int, img: PIL.Image) -> PIL.Image: + def _scale_image(height: int, width: int, img: PIL.Image) -> PIL.Image: if height != img.height or width != img.width: log.warning(f"Control image has different dimensions than the model. Resizing to {width}x{height}") img = img.resize((width, height), PIL.Image.LANCZOS) return img - - @staticmethod - def save_canny_image(control_image: PIL.Image, path: str): - from mflux import ImageUtil - - base, ext = os.path.splitext(path) - new_filename = f"{base}_controlnet_canny{ext}" - ImageUtil.save_image(control_image, new_filename) diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index eda6a4b..2642739 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -1,6 +1,7 @@ from pathlib import Path import mlx.core as mx +from mlx import nn from tqdm import tqdm from mflux.config.config import ConfigControlnet @@ -8,8 +9,8 @@ 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.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 from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder @@ -19,16 +20,16 @@ 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.tokenizer.clip_tokenizer import TokenizerCLIP -from mflux.tokenizer.t5_tokenizer import TokenizerT5 -from mflux.tokenizer.tokenizer_handler import TokenizerHandler from mflux.weights.model_saver import ModelSaver -from mflux.weights.weight_handler import WeightHandler -from mflux.weights.weight_handler_lora import WeightHandlerLoRA -from mflux.weights.weight_util import WeightUtil -class Flux1Controlnet: +class Flux1Controlnet(nn.Module): + vae: VAE + transformer: Transformer + transformer_controlnet: TransformerControlnet + t5_text_encoder: T5Encoder + clip_text_encoder: CLIPEncoder + def __init__( self, model_config: ModelConfig, @@ -38,45 +39,14 @@ class Flux1Controlnet: lora_scales: list[float] | None = None, controlnet_path: str | None = None, ): - self.lora_paths = lora_paths - self.lora_scales = lora_scales - self.model_config = model_config - - # Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface - tokenizers = TokenizerHandler(model_config.model_name, self.model_config.max_sequence_length, local_path) - self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=self.model_config.max_sequence_length) - self.clip_tokenizer = TokenizerCLIP(tokenizers.clip) - - # Load the weights - weights = WeightHandler.load_regular_weights(repo_id=model_config.model_name, local_path=local_path) - - # Initialize the models - self.vae = VAE() - self.transformer = Transformer(model_config, num_transformer_blocks=weights.num_transformer_blocks(), num_single_transformer_blocks=weights.num_single_transformer_blocks()) # fmt: off - self.t5_text_encoder = T5Encoder() - self.clip_text_encoder = CLIPEncoder() - - # Set the weights and quantize the model - self.bits = WeightUtil.set_weights_and_quantize( - quantize_arg=quantize, - weights=weights, - vae=self.vae, - transformer=self.transformer, - t5_text_encoder=self.t5_text_encoder, - clip_text_encoder=self.clip_text_encoder, - ) - - # Set LoRA weights - lora_weights = WeightHandlerLoRA.load_lora_weights(transformer=self.transformer, lora_files=lora_paths, lora_scales=lora_scales) # fmt:off - WeightHandlerLoRA.set_lora_weights(transformer=self.transformer, loras=lora_weights) - - # Set Controlnet weights - weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer() - self.transformer_controlnet = TransformerControlnet(model_config=model_config, num_transformer_blocks=weights_controlnet.num_transformer_blocks(), num_single_transformer_blocks=weights_controlnet.num_single_transformer_blocks()) # fmt:off - WeightUtil.set_controlnet_weights_and_quantize( - quantize_arg=quantize, - weights=weights_controlnet, - transformer_controlnet=self.transformer_controlnet, + super().__init__() + FluxInitializer.init_controlnet( + flux_model=self, + model_config=model_config, + quantize=quantize, + local_path=local_path, + lora_paths=lora_paths, + lora_scales=lora_scales, ) def generate_image( @@ -88,7 +58,7 @@ class Flux1Controlnet: controlnet_save_canny: bool = False, config: ConfigControlnet = ConfigControlnet(), stepwise_output_dir: Path = None, - ) -> GeneratedImage: # fmt: off + ) -> 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)) @@ -101,11 +71,21 @@ class Flux1Controlnet: output_dir=stepwise_output_dir, ) - # 0. Embed the controlnet reference image - controlnet_condition = self._embed_image(config, controlnet_image_path, controlnet_save_canny, output) + # 0. Encode the controlnet reference image + controlnet_condition = ControlnetUtil.encode_image( + vae=self.vae, + config=config, + controlnet_image_path=controlnet_image_path, + controlnet_save_canny=controlnet_save_canny, + output=output, + ) # 1. Create the initial latents - latents = LatentCreator.create(seed=seed, height=config.height, width=config.width) + latents = LatentCreator.create( + seed=seed, + height=config.height, + width=config.width + ) # fmt: off # 2. Embed the prompt t5_tokens = self.t5_tokenizer.tokenize(prompt) @@ -165,26 +145,6 @@ class Flux1Controlnet: controlnet_image_path=controlnet_image_path, ) - def _embed_image( - self, - config: RuntimeConfig, - controlnet_image_path: str, - controlnet_save_canny: bool, - output: str, - ): - control_image = ImageUtil.load_image(controlnet_image_path) - control_image = ControlnetUtil.scale_image(config.height, config.width, control_image) - control_image = ControlnetUtil.preprocess_canny(control_image) - - if controlnet_save_canny: - ControlnetUtil.save_canny_image(control_image, output) - - 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 = ArrayUtil.pack_latents(latents=controlnet_cond, height=config.height, width=config.width) - return controlnet_cond - def save_model(self, base_path: str) -> None: ModelSaver.save_model(self, self.bits, base_path) ModelSaver.save_weights(base_path, self.bits, self.transformer_controlnet, "transformer_controlnet") diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index 388e224..5dc6090 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -9,6 +9,7 @@ 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.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder @@ -18,16 +19,15 @@ 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.tokenizer.clip_tokenizer import TokenizerCLIP -from mflux.tokenizer.t5_tokenizer import TokenizerT5 -from mflux.tokenizer.tokenizer_handler import TokenizerHandler from mflux.weights.model_saver import ModelSaver -from mflux.weights.weight_handler import WeightHandler -from mflux.weights.weight_handler_lora import WeightHandlerLoRA -from mflux.weights.weight_util import WeightUtil class Flux1(nn.Module): + vae: VAE + transformer: Transformer + t5_text_encoder: T5Encoder + clip_text_encoder: CLIPEncoder + def __init__( self, model_config: ModelConfig, @@ -37,38 +37,15 @@ class Flux1(nn.Module): lora_scales: list[float] | None = None, ): super().__init__() - self.lora_paths = lora_paths - self.lora_scales = lora_scales - self.model_config = model_config - - # Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface - tokenizers = TokenizerHandler(model_config.model_name, self.model_config.max_sequence_length, local_path) - self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=self.model_config.max_sequence_length) - self.clip_tokenizer = TokenizerCLIP(tokenizers.clip) - - # Load the weights - weights = WeightHandler.load_regular_weights(repo_id=model_config.model_name, local_path=local_path) - - # Initialize the models - self.vae = VAE() - self.transformer = Transformer(model_config, num_transformer_blocks=weights.num_transformer_blocks(), num_single_transformer_blocks=weights.num_single_transformer_blocks()) # fmt: off - self.t5_text_encoder = T5Encoder() - self.clip_text_encoder = CLIPEncoder() - - # Set the weights and quantize the model - self.bits = WeightUtil.set_weights_and_quantize( - quantize_arg=quantize, - weights=weights, - vae=self.vae, - transformer=self.transformer, - t5_text_encoder=self.t5_text_encoder, - clip_text_encoder=self.clip_text_encoder, + FluxInitializer.init( + flux_model=self, + model_config=model_config, + quantize=quantize, + local_path=local_path, + lora_paths=lora_paths, + lora_scales=lora_scales, ) - # Set LoRA weights - lora_weights = WeightHandlerLoRA.load_lora_weights(transformer=self.transformer, lora_files=lora_paths, lora_scales=lora_scales) # fmt:off - WeightHandlerLoRA.set_lora_weights(transformer=self.transformer, loras=lora_weights) - def generate_image( self, seed: int, @@ -89,7 +66,11 @@ class Flux1(nn.Module): ) # 1. Create the initial latents - latents = LatentCreator.create_for_txt2img_or_img2img(seed, config, self.vae) + latents = LatentCreator.create_for_txt2img_or_img2img( + seed=seed, + vae=self.vae, + runtime_conf=config, + ) # 2. Embed the prompt t5_tokens = self.t5_tokenizer.tokenize(prompt) diff --git a/src/mflux/flux/flux_initializer.py b/src/mflux/flux/flux_initializer.py new file mode 100644 index 0000000..2e57ab1 --- /dev/null +++ b/src/mflux/flux/flux_initializer.py @@ -0,0 +1,111 @@ +from mflux.controlnet.transformer_controlnet import TransformerControlnet +from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet +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.tokenizer.clip_tokenizer import TokenizerCLIP +from mflux.tokenizer.t5_tokenizer import TokenizerT5 +from mflux.tokenizer.tokenizer_handler import TokenizerHandler +from mflux.weights.weight_handler import WeightHandler +from mflux.weights.weight_handler_lora import WeightHandlerLoRA +from mflux.weights.weight_util import WeightUtil + + +class FluxInitializer: + @staticmethod + def init( + flux_model, + model_config, + quantize: int | None, + local_path: str | None, + lora_paths: list[str] | None, + lora_scales: list[float] | None, + ) -> None: + # 0. Set paths and config for later + flux_model.lora_paths = lora_paths + flux_model.lora_scales = lora_scales + flux_model.model_config = model_config + + # 1. Initialize tokenizers + tokenizers = TokenizerHandler( + repo_id=model_config.model_name, + max_t5_length=model_config.max_sequence_length, + local_path=local_path, + ) + flux_model.t5_tokenizer = TokenizerT5( + tokenizer=tokenizers.t5, + max_length=model_config.max_sequence_length + ) # fmt: off + flux_model.clip_tokenizer = TokenizerCLIP( + tokenizer=tokenizers.clip, + ) + + # 2. Load the regular weights + weights = WeightHandler.load_regular_weights( + repo_id=model_config.model_name, + local_path=local_path + ) # fmt: off + + # 3. Initialize all models + flux_model.vae = VAE() + flux_model.transformer = Transformer( + model_config=model_config, + num_transformer_blocks=weights.num_transformer_blocks(), + num_single_transformer_blocks=weights.num_single_transformer_blocks(), + ) + flux_model.t5_text_encoder = T5Encoder() + flux_model.clip_text_encoder = CLIPEncoder() + + # 4. Apply weights and quantize the models + flux_model.bits = WeightUtil.set_weights_and_quantize( + quantize_arg=quantize, + weights=weights, + vae=flux_model.vae, + transformer=flux_model.transformer, + t5_text_encoder=flux_model.t5_text_encoder, + clip_text_encoder=flux_model.clip_text_encoder, + ) + + # 5. Set LoRA weights + lora_weights = WeightHandlerLoRA.load_lora_weights( + transformer=flux_model.transformer, + lora_files=lora_paths, + lora_scales=lora_scales, + ) + WeightHandlerLoRA.set_lora_weights( + transformer=flux_model.transformer, + loras=lora_weights + ) # fmt: off + + @staticmethod + def init_controlnet( + flux_model, + model_config, + quantize: int | None, + local_path: str | None, + lora_paths: list[str] | None, + lora_scales: list[float] | None, + ) -> None: + # 1. Start with same init as regular Flux + FluxInitializer.init( + flux_model=flux_model, + model_config=model_config, + quantize=quantize, + local_path=local_path, + lora_paths=lora_paths, + lora_scales=lora_scales, + ) + + # 2. Apply ControlNet-specific initialization + weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer() + flux_model.transformer_controlnet = TransformerControlnet( + model_config=model_config, + num_transformer_blocks=weights_controlnet.num_transformer_blocks(), + num_single_transformer_blocks=weights_controlnet.num_single_transformer_blocks(), + ) + WeightUtil.set_controlnet_weights_and_quantize( + quantize_arg=quantize, + weights=weights_controlnet, + transformer_controlnet=flux_model.transformer_controlnet, + ) From e4af40e4eeb647258def799bae58660eaeb9798f Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 8 Feb 2025 13:29:32 +0100 Subject: [PATCH 2/3] Refactor: Remove ConfigControlnet --- src/mflux/__init__.py | 3 +-- src/mflux/config/config.py | 13 +------------ src/mflux/config/runtime_config.py | 19 ++++++++----------- src/mflux/controlnet/flux_controlnet.py | 10 +++++----- src/mflux/flux/flux.py | 6 +++--- src/mflux/generate.py | 6 +++--- src/mflux/generate_controlnet.py | 10 +++++----- src/mflux/post_processing/image_util.py | 8 +++----- src/mflux/post_processing/stepwise_handler.py | 4 ++-- ...image_generation_controlnet_test_helper.py | 4 ++-- 10 files changed, 33 insertions(+), 50 deletions(-) diff --git a/src/mflux/__init__.py b/src/mflux/__init__.py index 83273fc..67b368f 100644 --- a/src/mflux/__init__.py +++ b/src/mflux/__init__.py @@ -1,4 +1,4 @@ -from mflux.config.config import Config, ConfigControlnet +from mflux.config.config import Config from mflux.config.model_config import ModelConfig, ModelLookup from mflux.controlnet.flux_controlnet import Flux1Controlnet from mflux.error.exceptions import StopImageGenerationException @@ -9,7 +9,6 @@ __all__ = [ "Flux1", "Flux1Controlnet", "Config", - "ConfigControlnet", "ModelConfig", "ModelLookup", "ImageUtil", diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py index d2b0584..c77147c 100644 --- a/src/mflux/config/config.py +++ b/src/mflux/config/config.py @@ -17,6 +17,7 @@ class Config: guidance: float = 4.0, init_image_path: Path | None = None, init_image_strength: float | None = None, + controlnet_strength: float | None = None, ): if width % 16 != 0 or height % 16 != 0: log.warning("Width and height should be multiples of 16. Rounding down.") @@ -26,16 +27,4 @@ class Config: self.guidance = guidance self.init_image_path = init_image_path self.init_image_strength = init_image_strength - - -class ConfigControlnet(Config): - def __init__( - self, - num_inference_steps: int = 4, - width: int = 1024, - height: int = 1024, - guidance: float = 4.0, - controlnet_strength: float = 1.0, - ): - super().__init__(num_inference_steps=num_inference_steps, width=width, height=height, guidance=guidance) self.controlnet_strength = controlnet_strength diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py index 3c0cbc1..096338d 100644 --- a/src/mflux/config/runtime_config.py +++ b/src/mflux/config/runtime_config.py @@ -3,7 +3,7 @@ import logging import mlx.core as mx import numpy as np -from mflux.config.config import Config, ConfigControlnet +from mflux.config.config import Config from mflux.config.model_config import ModelConfig logger = logging.getLogger(__name__) @@ -12,7 +12,7 @@ logger = logging.getLogger(__name__) class RuntimeConfig: def __init__( self, - config: Config | ConfigControlnet, + config: Config, model_config: ModelConfig, ): self.config = config @@ -54,23 +54,20 @@ class RuntimeConfig: @property def init_time_step(self) -> int: if self.config.init_image_path is None: - # text to image, always begin at time step 0 + # For text-to-image, always begin at time step 0. return 0 else: - # we skip to the time step as informed by the init_image_strength - # the higher the strength number, the more time steps we skip + # For image-to-image: higher strength means we skip more steps. strength = max(0.0, min(1.0, self.config.init_image_strength)) - # if the strength is too small to even influence the image - # help the user round up so the init_image has influence at step 1 t = max(1, int(self.num_inference_steps * strength)) return t @property - def controlnet_strength(self) -> float: - if isinstance(self.config, ConfigControlnet): + def controlnet_strength(self) -> float | None: + if self.config.controlnet_strength is not None: return self.config.controlnet_strength - else: - raise NotImplementedError("Controlnet conditioning scale is only available for ConfigControlnet") + + return None @staticmethod def _create_sigmas(config, model) -> mx.array: diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index 2642739..41c9b03 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -4,7 +4,7 @@ import mlx.core as mx from mlx import nn from tqdm import tqdm -from mflux.config.config import ConfigControlnet +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 @@ -56,10 +56,10 @@ class Flux1Controlnet(nn.Module): output: str, controlnet_image_path: str, controlnet_save_canny: bool = False, - config: ConfigControlnet = ConfigControlnet(), + config: Config = Config(), stepwise_output_dir: Path = None, ) -> GeneratedImage: - # Create a new runtime config based on the model type and input parameters + # Convert the user config to a runtime config with derived parameters. config = RuntimeConfig(config, self.model_config) time_steps = tqdm(range(config.num_inference_steps)) stepwise_handler = StepwiseHandler( @@ -135,14 +135,14 @@ class Flux1Controlnet(nn.Module): decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, + config=config, 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, controlnet_image_path=controlnet_image_path, + generation_time=time_steps.format_dict["elapsed"], ) def save_model(self, base_path: str) -> None: diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index 5dc6090..15ed094 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -53,7 +53,7 @@ class Flux1(nn.Module): config: Config = Config(), stepwise_output_dir: Path = None, ) -> GeneratedImage: - # Create a new runtime config based on the model type and input parameters + # Convert the user config to a runtime config with derived parameters. config = RuntimeConfig(config, self.model_config) time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) stepwise_handler = StepwiseHandler( @@ -108,15 +108,15 @@ class Flux1(nn.Module): decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, + config=config, seed=seed, prompt=prompt, quantization=self.bits, - generation_time=time_steps.format_dict["elapsed"], lora_paths=self.lora_paths, lora_scales=self.lora_scales, init_image_path=config.init_image_path, init_image_strength=config.init_image_strength, - config=config, + generation_time=time_steps.format_dict["elapsed"], ) @staticmethod diff --git a/src/mflux/generate.py b/src/mflux/generate.py index 653380e..72f5a49 100644 --- a/src/mflux/generate.py +++ b/src/mflux/generate.py @@ -14,7 +14,7 @@ def main(): parser.add_output_arguments() args = parser.parse_args() - # Load the model + # 1. Load the model flux = Flux1( model_config=ModelLookup.from_name(model_name=args.model, base_model=args.base_model), quantize=args.quantize, @@ -25,7 +25,7 @@ def main(): try: for seed_value in args.seed: - # Generate an image for each seed value + # 2. Generate an image for each seed value image = flux.generate_image( seed=seed_value, prompt=args.prompt, @@ -39,7 +39,7 @@ def main(): init_image_strength=args.init_image_strength, ), ) - # Save the image + # 3. Save the image image.save(path=args.output.format(seed=seed_value), 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 03e3b83..ccbf394 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/generate_controlnet.py @@ -1,6 +1,6 @@ from pathlib import Path -from mflux import ConfigControlnet, Flux1Controlnet, ModelLookup, StopImageGenerationException +from mflux import Config, Flux1Controlnet, ModelLookup, StopImageGenerationException from mflux.ui.cli.parsers import CommandLineParser @@ -13,7 +13,7 @@ def main(): parser.add_output_arguments() args = parser.parse_args() - # Load the model + # 1. Load the model flux = Flux1Controlnet( model_config=ModelLookup.from_name(model_name=args.model, base_model=args.base_model), quantize=args.quantize, @@ -24,7 +24,7 @@ def main(): try: for seed_value in args.seed: - # Generate an image for each seed value + # 2. Generate an image for each seed value image = flux.generate_image( seed=seed_value, prompt=args.prompt, @@ -32,7 +32,7 @@ def main(): 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( + config=Config( num_inference_steps=args.steps, height=args.height, width=args.width, @@ -41,7 +41,7 @@ def main(): ), ) - # Save the image + # 3. Save the image image.save(path=args.output.format(seed=seed_value), export_json_metadata=args.metadata) except StopImageGenerationException as stop_exc: print(stop_exc) diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index c7ccca4..a21da15 100644 --- a/src/mflux/post_processing/image_util.py +++ b/src/mflux/post_processing/image_util.py @@ -8,25 +8,23 @@ import numpy as np import piexif import PIL.Image -from mflux.config.config import ConfigControlnet +from mflux.config.runtime_config import RuntimeConfig from mflux.post_processing.generated_image import GeneratedImage log = logging.getLogger(__name__) -RuntimeConfig: t.TypeAlias = "mflux.config.runtime_config.RuntimeConfig" # noqa: F821 - class ImageUtil: @staticmethod def to_image( decoded_latents: mx.array, + config: RuntimeConfig, seed: int, prompt: str, quantization: int, generation_time: float, lora_paths: list[str], lora_scales: list[float], - config: RuntimeConfig, controlnet_image_path: str | None = None, init_image_path: str | None = None, init_image_strength: float | None = None, @@ -49,7 +47,7 @@ class ImageUtil: init_image_path=init_image_path, init_image_strength=init_image_strength, controlnet_image_path=controlnet_image_path, - controlnet_strength=config.controlnet_strength if isinstance(config.config, ConfigControlnet) else None, + controlnet_strength=config.controlnet_strength, ) @staticmethod diff --git a/src/mflux/post_processing/stepwise_handler.py b/src/mflux/post_processing/stepwise_handler.py index 60332bb..59dcc25 100644 --- a/src/mflux/post_processing/stepwise_handler.py +++ b/src/mflux/post_processing/stepwise_handler.py @@ -40,13 +40,13 @@ class StepwiseHandler: 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, - generation_time=self.time_steps.format_dict["elapsed"], lora_paths=self.flux.lora_paths, lora_scales=self.flux.lora_scales, - config=self.config, + generation_time=self.time_steps.format_dict["elapsed"], ) self.step_wise_images.append(stepwise_img) 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 9a47cd3..e2794a4 100644 --- a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py +++ b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py @@ -3,7 +3,7 @@ import os import numpy as np from PIL import Image -from mflux import ConfigControlnet, Flux1Controlnet, ModelConfig +from mflux import Config, Flux1Controlnet, ModelConfig from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper @@ -43,7 +43,7 @@ class ImageGeneratorControlnetTestHelper: output=str(output_image_path), controlnet_image_path=controlnet_image_path, controlnet_save_canny=False, - config=ConfigControlnet( + config=Config( num_inference_steps=steps, height=768, width=493, From 3601baf75ed07845c166eebf4a94064bd14e9ea7 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sun, 9 Feb 2025 11:29:50 +0100 Subject: [PATCH 3/3] 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,