From d799a924c430d74b02b71e8c1905946d96b2c055 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sun, 22 Sep 2024 19:30:21 +0200 Subject: [PATCH] Light refactor of controlnet class --- src/mflux/controlnet/controlnet_util.py | 11 +++++ src/mflux/controlnet/flux_controlnet.py | 54 +++++++++++++------------ src/mflux/generate_controlnet.py | 4 +- 3 files changed, 41 insertions(+), 28 deletions(-) diff --git a/src/mflux/controlnet/controlnet_util.py b/src/mflux/controlnet/controlnet_util.py index 6e38c75..e4eadf8 100644 --- a/src/mflux/controlnet/controlnet_util.py +++ b/src/mflux/controlnet/controlnet_util.py @@ -1,7 +1,11 @@ +import logging + import cv2 import numpy as np import PIL +log = logging.getLogger(__name__) + class ControlnetUtil: @staticmethod @@ -11,3 +15,10 @@ class ControlnetUtil: 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) + + @staticmethod + 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 diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index 0fba545..a03e883 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -1,6 +1,5 @@ import logging -import PIL.Image import mlx.core as mx from mlx import nn from tqdm import tqdm @@ -80,51 +79,51 @@ class Flux1Controlnet: if weights.quantization_level is not None: self._set_model_weights(weights) - weights_controlnet, ctrlnet_quantization_level, controlnet_config = WeightHandlerControlnet.load_controlnet_transformer( + weights_controlnet, controlnet_quantization_level, controlnet_config = WeightHandlerControlnet.load_controlnet_transformer( controlnet_id=CONTROLNET_ID - ) + ) # fmt: off self.transformer_controlnet = TransformerControlnet( model_config=model_config, num_blocks=controlnet_config["num_layers"], num_single_blocks=controlnet_config["num_single_layers"], ) - if ctrlnet_quantization_level is None: + if controlnet_quantization_level is None: self.transformer_controlnet.update(weights_controlnet) self.bits = None - 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 - # fmt: off - 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) - # fmt: on + if quantize is not None or controlnet_quantization_level is not None: + self.bits = controlnet_quantization_level if 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]) > 128, group_size=128, bits=self.bits) # fmt: off - if ctrlnet_quantization_level is not None: + if controlnet_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: # fmt: off + def generate_image( + self, + seed: int, + prompt: str, + control_image_path: str, + config: ConfigControlnet = ConfigControlnet() + ) -> 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)) - if config.height != control_image.height or config.width != control_image.width: - log.warning( - f"Control image has different dimensions than the model. Resizing to {config.width}x{config.height}" - ) - control_image = control_image.resize((config.width, config.height), PIL.Image.LANCZOS) + # Embedd the controlnet reference image + control_image = ImageUtil.load_image(control_image_path) + control_image = ControlnetUtil.scale_image(config.height, config.width, control_image) + control_image = ControlnetUtil.preprocess_canny(control_image) + 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 = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width) # 1. Create the initial latents latents = mx.random.normal( shape=[1, (config.height // 16) * (config.width // 16), 64], key=mx.random.key(seed) ) # fmt: off - control_image = ControlnetUtil.preprocess_canny(control_image) - controlnet_cond = ImageUtil.to_array(control_image) - controlnet_cong = self.vae.encode(controlnet_cond) - # the rescaling in the next line is not in the huggingface code, but without it the images from - # the chosen controlnet model are very bad - controlnet_cond = (controlnet_cong / self.vae.scaling_factor) + self.vae.shift_factor - controlnet_cond = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width) # 2. Embedd the prompt t5_tokens = self.t5_tokenizer.tokenize(prompt) @@ -133,7 +132,8 @@ class Flux1Controlnet: pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens) for t in time_steps: - ctrlnet_block_samples, ctrlnet_single_block_samples = self.transformer_controlnet.forward( + # Compute controlnet samples + controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet.forward( t=t, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, @@ -141,6 +141,7 @@ class Flux1Controlnet: controlnet_cond=controlnet_cond, config=config, ) + # 3.t Predict the noise noise = self.transformer.predict( t=t, @@ -148,8 +149,8 @@ class Flux1Controlnet: pooled_prompt_embeds=pooled_prompt_embeds, hidden_states=latents, config=config, - controlnet_block_samples=ctrlnet_block_samples, - controlnet_single_block_samples=ctrlnet_single_block_samples, + controlnet_block_samples=controlnet_block_samples, + controlnet_single_block_samples=controlnet_single_block_samples, ) # 4.t Take one denoise step @@ -171,6 +172,7 @@ class Flux1Controlnet: lora_paths=self.lora_paths, lora_scales=self.lora_scales, config=config, + controlnet_image_path=control_image_path, ) @staticmethod diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index 0e077f3..a8ff54e 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/generate_controlnet.py @@ -1,7 +1,7 @@ import argparse import time -from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig, ImageUtil +from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig def main(): @@ -45,7 +45,7 @@ def main(): image = flux.generate_image( seed=int(time.time()) if args.seed is None else args.seed, prompt=args.prompt, - control_image=ImageUtil.load_image(args.controlnet_image_path), + control_image_path=args.controlnet_image_path, config=ConfigControlnet( num_inference_steps=args.steps, height=args.height,