From fa99cf0cc8d5abacaab2e2cf122d74504b5af906 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 1 Feb 2025 21:22:47 +0100 Subject: [PATCH] Refactor: Controlnet transformer & Flux controlnet --- src/mflux/controlnet/flux_controlnet.py | 89 +++++------ .../controlnet/transformer_controlnet.py | 143 +++++++++++------- .../controlnet/weight_handler_controlnet.py | 12 +- 3 files changed, 144 insertions(+), 100 deletions(-) diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index b961514..eda6a4b 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -1,6 +1,4 @@ -import logging from pathlib import Path -from typing import TYPE_CHECKING import mlx.core as mx from tqdm import tqdm @@ -18,6 +16,7 @@ from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder from mflux.models.transformer.transformer import Transformer from mflux.models.vae.vae import VAE from mflux.post_processing.array_util import ArrayUtil +from mflux.post_processing.generated_image import GeneratedImage from mflux.post_processing.image_util import ImageUtil from mflux.post_processing.stepwise_handler import StepwiseHandler from mflux.tokenizer.clip_tokenizer import TokenizerCLIP @@ -28,14 +27,6 @@ from mflux.weights.weight_handler import WeightHandler from mflux.weights.weight_handler_lora import WeightHandlerLoRA from mflux.weights.weight_util import WeightUtil -if TYPE_CHECKING: - from mflux.post_processing.generated_image import GeneratedImage - - -log = logging.getLogger(__name__) - -CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny" - class Flux1Controlnet: def __init__( @@ -61,7 +52,7 @@ class Flux1Controlnet: # Initialize the models self.vae = VAE() - self.transformer = Transformer(model_config, num_transformer_blocks=weights.num_transformer_blocks()) + 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() @@ -80,8 +71,8 @@ class Flux1Controlnet: WeightHandlerLoRA.set_lora_weights(transformer=self.transformer, loras=lora_weights) # Set Controlnet weights - weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer(controlnet_id=CONTROLNET_ID) - self.transformer_controlnet = TransformerControlnet(model_config=model_config, num_blocks=weights_controlnet.config["num_layers"], num_single_blocks=weights_controlnet.config["num_single_layers"]) # fmt:off + 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, @@ -89,15 +80,15 @@ class Flux1Controlnet: ) def generate_image( - self, - seed: int, - prompt: str, - output: str, - controlnet_image_path: str, - controlnet_save_canny: bool = False, - config: ConfigControlnet = ConfigControlnet(), - stepwise_output_dir: Path = None, - ) -> "GeneratedImage": # fmt: off + self, + seed: int, + prompt: str, + output: str, + controlnet_image_path: str, + controlnet_save_canny: bool = False, + config: ConfigControlnet = ConfigControlnet(), + stepwise_output_dir: Path = None, + ) -> 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)) @@ -110,16 +101,8 @@ class Flux1Controlnet: output_dir=stepwise_output_dir, ) - # Embed the controlnet reference image - 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) + # 0. Embed the controlnet reference image + controlnet_condition = self._embed_image(config, controlnet_image_path, controlnet_save_canny, output) # 1. Create the initial latents latents = LatentCreator.create(seed=seed, height=config.height, width=config.width) @@ -130,35 +113,35 @@ class Flux1Controlnet: prompt_embeds = self.t5_text_encoder(t5_tokens) pooled_prompt_embeds = self.clip_text_encoder(clip_tokens) - for t in time_steps: + for gen_step, t in enumerate(time_steps, 1): try: - # Compute controlnet samples + # 3.t Compute controlnet samples controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet( t=t, + config=config, + hidden_states=latents, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, - hidden_states=latents, - controlnet_cond=controlnet_cond, - config=config, + controlnet_condition=controlnet_condition, ) - # 3.t Predict the noise - noise = self.transformer.predict( + # 4.t Predict the noise + noise = self.transformer( t=t, + config=config, + hidden_states=latents, 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 + # 5.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(t, latents) + stepwise_handler.process_step(gen_step, latents) # Evaluate to enable progress tracking mx.eval(latents) @@ -182,6 +165,26 @@ 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/controlnet/transformer_controlnet.py b/src/mflux/controlnet/transformer_controlnet.py index 6d9d4ad..8558965 100644 --- a/src/mflux/controlnet/transformer_controlnet.py +++ b/src/mflux/controlnet/transformer_controlnet.py @@ -4,12 +4,8 @@ from mlx import nn from mflux.config.model_config import ModelConfig from mflux.config.runtime_config import RuntimeConfig from mflux.models.transformer.embed_nd import EmbedND -from mflux.models.transformer.joint_transformer_block import ( - JointTransformerBlock, -) -from mflux.models.transformer.single_transformer_block import ( - SingleTransformerBlock, -) +from mflux.models.transformer.joint_transformer_block import JointTransformerBlock +from mflux.models.transformer.single_transformer_block import SingleTransformerBlock from mflux.models.transformer.time_text_embed import TimeTextEmbed from mflux.models.transformer.transformer import Transformer @@ -18,80 +14,117 @@ class TransformerControlnet(nn.Module): def __init__( self, model_config: ModelConfig, - num_blocks: int, - num_single_blocks: int, + num_transformer_blocks: int = 5, + num_single_transformer_blocks: int = 0, ): super().__init__() self.pos_embed = EmbedND() self.x_embedder = nn.Linear(64, 3072) self.time_text_embed = TimeTextEmbed(model_config=model_config) self.context_embedder = nn.Linear(4096, 3072) - self.transformer_blocks = [JointTransformerBlock(i) for i in range(num_blocks)] - self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(num_single_blocks)] + self.transformer_blocks = [JointTransformerBlock(i) for i in range(num_transformer_blocks)] + self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(num_single_transformer_blocks)] - zero_init = nn.init.constant(0) - self.controlnet_x_embedder = nn.Linear(64, 3072).apply(zero_init) - self.controlnet_blocks = [nn.Linear(3072, 3072).apply(zero_init) for _ in range(num_blocks)] - - self.controlnet_single_blocks = [nn.Linear(3072, 3072) for _ in range(num_single_blocks)] + self.controlnet_x_embedder = nn.Linear(64, 3072).apply(nn.init.constant(0)) + self.controlnet_blocks = [nn.Linear(3072, 3072).apply(nn.init.constant(0)) for _ in range(num_transformer_blocks)] # fmt: off + self.controlnet_single_blocks = [nn.Linear(3072, 3072) for _ in range(num_single_transformer_blocks)] def __call__( self, t: int, + config: RuntimeConfig, + hidden_states: mx.array, prompt_embeds: mx.array, pooled_prompt_embeds: mx.array, - hidden_states: mx.array, - controlnet_cond: mx.array, - config: RuntimeConfig, + controlnet_condition: mx.array, ) -> (list[mx.array], list[mx.array]): - time_step = config.sigmas[t] * config.num_train_steps - time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision) - hidden_states = self.x_embedder(hidden_states) - hidden_states = hidden_states + self.controlnet_x_embedder(controlnet_cond) - conditioning_scale = config.config.controlnet_strength - - guidance = mx.broadcast_to(config.guidance * config.num_train_steps, (1,)).astype(config.precision) - text_embeddings = self.time_text_embed(time_step, pooled_prompt_embeds, guidance) + # 1. Create embeddings + hidden_states = self.x_embedder(hidden_states) + self.controlnet_x_embedder(controlnet_condition) encoder_hidden_states = self.context_embedder(prompt_embeds) - txt_ids = Transformer.prepare_text_ids(seq_len=prompt_embeds.shape[1]) - img_ids = Transformer.prepare_latent_image_ids(height=config.height, width=config.width) - ids = mx.concatenate((txt_ids, img_ids), axis=1) - image_rotary_emb = self.pos_embed(ids) + text_embeddings = Transformer.compute_text_embeddings(t, pooled_prompt_embeds, self.time_text_embed, config) + image_rotary_embeddings = Transformer.compute_rotary_embeddings(prompt_embeds, self.pos_embed, config) - block_samples = () - for block in self.transformer_blocks: - encoder_hidden_states, hidden_states = block( + # 2. Run the joint transformer blocks + controlnet_block_samples = [] + for idx, block in enumerate(self.transformer_blocks): + encoder_hidden_states, hidden_states = self._apply_joint_transformer_block( + idx=idx, + block=block, + config=config, hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, text_embeddings=text_embeddings, - rotary_embeddings=image_rotary_emb, + image_rotary_embeddings=image_rotary_embeddings, + controlnet_block_samples=controlnet_block_samples, ) - block_samples = block_samples + (hidden_states,) + # 3. Concat the hidden states hidden_states = mx.concatenate([encoder_hidden_states, hidden_states], axis=1) - # controlnet block - controlnet_block_samples = () - for block_sample, controlnet_block in zip(block_samples, self.controlnet_blocks): - block_sample = controlnet_block(block_sample) - controlnet_block_samples = controlnet_block_samples + (block_sample,) - - single_block_samples = () - for block in self.single_transformer_blocks: - hidden_states = block( + # 4. Run the single transformer blocks + controlnet_single_block_samples = [] + for idx, block in enumerate(self.single_transformer_blocks): + hidden_states = self._apply_single_transformer_block( + idx=idx, + block=block, + config=config, hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, text_embeddings=text_embeddings, - rotary_embeddings=image_rotary_emb, + image_rotary_embeddings=image_rotary_embeddings, + controlnet_single_block_samples=controlnet_single_block_samples, ) - single_block_samples = single_block_samples + (hidden_states[:, encoder_hidden_states.shape[1] :],) - - controlnet_single_block_samples = () - for single_block_sample, controlnet_block in zip(single_block_samples, self.controlnet_single_blocks): - single_block_sample = controlnet_block(single_block_sample) - controlnet_single_block_samples = controlnet_single_block_samples + (single_block_sample,) - - # scaling - controlnet_block_samples = [sample * conditioning_scale for sample in controlnet_block_samples] - controlnet_single_block_samples = [sample * conditioning_scale for sample in controlnet_single_block_samples] return controlnet_block_samples, controlnet_single_block_samples + + def _apply_single_transformer_block( + self, + idx: int, + config: RuntimeConfig, + block: SingleTransformerBlock, + hidden_states: mx.array, + encoder_hidden_states: mx.array, + text_embeddings: mx.array, + image_rotary_embeddings: mx.array, + controlnet_single_block_samples: list[mx.array], + ) -> mx.array: + # 1. Apply single transformer block + hidden_states = block( + hidden_states=hidden_states, + text_embeddings=text_embeddings, + rotary_embeddings=image_rotary_embeddings, + ) + + # 2. Apply controlnet block + states = hidden_states[:, encoder_hidden_states.shape[1] :] + controlnet_sample = self.controlnet_single_blocks[idx](states) + scaled_controlnet_sample = controlnet_sample * config.config.controlnet_strength + controlnet_single_block_samples.append(scaled_controlnet_sample) + + return hidden_states + + def _apply_joint_transformer_block( + self, + idx: int, + config: RuntimeConfig, + block: JointTransformerBlock, + hidden_states: mx.array, + encoder_hidden_states: mx.array, + text_embeddings: mx.array, + image_rotary_embeddings: mx.array, + controlnet_block_samples: list[mx.array], + ) -> tuple[mx.array, mx.array]: + # 1. Apply joint transformer block + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + text_embeddings=text_embeddings, + rotary_embeddings=image_rotary_embeddings, + ) + + # 2. Apply controlnet block + controlnet_sample = self.controlnet_blocks[idx](hidden_states) + scaled_controlnet_example = controlnet_sample * config.config.controlnet_strength + controlnet_block_samples.append(scaled_controlnet_example) + + return encoder_hidden_states, hidden_states diff --git a/src/mflux/controlnet/weight_handler_controlnet.py b/src/mflux/controlnet/weight_handler_controlnet.py index 0cdfd70..f8b80c8 100644 --- a/src/mflux/controlnet/weight_handler_controlnet.py +++ b/src/mflux/controlnet/weight_handler_controlnet.py @@ -8,6 +8,8 @@ from mlx.utils import tree_unflatten from mflux.weights.weight_handler import MetaData from mflux.weights.weight_util import WeightUtil +CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny" + class WeightHandlerControlnet: def __init__(self, meta_data: MetaData, config: dict, controlnet_transformer: dict | None = None): @@ -16,8 +18,8 @@ class WeightHandlerControlnet: self.config = config @staticmethod - def load_controlnet_transformer(controlnet_id: str) -> "WeightHandlerControlnet": - controlnet_path = Path(snapshot_download(repo_id=controlnet_id, allow_patterns=["*.safetensors", "config.json"])) # fmt:off + def load_controlnet_transformer() -> "WeightHandlerControlnet": + controlnet_path = Path(snapshot_download(repo_id=CONTROLNET_ID, allow_patterns=["*.safetensors", "config.json"])) # fmt:off file = next(controlnet_path.glob("diffusion_pytorch_model.safetensors")) quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level") weights = list(mx.load(str(file)).items()) @@ -60,3 +62,9 @@ class WeightHandlerControlnet: controlnet_transformer=weights, meta_data=MetaData(quantization_level=quantization_level) ) # fmt:off + + def num_transformer_blocks(self) -> int: + return self.config["num_layers"] + + def num_single_transformer_blocks(self) -> int: + return self.config["num_single_layers"]