Refactor: Transformer & Flux

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/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
+        # 1. Create embeddings
-        time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision)
         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])
+        text_embeddings = Transformer.compute_text_embeddings(t, pooled_prompt_embeds, self.time_text_embed, config)
-        img_ids = Transformer.prepare_latent_image_ids(config.height, 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)
 
+        # 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 hidden_states
-        return noise
+
+    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)