From 1b9c6455030f0c2eaef22e6d2e8cdf0840bcbef7 Mon Sep 17 00:00:00 2001
From: Fabio <f.xuzzo@gmail.com>
Date: Fri, 13 Sep 2024 00:28:43 +0200
Subject: [PATCH 01/16] WIP controlnet

---
 src/mflux/config/config.py                  |  13 +
 src/mflux/flux/controlnet.py                | 281 ++++++++++++++++++++
 src/mflux/flux/flux.py                      |   4 +-
 src/mflux/models/transformer/transformer.py |   2 +
 src/mflux/post_processing/image.py          |   2 +-
 src/mflux/post_processing/image_util.py     |  12 +-
 trial.py                                    |  37 +++
 7 files changed, 339 insertions(+), 12 deletions(-)
 create mode 100644 src/mflux/flux/controlnet.py
 create mode 100644 trial.py

diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py
index b892e21..7686d3d 100644
--- a/src/mflux/config/config.py
+++ b/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_conditioning_scale: float = 1.0,
+    ):
+        super().__init__(num_inference_steps, width, height, guidance)
+        self.controlnet_conditioning_scale = controlnet_conditioning_scale
diff --git a/src/mflux/flux/controlnet.py b/src/mflux/flux/controlnet.py
new file mode 100644
index 0000000..54b4a22
--- /dev/null
+++ b/src/mflux/flux/controlnet.py
@@ -0,0 +1,281 @@
+from pathlib import Path
+from typing import Tuple
+
+import PIL.Image
+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, ConfigControlnet
+from mflux.config.model_config import ModelConfig
+from mflux.config.runtime_config import RuntimeConfig
+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 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.weight_handler import WeightHandler
+
+from mflux.config.model_config import ModelConfig
+from mflux.config.runtime_config import RuntimeConfig
+from mflux.models.transformer.ada_layer_norm_continous import AdaLayerNormContinuous
+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
+
+import logging
+
+log = logging.getLogger(__name__)
+
+
+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)
+
+        self.transformer_controlnet = TransformerControlnet(model_config=model_config, local_path=controlnet_path, quantize=quantize)
+
+        weights_controlnet = WeightHandler.load_transformer(root_path=controlnet_path)
+        if weights_controlnet.quantization_level is None:
+            self.transformer_controlnet.update(weights_controlnet)
+
+        self.bits = None
+        if quantize is not None or weights.quantization_level is not None:
+            self.bits = weights_controlnet.quantization_level if weights_controlnet.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]) > 64, group_size=64, bits=self.bits)
+
+        if weights_controlnet.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_cond = ImageUtil.to_array(control_image)
+        control_cond = self.vae.encode(control_cond)
+        control_cond = (control_cond - self.vae.shift_factor) * self.vae.scaling_factor
+        control_cond = Flux1Controlnet._pack_latents(control_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:
+            controlnet_samples = self.transformer_controlnet(
+                t=t,
+                prompt_embeds=prompt_embeds,
+                pooled_prompt_embeds=pooled_prompt_embeds,
+                hidden_states=latents,
+                control_cond=control_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_samples=controlnet_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):
+        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")
+        
+
+class ControlNetOutput:
+    controlnet_block_samples: Tuple[mx.array]
+    controlnet_single_block_samples: Tuple[mx.array]
+
+class TransformerControlnet(nn.Module):
+
+    def __init__(self, model_config: ModelConfig):
+        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(19)]
+        self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(38)]
+        self.norm_out = AdaLayerNormContinuous(3072, 3072)
+        self.proj_out = nn.Linear(3072, 64)
+
+    def forward(
+            self,
+            t: int,
+            prompt_embeds: mx.array,
+            pooled_prompt_embeds: mx.array,
+            hidden_states: 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)
+        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)
+
+        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
+            )
+
+        hidden_states = mx.concatenate([encoder_hidden_states, hidden_states], axis=1)
+
+        for block in self.single_transformer_blocks:
+            hidden_states = block.forward(
+                hidden_states=hidden_states,
+                text_embeddings=text_embeddings,
+                rotary_embeddings=image_rotary_emb
+            )
+
+        hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...]
+        hidden_states = self.norm_out.forward(hidden_states, text_embeddings)
+        hidden_states = self.proj_out(hidden_states)
+        noise = hidden_states
+        return noise
\ No newline at end of file
diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py
index 96964c9..e7d52c4 100644
--- a/src/mflux/flux/flux.py
+++ b/src/mflux/flux/flux.py
@@ -12,7 +12,7 @@ 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.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
@@ -70,7 +70,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))
diff --git a/src/mflux/models/transformer/transformer.py b/src/mflux/models/transformer/transformer.py
index ee7d095..dfc4979 100644
--- a/src/mflux/models/transformer/transformer.py
+++ b/src/mflux/models/transformer/transformer.py
@@ -3,6 +3,7 @@ from mlx import nn
 
 from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
+from mflux.flux.controlnet import ControlNetOutput
 from mflux.models.transformer.ada_layer_norm_continous import AdaLayerNormContinuous
 from mflux.models.transformer.embed_nd import EmbedND
 from mflux.models.transformer.joint_transformer_block import JointTransformerBlock
@@ -30,6 +31,7 @@ class Transformer(nn.Module):
             pooled_prompt_embeds: mx.array,
             hidden_states: mx.array,
             config: RuntimeConfig,
+            controlnet_samples: ControlNetOutput | None = None
     ) -> mx.array:
         time_step = config.sigmas[t] * config.num_train_steps
         time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision)
diff --git a/src/mflux/post_processing/image.py b/src/mflux/post_processing/image.py
index 5d48ffe..fe69cd6 100644
--- a/src/mflux/post_processing/image.py
+++ b/src/mflux/post_processing/image.py
@@ -11,7 +11,7 @@ from mflux.config.model_config import ModelConfig
 log = logging.getLogger(__name__)
 
 
-class Image:
+class GeneratedImage:
 
     def __init__(
             self,
diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py
index f6f9ea6..5c60089 100644
--- a/src/mflux/post_processing/image_util.py
+++ b/src/mflux/post_processing/image_util.py
@@ -4,7 +4,7 @@ import mlx.core as mx
 import numpy as np
 
 from mflux.config.runtime_config import RuntimeConfig
-from mflux.post_processing.image import Image
+from mflux.post_processing.image import GeneratedImage
 
 
 class ImageUtil:
@@ -19,11 +19,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,
@@ -66,14 +66,8 @@ 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
diff --git a/trial.py b/trial.py
new file mode 100644
index 0000000..53df634
--- /dev/null
+++ b/trial.py
@@ -0,0 +1,37 @@
+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 Config
+from mflux.flux.flux import Flux1
+
+
+prompt = "Luxury food photograph"
+
+# Load the model
+flux = Flux1(
+    model_config=ModelConfig.from_alias("dev"),
+    quantize=None,
+    local_path=None,
+    lora_paths=["diffusion_pytorch_model.safetensors"],
+    lora_scales=None,
+)
+
+# Generate an image
+image = flux.generate_image(
+    seed=3,
+    prompt=prompt,
+    config=Config(
+        num_inference_steps=10,
+        height=256,
+        width=512,
+        guidance=3.5,
+    )
+)
+
+# Save the image
+image.save(path="image.png", export_json_metadata=False)
\ No newline at end of file

From 3dba1e38e17c75af936f80d7078c8b3674b059ac Mon Sep 17 00:00:00 2001
From: Fabio <f.xuzzo@gmail.com>
Date: Sat, 14 Sep 2024 14:37:30 +0200
Subject: [PATCH 02/16] controlnet - first working version

---
 src/mflux/config/config.py                    |  2 +-
 src/mflux/config/runtime_config.py            | 13 ++-
 .../flux_controlnet.py}                       | 91 ++++++++++++-------
 src/mflux/controlnet/utils_controlnet.py      | 10 ++
 src/mflux/models/transformer/feed_forward.py  |  4 +-
 src/mflux/models/transformer/transformer.py   | 20 +++-
 src/mflux/post_processing/image_util.py       |  4 +
 src/mflux/weights/weight_handler.py           | 28 +++++-
 trial.py                                      | 26 +++---
 9 files changed, 141 insertions(+), 57 deletions(-)
 rename src/mflux/{flux/controlnet.py => controlnet/flux_controlnet.py} (77%)
 create mode 100644 src/mflux/controlnet/utils_controlnet.py

diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py
index 7686d3d..271628d 100644
--- a/src/mflux/config/config.py
+++ b/src/mflux/config/config.py
@@ -6,7 +6,7 @@ log = logging.getLogger(__name__)
 
 
 class Config:
-    precision: mx.Dtype = mx.bfloat16
+    precision: mx.Dtype = mx.float16
 
     def __init__(
             self,
diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py
index 33cd634..dabd8f5 100644
--- a/src/mflux/config/runtime_config.py
+++ b/src/mflux/config/runtime_config.py
@@ -1,13 +1,15 @@
+from math import e
+from cycler import V
 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 +37,13 @@ class RuntimeConfig:
     @property
     def num_train_steps(self) -> int:
         return self.model_config.num_train_steps
+    
+    @property
+    def controlnet_conditioning_scale(self) -> float:
+        if isinstance(self.config, ConfigControlnet):
+            return self.config.controlnet_strength
+        else:
+            return ValueError("Controlnet conditioning scale is only available for ConfigControlnet")
 
     @staticmethod
     def _create_sigmas(config, model) -> mx.array:
diff --git a/src/mflux/flux/controlnet.py b/src/mflux/controlnet/flux_controlnet.py
similarity index 77%
rename from src/mflux/flux/controlnet.py
rename to src/mflux/controlnet/flux_controlnet.py
index 54b4a22..ebb6b15 100644
--- a/src/mflux/flux/controlnet.py
+++ b/src/mflux/controlnet/flux_controlnet.py
@@ -10,6 +10,7 @@ from tqdm import tqdm
 from mflux.config.config import Config, ConfigControlnet
 from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
+from mflux.controlnet.utils_controlnet import preprocess_canny
 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
@@ -28,6 +29,8 @@ 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
+import numpy as np
+import cv2
 
 import logging
 
@@ -84,18 +87,18 @@ class Flux1Controlnet:
         if weights.quantization_level is not None:
             self._set_model_weights(weights)
 
-        self.transformer_controlnet = TransformerControlnet(model_config=model_config, local_path=controlnet_path, quantize=quantize)
+        self.transformer_controlnet = TransformerControlnet(model_config=model_config)
 
-        weights_controlnet = WeightHandler.load_transformer(root_path=controlnet_path)
-        if weights_controlnet.quantization_level is None:
+        weights_controlnet, ctrlnet_quantization_level = WeightHandler.load_controlnet_transformer(controlnet_path=controlnet_path)
+        if ctrlnet_quantization_level is None:
             self.transformer_controlnet.update(weights_controlnet)
 
         self.bits = None
-        if quantize is not None or weights.quantization_level is not None:
-            self.bits = weights_controlnet.quantization_level if weights_controlnet.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]) > 64, group_size=64, bits=self.bits)
+        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 weights_controlnet.quantization_level is not None:
+        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:
@@ -112,10 +115,10 @@ class Flux1Controlnet:
             shape=[1, (config.height // 16) * (config.width // 16), 64],
             key=mx.random.key(seed)
         )
-        control_cond = ImageUtil.to_array(control_image)
-        control_cond = self.vae.encode(control_cond)
-        control_cond = (control_cond - self.vae.shift_factor) * self.vae.scaling_factor
-        control_cond = Flux1Controlnet._pack_latents(control_cond, config.height, config.width)
+        control_image = preprocess_canny(control_image)
+        controlnet_cond = ImageUtil.to_array(control_image)
+        controlnet_cond = self.vae.encode(controlnet_cond)
+        controlnet_cond = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width)
 
         # 2. Embedd the prompt
         t5_tokens = self.t5_tokenizer.tokenize(prompt)
@@ -124,12 +127,12 @@ class Flux1Controlnet:
         pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens)
 
         for t in time_steps:
-            controlnet_samples = self.transformer_controlnet(
+            controlnet_block_samples = self.transformer_controlnet.forward(
                 t=t,
                 prompt_embeds=prompt_embeds,
                 pooled_prompt_embeds=pooled_prompt_embeds,
                 hidden_states=latents,
-                control_cond=control_cond,
+                controlnet_cond=controlnet_cond,
                 config=config,
             )
             # 3.t Predict the noise
@@ -139,7 +142,8 @@ class Flux1Controlnet:
                 pooled_prompt_embeds=pooled_prompt_embeds,
                 hidden_states=latents,
                 config=config,
-                controlnet_samples=controlnet_samples,
+                controlnet_block_samples=controlnet_block_samples,
+                # controlnet_single_block_samples=controlnet_single_block_samples,
             )
 
             # 4.t Take one denoise step
@@ -219,11 +223,10 @@ class Flux1Controlnet:
         _save_weights(self.transformer, "transformer")
         _save_weights(self.clip_text_encoder, "text_encoder")
         _save_weights(self.t5_text_encoder, "text_encoder_2")
+        _save_weights(self.transformer_controlnet, "transformer_controlnet")
         
 
-class ControlNetOutput:
-    controlnet_block_samples: Tuple[mx.array]
-    controlnet_single_block_samples: Tuple[mx.array]
+ControlNetOutput = Tuple[list[mx.array], list[mx.array]]
 
 class TransformerControlnet(nn.Module):
 
@@ -233,10 +236,14 @@ class TransformerControlnet(nn.Module):
         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(19)]
-        self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(38)]
-        self.norm_out = AdaLayerNormContinuous(3072, 3072)
-        self.proj_out = nn.Linear(3072, 64)
+        self.transformer_blocks = [JointTransformerBlock(i) for i in range(5)]
+        # self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(38)]
+
+        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(5)]
+
+        # self.controlnet_single_blocks = [nn.Linear(3072, 3072) for _ in range(38)]
 
     def forward(
             self,
@@ -244,11 +251,15 @@ class TransformerControlnet(nn.Module):
             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_conditioning_scale
+
         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)
@@ -257,6 +268,7 @@ class TransformerControlnet(nn.Module):
         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,
@@ -264,18 +276,33 @@ class TransformerControlnet(nn.Module):
                 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)
 
-        for block in self.single_transformer_blocks:
-            hidden_states = block.forward(
-                hidden_states=hidden_states,
-                text_embeddings=text_embeddings,
-                rotary_embeddings=image_rotary_emb
-            )
+        # 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,)
 
