Move new controlnet module to its own file

src/mflux/controlnet/flux_controlnet.py

@@ -1,6 +1,5 @@
 import logging
 from pathlib import Path
-from typing import Tuple
 
 import PIL.Image
 import mlx.core as mx
@@ -12,12 +11,9 @@ from mflux.config.config import ConfigControlnet
 from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.controlnet.controlnet_util import ControlnetUtil
+from mflux.controlnet.transformer_controlnet import TransformerControlnet
 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.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
 from mflux.models.transformer.transformer import Transformer
 from mflux.models.vae.vae import VAE
 from mflux.post_processing.image import GeneratedImage
@@ -228,88 +224,3 @@ class Flux1Controlnet:
         _save_weights(self.transformer_controlnet, "transformer_controlnet")
         
 
-ControlNetOutput = Tuple[list[mx.array], list[mx.array]]
-
-
-class TransformerControlnet(nn.Module):
-
-    def __init__(
-            self,
-            model_config: ModelConfig,
-            num_blocks: int,
-            num_single_blocks: int,
-    ):
-        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)]
-
-        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)]
-
-    def forward(
-            self,
-            t: int,
-            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_strength
-
-        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)
-        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.forward(ids)
-
-        block_samples = ()
-        for block in 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
-            )
-            block_samples = block_samples + (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:
-            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, controlnet_single_block_samples

src/mflux/controlnet/transformer_controlnet.py

@@ -0,0 +1,96 @@
+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.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
+from mflux.models.transformer.transformer import Transformer
+
+
+class TransformerControlnet(nn.Module):
+
+    def __init__(
+            self,
+            model_config: ModelConfig,
+            num_blocks: int,
+            num_single_blocks: int,
+    ):
+        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)]
+
+        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)]
+
+    def forward(
+            self,
+            t: int,
+            prompt_embeds: mx.array,
+            pooled_prompt_embeds: mx.array,
+            hidden_states: mx.array,
+            controlnet_cond: mx.array,
+            config: RuntimeConfig,
+    ) -> (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.forward(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.forward(ids)
+
+        block_samples = ()
+        for block in 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
+            )
+            block_samples = block_samples + (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:
+            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, controlnet_single_block_samples