Merge pull request #122 from filipstrand/transformer-refactor

Transformer refactor

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)
-        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
+        weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer()
+        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,
-            seed: int,
-            prompt: str,
-            output: str,
-            controlnet_image_path: str,
-            controlnet_save_canny: bool = False,
-            config: ConfigControlnet = ConfigControlnet(),
-            stepwise_output_dir: Path = None,
-    ) -> "GeneratedImage":  # fmt: off
+        self,
+        seed: int,
+        prompt: str,
+        output: str,
+        controlnet_image_path: str,
+        controlnet_save_canny: bool = False,
+        config: ConfigControlnet = ConfigControlnet(),
+        stepwise_output_dir: Path = None,
+    ) -> 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
-        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)
+        # 0. Embed the controlnet reference image
+        controlnet_condition = self._embed_image(config, controlnet_image_path, controlnet_save_canny, output)
 
         # 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_cond=controlnet_cond,
-                    config=config,
+                    controlnet_condition=controlnet_condition,
                 )
 
-                # 3.t Predict the noise
-                noise = self.transformer.predict(
+                # 4.t Predict the noise
+                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 (
-    JointTransformerBlock,
-)
-from mflux.models.transformer.single_transformer_block import (
-    SingleTransformerBlock,
-)
+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
 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_single_blocks: int,
+        num_transformer_blocks: int = 5,
+        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.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(num_single_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_transformer_blocks)]
 
-        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(num_blocks)]
-
-        self.controlnet_single_blocks = [nn.Linear(3072, 3072) for _ in range(num_single_blocks)]
+        self.controlnet_x_embedder = nn.Linear(64, 3072).apply(nn.init.constant(0))
+        self.controlnet_blocks = [nn.Linear(3072, 3072).apply(nn.init.constant(0)) for _ in range(num_transformer_blocks)]  # fmt: off
+        self.controlnet_single_blocks = [nn.Linear(3072, 3072) for _ in range(num_single_transformer_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_cond: mx.array,
-        config: RuntimeConfig,
+        controlnet_condition: mx.array,
     ) -> (list[mx.array], list[mx.array]):
-        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_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)
+        # 1. Create embeddings
+        hidden_states = self.x_embedder(hidden_states) + self.controlnet_x_embedder(controlnet_condition)
         encoder_hidden_states = self.context_embedder(prompt_embeds)
-        txt_ids = Transformer.prepare_text_ids(seq_len=prompt_embeds.shape[1])
-        img_ids = Transformer.prepare_latent_image_ids(height=config.height, width=config.width)
-        ids = mx.concatenate((txt_ids, img_ids), axis=1)
-        image_rotary_emb = self.pos_embed(ids)
+        text_embeddings = Transformer.compute_text_embeddings(t, pooled_prompt_embeds, self.time_text_embed, config)
+        image_rotary_embeddings = Transformer.compute_rotary_embeddings(prompt_embeds, self.pos_embed, config)
 
-        block_samples = ()
-        for block in self.transformer_blocks:
-            encoder_hidden_states, hidden_states = block(
+        # 2. Run the joint transformer blocks
+        controlnet_block_samples = []
+        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
-        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,)
-
-        single_block_samples = ()
-        for block in self.single_transformer_blocks:
-            hidden_states = block(
+        # 4. Run the single transformer blocks
+        controlnet_single_block_samples = []
+        for idx, block in enumerate(self.single_transformer_blocks):
+            hidden_states = self._apply_single_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_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":
-        controlnet_path = Path(snapshot_download(repo_id=controlnet_id, allow_patterns=["*.safetensors", "config.json"]))  # fmt:off
+    def load_controlnet_transformer() -> "WeightHandlerControlnet":
+        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"]

src/mflux/dreambooth/optimization/dreambooth_loss.py

@@ -53,12 +53,12 @@ class DreamBoothLoss:
         )  # fmt: off
 
         # Predict the noise from timestep t
-        predicted_noise = flux.transformer.predict(
+        predicted_noise = flux.transformer(
             t=t,
+            config=config,
+            hidden_states=latents_t,
             prompt_embeds=example.prompt_embeds,
             pooled_prompt_embeds=example.pooled_prompt_embeds,
-            hidden_states=latents_t,
-            config=config,
         )
 
         # Construct the loss (derivation in src/mflux/dreambooth/optimization/_loss_derivation)

src/mflux/flux/flux.py

@@ -51,7 +51,7 @@ class Flux1(nn.Module):
 
         # 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()
 
@@ -100,12 +100,12 @@ class Flux1(nn.Module):
         for gen_step, t in enumerate(time_steps, 1):
             try:
                 # 3.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,
                 )
 
                 # 4.t Take one denoise step