-        hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...]
-        hidden_states = self.norm_out.forward(hidden_states, text_embeddings)
-        hidden_states = self.proj_out(hidden_states)
-        noise = hidden_states
-        return noise
\ No newline at end of file
+
+        # 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
\ No newline at end of file
diff --git a/src/mflux/controlnet/utils_controlnet.py b/src/mflux/controlnet/utils_controlnet.py
new file mode 100644
index 0000000..5de6b06
--- /dev/null
+++ b/src/mflux/controlnet/utils_controlnet.py
@@ -0,0 +1,10 @@
+import cv2
+import numpy as np
+import PIL
+
+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)
diff --git a/src/mflux/models/transformer/feed_forward.py b/src/mflux/models/transformer/feed_forward.py
index c5536a0..cb01d8e 100644
--- a/src/mflux/models/transformer/feed_forward.py
+++ b/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:
diff --git a/src/mflux/models/transformer/transformer.py b/src/mflux/models/transformer/transformer.py
index dfc4979..6afe28d 100644
--- a/src/mflux/models/transformer/transformer.py
+++ b/src/mflux/models/transformer/transformer.py
@@ -1,15 +1,17 @@
+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.flux.controlnet import ControlNetOutput
 from mflux.models.transformer.ada_layer_norm_continous import AdaLayerNormContinuous
 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
 
+import numpy as np
+
 
 class Transformer(nn.Module):
 
@@ -31,7 +33,8 @@ class Transformer(nn.Module):
             pooled_prompt_embeds: mx.array,
             hidden_states: mx.array,
             config: RuntimeConfig,
-            controlnet_samples: ControlNetOutput | None = None
+            controlnet_block_samples: Tuple[mx.array] | None = None,
+            controlnet_single_block_samples: Tuple[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)
@@ -44,22 +47,31 @@ class Transformer(nn.Module):
         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:
+                interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                interval_control = int(np.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:
+                hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                    hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+                    + controlnet_single_block_samples[idx]
+                )
 
         hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...]
         hidden_states = self.norm_out.forward(hidden_states, text_embeddings)
diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py
index 5c60089..2ef5b5e 100644
--- a/src/mflux/post_processing/image_util.py
+++ b/src/mflux/post_processing/image_util.py
@@ -71,3 +71,7 @@ class ImageUtil:
         array = mx.transpose(array, (0, 3, 1, 2))
         array = ImageUtil._normalize(array)
         return array
+    
+    @staticmethod
+    def load_image(path: str) -> Image.Image:
+        return Image.open(path)
diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py
index fc485ca..2c2acfe 100644
--- a/src/mflux/weights/weight_handler.py
+++ b/src/mflux/weights/weight_handler.py
@@ -85,6 +85,28 @@ class WeightHandler:
                         "linear2": block["ff_context"]["net"][2]
                     }
         return weights, quantization_level
+    
+    @staticmethod
+    def load_controlnet_transformer(controlnet_path: Path | None = None) -> (dict, int):
+        weights, quantization_level = WeightHandler._get_weights("transformer", weights_path=controlnet_path)
+
+        # 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]
+                    }
+        return weights, quantization_level
 
     @staticmethod
     def load_vae(root_path: Path) -> (dict, int):
@@ -104,7 +126,7 @@ class WeightHandler:
         return weights, quantization_level
 
     @staticmethod
-    def _get_weights(model_name: str, root_path: Path | None = None, lora_path: str | None = None) -> (dict, int):
+    def _get_weights(model_name: str, root_path: Path | None = None, weights_path: str | None = None) -> (dict, int):
         weights = []
         quantization_level = None
 
@@ -114,8 +136,8 @@ class WeightHandler:
                 weight = list(mx.load(str(file)).items())
                 weights.extend(weight)
 
-        if lora_path and root_path is None:
-            weight = list(mx.load(lora_path).items())
+        if weights_path and root_path is None:
+            weight = list(mx.load(weights_path).items())
             weights.extend(weight)
 
         # Non huggingface weights (i.e. ones exported from this project) don't need any reshaping.
diff --git a/trial.py b/trial.py
index 53df634..533905f 100644
--- a/trial.py
+++ b/trial.py
@@ -1,35 +1,35 @@
-import argparse
 import os
 import sys
-import time
+from mflux.post_processing.image_util import ImageUtil
 
 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 Config
-from mflux.flux.flux import Flux1
+from mflux.config.config import ConfigControlnet
+from mflux.controlnet.flux_controlnet import Flux1Controlnet
 
-
-prompt = "Luxury food photograph"
+prompt = "A girl in Venice, 25 years old"
 
 # Load the model
