Merge pull request #48 from Xuzzo/feature/add_controlnet

Add controlnet

.gitignore

@@ -13,3 +13,5 @@
 *.pyc
 *.safetensors
 *.json
+
+*.egg-info

pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
     "huggingface-hub>=0.24.5",
     "safetensors>=0.4.4",
     "piexif>=1.1.3",
+    "opencv-python>=4.10.0",
 ]
 
 [project.urls]
@@ -33,6 +34,7 @@ homepage = "https://github.com/filipstrand/mflux"
 [project.scripts]
 mflux-generate = "mflux.generate:main"
 mflux-save = "mflux.save:main"
+mflux-generate-controlnet = "mflux.generate_controlnet:main"
 
 [tool.setuptools.packages.find]
 where = ["src"]

requirements.txt

@@ -7,4 +7,5 @@ torch>=2.3.1
 tqdm>=4.66.5
 huggingface-hub>=0.24.5
 safetensors>=0.4.4
-piexif>=1.1.3
+piexif>=1.1.3
+opencv-python>=4.10.0

src/mflux/config/config.py

@@ -21,3 +21,16 @@ class Config:
         self.height = 16 * (height // 16)
         self.num_inference_steps = num_inference_steps
         self.guidance = guidance
+
+
+class ConfigControlnet(Config):
+    def __init__(
+            self,
+            num_inference_steps: int = 4,
+            width: int = 1024,
+            height: int = 1024,
+            guidance: float = 4.0,
+            controlnet_strength: float = 1.0,
+    ):
+        super().__init__(num_inference_steps, width, height, guidance)
+        self.controlnet_strength = controlnet_strength

src/mflux/config/runtime_config.py

@@ -1,13 +1,13 @@
 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 +35,13 @@ class RuntimeConfig:
     @property
     def num_train_steps(self) -> int:
         return self.model_config.num_train_steps
+    
+    @property
+    def controlnet_strength(self) -> float:
+        if isinstance(self.config, ConfigControlnet):
+            return self.config.controlnet_strength
+        else:
+            raise NotImplementedError("Controlnet conditioning scale is only available for ConfigControlnet")
 
     @staticmethod
     def _create_sigmas(config, model) -> mx.array:

src/mflux/controlnet/controlnet_util.py

@@ -0,0 +1,14 @@
+import cv2
+import numpy as np
+import PIL
+
+
+class ControlnetUtil:
+
+    @staticmethod
+    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/controlnet/flux_controlnet.py

@@ -0,0 +1,189 @@
+import logging
+
+import PIL.Image
+import mlx.core as mx
+from mlx import nn
+from tqdm import tqdm
+
+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.transformer import Transformer
+from mflux.models.vae.vae import VAE
+from mflux.post_processing.generated_image import GeneratedImage
+from mflux.post_processing.image_util import ImageUtil
+from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
+from mflux.tokenizer.t5_tokenizer import TokenizerT5
+from mflux.tokenizer.tokenizer_handler import TokenizerHandler
+from mflux.weights.model_saver import ModelSaver
+from mflux.weights.weight_handler import WeightHandler
+
+log = logging.getLogger(__name__)
+
+CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny"
+
+
+class Flux1Controlnet:
+    def __init__(
+            self,
+            model_config: ModelConfig,
+            quantize: int | None = None,
+            local_path: str | None = None,
+            lora_paths: list[str] | None = None,
+            lora_scales: list[float] | None = None,
+            controlnet_path: str | None = None,
+    ):
+        self.lora_paths = lora_paths
+        self.lora_scales = lora_scales
+        self.model_config = model_config
+
+        # Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface
+        tokenizers = TokenizerHandler(model_config.model_name, self.model_config.max_sequence_length, local_path)
+        self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=self.model_config.max_sequence_length)
+        self.clip_tokenizer = TokenizerCLIP(tokenizers.clip)
+
+        # Initialize the models
+        self.vae = VAE()
+        self.transformer = Transformer(model_config)
+        self.t5_text_encoder = T5Encoder()
+        self.clip_text_encoder = CLIPEncoder()
+
+        # Load the weights from disk, huggingface cache, or download from huggingface
+        weights = WeightHandler(
+            repo_id=model_config.model_name,
+            local_path=local_path,
+            lora_paths=lora_paths,
+            lora_scales=lora_scales
+        )
+
+        # Set the loaded weights if they are not quantized
+        if weights.quantization_level is None:
+            self._set_model_weights(weights)
+
+        # Optionally quantize the model here at initialization (also required if about to load quantized weights)
+        self.bits = None
+        if quantize is not None or weights.quantization_level is not None:
+            self.bits = weights.quantization_level if weights.quantization_level is not None else quantize
+            nn.quantize(self.vae, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=self.bits)
+            nn.quantize(self.transformer, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 64, group_size=64, bits=self.bits)
+            nn.quantize(self.t5_text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=self.bits)
+            nn.quantize(self.clip_text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=self.bits)
+
+        # If loading previously saved quantized weights, the weights must be set after modules have been quantized
+        if weights.quantization_level is not None:
+            self._set_model_weights(weights)
+
+        weights_controlnet, ctrlnet_quantization_level, controlnet_config = WeightHandler.load_controlnet_transformer(controlnet_id=CONTROLNET_ID)
+        self.transformer_controlnet = TransformerControlnet(
+            model_config=model_config,
+            num_blocks=controlnet_config["num_layers"],
+            num_single_blocks=controlnet_config["num_single_layers"],
+        )
+
+        if ctrlnet_quantization_level is None:
+            self.transformer_controlnet.update(weights_controlnet)
+
+        self.bits = None
+        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 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:
+        # 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))
+
+        if config.height != control_image.height or config.width != control_image.width:
+            log.warning(f"Control image has different dimensions than the model. Resizing to {config.width}x{config.height}")
+            control_image = control_image.resize((config.width, config.height), PIL.Image.LANCZOS)
+
+        # 1. Create the initial latents
+        latents = mx.random.normal(
+            shape=[1, (config.height // 16) * (config.width // 16), 64],
+            key=mx.random.key(seed)
+        )
+        control_image = ControlnetUtil.preprocess_canny(control_image)
+        controlnet_cond = ImageUtil.to_array(control_image)
+        controlnet_cong = self.vae.encode(controlnet_cond)
+        # the rescaling in the next line is not in the huggingface code, but without it the images from 
+        # the chosen controlnet model are very bad
+        controlnet_cond = (controlnet_cong / self.vae.scaling_factor) + self.vae.shift_factor
+        controlnet_cond = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width)
+
+        # 2. Embedd the prompt
+        t5_tokens = self.t5_tokenizer.tokenize(prompt)
+        clip_tokens = self.clip_tokenizer.tokenize(prompt)
+        prompt_embeds = self.t5_text_encoder.forward(t5_tokens)
+        pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens)
+
+        for t in time_steps:
+            ctrlnet_block_samples, ctrlnet_single_block_samples = self.transformer_controlnet.forward(
+                t=t,
+                prompt_embeds=prompt_embeds,
+                pooled_prompt_embeds=pooled_prompt_embeds,
+                hidden_states=latents,
+                controlnet_cond=controlnet_cond,
+                config=config,
+            )
+            # 3.t Predict the noise
+            noise = self.transformer.predict(
+                t=t,
+                prompt_embeds=prompt_embeds,
+                pooled_prompt_embeds=pooled_prompt_embeds,
+                hidden_states=latents,
+                config=config,
+                controlnet_block_samples=ctrlnet_block_samples,
+                controlnet_single_block_samples=ctrlnet_single_block_samples,
+            )
+
+            # 4.t Take one denoise step
+            dt = config.sigmas[t + 1] - config.sigmas[t]
+            latents += noise * dt
+
+            # Evaluate to enable progress tracking
+            mx.eval(latents)
+
+        # 5. Decode the latent array and return the image
+        latents = Flux1Controlnet._unpack_latents(latents, config.height, config.width)
+        decoded = self.vae.decode(latents)
+        return ImageUtil.to_image(
+            decoded_latents=decoded,
+            seed=seed,
+            prompt=prompt,
+            quantization=self.bits,
+            generation_time=time_steps.format_dict['elapsed'],
+            lora_paths=self.lora_paths,
+            lora_scales=self.lora_scales,
+            config=config,
+        )
+
+    @staticmethod
+    def _unpack_latents(latents: mx.array, height: int, width: int) -> mx.array:
+        latents = mx.reshape(latents, (1, height // 16, width // 16, 16, 2, 2))
+        latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5))
+        latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2))
+        return latents
+    
+    @staticmethod
+    def _pack_latents(latents: mx.array, height: int, width: int) -> mx.array:
+        latents = mx.reshape(latents, (1, 16, height // 16, 2, width // 16, 2))
+        latents = mx.transpose(latents, (0, 2, 4, 1, 3, 5))
+        latents = mx.reshape(latents, (1, (width // 16) * (height // 16), 64))
+        return latents
+
+    def _set_model_weights(self, weights):
+        self.vae.update(weights.vae)
+        self.transformer.update(weights.transformer)
+        self.t5_text_encoder.update(weights.t5_encoder)
+        self.clip_text_encoder.update(weights.clip_encoder)
+
+    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

@@ -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

src/mflux/flux/flux.py

@@ -1,8 +1,5 @@
-from pathlib import Path
-
 import mlx.core as mx
 from mlx import nn
-from mlx.utils import tree_flatten
 from tqdm import tqdm
 
 from mflux.config.config import Config
@@ -12,11 +9,12 @@ 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
 from mflux.models.vae.vae import VAE
-from mflux.post_processing.image import Image
+from mflux.post_processing.generated_image import GeneratedImage
 from mflux.post_processing.image_util import ImageUtil
 from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
 from mflux.tokenizer.t5_tokenizer import TokenizerT5
 from mflux.tokenizer.tokenizer_handler import TokenizerHandler
+from mflux.weights.model_saver import ModelSaver
 from mflux.weights.weight_handler import WeightHandler
 
 
@@ -70,7 +68,7 @@ class Flux1:
         if weights.quantization_level is not None:
             self._set_model_weights(weights)
 
-    def generate_image(self, seed: int, prompt: str, config: Config = Config()) -> Image:
+    def generate_image(self, seed: int, prompt: str, config: Config = Config()) -> GeneratedImage:
         # 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))
@@ -138,39 +136,5 @@ class Flux1:
         self.t5_text_encoder.update(weights.t5_encoder)
         self.clip_text_encoder.update(weights.clip_encoder)
 
-    def save_model(self, base_path: str):
-        def _save_tokenizer(tokenizer, subdir: str):
-            path = Path(base_path) / subdir
-            path.mkdir(parents=True, exist_ok=True)
-            tokenizer.save_pretrained(path)
-
-        def _save_weights(model, subdir: str):
-            path = Path(base_path) / subdir
-            path.mkdir(parents=True, exist_ok=True)
-            weights = _split_weights(dict(tree_flatten(model.parameters())))
-            for i, weight in enumerate(weights):
-                mx.save_safetensors(str(path / f"{i}.safetensors"), weight, {"quantization_level": str(self.bits)})
-
-        def _split_weights(weights: dict, max_file_size_gb: int = 2) -> list:
-            # Copied from mlx-examples repo
-            max_file_size_bytes = max_file_size_gb << 30
-            shards = []
-            shard, shard_size = {}, 0
-            for k, v in weights.items():
-                if shard_size + v.nbytes > max_file_size_bytes:
-                    shards.append(shard)
-                    shard, shard_size = {}, 0
-                shard[k] = v
-                shard_size += v.nbytes
-            shards.append(shard)
-            return shards
-
-        # Save the tokenizers
-        _save_tokenizer(self.clip_tokenizer.tokenizer, "tokenizer")
-        _save_tokenizer(self.t5_tokenizer.tokenizer, "tokenizer_2")
-
-        # Save the models
-        _save_weights(self.vae, "vae")
-        _save_weights(self.transformer, "transformer")
-        _save_weights(self.clip_text_encoder, "text_encoder")
-        _save_weights(self.t5_text_encoder, "text_encoder_2")
+    def save_model(self, base_path: str) -> None:
+        ModelSaver.save_model(self, self.bits, base_path)

src/mflux/generate.py

@@ -18,7 +18,7 @@ def main():
     parser.add_argument('--seed', type=int, default=None, help='Entropy Seed (Default is time-based random-seed)')
     parser.add_argument('--height', type=int, default=1024, help='Image height (Default is 1024)')
     parser.add_argument('--width', type=int, default=1024, help='Image width (Default is 1024)')
-    parser.add_argument('--steps', type=int, default=4, help='Inference Steps')
+    parser.add_argument('--steps', type=int, default=None, help='Inference Steps')
     parser.add_argument('--guidance', type=float, default=3.5, help='Guidance Scale (Default is 3.5)')
     parser.add_argument('--quantize',  "-q", type=int, choices=[4, 8], default=None, help='Quantize the model (4 or 8, Default is None)')
     parser.add_argument('--path', type=str, default=None, help='Local path for loading a model from disk')
@@ -31,6 +31,9 @@ def main():
     if args.path and args.model is None:
         parser.error("--model must be specified when using --path")
 
+    if args.steps is None:
+        args.steps = 4 if args.model == "schnell" else 14
+
     # Load the model
     flux = Flux1(
         model_config=ModelConfig.from_alias(args.model),

src/mflux/generate_controlnet.py

@@ -0,0 +1,68 @@
+import argparse
+import os
+import sys
+import time
+
+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 ConfigControlnet
+from mflux.controlnet.flux_controlnet import Flux1Controlnet
+from mflux.post_processing.image_util import ImageUtil
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Generate an image based on a prompt.')
+    parser.add_argument('--prompt', type=str, required=True, help='The textual description of the image to generate.')
+    parser.add_argument('--control-image-path', type=str, required=True, help='Local path of the image to use as input for controlnet.')
+    parser.add_argument('--output', type=str, default="image.png", help='The filename for the output image. Default is "image.png".')
+    parser.add_argument('--model', "-m", type=str, required=True, choices=["dev", "schnell"], help='The model to use ("schnell" or "dev").')
+    parser.add_argument('--seed', type=int, default=None, help='Entropy Seed (Default is time-based random-seed)')
+    parser.add_argument('--height', type=int, default=1024, help='Image height (Default is 1024)')
+    parser.add_argument('--width', type=int, default=1024, help='Image width (Default is 1024)')
+    parser.add_argument('--steps', type=int, default=None, help='Inference Steps')
+    parser.add_argument('--guidance', type=float, default=3.5, help='Guidance Scale (Default is 3.5)')
+    parser.add_argument('--controlnet-strength', type=float, default=0.7, help='Controls how strongly the control image influences the output image. A value of 0.0 means no influence. (Default is 0.7)')
+    parser.add_argument('--quantize',  "-q", type=int, choices=[4, 8], default=None, help='Quantize the model (4 or 8, Default is None)')
+    parser.add_argument('--path', type=str, default=None, help='Local path for loading a model from disk')
+    parser.add_argument('--lora-paths', type=str, nargs='*', default=None, help='Local safetensors for applying LORA from disk')
+    parser.add_argument('--lora-scales', type=float, nargs='*', default=None, help='Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.')
+    parser.add_argument('--metadata', action='store_true', help='Export image metadata as a JSON file.')
+
+    args = parser.parse_args()
+
+    if args.path and args.model is None:
+        parser.error("--model must be specified when using --path")
+
+    if args.steps is None:
+        args.steps = 4 if args.model == "schnell" else 14
+
+    # Load the model
+    flux = Flux1Controlnet(
+        model_config=ModelConfig.from_alias(args.model),
+        quantize=args.quantize,
+        local_path=args.path,
+        lora_paths=args.lora_paths,
+        lora_scales=args.lora_scales
+    )
+
+    # Generate an image
+    image = flux.generate_image(
+        seed=int(time.time()) if args.seed is None else args.seed,
+        prompt=args.prompt,
+        control_image=ImageUtil.load_image(args.control_image_path),
+        config=ConfigControlnet(
+            num_inference_steps=args.steps,
+            height=args.height,
+            width=args.width,
+            guidance=args.guidance,
+            controlnet_strength=args.controlnet_strength
+        )
+    )
+
+    # Save the image
+    image.save(path=args.output, export_json_metadata=args.metadata)
+
+
+if __name__ == '__main__':
+    main()

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,3 +1,5 @@
+import math
+
 import mlx.core as mx
 from mlx import nn
 
@@ -30,6 +32,8 @@ class Transformer(nn.Module):
             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)
@@ -37,27 +41,38 @@ class Transformer(nn.Module):
         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)
+        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)
 
-        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 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]
 
         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 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]
+                )
 
         hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...]
         hidden_states = self.norm_out.forward(hidden_states, text_embeddings)