src/mflux/models/transformer/ada_layer_norm_zero.py

@@ -8,7 +8,7 @@ class AdaLayerNormZero(nn.Module):
         self.linear = nn.Linear(3072, 18432)
         self.norm = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
 
-    def __call__(self, x: mx.array, text_embeddings: mx.array):
+    def __call__(self, hidden_states: mx.array, text_embeddings: mx.array):
         text_embeddings = self.linear(nn.silu(text_embeddings))
         chunk_size = 18432 // 6
         shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
@@ -17,5 +17,5 @@ class AdaLayerNormZero(nn.Module):
         shift_mlp = text_embeddings[:, 3 * chunk_size : 4 * chunk_size]
         scale_mlp = text_embeddings[:, 4 * chunk_size : 5 * chunk_size]
         gate_mlp = text_embeddings[:, 5 * chunk_size : 6 * chunk_size]
-        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
-        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+        hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp

src/mflux/models/transformer/ada_layer_norm_zero_single.py

@@ -8,11 +8,11 @@ class AdaLayerNormZeroSingle(nn.Module):
         self.linear = nn.Linear(3072, 3 * 3072)
         self.norm = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
 
-    def __call__(self, x: mx.array, text_embeddings: mx.array):
+    def __call__(self, hidden_states: mx.array, text_embeddings: mx.array) -> mx.array:
         text_embeddings = self.linear(nn.silu(text_embeddings))
         chunk_size = 9216 // 3
         shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
         scale_msa = text_embeddings[:, 1 * chunk_size : 2 * chunk_size]
         gate_msa = text_embeddings[:, 2 * chunk_size : 3 * chunk_size]
-        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
-        return x, gate_msa
+        hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return hidden_states, gate_msa

src/mflux/models/transformer/joint_attention.py

@@ -30,6 +30,7 @@ class JointAttention(nn.Module):
         encoder_hidden_states: mx.array,
         image_rotary_emb: mx.array,
     ) -> (mx.array, mx.array):
+        # 1a. Compute Q,K,V for hidden_states
         query, key, value = AttentionUtils.process_qkv(
             hidden_states=hidden_states,
             to_q=self.to_q,
@@ -40,6 +41,8 @@ class JointAttention(nn.Module):
             num_heads=self.num_heads,
             head_dim=self.head_dimension,
         )
+
+        # 1b. Compute Q,K,V for encoder_hidden_states
         enc_query, enc_key, enc_value = AttentionUtils.process_qkv(
             hidden_states=encoder_hidden_states,
             to_q=self.add_q_proj,
@@ -51,12 +54,15 @@ class JointAttention(nn.Module):
             head_dim=self.head_dimension,
         )
 
+        # 1c. Concatenate results
         query = mx.concatenate([enc_query, query], axis=2)
         key = mx.concatenate([enc_key, key], axis=2)
         value = mx.concatenate([enc_value, value], axis=2)
 
+        # 1d. Apply rope to Q,K
         query, key = AttentionUtils.apply_rope(xq=query, xk=key, freqs_cis=image_rotary_emb)
 
+        # 2. Compute attention
         hidden_states = AttentionUtils.compute_attention(
             query=query,
             key=key,
@@ -66,12 +72,13 @@ class JointAttention(nn.Module):
             head_dim=self.head_dimension,
         )
 
+        # 3. Separate the results
         encoder_hidden_states, hidden_states = (
             hidden_states[:, : encoder_hidden_states.shape[1]],
             hidden_states[:, encoder_hidden_states.shape[1] :],
         )
 
+        # 4. Project the output
         hidden_states = self.to_out[0](hidden_states)
         encoder_hidden_states = self.to_add_out(encoder_hidden_states)
-
         return hidden_states, encoder_hidden_states

src/mflux/models/transformer/joint_transformer_block.py

@@ -11,12 +11,12 @@ class JointTransformerBlock(nn.Module):
         super().__init__()
         self.layer = layer
         self.norm1 = AdaLayerNormZero()
-        self.norm2 = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
-        self.ff = FeedForward(activation_function=nn.gelu)
-        self.attn = JointAttention()
         self.norm1_context = AdaLayerNormZero()
