controlnet - first working version

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,

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:

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
+
+        # 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

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)

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:

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)

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)

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.

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