From 3dba1e38e17c75af936f80d7078c8b3674b059ac Mon Sep 17 00:00:00 2001 From: Fabio Date: Sat, 14 Sep 2024 14:37:30 +0200 Subject: [PATCH] controlnet - first working version --- src/mflux/config/config.py | 2 +- src/mflux/config/runtime_config.py | 13 ++- .../flux_controlnet.py} | 91 ++++++++++++------- src/mflux/controlnet/utils_controlnet.py | 10 ++ src/mflux/models/transformer/feed_forward.py | 4 +- src/mflux/models/transformer/transformer.py | 20 +++- src/mflux/post_processing/image_util.py | 4 + src/mflux/weights/weight_handler.py | 28 +++++- trial.py | 26 +++--- 9 files changed, 141 insertions(+), 57 deletions(-) rename src/mflux/{flux/controlnet.py => controlnet/flux_controlnet.py} (77%) create mode 100644 src/mflux/controlnet/utils_controlnet.py diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py index 7686d3d..271628d 100644 --- a/src/mflux/config/config.py +++ b/src/mflux/config/config.py @@ -6,7 +6,7 @@ log = logging.getLogger(__name__) class Config: - precision: mx.Dtype = mx.bfloat16 + precision: mx.Dtype = mx.float16 def __init__( self, diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py index 33cd634..dabd8f5 100644 --- a/src/mflux/config/runtime_config.py +++ b/src/mflux/config/runtime_config.py @@ -1,13 +1,15 @@ +from math import e +from cycler import V import mlx.core as mx import numpy as np -from mflux.config.config import Config +from mflux.config.config import Config, ConfigControlnet from mflux.config.model_config import ModelConfig class RuntimeConfig: - def __init__(self, config: Config, model_config: ModelConfig): + 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) @@ -35,6 +37,13 @@ class RuntimeConfig: @property def num_train_steps(self) -> int: return self.model_config.num_train_steps + + @property + def controlnet_conditioning_scale(self) -> float: + if isinstance(self.config, ConfigControlnet): + return self.config.controlnet_strength + else: + return ValueError("Controlnet conditioning scale is only available for ConfigControlnet") @staticmethod def _create_sigmas(config, model) -> mx.array: diff --git a/src/mflux/flux/controlnet.py b/src/mflux/controlnet/flux_controlnet.py similarity index 77% rename from src/mflux/flux/controlnet.py rename to src/mflux/controlnet/flux_controlnet.py index 54b4a22..ebb6b15 100644 --- a/src/mflux/flux/controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -10,6 +10,7 @@ from tqdm import tqdm from mflux.config.config import Config, ConfigControlnet from mflux.config.model_config import ModelConfig from mflux.config.runtime_config import RuntimeConfig +from mflux.controlnet.utils_controlnet import preprocess_canny 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 @@ -28,6 +29,8 @@ 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.time_text_embed import TimeTextEmbed +import numpy as np +import cv2 import logging @@ -84,18 +87,18 @@ class Flux1Controlnet: if weights.quantization_level is not None: self._set_model_weights(weights) - self.transformer_controlnet = TransformerControlnet(model_config=model_config, local_path=controlnet_path, quantize=quantize) + self.transformer_controlnet = TransformerControlnet(model_config=model_config) - weights_controlnet = WeightHandler.load_transformer(root_path=controlnet_path) - if weights_controlnet.quantization_level is None: + weights_controlnet, ctrlnet_quantization_level = WeightHandler.load_controlnet_transformer(controlnet_path=controlnet_path) + if ctrlnet_quantization_level is None: self.transformer_controlnet.update(weights_controlnet) self.bits = None - if quantize is not None or weights.quantization_level is not None: - self.bits = weights_controlnet.quantization_level if weights_controlnet.quantization_level is not None else quantize - nn.quantize(self.transformer_controlnet, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 64, group_size=64, bits=self.bits) + if quantize is not None or ctrlnet_quantization_level is not None: + self.bits = ctrlnet_quantization_level if ctrlnet_quantization_level is not None else quantize + nn.quantize(self.transformer_controlnet, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 128, group_size=128, bits=self.bits) - if weights_controlnet.quantization_level is not None: + if ctrlnet_quantization_level is not None: self.transformer_controlnet.update(weights_controlnet) def generate_image(self, seed: int, prompt: str, control_image: PIL.Image.Image, config: ConfigControlnet = ConfigControlnet()) -> GeneratedImage: @@ -112,10 +115,10 @@ class Flux1Controlnet: shape=[1, (config.height // 16) * (config.width // 16), 64], key=mx.random.key(seed) ) - control_cond = ImageUtil.to_array(control_image) - control_cond = self.vae.encode(control_cond) - control_cond = (control_cond - self.vae.shift_factor) * self.vae.scaling_factor - control_cond = Flux1Controlnet._pack_latents(control_cond, config.height, config.width) + control_image = preprocess_canny(control_image) + controlnet_cond = ImageUtil.to_array(control_image) + controlnet_cond = self.vae.encode(controlnet_cond) + controlnet_cond = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width) # 2. Embedd the prompt t5_tokens = self.t5_tokenizer.tokenize(prompt) @@ -124,12 +127,12 @@ class Flux1Controlnet: pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens) for t in time_steps: - controlnet_samples = self.transformer_controlnet( + controlnet_block_samples = self.transformer_controlnet.forward( t=t, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, hidden_states=latents, - control_cond=control_cond, + controlnet_cond=controlnet_cond, config=config, ) # 3.t Predict the noise @@ -139,7 +142,8 @@ class Flux1Controlnet: pooled_prompt_embeds=pooled_prompt_embeds, hidden_states=latents, config=config, - controlnet_samples=controlnet_samples, + controlnet_block_samples=controlnet_block_samples, + # controlnet_single_block_samples=controlnet_single_block_samples, ) # 4.t Take one denoise step @@ -219,11 +223,10 @@ class Flux1Controlnet: _save_weights(self.transformer, "transformer") _save_weights(self.clip_text_encoder, "text_encoder") _save_weights(self.t5_text_encoder, "text_encoder_2") + _save_weights(self.transformer_controlnet, "transformer_controlnet") -class ControlNetOutput: - controlnet_block_samples: Tuple[mx.array] - controlnet_single_block_samples: Tuple[mx.array] +ControlNetOutput = Tuple[list[mx.array], list[mx.array]] class TransformerControlnet(nn.Module): @@ -233,10 +236,14 @@ class TransformerControlnet(nn.Module): 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(19)] - self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(38)] - self.norm_out = AdaLayerNormContinuous(3072, 3072) - self.proj_out = nn.Linear(3072, 64) + self.transformer_blocks = [JointTransformerBlock(i) for i in range(5)] + # self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(38)] + + 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(5)] + + # self.controlnet_single_blocks = [nn.Linear(3072, 3072) for _ in range(38)] def forward( self, @@ -244,11 +251,15 @@ class TransformerControlnet(nn.Module): prompt_embeds: mx.array, pooled_prompt_embeds: mx.array, hidden_states: mx.array, + controlnet_cond: mx.array, config: RuntimeConfig, ) -> ControlNetOutput: 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_conditioning_scale + guidance = mx.broadcast_to(config.guidance * config.num_train_steps, (1,)).astype(config.precision) text_embeddings = self.time_text_embed.forward(time_step, pooled_prompt_embeds, guidance) encoder_hidden_states = self.context_embedder(prompt_embeds) @@ -257,6 +268,7 @@ class TransformerControlnet(nn.Module): ids = mx.concatenate((txt_ids, img_ids), axis=1) image_rotary_emb = self.pos_embed.forward(ids) + block_samples = () for block in self.transformer_blocks: encoder_hidden_states, hidden_states = block.forward( hidden_states=hidden_states, @@ -264,18 +276,33 @@ class TransformerControlnet(nn.Module): text_embeddings=text_embeddings, rotary_embeddings=image_rotary_emb ) + block_samples = block_samples + (hidden_states,) hidden_states = mx.concatenate([encoder_hidden_states, hidden_states], axis=1) - for block in self.single_transformer_blocks: - hidden_states = block.forward( - hidden_states=hidden_states, - text_embeddings=text_embeddings, - rotary_embeddings=image_rotary_emb - ) + # 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,) - hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...] - hidden_states = self.norm_out.forward(hidden_states, text_embeddings) - hidden_states = self.proj_out(hidden_states) - noise = hidden_states - return noise \ No newline at end of file + + # single_block_samples = () + # for block in self.single_transformer_blocks: + # ctrlnet_hidden_states = block.forward( + # hidden_states=ctrlnet_hidden_states, + # text_embeddings=text_embeddings, + # rotary_embeddings=image_rotary_emb + # ) + # single_block_samples = single_block_samples + (ctrlnet_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 \ No newline at end of file diff --git a/src/mflux/controlnet/utils_controlnet.py b/src/mflux/controlnet/utils_controlnet.py new file mode 100644 index 0000000..5de6b06 --- /dev/null +++ b/src/mflux/controlnet/utils_controlnet.py @@ -0,0 +1,10 @@ +import cv2 +import numpy as np +import PIL + +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]) + image_to_canny = np.concatenate([image_to_canny, image_to_canny, image_to_canny], axis=2) + return PIL.Image.fromarray(image_to_canny) diff --git a/src/mflux/models/transformer/feed_forward.py b/src/mflux/models/transformer/feed_forward.py index c5536a0..cb01d8e 100644 --- a/src/mflux/models/transformer/feed_forward.py +++ b/src/mflux/models/transformer/feed_forward.py @@ -6,8 +6,8 @@ class FeedForward(nn.Module): def __init__(self, activation_function): super().__init__() - self.linear1 = nn.Linear(3072, 6144) - self.linear2 = nn.Linear(6144, 3072) + self.linear1 = nn.Linear(3072, 12288) + self.linear2 = nn.Linear(12288, 3072) self.activation_function = activation_function def forward(self, hidden_states: mx.array) -> mx.array: diff --git a/src/mflux/models/transformer/transformer.py b/src/mflux/models/transformer/transformer.py index dfc4979..6afe28d 100644 --- a/src/mflux/models/transformer/transformer.py +++ b/src/mflux/models/transformer/transformer.py @@ -1,15 +1,17 @@ +from typing import Tuple import mlx.core as mx from mlx import nn from mflux.config.model_config import ModelConfig from mflux.config.runtime_config import RuntimeConfig -from mflux.flux.controlnet import ControlNetOutput from mflux.models.transformer.ada_layer_norm_continous import AdaLayerNormContinuous 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.time_text_embed import TimeTextEmbed +import numpy as np + class Transformer(nn.Module): @@ -31,7 +33,8 @@ class Transformer(nn.Module): pooled_prompt_embeds: mx.array, hidden_states: mx.array, config: RuntimeConfig, - controlnet_samples: ControlNetOutput | None = None + controlnet_block_samples: Tuple[mx.array] | None = None, + controlnet_single_block_samples: Tuple[mx.array] | None = None, ) -> mx.array: time_step = config.sigmas[t] * config.num_train_steps time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision) @@ -44,22 +47,31 @@ class Transformer(nn.Module): ids = mx.concatenate((txt_ids, img_ids), axis=1) image_rotary_emb = self.pos_embed.forward(ids) - for block in self.transformer_blocks: + for idx, block in enumerate(self.transformer_blocks): encoder_hidden_states, hidden_states = block.forward( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, text_embeddings=text_embeddings, rotary_embeddings=image_rotary_emb ) + if controlnet_block_samples is not None: + interval_control = len(self.transformer_blocks) / len(controlnet_block_samples) + interval_control = int(np.ceil(interval_control)) + hidden_states = hidden_states + controlnet_block_samples[idx // interval_control] hidden_states = mx.concatenate([encoder_hidden_states, hidden_states], axis=1) - for block in self.single_transformer_blocks: + for idx, block in enumerate(self.single_transformer_blocks): hidden_states = block.forward( hidden_states=hidden_states, text_embeddings=text_embeddings, rotary_embeddings=image_rotary_emb ) + if controlnet_single_block_samples is not None: + hidden_states[:, encoder_hidden_states.shape[1] :, ...] = ( + hidden_states[:, encoder_hidden_states.shape[1] :, ...] + + controlnet_single_block_samples[idx] + ) hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...] hidden_states = self.norm_out.forward(hidden_states, text_embeddings) diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index 5c60089..2ef5b5e 100644 --- a/src/mflux/post_processing/image_util.py +++ b/src/mflux/post_processing/image_util.py @@ -71,3 +71,7 @@ class ImageUtil: array = mx.transpose(array, (0, 3, 1, 2)) array = ImageUtil._normalize(array) return array + + @staticmethod + def load_image(path: str) -> Image.Image: + return Image.open(path) diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py index fc485ca..2c2acfe 100644 --- a/src/mflux/weights/weight_handler.py +++ b/src/mflux/weights/weight_handler.py @@ -85,6 +85,28 @@ class WeightHandler: "linear2": block["ff_context"]["net"][2] } return weights, quantization_level + + @staticmethod + def load_controlnet_transformer(controlnet_path: Path | None = None) -> (dict, int): + weights, quantization_level = WeightHandler._get_weights("transformer", weights_path=controlnet_path) + + # Quantized weights (i.e. ones exported from this project) don't need any post-processing. + if quantization_level is not None: + return weights, quantization_level + + # Reshape and process the huggingface weights + if "transformer_blocks" in weights: + for block in weights["transformer_blocks"]: + block["ff"] = { + "linear1": block["ff"]["net"][0]["proj"], + "linear2": block["ff"]["net"][2] + } + if block.get("ff_context") is not None: + block["ff_context"] = { + "linear1": block["ff_context"]["net"][0]["proj"], + "linear2": block["ff_context"]["net"][2] + } + return weights, quantization_level @staticmethod def load_vae(root_path: Path) -> (dict, int): @@ -104,7 +126,7 @@ class WeightHandler: return weights, quantization_level @staticmethod - def _get_weights(model_name: str, root_path: Path | None = None, lora_path: str | None = None) -> (dict, int): + def _get_weights(model_name: str, root_path: Path | None = None, weights_path: str | None = None) -> (dict, int): weights = [] quantization_level = None @@ -114,8 +136,8 @@ class WeightHandler: weight = list(mx.load(str(file)).items()) weights.extend(weight) - if lora_path and root_path is None: - weight = list(mx.load(lora_path).items()) + if weights_path and root_path is None: + weight = list(mx.load(weights_path).items()) weights.extend(weight) # Non huggingface weights (i.e. ones exported from this project) don't need any reshaping. diff --git a/trial.py b/trial.py index 53df634..533905f 100644 --- a/trial.py +++ b/trial.py @@ -1,35 +1,35 @@ -import argparse import os import sys -import time +from mflux.post_processing.image_util import ImageUtil sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) from mflux.config.model_config import ModelConfig -from mflux.config.config import Config -from mflux.flux.flux import Flux1 +from mflux.config.config import ConfigControlnet +from mflux.controlnet.flux_controlnet import Flux1Controlnet - -prompt = "Luxury food photograph" +prompt = "A girl in Venice, 25 years old" # Load the model -flux = Flux1( +flux = Flux1Controlnet( model_config=ModelConfig.from_alias("dev"), - quantize=None, - local_path=None, - lora_paths=["diffusion_pytorch_model.safetensors"], - lora_scales=None, + quantize=8, + controlnet_path="diffusion_pytorch_model.safetensors", ) +control_image = ImageUtil.load_image("image_51.png") + # Generate an image image = flux.generate_image( - seed=3, + seed=4, prompt=prompt, - config=Config( + control_image=control_image, + config=ConfigControlnet( num_inference_steps=10, height=256, width=512, guidance=3.5, + controlnet_conditioning_scale=1.6, ) )