Merge pull request #159 from filipstrand/redux-and-depth

FLUX.1 Tools | Redux & Depth
2025-04-24 10:05:52 +02:00 · 2025-04-24 10:05:52 +02:00 · c9603486ac
commit c9603486ac
parent c1df171d85 d376f9c931
96 changed files with 2228 additions and 106 deletions
--- a/1
+++ b/1
@ -93,6 +93,7 @@ check: ensure-ruff
 .PHONY: test
 test: ensure-pytest
 	# 🏗️ Running tests...
+	uv pip install mlx==0.25.0  # Install pinned MLX version specifically for testing
 	$(PYTHON) -m pytest
 	# ✅ Tests completed

--- a/README.md
+++ b/README.md
@ -1,4 +1,3 @@
-
 ![image](src/mflux/assets/logo.png)
 *A MLX port of FLUX based on the Huggingface Diffusers implementation.*

@ -35,6 +34,9 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
  * [🖌️ Fill](#%EF%B8%8F-fill)
    + [Inpainting](#inpainting)
    + [Outpainting](#outpainting)
+  * [🔍 Depth](#-depth)
+  * [🔄 Redux](#-redux)
+- [🕹️ Controlnet](#%EF%B8%8F-controlnet)
 - [🎛️ Dreambooth fine-tuning](#-dreambooth-fine-tuning)
  * [Training configuration](#training-configuration)
  * [Training example](#training-example)
@ -42,13 +44,12 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
  * [Configuration details](#configuration-details)
  * [Memory issues](#memory-issues)
  * [Misc](#misc)
- [🕹️ Controlnet](#%EF%B8%8F-controlnet)
 - [🚧 Current limitations](#-current-limitations)
 - [💡Workflow tips](#workflow-tips)
- [✅ TODO](#-todo)
 - [🔬 Cool research / features to support](#-cool-research--features-to-support-)
 - [🌱‍ Related projects](#-related-projects)
- [License](#license)
+- [🙏 Acknowledgements](#-acknowledgements)
+- [⚖️ License](#-license)

 <!-- TOC end -->

@ -869,6 +870,119 @@ mflux-generate-fill \

 ⚠️ *Note: Using the Fill tool requires an additional 33.92 GB download from [black-forest-labs/FLUX.1-Fill-dev](https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev). The download happens automatically on first use.*

+#### 🔍 Depth
+
+Using the depth tool, you can generate high-quality images constrained on a depth map from a reference image. These maps represent the estimated distance of each pixel from the camera, with brighter areas representing objects closer to the camera and darker areas representing objects farther away.
+
+For state-of-the-art depth extraction, MFLUX now supports the [Depth Pro](https://github.com/apple/ml-depth-pro) model, natively implemented in MLX based on the reference implementation. 
+
+![depth example](src/mflux/assets/depth_example.jpg)
+*Original hallway image credit: [Yuliya Matuzava on Unsplash](https://unsplash.com/photos/a-long-hallway-in-a-building-with-arches-and-arches-39NJt9jfGSU)*
+
+##### Example
+
+Using the `mflux-generate-depth` we can let MFLUX generate a new image conditioned on the depth map of the input image:
+
+```bash
+mflux-generate-depth \
+  --prompt "A hallway made of white carrara marble" \
+  -q 8 \
+  --height 1680 \
+  --width 1200 \
+  --image-path "original_image.png" \
+  --steps 20 \
+  --seed 7102725
+```
+
+Instead of automatically generating the depth map based on a reference image, you can instead input the depth map yourself using the `--depth-image-path`, like so:
+
+```bash
+mflux-generate-depth \
+  --prompt "A hallway made of white carrara marble" \
+  -q 8 \
+  --height 1680 \
+  --width 1200 \
+  --depth-image-path "depth_image.png" \
+  --steps 20
+```
+
+To generate and export the depth map from an image without triggering the image generation, run the following command:
+
+```bash
+mflux-save-depth --image-path "your_image.jpg" -q 8
+```
+
+This will create a depth map and save it with the same name as your image but with a "_depth" suffix (e.g., "your_image_depth.png").
+Quantization is supported for the Depth Pro model, however, output quality can very depend on the input image. 
+
+⚠️ *Note: The Depth Pro model requires an additional 1.9GB download from Apple. The download happens automatically on first use.*
+
+⚠️ *Note: Using the Depth tool requires an additional 33.92 GB download from [black-forest-labs/FLUX.1-Depth-dev](https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev). The download happens automatically on first use.*
+
+#### 🔄 Redux
+
+The Redux tool is an adapter for FLUX.1 models that enables image variation generation similar to the [image-to-image](#-image-to-image) technique. 
+It can reproduce an input image with slight variations, allowing you to refine and enhance existing images. 
+Unlike the image-to-image technique which resumes the denoising process mid-way starting from the input image, the redux tool instead embeds the image and joins it with the T5 text encodings.
+
+![redux example](src/mflux/assets/redux_example.jpg)
+*Image credit: [Paris Bilal on Unsplash](https://unsplash.com/photos/a-cat-statue-sitting-on-top-of-a-white-base-Bedd4qHGWCg)*
+
+##### Example
+
+The following examples show how to use the Redux tool to generate variations:
+
+```bash
+mflux-generate-redux \
+  --prompt "a grey statue of a cat on a white platform in front of a blue background" \
+  --redux-image-paths "original.png" \
+  --steps 20 \
+  --height 1654 \
+  --width 1154 \
+  -q 8
+```
+There is a tendency for the reference image to dominate over the input prompt.
+
+⚠️ *Note: Using the Redux tool requires an additional 1.1GB download from [black-forest-labs/FLUX.1-Redux-dev](https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev). The download happens automatically on first use.*
+
+---
+
+### 🕹️ 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)
+
 ---

 ### 🎛️ Dreambooth fine-tuning
@ -1027,44 +1141,6 @@ The aim is to also gradually expand the scope of this feature with alternative t
  - The original fine-tuning script in [Diffusers](https://huggingface.co/docs/diffusers/v0.11.0/en/training/dreambooth)


---
-
-### 🕹️ 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
@ -1107,14 +1183,9 @@ See `uv run tools/rename_images.py --help` for full CLI usage help.
 - When generating multiple images with different seeds, use `--seed` with multiple values or `--auto-seeds` to automatically generate a series of random seeds
 - Use `--stepwise-image-output-dir` to save intermediate images at each denoising step, which can be useful for debugging or creating animations of the generation process

-### ✅ TODO
-
- [ ] [FLUX.1 Tools](https://blackforestlabs.ai/flux-1-tools/)
-
 ### 🔬 Cool research / features to support
 - [ ] [ConceptAttention](https://github.com/helblazer811/ConceptAttention)
 - [ ] [PuLID](https://github.com/ToTheBeginning/PuLID)
- [ ] [depth based controlnet](https://huggingface.co/InstantX/SD3-Controlnet-Depth) via [ml-depth-pro](https://github.com/apple/ml-depth-pro) or similar?
 - [ ] [RF-Inversion](https://github.com/filipstrand/mflux/issues/91)
 - [ ] [catvton-flux](https://github.com/nftblackmagic/catvton-flux)

@ -1125,6 +1196,16 @@ See `uv run tools/rename_images.py --help` for full CLI usage help.
 - [mflux-fasthtml](https://github.com/anthonywu/mflux-fasthtml) by [@anthonywu](https://github.com/anthonywu)
 - [mflux-streamlit](https://github.com/elitexp/mflux-streamlit) by [@elitexp](https://github.com/elitexp)

-### License
+### 🙏 Acknowledgements
+
+MFLUX would not be possible without the great work of:
+
+- The MLX Team for [MLX](https://github.com/ml-explore/mlx) and [MLX examples](https://github.com/ml-explore/mlx-examples)
+- Black Forest Labs for the [FLUX project](https://github.com/black-forest-labs/flux)
+- Hugging Face for the [Diffusers library implementation of Flux](https://huggingface.co/black-forest-labs/FLUX.1-dev) 
+- Depth Pro authors for the [Depth Pro model](https://github.com/apple/ml-depth-pro?tab=readme-ov-file#citation)
+- The MLX community and all [contributors and testers](https://github.com/filipstrand/mflux/graphs/contributors)
+
+### ⚖️ License

 This project is licensed under the [MIT License](LICENSE).
--- a/pyproject.toml
+++ b/pyproject.toml
@ -29,6 +29,7 @@ dependencies = [
    "tokenizers>=0.20.3; python_version>='3.13'",  # transformers -> tokenizers
    "toml>=0.10.2,<1.0",
    "torch>=2.3.1,<3.0; python_version<'3.13'",
+    "torchvision>=0.18.1,<0.22.0",
    # torch dev builds: pip install --pre --index-url https://download.pytorch.org/whl/nightly
    "torch>=2.6.0.dev20241106; python_version>='3.13'",
    "tqdm>=4.66.5,<5.0",
@ -48,6 +49,7 @@ classifiers = [
 dev = [
    "pytest>=8.3.0,<9.0",
    "pytest-timer>=1.0,<2.0",
+    "mlx==0.25.0",  # Used ONLY during test runs to ensure deterministic test results
 ]

 [project.urls]
@ -58,7 +60,10 @@ mflux-generate = "mflux.generate:main"
 mflux-generate-controlnet = "mflux.generate_controlnet:main"
 mflux-generate-in-context = "mflux.generate_in_context:main"
 mflux-generate-fill = "mflux.generate_fill:main"
+mflux-generate-depth = "mflux.generate_depth:main"
+mflux-generate-redux = "mflux.generate_redux:main"
 mflux-save = "mflux.save:main"
+mflux-save-depth = "mflux.save_depth:main"
 mflux-train = "mflux.train:main"

 [tool.setuptools.packages.find]
--- a/src/mflux/assets/depth_example.jpg
+++ b/src/mflux/assets/depth_example.jpg
--- a/src/mflux/assets/redux_example.jpg
+++ b/src/mflux/assets/redux_example.jpg
--- a/src/mflux/callbacks/callback.py
+++ b/src/mflux/callbacks/callback.py
@ -15,6 +15,7 @@ class BeforeLoopCallback(Protocol):
        latents: mx.array,
        config: RuntimeConfig,
        canny_image: PIL.Image.Image | None = None,
+        depth_image: PIL.Image.Image | None = None,
    ) -> None: ...


--- a/src/mflux/callbacks/callbacks.py
+++ b/src/mflux/callbacks/callbacks.py
@ -14,6 +14,7 @@ class Callbacks:
        latents: mx.array,
        config: RuntimeConfig,
        canny_image: PIL.Image.Image | None = None,
+        depth_image: PIL.Image.Image | None = None,
    ):
        for subscriber in CallbackRegistry.before_loop_callbacks():
            subscriber.call_before_loop(
@ -22,6 +23,7 @@ class Callbacks:
                latents=latents,
                config=config,
                canny_image=canny_image,
+                depth_image=depth_image,
            )

    @staticmethod
--- a/src/mflux/callbacks/instances/canny_saver.py
+++ b/src/mflux/callbacks/instances/canny_saver.py
@ -20,6 +20,7 @@ class CannyImageSaver(BeforeLoopCallback):
        latents: mx.array,
        config: RuntimeConfig,
        canny_image: PIL.Image.Image | None = None,
+        depth_image: PIL.Image.Image | None = None,
    ) -> None:
        base, ext = os.path.splitext(self.path)
        ImageUtil.save_image(
--- a/src/mflux/callbacks/instances/depth_saver.py
+++ b/src/mflux/callbacks/instances/depth_saver.py
@ -0,0 +1,32 @@
+import os
+from pathlib import Path
+
+import mlx.core as mx
+import PIL.Image
+
+from mflux import ImageUtil
+from mflux.callbacks.callback import BeforeLoopCallback
+from mflux.config.runtime_config import RuntimeConfig
+
+
+class DepthImageSaver(BeforeLoopCallback):
+    def __init__(self, path: str):
+        self.path = Path(path)
+
+    def call_before_loop(
+        self,
+        seed: int,
+        prompt: str,
+        latents: mx.array,
+        config: RuntimeConfig,
+        canny_image: PIL.Image.Image | None = None,
+        depth_image: PIL.Image.Image | None = None,
+    ) -> None:
+        if depth_image is None:
+            return
+
+        base, ext = os.path.splitext(self.path)
+        ImageUtil.save_image(
+            image=depth_image,
+            path=f"{base}_depth_map{ext}"
+        )  # fmt: off
--- a/src/mflux/callbacks/instances/memory_saver.py
+++ b/src/mflux/callbacks/instances/memory_saver.py
@ -29,6 +29,7 @@ class MemorySaver(BeforeLoopCallback, InLoopCallback, AfterLoopCallback):
        latents: mx.array,
        config: RuntimeConfig,
        canny_image: PIL.Image.Image | None = None,
+        depth_image: PIL.Image.Image | None = None,
    ) -> None:
        self.peak_memory = mx.metal.get_peak_memory()
        self._delete_encoders()
--- a/src/mflux/callbacks/instances/stepwise_handler.py
+++ b/src/mflux/callbacks/instances/stepwise_handler.py
@ -30,6 +30,7 @@ class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback):
        latents: mx.array,
        config: RuntimeConfig,
        canny_image: PIL.Image.Image | None = None,
+        depth_image: PIL.Image.Image | None = None,
    ) -> None:
        self._save_image(
            step=config.init_time_step,
--- a/src/mflux/community/in_context_lora/flux_in_context_lora.py
+++ b/src/mflux/community/in_context_lora/flux_in_context_lora.py
@ -8,7 +8,7 @@ from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.error.exceptions import StopImageGenerationException
 from mflux.flux.flux_initializer import FluxInitializer
-from mflux.latent_creator.latent_creator import Img2Img, LatentCreator
+from mflux.latent_creator.latent_creator import LatentCreator
 from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
 from mflux.models.text_encoder.prompt_encoder import PromptEncoder
 from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
@ -59,14 +59,10 @@ class Flux1InContextLoRA(nn.Module):

        # 1. Encode the reference image
        encoded_image = LatentCreator.encode_image(
+            vae=self.vae,
+            image_path=config.image_path,
            height=config.height,
            width=config.width,
-            img2img=Img2Img(
-                vae=self.vae,
-                sigmas=config.sigmas,
-                init_time_step=config.init_time_step,
-                image_path=config.image_path,
-            ),
        )

        # 2. Create the initial latents and keep the initial static noise for later blending
--- a/src/mflux/config/config.py
+++ b/src/mflux/config/config.py
@ -17,6 +17,8 @@ class Config:
        guidance: float = 4.0,
        image_path: Path | None = None,
        image_strength: float | None = None,
+        depth_image_path: Path | None = None,
+        redux_image_paths: list[Path] | None = None,
        masked_image_path: Path | None = None,
        controlnet_strength: float | None = None,
    ):
@ -28,5 +30,7 @@ class Config:
        self.guidance = guidance
        self.image_path = image_path
        self.image_strength = image_strength
+        self.depth_image_path = depth_image_path
+        self.redux_image_paths = redux_image_paths
        self.masked_image_path = masked_image_path
        self.controlnet_strength = controlnet_strength
--- a/src/mflux/config/model_config.py
+++ b/src/mflux/config/model_config.py
@ -35,11 +35,29 @@ class ModelConfig:
    def dev_fill() -> "ModelConfig":
        return AVAILABLE_MODELS["dev-fill"]

+    @staticmethod
+    @lru_cache
+    def dev_depth() -> "ModelConfig":
+        return AVAILABLE_MODELS["dev-depth"]
+
+    @staticmethod
+    @lru_cache
+    def dev_redux() -> "ModelConfig":
+        return AVAILABLE_MODELS["dev-redux"]
+
    @staticmethod
    @lru_cache
    def schnell() -> "ModelConfig":
        return AVAILABLE_MODELS["schnell"]

+    def x_embedder_input_dim(self) -> int:
+        if self.alias == "dev-fill":
+            return 384
+        if self.alias == "dev-depth":
+            return 128
+        else:
+            return 64
+
    @staticmethod
    def from_name(
        model_name: str,
@ -116,4 +134,24 @@ AVAILABLE_MODELS = {
        requires_sigma_shift=True,
        priority=0,
    ),
+    "dev-depth": ModelConfig(
+        alias="dev-depth",
+        model_name="black-forest-labs/FLUX.1-Depth-dev",
+        base_model=None,
+        num_train_steps=1000,
+        max_sequence_length=512,
+        supports_guidance=True,
+        requires_sigma_shift=True,
+        priority=4,
+    ),
+    "dev-redux": ModelConfig(
+        alias="dev-redux",
+        model_name="black-forest-labs/FLUX.1-Redux-dev",
+        base_model=None,
+        num_train_steps=1000,
+        max_sequence_length=512,
+        supports_guidance=True,
+        requires_sigma_shift=True,
+        priority=3,
+    ),
 }
--- a/src/mflux/config/runtime_config.py
+++ b/src/mflux/config/runtime_config.py
@ -55,6 +55,14 @@ class RuntimeConfig:
    def image_strength(self) -> float | None:
        return self.config.image_strength

+    @property
+    def depth_image_path(self) -> str | None:
+        return self.config.depth_image_path
+
+    @property
+    def redux_image_paths(self) -> str | None:
+        return self.config.redux_image_paths
+
    @property
    def masked_image_path(self) -> str | None:
        return self.config.masked_image_path
--- a/src/mflux/flux/flux.py
+++ b/src/mflux/flux/flux.py
@ -61,9 +61,9 @@ class Flux1(nn.Module):
            width=config.width,
            img2img=Img2Img(
                vae=self.vae,
+                image_path=config.image_path,
                sigmas=config.sigmas,
                init_time_step=config.init_time_step,
-                image_path=config.image_path,
            ),
        )

--- a/src/mflux/flux/flux_initializer.py
+++ b/src/mflux/flux/flux_initializer.py
@ -1,6 +1,10 @@
 from mflux import ModelConfig
 from mflux.controlnet.transformer_controlnet import TransformerControlnet
 from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet
+from mflux.flux_tools.redux.weight_handler_redux import WeightHandlerRedux
+from mflux.models.depth_pro.depth_pro_initializer import DepthProInitializer
+from mflux.models.redux_encoder.redux_encoder import ReduxEncoder
+from mflux.models.siglip_vision_transformer.siglip_vision_transformer import SiglipVisionTransformer
 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
@ -26,7 +30,7 @@ class FluxInitializer:
        lora_names: list[str] | None = None,
        lora_repo_id: str | None = None,
    ) -> None:
-        # 0. Set paths, configs and prompt_cache for later
+        # 0. Set paths, configs, and prompt_cache for later
        lora_paths = lora_paths or []
        flux_model.prompt_cache = {}
        flux_model.lora_paths = lora_paths
@ -88,6 +92,65 @@ class FluxInitializer:
            loras=lora_weights,
        )

+    @staticmethod
+    def init_depth(
+        flux_model,
+        model_config: ModelConfig,
+        quantize: int | None,
+        local_path: str | None,
+        lora_paths: list[str] | None = None,
+        lora_scales: list[float] | None = None,
+        lora_names: list[str] | None = None,
+        lora_repo_id: str | None = None,
+    ):
+        # 1. Start with the same init as regular Flux
+        FluxInitializer.init(
+            flux_model=flux_model,
+            model_config=model_config,
+            quantize=quantize,
+            local_path=local_path,
+            lora_paths=lora_paths,
+            lora_scales=lora_scales,
+            lora_names=lora_names,
+            lora_repo_id=lora_repo_id,
+        )
+
+        # 2. Initialize the DepthPro model and assign the weights
+        flux_model.depth_pro = DepthProInitializer.init()
+
+    @staticmethod
+    def init_redux(
+        flux_model,
+        quantize: int | None,
+        local_path: str | None,
+        lora_paths: list[str] | None = None,
+        lora_scales: list[float] | None = None,
+        lora_names: list[str] | None = None,
+        lora_repo_id: str | None = None,
+    ):
+        # 1. Start with the same init as regular Flux dev
+        FluxInitializer.init(
+            flux_model=flux_model,
+            model_config=ModelConfig.dev(),
+            quantize=quantize,
+            local_path=local_path,
+            lora_paths=lora_paths,
+            lora_scales=lora_scales,
+            lora_names=lora_names,
+            lora_repo_id=lora_repo_id,
+        )
+
+        # 2. Initialize the redux specific addons
+        redux_weights = WeightHandlerRedux.load_weights()
+        flux_model.image_embedder = ReduxEncoder()
+        flux_model.image_encoder = SiglipVisionTransformer()
+        WeightUtil.set_redux_weights_and_quantize(
+            quantize_arg=quantize,
+            weights=redux_weights,
+            redux_encoder=flux_model.image_embedder,
+            siglip_vision_transformer=flux_model.image_encoder,
+        )
+
    @staticmethod
    def init_controlnet(
        flux_model,
@ -99,7 +162,7 @@ class FluxInitializer:
        lora_names: list[str] | None = None,
        lora_repo_id: str | None = None,
    ) -> None:
-        # 1. Start with same init as regular Flux
+        # 1. Start with the same init as regular Flux
        FluxInitializer.init(
            flux_model=flux_model,
            model_config=model_config,
--- a/src/mflux/flux_tools/depth/init.py
+++ b/src/mflux/flux_tools/depth/init.py
--- a/src/mflux/flux_tools/depth/depth_util.py
+++ b/src/mflux/flux_tools/depth/depth_util.py
@ -0,0 +1,66 @@
+import logging
+import os
+
+import mlx.core as mx
+import PIL.Image
+
+from mflux.config.runtime_config import RuntimeConfig
+from mflux.models.depth_pro.depth_pro import DepthPro
+from mflux.models.vae.vae import VAE
+from mflux.post_processing.array_util import ArrayUtil
+from mflux.post_processing.image_util import ImageUtil
+
+logger = logging.getLogger(__name__)
+
+
+class DepthUtil:
+    @staticmethod
+    def encode_depth_map(
+        vae: VAE,
+        depth_pro: DepthPro,
+        config: RuntimeConfig,
+        image_path: str | None = None,
+        depth_image_path: str | None = None,
+    ) -> (mx.array, PIL.Image.Image):
+        # 1. Create the depth map or use existing one
+        depth_image_path, depth_image = DepthUtil.get_or_create_depth_map(
+            depth_pro=depth_pro,
+            image_path=image_path,
+            depth_map_path=depth_image_path,
+        )
+
+        # 2. Encode the depth map
+        scaled_depth_map = ImageUtil.scale_to_dimensions(
+            image=ImageUtil.load_image(depth_image_path).convert("RGB"),
+            target_width=config.width,
+            target_height=config.height,
+        )
+        depth_map_array = ImageUtil.to_array(scaled_depth_map)
+        encoded_depth = vae.encode(depth_map_array)
+        depth_latents = ArrayUtil.pack_latents(latents=encoded_depth, height=config.height, width=config.width)
+
+        return depth_latents, depth_image
+
+    @staticmethod
+    def get_or_create_depth_map(
+        depth_pro: DepthPro,
+        image_path: str | None = None,
+        depth_map_path: str | None = None,
+    ) -> tuple[str, PIL.Image.Image | None]:
+        # 1. If a depth map path is provided, use it directly
+        if depth_map_path:
+            if not os.path.exists(depth_map_path):
+                raise FileNotFoundError(f"Depth map file not found: {depth_map_path}")
+            return depth_map_path, None
+        if not image_path:
+            raise ValueError("Either depth_map_path or image_path must be provided")
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+
+        # 2. Generate a depth map from the image using the DepthPro model
+        depth_result = depth_pro(image_path)
+
+        # 3. Save the depth map to a file with the same name + _depth suffix
+        generated_depth_path = str(image_path).rsplit(".", 1)[0] + "_depth.png"
+        depth_result.depth_image.save(generated_depth_path)
+        return generated_depth_path, depth_result.depth_image
--- a/src/mflux/flux_tools/depth/flux_depth.py
+++ b/src/mflux/flux_tools/depth/flux_depth.py
@ -0,0 +1,157 @@
+import mlx.core as mx
+from mlx import nn
+from tqdm import tqdm
+
+from mflux.callbacks.callbacks import Callbacks
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.config.runtime_config import RuntimeConfig
+from mflux.error.exceptions import StopImageGenerationException
+from mflux.flux.flux_initializer import FluxInitializer
+from mflux.flux_tools.depth.depth_util import DepthUtil
+from mflux.latent_creator.latent_creator import LatentCreator
+from mflux.models.depth_pro.depth_pro import DepthPro
+from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
+from mflux.models.text_encoder.prompt_encoder import PromptEncoder
+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.array_util import ArrayUtil
+from mflux.post_processing.generated_image import GeneratedImage
+from mflux.post_processing.image_util import ImageUtil
+
+
+class Flux1Depth(nn.Module):
+    vae: VAE
+    depth_pro: DepthPro
+    transformer: Transformer
+    t5_text_encoder: T5Encoder
+    clip_text_encoder: CLIPEncoder
+
+    def __init__(
+        self,
+        quantize: int | None = None,
+        local_path: str | None = None,
+        lora_paths: list[str] | None = None,
+        lora_scales: list[float] | None = None,
+    ):
+        super().__init__()
+        FluxInitializer.init_depth(
+            flux_model=self,
+            model_config=ModelConfig.dev_depth(),
+            quantize=quantize,
+            local_path=local_path,
+            lora_paths=lora_paths,
+            lora_scales=lora_scales,
+        )
+
+    def generate_image(
+        self,
+        seed: int,
+        prompt: str,
+        config: Config,
+    ) -> GeneratedImage:
+        # 0. Create a new runtime config based on the model type and input parameters
+        config = RuntimeConfig(config, self.model_config)
+        time_steps = tqdm(range(config.init_time_step, config.num_inference_steps))
+
+        # 1. Create the initial latents
+        latents = LatentCreator.create(
+            seed=seed,
+            height=config.height,
+            width=config.width,
+        )
+
+        # 2. Encode the prompt
+        prompt_embeds, pooled_prompt_embeds = PromptEncoder.encode_prompt(
+            prompt=prompt,
+            prompt_cache=self.prompt_cache,
+            t5_tokenizer=self.t5_tokenizer,
+            clip_tokenizer=self.clip_tokenizer,
+            t5_text_encoder=self.t5_text_encoder,
+            clip_text_encoder=self.clip_text_encoder,
+        )
+
+        # 3. Encode the depth map
+        static_depth_map, depth_image = DepthUtil.encode_depth_map(
+            vae=self.vae,
+            depth_pro=self.depth_pro,
+            config=config,
+            image_path=config.image_path,
+            depth_image_path=config.depth_image_path,
+        )
+
+        # (Optional) Call subscribers for beginning of loop
+        Callbacks.before_loop(
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=config,
+            depth_image=depth_image,
+        )  # fmt: off
+
+        for t in time_steps:
+            try:
+                # 4.t Concatenate the updated latents with the (static) depth map
+                hidden_states = mx.concatenate([latents, static_depth_map], axis=-1)
+
+                # 5.t Predict the noise
+                noise = self.transformer(
+                    t=t,
+                    config=config,
+                    hidden_states=hidden_states,
+                    prompt_embeds=prompt_embeds,
+                    pooled_prompt_embeds=pooled_prompt_embeds,
+                )
+
+                # 6.t Take one denoise step
+                dt = config.sigmas[t + 1] - config.sigmas[t]
+                latents += noise * dt
+
+                # (Optional) Call subscribers in-loop
+                Callbacks.in_loop(
+                    t=t,
+                    seed=seed,
+                    prompt=prompt,
+                    latents=latents,
+                    config=config,
+                    time_steps=time_steps,
+                )  # fmt: off
+
+                # (Optional) Evaluate to enable progress tracking
+                mx.eval(latents)
+
+            except KeyboardInterrupt:  # noqa: PERF203
+                Callbacks.interruption(
+                    t=t,
+                    seed=seed,
+                    prompt=prompt,
+                    latents=latents,
+                    config=config,
+                    time_steps=time_steps,
+                )
+                raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
+
+        # (Optional) Call subscribers after loop
+        Callbacks.after_loop(
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=config,
+        )  # fmt: off
+
+        # 7. Decode the latent array and return the image
+        latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
+        decoded = self.vae.decode(latents)
+        return ImageUtil.to_image(
+            decoded_latents=decoded,
+            config=config,
+            seed=seed,
+            prompt=prompt,
+            quantization=self.bits,
+            lora_paths=self.lora_paths,
+            lora_scales=self.lora_scales,
+            image_path=config.image_path,
+            depth_image_path=config.depth_image_path,
+            generation_time=time_steps.format_dict["elapsed"],
+        )
--- a/src/mflux/flux_tools/fill/flux_fill.py
+++ b/src/mflux/flux_tools/fill/flux_fill.py
@ -18,7 +18,6 @@ from mflux.models.vae.vae import VAE
 from mflux.post_processing.array_util import ArrayUtil
 from mflux.post_processing.generated_image import GeneratedImage
 from mflux.post_processing.image_util import ImageUtil
-from mflux.weights.model_saver import ModelSaver


 class Flux1Fill(nn.Module):
@ -29,7 +28,6 @@ class Flux1Fill(nn.Module):

    def __init__(
        self,
-        model_config: ModelConfig,
        quantize: int | None = None,
        local_path: str | None = None,
        lora_paths: list[str] | None = None,
@ -38,7 +36,7 @@ class Flux1Fill(nn.Module):
        super().__init__()
        FluxInitializer.init(
            flux_model=self,
-            model_config=model_config,
+            model_config=ModelConfig.dev_fill(),
            quantize=quantize,
            local_path=local_path,
            lora_paths=lora_paths,
@ -155,6 +153,3 @@ class Flux1Fill(nn.Module):
            masked_image_path=config.masked_image_path,
            generation_time=time_steps.format_dict["elapsed"],
        )
-
-    def save_model(self, base_path: str) -> None:
-        ModelSaver.save_model(self, self.bits, base_path)
--- a/src/mflux/flux_tools/redux/init.py
+++ b/src/mflux/flux_tools/redux/init.py
--- a/src/mflux/flux_tools/redux/flux_redux.py
+++ b/src/mflux/flux_tools/redux/flux_redux.py
@ -0,0 +1,185 @@
+import mlx.core as mx
+from mlx import nn
+from tqdm import tqdm
+
+from mflux.callbacks.callbacks import Callbacks
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.config.runtime_config import RuntimeConfig
+from mflux.error.exceptions import StopImageGenerationException
+from mflux.flux.flux_initializer import FluxInitializer
+from mflux.flux_tools.redux.redux_util import ReduxUtil
+from mflux.latent_creator.latent_creator import LatentCreator
+from mflux.models.redux_encoder.redux_encoder import ReduxEncoder
+from mflux.models.siglip_vision_transformer.siglip_vision_transformer import SiglipVisionTransformer
+from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
+from mflux.models.text_encoder.prompt_encoder import PromptEncoder
+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.array_util import ArrayUtil
+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
+
+
+class Flux1Redux(nn.Module):
+    vae: VAE
+    image_encoder: SiglipVisionTransformer
+    image_embedder: ReduxEncoder
+    transformer: Transformer
+    t5_text_encoder: T5Encoder
+    clip_text_encoder: CLIPEncoder
+
+    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,
+    ):
+        super().__init__()
+        FluxInitializer.init_redux(
+            flux_model=self,
+            quantize=quantize,
+            local_path=local_path,
+            lora_paths=lora_paths,
+            lora_scales=lora_scales,
+        )
+
+    def generate_image(
+        self,
+        seed: int,
+        prompt: str,
+        config: Config,
+    ) -> GeneratedImage:
+        # 0. Create a new runtime config based on the model type and input parameters
+        config = RuntimeConfig(config, self.model_config)
+        time_steps = tqdm(range(config.init_time_step, config.num_inference_steps))
+
+        # 1. Create the initial latents
+        latents = LatentCreator.create(
+            seed=seed,
+            height=config.height,
+            width=config.width,
+        )
+
+        # 2. Get prompt embeddings by fusing the prompt and image embeddings
+        prompt_embeds, pooled_prompt_embeds = Flux1Redux._get_prompt_embeddings(
+            prompt=prompt,
+            prompt_cache=self.prompt_cache,
+            t5_tokenizer=self.t5_tokenizer,
+            clip_tokenizer=self.clip_tokenizer,
+            t5_text_encoder=self.t5_text_encoder,
+            clip_text_encoder=self.clip_text_encoder,
+            image_paths=config.redux_image_paths,
+            image_encoder=self.image_encoder,
+            image_embedder=self.image_embedder,
+        )  # fmt: off
+
+        # (Optional) Call subscribers for beginning of loop
+        Callbacks.before_loop(
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=config,
+        )  # fmt: off
+
+        for t in time_steps:
+            try:
+                # 3.t Predict the noise
+                noise = self.transformer(
+                    t=t,
+                    config=config,
+                    hidden_states=latents,
+                    prompt_embeds=prompt_embeds,
+                    pooled_prompt_embeds=pooled_prompt_embeds,
+                )
+
+                # 4.t Take one denoise step
+                dt = config.sigmas[t + 1] - config.sigmas[t]
+                latents += noise * dt
+
+                # (Optional) Call subscribers in-loop
+                Callbacks.in_loop(
+                    t=t,
+                    seed=seed,
+                    prompt=prompt,
+                    latents=latents,
+                    config=config,
+                    time_steps=time_steps,
+                )  # fmt: off
+
+                # (Optional) Evaluate to enable progress tracking
+                mx.eval(latents)
+
+            except KeyboardInterrupt:  # noqa: PERF203
+                Callbacks.interruption(
+                    t=t,
+                    seed=seed,
+                    prompt=prompt,
+                    latents=latents,
+                    config=config,
+                    time_steps=time_steps,
+                )
+                raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
+
+        # (Optional) Call subscribers after loop
+        Callbacks.after_loop(
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=config,
+        )  # fmt: off
+
+        # 7. Decode the latent array and return the image
+        latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
+        decoded = self.vae.decode(latents)
+        return ImageUtil.to_image(
+            decoded_latents=decoded,
+            config=config,
+            seed=seed,
+            prompt=prompt,
+            quantization=self.bits,
+            lora_paths=self.lora_paths,
+            lora_scales=self.lora_scales,
+            redux_image_paths=config.redux_image_paths,
+            image_strength=config.image_strength,
+            generation_time=time_steps.format_dict["elapsed"],
+        )
+
+    @staticmethod
+    def _get_prompt_embeddings(
+        prompt: str,
+        prompt_cache: dict[str, tuple[mx.array, mx.array]],
+        t5_tokenizer: TokenizerT5,
+        clip_tokenizer: TokenizerCLIP,
+        t5_text_encoder: T5Encoder,
+        clip_text_encoder: CLIPEncoder,
+        image_paths: list[str],
+        image_encoder: SiglipVisionTransformer,
+        image_embedder: ReduxEncoder,
+    ) -> (mx.array, mx.array):
+        # 1. Encode the prompt
+        prompt_embeds_txt, pooled_prompt_embeds = PromptEncoder.encode_prompt(
+            prompt=prompt,
+            prompt_cache=prompt_cache,
+            t5_tokenizer=t5_tokenizer,
+            clip_tokenizer=clip_tokenizer,
+            t5_text_encoder=t5_text_encoder,
+            clip_text_encoder=clip_text_encoder,
+        )
+
+        # 2. Encode the image(s) using the Siglip and Redux encoder
+        image_embeds = ReduxUtil.embed_images(
+            image_paths=image_paths,
+            image_encoder=image_encoder,
+            image_embedder=image_embedder,
+        )  # fmt:off
+
+        # 3. Join text embeddings with all image embeddings
+        prompt_embeds = mx.concatenate([prompt_embeds_txt] + image_embeds, axis=1)
+
+        return prompt_embeds, pooled_prompt_embeds
--- a/src/mflux/flux_tools/redux/redux_util.py
+++ b/src/mflux/flux_tools/redux/redux_util.py
@ -0,0 +1,35 @@
+import mlx.core as mx
+
+from mflux import ImageUtil
+from mflux.models.redux_encoder.redux_encoder import ReduxEncoder
+from mflux.models.siglip_vision_transformer.siglip_vision_transformer import SiglipVisionTransformer
+
+
+class ReduxUtil:
+    @staticmethod
+    def embed_images(
+        image_paths: list[str],
+        image_encoder: SiglipVisionTransformer,
+        image_embedder: ReduxEncoder,
+    ) -> list[mx.array]:  # fmt:off
+        image_embeds_list = []
+        for image_path in image_paths:
+            image_embeds = ReduxUtil._embed_single_image(
+                image_path=image_path,
+                image_encoder=image_encoder,
+                image_embedder=image_embedder,
+            )
+            image_embeds_list.append(image_embeds)
+        return image_embeds_list
+
+    @staticmethod
+    def _embed_single_image(
+        image_path: str,
+        image_encoder: SiglipVisionTransformer,
+        image_embedder: ReduxEncoder,
+    ) -> mx.array:  # fmt:off
+        image = ImageUtil.load_image(image_path).convert("RGB")
+        image = ImageUtil.preprocess_for_model(image=image)
+        image_latents, pooler_output = image_encoder(image)
+        image_embeds = image_embedder(image_latents)
+        return image_embeds
--- a/src/mflux/flux_tools/redux/weight_handler_redux.py
+++ b/src/mflux/flux_tools/redux/weight_handler_redux.py
@ -0,0 +1,36 @@
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+from mflux import ModelConfig
+from mflux.weights.weight_handler import MetaData, WeightHandler
+
+
+class WeightHandlerRedux:
+    def __init__(self, siglip: dict, redux_encoder: dict, meta_data: MetaData):
+        self.siglip = siglip
+        self.redux_encoder = redux_encoder
+        self.meta_data = meta_data
+
+    @staticmethod
+    def load_weights() -> "WeightHandlerRedux":
+        root_path = Path(snapshot_download(repo_id=ModelConfig.dev_redux().model_name, allow_patterns=["*.safetensors", "config.json"]))  # fmt:off
+
+        siglip, _, _ = WeightHandlerRedux._load_siglip_weights(root_path=root_path)
+        redux_encoder, _, _ = WeightHandlerRedux._load_redux_encoder_weights(root_path=root_path)
+
+        return WeightHandlerRedux(
+            siglip=siglip,
+            redux_encoder=redux_encoder,
+            meta_data=MetaData(quantization_level=None)
+        )  # fmt:off
+
+    @staticmethod
+    def _load_siglip_weights(root_path: Path) -> (dict, int, str | None):
+        weights, _, _ = WeightHandler.get_weights("image_encoder", root_path)
+        return weights, _, _
+
+    @staticmethod
+    def _load_redux_encoder_weights(root_path: Path) -> (dict, int, str | None):
+        weights, _, _ = WeightHandler.get_weights("image_embedder", root_path)
+        return weights, _, _
--- a/src/mflux/generate_depth.py
+++ b/src/mflux/generate_depth.py
@ -0,0 +1,75 @@
+from mflux import Config, StopImageGenerationException
+from mflux.callbacks.callback_registry import CallbackRegistry
+from mflux.callbacks.instances.depth_saver import DepthImageSaver
+from mflux.callbacks.instances.memory_saver import MemorySaver
+from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
+from mflux.flux_tools.depth.flux_depth import Flux1Depth
+from mflux.ui.cli.parsers import CommandLineParser
+
+
+def main():
+    # 0. Parse command line arguments
+    parser = CommandLineParser(description="Generate an image using the depth tool.")
+    parser.add_general_arguments()
+    parser.add_model_arguments(require_model_arg=False)
+    parser.add_lora_arguments()
+    parser.add_image_generator_arguments(supports_metadata_config=False)
+    parser.add_depth_arguments()
+    parser.add_output_arguments()
+    args = parser.parse_args()
+
+    # 0. Default to a medium guidance value for depth related tasks.
+    if args.guidance is None:
+        args.guidance = 10
+
+    # 1. Load the model
+    flux = Flux1Depth(
+        quantize=args.quantize,
+        local_path=args.path,
+        lora_paths=args.lora_paths,
+        lora_scales=args.lora_scales,
+    )
+
+    # 2. Register the optional callbacks
+    if args.save_depth_map:
+        CallbackRegistry.register_before_loop(DepthImageSaver(path=args.output))
+    if args.stepwise_image_output_dir:
+        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
+        CallbackRegistry.register_before_loop(handler)
+        CallbackRegistry.register_in_loop(handler)
+        CallbackRegistry.register_interrupt(handler)
+
+    memory_saver = None
+    if args.low_ram:
+        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
+        CallbackRegistry.register_before_loop(memory_saver)
+        CallbackRegistry.register_in_loop(memory_saver)
+        CallbackRegistry.register_after_loop(memory_saver)
+
+    try:
+        for seed in args.seed:
+            # 3. Generate an image for each seed value
+            image = flux.generate_image(
+                seed=seed,
+                prompt=args.prompt,
+                config=Config(
+                    num_inference_steps=args.steps,
+                    height=args.height,
+                    width=args.width,
+                    guidance=args.guidance,
+                    image_path=args.image_path,
+                    depth_image_path=args.depth_image_path,
+                ),
+            )
+
+            # 4. Save the image
+            image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata)
+    except StopImageGenerationException as stop_exc:
+        print(stop_exc)
+    finally:
+        if memory_saver:
+            print(memory_saver.memory_stats())
+
+
+if __name__ == "__main__":
+    main()
--- a/src/mflux/generate_fill.py
+++ b/src/mflux/generate_fill.py
@ -1,4 +1,4 @@
-from mflux import Config, ModelConfig, StopImageGenerationException
+from mflux import Config, StopImageGenerationException
 from mflux.callbacks.callback_registry import CallbackRegistry
 from mflux.callbacks.instances.memory_saver import MemorySaver
 from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
@ -23,7 +23,6 @@ def main():

    # 1. Load the model
    flux = Flux1Fill(
-        model_config=ModelConfig.dev_fill(),
        quantize=args.quantize,
        local_path=args.path,
        lora_paths=args.lora_paths,
@ -39,7 +38,7 @@ def main():

    memory_saver = None
    if args.low_ram:
-        memory_saver = MemorySaver(flux, keep_transformer=len(args.seed) > 1)
+        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
        CallbackRegistry.register_before_loop(memory_saver)
        CallbackRegistry.register_in_loop(memory_saver)
        CallbackRegistry.register_after_loop(memory_saver)
--- a/src/mflux/generate_in_context.py
+++ b/src/mflux/generate_in_context.py
@ -38,7 +38,7 @@ def main():

    memory_saver = None
    if args.low_ram:
-        memory_saver = MemorySaver(flux, keep_transformer=len(args.seed) > 1)
+        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
        CallbackRegistry.register_before_loop(memory_saver)
        CallbackRegistry.register_in_loop(memory_saver)
        CallbackRegistry.register_after_loop(memory_saver)
--- a/src/mflux/generate_redux.py
+++ b/src/mflux/generate_redux.py
@ -0,0 +1,68 @@
+from mflux import Config, ModelConfig, StopImageGenerationException
+from mflux.callbacks.callback_registry import CallbackRegistry
+from mflux.callbacks.instances.memory_saver import MemorySaver
+from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
+from mflux.flux_tools.redux.flux_redux import Flux1Redux
+from mflux.ui.cli.parsers import CommandLineParser
+
+
+def main():
+    # 0. Parse command line arguments
+    parser = CommandLineParser(description="Generate an image using the fill tool to complete masked areas.")
+    parser.add_general_arguments()
+    parser.add_model_arguments(require_model_arg=False)
+    parser.add_lora_arguments()
+    parser.add_image_generator_arguments(supports_metadata_config=False)
+    parser.add_redux_arguments()
+    parser.add_output_arguments()
+    args = parser.parse_args()
+
+    # 1. Load the model
+    flux = Flux1Redux(
+        model_config=ModelConfig.dev_redux(),
+        quantize=args.quantize,
+        local_path=args.path,
+        lora_paths=args.lora_paths,
+        lora_scales=args.lora_scales,
+    )
+
+    # 2. Register the optional callbacks
+    if args.stepwise_image_output_dir:
+        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
+        CallbackRegistry.register_before_loop(handler)
+        CallbackRegistry.register_in_loop(handler)
+        CallbackRegistry.register_interrupt(handler)
+
+    memory_saver = None
+    if args.low_ram:
+        memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
+        CallbackRegistry.register_before_loop(memory_saver)
+        CallbackRegistry.register_in_loop(memory_saver)
+        CallbackRegistry.register_after_loop(memory_saver)
+
+    try:
+        for seed in args.seed:
+            # 3. Generate an image for each seed value
+            image = flux.generate_image(
+                seed=seed,
+                prompt=args.prompt,
+                config=Config(
+                    num_inference_steps=args.steps,
+                    height=args.height,
+                    width=args.width,
+                    guidance=args.guidance,
+                    redux_image_paths=args.redux_image_paths,
+                ),
+            )
+
+            # 4. Save the image
+            image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata)
+    except StopImageGenerationException as stop_exc:
+        print(stop_exc)
+    finally:
+        if memory_saver:
+            print(memory_saver.memory_stats())
+
+
+if __name__ == "__main__":
+    main()
--- a/src/mflux/latent_creator/latent_creator.py
+++ b/src/mflux/latent_creator/latent_creator.py
@ -57,7 +57,7 @@ class LatentCreator:
            )

            # 2. Encode the image
-            encoded = LatentCreator.encode_image(height=height, width=width, img2img=img2img)
+            encoded = LatentCreator.encode_image(vae=img2img.vae, image_path=img2img.image_path, height=height, width=width)  # fmt: off
            latents = ArrayUtil.pack_latents(latents=encoded, height=height, width=width)

            # 3. Find the appropriate sigma value
@ -71,13 +71,13 @@ class LatentCreator:
            )

    @staticmethod
-    def encode_image(height: int, width: int, img2img: Img2Img):
+    def encode_image(vae: VAE, image_path: str, height: int, width: int):
        scaled_user_image = ImageUtil.scale_to_dimensions(
-            image=ImageUtil.load_image(img2img.image_path).convert("RGB"),
+            image=ImageUtil.load_image(image_path).convert("RGB"),
            target_width=width,
            target_height=height,
        )
-        encoded = img2img.vae.encode(ImageUtil.to_array(scaled_user_image))
+        encoded = vae.encode(ImageUtil.to_array(scaled_user_image))
        return encoded

    @staticmethod
--- a/src/mflux/models/depth_pro/init.py
+++ b/src/mflux/models/depth_pro/init.py
--- a/src/mflux/models/depth_pro/conv_utils.py
+++ b/src/mflux/models/depth_pro/conv_utils.py
@ -0,0 +1,23 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class ConvUtils:
+    @staticmethod
+    def apply_conv(x: mx.array, conv_module: nn.Module) -> mx.array:
+        """Apply a convolution with channel format conversion.
+
+        MLX expects channels-last format (B,H,W,C) for convolutions,
+        but tensors are generally in channels-first format (B,C,H,W).
+        This helper handles the conversion automatically.
+
+        Args:
+            x: Input tensor in channels-first format (B,C,H,W)
+            conv_module: Convolution module to apply
+
+        Returns:
+            Output tensor in channels-first format (B,C,H,W)
+        """
+        x = mx.transpose(x, (0, 2, 3, 1))
+        x = conv_module(x)
+        return mx.transpose(x, (0, 3, 1, 2))
--- a/src/mflux/models/depth_pro/depth_pro.py
+++ b/src/mflux/models/depth_pro/depth_pro.py
@ -0,0 +1,66 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+import torch
+from PIL import Image
+
+from mflux.models.depth_pro.depth_pro_model import DepthProModel
+from mflux.post_processing.image_util import ImageUtil
+
+
+@dataclass
+class DepthResult:
+    depth_image: Image.Image
+    depth_array: mx.array
+    min_depth: float
+    max_depth: float
+
+
+class DepthPro(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.depth_pro_model = DepthProModel()
+
+    def __call__(self, image_path: str | Path, resize: bool = True) -> DepthResult:
+        input_array, height, width = self._pre_process(image_path)
+        depth = self.depth_pro_model(input_array)
+        return self.post_process(depth, height=height, width=width)
+
+    @staticmethod
+    def _pre_process(image_path):
+        image = Image.open(image_path).convert("RGB")
+        input_array = ImageUtil.preprocess_for_depth_pro(image)
+        input_array = DepthPro._resize(input_array)
+        return input_array, image.height, image.width
+
+    @staticmethod
+    def post_process(depth: mx.array, height: int, width: int):
+        depth_min = mx.min(depth)
+        depth_max = mx.max(depth)
+        normalized_depth = (depth - depth_min) / (depth_max - depth_min)
+        depth_np = np.asarray(normalized_depth.squeeze())
+        depth_image = Image.fromarray((depth_np * 255).astype(np.uint8))
+        depth_image = depth_image.resize((width, height))
+        depth_np = np.array(depth_image) / 255.0
+        return DepthResult(
+            depth_image=depth_image,
+            depth_array=depth,
+            min_depth=depth_min.item(),
+            max_depth=depth_max.item(),
+        )
+
+    @staticmethod
+    def _resize(x: mx.array) -> mx.array:
+        x_np = np.array(x)
+        x_torch = torch.from_numpy(x_np)
+        x_torch = x_torch.unsqueeze(0)
+        x_torch = torch.nn.functional.interpolate(
+            x_torch,
+            size=(1536, 1536),
+            mode="bilinear",
+            align_corners=False,
+        )
+        return mx.array(x_torch)
--- a/src/mflux/models/depth_pro/depth_pro_encoder.py
+++ b/src/mflux/models/depth_pro/depth_pro_encoder.py
@ -0,0 +1,115 @@
+import math
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mflux.models.depth_pro.conv_utils import ConvUtils
+from mflux.models.depth_pro.depth_pro_util import DepthProUtil
+from mflux.models.depth_pro.dino_v2.dino_vision_transformer import DinoVisionTransformer
+from mflux.models.depth_pro.upsample_block import UpSampleBlock
+
+
+class DepthProEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.patch_encoder = DinoVisionTransformer()
+        self.image_encoder = DinoVisionTransformer()
+        self.upsample_latent0 = UpSampleBlock(dim_in=1024, dim_int=256, dim_out=256, upsample_layers=3)
+        self.upsample_latent1 = UpSampleBlock(dim_in=1024, dim_out=256, upsample_layers=2)
+        self.upsample0 = UpSampleBlock(dim_in=1024, dim_out=512, upsample_layers=1)
+        self.upsample1 = UpSampleBlock(dim_in=1024, dim_out=1024, upsample_layers=1)
+        self.upsample2 = UpSampleBlock(dim_in=1024, dim_out=1024, upsample_layers=1)
+        self.upsample_lowres = nn.ConvTranspose2d(in_channels=1024, out_channels=1024, kernel_size=2, stride=2, padding=0, bias=True)  # fmt: off
+        self.fuse_lowres = nn.Conv2d(in_channels=1024 * 2, out_channels=1024, kernel_size=1, stride=1, padding=0, bias=True)  # fmt: off
+
+    def __call__(self, x: mx.array) -> list[mx.array]:
+        # 1. Create the image pyramid
+        x0, x1, x2 = DepthProUtil.create_pyramid(x)
+
+        # 2: Split to create batched overlapped mini-images at the backbone (BeiT/ViT/Dino) resolution.
+        x0_patches = DepthProUtil.split(x0, overlap_ratio=0.25)
+        x1_patches = DepthProUtil.split(x1, overlap_ratio=0.5)
+        x2_patches = x2
+
+        # 3: Run the backbone (BeiT) model and get the result of large batch size.
+        x_pyramid_patches = mx.concatenate((x0_patches, x1_patches, x2_patches), axis=0)
+        x_pyramid_encodings, backbone_highres_hook0, backbone_highres_hook1 = self.patch_encoder(x_pyramid_patches)
+        x_pyramid_encodings = DepthProEncoder._reshape_feature(x_pyramid_encodings, width=24, height=24)
+
+        # Calculate indices for splitting
+        x0_encodings = x_pyramid_encodings[: len(x0_patches)]
+        x1_encodings = x_pyramid_encodings[len(x0_patches) : len(x0_patches) + len(x1_patches)]
+        x2_encodings = x_pyramid_encodings[len(x0_patches) + len(x1_patches) :]
+
+        # 4. Merging
+        x_latent0_encodings = DepthProEncoder._reshape_feature(backbone_highres_hook0, width=24, height=24)
+        x_latent1_encodings = DepthProEncoder._reshape_feature(backbone_highres_hook1, width=24, height=24)
+        x_latent0_features = DepthProEncoder._merge(x_latent0_encodings[: 1 * 5 * 5], batch_size=1, padding=3)
+        x_latent1_features = DepthProEncoder._merge(x_latent1_encodings[: 1 * 5 * 5], batch_size=1, padding=3)
+        x0_features = DepthProEncoder._merge(x0_encodings, batch_size=1, padding=3)
+        x1_features = DepthProEncoder._merge(x1_encodings, batch_size=1, padding=6)
+        x2_features = x2_encodings
+
+        # 5. Upsample feature maps.
+        x_latent0_features = self.upsample_latent0(x_latent0_features)
+        x_latent1_features = self.upsample_latent1(x_latent1_features)
+        x0_features = self.upsample0(x0_features)
+        x1_features = self.upsample1(x1_features)
+        x2_features = self.upsample2(x2_features)
+
+        # 6. Apply the image encoder model.
+        x_global_features, _, _ = self.image_encoder(x2_patches)
+        x_global_features = DepthProEncoder._reshape_feature(embeddings=x_global_features, width=24, height=24)
+        x_global_features = ConvUtils.apply_conv(x_global_features, self.upsample_lowres)
+        x_global_features = mx.concatenate((x2_features, x_global_features), axis=1)
+        x_global_features = ConvUtils.apply_conv(x_global_features, self.fuse_lowres)
+
+        return [
+            x_latent0_features,
+            x_latent1_features,
+            x0_features,
+            x1_features,
+            x_global_features,
+        ]
+
+    @staticmethod
+    def _reshape_feature(
+        embeddings: mx.array,
+        width: int,
+        height: int,
+        cls_token_offset: int = 1,
+    ) -> mx.array:
+        b, hw, c = embeddings.shape
+        if cls_token_offset > 0:
+            embeddings = embeddings[:, cls_token_offset:, :]
+        embeddings = embeddings.reshape(b, height, width, c).transpose(0, 3, 1, 2)
+        return embeddings
+
+    @staticmethod
+    def _merge(x: mx.array, batch_size: int, padding: int = 3) -> mx.array:
+        steps = int(math.sqrt(x.shape[0] // batch_size))
+
+        idx = 0
+
+        output_list = []
+        for j in range(steps):
+            output_row_list = []
+            for i in range(steps):
+                output = x[batch_size * idx : batch_size * (idx + 1)]
+
+                if j != 0:
+                    output = output[..., padding:, :]
+                if i != 0:
+                    output = output[..., :, padding:]
+                if j != steps - 1:
+                    output = output[..., :-padding, :]
+                if i != steps - 1:
+                    output = output[..., :, :-padding]
+
+                output_row_list.append(output)
+                idx += 1
+
+            output_row = mx.concatenate(output_row_list, axis=-1)
+            output_list.append(output_row)
+        output = mx.concatenate(output_list, axis=-2)
+        return output
--- a/src/mflux/models/depth_pro/depth_pro_initializer.py
+++ b/src/mflux/models/depth_pro/depth_pro_initializer.py
@ -0,0 +1,30 @@
+import mlx.nn as nn
+
+from mflux.models.depth_pro.depth_pro import DepthPro
+from mflux.models.depth_pro.weight_handler_depth_pro import WeightHandlerDepthPro
+
+
+class DepthProInitializer:
+    @staticmethod
+    def init(quantize: int | None = None) -> DepthPro:
+        # 1. Initialize the model
+        depth_pro = DepthPro()
+
+        # 2. Load the weights
+        depth_pro_weights = WeightHandlerDepthPro.load_weights()
+        WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample_latent0")
+        WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample_latent1")
+        WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample0")
+        WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample1")
+        WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample2")
+        WeightHandlerDepthPro.reposition_head_weights(depth_pro_weights)
+        WeightHandlerDepthPro.reshape_transposed_convolution_weights(depth_pro_weights)
+
+        # 3. Assign the weights to the model
+        depth_pro.depth_pro_model.update(depth_pro_weights.weights)
+
+        # 4. Optionally quantize the model
+        if quantize:
+            nn.quantize(depth_pro.depth_pro_model, bits=quantize)
+
+        return depth_pro
--- a/src/mflux/models/depth_pro/depth_pro_model.py
+++ b/src/mflux/models/depth_pro/depth_pro_model.py
@ -0,0 +1,19 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+from mflux.models.depth_pro.depth_pro_encoder import DepthProEncoder
+from mflux.models.depth_pro.fov_head import FOVHead
+from mflux.models.depth_pro.multires_conv_decoder import MultiresConvDecoder
+
+
+class DepthProModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = DepthProEncoder()
+        self.decoder = MultiresConvDecoder()
+        self.head = FOVHead()
+
+    def __call__(self, x: mx.array) -> (mx.array, mx.array):
+        encodings = self.encoder(x)
+        features = self.decoder(encodings)
+        return self.head(features)
--- a/src/mflux/models/depth_pro/depth_pro_util.py
+++ b/src/mflux/models/depth_pro/depth_pro_util.py
@ -0,0 +1,39 @@
+import math
+
+import mlx.core as mx
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+class DepthProUtil:
+    @staticmethod
+    def create_pyramid(x: mx.array) -> (mx.array, mx.array, mx.array):
+        x0 = x
+        x_np = np.array(x)
+        x_torch = torch.from_numpy(x_np)
+        x1_torch = F.interpolate(x_torch, size=None, scale_factor=0.5, mode="bilinear", align_corners=False)
+        x2_torch = F.interpolate(x_torch, size=None, scale_factor=0.25, mode="bilinear", align_corners=False)
+        x1 = mx.array(x1_torch.numpy())
+        x2 = mx.array(x2_torch.numpy())
+        return x0, x1, x2
+
+    @staticmethod
+    def split(x: mx.array, overlap_ratio: float = 0.25) -> mx.array:
+        patch_size = 384
+        patch_stride = int(patch_size * (1 - overlap_ratio))
+
+        image_size = x.shape[-1]
+        steps = int(math.ceil((image_size - patch_size) / patch_stride)) + 1
+
+        x_patch_list = []
+        for j in range(steps):
+            j0 = j * patch_stride
+            j1 = j0 + patch_size
+
+            for i in range(steps):
+                i0 = i * patch_stride
+                i1 = i0 + patch_size
+                x_patch_list.append(x[..., j0:j1, i0:i1])
+
+        return mx.concatenate(x_patch_list, axis=0)
--- a/src/mflux/models/depth_pro/dino_v2/init.py
+++ b/src/mflux/models/depth_pro/dino_v2/init.py
--- a/src/mflux/models/depth_pro/dino_v2/attention.py
+++ b/src/mflux/models/depth_pro/dino_v2/attention.py
@ -0,0 +1,35 @@
+import mlx.core as mx
+from mlx import nn
+from mlx.core.fast import scaled_dot_product_attention
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int = 1024,
+        head_dim: int = 64,
+        num_heads: int = 16,
+    ):
+        super().__init__()
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.proj = nn.Linear(dim, dim, bias=True)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        # Produce queries, keys and values
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
+        qkv = mx.transpose(qkv, (2, 0, 3, 1, 4))
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        # Scaled dot product attention
+        scale = 1 / mx.sqrt(q.shape[-1])
+        attn_output = scaled_dot_product_attention(q, k, v, scale=scale)
+
+        # Reshape back
+        attn_output = mx.transpose(attn_output, (0, 2, 1, 3))
+        attn_output = mx.reshape(attn_output, (B, -1, self.num_heads * self.head_dim))
+
+        # Final projection
+        return self.proj(attn_output)
--- a/src/mflux/models/depth_pro/dino_v2/dino_vision_transformer.py
+++ b/src/mflux/models/depth_pro/dino_v2/dino_vision_transformer.py
@ -0,0 +1,41 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+from mflux.models.depth_pro.dino_v2.patch_embed import PatchEmbed
+from mflux.models.depth_pro.dino_v2.transformer_block import TransformerBlock
+
+
+class DinoVisionTransformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.cls_token = mx.random.normal(shape=(1, 1, 1024))
+        self.pos_embed = mx.random.normal(shape=(1, 577, 1024))
+        self.patch_embed = PatchEmbed()
+        self.blocks = [TransformerBlock() for i in range(24)]
+        self.norm = nn.LayerNorm(dims=1024, eps=1e-6, bias=True)
+
+    def __call__(self, x: mx.array) -> (mx.array, mx.array, mx.array):
+        backbone_highres_hook0 = None
+        backbone_highres_hook1 = None
+
+        x = self.patch_embed(x)
+        x = self._pos_embed(x)
+        for i, block in enumerate(self.blocks):
+            x = block(x)
+
+            # Save intermediary results for later
+            if i == 5:
+                backbone_highres_hook0 = x
+            if i == 11:
+                backbone_highres_hook1 = x
+
+        x = self.norm(x)
+        return x, backbone_highres_hook0, backbone_highres_hook1
+
+    def _pos_embed(self, x: mx.array) -> mx.array:
+        B, H, W, C = x.shape
+        x = x.reshape((B, -1, 1024))
+        to_cat = [mx.broadcast_to(self.cls_token, (B,) + self.cls_token.shape[1:])]
+        x = mx.concatenate(to_cat + [x], axis=1)
+        x = x + self.pos_embed
+        return x
--- a/src/mflux/models/depth_pro/dino_v2/layer_scale.py
+++ b/src/mflux/models/depth_pro/dino_v2/layer_scale.py
@ -0,0 +1,15 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        init_values: float = 1e-5,
+    ) -> None:
+        super().__init__()
+        self.gamma = init_values * mx.ones((dims,))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return x * self.gamma
--- a/src/mflux/models/depth_pro/dino_v2/mlp.py
+++ b/src/mflux/models/depth_pro/dino_v2/mlp.py
@ -0,0 +1,15 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class MLP(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(1024, 4096, bias=True)
+        self.fc2 = nn.Linear(4096, 1024, bias=True)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = self.fc1(x)
+        x = nn.gelu(x)
+        x = self.fc2(x)
+        return x
--- a/src/mflux/models/depth_pro/dino_v2/patch_embed.py
+++ b/src/mflux/models/depth_pro/dino_v2/patch_embed.py
@ -0,0 +1,13 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.proj = nn.Conv2d(in_channels=3, out_channels=1024, kernel_size=16, stride=16, bias=True)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = mx.transpose(x, (0, 2, 3, 1))
+        x = self.proj(x)
+        return x
--- a/src/mflux/models/depth_pro/dino_v2/transformer_block.py
+++ b/src/mflux/models/depth_pro/dino_v2/transformer_block.py
@ -0,0 +1,22 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+from mflux.models.depth_pro.dino_v2.attention import Attention
+from mflux.models.depth_pro.dino_v2.layer_scale import LayerScale
+from mflux.models.depth_pro.dino_v2.mlp import MLP
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dims=1024, eps=1e-6, bias=True)
+        self.attn = Attention()
+        self.ls1 = LayerScale(dims=1024, init_values=1e-5)
+        self.norm2 = nn.LayerNorm(dims=1024, eps=1e-6, bias=True)
+        self.mlp = MLP()
+        self.ls2 = LayerScale(dims=1024, init_values=1e-5)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = x + self.ls1(self.attn(self.norm1(x)))
+        x = x + self.ls2(self.mlp(self.norm2(x)))
+        return x
--- a/src/mflux/models/depth_pro/feature_fusion_block_2d.py
+++ b/src/mflux/models/depth_pro/feature_fusion_block_2d.py
@ -0,0 +1,26 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+from mflux.models.depth_pro.conv_utils import ConvUtils
+from mflux.models.depth_pro.residual_block import ResidualBlock
+
+
+class FeatureFusionBlock2d(nn.Module):
+    def __init__(self, num_features: int, deconv: bool = False):
+        super().__init__()
+        self.use_deconv = deconv
+        self.resnet1 = ResidualBlock(num_features)
+        self.resnet2 = ResidualBlock(num_features)
+        self.deconv = nn.ConvTranspose2d(in_channels=num_features, out_channels=num_features, kernel_size=2, stride=2, padding=0, bias=False)  # fmt: off
+        self.out_conv = nn.Conv2d(in_channels=num_features, out_channels=num_features, kernel_size=1, stride=1, padding=0, bias=True)  # fmt: off
+
+    def __call__(self, x0: mx.array, x1: mx.array = None) -> mx.array:
+        x = x0
+        if x1 is not None:
+            res = self.resnet1(x1)
+            x = x + res
+        x = self.resnet2(x)
+        if self.use_deconv:
+            x = ConvUtils.apply_conv(x, self.deconv)
+        x = ConvUtils.apply_conv(x, self.out_conv)
+        return x
--- a/src/mflux/models/depth_pro/fov_head.py
+++ b/src/mflux/models/depth_pro/fov_head.py
@ -0,0 +1,25 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+from mflux.models.depth_pro.conv_utils import ConvUtils
+
+
+class FOVHead(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.convs = [
+            nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3, stride=1, padding=1),
+            nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=2, stride=2, padding=0, bias=True),
+            nn.Conv2d(in_channels=128, out_channels=32, kernel_size=3, stride=1, padding=1),
+            nn.Identity(),
+            nn.Conv2d(in_channels=32, out_channels=1, kernel_size=1, stride=1, padding=0),
+        ]
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = ConvUtils.apply_conv(x, self.convs[0])
+        x = ConvUtils.apply_conv(x, self.convs[1])
+        x = ConvUtils.apply_conv(x, self.convs[2])
+        x = nn.relu(x)
+        x = ConvUtils.apply_conv(x, self.convs[4])
+        x = nn.relu(x)
+        return x
--- a/src/mflux/models/depth_pro/multires_conv_decoder.py
+++ b/src/mflux/models/depth_pro/multires_conv_decoder.py
@ -0,0 +1,38 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+from mflux.models.depth_pro.conv_utils import ConvUtils
+from mflux.models.depth_pro.feature_fusion_block_2d import FeatureFusionBlock2d
+
+
+class MultiresConvDecoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.convs = [
+            nn.Identity(),
+            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.Conv2d(in_channels=512, out_channels=256, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.Conv2d(in_channels=1024, out_channels=256, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.Conv2d(in_channels=1024, out_channels=256, kernel_size=3, stride=1, padding=1, bias=False),
+        ]
+        self.fusions = [
+            FeatureFusionBlock2d(num_features=256, deconv=False),
+            FeatureFusionBlock2d(num_features=256, deconv=True),
+            FeatureFusionBlock2d(num_features=256, deconv=True),
+            FeatureFusionBlock2d(num_features=256, deconv=True),
+            FeatureFusionBlock2d(num_features=256, deconv=True),
+        ]
+
+    def __call__(self, encodings: list[mx.array]) -> tuple[mx.array, mx.array]:
+        # Process last layer
+        encodings_last = encodings[4]
+        features = ConvUtils.apply_conv(encodings_last, self.convs[4])
+        features = self.fusions[4](features)
+
+        # Process remaining levels with skip connections
+        for i in [3, 2, 1, 0]:
+            enc = encodings[i]
+            features_i = ConvUtils.apply_conv(enc, self.convs[i])
+            features = self.fusions[i](features, features_i)
+
+        return features
--- a/src/mflux/models/depth_pro/residual_block.py
+++ b/src/mflux/models/depth_pro/residual_block.py
@ -0,0 +1,36 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+from mflux.models.depth_pro.conv_utils import ConvUtils
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, num_features: int):
+        super().__init__()
+        self.residual = [
+            nn.Identity(),
+            nn.Conv2d(
+                in_channels=num_features,
+                out_channels=num_features,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=True,
+            ),
+            nn.Identity(),
+            nn.Conv2d(
+                in_channels=num_features,
+                out_channels=num_features,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=True,
+            ),
+        ]
+
+    def __call__(self, x: mx.array) -> mx.array:
+        delta_x = nn.relu(x)
+        delta_x = ConvUtils.apply_conv(delta_x, self.residual[1])
+        delta_x = nn.relu(delta_x)
+        delta_x = ConvUtils.apply_conv(delta_x, self.residual[3])
+        return x + delta_x
--- a/src/mflux/models/depth_pro/upsample_block.py
+++ b/src/mflux/models/depth_pro/upsample_block.py
@ -0,0 +1,51 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+from mflux.models.depth_pro.conv_utils import ConvUtils
+
+
+class UpSampleBlock(nn.Module):
+    def __init__(
+        self,
+        dim_in: int = 1152,
+        dim_int: int = 256,
+        dim_out: int = 256,
+        upsample_layers: int = 3,
+    ):
+        super().__init__()
+        self.layers = UpSampleBlock._create_layers(
+            dim_in=dim_in,
+            dim_int=dim_int,
+            dim_out=dim_out,
+            upsample_layers=upsample_layers
+        )  # fmt: off
+
+    @staticmethod
+    def _create_layers(dim_in: int, dim_out: int, upsample_layers: int, dim_int: int = None) -> list[nn.Module]:
+        if dim_int is None:
+            dim_int = dim_out
+
+        # Create projection layer
+        layers = [nn.Conv2d(in_channels=dim_in, out_channels=dim_int, kernel_size=1, stride=1, padding=0, bias=False)]
+
+        # Add upsampling layers
+        layers.extend(
+            [
+                nn.ConvTranspose2d(
+                    in_channels=dim_int if i == 0 else dim_out,
+                    out_channels=dim_out,
+                    kernel_size=2,
+                    stride=2,
+                    padding=0,
+                    bias=False,
+                )
+                for i in range(upsample_layers)
+            ]
+        )
+
+        return layers
+
+    def __call__(self, x: mx.array) -> mx.array:
+        for layer in self.layers:
+            x = ConvUtils.apply_conv(x, layer)
+        return x
--- a/src/mflux/models/depth_pro/weight_handler_depth_pro.py
+++ b/src/mflux/models/depth_pro/weight_handler_depth_pro.py
@ -0,0 +1,129 @@
+import logging
+import os
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+import mlx.core as mx
+import torch
+from mlx.utils import tree_unflatten
+
+from mflux.weights.weight_handler import MetaData
+from mflux.weights.weight_util import WeightUtil
+
+
+class WeightHandlerDepthPro:
+    def __init__(self, weights: dict, meta_data: MetaData):
+        self.weights = weights
+        self.meta_data = meta_data
+
+    @staticmethod
+    def load_weights() -> "WeightHandlerDepthPro":
+        model_path = WeightHandlerDepthPro._download_or_get_cached_weights()
+        pt_weights = torch.load(model_path, map_location="cpu")
+        weights = WeightHandlerDepthPro._to_mlx_weights(pt_weights)
+        weights = [WeightUtil.reshape_weights(k, v) for k, v in weights.items()]
+        weights = WeightUtil.flatten(weights)
+        weights = tree_unflatten(weights)
+        return WeightHandlerDepthPro(
+            weights=weights,
+            meta_data=MetaData(quantization_level=None)
+        )  # fmt:off
+
+    @staticmethod
+    def _to_mlx_weights(pt_weights) -> dict:
+        mlx_weights = {}
+        for key, value in pt_weights.items():
+            if isinstance(value, torch.Tensor):
+                mlx_weights[key] = mx.array(value.numpy())
+            else:
+                mlx_weights[key] = value
+        return mlx_weights
+
+    @staticmethod
+    def _download_or_get_cached_weights():
+        APPLE_MODEL_URL = "https://ml-site.cdn-apple.com/models/depth-pro/depth_pro.pt"
+
+        # 1. Create cache directory for the model
+        cache_dir = Path(os.path.expanduser("~/.cache/mflux/depth_pro"))
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        model_path = cache_dir / "depth_pro.pt"
+
+        # 2. Download if model doesn't exist
+        if not model_path.exists():
+            logging.info("Downloading Depth Pro model from Apple...")
+            try:
+                urllib.request.urlretrieve(APPLE_MODEL_URL, model_path)
+                logging.info(f"Downloaded model to {model_path}")
+            except (urllib.error.URLError, urllib.error.HTTPError) as e:
+                logging.error(f"Failed to download model: {e}")
+                logging.info(f"Please manually download from: {APPLE_MODEL_URL}")
+                if not model_path.exists():
+                    raise FileNotFoundError(f"Model file not found at {model_path}")
+
+        return model_path
+
+    @staticmethod
+    def reposition_encoder_weights(depth_pro_weights, name):
+        tmp = depth_pro_weights.weights["encoder"][name]
+        depth_pro_weights.weights["encoder"][name] = {}
+        depth_pro_weights.weights["encoder"][name]["layers"] = tmp
+
+    @staticmethod
+    def reposition_head_weights(depth_pro_weights):
+        tmp = depth_pro_weights.weights["head"]
+        depth_pro_weights.weights["head"] = {}
+        depth_pro_weights.weights["head"]["convs"] = tmp
+
+    @staticmethod
+    def reshape_transposed_convolution_weights(depth_pro_weights):
+        WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent0", 1)
+        WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent0", 2)
+        WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent0", 3)
+
+        WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent1", 1)
+        WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent1", 2)
+
+        WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample0", 1)
+        WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample1", 1)
+        WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample2", 1)
+
+        WeightHandlerDepthPro._reshape_upsample_lowres(depth_pro_weights)
+
+        WeightHandlerDepthPro._reshape_deconv(depth_pro_weights, 1)
+        WeightHandlerDepthPro._reshape_deconv(depth_pro_weights, 2)
+        WeightHandlerDepthPro._reshape_deconv(depth_pro_weights, 3)
+        WeightHandlerDepthPro._reshape_deconv(depth_pro_weights, 4)
+
+        WeightHandlerDepthPro._reshape_head(depth_pro_weights, 1)
+
+    @staticmethod
+    def _reshape_upsample(depth_pro_weights, name, layer):
+        tmp = depth_pro_weights.weights["encoder"][name]["layers"][layer]["weight"]
+        tmp = WeightHandlerDepthPro._reshape(tmp)
+        depth_pro_weights.weights["encoder"][name]["layers"][layer]["weight"] = tmp
+
+    @staticmethod
+    def _reshape_upsample_lowres(depth_pro_weights):
+        tmp = depth_pro_weights.weights["encoder"]["upsample_lowres"]["weight"]
+        tmp = WeightHandlerDepthPro._reshape(tmp)
+        depth_pro_weights.weights["encoder"]["upsample_lowres"]["weight"] = tmp
+
+    @staticmethod
+    def _reshape_deconv(depth_pro_weights, layer):
+        tmp = depth_pro_weights.weights["decoder"]["fusions"][layer]["deconv"]["weight"]
+        tmp = WeightHandlerDepthPro._reshape(tmp)
+        depth_pro_weights.weights["decoder"]["fusions"][layer]["deconv"]["weight"] = tmp
+
+    @staticmethod
+    def _reshape_head(depth_pro_weights, layer):
+        tmp = depth_pro_weights.weights["head"]["convs"][layer]["weight"]
+        tmp = WeightHandlerDepthPro._reshape(tmp)
+        depth_pro_weights.weights["head"]["convs"][layer]["weight"] = tmp
+
+    @staticmethod
+    def _reshape(tensor):
+        tensor = tensor.transpose(0, 3, 1, 2)
+        tensor = tensor.transpose(1, 0, 2, 3)
+        tensor = tensor.transpose(0, 2, 3, 1)
+        return tensor
--- a/src/mflux/models/redux_encoder/init.py
+++ b/src/mflux/models/redux_encoder/init.py
--- a/src/mflux/models/redux_encoder/redux_encoder.py
+++ b/src/mflux/models/redux_encoder/redux_encoder.py
@ -0,0 +1,12 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class ReduxEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.redux_up = nn.Linear(1152, 4096 * 3)
+        self.redux_down = nn.Linear(4096 * 3, 4096)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.redux_down(nn.silu(self.redux_up(x)))
--- a/src/mflux/models/siglip_vision_transformer/init.py
+++ b/src/mflux/models/siglip_vision_transformer/init.py
--- a/src/mflux/models/siglip_vision_transformer/siglip_encoder.py
+++ b/src/mflux/models/siglip_vision_transformer/siglip_encoder.py
@ -0,0 +1,16 @@
+import mlx.core as mx
+from mlx import nn
+
+from mflux.models.siglip_vision_transformer.siglip_encoder_layer import SiglipEncoderLayer
+
+
+class SiglipEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.layers = [SiglipEncoderLayer() for _ in range(27)]
+
+    def __call__(self, inputs_embeds: mx.array) -> mx.array:
+        hidden_states = inputs_embeds
+        for layer in self.layers:
+            hidden_states = layer(hidden_states)
+        return hidden_states
--- a/src/mflux/models/siglip_vision_transformer/siglip_encoder_layer.py
+++ b/src/mflux/models/siglip_vision_transformer/siglip_encoder_layer.py
@ -0,0 +1,29 @@
+import mlx.core as mx
+from mlx import nn
+
+from mflux.models.siglip_vision_transformer.siglip_mlp import SiglipMLP
+from mflux.models.siglip_vision_transformer.siglip_sdpa_attention import SiglipSdpaAttention
+
+
+class SiglipEncoderLayer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.self_attn = SiglipSdpaAttention()
+        self.layer_norm1 = nn.LayerNorm(1152, eps=1e-6)
+        self.mlp = SiglipMLP()
+        self.layer_norm2 = nn.LayerNorm(1152, eps=1e-6)
+
+    def __call__(self, hidden_states: mx.array) -> mx.array:
+        # Self-attention
+        residual = hidden_states
+        hidden_states = self.layer_norm1(hidden_states)
+        attention_output = self.self_attn(hidden_states)
+        hidden_states = residual + attention_output
+
+        # MLP
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
--- a/src/mflux/models/siglip_vision_transformer/siglip_mlp.py
+++ b/src/mflux/models/siglip_vision_transformer/siglip_mlp.py
@ -0,0 +1,15 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class SiglipMLP(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(1152, 4304)
+        self.fc2 = nn.Linear(4304, 1152)
+
+    def __call__(self, hidden_states: mx.array) -> mx.array:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = nn.gelu_approx(hidden_states)  # Check that this is equivalent to gelu_pytorch_tanh
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
--- a/src/mflux/models/siglip_vision_transformer/siglip_multi_head_attention_pooling_head.py
+++ b/src/mflux/models/siglip_vision_transformer/siglip_multi_head_attention_pooling_head.py
@ -0,0 +1,22 @@
+import mlx.core as mx
+from mlx import nn
+
+from mflux.models.siglip_vision_transformer.siglip_mlp import SiglipMLP
+
+
+class SiglipMultiHeadAttentionPoolingHead(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.probe = mx.random.normal(shape=(1, 1, 1152))
+        self.attention = nn.MultiHeadAttention(dims=1152, num_heads=16, bias=True)
+        self.layernorm = nn.LayerNorm(1152, eps=1e-6)
+        self.mlp = SiglipMLP()
+
+    def __call__(self, hidden_states: mx.array) -> mx.array:
+        query = mx.broadcast_to(self.probe, (1, 1, 1152))
+        pooled_output = self.attention(query, hidden_states, hidden_states)
+        residual = pooled_output
+        pooled_output = self.layernorm(pooled_output)
+        pooled_output = residual + self.mlp(pooled_output)
+        pooled_output = pooled_output[:, 0]
+        return pooled_output
--- a/src/mflux/models/siglip_vision_transformer/siglip_sdpa_attention.py
+++ b/src/mflux/models/siglip_vision_transformer/siglip_sdpa_attention.py
@ -0,0 +1,38 @@
+import mlx.core as mx
+from mlx import nn
+from mlx.core.fast import scaled_dot_product_attention
+
+
+class SiglipSdpaAttention(nn.Module):
+    head_dimension = 72
+    batch_size = 1
+    num_heads = 16
+
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(input_dims=768, output_dims=768)
+        self.k_proj = nn.Linear(input_dims=768, output_dims=768)
+        self.v_proj = nn.Linear(input_dims=768, output_dims=768)
+        self.out_proj = nn.Linear(input_dims=768, output_dims=768)
+
+    def __call__(self, hidden_states: mx.array) -> mx.array:
+        query = self.q_proj(hidden_states)
+        key = self.k_proj(hidden_states)
+        value = self.v_proj(hidden_states)
+
+        query = SiglipSdpaAttention.reshape_and_transpose(query, self.batch_size, self.num_heads, self.head_dimension)
+        key = SiglipSdpaAttention.reshape_and_transpose(key, self.batch_size, self.num_heads, self.head_dimension)
+        value = SiglipSdpaAttention.reshape_and_transpose(value, self.batch_size, self.num_heads, self.head_dimension)
+
+        scale = 1 / mx.sqrt(query.shape[-1])
+        hidden_states = scaled_dot_product_attention(query, key, value, scale=scale)
+
+        hidden_states = mx.transpose(hidden_states, (0, 2, 1, 3))
+        hidden_states = mx.reshape(hidden_states, (self.batch_size, -1, self.num_heads * self.head_dimension))
+
+        hidden_states = self.out_proj(hidden_states)
+        return hidden_states
+
+    @staticmethod
+    def reshape_and_transpose(x, batch_size, num_heads, head_dim):
+        return mx.transpose(mx.reshape(x, (batch_size, -1, num_heads, head_dim)), (0, 2, 1, 3))
--- a/src/mflux/models/siglip_vision_transformer/siglip_vision_embeddings.py
+++ b/src/mflux/models/siglip_vision_transformer/siglip_vision_embeddings.py
@ -0,0 +1,35 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class SiglipVisionEmbeddings(nn.Module):
+    embed_dim = 1152
+    image_size = 384
+    patch_size = 14
+
+    def __init__(self):
+        super().__init__()
+        self.patch_embedding = nn.Conv2d(
+            in_channels=3,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            padding=0,
+        )
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.num_positions = self.num_patches
+        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+
+    def __call__(self, pixel_values: mx.array) -> mx.array:
+        # Get embeddings
+        pixel_values = mx.transpose(pixel_values, (0, 2, 3, 1))
+        embeddings = self.patch_embedding(pixel_values)
+        embeddings = mx.transpose(embeddings, (0, 3, 1, 2))
+        embeddings = embeddings.reshape(1, 1152, 27 * 27)
+        embeddings = mx.transpose(embeddings, (0, 2, 1))
+
+        # Get position embeddings
+        position_ids = mx.arange(self.num_positions)
+        position_embeddings = self.position_embedding(position_ids)
+
+        return embeddings + position_embeddings
--- a/src/mflux/models/siglip_vision_transformer/siglip_vision_transformer.py
+++ b/src/mflux/models/siglip_vision_transformer/siglip_vision_transformer.py
@ -0,0 +1,24 @@
+import mlx.core as mx
+from mlx import nn
+
+from mflux.models.siglip_vision_transformer.siglip_encoder import SiglipEncoder
+from mflux.models.siglip_vision_transformer.siglip_multi_head_attention_pooling_head import (
+    SiglipMultiHeadAttentionPoolingHead,
+)
+from mflux.models.siglip_vision_transformer.siglip_vision_embeddings import SiglipVisionEmbeddings
+
+
+class SiglipVisionTransformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embeddings = SiglipVisionEmbeddings()
+        self.encoder = SiglipEncoder()
+        self.post_layernorm = nn.LayerNorm(dims=1152, eps=1e-6)
+        self.head = SiglipMultiHeadAttentionPoolingHead()
+
+    def __call__(self, pixel_values: mx.array) -> mx.array:
+        hidden_states = self.embeddings(pixel_values)
+        encoder_outputs = self.encoder(hidden_states)
+        hidden_state = self.post_layernorm(encoder_outputs)
+        pooler_output = self.head(hidden_state)
+        return hidden_state, pooler_output
--- a/src/mflux/models/text_encoder/prompt_encoder.py
+++ b/src/mflux/models/text_encoder/prompt_encoder.py
@ -10,7 +10,7 @@ class PromptEncoder:
    @staticmethod
    def encode_prompt(
        prompt: str,
-        prompt_cache: dict[str, (mx.array, mx.array)],
+        prompt_cache: dict[str, tuple[mx.array, mx.array]],
        t5_tokenizer: TokenizerT5,
        clip_tokenizer: TokenizerCLIP,
        t5_text_encoder: T5Encoder,
--- a/src/mflux/models/transformer/transformer.py
+++ b/src/mflux/models/transformer/transformer.py
@ -27,7 +27,7 @@ class Transformer(nn.Module):
    ):
        super().__init__()
        self.pos_embed = EmbedND()
-        self.x_embedder = nn.Linear(64, 3072)
+        self.x_embedder = nn.Linear(model_config.x_embedder_input_dim(), 3072)
        self.time_text_embed = TimeTextEmbed(model_config=model_config)
        self.context_embedder = nn.Linear(4096, 3072)
        self.transformer_blocks = [JointTransformerBlock(i) for i in range(num_transformer_blocks)]
--- a/src/mflux/post_processing/generated_image.py
+++ b/src/mflux/post_processing/generated_image.py
@ -28,6 +28,8 @@ class GeneratedImage:
        image_path: str | pathlib.Path | None = None,
        image_strength: float | None = None,
        masked_image_path: str | pathlib.Path | None = None,
+        depth_image_path: str | pathlib.Path | None = None,
+        redux_image_paths: list[str] | list[pathlib.Path] | None = None,
    ):
        self.image = image
        self.model_config = model_config
@ -45,6 +47,8 @@ class GeneratedImage:
        self.image_path = image_path
        self.image_strength = image_strength
        self.masked_image_path = masked_image_path
+        self.depth_image_path = depth_image_path
+        self.redux_image_paths = redux_image_paths

    def get_right_half(self) -> "GeneratedImage":
        # Calculate the coordinates for the right half
@ -69,6 +73,7 @@ class GeneratedImage:
            image_path=self.image_path,
            image_strength=self.image_strength,
            masked_image_path=self.masked_image_path,
+            depth_image_path=self.depth_image_path,
        )

    def save(
@ -104,6 +109,8 @@ class GeneratedImage:
            "controlnet_image_path": str(self.controlnet_image_path) if self.controlnet_image_path else None,
            "controlnet_strength": round(self.controlnet_strength, 2) if self.controlnet_strength else None,
            "masked_image_path": str(self.masked_image_path) if self.masked_image_path else None,
+            "depth_image_path": str(self.depth_image_path) if self.depth_image_path else None,
+            "redux_image_paths": str(self.redux_image_paths) if self.redux_image_paths else None,
            "prompt": self.prompt,
        }

@ -113,7 +120,7 @@ class GeneratedImage:
        if version:
            return version

-        # Fallback to installed package version
+        # Fallback to an installed package version
        try:
            return str(importlib.metadata.version("mflux"))
        except importlib.metadata.PackageNotFoundError:
--- a/src/mflux/post_processing/image_util.py
+++ b/src/mflux/post_processing/image_util.py
@ -8,6 +8,7 @@ import numpy as np
 import piexif
 import PIL.Image
 import PIL.ImageDraw
+import torchvision.transforms.functional as F

 from mflux.config.runtime_config import RuntimeConfig
 from mflux.post_processing.generated_image import GeneratedImage
@ -29,8 +30,10 @@ class ImageUtil:
        lora_scales: list[float],
        controlnet_image_path: str | None = None,
        image_path: str | None = None,
+        redux_image_paths: list[str] | None = None,
        image_strength: float | None = None,
        masked_image_path: str | None = None,
+        depth_image_path: str | None = None,
    ) -> GeneratedImage:
        normalized = ImageUtil._denormalize(decoded_latents)
        normalized_numpy = ImageUtil._to_numpy(normalized)
@ -52,6 +55,8 @@ class ImageUtil:
            controlnet_image_path=controlnet_image_path,
            controlnet_strength=config.controlnet_strength,
            masked_image_path=masked_image_path,
+            depth_image_path=depth_image_path,
+            redux_image_paths=redux_image_paths,
        )

    @staticmethod
@ -260,3 +265,43 @@ class ImageUtil:

        except Exception as e:  # noqa: BLE001
            log.error(f"Error embedding metadata: {e}")
+
+    @staticmethod
+    def preprocess_for_model(
+        image: PIL.Image.Image,
+        target_size: tuple = (384, 384),
+        mean: list = [0.5, 0.5, 0.5],
+        std: list = [0.5, 0.5, 0.5],
+        resample: int = PIL.Image.LANCZOS,
+    ) -> mx.array:
+        # Resize the image to target size
+        image = image.resize(target_size, resample=resample)
+
+        # Convert PIL image to numpy array and normalize to [0, 1]
+        image_np = np.array(image).astype(np.float32) / 255.0
+
+        # Normalize using specified mean and std
+        mean_np = np.array(mean)
+        std_np = np.array(std)
+        image_np = (image_np - mean_np) / std_np
+
+        # Convert from HWC to CHW format
+        image_np = image_np.transpose(2, 0, 1)
+
+        # Convert to MLX array and add batch dimension
+        image_mx = mx.array(image_np)
+        image_mx = mx.expand_dims(image_mx, axis=0)
+
+        return image_mx
+
+    @staticmethod
+    def preprocess_for_depth_pro(
+        image: PIL.Image.Image,
+        target_size: tuple = (384, 384),
+        mean: list = [0.5, 0.5, 0.5],
+        std: list = [0.5, 0.5, 0.5],
+        resample: int = PIL.Image.LANCZOS,
+    ) -> mx.array:
+        image_np = F.to_tensor(image)
+        tensor = F.normalize(image_np, mean, std, False)
+        return mx.array(tensor)
--- a/src/mflux/save.py
+++ b/src/mflux/save.py
@ -3,24 +3,21 @@ from mflux.ui.cli.parsers import CommandLineParser


 def main():
+    # 0. Parse command line arguments
    parser = CommandLineParser(description="Save a quantized version of Flux.1 to disk.")  # fmt: off
    parser.add_model_arguments(path_type="save", require_model_arg=True)
    parser.add_lora_arguments()
    args = parser.parse_args()

-    print(f"Saving model {args.model} with quantization level {args.quantize}\n")
-
+    # 1. Load, quantize and save the model
    flux = Flux1(
        model_config=ModelConfig.from_name(args.model, base_model=args.base_model),
        quantize=args.quantize,
        lora_paths=args.lora_paths,
        lora_scales=args.lora_scales,
    )
-
    flux.save_model(args.path)

-    print(f"Model saved at {args.path}\n")
-

 if __name__ == "__main__":
    main()
--- a/src/mflux/save_depth.py
+++ b/src/mflux/save_depth.py
@ -0,0 +1,18 @@
+from mflux.flux_tools.depth.depth_util import DepthUtil
+from mflux.models.depth_pro.depth_pro_initializer import DepthProInitializer
+from mflux.ui.cli.parsers import CommandLineParser
+
+
+def main():
+    # 0. Parse command line arguments
+    parser = CommandLineParser(description="Save a depth map of an input image")
+    parser.add_depth_arguments()
+    args = parser.parse_args()
+
+    # 1. Create and save the depth map
+    depth_pro = DepthProInitializer.init(quantize=args.quantize)
+    DepthUtil.get_or_create_depth_map(depth_pro=depth_pro, image_path=args.image_path)
+
+
+if __name__ == "__main__":
+    main()
--- a/src/mflux/train.py
+++ b/src/mflux/train.py
@ -5,17 +5,20 @@ from mflux.ui.cli.parsers import CommandLineParser


 def main():
+    # 0. Parse command line arguments
    parser = CommandLineParser(description="Finetune a LoRA adapter")
    parser.add_general_arguments()
    parser.add_model_arguments(require_model_arg=False)
    parser.add_training_arguments()
    args = parser.parse_args()

+    # 1. Initialize the required resources
    flux, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize(
        config_path=args.train_config,
        checkpoint_path=args.train_checkpoint,
    )

+    # 2. Start the training process
    try:
        DreamBooth.train(
            flux=flux,
--- a/src/mflux/ui/cli/parsers.py
+++ b/src/mflux/ui/cli/parsers.py
@ -92,6 +92,15 @@ class CommandLineParser(argparse.ArgumentParser):
        self.add_argument("--image-path", type=Path, required=True, help="Local path to the source image")
        self.add_argument("--masked-image-path", type=Path, required=True, help="Local path to the mask image")

+    def add_depth_arguments(self) -> None:
+        self.add_argument("--image-path", type=Path, required=False, help="Local path to the source image")
+        self.add_argument("--depth-image-path", type=Path, required=False, help="Local path to the depth image")
+        self.add_argument("--save-depth-map", action="store_true", required=False, help="If set, save the depth map created from the source image.")
+        self.add_argument("--quantize",  "-q", type=int, choices=ui_defaults.QUANTIZE_CHOICES, default=None, help=f"Quantize the model ({' or '.join(map(str, ui_defaults.QUANTIZE_CHOICES))}, Default is None)")
+
+    def add_redux_arguments(self) -> None:
+        self.add_argument("--redux-image-paths", type=Path, nargs="*", required=True, help="Local path to the source image")
+
    def add_output_arguments(self) -> None:
        self.add_argument("--metadata", action="store_true", help="Export image metadata as a JSON file.")
        self.add_argument("--output", type=str, default="image.png", help="The filename for the output image. Default is \"image.png\".")
@ -212,7 +221,7 @@ class CommandLineParser(argparse.ArgumentParser):
            self.error("--prompt argument required or 'prompt' required in metadata config file")

        if self.supports_image_generation and namespace.steps is None:
-            namespace.steps = ui_defaults.MODEL_INFERENCE_STEPS.get(namespace.model, None)
+            namespace.steps = ui_defaults.MODEL_INFERENCE_STEPS.get(namespace.model, 14)

        if self.supports_image_outpaint and namespace.image_outpaint_padding is not None:
            # parse and normalize any acceptable 1,2,3,4-tuple box value to 4-tuple
--- a/src/mflux/ui/defaults.py
+++ b/src/mflux/ui/defaults.py
@ -6,6 +6,8 @@ MODEL_CHOICES = ["dev", "dev-fill", "schnell"]
 MODEL_INFERENCE_STEPS = {
    "dev": 14,
    "dev-fill": 14,
+    "dev-depth": 14,
+    "dev-redux": 14,
    "schnell": 4,
 }
 QUANTIZE_CHOICES = [3, 4, 6, 8]
--- a/src/mflux/weights/quantization_util.py
+++ b/src/mflux/weights/quantization_util.py
@ -31,9 +31,40 @@ class QuantizationUtil:
        quantize: int,
        weights: "WeightHandlerControlnet",
        transformer_controlnet: nn.Module,
-    ):
+    ) -> None:
        q_level = weights.meta_data.quantization_level

        if quantize is not None or q_level is not None:
            bits = int(q_level) if q_level is not None else quantize
            nn.quantize(transformer_controlnet, bits=bits)
+
+    @staticmethod
+    def quantize_redux_models(
+        quantize: int,
+        weights: "WeightHandler",
+        redux_encoder: nn.Module,
+        siglip_vision_transformer: nn.Module,
+    ) -> None:
+        q_level = weights.meta_data.quantization_level
+
+        if quantize is not None or q_level is not None:
+            bits = int(q_level) if q_level is not None else quantize
+            nn.quantize(redux_encoder, class_predicate=QuantizationUtil.quantization_predicate, bits=bits)
+            nn.quantize(siglip_vision_transformer, class_predicate=QuantizationUtil.quantization_predicate, bits=bits)
+
+    @staticmethod
+    def quantization_predicate(path, m):
+        # 1. Skip Conv2d layers
+        if isinstance(m, nn.Conv2d):
+            return False
+
+        # 2. Skip any layer with incompatible dimensions
+        if hasattr(m, "weight") and hasattr(m.weight, "shape"):
+            if m.weight.shape == (1152, 4304):
+                return False
+
+            if m.weight.shape[-1] % 64 != 0:
+                return False
+
+        # Only quantize layers that have to_quantized method
+        return hasattr(m, "to_quantized")
--- a/src/mflux/weights/weight_handler.py
+++ b/src/mflux/weights/weight_handler.py
@ -66,12 +66,12 @@ class WeightHandler:

    @staticmethod
    def _load_clip_encoder(root_path: Path) -> (dict, int, str | None):
-        weights, quantization_level, mflux_version = WeightHandler._get_weights("text_encoder", root_path)
+        weights, quantization_level, mflux_version = WeightHandler.get_weights("text_encoder", root_path)
        return weights, quantization_level, mflux_version

    @staticmethod
    def _load_t5_encoder(root_path: Path) -> (dict, int, str | None):
-        weights, quantization_level, mflux_version = WeightHandler._get_weights("text_encoder_2", root_path)
+        weights, quantization_level, mflux_version = WeightHandler.get_weights("text_encoder_2", root_path)

        # Quantized weights (i.e. ones exported from this project) don't need any post-processing.
        if quantization_level is not None:
@ -98,7 +98,7 @@ class WeightHandler:

    @staticmethod
    def load_transformer(root_path: Path | None = None, lora_path: str | None = None) -> (dict, int, str | None):
-        weights, quantization_level, mflux_version = WeightHandler._get_weights("transformer", root_path, lora_path)
+        weights, quantization_level, mflux_version = WeightHandler.get_weights("transformer", root_path, lora_path)

        if lora_path:
            if "transformer" not in weights:
@ -126,7 +126,7 @@ class WeightHandler:

    @staticmethod
    def _load_vae(root_path: Path) -> (dict, int, str | None):
-        weights, quantization_level, mflux_version = WeightHandler._get_weights("vae", root_path)
+        weights, quantization_level, mflux_version = WeightHandler.get_weights("vae", root_path)

        # Quantized weights (i.e. ones exported from this project) don't need any post-processing.
        if quantization_level is not None:
@ -142,7 +142,7 @@ class WeightHandler:
        return weights, quantization_level, mflux_version

    @staticmethod
-    def _get_weights(
+    def get_weights(
        model_name: str,
        root_path: Path | None = None,
        lora_path: str | None = None,
--- a/src/mflux/weights/weight_util.py
+++ b/src/mflux/weights/weight_util.py
@ -7,6 +7,7 @@ from mflux.weights.quantization_util import QuantizationUtil

 if TYPE_CHECKING:
    from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet
+    from mflux.flux_tools.redux.weight_handler_redux import WeightHandlerRedux
    from mflux.weights.weight_handler import WeightHandler


@ -83,3 +84,35 @@ class WeightUtil:
        transformer.update(weights.transformer)
        t5_text_encoder.update(weights.t5_encoder)
        clip_text_encoder.update(weights.clip_encoder)
+
+    @staticmethod
+    def _set_redux_model_weights(
+        weights: "WeightHandlerRedux",
+        redux_encoder: nn.Module,
+        siglip_vision_transformer: nn.Module,
+    ):
+        redux_encoder.update(weights.redux_encoder)
+        siglip_vision_transformer.update(weights.siglip["vision_model"])
+
+    @staticmethod
+    def set_redux_weights_and_quantize(
+        quantize_arg: int | None,
+        weights: "WeightHandlerRedux",
+        redux_encoder: nn.Module,
+        siglip_vision_transformer: nn.Module,
+    ) -> int | None:
+        if weights.meta_data.quantization_level is None and quantize_arg is None:
+            WeightUtil._set_redux_model_weights(weights, redux_encoder, siglip_vision_transformer)
+            return None
+
+        if weights.meta_data.quantization_level is None and quantize_arg is not None:
+            bits = quantize_arg
+            WeightUtil._set_redux_model_weights(weights, redux_encoder, siglip_vision_transformer)
+            QuantizationUtil.quantize_redux_models(bits, weights, redux_encoder, siglip_vision_transformer)
+            return bits
+
+        if weights.meta_data.quantization_level is not None:
+            bits = weights.meta_data.quantization_level
+            QuantizationUtil.quantize_redux_models(bits, weights, redux_encoder, siglip_vision_transformer)
+            WeightUtil._set_redux_model_weights(weights, redux_encoder, siglip_vision_transformer)
+            return bits
--- a/tests/depth/init.py
+++ b/tests/depth/init.py
--- a/tests/depth/test_depth_pro.py
+++ b/tests/depth/test_depth_pro.py
@ -0,0 +1,38 @@
+import os
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+
+from mflux.models.depth_pro.depth_pro_initializer import DepthProInitializer
+
+
+class TestDepthPro:
+    def test_depth_pro_generation(self):
+        # Resolve paths
+        resource_dir = Path(__file__).parent.parent / "resources"
+        input_image_path = resource_dir / "reference_controlnet_dev_lora.png"
+        reference_image_path = resource_dir / "reference_controlnet_dev_lora_depth.png"
+        output_image_path = resource_dir / "depth_output.png"
+
+        try:
+            # Initialize DepthPro model
+            depth_pro = DepthProInitializer.init()
+
+            # Process the image to generate a depth map
+            depth_result = depth_pro(str(input_image_path))
+
+            # Save the output image
+            depth_result.depth_image.save(output_image_path)
+
+            # Assert that the generated depth image matches the reference image
+            np.testing.assert_array_almost_equal(
+                np.array(Image.open(output_image_path)),
+                np.array(Image.open(reference_image_path)),
+                err_msg=f"Generated depth image doesn't match reference depth image. Check {output_image_path} vs {reference_image_path}",
+            )
+
+        finally:
+            # Cleanup
+            if os.path.exists(output_image_path):
+                os.remove(output_image_path)
--- a/tests/image_generation/helpers/image_generation_depth_test_helper.py
+++ b/tests/image_generation/helpers/image_generation_depth_test_helper.py
@ -0,0 +1,64 @@
+import os
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+
+from mflux import Config, ModelConfig
+from mflux.flux_tools.depth.flux_depth import Flux1Depth
+
+
+class ImageGeneratorDepthTestHelper:
+    @staticmethod
+    def assert_matches_reference_image(
+        reference_image_path: str,
+        output_image_path: str,
+        model_config: ModelConfig,
+        prompt: str,
+        steps: int,
+        seed: int,
+        height: int = None,
+        width: int = None,
+        image_path: str | None = None,
+        depth_image_path: str | None = None,
+    ):
+        # resolve paths
+        reference_image_path = ImageGeneratorDepthTestHelper.resolve_path(reference_image_path)
+        output_image_path = ImageGeneratorDepthTestHelper.resolve_path(output_image_path)
+        depth_image_path = ImageGeneratorDepthTestHelper.resolve_path(depth_image_path)
+        image_path = ImageGeneratorDepthTestHelper.resolve_path(image_path)
+
+        try:
+            # given
+            flux = Flux1Depth(quantize=8)
+            image = flux.generate_image(
+                seed=seed,
+                prompt=prompt,
+                config=Config(
+                    num_inference_steps=steps,
+                    height=height,
+                    width=width,
+                    image_path=image_path,
+                    depth_image_path=depth_image_path,
+                ),
+            )
+
+            image.save(path=output_image_path, overwrite=True)
+
+            # then
+            np.testing.assert_array_equal(
+                np.array(Image.open(output_image_path)),
+                np.array(Image.open(reference_image_path)),
+                err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}",
+            )
+
+        finally:
+            # cleanup
+            if os.path.exists(output_image_path):
+                os.remove(output_image_path)
+
+    @staticmethod
+    def resolve_path(path) -> Path | None:
+        if path is None:
+            return None
+        return Path(__file__).parent.parent.parent / "resources" / path
--- a/tests/image_generation/helpers/image_generation_fill_test_helper.py
+++ b/tests/image_generation/helpers/image_generation_fill_test_helper.py
@ -34,7 +34,6 @@ class ImageGeneratorFillTestHelper:
        try:
            # given
            flux = Flux1Fill(
-                model_config=model_config,
                quantize=8,
                lora_paths=lora_paths,
                lora_scales=lora_scales,
--- a/tests/image_generation/helpers/image_generation_redux_test_helper.py
+++ b/tests/image_generation/helpers/image_generation_redux_test_helper.py
@ -0,0 +1,65 @@
+import os
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+
+from mflux import Config, ModelConfig
+from mflux.flux_tools.redux.flux_redux import Flux1Redux
+
+
+class ImageGeneratorReduxTestHelper:
+    @staticmethod
+    def assert_matches_reference_image(
+        reference_image_path: str,
+        output_image_path: str,
+        model_config: ModelConfig,
+        steps: int,
+        seed: int,
+        height: int,
+        width: int,
+        prompt: str,
+        redux_image_path: str,
+    ):
+        # resolve paths
+        reference_image_path = ImageGeneratorReduxTestHelper.resolve_path(reference_image_path)
+        output_image_path = ImageGeneratorReduxTestHelper.resolve_path(output_image_path)
+        redux_image_path = ImageGeneratorReduxTestHelper.resolve_path(redux_image_path)
+
+        try:
+            # given
+            flux = Flux1Redux(
+                model_config=model_config,
+                quantize=8,
+            )
+
+            # when
+            image = flux.generate_image(
+                seed=seed,
+                prompt=prompt,
+                config=Config(
+                    num_inference_steps=steps,
+                    height=height,
+                    width=width,
+                    redux_image_paths=[redux_image_path],
+                ),
+            )
+            image.save(path=output_image_path)
+
+            # then
+            np.testing.assert_array_equal(
+                np.array(Image.open(output_image_path)),
+                np.array(Image.open(reference_image_path)),
+                err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}",
+            )
+
+        finally:
+            # cleanup
+            if os.path.exists(output_image_path):
+                os.remove(output_image_path)
+
+    @staticmethod
+    def resolve_path(path) -> Path | None:
+        if path is None:
+            return None
+        return Path(__file__).parent.parent.parent / "resources" / path
--- a/tests/image_generation/test_generate_image.py
+++ b/tests/image_generation/test_generate_image.py
@ -3,12 +3,10 @@ from tests.image_generation.helpers.image_generation_test_helper import ImageGen


 class TestImageGenerator:
-    OUTPUT_IMAGE_FILENAME = "output.png"
-
    def test_image_generation_schnell(self):
        ImageGeneratorTestHelper.assert_matches_reference_image(
            reference_image_path="reference_schnell.png",
-            output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
+            output_image_path="output_schnell.png",
            model_config=ModelConfig.schnell(),
            steps=2,
            seed=42,
@ -20,7 +18,7 @@ class TestImageGenerator:
    def test_image_generation_dev(self):
        ImageGeneratorTestHelper.assert_matches_reference_image(
            reference_image_path="reference_dev.png",
-            output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
+            output_image_path="output_dev.png",
            model_config=ModelConfig.dev(),
            steps=15,
            seed=42,
@ -32,7 +30,7 @@ class TestImageGenerator:
    def test_image_generation_dev_lora(self):
        ImageGeneratorTestHelper.assert_matches_reference_image(
            reference_image_path="reference_dev_lora.png",
-            output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
+            output_image_path="output_dev_lora.png",
            model_config=ModelConfig.dev(),
            steps=15,
            seed=42,
@ -46,7 +44,7 @@ class TestImageGenerator:
    def test_image_generation_dev_multiple_loras(self):
        ImageGeneratorTestHelper.assert_matches_reference_image(
            reference_image_path="reference_dev_lora_multiple.png",
-            output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
+            output_image_path="output_dev_lora_multiple.png",
            model_config=ModelConfig.dev(),
            steps=15,
            seed=42,
@ -60,13 +58,13 @@ class TestImageGenerator:
    def test_image_generation_dev_image_to_image(self):
        ImageGeneratorTestHelper.assert_matches_reference_image(
            reference_image_path="reference_dev_image_to_image_result.png",
-            image_path="reference_dev_lora.png",
+            output_image_path="output_dev_image_to_image_result.png",
            image_strength=0.4,
-            output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
            model_config=ModelConfig.dev(),
            steps=8,
            seed=44,
            height=341,
            width=768,
+            image_path="reference_dev_lora.png",
            prompt="Luxury food photograph of a burger",
        )
--- a/tests/image_generation/test_generate_image_controlnet.py
+++ b/tests/image_generation/test_generate_image_controlnet.py
@ -3,13 +3,12 @@ from tests.image_generation.helpers.image_generation_controlnet_test_helper impo


 class TestImageGeneratorControlnet:
-    OUTPUT_IMAGE_FILENAME = "output.png"
    CONTROLNET_REFERENCE_FILENAME = "controlnet_reference.png"

    def test_image_generation_schnell_controlnet(self):
        ImageGeneratorControlnetTestHelper.assert_matches_reference_image(
            reference_image_path="reference_controlnet_schnell.png",
-            output_image_path=TestImageGeneratorControlnet.OUTPUT_IMAGE_FILENAME,
+            output_image_path="output_controlnet_schnell.png",
            controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
            model_config=ModelConfig.schnell(),
            steps=2,
@ -21,7 +20,7 @@ class TestImageGeneratorControlnet:
    def test_image_generation_dev_controlnet(self):
        ImageGeneratorControlnetTestHelper.assert_matches_reference_image(
            reference_image_path="reference_controlnet_dev.png",
-            output_image_path=TestImageGeneratorControlnet.OUTPUT_IMAGE_FILENAME,
+            output_image_path="output_controlnet_dev.png",
            controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
            model_config=ModelConfig.dev(),
            steps=15,
@ -33,7 +32,7 @@ class TestImageGeneratorControlnet:
    def test_image_generation_dev_lora_controlnet(self):
        ImageGeneratorControlnetTestHelper.assert_matches_reference_image(
            reference_image_path="reference_controlnet_dev_lora.png",
-            output_image_path=TestImageGeneratorControlnet.OUTPUT_IMAGE_FILENAME,
+            output_image_path="output_controlnet_dev_lora.png",
            controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
            model_config=ModelConfig.dev(),
            steps=15,
--- a/tests/image_generation/test_generate_image_depth.py
+++ b/tests/image_generation/test_generate_image_depth.py
@ -0,0 +1,19 @@
+from mflux import ModelConfig
+from tests.image_generation.helpers.image_generation_depth_test_helper import ImageGeneratorDepthTestHelper
+
+
+class TestImageGeneratorDepth:
+    SOURCE_IMAGE_FILENAME = "reference_controlnet_dev_lora.png"
+
+    def test_image_generation_with_reference_image(self):
+        ImageGeneratorDepthTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_depth_dev_from_image.png",
+            output_image_path="output_depth_dev_from_image.png",
+            image_path=TestImageGeneratorDepth.SOURCE_IMAGE_FILENAME,
+            model_config=ModelConfig.dev_depth(),
+            steps=15,
+            seed=42,
+            height=512,
+            width=320,
+            prompt="Cartoon picture of einstein with a cane",
+        )
--- a/tests/image_generation/test_generate_image_fill.py
+++ b/tests/image_generation/test_generate_image_fill.py
@ -3,14 +3,13 @@ from tests.image_generation.helpers.image_generation_fill_test_helper import Ima


 class TestImageGeneratorFill:
-    OUTPUT_IMAGE_FILENAME = "output.png"
    SOURCE_IMAGE_FILENAME = "reference_dev_image_to_image_result.png"
    MASK_IMAGE_FILENAME = "mask.png"

-    def test_inpaint(self):
+    def test_fill(self):
        ImageGeneratorFillTestHelper.assert_matches_reference_image(
            reference_image_path="reference_dev_image_to_image_result_inpaint.png",
-            output_image_path=TestImageGeneratorFill.OUTPUT_IMAGE_FILENAME,
+            output_image_path="output_dev_image_to_image_result_inpaint.png",
            model_config=ModelConfig.dev_fill(),
            steps=15,
            seed=42,
--- a/tests/image_generation/test_generate_image_in_context.py
+++ b/tests/image_generation/test_generate_image_in_context.py
@ -3,12 +3,10 @@ from tests.image_generation.helpers.image_generation_in_context_test_helper impo


 class TestImageGeneratorInContext:
-    OUTPUT_IMAGE_FILENAME = "output.png"
-
    def test_in_context_lora_identity(self):
        ImageGeneratorInContextTestHelper.assert_matches_reference_image(
            reference_image_path="ic_lora_reference_in_context_identity.png",
-            output_image_path=TestImageGeneratorInContext.OUTPUT_IMAGE_FILENAME,
+            output_image_path="output_ic_lora_reference_in_context_identity.png",
            model_config=ModelConfig.dev(),
            steps=25,
            seed=42,
--- a/tests/image_generation/test_generate_image_redux.py
+++ b/tests/image_generation/test_generate_image_redux.py
@ -0,0 +1,19 @@
+from mflux import ModelConfig
+from tests.image_generation.helpers.image_generation_redux_test_helper import ImageGeneratorReduxTestHelper
+
+
+class TestImageGeneratorRedux:
+    SOURCE_IMAGE_FILENAME = "reference_dev_image_to_image_result.png"
+
+    def test_image_generation_redux(self):
+        ImageGeneratorReduxTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_redux_dev.png",
+            output_image_path="output_redux_dev.png",
+            model_config=ModelConfig.dev_redux(),
+            steps=15,
+            seed=42,
+            height=341,
+            width=768,
+            prompt="A delicious burger",
+            redux_image_path=TestImageGeneratorRedux.SOURCE_IMAGE_FILENAME,
+        )
--- a/tests/resources/ic_lora_reference_in_context_identity.png
+++ b/tests/resources/ic_lora_reference_in_context_identity.png
--- a/tests/resources/reference_controlnet_dev.png
+++ b/tests/resources/reference_controlnet_dev.png
--- a/tests/resources/reference_controlnet_dev_lora.png
+++ b/tests/resources/reference_controlnet_dev_lora.png
--- a/tests/resources/reference_controlnet_dev_lora_depth.png
+++ b/tests/resources/reference_controlnet_dev_lora_depth.png
--- a/tests/resources/reference_controlnet_schnell.png
+++ b/tests/resources/reference_controlnet_schnell.png
--- a/tests/resources/reference_depth_dev_from_image.png
+++ b/tests/resources/reference_depth_dev_from_image.png
--- a/tests/resources/reference_dev.png
+++ b/tests/resources/reference_dev.png
--- a/tests/resources/reference_dev_image_to_image_result.png
+++ b/tests/resources/reference_dev_image_to_image_result.png
--- a/tests/resources/reference_dev_image_to_image_result_inpaint.png
+++ b/tests/resources/reference_dev_image_to_image_result_inpaint.png
--- a/tests/resources/reference_dev_lora.png
+++ b/tests/resources/reference_dev_lora.png
--- a/tests/resources/reference_dev_lora_multiple.png
+++ b/tests/resources/reference_dev_lora_multiple.png
--- a/tests/resources/reference_redux_dev.png
+++ b/tests/resources/reference_redux_dev.png
--- a/tests/resources/reference_schnell.png
+++ b/tests/resources/reference_schnell.png