-        self.ff_context = FeedForward(activation_function=nn.gelu_approx)
+        self.attn = JointAttention()
+        self.norm2 = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
         self.norm2_context = nn.LayerNorm(dims=1536, eps=1e-6, affine=False)
+        self.ff = FeedForward(activation_function=nn.gelu)
+        self.ff_context = FeedForward(activation_function=nn.gelu_approx)
 
     def __call__(
         self,
@@ -25,30 +25,67 @@ class JointTransformerBlock(nn.Module):
         text_embeddings: mx.array,
         rotary_embeddings: mx.array,
     ) -> (mx.array, mx.array):
-        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, text_embeddings)
+        # 1a. Compute norm for hidden_states
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+            hidden_states=hidden_states,
+            text_embeddings=text_embeddings
+        )  # fmt: off
 
+        # 1b. Compute norm for encoder_hidden_states
         norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
-            x=encoder_hidden_states, text_embeddings=text_embeddings
-        )
+            hidden_states=encoder_hidden_states,
+            text_embeddings=text_embeddings
+        )  # fmt: off
 
+        # 2. Compute attention
         attn_output, context_attn_output = self.attn(
             hidden_states=norm_hidden_states,
             encoder_hidden_states=norm_encoder_hidden_states,
             image_rotary_emb=rotary_embeddings,
         )
 
+        # 3a. Apply norm and feed forward for hidden states
+        hidden_states = JointTransformerBlock._apply_norm_and_feed_forward(
+            hidden_states=hidden_states,
+            attn_output=attn_output,
+            gate_mlp=gate_mlp,
+            gate_msa=gate_msa,
+            scale_mlp=scale_mlp,
+            shift_mlp=shift_mlp,
+            norm_layer=self.norm2,
+            ff_layer=self.ff,
+        )
+
+        # 3b. Apply norm and feed forward for encoder hidden states
+        encoder_hidden_states = JointTransformerBlock._apply_norm_and_feed_forward(
+            hidden_states=encoder_hidden_states,
+            attn_output=context_attn_output,
+            gate_mlp=c_gate_mlp,
+            gate_msa=c_gate_msa,
+            scale_mlp=c_scale_mlp,
+            shift_mlp=c_shift_mlp,
+            norm_layer=self.norm2_context,
+            ff_layer=self.ff_context,
+        )
+
+        return encoder_hidden_states, hidden_states
+
+    @staticmethod
+    def _apply_norm_and_feed_forward(
+        hidden_states: mx.array,
+        attn_output: mx.array,
+        gate_mlp: mx.array,
+        gate_msa: mx.array,
+        scale_mlp: mx.array,
+        shift_mlp: mx.array,
+        norm_layer: nn.Module,
+        ff_layer: nn.Module,
+    ) -> mx.array:
         attn_output = mx.expand_dims(gate_msa, axis=1) * attn_output
         hidden_states = hidden_states + attn_output
-        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_layer(hidden_states)
         norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
-        ff_output = self.ff(norm_hidden_states)
+        ff_output = ff_layer(norm_hidden_states)
         ff_output = mx.expand_dims(gate_mlp, axis=1) * ff_output
         hidden_states = hidden_states + ff_output
-
-        context_attn_output = mx.expand_dims(c_gate_msa, axis=1) * context_attn_output
-        encoder_hidden_states = encoder_hidden_states + context_attn_output
-        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
-        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
-        context_ff_output = self.ff_context(norm_encoder_hidden_states)
-        encoder_hidden_states = encoder_hidden_states + mx.expand_dims(c_gate_mlp, axis=1) * context_ff_output
-        return encoder_hidden_states, hidden_states
+        return hidden_states

src/mflux/models/transformer/single_block_attention.py

@@ -18,6 +18,7 @@ class SingleBlockAttention(nn.Module):
         self.norm_k = nn.RMSNorm(128)
 
     def __call__(self, hidden_states: mx.array, image_rotary_emb: mx.array) -> mx.array:
+        # 1a. Compute Q,K,V for hidden_states
         query, key, value = AttentionUtils.process_qkv(
             hidden_states=hidden_states,
             to_q=self.to_q,
@@ -29,8 +30,10 @@ class SingleBlockAttention(nn.Module):
             head_dim=self.head_dimension,
         )
 
+        # 1b. Apply rope to Q,K
         query, key = AttentionUtils.apply_rope(xq=query, xk=key, freqs_cis=image_rotary_emb)
 
