Light refactor of controlnet class

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

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:
+        if quantize is not None or controlnet_quantization_level is not None:
-            self.bits = ctrlnet_quantization_level if ctrlnet_quantization_level is not None else quantize
+            self.bits = controlnet_quantization_level if controlnet_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: 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 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:
+        # Embedd the controlnet reference image
-            log.warning(
+        control_image = ImageUtil.load_image(control_image_path)
-                f"Control image has different dimensions than the model. Resizing to {config.width}x{config.height}"
+        control_image = ControlnetUtil.scale_image(config.height, config.width, control_image)
-            )
+        control_image = ControlnetUtil.preprocess_canny(control_image)
-            control_image = control_image.resize((config.width, config.height), PIL.Image.LANCZOS)
+        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_block_samples=controlnet_block_samples,
-                controlnet_single_block_samples=ctrlnet_single_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

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,