@@ -66,7 +81,7 @@ class Transformer(nn.Module):
         return noise
 
     @staticmethod
-    def _prepare_latent_image_ids(height: int, width: int) -> mx.array:
+    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))
@@ -77,5 +92,5 @@ 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))

src/mflux/post_processing/generated_image.py

@@ -11,7 +11,7 @@ from mflux.config.model_config import ModelConfig
 log = logging.getLogger(__name__)
 
 
-class Image:
+class GeneratedImage:
 
     def __init__(
             self,
@@ -26,6 +26,7 @@ class Image:
             generation_time: float,
             lora_paths: list[str],
             lora_scales: list[float],
+            controlnet_strength: float | None = None,
     ):
         self.image = image
         self.model_config = model_config
@@ -38,6 +39,7 @@ class Image:
         self.generation_time = generation_time
         self.lora_paths = lora_paths
         self.lora_scales = lora_scales
+        self.controlnet_strength = controlnet_strength
 
     def save(self, path: str, export_json_metadata: bool = False) -> None:
         file_path = Path(path)
@@ -123,4 +125,5 @@ class Image:
             'lora_paths': ', '.join(self.lora_paths) if self.lora_paths else '',
             'lora_scales': ', '.join([f"{scale:.2f}" for scale in self.lora_scales]) if self.lora_scales else '',
             'prompt': self.prompt,
