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, + )