+        # 2. Compute attention
         return AttentionUtils.compute_attention(
             query=query,
             key=key,

src/mflux/models/transformer/single_transformer_block.py

@@ -12,8 +12,8 @@ class SingleTransformerBlock(nn.Module):
         super().__init__()
         self.layer = layer
         self.norm = AdaLayerNormZeroSingle()
-        self.proj_mlp = nn.Linear(3072, 4 * 3072)
         self.attn = SingleBlockAttention()
+        self.proj_mlp = nn.Linear(3072, 4 * 3072)
         self.proj_out = nn.Linear(3072 + 4 * 3072, 3072)
 
     def __call__(
@@ -22,15 +22,39 @@ class SingleTransformerBlock(nn.Module):
         text_embeddings: mx.array,
         rotary_embeddings: mx.array,
     ) -> (mx.array, mx.array):
+        # 0. Establish residual connection
         residual = hidden_states
-        norm_hidden_states, gate = self.norm(x=hidden_states, text_embeddings=text_embeddings)
-        mlp_hidden_states = nn.gelu_approx(self.proj_mlp(norm_hidden_states))
+
+        # 1. Compute norm for hidden_states
+        norm_hidden_states, gate = self.norm(
+            hidden_states=hidden_states,
+            text_embeddings=text_embeddings
+        )  # fmt: off
+
+        # 2. Compute attention
         attn_output = self.attn(
             hidden_states=norm_hidden_states,
             image_rotary_emb=rotary_embeddings,
         )
+
+        # 3. Apply norm and feed forward for hidden states
+        hidden_states = self._apply_feed_forward_and_projection(
+            norm_hidden_states=norm_hidden_states,
+            attn_output=attn_output,
+            gate=gate,
+        )
+
+        return residual + hidden_states
+
+    def _apply_feed_forward_and_projection(
+        self,
+        norm_hidden_states: mx.array,
+        attn_output: mx.array,
+        gate: mx.array,
+    ) -> mx.array:
+        feed_forward = self.proj_mlp(norm_hidden_states)
+        mlp_hidden_states = nn.gelu_approx(feed_forward)
         hidden_states = mx.concatenate([attn_output, mlp_hidden_states], axis=2)
         gate = mx.expand_dims(gate, axis=1)
         hidden_states = gate * self.proj_out(hidden_states)
-        hidden_states = residual + hidden_states
         return hidden_states

src/mflux/models/transformer/transformer.py

@@ -19,74 +19,141 @@ from mflux.models.transformer.time_text_embed import TimeTextEmbed
 
 
 class Transformer(nn.Module):
-    def __init__(self, model_config: ModelConfig, num_transformer_blocks: int):
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        num_transformer_blocks: int = 19,
+        num_single_transformer_blocks: int = 38,
+    ):
         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_transformer_blocks)]
-        self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(38)]
+        self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(num_single_transformer_blocks)]
         self.norm_out = AdaLayerNormContinuous(3072, 3072)
         self.proj_out = nn.Linear(3072, 64)
 
-    def predict(
+    def __call__(
         self,
         t: int,
+        config: RuntimeConfig,
+        hidden_states: mx.array,
         prompt_embeds: mx.array,
         pooled_prompt_embeds: mx.array,
-        hidden_states: mx.array,
-        config: RuntimeConfig,
         controlnet_block_samples: list[mx.array] | None = None,
         controlnet_single_block_samples: list[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)
+        # 1. Create embeddings
         hidden_states = self.x_embedder(hidden_states)
-        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])
-        img_ids = Transformer.prepare_latent_image_ids(config.height, config.width)
-        ids = mx.concatenate((txt_ids, img_ids), axis=1)
-        image_rotary_emb = self.pos_embed(ids)
+        text_embeddings = Transformer.compute_text_embeddings(t, pooled_prompt_embeds, self.time_text_embed, config)
+        image_rotary_embeddings = Transformer.compute_rotary_embeddings(prompt_embeds, self.pos_embed, config)
 
+        # 2. Run the joint transformer blocks
         for idx, block in enumerate(self.transformer_blocks):