-flux = Flux1(
+flux = Flux1Controlnet(
     model_config=ModelConfig.from_alias("dev"),
-    quantize=None,
-    local_path=None,
-    lora_paths=["diffusion_pytorch_model.safetensors"],
-    lora_scales=None,
+    quantize=8,
+    controlnet_path="diffusion_pytorch_model.safetensors",
 )
 
+control_image = ImageUtil.load_image("image_51.png")
+
 # Generate an image
 image = flux.generate_image(
-    seed=3,
+    seed=4,
     prompt=prompt,
-    config=Config(
+    control_image=control_image,
+    config=ConfigControlnet(
         num_inference_steps=10,
         height=256,
         width=512,
         guidance=3.5,
+        controlnet_conditioning_scale=1.6,
     )
 )
 

From 369bfc0b5371ab792db3a65577d595d0d8e6796d Mon Sep 17 00:00:00 2001
From: Fabio <f.xuzzo@gmail.com>
Date: Sat, 14 Sep 2024 19:58:12 +0200
Subject: [PATCH 03/16] add cli, controlnet config and tidy up

---
 pyproject.toml                              |   1 +
 src/mflux.egg-info/PKG-INFO                 | 402 ++++++++++++++++++++
 src/mflux.egg-info/SOURCES.txt              |  89 +++++
 src/mflux.egg-info/dependency_links.txt     |   1 +
 src/mflux.egg-info/entry_points.txt         |   4 +
 src/mflux.egg-info/requires.txt             |  10 +
 src/mflux.egg-info/top_level.txt            |   1 +
 src/mflux/config/config.py                  |   4 +-
 src/mflux/controlnet/flux_controlnet.py     |  55 ++-
 src/mflux/generate_controlnet.py            |  65 ++++
 src/mflux/models/transformer/transformer.py |   8 +-
 src/mflux/weights/weight_handler.py         |  24 +-
 trial.py                                    |  11 +-
 13 files changed, 638 insertions(+), 37 deletions(-)
 create mode 100644 src/mflux.egg-info/PKG-INFO
 create mode 100644 src/mflux.egg-info/SOURCES.txt
 create mode 100644 src/mflux.egg-info/dependency_links.txt
 create mode 100644 src/mflux.egg-info/entry_points.txt
 create mode 100644 src/mflux.egg-info/requires.txt
 create mode 100644 src/mflux.egg-info/top_level.txt
 create mode 100644 src/mflux/generate_controlnet.py

diff --git a/pyproject.toml b/pyproject.toml
index a168baa..08d7271 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -33,6 +33,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"]
diff --git a/src/mflux.egg-info/PKG-INFO b/src/mflux.egg-info/PKG-INFO
new file mode 100644
index 0000000..1abb1e0
--- /dev/null
+++ b/src/mflux.egg-info/PKG-INFO
@@ -0,0 +1,402 @@
+Metadata-Version: 2.1
+Name: mflux
+Version: 0.2.0
+Summary: A MLX port of FLUX based on the Huggingface Diffusers implementation.
+Author-email: Filip Strand <strand.filip@gmail.com>
+Project-URL: homepage, https://github.com/filipstrand/mflux
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: MacOS
+Description-Content-Type: text/markdown
+Requires-Dist: mlx>=0.16.0
+Requires-Dist: numpy>=2.0.0
+Requires-Dist: pillow>=10.4.0
+Requires-Dist: transformers>=4.44.0
+Requires-Dist: sentencepiece>=0.2.0
+Requires-Dist: torch>=2.3.1
+Requires-Dist: tqdm>=4.66.5
+Requires-Dist: huggingface-hub>=0.24.5
+Requires-Dist: safetensors>=0.4.4
+Requires-Dist: piexif>=1.1.3
+
+
+![image](src/mflux/assets/logo.png)
+*A MLX port of FLUX based on the Huggingface Diffusers implementation.*
+
+### About
+
+Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black Forest Labs](https://blackforestlabs.ai) locally on your Mac!
+
+### Philosophy
+
+MFLUX is a line-by-line port of the FLUX implementation in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) library to [Apple MLX](https://github.com/ml-explore/mlx). 
+MFLUX is purposefully kept minimal and explicit - Network architectures are hardcoded and no config files are used
+except for the tokenizers. The aim is to have a tiny codebase with the single purpose of expressing these models 
+(thereby avoiding too many abstractions). While MFLUX priorities readability over generality and performance, [it can still be quite fast](#image-generation-speed-updated), [and even faster quantized](#quantization).
+
+All models are implemented from scratch in MLX and only the tokenizers are used via the 
+[Huggingface Transformers](https://github.com/huggingface/transformers) library. Other than that, there are only minimal dependencies 
+like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) for simple image post-processing.
+
+### Models
+
+- [x] FLUX.1-Scnhell
+- [x] FLUX.1-Dev
+
+### Installation
+For users, the easiest way to install MFLUX is to first create a new virtual environment to isolate dependencies:
+   ```
+   mkdir -p mflux && cd mflux && python3 -m venv .venv && source .venv/bin/activate
+   ```
+This creates and activates a virtual environment in the `mflux` folder. After that, install MFLUX via pip:
+   ```
+   pip install -U mflux
+   ```
+<details>
+<summary>For contributors (click to expand)</summary>
+
+1. Clone the repo:
+ ```
+ git clone git@github.com:filipstrand/mflux.git
+ ```
+2. Navigate to the project and set up a virtual environment:
+ ```
+ cd mflux && python3 -m venv .venv && source .venv/bin/activate
+ ```
+3. Install the required dependencies:
+ ```
+ pip install -r requirements.txt
+ ```
+</details>
+
+### Generating an image
+
+Run the command `mflux-generate` by specifying a prompt and the model and some optional arguments. For example, here we use a quantized version of the `schnell` model for 2 steps:
+
+```
+mflux-generate --model schnell --prompt "Luxury food photograph" --steps 2 --seed 2 -q 8
+```
+
+This example uses the more powerful `dev` model with 25 time steps:
+
+```
+mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2 -q 8
+```
+
+⚠️ *If the specific model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell or Dev model respectively). See the [quantization](#quantization) section for running compressed versions of the model.* ⚠️
+
+*By default, model files are downloaded to the `.cache` folder within your home directory. For example, in my setup, the path looks like this:* 
+
+```
+/Users/filipstrand/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-dev
+```
+
+*To change this default behavior, you can do so by modifying the `HF_HOME` environment variable. For more details on how to adjust this setting, please refer to the [Hugging Face documentation](https://huggingface.co/docs/huggingface_hub/en/package_reference/environment_variables)*.
+
+🔒 [FLUX.1-dev currently requires granted access to its Huggingface repo. For troubleshooting, see the issue tracker](https://github.com/filipstrand/mflux/issues/14) 🔒 
+
+#### Full list of Command-Line Arguments 
+
+- **`--prompt`** (required, `str`): Text description of the image to generate.
+
+- **`--model`** or **`-m`** (required, `str`): Model to use for generation (`"schnell"` or `"dev"`).
+
+- **`--output`** (optional, `str`, default: `"image.png"`): Output image filename.
+
+- **`--seed`** (optional, `int`, default: `None`): Seed for random number generation. Default is time-based.
+
+- **`--height`** (optional, `int`, default: `1024`): Height of the output image in pixels.
+
+- **`--width`** (optional, `int`, default: `1024`): Width of the output image in pixels.
+
+- **`--steps`** (optional, `int`, default: `4`): Number of inference steps.
+
+- **`--guidance`** (optional, `float`, default: `3.5`): Guidance scale (only used for `"dev"` model).
+
+- **`--path`** (optional, `str`, default: `None`): Path to a local model on disk.
+
+- **`--quantize`** or **`-q`** (optional, `int`, default: `None`): [Quantization](#quantization) (choose between `4` or `8`).
+
+- **`--lora-paths`** (optional, `[str]`, default: `None`): The paths to the [LoRA](#LoRA) weights.
+
+- **`--lora-scales`** (optional, `[float]`, default: `None`): The scale for each respective [LoRA](#LoRA) (will default to `1.0` if not specified and only one LoRA weight is loaded.)
+ 
+- **`--metadata`** (optional): Exports a `.json` file containing the metadata for the image with the same name. (Even without this flag, the image metadata is saved and can be viewed using `exiftool image.png`)
+
+Or, with the correct python environment active, create and run a separate script like the following:
+
+```python
+from mflux.flux.flux import Flux1
+from mflux.config.config import Config
+
+# Load the model
+flux = Flux1.from_alias(
+   alias="schnell",  # "schnell" or "dev"
+   quantize=8,       # 4 or 8
+)  
+
+# Generate an image
+image = flux.generate_image(
+   seed=2,
+   prompt="Luxury food photograph",
+   config=Config(
+      num_inference_steps=2,  # "schnell" works well with 2-4 steps, "dev" works well with 20-25 steps
+      height=1024,
+      width=1024,
+   )
+)
+
+image.save(path="image.png")
+```
+
+For more options on how to configure MFLUX, please see [generate.py](src/mflux/generate.py).
+
+### Image generation speed (updated)
+
+These numbers are based on the non-quantized `schnell` model, with the configuration provided in the code snippet below. 
+To time your machine, run the following:
+```
+time mflux-generate \
+--prompt "Luxury food photograph" \
+--model schnell \
+--steps 2 \
+--seed 2 \
+--height 1024 \
+--width 1024
+```
+
+| Device             | User                                                                                                                        | Reported Time | Notes                     |
+|--------------------|-----------------------------------------------------------------------------------------------------------------------------|---------------|---------------------------|
+| M3 Max             | [@karpathy](https://gist.github.com/awni/a67d16d50f0f492d94a10418e0592bde?permalink_comment_id=5153531#gistcomment-5153531) | ~20s          |                           |
+| M2 Ultra           | [@awni](https://x.com/awnihannun/status/1823515121827897385)                                                                | <15s          |                           |
+| 2023 M2 Max (96GB) | [@explorigin](https://github.com/filipstrand/mflux/issues/6)                                                                | ~25s          |                           |
+| 2021 M1 Pro (16GB) | [@qw-in](https://github.com/filipstrand/mflux/issues/7)                                                                     | ~175s         | Might freeze your mac     |
+| 2023 M3 Pro (36GB) | [@kush-gupt](https://github.com/filipstrand/mflux/issues/11)                                                                | ~80s          |                           |
+| 2020 M1 (8GB)      | [@mbvillaverde](https://github.com/filipstrand/mflux/issues/13)                                                             | ~335s         | With resolution 512 x 512 |
+| 2022 M1 MAX (64GB) | [@BosseParra](https://x.com/BosseParra/status/1826191780812877968)                                                          | ~55s          |                           |
+| 2021 M1 Pro (32GB) | @filipstrand                                                                                                                | ~160s         |                           |
+| 2023 M2 Max (32GB) | @filipstrand                                                                                                                | ~70s          |                           |
+
+*Note that these numbers includes starting the application from scratch, which means doing model i/o, setting/quantizing weights etc. 
+If we assume that the model is already loaded, you can inspect the image metadata using `exiftool image.png` and see the total duration of the denoising loop (excluding text embedding).*
+
+### Equivalent to Diffusers implementation 
+
+There is only a single source of randomness when generating an image: The initial latent array. 
+In this implementation, this initial latent is fully deterministically controlled by the input `seed` parameter. 
+However, if we were to import a fixed instance of this latent array saved from the Diffusers implementation, then MFLUX will produce an identical image to the Diffusers implementation (assuming a fixed prompt and using the default parameter settings in the Diffusers setup).
+
+
+The images below illustrate this equivalence. 
+In all cases the Schnell model was run for 2 time steps. 
+The Diffusers implementation ran in CPU mode. 
+The precision for MFLUX can be set in the [Config](src/mflux/config/config.py) class. 
+There is typically a noticeable but very small difference in the final image when switching between 16bit and 32bit precision.
+
+---
+```
+Luxury food photograph
+```
+![image](src/mflux/assets/comparison1.jpg)
+
+---
+```
+detailed cinematic dof render of an old dusty detailed CRT monitor on a wooden desk in a dim room with items around, messy dirty room. On the screen are the letters "FLUX" glowing softly. High detail hard surface render
+```
+![image](src/mflux/assets/comparison2.jpg)
+
+---
+
+```
+photorealistic, lotr, A tiny red dragon curled up asleep inside a nest, (Soft Focus) , (f_stop 2.8) , (focal_length 50mm) macro lens f/2. 8, medieval wizard table, (pastel) colors, (cozy) morning light filtering through a nearby window, (whimsical) steam shapes, captured with a (Canon EOS R5) , highlighting (serene) comfort, medieval, dnd, rpg, 3d, 16K, 8K
+```
+![image](src/mflux/assets/comparison3.jpg)
+
+---
+
+
+```
+A weathered fisherman in his early 60s stands on the deck of his boat, gazing out at a stormy sea. He has a thick, salt-and-pepper beard, deep-set blue eyes, and skin tanned and creased from years of sun exposure. He's wearing a yellow raincoat and hat, with water droplets clinging to the fabric. Behind him, dark clouds loom ominously, and waves crash against the side of the boat. The overall atmosphere is one of tension and respect for the power of nature.
+```
+![image](src/mflux/assets/comparison4.jpg)
+
+---
+
+```
+Luxury food photograph of an italian Linguine pasta alle vongole dish with lots of clams. It has perfect lighting and a cozy background with big bokeh and shallow depth of field. The mood is a sunset balcony in tuscany.  The photo is taken from the side of the plate. The pasta is shiny with sprinkled parmesan cheese and basil leaves on top. The scene is complemented by a warm, inviting light that highlights the textures and colors of the ingredients, giving it an appetizing and elegant look.
+```
+![image](src/mflux/assets/comparison5.jpg)
+
+---
+
+### Quantization 
+
+MFLUX supports running FLUX in 4-bit or 8-bit quantized mode. Running a quantized version can greatly speed up the
+generation process and reduce the memory consumption by several gigabytes. [Quantized models also take up less disk space](#size-comparisons-for-quantized-models). 
+
+```
+mflux-generate \
+    --model schnell \
+    --steps 2 \
+    --seed 2 \
+    --quantize 8 \
+    --height 1920 \
+    --width 1024 \
+    --prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond’s calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier" 
+```
+![image](src/mflux/assets/comparison6.jpg)
+
+*In this example, weights are quantized at **runtime** - this is convenient if you don't want to [save a quantized copy of the weights to disk](#saving-a-quantized-version-to-disk), but still want to benefit from the potential speedup and RAM reduction quantization might bring.*
+
+
+By selecting the `--quantize` or `-q` flag to be `4`, `8`, or removing it entirely, we get all 3 images above. As can be seen, there is very little difference between the images (especially between the 8-bit, and the non-quantized result).
+Image generation times in this example are based on a 2021 M1 Pro (32GB) machine. Even though the images are almost identical, there is a ~2x speedup by
+running the 8-bit quantized version on this particular machine. Unlike the non-quantized version, for the 8-bit version the swap memory usage is drastically reduced and GPU utilization is close to 100% during the whole generation. Results here can vary across different machines.
+
+#### Size comparisons for quantized models 
+
+The model sizes for both `schnell` and `dev` at various quantization levels are as follows:
+
+| 4 bit  | 8 bit   | Original (16 bit) |
+|--------|---------|-------------------|
+| 9.85GB | 18.16GB | 33.73GB           |
+
+The reason weights sizes are not fully cut in half is because a small number of weights are not quantized and kept at full precision.
+
+#### Saving a quantized version to disk
+
+To save a local copy of the quantized weights, run the `mflux-save` command like so:
+
+```
+mflux-save \
+    --path "/Users/filipstrand/Desktop/schnell_8bit" \
+    --model schnell \
+    --quantize 8
+```
+
+*Note that when saving a quantized version, you will need the original huggingface weights.*
+
+#### Loading and running a quantized version from disk
+
+To generate a new image from the quantized model, simply provide a `--path` to where it was saved: 
+
+```
+mflux-generate \
+    --path "/Users/filipstrand/Desktop/schnell_8bit" \
+    --model schnell \
+    --steps 2 \
+    --seed 2 \
+    --height 1920 \
+    --width 1024 \
+    --prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond’s calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier" 
+```
+
+*Note: When loading a quantized model from disk, there is no need to pass in `-q` flag, since we can infer this from the weight metadata.*
+
+*Also Note: Once we have a local model (quantized [or not](#running-a-non-quantized-model-directly-from-disk)) specified via the `--path` argument, the huggingface cache models are not required to launch the model. 
+In other words, you can reclaim the 34GB diskspace (per model) by deleting the full 16-bit model from the [Huggingface cache](#generating-an-image) if you choose.* 
+
+*If you don't want to download the full models and quantize them yourself, the 4-bit weights are available here for a direct download:* 
+- [madroid/flux.1-schnell-mflux-4bit](https://huggingface.co/madroid/flux.1-schnell-mflux-4bit)
+- [madroid/flux.1-dev-mflux-4bit](https://huggingface.co/madroid/flux.1-dev-mflux-4bit)
+
+### Running a non-quantized model directly from disk
+
+MFLUX also supports running a non-quantized model directly from a custom location.  
+In the example below, the model is placed in `/Users/filipstrand/Desktop/schnell`:
+
+```
+mflux-generate \
+    --path "/Users/filipstrand/Desktop/schnell" \
+    --model schnell \
+    --steps 2 \
+    --seed 2 \
+    --prompt "Luxury food photograph" 
+```
+
+Note that the `--model` flag must be set when loading a model from disk. 
+
+Also note that unlike when using the typical `alias` way of initializing the model (which internally handles that the required resources are downloaded),
+when loading a model directly from disk, we require the downloaded models to look like the following:
+
+```
+.
+├── text_encoder
+│   └── model.safetensors
+├── text_encoder_2
+│   ├── model-00001-of-00002.safetensors
+│   └── model-00002-of-00002.safetensors
+├── tokenizer
+│   ├── merges.txt
+│   ├── special_tokens_map.json
+│   ├── tokenizer_config.json
+│   └── vocab.json
+├── tokenizer_2
+│   ├── special_tokens_map.json
+│   ├── spiece.model
+│   ├── tokenizer.json
+│   └── tokenizer_config.json
+├── transformer
+│   ├── diffusion_pytorch_model-00001-of-00003.safetensors
+│   ├── diffusion_pytorch_model-00002-of-00003.safetensors
+│   └── diffusion_pytorch_model-00003-of-00003.safetensors
+└── vae
+    └── diffusion_pytorch_model.safetensors
+```
+This mirrors how the resources are placed in the [HuggingFace Repo](https://huggingface.co/black-forest-labs/FLUX.1-schnell/tree/main) for FLUX.1.
+*Huggingface weights, unlike quantized ones exported directly from this project, have to be
+processed a bit differently, which is why we require this structure above.* 
+
+
+### LoRA
+
+MFLUX support loading trained [LoRA](https://huggingface.co/docs/diffusers/en/training/lora) adapters (actual training support is coming).
+
+The following example [The_Hound](https://huggingface.co/TheLastBen/The_Hound) LoRA from [@TheLastBen](https://github.com/TheLastBen): 
+
+```
+mflux-generate --prompt "sandor clegane" --model dev --steps 20 --seed 43 -q 8 --lora-paths "sandor_clegane_single_layer.safetensors"
+```
+
+![image](src/mflux/assets/lora1.jpg)
+---
+
+The following example is [Flux_1_Dev_LoRA_Paper-Cutout-Style](https://huggingface.co/Norod78/Flux_1_Dev_LoRA_Paper-Cutout-Style) LoRA from [@Norod78](https://huggingface.co/Norod78):
+
+```
+mflux-generate --prompt "pikachu, Paper Cutout Style" --model schnell --steps 4 --seed 43 -q 8 --lora-paths "Flux_1_Dev_LoRA_Paper-Cutout-Style.safetensors"
+```
+![image](src/mflux/assets/lora2.jpg)
+
+*Note that LoRA trained weights are typically trained with a **trigger word or phrase**. For example, in the latter case, the sentence should include the phrase **"Paper Cutout Style"**.*
+
+*Also note that the same LoRA weights can work well with both the `schnell` and `dev` models. Refer to the original LoRA repository to see what mode it was trained for.*
+
+#### Multi-LoRA
+
+Multiple LoRAs can be sent in to combine the effects of the individual adapters. The following example combines both of the above LoRAs:
+
+```
+mflux-generate \
+   --prompt "sandor clegane in a forest, Paper Cutout Style" \
+   --model dev \
+   --steps 20 \
+   --seed 43 \
+   --lora-paths sandor_clegane_single_layer.safetensors Flux_1_Dev_LoRA_Paper-Cutout-Style.safetensors \
+   --lora-scales 1.0 1.0 \
+   -q 8
+```
+![image](src/mflux/assets/lora3.jpg)
+
+Just to see the difference, this image displays the four cases: One of having both adapters fully active, partially active and no LoRA at all. 
+The example above also show the usage of `--lora-scales` flag. 
+
+### Current limitations
+
+- Images are generated one by one.
+- Negative prompts not supported.
+- LoRA weights are only supported for the transformer part of the network.
+
+### TODO
+
+- [ ] LoRA fine-tuning
+- [ ] Frontend support (Gradio/Streamlit/Other?)
diff --git a/src/mflux.egg-info/SOURCES.txt b/src/mflux.egg-info/SOURCES.txt
new file mode 100644
index 0000000..d3e63de
--- /dev/null
+++ b/src/mflux.egg-info/SOURCES.txt
@@ -0,0 +1,89 @@
+README.md
+pyproject.toml
+src/mflux/__init__.py
+src/mflux/generate.py
+src/mflux/generate_controlnet.py
+src/mflux/save.py
+src/mflux.egg-info/PKG-INFO
+src/mflux.egg-info/SOURCES.txt
+src/mflux.egg-info/dependency_links.txt
+src/mflux.egg-info/entry_points.txt
+src/mflux.egg-info/requires.txt
+src/mflux.egg-info/top_level.txt
+src/mflux/config/__init__.py
+src/mflux/config/config.py
+src/mflux/config/model_config.py
+src/mflux/config/runtime_config.py
+src/mflux/controlnet/flux_controlnet.py
+src/mflux/controlnet/utils_controlnet.py
+src/mflux/flux/__init__.py
+src/mflux/flux/flux.py
+src/mflux/models/__init__.py
+src/mflux/models/text_encoder/__init__.py
+src/mflux/models/text_encoder/clip_encoder/__init__.py
+src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py
+src/mflux/models/text_encoder/clip_encoder/clip_encoder.py
+src/mflux/models/text_encoder/clip_encoder/clip_encoder_layer.py
+src/mflux/models/text_encoder/clip_encoder/clip_mlp.py
+src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py
+src/mflux/models/text_encoder/clip_encoder/clip_text_model.py
+src/mflux/models/text_encoder/clip_encoder/encoder_clip.py
+src/mflux/models/text_encoder/t5_encoder/__init__.py
+src/mflux/models/text_encoder/t5_encoder/t5_attention.py
+src/mflux/models/text_encoder/t5_encoder/t5_block.py
+src/mflux/models/text_encoder/t5_encoder/t5_dense_relu_dense.py
+src/mflux/models/text_encoder/t5_encoder/t5_encoder.py
+src/mflux/models/text_encoder/t5_encoder/t5_feed_forward.py
+src/mflux/models/text_encoder/t5_encoder/t5_layer_norm.py
+src/mflux/models/text_encoder/t5_encoder/t5_self_attention.py
+src/mflux/models/transformer/__init__.py
+src/mflux/models/transformer/ada_layer_norm_continous.py
+src/mflux/models/transformer/ada_layer_norm_zero.py
+src/mflux/models/transformer/ada_layer_norm_zero_single.py
+src/mflux/models/transformer/embed_nd.py
+src/mflux/models/transformer/feed_forward.py
+src/mflux/models/transformer/guidance_embedder.py
+src/mflux/models/transformer/joint_attention.py
+src/mflux/models/transformer/joint_transformer_block.py
+src/mflux/models/transformer/single_block_attention.py
+src/mflux/models/transformer/single_transformer_block.py
+src/mflux/models/transformer/text_embedder.py
+src/mflux/models/transformer/time_text_embed.py
+src/mflux/models/transformer/timestep_embedder.py
+src/mflux/models/transformer/transformer.py
+src/mflux/models/vae/__init__.py
+src/mflux/models/vae/vae.py
+src/mflux/models/vae/common/__init__.py
+src/mflux/models/vae/common/attention.py
+src/mflux/models/vae/common/resnet_block_2d.py
+src/mflux/models/vae/common/unet_mid_block.py
+src/mflux/models/vae/decoder/__init__.py
+src/mflux/models/vae/decoder/conv_in.py
+src/mflux/models/vae/decoder/conv_norm_out.py
+src/mflux/models/vae/decoder/conv_out.py
+src/mflux/models/vae/decoder/decoder.py
+src/mflux/models/vae/decoder/up_block_1_or_2.py
+src/mflux/models/vae/decoder/up_block_3.py
+src/mflux/models/vae/decoder/up_block_4.py
+src/mflux/models/vae/decoder/up_sampler.py
+src/mflux/models/vae/encoder/__init__.py
+src/mflux/models/vae/encoder/conv_in.py
+src/mflux/models/vae/encoder/conv_norm_out.py
+src/mflux/models/vae/encoder/conv_out.py
+src/mflux/models/vae/encoder/down_block_1.py
+src/mflux/models/vae/encoder/down_block_2.py
+src/mflux/models/vae/encoder/down_block_3.py
+src/mflux/models/vae/encoder/down_block_4.py
+src/mflux/models/vae/encoder/down_sampler.py
+src/mflux/models/vae/encoder/encoder.py
+src/mflux/post_processing/__init__.py
+src/mflux/post_processing/image.py
+src/mflux/post_processing/image_util.py
+src/mflux/tokenizer/__init__.py
+src/mflux/tokenizer/clip_tokenizer.py
+src/mflux/tokenizer/t5_tokenizer.py
+src/mflux/tokenizer/tokenizer_handler.py
+src/mflux/weights/__init__.py
+src/mflux/weights/lora_util.py
+src/mflux/weights/weight_handler.py
+src/mflux/weights/weight_util.py
\ No newline at end of file
diff --git a/src/mflux.egg-info/dependency_links.txt b/src/mflux.egg-info/dependency_links.txt
new file mode 100644
index 0000000..8b13789
--- /dev/null
+++ b/src/mflux.egg-info/dependency_links.txt
@@ -0,0 +1 @@
+
diff --git a/src/mflux.egg-info/entry_points.txt b/src/mflux.egg-info/entry_points.txt
new file mode 100644
index 0000000..e9824aa
--- /dev/null
+++ b/src/mflux.egg-info/entry_points.txt
@@ -0,0 +1,4 @@
+[console_scripts]
+mflux-generate = mflux.generate:main
+mflux-generate-controlnet = mflux.generate_controlnet:main
+mflux-save = mflux.save:main
diff --git a/src/mflux.egg-info/requires.txt b/src/mflux.egg-info/requires.txt
new file mode 100644
index 0000000..149534c
--- /dev/null
+++ b/src/mflux.egg-info/requires.txt
@@ -0,0 +1,10 @@
+mlx>=0.16.0
+numpy>=2.0.0
+pillow>=10.4.0
+transformers>=4.44.0
+sentencepiece>=0.2.0
+torch>=2.3.1
+tqdm>=4.66.5
+huggingface-hub>=0.24.5
+safetensors>=0.4.4
+piexif>=1.1.3
diff --git a/src/mflux.egg-info/top_level.txt b/src/mflux.egg-info/top_level.txt
new file mode 100644
index 0000000..4bd59a6
--- /dev/null
+++ b/src/mflux.egg-info/top_level.txt
@@ -0,0 +1 @@
+mflux
diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py
index 271628d..88566ae 100644
--- a/src/mflux/config/config.py
+++ b/src/mflux/config/config.py
@@ -30,7 +30,7 @@ class ConfigControlnet(Config):
             width: int = 1024,
             height: int = 1024,
             guidance: float = 4.0,
-            controlnet_conditioning_scale: float = 1.0,
+            controlnet_strength: float = 1.0,
     ):
         super().__init__(num_inference_steps, width, height, guidance)
-        self.controlnet_conditioning_scale = controlnet_conditioning_scale
+        self.controlnet_conditioning_scale = controlnet_strength
diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py
index ebb6b15..5a5a71b 100644
--- a/src/mflux/controlnet/flux_controlnet.py
+++ b/src/mflux/controlnet/flux_controlnet.py
@@ -36,6 +36,7 @@ import logging
 
 log = logging.getLogger(__name__)
 
+CONTOLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny"
 
 class Flux1Controlnet:
     def __init__(
@@ -87,9 +88,13 @@ class Flux1Controlnet:
         if weights.quantization_level is not None:
             self._set_model_weights(weights)
 
-        self.transformer_controlnet = TransformerControlnet(model_config=model_config)
+        weights_controlnet, ctrlnet_quantization_level, controlnet_config = WeightHandler.load_controlnet_transformer(controlnet_id=CONTOLNET_ID)
+        self.transformer_controlnet = TransformerControlnet(
+            model_config=model_config,
+            num_blocks= controlnet_config["num_layers"],
+            num_single_blocks= controlnet_config["num_single_layers"],
+        )
 
-        weights_controlnet, ctrlnet_quantization_level = WeightHandler.load_controlnet_transformer(controlnet_path=controlnet_path)
         if ctrlnet_quantization_level is None:
             self.transformer_controlnet.update(weights_controlnet)
 
@@ -117,7 +122,10 @@ class Flux1Controlnet:
         )
         control_image = preprocess_canny(control_image)
         controlnet_cond = ImageUtil.to_array(control_image)
-        controlnet_cond = self.vae.encode(controlnet_cond)
+        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
@@ -230,20 +238,25 @@ ControlNetOutput = Tuple[list[mx.array], list[mx.array]]
 
 class TransformerControlnet(nn.Module):
 
-    def __init__(self, model_config: ModelConfig):
+    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(5)]
-        # self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(38)]
+        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(5)]
+        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(38)]
+        self.controlnet_single_blocks = [nn.Linear(3072, 3072) for _ in range(num_single_blocks)]
 
     def forward(
             self,
@@ -287,22 +300,22 @@ class TransformerControlnet(nn.Module):
             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] :],)
+        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,)
+        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]
+        controlnet_single_block_samples = [sample * conditioning_scale for sample in controlnet_single_block_samples]
 
         return controlnet_block_samples
\ No newline at end of file
diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py
new file mode 100644
index 0000000..a8b3c57
--- /dev/null
+++ b/src/mflux/generate_controlnet.py
@@ -0,0 +1,65 @@
+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=4, 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")
+
+    # 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()
diff --git a/src/mflux/models/transformer/transformer.py b/src/mflux/models/transformer/transformer.py
index 6afe28d..10af56e 100644
--- a/src/mflux/models/transformer/transformer.py
+++ b/src/mflux/models/transformer/transformer.py
@@ -1,6 +1,7 @@
 from typing import Tuple
 import mlx.core as mx
 from mlx import nn
+import math
 
 from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
@@ -10,7 +11,6 @@ from mflux.models.transformer.joint_transformer_block import JointTransformerBlo
 from mflux.models.transformer.single_transformer_block import SingleTransformerBlock
 from mflux.models.transformer.time_text_embed import TimeTextEmbed
 
-import numpy as np
 
 
 class Transformer(nn.Module):
@@ -56,7 +56,7 @@ class Transformer(nn.Module):
             )
             if controlnet_block_samples is not None:
                 interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
-                interval_control = int(np.ceil(interval_control))
+                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)
@@ -68,9 +68,11 @@ class Transformer(nn.Module):
                 rotary_embeddings=image_rotary_emb
             )
             if controlnet_single_block_samples is not None:
+                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]
+                    + controlnet_single_block_samples[idx // interval_control]
                 )
 
         hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...]
diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py
index 2c2acfe..12e5a95 100644
--- a/src/mflux/weights/weight_handler.py
+++ b/src/mflux/weights/weight_handler.py
@@ -1,3 +1,4 @@
+import json
 from pathlib import Path
 
 import mlx.core as mx
@@ -87,8 +88,18 @@ class WeightHandler:
         return weights, quantization_level
     
     @staticmethod
-    def load_controlnet_transformer(controlnet_path: Path | None = None) -> (dict, int):
-        weights, quantization_level = WeightHandler._get_weights("transformer", weights_path=controlnet_path)
+    def load_controlnet_transformer(controlnet_id: Path | None = None) -> (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:
@@ -106,7 +117,8 @@ class WeightHandler:
                         "linear1": block["ff_context"]["net"][0]["proj"],
                         "linear2": block["ff_context"]["net"][2]
                     }
-        return weights, quantization_level
+        config = json.load(open(controlnet_path / "config.json"))
+        return weights, quantization_level, config
 
     @staticmethod
     def load_vae(root_path: Path) -> (dict, int):
@@ -126,7 +138,7 @@ class WeightHandler:
         return weights, quantization_level
 
     @staticmethod
-    def _get_weights(model_name: str, root_path: Path | None = None, weights_path: str | None = None) -> (dict, int):
+    def _get_weights(model_name: str, root_path: Path | None = None, lora_path: str | None = None) -> (dict, int):
         weights = []
         quantization_level = None
 
@@ -136,8 +148,8 @@ class WeightHandler:
                 weight = list(mx.load(str(file)).items())
                 weights.extend(weight)
 
-        if weights_path and root_path is None:
-            weight = list(mx.load(weights_path).items())
+        if lora_path and root_path is None:
+            weight = list(mx.load(lora_path).items())
             weights.extend(weight)
 
         # Non huggingface weights (i.e. ones exported from this project) don't need any reshaping.
diff --git a/trial.py b/trial.py
index 533905f..ee933b2 100644
--- a/trial.py
+++ b/trial.py
@@ -8,20 +8,21 @@ from mflux.config.model_config import ModelConfig
 from mflux.config.config import ConfigControlnet
 from mflux.controlnet.flux_controlnet import Flux1Controlnet
 
-prompt = "A girl in Venice, 25 years old"
+prompt = "Luxury picture of food"
+
 
 # Load the model
 flux = Flux1Controlnet(
     model_config=ModelConfig.from_alias("dev"),
     quantize=8,
-    controlnet_path="diffusion_pytorch_model.safetensors",
+    lora_paths=["Flux_1_Dev_LoRA_Paper-Cutout-Style.safetensors"]
 )
 
-control_image = ImageUtil.load_image("image_51.png")
+control_image = ImageUtil.load_image("image_1.png")
 
 # Generate an image
 image = flux.generate_image(
-    seed=4,
+    seed=2500,
     prompt=prompt,
     control_image=control_image,
     config=ConfigControlnet(
@@ -29,7 +30,7 @@ image = flux.generate_image(
         height=256,
         width=512,
         guidance=3.5,
-        controlnet_conditioning_scale=1.6,
+        controlnet_strength=0.7,
     )
 )
 

From 2a2a3fff638e696fd166ecca46984aea17310a28 Mon Sep 17 00:00:00 2001
From: Fabio <f.xuzzo@gmail.com>
Date: Sat, 14 Sep 2024 20:01:33 +0200
Subject: [PATCH 04/16] delete egg-info

---
 src/mflux.egg-info/PKG-INFO             | 402 ------------------------
 src/mflux.egg-info/SOURCES.txt          |  89 ------
 src/mflux.egg-info/dependency_links.txt |   1 -
 src/mflux.egg-info/entry_points.txt     |   4 -
 src/mflux.egg-info/requires.txt         |  10 -
 src/mflux.egg-info/top_level.txt        |   1 -
 6 files changed, 507 deletions(-)
 delete mode 100644 src/mflux.egg-info/PKG-INFO
 delete mode 100644 src/mflux.egg-info/SOURCES.txt
 delete mode 100644 src/mflux.egg-info/dependency_links.txt
 delete mode 100644 src/mflux.egg-info/entry_points.txt
 delete mode 100644 src/mflux.egg-info/requires.txt
 delete mode 100644 src/mflux.egg-info/top_level.txt

diff --git a/src/mflux.egg-info/PKG-INFO b/src/mflux.egg-info/PKG-INFO
deleted file mode 100644
index 1abb1e0..0000000
--- a/src/mflux.egg-info/PKG-INFO
+++ /dev/null
@@ -1,402 +0,0 @@
-Metadata-Version: 2.1
-Name: mflux
-Version: 0.2.0
-Summary: A MLX port of FLUX based on the Huggingface Diffusers implementation.
-Author-email: Filip Strand <strand.filip@gmail.com>
-Project-URL: homepage, https://github.com/filipstrand/mflux
-Classifier: Programming Language :: Python :: 3
-Classifier: Operating System :: MacOS
-Description-Content-Type: text/markdown
-Requires-Dist: mlx>=0.16.0
-Requires-Dist: numpy>=2.0.0
-Requires-Dist: pillow>=10.4.0
-Requires-Dist: transformers>=4.44.0
-Requires-Dist: sentencepiece>=0.2.0
-Requires-Dist: torch>=2.3.1
-Requires-Dist: tqdm>=4.66.5
-Requires-Dist: huggingface-hub>=0.24.5
-Requires-Dist: safetensors>=0.4.4
-Requires-Dist: piexif>=1.1.3
-
-
-![image](src/mflux/assets/logo.png)
-*A MLX port of FLUX based on the Huggingface Diffusers implementation.*
-
-### About
-
-Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black Forest Labs](https://blackforestlabs.ai) locally on your Mac!
-
-### Philosophy
-
-MFLUX is a line-by-line port of the FLUX implementation in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) library to [Apple MLX](https://github.com/ml-explore/mlx). 
-MFLUX is purposefully kept minimal and explicit - Network architectures are hardcoded and no config files are used
-except for the tokenizers. The aim is to have a tiny codebase with the single purpose of expressing these models 
-(thereby avoiding too many abstractions). While MFLUX priorities readability over generality and performance, [it can still be quite fast](#image-generation-speed-updated), [and even faster quantized](#quantization).
-
-All models are implemented from scratch in MLX and only the tokenizers are used via the 
-[Huggingface Transformers](https://github.com/huggingface/transformers) library. Other than that, there are only minimal dependencies 
-like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) for simple image post-processing.
-
-### Models
-
-- [x] FLUX.1-Scnhell
-- [x] FLUX.1-Dev
-
-### Installation
-For users, the easiest way to install MFLUX is to first create a new virtual environment to isolate dependencies:
-   ```
-   mkdir -p mflux && cd mflux && python3 -m venv .venv && source .venv/bin/activate
-   ```
-This creates and activates a virtual environment in the `mflux` folder. After that, install MFLUX via pip:
-   ```
-   pip install -U mflux
-   ```
-<details>
-<summary>For contributors (click to expand)</summary>
-
-1. Clone the repo:
- ```
- git clone git@github.com:filipstrand/mflux.git
- ```
-2. Navigate to the project and set up a virtual environment:
- ```
- cd mflux && python3 -m venv .venv && source .venv/bin/activate
- ```
-3. Install the required dependencies:
- ```
- pip install -r requirements.txt
- ```
-</details>
-
-### Generating an image
-
-Run the command `mflux-generate` by specifying a prompt and the model and some optional arguments. For example, here we use a quantized version of the `schnell` model for 2 steps:
-
-```
-mflux-generate --model schnell --prompt "Luxury food photograph" --steps 2 --seed 2 -q 8
-```
-
-This example uses the more powerful `dev` model with 25 time steps:
-
-```
-mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2 -q 8
-```
-
-⚠️ *If the specific model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell or Dev model respectively). See the [quantization](#quantization) section for running compressed versions of the model.* ⚠️
-
-*By default, model files are downloaded to the `.cache` folder within your home directory. For example, in my setup, the path looks like this:* 
-
-```
-/Users/filipstrand/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-dev
-```
-
-*To change this default behavior, you can do so by modifying the `HF_HOME` environment variable. For more details on how to adjust this setting, please refer to the [Hugging Face documentation](https://huggingface.co/docs/huggingface_hub/en/package_reference/environment_variables)*.
-
-🔒 [FLUX.1-dev currently requires granted access to its Huggingface repo. For troubleshooting, see the issue tracker](https://github.com/filipstrand/mflux/issues/14) 🔒 
-
-#### Full list of Command-Line Arguments 
-
-- **`--prompt`** (required, `str`): Text description of the image to generate.
-
-- **`--model`** or **`-m`** (required, `str`): Model to use for generation (`"schnell"` or `"dev"`).
-
-- **`--output`** (optional, `str`, default: `"image.png"`): Output image filename.
-
-- **`--seed`** (optional, `int`, default: `None`): Seed for random number generation. Default is time-based.
-
-- **`--height`** (optional, `int`, default: `1024`): Height of the output image in pixels.
-
-- **`--width`** (optional, `int`, default: `1024`): Width of the output image in pixels.
-
-- **`--steps`** (optional, `int`, default: `4`): Number of inference steps.
-
-- **`--guidance`** (optional, `float`, default: `3.5`): Guidance scale (only used for `"dev"` model).
-
-- **`--path`** (optional, `str`, default: `None`): Path to a local model on disk.
-
-- **`--quantize`** or **`-q`** (optional, `int`, default: `None`): [Quantization](#quantization) (choose between `4` or `8`).
-
-- **`--lora-paths`** (optional, `[str]`, default: `None`): The paths to the [LoRA](#LoRA) weights.
-
-- **`--lora-scales`** (optional, `[float]`, default: `None`): The scale for each respective [LoRA](#LoRA) (will default to `1.0` if not specified and only one LoRA weight is loaded.)
- 
-- **`--metadata`** (optional): Exports a `.json` file containing the metadata for the image with the same name. (Even without this flag, the image metadata is saved and can be viewed using `exiftool image.png`)
-
-Or, with the correct python environment active, create and run a separate script like the following:
-
-```python
-from mflux.flux.flux import Flux1
-from mflux.config.config import Config
-
-# Load the model
-flux = Flux1.from_alias(
-   alias="schnell",  # "schnell" or "dev"
-   quantize=8,       # 4 or 8
-)  
-
-# Generate an image
-image = flux.generate_image(
-   seed=2,
-   prompt="Luxury food photograph",
-   config=Config(
-      num_inference_steps=2,  # "schnell" works well with 2-4 steps, "dev" works well with 20-25 steps
-      height=1024,
-      width=1024,
-   )
-)
-
-image.save(path="image.png")
-```
-
-For more options on how to configure MFLUX, please see [generate.py](src/mflux/generate.py).
-
-### Image generation speed (updated)
-
-These numbers are based on the non-quantized `schnell` model, with the configuration provided in the code snippet below. 
-To time your machine, run the following:
-```
-time mflux-generate \
---prompt "Luxury food photograph" \
---model schnell \
---steps 2 \
---seed 2 \
---height 1024 \
---width 1024
-```
-
-| Device             | User                                                                                                                        | Reported Time | Notes                     |
-|--------------------|-----------------------------------------------------------------------------------------------------------------------------|---------------|---------------------------|
-| M3 Max             | [@karpathy](https://gist.github.com/awni/a67d16d50f0f492d94a10418e0592bde?permalink_comment_id=5153531#gistcomment-5153531) | ~20s          |                           |
-| M2 Ultra           | [@awni](https://x.com/awnihannun/status/1823515121827897385)                                                                | <15s          |                           |
-| 2023 M2 Max (96GB) | [@explorigin](https://github.com/filipstrand/mflux/issues/6)                                                                | ~25s          |                           |
-| 2021 M1 Pro (16GB) | [@qw-in](https://github.com/filipstrand/mflux/issues/7)                                                                     | ~175s         | Might freeze your mac     |
-| 2023 M3 Pro (36GB) | [@kush-gupt](https://github.com/filipstrand/mflux/issues/11)                                                                | ~80s          |                           |
-| 2020 M1 (8GB)      | [@mbvillaverde](https://github.com/filipstrand/mflux/issues/13)                                                             | ~335s         | With resolution 512 x 512 |
-| 2022 M1 MAX (64GB) | [@BosseParra](https://x.com/BosseParra/status/1826191780812877968)                                                          | ~55s          |                           |
-| 2021 M1 Pro (32GB) | @filipstrand                                                                                                                | ~160s         |                           |
-| 2023 M2 Max (32GB) | @filipstrand                                                                                                                | ~70s          |                           |
-
-*Note that these numbers includes starting the application from scratch, which means doing model i/o, setting/quantizing weights etc. 
-If we assume that the model is already loaded, you can inspect the image metadata using `exiftool image.png` and see the total duration of the denoising loop (excluding text embedding).*
-
-### Equivalent to Diffusers implementation 
-
-There is only a single source of randomness when generating an image: The initial latent array. 
-In this implementation, this initial latent is fully deterministically controlled by the input `seed` parameter. 
-However, if we were to import a fixed instance of this latent array saved from the Diffusers implementation, then MFLUX will produce an identical image to the Diffusers implementation (assuming a fixed prompt and using the default parameter settings in the Diffusers setup).
-
-
-The images below illustrate this equivalence. 
-In all cases the Schnell model was run for 2 time steps. 
-The Diffusers implementation ran in CPU mode. 
-The precision for MFLUX can be set in the [Config](src/mflux/config/config.py) class. 
-There is typically a noticeable but very small difference in the final image when switching between 16bit and 32bit precision.
-
----
-```
-Luxury food photograph
-```
-![image](src/mflux/assets/comparison1.jpg)
-
----
-```
-detailed cinematic dof render of an old dusty detailed CRT monitor on a wooden desk in a dim room with items around, messy dirty room. On the screen are the letters "FLUX" glowing softly. High detail hard surface render
-```
-![image](src/mflux/assets/comparison2.jpg)
-
----
-
-```
-photorealistic, lotr, A tiny red dragon curled up asleep inside a nest, (Soft Focus) , (f_stop 2.8) , (focal_length 50mm) macro lens f/2. 8, medieval wizard table, (pastel) colors, (cozy) morning light filtering through a nearby window, (whimsical) steam shapes, captured with a (Canon EOS R5) , highlighting (serene) comfort, medieval, dnd, rpg, 3d, 16K, 8K
-```
-![image](src/mflux/assets/comparison3.jpg)
-
----
-
-
-```
-A weathered fisherman in his early 60s stands on the deck of his boat, gazing out at a stormy sea. He has a thick, salt-and-pepper beard, deep-set blue eyes, and skin tanned and creased from years of sun exposure. He's wearing a yellow raincoat and hat, with water droplets clinging to the fabric. Behind him, dark clouds loom ominously, and waves crash against the side of the boat. The overall atmosphere is one of tension and respect for the power of nature.
-```
-![image](src/mflux/assets/comparison4.jpg)
-
----
-
-```
-Luxury food photograph of an italian Linguine pasta alle vongole dish with lots of clams. It has perfect lighting and a cozy background with big bokeh and shallow depth of field. The mood is a sunset balcony in tuscany.  The photo is taken from the side of the plate. The pasta is shiny with sprinkled parmesan cheese and basil leaves on top. The scene is complemented by a warm, inviting light that highlights the textures and colors of the ingredients, giving it an appetizing and elegant look.
-```
-![image](src/mflux/assets/comparison5.jpg)
-
----
-
-### Quantization 
-
-MFLUX supports running FLUX in 4-bit or 8-bit quantized mode. Running a quantized version can greatly speed up the
-generation process and reduce the memory consumption by several gigabytes. [Quantized models also take up less disk space](#size-comparisons-for-quantized-models). 
-
-```
-mflux-generate \
-    --model schnell \
-    --steps 2 \
-    --seed 2 \
-    --quantize 8 \
-    --height 1920 \
-    --width 1024 \
-    --prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond’s calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier" 
-```
-![image](src/mflux/assets/comparison6.jpg)
-
-*In this example, weights are quantized at **runtime** - this is convenient if you don't want to [save a quantized copy of the weights to disk](#saving-a-quantized-version-to-disk), but still want to benefit from the potential speedup and RAM reduction quantization might bring.*
-
-
-By selecting the `--quantize` or `-q` flag to be `4`, `8`, or removing it entirely, we get all 3 images above. As can be seen, there is very little difference between the images (especially between the 8-bit, and the non-quantized result).
-Image generation times in this example are based on a 2021 M1 Pro (32GB) machine. Even though the images are almost identical, there is a ~2x speedup by
-running the 8-bit quantized version on this particular machine. Unlike the non-quantized version, for the 8-bit version the swap memory usage is drastically reduced and GPU utilization is close to 100% during the whole generation. Results here can vary across different machines.
-
-#### Size comparisons for quantized models 
-
-The model sizes for both `schnell` and `dev` at various quantization levels are as follows:
-
-| 4 bit  | 8 bit   | Original (16 bit) |
-|--------|---------|-------------------|
-| 9.85GB | 18.16GB | 33.73GB           |
-
-The reason weights sizes are not fully cut in half is because a small number of weights are not quantized and kept at full precision.
-
-#### Saving a quantized version to disk
-
-To save a local copy of the quantized weights, run the `mflux-save` command like so:
-
-```
-mflux-save \
-    --path "/Users/filipstrand/Desktop/schnell_8bit" \
-    --model schnell \
-    --quantize 8
-```
-
-*Note that when saving a quantized version, you will need the original huggingface weights.*
-
-#### Loading and running a quantized version from disk
-
-To generate a new image from the quantized model, simply provide a `--path` to where it was saved: 
-
-```
-mflux-generate \
-    --path "/Users/filipstrand/Desktop/schnell_8bit" \
-    --model schnell \
-    --steps 2 \
-    --seed 2 \
-    --height 1920 \
-    --width 1024 \
-    --prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond’s calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier" 
-```
-
-*Note: When loading a quantized model from disk, there is no need to pass in `-q` flag, since we can infer this from the weight metadata.*
-
-*Also Note: Once we have a local model (quantized [or not](#running-a-non-quantized-model-directly-from-disk)) specified via the `--path` argument, the huggingface cache models are not required to launch the model. 
-In other words, you can reclaim the 34GB diskspace (per model) by deleting the full 16-bit model from the [Huggingface cache](#generating-an-image) if you choose.* 
-
-*If you don't want to download the full models and quantize them yourself, the 4-bit weights are available here for a direct download:* 
-- [madroid/flux.1-schnell-mflux-4bit](https://huggingface.co/madroid/flux.1-schnell-mflux-4bit)
-- [madroid/flux.1-dev-mflux-4bit](https://huggingface.co/madroid/flux.1-dev-mflux-4bit)
-
-### Running a non-quantized model directly from disk
-
-MFLUX also supports running a non-quantized model directly from a custom location.  
-In the example below, the model is placed in `/Users/filipstrand/Desktop/schnell`:
-
-```
-mflux-generate \
-    --path "/Users/filipstrand/Desktop/schnell" \
-    --model schnell \
-    --steps 2 \
-    --seed 2 \
-    --prompt "Luxury food photograph" 
-```
-
-Note that the `--model` flag must be set when loading a model from disk. 
-
-Also note that unlike when using the typical `alias` way of initializing the model (which internally handles that the required resources are downloaded),
-when loading a model directly from disk, we require the downloaded models to look like the following:
-
-```
-.
-├── text_encoder
-│   └── model.safetensors
-├── text_encoder_2
-│   ├── model-00001-of-00002.safetensors
-│   └── model-00002-of-00002.safetensors
-├── tokenizer
-│   ├── merges.txt
-│   ├── special_tokens_map.json
-│   ├── tokenizer_config.json
-│   └── vocab.json
-├── tokenizer_2
-│   ├── special_tokens_map.json
-│   ├── spiece.model
-│   ├── tokenizer.json
-│   └── tokenizer_config.json
-├── transformer
-│   ├── diffusion_pytorch_model-00001-of-00003.safetensors
-│   ├── diffusion_pytorch_model-00002-of-00003.safetensors
-│   └── diffusion_pytorch_model-00003-of-00003.safetensors
-└── vae
-    └── diffusion_pytorch_model.safetensors
-```
-This mirrors how the resources are placed in the [HuggingFace Repo](https://huggingface.co/black-forest-labs/FLUX.1-schnell/tree/main) for FLUX.1.
-*Huggingface weights, unlike quantized ones exported directly from this project, have to be
-processed a bit differently, which is why we require this structure above.* 
-
-
-### LoRA
-
-MFLUX support loading trained [LoRA](https://huggingface.co/docs/diffusers/en/training/lora) adapters (actual training support is coming).
-
-The following example [The_Hound](https://huggingface.co/TheLastBen/The_Hound) LoRA from [@TheLastBen](https://github.com/TheLastBen): 
-
-```
-mflux-generate --prompt "sandor clegane" --model dev --steps 20 --seed 43 -q 8 --lora-paths "sandor_clegane_single_layer.safetensors"
-```
-
-![image](src/mflux/assets/lora1.jpg)
----
-
-The following example is [Flux_1_Dev_LoRA_Paper-Cutout-Style](https://huggingface.co/Norod78/Flux_1_Dev_LoRA_Paper-Cutout-Style) LoRA from [@Norod78](https://huggingface.co/Norod78):
-
-```
-mflux-generate --prompt "pikachu, Paper Cutout Style" --model schnell --steps 4 --seed 43 -q 8 --lora-paths "Flux_1_Dev_LoRA_Paper-Cutout-Style.safetensors"
-```
-![image](src/mflux/assets/lora2.jpg)
-
-*Note that LoRA trained weights are typically trained with a **trigger word or phrase**. For example, in the latter case, the sentence should include the phrase **"Paper Cutout Style"**.*
-
-*Also note that the same LoRA weights can work well with both the `schnell` and `dev` models. Refer to the original LoRA repository to see what mode it was trained for.*
-
-#### Multi-LoRA
-
-Multiple LoRAs can be sent in to combine the effects of the individual adapters. The following example combines both of the above LoRAs:
-
-```
-mflux-generate \
-   --prompt "sandor clegane in a forest, Paper Cutout Style" \
-   --model dev \
-   --steps 20 \
-   --seed 43 \
-   --lora-paths sandor_clegane_single_layer.safetensors Flux_1_Dev_LoRA_Paper-Cutout-Style.safetensors \
-   --lora-scales 1.0 1.0 \
-   -q 8
-```
-![image](src/mflux/assets/lora3.jpg)
-
-Just to see the difference, this image displays the four cases: One of having both adapters fully active, partially active and no LoRA at all. 
-The example above also show the usage of `--lora-scales` flag. 
-
-### Current limitations
-
-- Images are generated one by one.
-- Negative prompts not supported.
-- LoRA weights are only supported for the transformer part of the network.
-
-### TODO
-
-- [ ] LoRA fine-tuning
-- [ ] Frontend support (Gradio/Streamlit/Other?)
diff --git a/src/mflux.egg-info/SOURCES.txt b/src/mflux.egg-info/SOURCES.txt
deleted file mode 100644
index d3e63de..0000000
--- a/src/mflux.egg-info/SOURCES.txt
+++ /dev/null
@@ -1,89 +0,0 @@
-README.md
-pyproject.toml
-src/mflux/__init__.py
-src/mflux/generate.py
-src/mflux/generate_controlnet.py
-src/mflux/save.py
-src/mflux.egg-info/PKG-INFO
-src/mflux.egg-info/SOURCES.txt
-src/mflux.egg-info/dependency_links.txt
-src/mflux.egg-info/entry_points.txt
-src/mflux.egg-info/requires.txt
-src/mflux.egg-info/top_level.txt
-src/mflux/config/__init__.py
-src/mflux/config/config.py
-src/mflux/config/model_config.py
-src/mflux/config/runtime_config.py
-src/mflux/controlnet/flux_controlnet.py
-src/mflux/controlnet/utils_controlnet.py
-src/mflux/flux/__init__.py
-src/mflux/flux/flux.py
-src/mflux/models/__init__.py
-src/mflux/models/text_encoder/__init__.py
-src/mflux/models/text_encoder/clip_encoder/__init__.py
-src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py
-src/mflux/models/text_encoder/clip_encoder/clip_encoder.py
-src/mflux/models/text_encoder/clip_encoder/clip_encoder_layer.py
-src/mflux/models/text_encoder/clip_encoder/clip_mlp.py
-src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py
-src/mflux/models/text_encoder/clip_encoder/clip_text_model.py
-src/mflux/models/text_encoder/clip_encoder/encoder_clip.py
-src/mflux/models/text_encoder/t5_encoder/__init__.py
-src/mflux/models/text_encoder/t5_encoder/t5_attention.py
-src/mflux/models/text_encoder/t5_encoder/t5_block.py
-src/mflux/models/text_encoder/t5_encoder/t5_dense_relu_dense.py
-src/mflux/models/text_encoder/t5_encoder/t5_encoder.py
-src/mflux/models/text_encoder/t5_encoder/t5_feed_forward.py
-src/mflux/models/text_encoder/t5_encoder/t5_layer_norm.py
-src/mflux/models/text_encoder/t5_encoder/t5_self_attention.py
-src/mflux/models/transformer/__init__.py
-src/mflux/models/transformer/ada_layer_norm_continous.py
-src/mflux/models/transformer/ada_layer_norm_zero.py
-src/mflux/models/transformer/ada_layer_norm_zero_single.py
-src/mflux/models/transformer/embed_nd.py
-src/mflux/models/transformer/feed_forward.py
-src/mflux/models/transformer/guidance_embedder.py
-src/mflux/models/transformer/joint_attention.py
-src/mflux/models/transformer/joint_transformer_block.py
-src/mflux/models/transformer/single_block_attention.py
-src/mflux/models/transformer/single_transformer_block.py
-src/mflux/models/transformer/text_embedder.py
-src/mflux/models/transformer/time_text_embed.py
-src/mflux/models/transformer/timestep_embedder.py
-src/mflux/models/transformer/transformer.py
-src/mflux/models/vae/__init__.py
-src/mflux/models/vae/vae.py
-src/mflux/models/vae/common/__init__.py
-src/mflux/models/vae/common/attention.py
-src/mflux/models/vae/common/resnet_block_2d.py
-src/mflux/models/vae/common/unet_mid_block.py
-src/mflux/models/vae/decoder/__init__.py
-src/mflux/models/vae/decoder/conv_in.py
-src/mflux/models/vae/decoder/conv_norm_out.py
-src/mflux/models/vae/decoder/conv_out.py
-src/mflux/models/vae/decoder/decoder.py
-src/mflux/models/vae/decoder/up_block_1_or_2.py
-src/mflux/models/vae/decoder/up_block_3.py
-src/mflux/models/vae/decoder/up_block_4.py
-src/mflux/models/vae/decoder/up_sampler.py
-src/mflux/models/vae/encoder/__init__.py
-src/mflux/models/vae/encoder/conv_in.py
-src/mflux/models/vae/encoder/conv_norm_out.py
-src/mflux/models/vae/encoder/conv_out.py
-src/mflux/models/vae/encoder/down_block_1.py
-src/mflux/models/vae/encoder/down_block_2.py
-src/mflux/models/vae/encoder/down_block_3.py
-src/mflux/models/vae/encoder/down_block_4.py
-src/mflux/models/vae/encoder/down_sampler.py
-src/mflux/models/vae/encoder/encoder.py
-src/mflux/post_processing/__init__.py
-src/mflux/post_processing/image.py
-src/mflux/post_processing/image_util.py
-src/mflux/tokenizer/__init__.py
-src/mflux/tokenizer/clip_tokenizer.py
-src/mflux/tokenizer/t5_tokenizer.py
-src/mflux/tokenizer/tokenizer_handler.py
-src/mflux/weights/__init__.py
-src/mflux/weights/lora_util.py
-src/mflux/weights/weight_handler.py
-src/mflux/weights/weight_util.py
\ No newline at end of file
diff --git a/src/mflux.egg-info/dependency_links.txt b/src/mflux.egg-info/dependency_links.txt
deleted file mode 100644
index 8b13789..0000000
--- a/src/mflux.egg-info/dependency_links.txt
+++ /dev/null
@@ -1 +0,0 @@
-
diff --git a/src/mflux.egg-info/entry_points.txt b/src/mflux.egg-info/entry_points.txt
deleted file mode 100644
index e9824aa..0000000
--- a/src/mflux.egg-info/entry_points.txt
+++ /dev/null
@@ -1,4 +0,0 @@
-[console_scripts]
-mflux-generate = mflux.generate:main
-mflux-generate-controlnet = mflux.generate_controlnet:main
-mflux-save = mflux.save:main
diff --git a/src/mflux.egg-info/requires.txt b/src/mflux.egg-info/requires.txt
deleted file mode 100644
index 149534c..0000000
--- a/src/mflux.egg-info/requires.txt
+++ /dev/null
@@ -1,10 +0,0 @@
-mlx>=0.16.0
-numpy>=2.0.0
-pillow>=10.4.0
-transformers>=4.44.0
-sentencepiece>=0.2.0
-torch>=2.3.1
-tqdm>=4.66.5
-huggingface-hub>=0.24.5
-safetensors>=0.4.4
-piexif>=1.1.3
diff --git a/src/mflux.egg-info/top_level.txt b/src/mflux.egg-info/top_level.txt
deleted file mode 100644
index 4bd59a6..0000000
--- a/src/mflux.egg-info/top_level.txt
+++ /dev/null
@@ -1 +0,0 @@
-mflux

From ff1086c06f78a4d0959e8be0098fba9a573f86bf Mon Sep 17 00:00:00 2001
From: Fabio <f.xuzzo@gmail.com>
Date: Sat, 14 Sep 2024 20:02:06 +0200
Subject: [PATCH 05/16] add egg-info to gitignore

---
 .gitignore | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/.gitignore b/.gitignore
index 4f1076d..db83963 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,3 +13,5 @@
 *.pyc
 *.safetensors
 *.json
+
+*.egg-info
\ No newline at end of file

From a04d139ff5b56ba2eb26e3e481a26b69431a7ee5 Mon Sep 17 00:00:00 2001
From: Fabio <f.xuzzo@gmail.com>
Date: Sat, 14 Sep 2024 20:04:08 +0200
Subject: [PATCH 06/16] further cleanup

---
 src/mflux/config/config.py         |  2 +-
 src/mflux/config/runtime_config.py |  2 --
 trial.py                           | 38 ------------------------------
 3 files changed, 1 insertion(+), 41 deletions(-)
 delete mode 100644 trial.py

diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py
index 88566ae..123bb44 100644
--- a/src/mflux/config/config.py
+++ b/src/mflux/config/config.py
@@ -6,7 +6,7 @@ log = logging.getLogger(__name__)
 
 
 class Config:
-    precision: mx.Dtype = mx.float16
+    precision: mx.Dtype = mx.bfloat16
 
     def __init__(
             self,
diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py
index dabd8f5..96bd89b 100644
--- a/src/mflux/config/runtime_config.py
+++ b/src/mflux/config/runtime_config.py
@@ -1,5 +1,3 @@
-from math import e
-from cycler import V
 import mlx.core as mx
 import numpy as np
 
diff --git a/trial.py b/trial.py
deleted file mode 100644
index ee933b2..0000000
--- a/trial.py
+++ /dev/null
@@ -1,38 +0,0 @@
-import os
-import sys
-from mflux.post_processing.image_util import ImageUtil
-
-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
-
-prompt = "Luxury picture of food"
-
-
-# Load the model
-flux = Flux1Controlnet(
-    model_config=ModelConfig.from_alias("dev"),
-    quantize=8,
-    lora_paths=["Flux_1_Dev_LoRA_Paper-Cutout-Style.safetensors"]
-)
-
-control_image = ImageUtil.load_image("image_1.png")
-
-# Generate an image
-image = flux.generate_image(
-    seed=2500,
-    prompt=prompt,
-    control_image=control_image,
-    config=ConfigControlnet(
-        num_inference_steps=10,
-        height=256,
-        width=512,
-        guidance=3.5,
-        controlnet_strength=0.7,
-    )
-)
-
-# Save the image
-image.save(path="image.png", export_json_metadata=False)
\ No newline at end of file

From dee452b3bb6f4a70336eb7506aa50b4459bd9c17 Mon Sep 17 00:00:00 2001
From: Fabio <f.xuzzo@gmail.com>
Date: Sun, 15 Sep 2024 21:37:29 +0200
Subject: [PATCH 07/16] fixes from review

---
 pyproject.toml                              |  1 +
 requirements.txt                            |  3 ++-
 src/mflux/config/runtime_config.py          |  2 +-
 src/mflux/controlnet/flux_controlnet.py     | 17 +++++++----------
 src/mflux/generate_controlnet.py            |  5 ++++-
 src/mflux/models/transformer/transformer.py |  4 ++--
 src/mflux/weights/weight_handler.py         |  2 +-
 7 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 60b6838..17bba62 100644
--- a/pyproject.toml
+++ b/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]
diff --git a/requirements.txt b/requirements.txt
index b9a4c8f..f6d689e 100644
--- a/requirements.txt
+++ b/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
\ No newline at end of file
+piexif>=1.1.3
+opencv-python>=4.10.0
\ No newline at end of file
diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py
index 96bd89b..67a357a 100644
--- a/src/mflux/config/runtime_config.py
+++ b/src/mflux/config/runtime_config.py
@@ -41,7 +41,7 @@ class RuntimeConfig:
         if isinstance(self.config, ConfigControlnet):
             return self.config.controlnet_strength
         else:
-            return ValueError("Controlnet conditioning scale is only available for ConfigControlnet")
+            return NotImplementedError("Controlnet conditioning scale is only available for ConfigControlnet")
 
     @staticmethod
     def _create_sigmas(config, model) -> mx.array:
diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py
index 5a5a71b..a1aafac 100644
--- a/src/mflux/controlnet/flux_controlnet.py
+++ b/src/mflux/controlnet/flux_controlnet.py
@@ -7,7 +7,7 @@ from mlx import nn
 from mlx.utils import tree_flatten
 from tqdm import tqdm
 
-from mflux.config.config import Config, ConfigControlnet
+from mflux.config.config import ConfigControlnet
 from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.controlnet.utils_controlnet import preprocess_canny
@@ -24,19 +24,16 @@ from mflux.weights.weight_handler import WeightHandler
 
 from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
-from mflux.models.transformer.ada_layer_norm_continous import AdaLayerNormContinuous
 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
-import numpy as np
-import cv2
 
 import logging
 
 log = logging.getLogger(__name__)
 
-CONTOLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny"
+CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny"
 
 class Flux1Controlnet:
     def __init__(
@@ -88,7 +85,7 @@ class Flux1Controlnet:
         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=CONTOLNET_ID)
+        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"],
@@ -135,7 +132,7 @@ class Flux1Controlnet:
         pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens)
 
         for t in time_steps:
-            controlnet_block_samples = self.transformer_controlnet.forward(
+            ctrlnet_block_samples, ctrlnet_single_block_samples = self.transformer_controlnet.forward(
                 t=t,
                 prompt_embeds=prompt_embeds,
                 pooled_prompt_embeds=pooled_prompt_embeds,
@@ -150,8 +147,8 @@ class Flux1Controlnet:
                 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,
+                controlnet_block_samples=ctrlnet_block_samples,
+                controlnet_single_block_samples=ctrlnet_single_block_samples,
             )
 
             # 4.t Take one denoise step
@@ -318,4 +315,4 @@ class TransformerControlnet(nn.Module):
         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
\ No newline at end of file
+        return controlnet_block_samples, controlnet_single_block_samples
\ No newline at end of file
diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py
index a8b3c57..4c5bb1a 100644
--- a/src/mflux/generate_controlnet.py
+++ b/src/mflux/generate_controlnet.py
@@ -20,7 +20,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('--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)')
@@ -34,6 +34,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 = Flux1Controlnet(
         model_config=ModelConfig.from_alias(args.model),
diff --git a/src/mflux/models/transformer/transformer.py b/src/mflux/models/transformer/transformer.py
index 10af56e..e57d514 100644
--- a/src/mflux/models/transformer/transformer.py
+++ b/src/mflux/models/transformer/transformer.py
@@ -54,7 +54,7 @@ class Transformer(nn.Module):
                 text_embeddings=text_embeddings,
                 rotary_embeddings=image_rotary_emb
             )
-            if controlnet_block_samples is not None:
+            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]
@@ -67,7 +67,7 @@ class Transformer(nn.Module):
                 text_embeddings=text_embeddings,
                 rotary_embeddings=image_rotary_emb
             )
-            if controlnet_single_block_samples is not None:
+            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] :, ...] = (
diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py
index bc3a7fe..dbd25f1 100644
--- a/src/mflux/weights/weight_handler.py
+++ b/src/mflux/weights/weight_handler.py
@@ -90,7 +90,7 @@ class WeightHandler:
         return weights, quantization_level
     
     @staticmethod
-    def load_controlnet_transformer(controlnet_id: Path | None = None) -> (dict, int):
+    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")

From a2900c3b134e28b5f6cfb7978791b376fb2c009a Mon Sep 17 00:00:00 2001
From: Fabio <f.xuzzo@gmail.com>
Date: Sun, 15 Sep 2024 21:54:17 +0200
Subject: [PATCH 08/16] update generated image annotations

---
 src/mflux/config/config.py              | 2 +-
 src/mflux/config/runtime_config.py      | 2 +-
 src/mflux/controlnet/__init__.py        | 0
 src/mflux/controlnet/flux_controlnet.py | 2 +-
 src/mflux/post_processing/image.py      | 3 +++
 src/mflux/post_processing/image_util.py | 2 ++
 6 files changed, 8 insertions(+), 3 deletions(-)
 create mode 100644 src/mflux/controlnet/__init__.py

diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py
index 123bb44..1cc86b4 100644
--- a/src/mflux/config/config.py
+++ b/src/mflux/config/config.py
@@ -33,4 +33,4 @@ class ConfigControlnet(Config):
             controlnet_strength: float = 1.0,
     ):
         super().__init__(num_inference_steps, width, height, guidance)
-        self.controlnet_conditioning_scale = controlnet_strength
+        self.controlnet_strength = controlnet_strength
diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py
index 67a357a..20205bc 100644
--- a/src/mflux/config/runtime_config.py
+++ b/src/mflux/config/runtime_config.py
@@ -37,7 +37,7 @@ class RuntimeConfig:
         return self.model_config.num_train_steps
     
     @property
-    def controlnet_conditioning_scale(self) -> float:
+    def controlnet_strength(self) -> float:
         if isinstance(self.config, ConfigControlnet):
             return self.config.controlnet_strength
         else:
diff --git a/src/mflux/controlnet/__init__.py b/src/mflux/controlnet/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py
index a1aafac..2ed4c70 100644
--- a/src/mflux/controlnet/flux_controlnet.py
+++ b/src/mflux/controlnet/flux_controlnet.py
@@ -268,7 +268,7 @@ class TransformerControlnet(nn.Module):
         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_conditioning_scale
+        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)
diff --git a/src/mflux/post_processing/image.py b/src/mflux/post_processing/image.py
index fe69cd6..91aab40 100644
--- a/src/mflux/post_processing/image.py
+++ b/src/mflux/post_processing/image.py
@@ -26,6 +26,7 @@ class GeneratedImage:
             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 GeneratedImage:
         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 GeneratedImage:
             '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}",
         }
diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py
index 2ef5b5e..dc54e7a 100644
--- a/src/mflux/post_processing/image_util.py
+++ b/src/mflux/post_processing/image_util.py
@@ -3,6 +3,7 @@ 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 GeneratedImage
 
@@ -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

From 25e375d45debe167d9e997d338f37c8b291f0802 Mon Sep 17 00:00:00 2001
From: filipstrand <strand.filip@gmail.com>
Date: Tue, 17 Sep 2024 06:35:06 +0200
Subject: [PATCH 09/16] Small additions and refactoring

---
 src/mflux/config/runtime_config.py      |  2 +-
 src/mflux/controlnet/flux_controlnet.py | 15 ++++++---------
 src/mflux/generate.py                   |  5 ++++-
 3 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py
index 20205bc..efb33cf 100644
--- a/src/mflux/config/runtime_config.py
+++ b/src/mflux/config/runtime_config.py
@@ -41,7 +41,7 @@ class RuntimeConfig:
         if isinstance(self.config, ConfigControlnet):
             return self.config.controlnet_strength
         else:
-            return NotImplementedError("Controlnet conditioning scale is only available for ConfigControlnet")
+            raise NotImplementedError("Controlnet conditioning scale is only available for ConfigControlnet")
 
     @staticmethod
     def _create_sigmas(config, model) -> mx.array:
diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py
index 2ed4c70..41174ba 100644
--- a/src/mflux/controlnet/flux_controlnet.py
+++ b/src/mflux/controlnet/flux_controlnet.py
@@ -1,3 +1,4 @@
+import logging
 from pathlib import Path
 from typing import Tuple
 
@@ -13,6 +14,10 @@ from mflux.config.runtime_config import RuntimeConfig
 from mflux.controlnet.utils_controlnet import preprocess_canny
 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
@@ -22,19 +27,11 @@ from mflux.tokenizer.t5_tokenizer import TokenizerT5
 from mflux.tokenizer.tokenizer_handler import TokenizerHandler
 from mflux.weights.weight_handler import WeightHandler
 
-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
-
-import logging
-
 log = logging.getLogger(__name__)
 
 CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny"
 
+
 class Flux1Controlnet:
     def __init__(
             self,
diff --git a/src/mflux/generate.py b/src/mflux/generate.py
index 1b8e59f..9bf281d 100644
--- a/src/mflux/generate.py
+++ b/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),

From 78bbf3d8fb6affc3adc1f2a8649ba6008e32ac29 Mon Sep 17 00:00:00 2001
From: filipstrand <strand.filip@gmail.com>
Date: Tue, 17 Sep 2024 06:41:05 +0200
Subject: [PATCH 10/16] Fix formatting and some warnings

---
 src/mflux/controlnet/flux_controlnet.py     | 14 +++++++-------
 src/mflux/models/transformer/transformer.py | 13 ++++++-------
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py
index 41174ba..36e92bf 100644
--- a/src/mflux/controlnet/flux_controlnet.py
+++ b/src/mflux/controlnet/flux_controlnet.py
@@ -85,8 +85,8 @@ class Flux1Controlnet:
         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"],
+            num_blocks=controlnet_config["num_layers"],
+            num_single_blocks=controlnet_config["num_single_layers"],
         )
 
         if ctrlnet_quantization_level is None:
@@ -230,6 +230,7 @@ class Flux1Controlnet:
 
 ControlNetOutput = Tuple[list[mx.array], list[mx.array]]
 
+
 class TransformerControlnet(nn.Module):
 
     def __init__(
@@ -237,7 +238,7 @@ class TransformerControlnet(nn.Module):
             model_config: ModelConfig,
             num_blocks: int,
             num_single_blocks: int,
-        ):
+    ):
         super().__init__()
         self.pos_embed = EmbedND()
         self.x_embedder = nn.Linear(64, 3072)
@@ -270,8 +271,8 @@ class TransformerControlnet(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)
 
@@ -293,7 +294,6 @@ class TransformerControlnet(nn.Module):
             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(
@@ -312,4 +312,4 @@ class TransformerControlnet(nn.Module):
         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
\ No newline at end of file
+        return controlnet_block_samples, controlnet_single_block_samples
diff --git a/src/mflux/models/transformer/transformer.py b/src/mflux/models/transformer/transformer.py
index e57d514..6bbcc53 100644
--- a/src/mflux/models/transformer/transformer.py
+++ b/src/mflux/models/transformer/transformer.py
@@ -12,7 +12,6 @@ from mflux.models.transformer.single_transformer_block import SingleTransformerB
 from mflux.models.transformer.time_text_embed import TimeTextEmbed
 
 
-
 class Transformer(nn.Module):
 
     def __init__(self, model_config: ModelConfig):
@@ -33,8 +32,8 @@ class Transformer(nn.Module):
             pooled_prompt_embeds: mx.array,
             hidden_states: mx.array,
             config: RuntimeConfig,
-            controlnet_block_samples: Tuple[mx.array] | None = None,
-            controlnet_single_block_samples: Tuple[mx.array] | None = None,
+            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)
@@ -42,8 +41,8 @@ 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)
 
@@ -82,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))
@@ -93,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))

From ad6f4ca8afd742d8a1fdeb22131d2fc01cf8086f Mon Sep 17 00:00:00 2001
From: filipstrand <strand.filip@gmail.com>
Date: Tue, 17 Sep 2024 06:43:13 +0200
Subject: [PATCH 11/16] Create util class for controlnet

---
 src/mflux/controlnet/controlnet_util.py  | 14 ++++++++++++++
 src/mflux/controlnet/flux_controlnet.py  |  4 ++--
 src/mflux/controlnet/utils_controlnet.py | 10 ----------
 3 files changed, 16 insertions(+), 12 deletions(-)
 create mode 100644 src/mflux/controlnet/controlnet_util.py
 delete mode 100644 src/mflux/controlnet/utils_controlnet.py

diff --git a/src/mflux/controlnet/controlnet_util.py b/src/mflux/controlnet/controlnet_util.py
new file mode 100644
index 0000000..ab6a9c5
--- /dev/null
+++ b/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)
diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py
index 36e92bf..6d36e4a 100644
--- a/src/mflux/controlnet/flux_controlnet.py
+++ b/src/mflux/controlnet/flux_controlnet.py
@@ -11,7 +11,7 @@ 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.utils_controlnet import preprocess_canny
+from mflux.controlnet.controlnet_util import ControlnetUtil
 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
@@ -114,7 +114,7 @@ class Flux1Controlnet:
             shape=[1, (config.height // 16) * (config.width // 16), 64],
             key=mx.random.key(seed)
         )
-        control_image = preprocess_canny(control_image)
+        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 
diff --git a/src/mflux/controlnet/utils_controlnet.py b/src/mflux/controlnet/utils_controlnet.py
deleted file mode 100644
index 5de6b06..0000000
--- a/src/mflux/controlnet/utils_controlnet.py
+++ /dev/null
@@ -1,10 +0,0 @@
-import cv2
-import numpy as np
-import PIL
-
-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)

From fc8099e90829da4abcbce5bdc56a12b6fa243eb3 Mon Sep 17 00:00:00 2001
From: filipstrand <strand.filip@gmail.com>
Date: Tue, 17 Sep 2024 06:50:02 +0200
Subject: [PATCH 12/16] Move new controlnet module to its own file

---
 src/mflux/controlnet/flux_controlnet.py       | 91 +-----------------
 .../controlnet/transformer_controlnet.py      | 96 +++++++++++++++++++
 2 files changed, 97 insertions(+), 90 deletions(-)
 create mode 100644 src/mflux/controlnet/transformer_controlnet.py

diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py
index 6d36e4a..36cc546 100644
--- a/src/mflux/controlnet/flux_controlnet.py
+++ b/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
diff --git a/src/mflux/controlnet/transformer_controlnet.py b/src/mflux/controlnet/transformer_controlnet.py
new file mode 100644
index 0000000..9841047
--- /dev/null
+++ b/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

From bce18747ffc5868778999a4899d2dc8698c1ff0e Mon Sep 17 00:00:00 2001
From: filipstrand <strand.filip@gmail.com>
Date: Tue, 17 Sep 2024 06:53:56 +0200
Subject: [PATCH 13/16] Fix formatting

---
 src/mflux/controlnet/transformer_controlnet.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/mflux/controlnet/transformer_controlnet.py b/src/mflux/controlnet/transformer_controlnet.py
index 9841047..33bbe84 100644
--- a/src/mflux/controlnet/transformer_controlnet.py
+++ b/src/mflux/controlnet/transformer_controlnet.py
@@ -89,7 +89,7 @@ class TransformerControlnet(nn.Module):
             single_block_sample = controlnet_block(single_block_sample)
             controlnet_single_block_samples = controlnet_single_block_samples + (single_block_sample,)
 
-        # # scaling
+        # 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]
 

From a30dc4d5cc7078e39cf60b1382f918bfefd45f20 Mon Sep 17 00:00:00 2001
From: filipstrand <strand.filip@gmail.com>
Date: Tue, 17 Sep 2024 06:56:08 +0200
Subject: [PATCH 14/16] Rename image.py to generated_image.py

---
 src/mflux/controlnet/flux_controlnet.py                    | 2 +-
 src/mflux/flux/flux.py                                     | 2 +-
 src/mflux/post_processing/{image.py => generated_image.py} | 0
 src/mflux/post_processing/image_util.py                    | 2 +-
 4 files changed, 3 insertions(+), 3 deletions(-)
 rename src/mflux/post_processing/{image.py => generated_image.py} (100%)

diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py
index 36cc546..1b8c8c0 100644
--- a/src/mflux/controlnet/flux_controlnet.py
+++ b/src/mflux/controlnet/flux_controlnet.py
@@ -16,7 +16,7 @@ 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 GeneratedImage
+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
diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py
index e7d52c4..7cd6524 100644
--- a/src/mflux/flux/flux.py
+++ b/src/mflux/flux/flux.py
@@ -12,7 +12,7 @@ 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 GeneratedImage
+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
diff --git a/src/mflux/post_processing/image.py b/src/mflux/post_processing/generated_image.py
similarity index 100%
rename from src/mflux/post_processing/image.py
rename to src/mflux/post_processing/generated_image.py
diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py
index dc54e7a..9764bc2 100644
--- a/src/mflux/post_processing/image_util.py
+++ b/src/mflux/post_processing/image_util.py
@@ -5,7 +5,7 @@ import numpy as np
 
 from mflux.config.config import ConfigControlnet
 from mflux.config.runtime_config import RuntimeConfig
-from mflux.post_processing.image import GeneratedImage
+from mflux.post_processing.generated_image import GeneratedImage
 
 
 class ImageUtil:

From 38c747b68000e164ee88c7e4b948ef2b02fadbc2 Mon Sep 17 00:00:00 2001
From: filipstrand <strand.filip@gmail.com>
Date: Tue, 17 Sep 2024 07:13:59 +0200
Subject: [PATCH 15/16] Extract model saving logic to separate class

---
 src/mflux/controlnet/flux_controlnet.py     | 45 ++-----------------
 src/mflux/flux/flux.py                      | 42 ++---------------
 src/mflux/models/transformer/transformer.py |  4 +-
 src/mflux/weights/model_saver.py            | 50 +++++++++++++++++++++
 4 files changed, 59 insertions(+), 82 deletions(-)
 create mode 100644 src/mflux/weights/model_saver.py

diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py
index 1b8c8c0..4edf228 100644
--- a/src/mflux/controlnet/flux_controlnet.py
+++ b/src/mflux/controlnet/flux_controlnet.py
@@ -1,10 +1,8 @@
 import logging
-from pathlib import Path
 
 import PIL.Image
 import mlx.core as mx
 from mlx import nn
-from mlx.utils import tree_flatten
 from tqdm import tqdm
 
 from mflux.config.config import ConfigControlnet
@@ -21,6 +19,7 @@ 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__)
@@ -185,42 +184,6 @@ class Flux1Controlnet:
         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")
-        _save_weights(self.transformer_controlnet, "transformer_controlnet")
-        
-
+    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")
diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py
index 7cd6524..fd45ee9 100644
--- a/src/mflux/flux/flux.py
+++ b/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
@@ -17,6 +14,7 @@ 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
 
 
@@ -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)
diff --git a/src/mflux/models/transformer/transformer.py b/src/mflux/models/transformer/transformer.py
index 6bbcc53..cd32d83 100644
--- a/src/mflux/models/transformer/transformer.py
+++ b/src/mflux/models/transformer/transformer.py
@@ -1,7 +1,7 @@
-from typing import Tuple
+import math
+
 import mlx.core as mx
 from mlx import nn
-import math
 
 from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
diff --git a/src/mflux/weights/model_saver.py b/src/mflux/weights/model_saver.py
new file mode 100644
index 0000000..0e7d372
--- /dev/null
+++ b/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

From c74279d3948b7b6e4cb31d54453e9abe28c2ce93 Mon Sep 17 00:00:00 2001
From: Anthony Wu <462072+anthonywu@users.noreply.github.com>
Date: Thu, 19 Sep 2024 11:20:06 -0700
Subject: [PATCH 16/16] pyproject.toml to be authoritative source of deps, soft
 introduce ruff linter/formatter

---
 README.md        | 84 ++++++++++++++++++++++++++----------------------
 pyproject.toml   | 35 ++++++++++++--------
 requirements.txt | 14 ++------
 3 files changed, 69 insertions(+), 64 deletions(-)

diff --git a/README.md b/README.md
index 8fd15b0..ab7398b 100644
--- a/README.md
+++ b/README.md
@@ -8,13 +8,13 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
 
 ### Philosophy
 
-MFLUX is a line-by-line port of the FLUX implementation in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) library to [Apple MLX](https://github.com/ml-explore/mlx). 
+MFLUX is a line-by-line port of the FLUX implementation in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) library to [Apple MLX](https://github.com/ml-explore/mlx).
 MFLUX is purposefully kept minimal and explicit - Network architectures are hardcoded and no config files are used
-except for the tokenizers. The aim is to have a tiny codebase with the single purpose of expressing these models 
+except for the tokenizers. The aim is to have a tiny codebase with the single purpose of expressing these models
 (thereby avoiding too many abstractions). While MFLUX priorities readability over generality and performance, [it can still be quite fast](#image-generation-speed-updated), [and even faster quantized](#quantization).
 
-All models are implemented from scratch in MLX and only the tokenizers are used via the 
-[Huggingface Transformers](https://github.com/huggingface/transformers) library. Other than that, there are only minimal dependencies 
+All models are implemented from scratch in MLX and only the tokenizers are used via the
+[Huggingface Transformers](https://github.com/huggingface/transformers) library. Other than that, there are only minimal dependencies
 like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) for simple image post-processing.
 
 ### Models
@@ -40,12 +40,18 @@ This creates and activates a virtual environment in the `mflux` folder. After th
  ```
 2. Navigate to the project and set up a virtual environment:
  ```
- cd mflux && python3 -m venv .venv && source .venv/bin/activate
+ cd mflux
+ python3 -m venv .venv  # alternative: `uv venv`
+ source .venv/bin/activate
  ```
-3. Install the required dependencies:
+3. Install the required dependencies in [development mode](https://setuptools.pypa.io/en/latest/userguide/development_mode.html) as defined in `pyproject.toml`:
  ```
- pip install -r requirements.txt
+ pip install -e .
  ```
+4. Follow format and lint checks prior to submitting Pull Requests. Install the `ruff` tool as a common resource somewhere else in your system external to this project.
+  - [`uv tool install ruff`](https://github.com/astral-sh/uv) OR [`pipx install ruff`](https://github.com/pypa/pipx).
+  - [`ruff check` and `ruff check --fix`](https://github.com/astral-sh/ruff?tab=readme-ov-file#usage) any linter warnings generated from your PR diff.
+  - [`ruff format`](https://github.com/astral-sh/ruff?tab=readme-ov-file#usage) any added or modified files in your PR diff.
 </details>
 
 ### Generating an image
@@ -64,7 +70,7 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
 
 ⚠️ *If the specific model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell or Dev model respectively). See the [quantization](#quantization) section for running compressed versions of the model.* ⚠️
 
-*By default, model files are downloaded to the `.cache` folder within your home directory. For example, in my setup, the path looks like this:* 
+*By default, model files are downloaded to the `.cache` folder within your home directory. For example, in my setup, the path looks like this:*
 
 ```
 /Users/filipstrand/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-dev
@@ -72,9 +78,9 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
 
 *To change this default behavior, you can do so by modifying the `HF_HOME` environment variable. For more details on how to adjust this setting, please refer to the [Hugging Face documentation](https://huggingface.co/docs/huggingface_hub/en/package_reference/environment_variables)*.
 
-🔒 [FLUX.1-dev currently requires granted access to its Huggingface repo. For troubleshooting, see the issue tracker](https://github.com/filipstrand/mflux/issues/14) 🔒 
+🔒 [FLUX.1-dev currently requires granted access to its Huggingface repo. For troubleshooting, see the issue tracker](https://github.com/filipstrand/mflux/issues/14) 🔒
 
-#### Full list of Command-Line Arguments 
+#### Full list of Command-Line Arguments
 
 - **`--prompt`** (required, `str`): Text description of the image to generate.
 
@@ -99,7 +105,7 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
 - **`--lora-paths`** (optional, `[str]`, default: `None`): The paths to the [LoRA](#LoRA) weights.
 
 - **`--lora-scales`** (optional, `[float]`, default: `None`): The scale for each respective [LoRA](#LoRA) (will default to `1.0` if not specified and only one LoRA weight is loaded.)
- 
+
 - **`--metadata`** (optional): Exports a `.json` file containing the metadata for the image with the same name. (Even without this flag, the image metadata is saved and can be viewed using `exiftool image.png`)
 
 Or, with the correct python environment active, create and run a separate script like the following:
@@ -112,7 +118,7 @@ from mflux.config.config import Config
 flux = Flux1.from_alias(
    alias="schnell",  # "schnell" or "dev"
    quantize=8,       # 4 or 8
-)  
+)
 
 # Generate an image
 image = flux.generate_image(
@@ -132,7 +138,7 @@ For more options on how to configure MFLUX, please see [generate.py](src/mflux/g
 
 ### Image generation speed (updated)
 
-These numbers are based on the non-quantized `schnell` model, with the configuration provided in the code snippet below. 
+These numbers are based on the non-quantized `schnell` model, with the configuration provided in the code snippet below.
 To time your machine, run the following:
 ```
 time mflux-generate \
@@ -156,20 +162,20 @@ time mflux-generate \
 | 2021 M1 Pro (32GB) | @filipstrand                                                                                                                | ~160s         |                           |
 | 2023 M2 Max (32GB) | @filipstrand                                                                                                                | ~70s          |                           |
 
-*Note that these numbers includes starting the application from scratch, which means doing model i/o, setting/quantizing weights etc. 
+*Note that these numbers includes starting the application from scratch, which means doing model i/o, setting/quantizing weights etc.
 If we assume that the model is already loaded, you can inspect the image metadata using `exiftool image.png` and see the total duration of the denoising loop (excluding text embedding).*
 
-### Equivalent to Diffusers implementation 
+### Equivalent to Diffusers implementation
 
-There is only a single source of randomness when generating an image: The initial latent array. 
-In this implementation, this initial latent is fully deterministically controlled by the input `seed` parameter. 
+There is only a single source of randomness when generating an image: The initial latent array.
+In this implementation, this initial latent is fully deterministically controlled by the input `seed` parameter.
 However, if we were to import a fixed instance of this latent array saved from the Diffusers implementation, then MFLUX will produce an identical image to the Diffusers implementation (assuming a fixed prompt and using the default parameter settings in the Diffusers setup).
 
 
-The images below illustrate this equivalence. 
-In all cases the Schnell model was run for 2 time steps. 
-The Diffusers implementation ran in CPU mode. 
-The precision for MFLUX can be set in the [Config](src/mflux/config/config.py) class. 
+The images below illustrate this equivalence.
+In all cases the Schnell model was run for 2 time steps.
+The Diffusers implementation ran in CPU mode.
+The precision for MFLUX can be set in the [Config](src/mflux/config/config.py) class.
 There is typically a noticeable but very small difference in the final image when switching between 16bit and 32bit precision.
 
 ---
@@ -208,10 +214,10 @@ Luxury food photograph of an italian Linguine pasta alle vongole dish with lots
 
 ---
 
-### Quantization 
+### Quantization
 
 MFLUX supports running FLUX in 4-bit or 8-bit quantized mode. Running a quantized version can greatly speed up the
-generation process and reduce the memory consumption by several gigabytes. [Quantized models also take up less disk space](#size-comparisons-for-quantized-models). 
+generation process and reduce the memory consumption by several gigabytes. [Quantized models also take up less disk space](#size-comparisons-for-quantized-models).
 
 ```
 mflux-generate \
@@ -221,7 +227,7 @@ mflux-generate \
     --quantize 8 \
     --height 1920 \
     --width 1024 \
-    --prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond’s calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier" 
+    --prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond’s calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier"
 ```
 ![image](src/mflux/assets/comparison6.jpg)
 
@@ -232,7 +238,7 @@ By selecting the `--quantize` or `-q` flag to be `4`, `8`, or removing it entire
 Image generation times in this example are based on a 2021 M1 Pro (32GB) machine. Even though the images are almost identical, there is a ~2x speedup by
 running the 8-bit quantized version on this particular machine. Unlike the non-quantized version, for the 8-bit version the swap memory usage is drastically reduced and GPU utilization is close to 100% during the whole generation. Results here can vary across different machines.
 
-#### Size comparisons for quantized models 
+#### Size comparisons for quantized models
 
 The model sizes for both `schnell` and `dev` at various quantization levels are as follows:
 
@@ -257,7 +263,7 @@ mflux-save \
 
 #### Loading and running a quantized version from disk
 
-To generate a new image from the quantized model, simply provide a `--path` to where it was saved: 
+To generate a new image from the quantized model, simply provide a `--path` to where it was saved:
 
 ```
 mflux-generate \
@@ -267,21 +273,21 @@ mflux-generate \
     --seed 2 \
     --height 1920 \
     --width 1024 \
-    --prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond’s calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier" 
+    --prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond’s calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier"
 ```
 
 *Note: When loading a quantized model from disk, there is no need to pass in `-q` flag, since we can infer this from the weight metadata.*
 
-*Also Note: Once we have a local model (quantized [or not](#running-a-non-quantized-model-directly-from-disk)) specified via the `--path` argument, the huggingface cache models are not required to launch the model. 
-In other words, you can reclaim the 34GB diskspace (per model) by deleting the full 16-bit model from the [Huggingface cache](#generating-an-image) if you choose.* 
+*Also Note: Once we have a local model (quantized [or not](#running-a-non-quantized-model-directly-from-disk)) specified via the `--path` argument, the huggingface cache models are not required to launch the model.
+In other words, you can reclaim the 34GB diskspace (per model) by deleting the full 16-bit model from the [Huggingface cache](#generating-an-image) if you choose.*
 
-*If you don't want to download the full models and quantize them yourself, the 4-bit weights are available here for a direct download:* 
+*If you don't want to download the full models and quantize them yourself, the 4-bit weights are available here for a direct download:*
 - [madroid/flux.1-schnell-mflux-4bit](https://huggingface.co/madroid/flux.1-schnell-mflux-4bit)
 - [madroid/flux.1-dev-mflux-4bit](https://huggingface.co/madroid/flux.1-dev-mflux-4bit)
 
 ### Running a non-quantized model directly from disk
 
-MFLUX also supports running a non-quantized model directly from a custom location.  
+MFLUX also supports running a non-quantized model directly from a custom location.
 In the example below, the model is placed in `/Users/filipstrand/Desktop/schnell`:
 
 ```
@@ -290,10 +296,10 @@ mflux-generate \
     --model schnell \
     --steps 2 \
     --seed 2 \
-    --prompt "Luxury food photograph" 
+    --prompt "Luxury food photograph"
 ```
 
-Note that the `--model` flag must be set when loading a model from disk. 
+Note that the `--model` flag must be set when loading a model from disk.
 
 Also note that unlike when using the typical `alias` way of initializing the model (which internally handles that the required resources are downloaded),
 when loading a model directly from disk, we require the downloaded models to look like the following:
@@ -324,14 +330,14 @@ when loading a model directly from disk, we require the downloaded models to loo
 ```
 This mirrors how the resources are placed in the [HuggingFace Repo](https://huggingface.co/black-forest-labs/FLUX.1-schnell/tree/main) for FLUX.1.
 *Huggingface weights, unlike quantized ones exported directly from this project, have to be
-processed a bit differently, which is why we require this structure above.* 
+processed a bit differently, which is why we require this structure above.*
 
 
 ### LoRA
 
 MFLUX support loading trained [LoRA](https://huggingface.co/docs/diffusers/en/training/lora) adapters (actual training support is coming).
 
-The following example [The_Hound](https://huggingface.co/TheLastBen/The_Hound) LoRA from [@TheLastBen](https://github.com/TheLastBen): 
+The following example [The_Hound](https://huggingface.co/TheLastBen/The_Hound) LoRA from [@TheLastBen](https://github.com/TheLastBen):
 
 ```
 mflux-generate --prompt "sandor clegane" --model dev --steps 20 --seed 43 -q 8 --lora-paths "sandor_clegane_single_layer.safetensors"
@@ -367,12 +373,12 @@ mflux-generate \
 ```
 ![image](src/mflux/assets/lora3.jpg)
 
-Just to see the difference, this image displays the four cases: One of having both adapters fully active, partially active and no LoRA at all. 
-The example above also show the usage of `--lora-scales` flag. 
+Just to see the difference, this image displays the four cases: One of having both adapters fully active, partially active and no LoRA at all.
+The example above also show the usage of `--lora-scales` flag.
 
 #### Supported LoRA formats (updated)
 
-Since different fine-tuning services can use different implementations of FLUX, the corresponding 
+Since different fine-tuning services can use different implementations of FLUX, the corresponding
 LoRA weights trained on these services can be different from one another. The aim of MFLUX is to support the most common ones.
 The following table show the current supported formats:
 
@@ -394,4 +400,4 @@ To report additional formats, examples or other any suggestions related to LoRA
 ### TODO
 
 - [ ] LoRA fine-tuning
-- [ ] Frontend support (Gradio/Streamlit/Other?)
\ No newline at end of file
+- [ ] Frontend support (Gradio/Streamlit/Other?)
diff --git a/pyproject.toml b/pyproject.toml
index 17bba62..6e9e1d6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,25 +7,32 @@ name = "mflux"
 version = "0.2.1"
 description = "A MLX port of FLUX based on the Huggingface Diffusers implementation."
 readme = "README.md"
-authors = [
-    { name = "Filip Strand", email = "strand.filip@gmail.com" }
-]
-classifiers = [
-    "Programming Language :: Python :: 3",
-    "Operating System :: MacOS",
-]
+keywords = ["diffusers", "flux", "mlx"]
+authors = [{ name = "Filip Strand", email = "strand.filip@gmail.com" }]
+maintainers = [{ name = "Filip Strand", email = "strand.filip@gmail.com" }]
+requires-python = ">=3.10"
 dependencies = [
+    "huggingface-hub>=0.24.5",
     "mlx>=0.16.0",
-    "numpy>=2.0.0",
+    "numpy>=2.0.1",
+    "opencv-python>=4.10.0",
+    "piexif>=1.1.3",
     "pillow>=10.4.0",
-    "transformers>=4.44.0",
+    "safetensors>=0.4.4",
     "sentencepiece>=0.2.0",
     "torch>=2.3.1",
     "tqdm>=4.66.5",
-    "huggingface-hub>=0.24.5",
-    "safetensors>=0.4.4",
-    "piexif>=1.1.3",
-    "opencv-python>=4.10.0",
+    "transformers>=4.44.0",
+]
+classifiers = [
+    "Framework :: MLX",
+    "Intended Audience :: Developers",
+    "Operating System :: MacOS",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
 ]
 
 [project.urls]
@@ -38,4 +45,4 @@ mflux-generate-controlnet = "mflux.generate_controlnet:main"
 
 [tool.setuptools.packages.find]
 where = ["src"]
-include = ["mflux*"]
\ No newline at end of file
+include = ["mflux*"]
diff --git a/requirements.txt b/requirements.txt
index f6d689e..0d16426 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,11 +1,3 @@
-mlx>=0.16.0
-numpy>=2.0.1
-pillow>=10.4.0
-transformers>=4.44.0
-sentencepiece>=0.2.0
-torch>=2.3.1
-tqdm>=4.66.5
-huggingface-hub>=0.24.5
-safetensors>=0.4.4
-piexif>=1.1.3
-opencv-python>=4.10.0
\ No newline at end of file
+# direct habitual `pip install -r requirements.txt`
+# users to `pyproject.toml` defined req list
+-e .