Refactor: Controlnet transformer & Flux controlnet

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)
+        weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer()
-        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
+        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,
+        self,
-            seed: int,
+        seed: int,
-            prompt: str,
+        prompt: str,
-            output: str,
+        output: str,
-            controlnet_image_path: str,
+        controlnet_image_path: str,
-            controlnet_save_canny: bool = False,
+        controlnet_save_canny: bool = False,
-            config: ConfigControlnet = ConfigControlnet(),
+        config: ConfigControlnet = ConfigControlnet(),
-            stepwise_output_dir: Path = None,
+        stepwise_output_dir: Path = None,
-    ) -> "GeneratedImage":  # fmt: off
+    ) -> 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
+        # 0. Embed the controlnet reference image
-        control_image = ImageUtil.load_image(controlnet_image_path)
+        controlnet_condition = self._embed_image(config, controlnet_image_path, controlnet_save_canny, output)
-        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)
 
         # 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_condition=controlnet_condition,
-                    controlnet_cond=controlnet_cond,
-                    config=config,
                 )
 
-                # 3.t Predict the noise
+                # 4.t Predict the noise
-                noise = self.transformer.predict(
+                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")

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 (
+from mflux.models.transformer.joint_transformer_block import JointTransformerBlock
-    JointTransformerBlock,
+from mflux.models.transformer.single_transformer_block import SingleTransformerBlock
-)
-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_transformer_blocks: int = 5,
-        num_single_blocks: int,
+        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.transformer_blocks = [JointTransformerBlock(i) for i in range(num_transformer_blocks)]
-        self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(num_single_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(nn.init.constant(0))
-        self.controlnet_x_embedder = nn.Linear(64, 3072).apply(zero_init)
+        self.controlnet_blocks = [nn.Linear(3072, 3072).apply(nn.init.constant(0)) for _ in range(num_transformer_blocks)]  # fmt: off
-        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_transformer_blocks)]
-
-        self.controlnet_single_blocks = [nn.Linear(3072, 3072) for _ in range(num_single_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_condition: mx.array,
-        controlnet_cond: mx.array,
-        config: RuntimeConfig,
     ) -> (list[mx.array], list[mx.array]):
-        time_step = config.sigmas[t] * config.num_train_steps
+        # 1. Create embeddings
-        time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision)
+        hidden_states = self.x_embedder(hidden_states) + self.controlnet_x_embedder(controlnet_condition)
-        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)
         encoder_hidden_states = self.context_embedder(prompt_embeds)
-        txt_ids = Transformer.prepare_text_ids(seq_len=prompt_embeds.shape[1])
+        text_embeddings = Transformer.compute_text_embeddings(t, pooled_prompt_embeds, self.time_text_embed, config)
-        img_ids = Transformer.prepare_latent_image_ids(height=config.height, width=config.width)
+        image_rotary_embeddings = Transformer.compute_rotary_embeddings(prompt_embeds, self.pos_embed, config)
-        ids = mx.concatenate((txt_ids, img_ids), axis=1)
-        image_rotary_emb = self.pos_embed(ids)
 
-        block_samples = ()
+        # 2. Run the joint transformer blocks
-        for block in self.transformer_blocks:
+        controlnet_block_samples = []
-            encoder_hidden_states, hidden_states = block(
+        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
+        # 4. Run the single transformer blocks
-        controlnet_block_samples = ()
+        controlnet_single_block_samples = []
-        for block_sample, controlnet_block in zip(block_samples, self.controlnet_blocks):
+        for idx, block in enumerate(self.single_transformer_blocks):
-            block_sample = controlnet_block(block_sample)
+            hidden_states = self._apply_single_transformer_block(
-            controlnet_block_samples = controlnet_block_samples + (block_sample,)
+                idx=idx,
-
+                block=block,
-        single_block_samples = ()
+                config=config,
-        for block in self.single_transformer_blocks:
-            hidden_states = block(
                 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

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":
+    def load_controlnet_transformer() -> "WeightHandlerControlnet":
-        controlnet_path = Path(snapshot_download(repo_id=controlnet_id, allow_patterns=["*.safetensors", "config.json"]))  # fmt:off
+        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"]