diff --git a/README.md b/README.md
index a88006f..c5fb192 100644
--- a/README.md
+++ b/README.md
@@ -2,30 +2,55 @@
 ![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!
 
+### Table of contents
+
+<!-- TOC start (generated with https://github.com/derlin/bitdowntoc) -->
+
+- [Philosophy](#philosophy)
+- [💿 Installation](#-installation)
+- [🖼️ Generating an image](#%EF%B8%8F-generating-an-image)
+  * [📜 Full list of Command-Line Arguments](#-full-list-of-command-line-arguments)
+- [⏱️ Image generation speed (updated)](#%EF%B8%8F-image-generation-speed-updated)
+- [↔️ Equivalent to Diffusers implementation](#%EF%B8%8F-equivalent-to-diffusers-implementation)
+- [🗜️ Quantization](#%EF%B8%8F-quantization)
+  * [📊 Size comparisons for quantized models](#-size-comparisons-for-quantized-models)
+  * [💾 Saving a quantized version to disk](#-saving-a-quantized-version-to-disk)
+  * [💽 Loading and running a quantized version from disk](#-loading-and-running-a-quantized-version-from-disk)
+- [💽 Running a non-quantized model directly from disk](#-running-a-non-quantized-model-directly-from-disk)
+- [🔌 LoRA](#-lora)
+  * [Multi-LoRA](#multi-lora)
+  * [Supported LoRA formats (updated)](#supported-lora-formats-updated)
+- [🕹️ Controlnet](#%EF%B8%8F-controlnet)
+- [🚧 Current limitations](#-current-limitations)
+- [✅ TODO](#-todo)
+
+<!-- TOC end -->
+
 ### 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).
+(thereby avoiding too many abstractions). While MFLUX priorities readability over generality and performance, [it can still be quite fast](#%EF%B8%8F-image-generation-speed-updated), [and even faster quantized](#%EF%B8%8F-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 use `uv tool`: If you have [installed `uv`](https://github.com/astral-sh/uv?tab=readme-ov-file#installation), simply: 
 
-### Installation
-For users, the easiest way to install MFLUX is to use `uv tool`:
+```sh
+uv tool install mflux
+``` 
 
-If you have [installed `uv`](https://github.com/astral-sh/uv?tab=readme-ov-file#installation), simply: `uv tool install mflux` to get the `mflux-generate` and related command line executables. You can skip to the usage guides below.
+to get the `mflux-generate` and related command line executables. You can skip to the usage guides below.
 
 <details>
 <summary>For the classic way to create a user virtual environment:</summary>
@@ -46,11 +71,19 @@ pip install -U mflux
 <summary>For contributors (click to expand)</summary>
 
 1. Clone the repo:
- ```sh
+```sh
  git clone git@github.com:filipstrand/mflux.git
  ```
-2. `make install` and `make test` (and `make clean` for venv resets)
-3. Follow format and lint checks prior to submitting Pull Requests. The recommended `make lint` and `make format` installs and uses [`ruff`](https://github.com/astral-sh/ruff). You can setup your editor/IDE to lint/format automatically, or use our provided `make` helpers:
+2. Install the application
+
+```sh
+ make install
+ ```
+3. To run the test suite 
+```sh
+ make test
+ ```
+4. Follow format and lint checks prior to submitting Pull Requests. The recommended `make lint` and `make format` installs and uses [`ruff`](https://github.com/astral-sh/ruff). You can setup your editor/IDE to lint/format automatically, or use our provided `make` helpers:
   - `make format` - formats your code
   - `make lint` - shows your lint errors and warnings, but does not auto fix
   - `make check` - via `pre-commit` hooks, formats your code **and** attempts to auto fix lint errors
@@ -58,7 +91,7 @@ pip install -U mflux
 
 </details>
 
-### Generating an image
+### 🖼️ 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:
 
@@ -72,7 +105,7 @@ 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.* ⚠️
+⚠️ *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](#%EF%B8%8F-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:*
 
@@ -84,7 +117,7 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
 
 🔒 [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.
 
@@ -104,14 +137,20 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
 
 - **`--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`).
+- **`--quantize`** or **`-q`** (optional, `int`, default: `None`): [Quantization](#%EF%B8%8F-quantization) (choose between `4` or `8`).
 
-- **`--lora-paths`** (optional, `[str]`, default: `None`): The paths to the [LoRA](#LoRA) weights.
+- **`--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.)
+- **`--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`)
 
+- **`--controlnet-image-path`** (required, `str`): Path to the local image used by ControlNet to guide output generation.
+
+- **`--controlnet-strength`** (optional, `float`, default: `0.4`): Degree of influence the control image has on the output. Ranges from `0.0` (no influence) to `1.0` (full influence).
+
+- **`--controlnet-save-canny`** (optional, bool, default: False): If set, saves the Canny edge detection reference image used by ControlNet.
+
 Or, with the correct python environment active, create and run a separate script like the following:
 
 ```python
@@ -139,7 +178,7 @@ 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)
+### ⏱️ 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:
@@ -168,7 +207,7 @@ time mflux-generate \
 *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.
@@ -217,10 +256,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).
 
 ```sh
 mflux-generate \
@@ -234,14 +273,14 @@ mflux-generate \
 ```
 ![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.*
+*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
+#### 📊 Size comparisons for quantized models
 
 The model sizes for both `schnell` and `dev` at various quantization levels are as follows:
 
@@ -251,7 +290,7 @@ The model sizes for both `schnell` and `dev` at various quantization levels are
 
 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
+#### 💾 Saving a quantized version to disk
 
 To save a local copy of the quantized weights, run the `mflux-save` command like so:
 
@@ -264,7 +303,7 @@ mflux-save \
 
 *Note that when saving a quantized version, you will need the original huggingface weights.*
 
-It is also possible to specify [LoRA](#lora) adapters when saving the model, e.g 
+It is also possible to specify [LoRA](#-lora) adapters when saving the model, e.g 
 
 ```sh
 mflux-save \
@@ -278,7 +317,7 @@ mflux-save \
 When generating images with a model like this, no LoRA adapter is needed to be specified since
 it is already baked into the saved quantized weights.
 
-#### Loading and running a quantized version from disk
+#### 💽 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:
 
@@ -295,14 +334,14 @@ mflux-generate \
 
 *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](#%EF%B8%8F-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
+### 💽 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`:
@@ -349,8 +388,9 @@ This mirrors how the resources are placed in the [HuggingFace Repo](https://hugg
 *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
+### 🔌 LoRA
 
 MFLUX support loading trained [LoRA](https://huggingface.co/docs/diffusers/en/training/lora) adapters (actual training support is coming).
 
@@ -407,14 +447,53 @@ The following table show the current supported formats:
 
 To report additional formats, examples or other any suggestions related to LoRA format support, please see [issue #47](https://github.com/filipstrand/mflux/issues/47).
 
-### Current limitations
+---
+
+### 🕹️ Controlnet
+
+MFLUX has [Controlnet](https://huggingface.co/docs/diffusers/en/using-diffusers/controlnet) support for an even more fine-grained control
+of the image generation. By providing a reference image via `--controlnet-image-path` and a strength parameter via `--controlnet-strength`, you can guide the generation toward the reference image.  
+
+```sh
+mflux-generate-controlnet \
+  --prompt "A comic strip with a joker in a purple suit" \
+  --model dev \
+  --steps 20 \
+  --seed 1727047657 \
+  --height 1066 \
+  --width 692 \
+  -q 8 \
+  --lora-paths "Dark Comic - s0_8 g4.safetensors" \
+  --controlnet-image-path "reference.png" \
+  --controlnet-strength 0.5 \
+  --controlnet-save-canny
+```
+![image](src/mflux/assets/controlnet1.jpg)
+
+*This example combines the controlnet reference image with the LoRA [Dark Comic Flux](https://civitai.com/models/742916/dark-comic-flux)*.
+
+⚠️ *Note: Controlnet requires an additional one-time download of ~3.58GB of weights from Huggingface. This happens automatically the first time you run the `generate-controlnet` command.
+At the moment, the Controlnet used is [InstantX/FLUX.1-dev-Controlnet-Canny](https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Canny), which was trained for the `dev` model. 
+It can work well with `schnell`, but performance is not guaranteed.* 
+
+⚠️ *Note: The output can be highly sensitive to the controlnet strength and is very much dependent on the reference image. 
+Too high settings will corrupt the image. A recommended starting point a value like 0.4 and to play around with the strength.*
+
+
+Controlnet can also work well together with [LoRA adapters](#-lora). In the example below the same reference image is used as a controlnet input
+with different prompts and LoRA adapters active.
+
+![image](src/mflux/assets/controlnet2.jpg)
+
+### 🚧 Current limitations
 
 - Images are generated one by one.
 - Negative prompts not supported.
 - LoRA weights are only supported for the transformer part of the network.
 - Some LoRA adapters does not work.
+- Currently, the supported controlnet is the [canny-only version](https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Canny). 
 
-### TODO
+### ✅ TODO
 
 - [ ] Establish unit test suite
 - [ ] LoRA fine-tuning
diff --git a/src/mflux/assets/controlnet1.jpg b/src/mflux/assets/controlnet1.jpg
new file mode 100644
index 0000000..9c40415
Binary files /dev/null and b/src/mflux/assets/controlnet1.jpg differ
diff --git a/src/mflux/assets/controlnet2.jpg b/src/mflux/assets/controlnet2.jpg
new file mode 100644
index 0000000..8464d2b
Binary files /dev/null and b/src/mflux/assets/controlnet2.jpg differ
diff --git a/src/mflux/controlnet/controlnet_util.py b/src/mflux/controlnet/controlnet_util.py
index 6e38c75..7f083cc 100644
--- a/src/mflux/controlnet/controlnet_util.py
+++ b/src/mflux/controlnet/controlnet_util.py
@@ -1,6 +1,12 @@
+import logging
+import os
+
 import cv2
 import numpy as np
-import PIL
+import PIL.Image
+
+
+log = logging.getLogger(__name__)
 
 
 class ControlnetUtil:
@@ -11,3 +17,18 @@ class ControlnetUtil:
         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)
+
+    @staticmethod
+    def scale_image(height: int, width: int, img: PIL.Image) -> PIL.Image:
+        if height != img.height or width != img.width:
+            log.warning(f"Control image has different dimensions than the model. Resizing to {width}x{height}")
+            img = img.resize((width, height), PIL.Image.LANCZOS)
+        return img
+
+    @staticmethod
+    def save_canny_image(control_image: PIL.Image, path: str):
+        from mflux import ImageUtil
+
+        base, ext = os.path.splitext(path)
+        new_filename = f"{base}_controlnet_canny{ext}"
+        ImageUtil.save_image(control_image, new_filename)
diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py
index c9e5b1c..fb68368 100644
--- a/src/mflux/controlnet/flux_controlnet.py
+++ b/src/mflux/controlnet/flux_controlnet.py
@@ -1,6 +1,6 @@
 import logging
+from typing import TYPE_CHECKING
 
-import PIL.Image
 import mlx.core as mx
 from mlx import nn
 from tqdm import tqdm
@@ -10,11 +10,11 @@ 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.controlnet.weight_handler_controlnet import WeightHandlerControlnet
 from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
 from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
 from mflux.models.transformer.transformer import Transformer
 from mflux.models.vae.vae import VAE
-from mflux.post_processing.generated_image import GeneratedImage
 from mflux.post_processing.image_util import ImageUtil
 from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
 from mflux.tokenizer.t5_tokenizer import TokenizerT5
@@ -22,6 +22,10 @@ from mflux.tokenizer.tokenizer_handler import TokenizerHandler
 from mflux.weights.model_saver import ModelSaver
 from mflux.weights.weight_handler import WeightHandler
 
+if TYPE_CHECKING:
+    from mflux.post_processing.generated_image import GeneratedImage
+
+
 log = logging.getLogger(__name__)
 
 CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny"
@@ -79,51 +83,55 @@ 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(
+        weights_controlnet, controlnet_quantization_level, controlnet_config = WeightHandlerControlnet.load_controlnet_transformer(
             controlnet_id=CONTROLNET_ID
-        )
+        )  # fmt: off
         self.transformer_controlnet = TransformerControlnet(
             model_config=model_config,
             num_blocks=controlnet_config["num_layers"],
             num_single_blocks=controlnet_config["num_single_layers"],
         )
 
-        if ctrlnet_quantization_level is None:
+        if controlnet_quantization_level is None:
             self.transformer_controlnet.update(weights_controlnet)
 
         self.bits = None
-        if quantize is not None or ctrlnet_quantization_level is not None:
-            self.bits = ctrlnet_quantization_level if ctrlnet_quantization_level is not None else quantize
-            # fmt: off
-            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)
-            # fmt: on
+        if quantize is not None or controlnet_quantization_level is not None:
+            self.bits = controlnet_quantization_level if 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]) > 128, group_size=128, bits=self.bits)  # fmt: off
 
-        if ctrlnet_quantization_level is not None:
+        if 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:  # fmt: off
+    def generate_image(
+            self,
+            seed: int,
+            prompt: str,
+            output: str,
+            controlnet_image_path: str,
+            controlnet_save_canny: bool = False,
+            config: ConfigControlnet = ConfigControlnet()
+    ) -> "GeneratedImage":  # fmt: off
         # Create a new runtime config based on the model type and input parameters
         config = RuntimeConfig(config, self.model_config)
         time_steps = tqdm(range(config.num_inference_steps))
 
-        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)
+        # Embedd the controlnet reference image
+        control_image = ImageUtil.load_image(controlnet_image_path)
+        control_image = ControlnetUtil.scale_image(config.height, config.width, control_image)
+        control_image = ControlnetUtil.preprocess_canny(control_image)
+        if controlnet_save_canny:
+            ControlnetUtil.save_canny_image(control_image, output)
+        controlnet_cond = ImageUtil.to_array(control_image)
+        controlnet_cond = self.vae.encode(controlnet_cond)
+        controlnet_cond = (controlnet_cond / self.vae.scaling_factor) + self.vae.shift_factor
+        controlnet_cond = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width)
 
         # 1. Create the initial latents
         latents = mx.random.normal(
             shape=[1, (config.height // 16) * (config.width // 16), 64],
             key=mx.random.key(seed)
         )  # fmt: off
-        control_image = ControlnetUtil.preprocess_canny(control_image)
-        controlnet_cond = ImageUtil.to_array(control_image)
-        controlnet_cong = self.vae.encode(controlnet_cond)
-        # the rescaling in the next line is not in the huggingface code, but without it the images from
-        # the chosen controlnet model are very bad
-        controlnet_cond = (controlnet_cong / self.vae.scaling_factor) + self.vae.shift_factor
-        controlnet_cond = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width)
 
         # 2. Embedd the prompt
         t5_tokens = self.t5_tokenizer.tokenize(prompt)
@@ -132,7 +140,8 @@ class Flux1Controlnet:
         pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens)
 
         for t in time_steps:
-            ctrlnet_block_samples, ctrlnet_single_block_samples = self.transformer_controlnet.forward(
+            # Compute controlnet samples
+            controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet.forward(
                 t=t,
                 prompt_embeds=prompt_embeds,
                 pooled_prompt_embeds=pooled_prompt_embeds,
@@ -140,6 +149,7 @@ class Flux1Controlnet:
                 controlnet_cond=controlnet_cond,
                 config=config,
             )
+
             # 3.t Predict the noise
             noise = self.transformer.predict(
                 t=t,
@@ -147,8 +157,8 @@ class Flux1Controlnet:
                 pooled_prompt_embeds=pooled_prompt_embeds,
                 hidden_states=latents,
                 config=config,
-                controlnet_block_samples=ctrlnet_block_samples,
-                controlnet_single_block_samples=ctrlnet_single_block_samples,
+                controlnet_block_samples=controlnet_block_samples,
+                controlnet_single_block_samples=controlnet_single_block_samples,
             )
 
             # 4.t Take one denoise step
@@ -170,6 +180,7 @@ class Flux1Controlnet:
             lora_paths=self.lora_paths,
             lora_scales=self.lora_scales,
             config=config,
+            controlnet_image_path=controlnet_image_path,
         )
 
     @staticmethod
diff --git a/src/mflux/controlnet/weight_handler_controlnet.py b/src/mflux/controlnet/weight_handler_controlnet.py
new file mode 100644
index 0000000..00654c9
--- /dev/null
+++ b/src/mflux/controlnet/weight_handler_controlnet.py
@@ -0,0 +1,45 @@
+import json
+from pathlib import Path
+
+import mlx.core as mx
+from huggingface_hub import snapshot_download
+from mlx.utils import tree_unflatten
+
+from mflux.weights.weight_util import WeightUtil
+
+
+class WeightHandlerControlnet:
+    @staticmethod
+    def load_controlnet_transformer(controlnet_id: str) -> (dict, int):
+        controlnet_path = Path(
+            snapshot_download(repo_id=controlnet_id, allow_patterns=["*.safetensors", "config.json"])
+        )
+        file = next(controlnet_path.glob("diffusion_pytorch_model.safetensors"))
+        quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level")
+        weights = list(mx.load(str(file)).items())
+
+        if quantization_level is not None:
+            return tree_unflatten(weights), quantization_level
+
+        weights = [WeightUtil.reshape_weights(k, v) for k, v in weights]
+        weights = WeightUtil.flatten(weights)
+        weights = tree_unflatten(weights)
+
+        # Quantized weights (i.e. ones exported from this project) don't need any post-processing.
+        if quantization_level is not None:
+            return weights, quantization_level
+
+        # Reshape and process the huggingface weights
+        if "transformer_blocks" in weights:
+            for block in weights["transformer_blocks"]:
+                block["ff"] = {
+                    "linear1": block["ff"]["net"][0]["proj"],
+                    "linear2": block["ff"]["net"][2]
+                }  # fmt: off
+                if block.get("ff_context") is not None:
+                    block["ff_context"] = {
+                        "linear1": block["ff_context"]["net"][0]["proj"],
+                        "linear2": block["ff_context"]["net"][2],
+                    }
+        config = json.load(open(controlnet_path / "config.json"))
+        return weights, quantization_level, config
diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py
index 257b806..5a2d861 100644
--- a/src/mflux/generate_controlnet.py
+++ b/src/mflux/generate_controlnet.py
@@ -1,14 +1,16 @@
 import argparse
 import time
 
-from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig, ImageUtil
+from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig
 
 
 def main():
     # fmt: off
     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("--controlnet-image-path", type=str, required=True, help="Local path of the image to use as input for controlnet.")
+    parser.add_argument("--controlnet-strength", type=float, default=0.4, help="Controls how strongly the control image influences the output image. A value of 0.0 means no influence. (Default is 0.4)")
+    parser.add_argument("--controlnet-save-canny", action="store_true", help="If set, save the Canny edge detection reference input image.")
     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)")
@@ -16,7 +18,6 @@ def main():
     parser.add_argument("--width", type=int, default=1024, help="Image width (Default is 1024)")
     parser.add_argument("--steps", type=int, default=None, help="Inference Steps")
     parser.add_argument("--guidance", type=float, default=3.5, help="Guidance Scale (Default is 3.5)")
-    parser.add_argument("--controlnet-strength", type=float, default=0.7, help="Controls how strongly the control image influences the output image. A value of 0.0 means no influence. (Default is 0.7)")
     parser.add_argument("--quantize",  "-q", type=int, choices=[4, 8], default=None, help="Quantize the model (4 or 8, Default is None)")
     parser.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk")
     parser.add_argument("--lora-paths", type=str, nargs="*", default=None, help="Local safetensors for applying LORA from disk")
@@ -45,7 +46,9 @@ def main():
     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),
+        output=args.output,
+        controlnet_image_path=args.controlnet_image_path,
+        controlnet_save_canny=args.controlnet_save_canny,
         config=ConfigControlnet(
             num_inference_steps=args.steps,
             height=args.height,
diff --git a/src/mflux/post_processing/generated_image.py b/src/mflux/post_processing/generated_image.py
index ad6c694..251e57b 100644
--- a/src/mflux/post_processing/generated_image.py
+++ b/src/mflux/post_processing/generated_image.py
@@ -1,15 +1,10 @@
-import json
-import logging
-from pathlib import Path
+import importlib
 
 import PIL.Image
 import mlx.core as mx
-import piexif
 
 from mflux.config.model_config import ModelConfig
 
-log = logging.getLogger(__name__)
-
 
 class GeneratedImage:
     def __init__(
@@ -25,6 +20,7 @@ class GeneratedImage:
         generation_time: float,
         lora_paths: list[str],
         lora_scales: list[float],
+        controlnet_image_path: str | None = None,
         controlnet_strength: float | None = None,
     ):
         self.image = image
@@ -38,67 +34,17 @@ class GeneratedImage:
         self.generation_time = generation_time
         self.lora_paths = lora_paths
         self.lora_scales = lora_scales
+        self.controlnet_image = controlnet_image_path
         self.controlnet_strength = controlnet_strength
 
     def save(self, path: str, export_json_metadata: bool = False) -> None:
-        file_path = Path(path)
-        file_path.parent.mkdir(parents=True, exist_ok=True)
-        file_name = file_path.stem
-        file_extension = file_path.suffix
+        from mflux import ImageUtil
 
-        # If a file already exists, create a new name with a counter
-        counter = 1
-        while file_path.exists():
-            new_name = f"{file_name}_{counter}{file_extension}"
-            file_path = file_path.with_name(new_name)
-            counter += 1
-
-        try:
-            # Save image without metadata first
-            self.image.save(file_path)
-            log.info(f"Image saved successfully at: {file_path}")
-
-            # Optionally save json metadata file
-            if export_json_metadata:
-                with open(f"{file_path.with_suffix('.json')}", "w") as json_file:
-                    json.dump(self._get_metadata(), json_file, indent=4)
-
-            # Embed metadata
-            self._embed_metadata(file_path)
-            log.info(f"Metadata embedded successfully at: {file_path}")
-        except Exception as e:
-            log.error(f"Error saving image: {e}")
-
-    def _embed_metadata(self, path: str) -> None:
-        try:
-            # Prepare metadata
-            metadata = self._get_metadata()
-
-            # Convert metadata dictionary to a string
-            metadata_str = str(metadata)
-
-            # Convert the string to bytes (using UTF-8 encoding)
-            user_comment_bytes = metadata_str.encode("utf-8")
-
-            # Define the UserComment tag ID
-            USER_COMMENT_TAG_ID = 0x9286
-
-            # Create a piexif-compatible dictionary structure
-            exif_piexif_dict = {"Exif": {USER_COMMENT_TAG_ID: user_comment_bytes}}
-
-            # Load the image and embed the EXIF data
-            image = PIL.Image.open(path)
-            exif_bytes = piexif.dump(exif_piexif_dict)
-            image.info["exif"] = exif_bytes
-
-            # Save the image with metadata
-            image.save(path, exif=exif_bytes)
-
-        except Exception as e:
-            log.error(f"Error embedding metadata: {e}")
+        ImageUtil.save_image(self.image, path, self._get_metadata(), export_json_metadata)
 
     def _get_metadata(self) -> dict:
         return {
+            "mflux_version": str(GeneratedImage.get_version()),
             "model": str(self.model_config.alias),
             "seed": str(self.seed),
             "steps": str(self.steps),
@@ -106,8 +52,16 @@ class GeneratedImage:
             "precision": f"{self.precision}",
             "quantization": "None" if self.quantization is None else f"{self.quantization} bit",
             "generation_time": f"{self.generation_time:.2f} seconds",
-            "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 "",
+            "lora_paths": ", ".join(self.lora_paths) if self.lora_paths else "None",
+            "lora_scales": ", ".join([f"{scale:.2f}" for scale in self.lora_scales]) if self.lora_scales else "None",
             "prompt": self.prompt,
+            "controlnet_image": "None" if self.controlnet_image is None else self.controlnet_image,
             "controlnet_strength": "None" if self.controlnet_strength is None else f"{self.controlnet_strength:.2f}",
         }
+
+    @staticmethod
+    def get_version():
+        try:
+            return importlib.metadata.version("mflux")
+        except importlib.metadata.PackageNotFoundError:
+            return "unknown"
diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py
index 8f83356..56b4145 100644
--- a/src/mflux/post_processing/image_util.py
+++ b/src/mflux/post_processing/image_util.py
@@ -1,12 +1,19 @@
+import json
+import logging
+from pathlib import Path
+
 import PIL
-from PIL import Image
+import PIL.Image
 import mlx.core as mx
 import numpy as np
-
+from PIL import Image
+import piexif
 from mflux.config.config import ConfigControlnet
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.post_processing.generated_image import GeneratedImage
 
+log = logging.getLogger(__name__)
+
 
 class ImageUtil:
     @staticmethod
@@ -19,6 +26,7 @@ class ImageUtil:
         lora_paths: list[str],
         lora_scales: list[float],
         config: RuntimeConfig,
+        controlnet_image_path: str | None = None,
     ) -> GeneratedImage:
         normalized = ImageUtil._denormalize(decoded_latents)
         normalized_numpy = ImageUtil._to_numpy(normalized)
@@ -35,6 +43,7 @@ class ImageUtil:
             generation_time=generation_time,
             lora_paths=lora_paths,
             lora_scales=lora_scales,
+            controlnet_image_path=controlnet_image_path,
             controlnet_strength=config.controlnet_strength if isinstance(config.config, ConfigControlnet) else None,
         )
 
@@ -76,3 +85,65 @@ class ImageUtil:
     @staticmethod
     def load_image(path: str) -> Image.Image:
         return Image.open(path)
+
+    @staticmethod
+    def save_image(
+            image: PIL.Image.Image,
+            path: str,
+            metadata: dict | None = None,
+            export_json_metadata: bool = False
+    ) -> None:  # fmt: off
+        file_path = Path(path)
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        file_name = file_path.stem
+        file_extension = file_path.suffix
+
+        # If a file already exists, create a new name with a counter
+        counter = 1
+        while file_path.exists():
+            new_name = f"{file_name}_{counter}{file_extension}"
+            file_path = file_path.with_name(new_name)
+            counter += 1
+
+        try:
+            # Save image without metadata first
+            image.save(file_path)
+            log.info(f"Image saved successfully at: {file_path}")
+
+            # Optionally save json metadata file
+            if export_json_metadata:
+                with open(f"{file_path.with_suffix('.json')}", "w") as json_file:
+                    json.dump(metadata, json_file, indent=4)
+
+            # Embed metadata
+            if metadata is not None:
+                ImageUtil._embed_metadata(metadata, file_path)
+                log.info(f"Metadata embedded successfully at: {file_path}")
+        except Exception as e:
+            log.error(f"Error saving image: {e}")
+
+    @staticmethod
+    def _embed_metadata(metadata: dict, path: str) -> None:
+        try:
+            # Convert metadata dictionary to a string
+            metadata_str = str(metadata)
+
+            # Convert the string to bytes (using UTF-8 encoding)
+            user_comment_bytes = metadata_str.encode("utf-8")
+
+            # Define the UserComment tag ID
+            USER_COMMENT_TAG_ID = 0x9286
+
+            # Create a piexif-compatible dictionary structure
+            exif_piexif_dict = {"Exif": {USER_COMMENT_TAG_ID: user_comment_bytes}}
+
+            # Load the image and embed the EXIF data
+            image = PIL.Image.open(path)
+            exif_bytes = piexif.dump(exif_piexif_dict)
+            image.info["exif"] = exif_bytes
+
+            # Save the image with metadata
+            image.save(path, exif=exif_bytes)
+
+        except Exception as e:
+            log.error(f"Error embedding metadata: {e}")
diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py
index 56d25cc..1db8f95 100644
--- a/src/mflux/weights/weight_handler.py
+++ b/src/mflux/weights/weight_handler.py
@@ -1,4 +1,3 @@
-import json
 from pathlib import Path
 
 import mlx.core as mx
@@ -88,44 +87,6 @@ class WeightHandler:
                     }
         return weights, quantization_level
 
-    @staticmethod
-    def load_controlnet_transformer(controlnet_id: str) -> (dict, int):
-        controlnet_path = Path(
-            snapshot_download(
-                repo_id=controlnet_id,
-                allow_patterns=["*.safetensors", "config.json"],
-            )
-        )
-        file = next(controlnet_path.glob("diffusion_pytorch_model.safetensors"))
-        quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level")
-        weights = list(mx.load(str(file)).items())
-
-        if quantization_level is not None:
-            return tree_unflatten(weights), quantization_level
-
-        weights = [WeightUtil.reshape_weights(k, v) for k, v in weights]
-        weights = WeightUtil.flatten(weights)
-        weights = tree_unflatten(weights)
-
-        # Quantized weights (i.e. ones exported from this project) don't need any post-processing.
-        if quantization_level is not None:
-            return weights, quantization_level
-
-        # Reshape and process the huggingface weights
-        if "transformer_blocks" in weights:
-            for block in weights["transformer_blocks"]:
-                block["ff"] = {
-                    "linear1": block["ff"]["net"][0]["proj"],
-                    "linear2": block["ff"]["net"][2],
-                }
-                if block.get("ff_context") is not None:
-                    block["ff_context"] = {
-                        "linear1": block["ff_context"]["net"][0]["proj"],
-                        "linear2": block["ff_context"]["net"][2],
-                    }
-        config = json.load(open(controlnet_path / "config.json"))
-        return weights, quantization_level, config
-
     @staticmethod
     def load_vae(root_path: Path) -> (dict, int):
         weights, quantization_level = WeightHandler._get_weights("vae", root_path)