+            'controlnet_strength': "None" if self.controlnet_strength is None else f"{self.controlnet_strength:.2f}",
         }

src/mflux/post_processing/image_util.py

@@ -3,8 +3,9 @@ from PIL import Image
 import mlx.core as mx
 import numpy as np
 
+from mflux.config.config import ConfigControlnet
 from mflux.config.runtime_config import RuntimeConfig
-from mflux.post_processing.image import Image
+from mflux.post_processing.generated_image import GeneratedImage
 
 
 class ImageUtil:
@@ -19,11 +20,11 @@ class ImageUtil:
             lora_paths: list[str],
             lora_scales: list[float],
             config: RuntimeConfig,
-    ) -> Image:
+    ) -> GeneratedImage:
         normalized = ImageUtil._denormalize(decoded_latents)
         normalized_numpy = ImageUtil._to_numpy(normalized)
         image = ImageUtil._numpy_to_pil(normalized_numpy)
-        return Image(
+        return GeneratedImage(
             image=image,
             model_config=config.model_config,
             seed=seed,
@@ -35,6 +36,7 @@ class ImageUtil:
             generation_time=generation_time,
             lora_paths=lora_paths,
             lora_scales=lora_scales,
+            controlnet_strength=config.controlnet_strength if isinstance(config.config, ConfigControlnet) else None,
         )
 
     @staticmethod
@@ -66,14 +68,12 @@ class ImageUtil:
 
     @staticmethod
     def to_array(image: PIL.Image.Image) -> mx.array:
-        image = ImageUtil._resize(image)
         image = ImageUtil._pil_to_numpy(image)
         array = mx.array(image)
         array = mx.transpose(array, (0, 3, 1, 2))
         array = ImageUtil._normalize(array)
         return array
-
+    
     @staticmethod
-    def _resize(image):
-        image = image.resize((1024, 1024), resample=PIL.Image.LANCZOS)
-        return image
+    def load_image(path: str) -> Image.Image:
+        return Image.open(path)

src/mflux/weights/model_saver.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+
+import mlx.core as mx
+from mlx import nn
+from mlx.utils import tree_flatten
+from transformers import CLIPTokenizer, T5Tokenizer
+
+
+class ModelSaver:
+
+    @staticmethod
+    def save_model(model, bits: int, base_path: str):
+        # Save the tokenizers
+        ModelSaver._save_tokenizer(base_path, model.clip_tokenizer.tokenizer, "tokenizer")
+        ModelSaver._save_tokenizer(base_path, model.t5_tokenizer.tokenizer, "tokenizer_2")
+
+        # Save the models
+        ModelSaver.save_weights(base_path, bits, model.vae, "vae")
+        ModelSaver.save_weights(base_path, bits, model.transformer, "transformer")
+        ModelSaver.save_weights(base_path, bits, model.clip_text_encoder, "text_encoder")
+        ModelSaver.save_weights(base_path, bits, model.t5_text_encoder, "text_encoder_2")
+
+    @staticmethod
+    def _save_tokenizer(base_path: str, tokenizer: CLIPTokenizer | T5Tokenizer, subdir: str):
+        path = Path(base_path) / subdir
+        path.mkdir(parents=True, exist_ok=True)
+        tokenizer.save_pretrained(path)
+
+    @staticmethod
+    def save_weights(base_path: str, bits: int, model: nn.Module, subdir: str):
+        path = Path(base_path) / subdir
+        path.mkdir(parents=True, exist_ok=True)
+        weights = ModelSaver._split_weights(base_path, dict(tree_flatten(model.parameters())))
+        for i, weight in enumerate(weights):
+            mx.save_safetensors(str(path / f"{i}.safetensors"), weight, {"quantization_level": str(bits)})
+
+    @staticmethod
+    def _split_weights(base_path: str, weights: dict, max_file_size_gb: int = 2) -> list:
+        # Copied from mlx-examples repo
+        max_file_size_bytes = max_file_size_gb << 30
+        shards = []
+        shard, shard_size = {}, 0
+        for k, v in weights.items():
+            if shard_size + v.nbytes > max_file_size_bytes:
+                shards.append(shard)
+                shard, shard_size = {}, 0
+            shard[k] = v
+            shard_size += v.nbytes
+        shards.append(shard)
+        return shards

src/mflux/weights/weight_handler.py

@@ -1,3 +1,4 @@
+import json
 from pathlib import Path
 
 import mlx.core as mx
@@ -87,6 +88,39 @@ class WeightHandler:
                         "linear2": block["ff_context"]["net"][2]
                     }
         return weights, quantization_level
+    
+    @staticmethod
+    def load_controlnet_transformer(controlnet_id: str) -> (dict, int):
+        controlnet_path = Path(snapshot_download(repo_id=controlnet_id,allow_patterns=["*.safetensors","config.json"]))
+        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())
+
+        if quantization_level is not None:
+            return tree_unflatten(weights), quantization_level
+
+        weights = [WeightUtil.reshape_weights(k, v) for k, v in weights]
+        weights = WeightUtil.flatten(weights)
+        weights = tree_unflatten(weights)
+
+        # 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]
+                    }
+        config = json.load(open(controlnet_path / "config.json"))
+        return weights, quantization_level, config
 
     @staticmethod
     def load_vae(root_path: Path) -> (dict, int):