-            encoder_hidden_states, hidden_states = block(
+            encoder_hidden_states, hidden_states = self._apply_joint_transformer_block(
+                idx=idx,
+                block=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_block_samples=controlnet_block_samples,
             )
-            if controlnet_block_samples is not None and len(controlnet_block_samples) > 0:
-                interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
-                interval_control = int(math.ceil(interval_control))
-                hidden_states = hidden_states + controlnet_block_samples[idx // interval_control]
 
+        # 3. Concat the hidden states
         hidden_states = mx.concatenate([encoder_hidden_states, hidden_states], axis=1)
 
+        # 4. Run the single transformer blocks
         for idx, block in enumerate(self.single_transformer_blocks):
-            hidden_states = block(
+            hidden_states = self._apply_single_transformer_block(
+                idx=idx,
+                block=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,
             )
-            if controlnet_single_block_samples is not None and len(controlnet_single_block_samples) > 0:
-                interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
-                interval_control = int(math.ceil(interval_control))
-                hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
-                    hidden_states[:, encoder_hidden_states.shape[1] :, ...]
-                    + controlnet_single_block_samples[idx // interval_control]
-                )
 
+        # 5. Project the final output
         hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
         hidden_states = self.norm_out(hidden_states, text_embeddings)
         hidden_states = self.proj_out(hidden_states)
-        noise = hidden_states
-        return noise
+        return hidden_states
+
+    def _apply_single_transformer_block(
+        self,
+        idx: int,
+        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 previously calculated controlnet result (if applicable)
+        sample = Transformer._get_controlnet_sample(idx, self.single_transformer_blocks, controlnet_single_block_samples)  # fmt: off
+        hidden_states[:, encoder_hidden_states.shape[1] :, ...] += sample if sample is not None else 0
+
+        return hidden_states
+
+    def _apply_joint_transformer_block(
+        self,
+        idx: int,
+        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],
+    ) -> 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 previously calculated controlnet result (if applicable)
+        sample = Transformer._get_controlnet_sample(idx, self.transformer_blocks, controlnet_block_samples)
+        hidden_states += sample if sample is not None else 0
+
+        return encoder_hidden_states, hidden_states
 
     @staticmethod
-    def prepare_latent_image_ids(height: int, width: int) -> mx.array:
+    def compute_rotary_embeddings(prompt_embeds: mx.array, pos_embed: EmbedND, config: RuntimeConfig) -> mx.array:
+        txt_ids = Transformer._prepare_text_ids(seq_len=prompt_embeds.shape[1])
+        img_ids = Transformer._prepare_latent_image_ids(height=config.height, width=config.width)
+        ids = mx.concatenate((txt_ids, img_ids), axis=1)
+        image_rotary_emb = pos_embed(ids)
+        return image_rotary_emb
+
+    @staticmethod
+    def compute_text_embeddings(
+        t: int,
+        pooled_prompt_embeds: mx.array,
+        time_text_embed: TimeTextEmbed,
+        config: RuntimeConfig,
+    ) -> mx.array:
+        time_step = config.sigmas[t] * config.num_train_steps
+        time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision)
+        guidance = mx.broadcast_to(config.guidance * config.num_train_steps, (1,)).astype(config.precision)
+        text_embeddings = time_text_embed(time_step, pooled_prompt_embeds, guidance)
+        return text_embeddings
+
+    @staticmethod
+    def _prepare_latent_image_ids(height: int, width: int) -> mx.array:
         latent_width = width // 16
         latent_height = height // 16
         latent_image_ids = mx.zeros((latent_height, latent_width, 3))
@@ -97,5 +164,23 @@ class Transformer(nn.Module):
         return latent_image_ids
 
     @staticmethod
-    def prepare_text_ids(seq_len: mx.array) -> mx.array:
+    def _prepare_text_ids(seq_len: mx.array) -> mx.array:
         return mx.zeros((1, seq_len, 3))
+
+    @staticmethod
+    def _get_controlnet_sample(
+        idx: int,
+        blocks: mx.array,
+        controlnet_samples: list[mx.array] | None,
+    ) -> mx.array | None:  # fmt: off
+        if controlnet_samples is None:
+            return None
+
+        if len(controlnet_samples) == 0:
+            return None
+
+        num_blocks = len(blocks)
+        num_samples = len(controlnet_samples)
+        interval_control = int(math.ceil(num_blocks / num_samples))
+        control_index = idx // interval_control
+        return controlnet_samples[control_index]

src/mflux/weights/weight_handler.py

@@ -60,6 +60,9 @@ class WeightHandler:
     def num_transformer_blocks(self) -> int:
         return len(self.transformer["transformer_blocks"])
 
+    def num_single_transformer_blocks(self) -> int:
+        return len(self.transformer["single_transformer_blocks"])
+
     @staticmethod
     def _load_clip_encoder(root_path: Path) -> (dict, int):
         weights, quantization_level, _ = WeightHandler._get_weights("text_encoder", root_path)