diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 5dc033a..a3a4c0d 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -48,6 +48,5 @@ jobs:
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
-          TEST_PYPI_API_TOKEN: ${{ secrets.TEST_PYPI_API_TOKEN }}
         run: |
           python -m mflux.release.release_script
diff --git a/CHANGELOG.md b/CHANGELOG.md
index da43955..80f20a8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,34 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.10.0] - 2025-08-02
+
+# MFLUX v.0.10.0 Release Notes
+
+### 🎨 Model Improvements
+
+- **FLUX.1 Krea [dev] Support!**
+
+### ✨ New Features
+
+- **5-bit Quantization Support**: Added support for 5-bit quantization as a new option alongside existing 3, 4, 6, and 8-bit quantization levels
+
+### 🔧 Improvements
+
+- **Enhanced Default Inference Steps**: Increased default inference steps for dev models from 14 to 25 for improved image quality
+- **Multiple Model Aliases Support**: Improved model configuration system to properly support multiple aliases per model, making model selection more flexible and robust
+
+### 🔧 Technical Requirements
+
+- **MLX Compatibility**: This release assumes MLX 0.27.0 and upwards for optimal performance and compatibility
+- **MLX Compatibility for test**: Fix MLX version to 0.27.1 for image generation tests
+- **Non-strict Weight Updates**: Explicitly added non-strict mode (`strict=False`) for weight updates to maintain compatibility with later MLX versions that enforce stricter weight validation by default
+
+### 👩‍💻 Developer Experience
+
+- **Streamlined Release Process**: Removed TestPyPi publishing step from release workflow for simplified deployment
+---
+
 ## [0.9.6] - 2025-07-20
 
 # MFLUX v.0.9.6 Release Notes
diff --git a/Makefile b/Makefile
index 09fe1b9..67e227e 100644
--- a/Makefile
+++ b/Makefile
@@ -93,7 +93,7 @@ check: ensure-ruff
 .PHONY: test
 test: ensure-pytest
 	# 🏗️ Running tests...
-	uv pip install mlx==0.26.1  # Install pinned MLX version specifically for testing
+	uv pip install mlx==0.27.1  # Install pinned MLX version specifically for testing
 	$(PYTHON) -m pytest
 	# ✅ Tests completed
 
diff --git a/README.md b/README.md
index 61c97ef..5344d7a 100644
--- a/README.md
+++ b/README.md
@@ -186,8 +186,8 @@ from mflux.config.config import Config
 
 # Load the model
 flux = Flux1.from_name(
-   model_name="schnell",  # "schnell" or "dev"
-   quantize=8,            # 4 or 8
+   model_name="schnell",  # "schnell", "dev", or "krea-dev"
+       quantize=8,            # 3, 4, 5, 6, or 8
 )
 
 # Generate an image
@@ -195,7 +195,7 @@ image = flux.generate_image(
    seed=2,
    prompt="Luxury food photograph",
    config=Config(
-      num_inference_steps=2,  # "schnell" works well with 2-4 steps, "dev" works well with 20-25 steps
+      num_inference_steps=2,  # "schnell" works well with 2-4 steps, "dev" and "krea-dev" work well with 20-25 steps
       height=1024,
       width=1024,
    )
@@ -214,6 +214,26 @@ For more advanced Python usage and additional configuration options, you can exp
 
 🔒 [FLUX.1-dev currently requires granted access to its Huggingface repo. For troubleshooting, see the issue tracker](https://github.com/filipstrand/mflux/issues/14) 🔒
 
+#### 🎨 FLUX.1 Krea [dev]: Enhanced Photorealism
+
+MFLUX now supports **FLUX.1 Krea [dev]**, an 'opinionated' text-to-image model developed in collaboration with [Krea AI](https://krea.ai). This model overcomes the oversaturated 'AI look' commonly found in generated images, achieving exceptional photorealism with distinctive aesthetics.
+This model can be used where the `dev` model is used, and it is available as `krea-dev` in MFLUX (also supports the alias `dev-krea`), this includes dreambooth fine-tuning.
+
+![image](src/mflux/assets/krea_dev_example.jpg)
+
+```sh
+mflux-generate \
+    --model krea-dev \
+    --prompt "A photo of a dog" \
+    --steps 25 \
+    --seed 2674888 \
+    -q 8 \
+    --height 1024 \
+    --width 1024
+```
+
+*Learn more about FLUX.1 Krea [dev] in the [official announcement](https://bfl.ai/announcements/flux-1-krea-dev).*
+
 #### 📜 Full list of Command-Line Arguments
 
 <details>
@@ -226,9 +246,9 @@ For more advanced Python usage and additional configuration options, you can exp
 
 - **`--prompt`** (required, `str`): Text description of the image to generate. Use `-` to read the prompt from stdin (e.g., `echo "A beautiful sunset" | mflux-generate --prompt -`).
 
-- **`--model`** or **`-m`** (required, `str`): Model to use for generation. Can be one of the official models (`"schnell"` or `"dev"`) or a HuggingFace repository ID for a compatible third-party model (e.g., `"Freepik/flux.1-lite-8B-alpha"`).
+- **`--model`** or **`-m`** (required, `str`): Model to use for generation. Can be one of the official models (`"schnell"`, `"dev"`, or `"krea-dev"`) or a HuggingFace repository ID for a compatible third-party model (e.g., `"Freepik/flux.1-lite-8B-alpha"`).
 
-- **`--base-model`** (optional, `str`, default: `None`): Specifies which base architecture a third-party model is derived from (`"schnell"` or `"dev"`). Required when using third-party models from HuggingFace.
+- **`--base-model`** (optional, `str`, default: `None`): Specifies which base architecture a third-party model is derived from (`"schnell"`, `"dev"`, or `"krea-dev"`). Required when using third-party models from HuggingFace.
 
 - **`--output`** (optional, `str`, default: `"image.png"`): Output image filename. If `--seed` or `--auto-seeds` establishes multiple seed values, the output filename will automatically be modified to include the seed value (e.g., `image_seed_42.png`).
 
@@ -246,7 +266,7 @@ For more advanced Python usage and additional configuration options, you can exp
 
 - **`--path`** (optional, `str`, default: `None`): Path to a local model on disk.
 
-- **`--quantize`** or **`-q`** (optional, `int`, default: `None`): [Quantization](#%EF%B8%8F-quantization) (choose between `3`, `4`, `6`, or `8` bits).
+- **`--quantize`** or **`-q`** (optional, `int`, default: `None`): [Quantization](#%EF%B8%8F-quantization) (choose between `3`, `4`, `5`, `6`, or `8` bits).
 
 - **`--lora-paths`** (optional, `[str]`, default: `None`): The paths to the [LoRA](#-LoRA) weights.
 
@@ -698,7 +718,7 @@ Luxury food photograph of an italian Linguine pasta alle vongole dish with lots
 
 ### 🗜️ Quantization
 
-MFLUX supports running FLUX in 3, 4, 6, or 8-bit quantized mode. Running a quantized version can greatly speed up the
+MFLUX supports running FLUX in 3, 4, 5, 6, or 8-bit quantized mode. Running a quantized version can greatly speed up the
 generation process and reduce the memory consumption by several gigabytes. [Quantized models also take up less disk space](#-size-comparisons-for-quantized-models).
 
 ```sh
@@ -726,9 +746,9 @@ For systems with limited RAM, you can also use the `--low-ram` option which redu
 
 The model sizes for both `schnell` and `dev` at various quantization levels are as follows:
 
-| 3 bit  | 4 bit  | 6 bit   | 8 bit   | Original (16 bit) |
-|--------|--------|---------|---------|-------------------|
-| 7.52GB | 9.61GB | 13.81GB | 18.01GB | 33.73GB           |
+| 3 bit  | 4 bit  | 5 bit   | 6 bit   | 8 bit   | Original (16 bit) |
+|--------|--------|---------|---------|---------|-------------------|
+| 7.52GB | 9.61GB | 11.71GB | 13.81GB | 18.01GB | 33.73GB           |
 
 
 #### 💾 Saving a quantized version to disk
diff --git a/pyproject.toml b/pyproject.toml
index 4e07f3a..d07d904 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -15,7 +15,7 @@ source-exclude = [
 
 [project]
 name = "mflux"
-version = "0.9.6"
+version = "0.10.0"
 description = "A MLX port of FLUX based on the Huggingface Diffusers implementation."
 readme = "README.md"
 keywords = ["diffusers", "flux", "mlx"]
@@ -26,7 +26,7 @@ requires-python = ">=3.10"
 dependencies = [
     "huggingface-hub>=0.24.5,<1.0",
     "matplotlib>=3.9.2,<4.0",
-    "mlx>=0.22.0,<=0.26.1",
+    "mlx>=0.27.0,<0.28.0",
     "numpy>=2.0.1,<3.0",
     "opencv-python>=4.10.0,<5.0",
     "piexif>=1.1.3,<2.0",
@@ -61,7 +61,7 @@ classifiers = [
 dev = [
     "pytest>=8.3.0,<9.0",
     "pytest-timer>=1.0,<2.0",
-    "mlx==0.26.1",  # Used ONLY during test runs to ensure deterministic test results
+    "mlx==0.27.1",  # Used ONLY during test runs to ensure deterministic test results
 ]
 
 [project.urls]
diff --git a/src/mflux/assets/krea_dev_example.jpg b/src/mflux/assets/krea_dev_example.jpg
new file mode 100644
index 0000000..4553586
Binary files /dev/null and b/src/mflux/assets/krea_dev_example.jpg differ
diff --git a/src/mflux/completions/README.md b/src/mflux/completions/README.md
index 39580a5..7d5eed0 100644
--- a/src/mflux/completions/README.md
+++ b/src/mflux/completions/README.md
@@ -24,7 +24,7 @@ This module provides ZSH shell completion support for all mflux CLI commands.
 - **Auto-completion for all 15 mflux commands**: Generate, train, save, upscale, and more
 - **Smart completions**: 
   - Model names (dev, schnell, dev-fill, or HuggingFace repos)
-  - Quantization levels (3, 4, 6, 8)
+  - Quantization levels (3, 4, 5, 6, 8)
   - LoRA styles (storyboard, portrait, etc.)
   - File paths with native completion
   - Common percentage values for battery limit
diff --git a/src/mflux/config/model_config.py b/src/mflux/config/model_config.py
index d5eb4fe..8136da3 100644
--- a/src/mflux/config/model_config.py
+++ b/src/mflux/config/model_config.py
@@ -7,7 +7,7 @@ from mflux.error.error import InvalidBaseModel, ModelConfigError
 class ModelConfig:
     def __init__(
         self,
-        alias: str | None,
+        aliases: list[str],
         model_name: str,
         base_model: str | None,
         controlnet_model: str | None,
@@ -18,7 +18,7 @@ class ModelConfig:
         requires_sigma_shift: bool,
         priority: int,
     ):
-        self.alias = alias
+        self.aliases = aliases
         self.model_name = model_name
         self.base_model = base_model
         self.controlnet_model = controlnet_model
@@ -79,49 +79,55 @@ class ModelConfig:
     def dev_fill_catvton() -> "ModelConfig":
         return AVAILABLE_MODELS["dev-fill-catvton"]
 
+    @staticmethod
+    @lru_cache
+    def krea_dev() -> "ModelConfig":
+        return AVAILABLE_MODELS["krea-dev"]
+
     def x_embedder_input_dim(self) -> int:
-        if self.alias and "dev-fill" in self.alias:
+        if "Fill" in self.model_name:
             return 384
-        if self.alias == "dev-depth":
+        if "Depth" in self.model_name:
             return 128
         else:
             return 64
 
     def is_canny(self) -> bool:
-        return self.alias == "dev-controlnet-canny" or self.alias == "schnell-controlnet-canny"
+        return self.controlnet_model is not None and "Canny" in self.controlnet_model
 
     @staticmethod
     def from_name(
         model_name: str,
-        base_model: Literal["dev", "schnell", "dev-fill"] | None = None,
+        base_model: Literal["dev", "schnell", "krea-dev"] | None = None,
     ) -> "ModelConfig":
         # 0. Get all base models (where base_model is None) sorted by priority
         base_models = sorted(
             [model for model in AVAILABLE_MODELS.values() if model.base_model is None], key=lambda x: x.priority
         )
 
-        # 1. Check if model_name matches any base model's alias or full name
+        # 1. Check if model_name matches any base model's aliases or full name
         for base in base_models:
-            if model_name in (base.alias, base.model_name):
+            if model_name == base.model_name or model_name in base.aliases:
                 return base
 
         # 2. Validate explicit base_model
         allowed_names = []
         for base in base_models:
-            allowed_names.extend([base.alias, base.model_name])
+            allowed_names.extend(base.aliases + [base.model_name])
         if base_model and base_model not in allowed_names:
             raise InvalidBaseModel(f"Invalid base_model. Choose one of {allowed_names}")
 
         # 3. Determine the base model (explicit or inferred)
         if base_model:
-            # Find by explicit base_model name
-            default_base = next((b for b in base_models if base_model in (b.alias, b.model_name)), None)
+            # Find by explicit base_model name (check all aliases)
+            default_base = next((b for b in base_models if base_model == b.model_name or base_model in b.aliases), None)
         else:
             # Infer from model_name substring - prefer longer matches (more specific)
-            matching_bases = [b for b in base_models if b.alias and b.alias in model_name]
+            matching_bases = [(b, alias) for b in base_models for alias in b.aliases if alias and alias in model_name]
+
             if matching_bases:
                 # Sort by alias length descending, then by priority ascending
-                default_base = sorted(matching_bases, key=lambda x: (-len(x.alias), x.priority))[0]
+                default_base = sorted(matching_bases, key=lambda x: (-len(x[1]), x[0].priority))[0][0]
             else:
                 default_base = None
             if not default_base:
@@ -129,7 +135,7 @@ class ModelConfig:
 
         # 4. Construct the config
         return ModelConfig(
-            alias=default_base.alias,
+            aliases=default_base.aliases,
             model_name=model_name,
             base_model=default_base.model_name,
             controlnet_model=default_base.controlnet_model,
@@ -144,7 +150,7 @@ class ModelConfig:
 
 AVAILABLE_MODELS = {
     "dev": ModelConfig(
-        alias="dev",
+        aliases=["dev"],
         model_name="black-forest-labs/FLUX.1-dev",
         base_model=None,
         controlnet_model=None,
@@ -156,7 +162,7 @@ AVAILABLE_MODELS = {
         priority=0,
     ),
     "schnell": ModelConfig(
-        alias="schnell",
+        aliases=["schnell"],
         model_name="black-forest-labs/FLUX.1-schnell",
         base_model=None,
         controlnet_model=None,
@@ -168,7 +174,7 @@ AVAILABLE_MODELS = {
         priority=1,
     ),
     "dev-kontext": ModelConfig(
-        alias="dev-kontext",
+        aliases=["dev-kontext"],
         model_name="black-forest-labs/FLUX.1-Kontext-dev",
         base_model=None,
         controlnet_model=None,
@@ -180,7 +186,7 @@ AVAILABLE_MODELS = {
         priority=2,
     ),
     "dev-fill": ModelConfig(
-        alias="dev-fill",
+        aliases=["dev-fill"],
         model_name="black-forest-labs/FLUX.1-Fill-dev",
         base_model=None,
         controlnet_model=None,
@@ -192,7 +198,7 @@ AVAILABLE_MODELS = {
         priority=3,
     ),
     "dev-redux": ModelConfig(
-        alias="dev-redux",
+        aliases=["dev-redux"],
         model_name="black-forest-labs/FLUX.1-Redux-dev",
         base_model=None,
         controlnet_model=None,
@@ -204,7 +210,7 @@ AVAILABLE_MODELS = {
         priority=4,
     ),
     "dev-depth": ModelConfig(
-        alias="dev-depth",
+        aliases=["dev-depth"],
         model_name="black-forest-labs/FLUX.1-Depth-dev",
         base_model=None,
         controlnet_model=None,
@@ -216,7 +222,7 @@ AVAILABLE_MODELS = {
         priority=5,
     ),
     "dev-controlnet-canny": ModelConfig(
-        alias="dev-controlnet-canny",
+        aliases=["dev-controlnet-canny"],
         model_name="black-forest-labs/FLUX.1-dev",
         base_model=None,
         controlnet_model="InstantX/FLUX.1-dev-Controlnet-Canny",
@@ -228,7 +234,7 @@ AVAILABLE_MODELS = {
         priority=6,
     ),
     "schnell-controlnet-canny": ModelConfig(
-        alias="schnell-controlnet-canny",
+        aliases=["schnell-controlnet-canny"],
         model_name="black-forest-labs/FLUX.1-schnell",
         base_model=None,
         controlnet_model="InstantX/FLUX.1-dev-Controlnet-Canny",
@@ -240,7 +246,7 @@ AVAILABLE_MODELS = {
         priority=7,
     ),
     "dev-controlnet-upscaler": ModelConfig(
-        alias="dev-controlnet-upscaler",
+        aliases=["dev-controlnet-upscaler"],
         model_name="black-forest-labs/FLUX.1-dev",
         base_model=None,
         controlnet_model="jasperai/Flux.1-dev-Controlnet-Upscaler",
@@ -252,7 +258,7 @@ AVAILABLE_MODELS = {
         priority=8,
     ),
     "dev-fill-catvton": ModelConfig(
-        alias="dev-fill-catvton",
+        aliases=["dev-fill-catvton"],
         model_name="black-forest-labs/FLUX.1-Fill-dev",
         base_model=None,
         controlnet_model=None,
@@ -263,4 +269,16 @@ AVAILABLE_MODELS = {
         requires_sigma_shift=False,  # Not sure why, but produced better results this way...
         priority=9,
     ),
+    "krea-dev": ModelConfig(
+        aliases=["krea-dev", "dev-krea"],
+        model_name="black-forest-labs/FLUX.1-Krea-dev",
+        base_model=None,
+        controlnet_model=None,
+        custom_transformer_model=None,
+        num_train_steps=1000,
+        max_sequence_length=512,
+        supports_guidance=True,
+        requires_sigma_shift=True,
+        priority=10,
+    ),
 }
diff --git a/src/mflux/models/depth_pro/depth_pro_initializer.py b/src/mflux/models/depth_pro/depth_pro_initializer.py
index 7d69527..6888fa4 100644
--- a/src/mflux/models/depth_pro/depth_pro_initializer.py
+++ b/src/mflux/models/depth_pro/depth_pro_initializer.py
@@ -18,7 +18,7 @@ class DepthProInitializer:
         WeightHandlerDepthPro.reshape_transposed_convolution_weights(depth_pro_weights)
 
         # 2. Assign the weights to the model
-        depth_pro_model.update(depth_pro_weights.weights)
+        depth_pro_model.update(depth_pro_weights.weights, strict=False)
 
         # 3. Optionally quantize the model
         if quantize:
diff --git a/src/mflux/release/pypi_publisher.py b/src/mflux/release/pypi_publisher.py
index 18cb99f..24071e9 100644
--- a/src/mflux/release/pypi_publisher.py
+++ b/src/mflux/release/pypi_publisher.py
@@ -2,7 +2,6 @@ import shutil
 import sys
 import time
 from pathlib import Path
-from typing import Optional
 
 import requests
 from twine.commands import upload
@@ -34,9 +33,9 @@ class PyPIPublisher:
         )
 
     @staticmethod
-    def version_exists_on_pypi(package_name: str, version: str, test_pypi: bool = False) -> bool:
+    def version_exists_on_pypi(package_name: str, version: str) -> bool:
         print("🔍 Checking if version already exists on PyPI...")
-        repo_url = "https://test.pypi.org/pypi" if test_pypi else "https://pypi.org/pypi"
+        repo_url = "https://pypi.org/pypi"
         url = f"{repo_url}/{package_name}/json"
 
         for attempt in range(3):
@@ -77,21 +76,6 @@ class PyPIPublisher:
 
         raise RuntimeError("Unexpected error in version_exists_on_pypi")
 
-    @staticmethod
-    def publish_to_test_pypi(test_pypi_token: Optional[str], package_name: str, version: str) -> None:
-        if not test_pypi_token:
-            print("⚠️  Test PyPI token not provided, skipping Test PyPI upload")
-            return
-
-        PyPIPublisher._upload_to_pypi(
-            token=test_pypi_token,
-            repository="testpypi",
-            display_name="Test PyPI",
-            package_name=package_name,
-            version=version,
-            optional=True,
-        )
-
     @staticmethod
     def publish_to_pypi(pypi_token: str, package_name: str, version: str) -> None:
         PyPIPublisher._upload_to_pypi(
diff --git a/src/mflux/release/release_manager.py b/src/mflux/release/release_manager.py
index 89229bd..c3b8248 100644
--- a/src/mflux/release/release_manager.py
+++ b/src/mflux/release/release_manager.py
@@ -1,5 +1,3 @@
-from typing import Optional
-
 import requests
 
 from mflux.release.changelog_parser import ChangelogParser
@@ -15,7 +13,6 @@ class ReleaseManager:
     def create_release(
         github_token: str,
         pypi_token: str,
-        test_pypi_token: Optional[str] = None,
         github_repo: str = "filipstrand/mflux",
         package_name: str = "mflux",
     ) -> None:
@@ -40,7 +37,6 @@ class ReleaseManager:
         # 4. Handle PyPI publishing FIRST (before creating git artifacts)
         if ReleaseManager._should_publish_to_pypi(git_tag_exists, github_release_exists, package_name, version):
             PyPIPublisher.build_and_verify_package()
-            PyPIPublisher.publish_to_test_pypi(test_pypi_token, package_name, version)
             PyPIPublisher.publish_to_pypi(pypi_token, package_name, version)
 
         # 5. Create git tag if needed
@@ -76,7 +72,7 @@ class ReleaseManager:
 
         # Check if version already exists on PyPI
         try:
-            if PyPIPublisher.version_exists_on_pypi(package_name, version, test_pypi=False):
+            if PyPIPublisher.version_exists_on_pypi(package_name, version):
                 return False
         except (requests.RequestException, ValueError, OSError) as e:
             print(f"❌ Failed to check PyPI version: {e}")
diff --git a/src/mflux/release/release_script.py b/src/mflux/release/release_script.py
index 74ef8ce..9117090 100644
--- a/src/mflux/release/release_script.py
+++ b/src/mflux/release/release_script.py
@@ -11,14 +11,12 @@ def main():
     parser = argparse.ArgumentParser(description="MFLUX Release Management")
     parser.add_argument("--github-token", help="GitHub token")
     parser.add_argument("--pypi-token", help="PyPI API token")
-    parser.add_argument("--test-pypi-token", help="Test PyPI API token (optional)")
     args = parser.parse_args()
 
     try:
         ReleaseManager.create_release(
             github_token=args.github_token or os.getenv("GITHUB_TOKEN"),
             pypi_token=args.pypi_token or os.getenv("PYPI_API_TOKEN"),
-            test_pypi_token=args.test_pypi_token or os.getenv("TEST_PYPI_API_TOKEN"),
         )
     except (ValueError, FileNotFoundError, requests.RequestException) as e:
         print(f"❌ Release failed: {e}")
diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py
index 4131b4a..f904b90 100644
--- a/src/mflux/ui/cli/parsers.py
+++ b/src/mflux/ui/cli/parsers.py
@@ -289,7 +289,7 @@ class CommandLineParser(argparse.ArgumentParser):
                 self.error("Either --prompt or --prompt-file argument is 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, 14)
+            namespace.steps = ui_defaults.MODEL_INFERENCE_STEPS.get(namespace.model, 25)
 
         # In-context edit specific validations
         if getattr(self, 'supports_in_context_edit', False):
diff --git a/src/mflux/ui/defaults.py b/src/mflux/ui/defaults.py
index 3d77f0f..76fe6ea 100644
--- a/src/mflux/ui/defaults.py
+++ b/src/mflux/ui/defaults.py
@@ -17,16 +17,13 @@ GUIDANCE_SCALE_KONTEXT = 2.5
 HEIGHT, WIDTH = 1024, 1024
 MAX_PIXELS_WARNING_THRESHOLD = 2048 * 2048
 IMAGE_STRENGTH = 0.4
-MODEL_CHOICES = ["dev", "schnell", "dev-kontext", "dev-fill"]
+MODEL_CHOICES = ["dev", "schnell", "krea-dev"]
 MODEL_INFERENCE_STEPS = {
-    "dev": 14,
-    "dev-fill": 14,
-    "dev-depth": 14,
-    "dev-redux": 14,
-    "dev-kontext": 14,
+    "dev": 25,
+    "krea-dev": 25,
     "schnell": 4,
 }
-QUANTIZE_CHOICES = [3, 4, 6, 8]
+QUANTIZE_CHOICES = [3, 5, 4, 6, 8]
 
 
 def _migrate_legacy_cache(new_cache_dir: Path) -> None:
diff --git a/src/mflux/weights/weight_util.py b/src/mflux/weights/weight_util.py
index eddf693..5745c47 100644
--- a/src/mflux/weights/weight_util.py
+++ b/src/mflux/weights/weight_util.py
@@ -57,19 +57,19 @@ class WeightUtil:
         transformer_controlnet: nn.Module,
     ) -> int | None:
         if weights.meta_data.quantization_level is None and quantize_arg is None:
-            transformer_controlnet.update(weights.controlnet_transformer)
+            transformer_controlnet.update(weights.controlnet_transformer, strict=False)
             return None
 
         if weights.meta_data.quantization_level is None and quantize_arg is not None:
             bits = quantize_arg
-            transformer_controlnet.update(weights.controlnet_transformer)
+            transformer_controlnet.update(weights.controlnet_transformer, strict=False)
             QuantizationUtil.quantize_controlnet(bits, weights, transformer_controlnet)
             return bits
 
         if weights.meta_data.quantization_level is not None:
             bits = weights.meta_data.quantization_level
             QuantizationUtil.quantize_controlnet(bits, weights, transformer_controlnet)
-            transformer_controlnet.update(weights.controlnet_transformer)
+            transformer_controlnet.update(weights.controlnet_transformer, strict=False)
             return bits
 
     @staticmethod
@@ -80,10 +80,10 @@ class WeightUtil:
         t5_text_encoder: nn.Module,
         clip_text_encoder: nn.Module,
     ):
-        vae.update(weights.vae)
-        transformer.update(weights.transformer)
-        t5_text_encoder.update(weights.t5_encoder)
-        clip_text_encoder.update(weights.clip_encoder)
+        vae.update(weights.vae, strict=False)
+        transformer.update(weights.transformer, strict=False)
+        t5_text_encoder.update(weights.t5_encoder, strict=False)
+        clip_text_encoder.update(weights.clip_encoder, strict=False)
 
     @staticmethod
     def _set_redux_model_weights(
@@ -91,8 +91,8 @@ class WeightUtil:
         redux_encoder: nn.Module,
         siglip_vision_transformer: nn.Module,
     ):
-        redux_encoder.update(weights.redux_encoder)
-        siglip_vision_transformer.update(weights.siglip["vision_model"])
+        redux_encoder.update(weights.redux_encoder, strict=False)
+        siglip_vision_transformer.update(weights.siglip["vision_model"], strict=False)
 
     @staticmethod
     def set_redux_weights_and_quantize(
diff --git a/tests/arg_parser/test_cli_argparser.py b/tests/arg_parser/test_cli_argparser.py
index 6a6b82f..039ec4a 100644
--- a/tests/arg_parser/test_cli_argparser.py
+++ b/tests/arg_parser/test_cli_argparser.py
@@ -410,7 +410,7 @@ def test_steps_arg(mflux_generate_parser, mflux_generate_minimal_argv, base_meta
     # test user default value for dev
     with patch("sys.argv", mflux_generate_minimal_argv + ["--model", "dev"]):
         args = mflux_generate_parser.parse_args()
-        assert args.steps == 14
+        assert args.steps == 25
 
     # test user default value for schnell
     with patch("sys.argv", mflux_generate_minimal_argv + ["--model", "schnell"]):
diff --git a/tests/image_generation/test_generate_image.py b/tests/image_generation/test_generate_image.py
index 89b8eab..b67c15b 100644
--- a/tests/image_generation/test_generate_image.py
+++ b/tests/image_generation/test_generate_image.py
@@ -57,8 +57,8 @@ 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",
-            output_image_path="output_dev_image_to_image_result.png",
+            reference_image_path="reference_dev_image_to_image.png",
+            output_image_path="output_dev_image_to_image.png",
             image_strength=0.4,
             model_config=ModelConfig.dev(),
             steps=8,
diff --git a/tests/image_generation/test_generate_image_controlnet.py b/tests/image_generation/test_generate_image_controlnet.py
index 311e168..6732205 100644
--- a/tests/image_generation/test_generate_image_controlnet.py
+++ b/tests/image_generation/test_generate_image_controlnet.py
@@ -3,13 +3,11 @@ from tests.image_generation.helpers.image_generation_controlnet_test_helper impo
 
 
 class TestImageGeneratorControlnet:
-    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="output_controlnet_schnell.png",
-            controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
+            controlnet_image_path="controlnet_reference.png",
             model_config=ModelConfig.schnell_controlnet_canny(),
             steps=2,
             seed=43,
@@ -23,7 +21,7 @@ class TestImageGeneratorControlnet:
         ImageGeneratorControlnetTestHelper.assert_matches_reference_image(
             reference_image_path="reference_controlnet_dev.png",
             output_image_path="output_controlnet_dev.png",
-            controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
+            controlnet_image_path="controlnet_reference.png",
             model_config=ModelConfig.dev_controlnet_canny(),
             steps=15,
             seed=42,
@@ -37,7 +35,7 @@ class TestImageGeneratorControlnet:
         ImageGeneratorControlnetTestHelper.assert_matches_reference_image(
             reference_image_path="reference_controlnet_dev_lora.png",
             output_image_path="output_controlnet_dev_lora.png",
-            controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
+            controlnet_image_path="controlnet_reference.png",
             model_config=ModelConfig.dev_controlnet_canny(),
             steps=15,
             seed=43,
@@ -52,7 +50,7 @@ class TestImageGeneratorControlnet:
     def test_image_upscaling(self):
         ImageGeneratorControlnetTestHelper.assert_matches_reference_image(
             reference_image_path="reference_upscaled.png",
-            output_image_path="output_upscaler.png",
+            output_image_path="output_upscaled.png",
             controlnet_image_path="low_res.jpg",
             model_config=ModelConfig.dev_controlnet_upscaler(),
             steps=20,
diff --git a/tests/image_generation/test_generate_image_depth.py b/tests/image_generation/test_generate_image_depth.py
index a4c0842..f8d4068 100644
--- a/tests/image_generation/test_generate_image_depth.py
+++ b/tests/image_generation/test_generate_image_depth.py
@@ -3,13 +3,11 @@ from tests.image_generation.helpers.image_generation_depth_test_helper import Im
 
 
 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,
+            image_path="reference_controlnet_dev_lora.png",
             model_config=ModelConfig.dev_depth(),
             steps=15,
             seed=42,
diff --git a/tests/image_generation/test_generate_image_fill.py b/tests/image_generation/test_generate_image_fill.py
index 9b849c8..d4fb718 100644
--- a/tests/image_generation/test_generate_image_fill.py
+++ b/tests/image_generation/test_generate_image_fill.py
@@ -3,19 +3,16 @@ from tests.image_generation.helpers.image_generation_fill_test_helper import Ima
 
 
 class TestImageGeneratorFill:
-    SOURCE_IMAGE_FILENAME = "reference_dev_image_to_image_result.png"
-    MASK_IMAGE_FILENAME = "mask.png"
-
     def test_fill(self):
         ImageGeneratorFillTestHelper.assert_matches_reference_image(
-            reference_image_path="reference_dev_image_to_image_result_inpaint.png",
-            output_image_path="output_dev_image_to_image_result_inpaint.png",
+            reference_image_path="reference_dev_image_to_image_fill.png",
+            output_image_path="output_dev_image_to_image_fill.png",
             model_config=ModelConfig.dev_fill(),
             steps=15,
             seed=42,
             height=341,
             width=768,
             prompt="A burger and crispy golden french fries next to it",
-            image_path=TestImageGeneratorFill.SOURCE_IMAGE_FILENAME,
-            masked_image_path=TestImageGeneratorFill.MASK_IMAGE_FILENAME,
+            image_path="reference_dev_image_to_image.png",
+            masked_image_path="mask.png",
         )
diff --git a/tests/image_generation/test_generate_image_in_context.py b/tests/image_generation/test_generate_image_in_context.py
index 0c0684f..44912ef 100644
--- a/tests/image_generation/test_generate_image_in_context.py
+++ b/tests/image_generation/test_generate_image_in_context.py
@@ -6,8 +6,8 @@ from tests.image_generation.helpers.image_generation_in_context_test_helper impo
 class TestImageGeneratorInContext:
     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="output_ic_lora_reference_in_context_identity.png",
+            reference_image_path="reference_ic_lora_identity.png",
+            output_image_path="output_ic_lora_identity.png",
             model_config=ModelConfig.dev(),
             steps=25,
             seed=42,
@@ -20,7 +20,7 @@ class TestImageGeneratorInContext:
 
     def test_ic_edit_with_instruction(self):
         ImageGeneratorICEditTestHelper.assert_matches_reference_image(
-            reference_image_path="ic_edit_reference.png",
+            reference_image_path="reference_ic_edit.png",
             output_image_path="output_ic_edit.png",
             steps=20,
             seed=2799665,
diff --git a/tests/image_generation/test_generate_image_redux.py b/tests/image_generation/test_generate_image_redux.py
index f1b5fa6..a4d9008 100644
--- a/tests/image_generation/test_generate_image_redux.py
+++ b/tests/image_generation/test_generate_image_redux.py
@@ -3,8 +3,6 @@ from tests.image_generation.helpers.image_generation_redux_test_helper import Im
 
 
 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",
@@ -15,5 +13,5 @@ class TestImageGeneratorRedux:
             height=341,
             width=768,
             prompt="A delicious burger",
-            redux_image_path=TestImageGeneratorRedux.SOURCE_IMAGE_FILENAME,
+            redux_image_path="reference_dev_image_to_image.png",
         )
diff --git a/tests/model_config/test_model_config.py b/tests/model_config/test_model_config.py
index 4006899..7345552 100644
--- a/tests/model_config/test_model_config.py
+++ b/tests/model_config/test_model_config.py
@@ -6,8 +6,7 @@ from mflux.error.error import InvalidBaseModel, ModelConfigError
 
 def test_bfl_dev():
     model = ModelConfig.from_name("dev")
-    assert model.alias == "dev"
-    assert model.model_name.startswith("black-forest-labs/")
+    assert model.model_name == "black-forest-labs/FLUX.1-dev"
     assert model.max_sequence_length == 512
     assert model.num_train_steps == 1000
     assert model.supports_guidance is True
@@ -16,8 +15,34 @@ def test_bfl_dev():
 
 def test_bfl_dev_full_name():
     model = ModelConfig.from_name("black-forest-labs/FLUX.1-dev")
-    assert model.alias == "dev"
-    assert model.model_name.startswith("black-forest-labs/")
+    assert model.model_name == "black-forest-labs/FLUX.1-dev"
+    assert model.max_sequence_length == 512
+    assert model.num_train_steps == 1000
+    assert model.supports_guidance is True
+    assert model.requires_sigma_shift is True
+
+
+def test_bfl_krea_dev():
+    model = ModelConfig.from_name("krea-dev")
+    assert model.model_name == "black-forest-labs/FLUX.1-Krea-dev"
+    assert model.max_sequence_length == 512
+    assert model.num_train_steps == 1000
+    assert model.supports_guidance is True
+    assert model.requires_sigma_shift is True
+
+
+def test_bfl_dev_krea_alias():
+    model = ModelConfig.from_name("dev-krea")
+    assert model.model_name == "black-forest-labs/FLUX.1-Krea-dev"
+    assert model.max_sequence_length == 512
+    assert model.num_train_steps == 1000
+    assert model.supports_guidance is True
+    assert model.requires_sigma_shift is True
+
+
+def test_bfl_krea_dev_full_name():
+    model = ModelConfig.from_name("black-forest-labs/FLUX.1-Krea-dev")
+    assert model.model_name == "black-forest-labs/FLUX.1-Krea-dev"
     assert model.max_sequence_length == 512
     assert model.num_train_steps == 1000
     assert model.supports_guidance is True
@@ -26,8 +51,7 @@ def test_bfl_dev_full_name():
 
 def test_bfl_schnell():
     model = ModelConfig.from_name("schnell")
-    assert model.alias == "schnell"
-    assert model.model_name.startswith("black-forest-labs/")
+    assert model.model_name == "black-forest-labs/FLUX.1-schnell"
     assert model.max_sequence_length == 256
     assert model.num_train_steps == 1000
     assert model.supports_guidance is False
@@ -36,8 +60,7 @@ def test_bfl_schnell():
 
 def test_bfl_schnell_full_name():
     model = ModelConfig.from_name("black-forest-labs/FLUX.1-schnell")
-    assert model.alias == "schnell"
-    assert model.model_name.startswith("black-forest-labs/")
+    assert model.model_name == "black-forest-labs/FLUX.1-schnell"
     assert model.max_sequence_length == 256
     assert model.num_train_steps == 1000
     assert model.supports_guidance is False
@@ -46,7 +69,6 @@ def test_bfl_schnell_full_name():
 
 def test_bfl_dev_fill():
     model = ModelConfig.from_name("dev-fill")
-    assert model.alias == "dev-fill"
     assert model.model_name == "black-forest-labs/FLUX.1-Fill-dev"
     assert model.max_sequence_length == 512
     assert model.num_train_steps == 1000
@@ -56,7 +78,6 @@ def test_bfl_dev_fill():
 
 def test_bfl_dev_fill_full_name():
     model = ModelConfig.from_name("black-forest-labs/FLUX.1-Fill-dev")
-    assert model.alias == "dev-fill"
     assert model.model_name == "black-forest-labs/FLUX.1-Fill-dev"
     assert model.max_sequence_length == 512
     assert model.num_train_steps == 1000
@@ -66,7 +87,6 @@ def test_bfl_dev_fill_full_name():
 
 def test_bfl_dev_depth():
     model = ModelConfig.from_name("dev-depth")
-    assert model.alias == "dev-depth"
     assert model.model_name == "black-forest-labs/FLUX.1-Depth-dev"
     assert model.max_sequence_length == 512
     assert model.num_train_steps == 1000
@@ -76,7 +96,6 @@ def test_bfl_dev_depth():
 
 def test_bfl_dev_depth_full_name():
     model = ModelConfig.from_name("black-forest-labs/FLUX.1-Depth-dev")
-    assert model.alias == "dev-depth"
     assert model.model_name == "black-forest-labs/FLUX.1-Depth-dev"
     assert model.max_sequence_length == 512
     assert model.num_train_steps == 1000
@@ -86,7 +105,6 @@ def test_bfl_dev_depth_full_name():
 
 def test_bfl_dev_redux():
     model = ModelConfig.from_name("dev-redux")
-    assert model.alias == "dev-redux"
     assert model.model_name == "black-forest-labs/FLUX.1-Redux-dev"
     assert model.max_sequence_length == 512
     assert model.num_train_steps == 1000
@@ -96,7 +114,6 @@ def test_bfl_dev_redux():
 
 def test_bfl_dev_redux_full_name():
     model = ModelConfig.from_name("black-forest-labs/FLUX.1-Redux-dev")
-    assert model.alias == "dev-redux"
     assert model.model_name == "black-forest-labs/FLUX.1-Redux-dev"
     assert model.max_sequence_length == 512
     assert model.num_train_steps == 1000
@@ -106,7 +123,6 @@ def test_bfl_dev_redux_full_name():
 
 def test_bfl_dev_controlnet_canny():
     model = ModelConfig.from_name("dev-controlnet-canny")
-    assert model.alias == "dev-controlnet-canny"
     assert model.model_name == "black-forest-labs/FLUX.1-dev"
     assert model.controlnet_model == "InstantX/FLUX.1-dev-Controlnet-Canny"
     assert model.max_sequence_length == 512
@@ -118,7 +134,6 @@ def test_bfl_dev_controlnet_canny():
 
 def test_bfl_schnell_controlnet_canny():
     model = ModelConfig.from_name("schnell-controlnet-canny")
-    assert model.alias == "schnell-controlnet-canny"
     assert model.model_name == "black-forest-labs/FLUX.1-schnell"
     assert model.controlnet_model == "InstantX/FLUX.1-dev-Controlnet-Canny"
     assert model.max_sequence_length == 256
@@ -130,7 +145,6 @@ def test_bfl_schnell_controlnet_canny():
 
 def test_bfl_dev_controlnet_upscaler():
     model = ModelConfig.from_name("dev-controlnet-upscaler")
-    assert model.alias == "dev-controlnet-upscaler"
     assert model.model_name == "black-forest-labs/FLUX.1-dev"
     assert model.controlnet_model == "jasperai/Flux.1-dev-Controlnet-Upscaler"
     assert model.max_sequence_length == 512
@@ -142,7 +156,7 @@ def test_bfl_dev_controlnet_upscaler():
 
 def test_community_dev_fill_implicit_base_model():
     model = ModelConfig.from_name("acme-lab/some-dev-fill-model")
-    assert model.alias == "dev-fill"
+    assert model.base_model == "black-forest-labs/FLUX.1-Fill-dev"
     assert model.max_sequence_length == 512
     assert model.num_train_steps == 1000
     assert model.supports_guidance is True
@@ -151,7 +165,7 @@ def test_community_dev_fill_implicit_base_model():
 
 def test_community_dev_fill_explicit_base_model():
     model = ModelConfig.from_name("acme-lab/some-model", base_model="dev-fill")
-    assert model.alias == "dev-fill"
+    assert model.base_model == "black-forest-labs/FLUX.1-Fill-dev"
     assert model.max_sequence_length == 512
     assert model.num_train_steps == 1000
     assert model.supports_guidance is True
@@ -160,7 +174,6 @@ def test_community_dev_fill_explicit_base_model():
 
 def test_implicit_base_model_prefers_dev_fill_over_dev():
     model = ModelConfig.from_name("acme-lab/dev-fill-based-model")
-    assert model.alias == "dev-fill"
     assert model.base_model == "black-forest-labs/FLUX.1-Fill-dev"
     assert model.max_sequence_length == 512
     assert model.num_train_steps == 1000
@@ -169,7 +182,7 @@ def test_implicit_base_model_prefers_dev_fill_over_dev():
 
 def test_community_dev_implicit_base_model():
     model = ModelConfig.from_name("acme-lab/some-awesome-dev-model")
-    assert model.alias == "dev"
+    assert model.base_model == "black-forest-labs/FLUX.1-dev"
     assert model.max_sequence_length == 512
     assert model.num_train_steps == 1000
     assert model.supports_guidance is True
@@ -178,7 +191,7 @@ def test_community_dev_implicit_base_model():
 
 def test_community_schnell_implicit_base_model():
     model = ModelConfig.from_name("acme-lab/some-quick-schnell-model")
-    assert model.alias == "schnell"
+    assert model.base_model == "black-forest-labs/FLUX.1-schnell"
     assert model.max_sequence_length == 256
     assert model.num_train_steps == 1000
     assert model.supports_guidance is False
@@ -187,7 +200,6 @@ def test_community_schnell_implicit_base_model():
 
 def test_community_dev_explicit_base_model():
     model = ModelConfig.from_name("acme-lab/some-awesome-model", base_model="dev")
-    assert model.alias == "dev"
     assert model.base_model == "black-forest-labs/FLUX.1-dev"
     assert model.max_sequence_length == 512
     assert model.num_train_steps == 1000
diff --git a/tests/resources/ic_edit_reference.png b/tests/resources/ic_edit_reference.png
deleted file mode 100644
index a483b6d..0000000
Binary files a/tests/resources/ic_edit_reference.png and /dev/null differ
diff --git a/tests/resources/ic_lora_reference_in_context_identity.png b/tests/resources/ic_lora_reference_in_context_identity.png
deleted file mode 100644
index 6698813..0000000
Binary files a/tests/resources/ic_lora_reference_in_context_identity.png and /dev/null differ
diff --git a/tests/resources/reference_concept_from_image_schnell_heatmap.png b/tests/resources/reference_concept_from_image_schnell_heatmap.png
index 6ca1e4d..9da995f 100644
Binary files a/tests/resources/reference_concept_from_image_schnell_heatmap.png and b/tests/resources/reference_concept_from_image_schnell_heatmap.png differ
diff --git a/tests/resources/reference_concept_schnell_heatmap.png b/tests/resources/reference_concept_schnell_heatmap.png
index b491e0c..c20fa89 100644
Binary files a/tests/resources/reference_concept_schnell_heatmap.png and b/tests/resources/reference_concept_schnell_heatmap.png differ
diff --git a/tests/resources/reference_controlnet_dev.png b/tests/resources/reference_controlnet_dev.png
index 8c51631..51aec6f 100644
Binary files a/tests/resources/reference_controlnet_dev.png and b/tests/resources/reference_controlnet_dev.png differ
diff --git a/tests/resources/reference_controlnet_dev_lora.png b/tests/resources/reference_controlnet_dev_lora.png
index cdcde61..ab14cfa 100644
Binary files a/tests/resources/reference_controlnet_dev_lora.png and b/tests/resources/reference_controlnet_dev_lora.png differ
diff --git a/tests/resources/reference_controlnet_dev_lora_depth.png b/tests/resources/reference_controlnet_dev_lora_depth.png
index b03c2dc..0cd2da4 100644
Binary files a/tests/resources/reference_controlnet_dev_lora_depth.png and b/tests/resources/reference_controlnet_dev_lora_depth.png differ
diff --git a/tests/resources/reference_controlnet_schnell.png b/tests/resources/reference_controlnet_schnell.png
index b087930..1c5a0aa 100644
Binary files a/tests/resources/reference_controlnet_schnell.png and b/tests/resources/reference_controlnet_schnell.png differ
diff --git a/tests/resources/reference_depth_dev_from_image.png b/tests/resources/reference_depth_dev_from_image.png
index 98a7948..29e5e00 100644
Binary files a/tests/resources/reference_depth_dev_from_image.png and b/tests/resources/reference_depth_dev_from_image.png differ
diff --git a/tests/resources/reference_dev.png b/tests/resources/reference_dev.png
index 72d6454..9e4e664 100644
Binary files a/tests/resources/reference_dev.png and b/tests/resources/reference_dev.png differ
diff --git a/tests/resources/reference_dev_image_to_image.png b/tests/resources/reference_dev_image_to_image.png
new file mode 100644
index 0000000..9a664e2
Binary files /dev/null and b/tests/resources/reference_dev_image_to_image.png differ
diff --git a/tests/resources/reference_dev_image_to_image_fill.png b/tests/resources/reference_dev_image_to_image_fill.png
new file mode 100644
index 0000000..272971e
Binary files /dev/null and b/tests/resources/reference_dev_image_to_image_fill.png differ
diff --git a/tests/resources/reference_dev_image_to_image_result.png b/tests/resources/reference_dev_image_to_image_result.png
deleted file mode 100644
index bf712a1..0000000
Binary files a/tests/resources/reference_dev_image_to_image_result.png and /dev/null differ
diff --git a/tests/resources/reference_dev_image_to_image_result_inpaint.png b/tests/resources/reference_dev_image_to_image_result_inpaint.png
deleted file mode 100644
index f860366..0000000
Binary files a/tests/resources/reference_dev_image_to_image_result_inpaint.png and /dev/null differ
diff --git a/tests/resources/reference_dev_kontext.png b/tests/resources/reference_dev_kontext.png
index e623a22..ea7bfeb 100644
Binary files a/tests/resources/reference_dev_kontext.png and b/tests/resources/reference_dev_kontext.png differ
diff --git a/tests/resources/reference_dev_lora.png b/tests/resources/reference_dev_lora.png
index deb9903..055b3c1 100644
Binary files a/tests/resources/reference_dev_lora.png and b/tests/resources/reference_dev_lora.png differ
diff --git a/tests/resources/reference_dev_lora_multiple.png b/tests/resources/reference_dev_lora_multiple.png
index 48179fd..f0494c2 100644
Binary files a/tests/resources/reference_dev_lora_multiple.png and b/tests/resources/reference_dev_lora_multiple.png differ
diff --git a/tests/resources/reference_ic_edit.png b/tests/resources/reference_ic_edit.png
new file mode 100644
index 0000000..0006554
Binary files /dev/null and b/tests/resources/reference_ic_edit.png differ
diff --git a/tests/resources/reference_ic_lora_identity.png b/tests/resources/reference_ic_lora_identity.png
new file mode 100644
index 0000000..9610703
Binary files /dev/null and b/tests/resources/reference_ic_lora_identity.png differ
diff --git a/tests/resources/reference_redux_dev.png b/tests/resources/reference_redux_dev.png
index 5c95711..ae146ea 100644
Binary files a/tests/resources/reference_redux_dev.png and b/tests/resources/reference_redux_dev.png differ
diff --git a/tests/resources/reference_schnell.png b/tests/resources/reference_schnell.png
index 87ffbb0..72d31d3 100644
Binary files a/tests/resources/reference_schnell.png and b/tests/resources/reference_schnell.png differ
diff --git a/tests/resources/reference_upscaled.png b/tests/resources/reference_upscaled.png
index b094181..bc5bfd0 100644
Binary files a/tests/resources/reference_upscaled.png and b/tests/resources/reference_upscaled.png differ
diff --git a/uv.lock b/uv.lock
index 8c567a1..693df04 100644
--- a/uv.lock
+++ b/uv.lock
@@ -712,7 +712,7 @@ wheels = [
 
 [[package]]
 name = "mflux"
-version = "0.9.6"
+version = "0.10.0"
 source = { editable = "." }
 dependencies = [
     { name = "huggingface-hub" },
@@ -746,8 +746,8 @@ dev = [
 requires-dist = [
     { name = "huggingface-hub", specifier = ">=0.24.5,<1.0" },
     { name = "matplotlib", specifier = ">=3.9.2,<4.0" },
-    { name = "mlx", specifier = ">=0.22.0,<=0.26.1" },
-    { name = "mlx", marker = "extra == 'dev'", specifier = "==0.26.1" },
+    { name = "mlx", specifier = ">=0.27.0,<0.28.0" },
+    { name = "mlx", marker = "extra == 'dev'", specifier = "==0.27.1" },
     { name = "numpy", specifier = ">=2.0.1,<3.0" },
     { name = "opencv-python", specifier = ">=4.10.0,<5.0" },
     { name = "piexif", specifier = ">=1.1.3,<2.0" },
@@ -770,25 +770,38 @@ provides-extras = ["dev"]
 
 [[package]]
 name = "mlx"
-version = "0.26.1"
+version = "0.27.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "mlx-metal", marker = "sys_platform == 'darwin'" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c9/ed/07a7145c500914ed8ab704a7a3828a5134e85f8b4458e1fee887306867d6/mlx-0.27.1-cp310-cp310-macosx_13_0_arm64.whl", hash = "sha256:a033b65fe46425ad5032867d5a71a556a5168108d89aa7092b457556c70d84fc", size = 557679, upload-time = "2025-07-25T22:55:31.942Z" },
+    { url = "https://files.pythonhosted.org/packages/aa/ae/b3e46904ea04b5badb77195169880733f09976a3dfc81c954b99f9adbbc4/mlx-0.27.1-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:2836c6fd9803dc0c6cd06f204e31b3e0191e6c5b6bc8570b28661d926908cba3", size = 530556, upload-time = "2025-07-25T22:56:13.171Z" },
+    { url = "https://files.pythonhosted.org/packages/6e/24/85f5a7dfcf0b4974aa43a2f6dc0d59dfc04e4a1051aacebd7e7810b823ff/mlx-0.27.1-cp310-cp310-macosx_15_0_arm64.whl", hash = "sha256:8b0566054c46d84c470cc99cda2afc3914ad6c7808fbb724dc1ec235e5b2a98c", size = 530559, upload-time = "2025-07-25T22:55:39.055Z" },
+    { url = "https://files.pythonhosted.org/packages/14/bd/91ad900837a3760ba056863c1350f0deabc3e462c3999b4d5ea875368177/mlx-0.27.1-cp310-cp310-manylinux_2_35_x86_64.whl", hash = "sha256:91ef93ce09900c9a8ca662cf34e3c39ab5af2762822ecd6b12fecae518be167f", size = 628183, upload-time = "2025-07-25T22:58:11.553Z" },
+    { url = "https://files.pythonhosted.org/packages/ac/63/88cb9a5019bea21244385eacc4c64deb4d074a8bc3b4b6d2d97cdacf97a2/mlx-0.27.1-cp311-cp311-macosx_13_0_arm64.whl", hash = "sha256:d2e5dedbfbcbe558e51a5c476ca6a18e307676f9e49854eb27e53778bc474699", size = 557582, upload-time = "2025-07-25T22:55:22.243Z" },
+    { url = "https://files.pythonhosted.org/packages/fa/87/c2d7343e7b054481c52e23a190911c40aed4f8c77a8dfeda1f48d5cb520a/mlx-0.27.1-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:9f04b9778897a879c9ca22e5413dfa1efc192d86d7211b184e079efec49dfb8b", size = 530820, upload-time = "2025-07-25T22:55:34.121Z" },
+    { url = "https://files.pythonhosted.org/packages/06/bf/20497ca7411028baa56372c20e94a3eaddac973155b415f08fc12592c2cf/mlx-0.27.1-cp311-cp311-macosx_15_0_arm64.whl", hash = "sha256:01d794f9e390438ab4f942a18d9a8ca65bef10c2c2007ef38ca988d039d6d9d3", size = 530821, upload-time = "2025-07-25T22:55:28.847Z" },
+    { url = "https://files.pythonhosted.org/packages/48/68/234a0d846576502ea1bf8ea478c764d87eec9042349e298d2aea58c4e8e9/mlx-0.27.1-cp311-cp311-manylinux_2_35_x86_64.whl", hash = "sha256:fae11432d0639789f1e172b19b35ac8987c8ab9716e55a23fc7a170d6545fc33", size = 628500, upload-time = "2025-07-25T22:57:59.556Z" },
+    { url = "https://files.pythonhosted.org/packages/65/43/125102bbb2be6825880ae2dc8d8702f99cfa7753407f574457b36e422218/mlx-0.27.1-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:0c570c9afb57c697bd864504115be8a7c4de97f0b80557a597d496ee426a6812", size = 549869, upload-time = "2025-07-25T22:55:32.698Z" },
+    { url = "https://files.pythonhosted.org/packages/f3/79/0bf681700fc8b165517e907f9ec777b5a5d628004a65a777148f68c6baa0/mlx-0.27.1-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:ccff7bbbd9df302b26e79013ef6d0c3531c9ba5963ead521e2d85856811b86a0", size = 531671, upload-time = "2025-07-25T22:57:03.392Z" },
+    { url = "https://files.pythonhosted.org/packages/ec/97/f1367b4892bef7f78e38737d3a28094e93124f11684a28a9e92ed5a13b2b/mlx-0.27.1-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:9ccadaed449c07dfeae484620992b904c17dfea7564f8df63095c60eed3af02b", size = 531672, upload-time = "2025-07-25T22:55:35.779Z" },
+    { url = "https://files.pythonhosted.org/packages/9c/0f/3ebec0806ed2c5ba8ba151394cd62a46b4d83ea347da33af91b86d8969d4/mlx-0.27.1-cp312-cp312-manylinux_2_35_x86_64.whl", hash = "sha256:742c413e75605b71db69379a176da63e32ba19b9e9ad03b8763adbd1fcfcd394", size = 621661, upload-time = "2025-07-25T22:58:00.96Z" },
+    { url = "https://files.pythonhosted.org/packages/86/f6/4324386b0764deb692e14a97282a348a9a938aa8b441bf8b6c7599f418d4/mlx-0.27.1-cp313-cp313-macosx_13_0_arm64.whl", hash = "sha256:803669a28031766c2b0fe94c0a3bfd030184e706092f0a831b33620c1e2ef865", size = 549847, upload-time = "2025-07-25T22:55:17.581Z" },
+    { url = "https://files.pythonhosted.org/packages/cf/4b/3194ccb03527a050c04d837d731a11599f8620e6ce16d3971798caae1d44/mlx-0.27.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:e9649b3d86ce564000797384510c9d07af38a9ce2a07df8e2f7c6a3e0f0f059e", size = 531664, upload-time = "2025-07-25T22:56:02.928Z" },
+    { url = "https://files.pythonhosted.org/packages/cc/57/a6e0d8dc6e7ba08a64d71fb89d743e77446040113ea1dbb7950be8f60f39/mlx-0.27.1-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:5c501ceec7c6aa2ea1c850dd4e14e679f5416a142264b9d5d405a4e0aeb991b2", size = 531663, upload-time = "2025-07-25T22:55:38.62Z" },
+    { url = "https://files.pythonhosted.org/packages/48/96/6932e0fa866d08cdf77a3fc37cf624e4cac8960d2e297c3fd80afb49417d/mlx-0.27.1-cp313-cp313-manylinux_2_35_x86_64.whl", hash = "sha256:42bfbabdcd0b7aa1c94748c8963e66680a3cf3f599b01036db14e496bf276eb8", size = 621609, upload-time = "2025-07-25T22:58:21.477Z" },
+]
+
+[[package]]
+name = "mlx-metal"
+version = "0.27.1"
 source = { registry = "https://pypi.org/simple" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/91/9d/645600137d0b7d4689e081de1c62247dae448dfadd8ebdac0fea8c70da33/mlx-0.26.1-cp310-cp310-macosx_13_0_arm64.whl", hash = "sha256:8f88900eaef76f8f23e957ef303f3664c273c73a504b4e57bb491e397d7a891e", size = 32399627, upload-time = "2025-06-04T01:02:30.075Z" },
-    { url = "https://files.pythonhosted.org/packages/50/8c/132c081c62ff3352f62bd25c9fc4f654d959402ff94de7c269b53479978d/mlx-0.26.1-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:794b52f2ae52969aac8fa5c271f55a668c0f802cbacda70652033db9c11887c6", size = 31869876, upload-time = "2025-06-04T01:03:04.508Z" },
-    { url = "https://files.pythonhosted.org/packages/1e/11/eba1bdbaa82b6178887ff22eeae5ab90ef6de73e5fdd518753b808bbbc21/mlx-0.26.1-cp310-cp310-macosx_15_0_arm64.whl", hash = "sha256:030b79376da1564d21af24ecfaa4e049bad823adfd7572376e1f81bd0f08af90", size = 31870531, upload-time = "2025-06-04T01:03:15.67Z" },
-    { url = "https://files.pythonhosted.org/packages/8d/3a/415cf516daa5cc6f95ac953d894346fc69c013f0c2348665223b130fccab/mlx-0.26.1-cp310-cp310-manylinux_2_31_x86_64.whl", hash = "sha256:70a286a1e9a5e9fb69914e552e563a8de33963a24496c8b69f3c084a88e7ef90", size = 10125716, upload-time = "2025-06-06T23:08:07.886Z" },
-    { url = "https://files.pythonhosted.org/packages/63/93/97c0a149cba1ad3fbfb01be4c6b293b4cb9d11e41d4b54735048e3d671bd/mlx-0.26.1-cp311-cp311-macosx_13_0_arm64.whl", hash = "sha256:1973b46817db50328fe8f74b2faa860e761eb08f4230461bb299057a411c070b", size = 32399548, upload-time = "2025-06-04T01:02:34.868Z" },
-    { url = "https://files.pythonhosted.org/packages/1d/4f/3d3b52f0040462ab53d66d65eefd062f577feb0647ec3aa916faa7f0282b/mlx-0.26.1-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:478abbe1de49cb1409a38c263bafe45e70e3c5b6cbefd4fc3f5c5b7db90d2908", size = 31870277, upload-time = "2025-06-04T01:01:29.932Z" },
-    { url = "https://files.pythonhosted.org/packages/80/75/ea19b3f1e9db337a8ef5ceded8f95ae0ed636c66ab40473c6d9eaad5ce01/mlx-0.26.1-cp311-cp311-macosx_15_0_arm64.whl", hash = "sha256:8d76d75d652f252f7887acebc6d1ef5d3aec16e72673332895584f1617c35151", size = 31871009, upload-time = "2025-06-04T01:03:09.963Z" },
-    { url = "https://files.pythonhosted.org/packages/72/a7/b4611911540e15ffd5c5b039ff4a83f43ef64a7ebc51df499b4456815470/mlx-0.26.1-cp311-cp311-manylinux_2_31_x86_64.whl", hash = "sha256:1cedf67d32d21e23e1639116cc24b62b6cc95a5f3bb0aa0ad6833e2910eb802c", size = 10125854, upload-time = "2025-06-06T23:07:39.23Z" },
-    { url = "https://files.pythonhosted.org/packages/73/3e/a099ca333bd9cd313c7792a18ad4dfc47b051957ba24edd59f3bcda87d8c/mlx-0.26.1-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:ebcc19c01d90b16462c8ff4255b00837bf45ce149d207170d95770b6699cc5ce", size = 32396393, upload-time = "2025-06-04T01:02:40.689Z" },
-    { url = "https://files.pythonhosted.org/packages/d1/e0/91707351c0becd33c477a1eba1a0fc124b3838cc65e3d2bd06dfa0a0533e/mlx-0.26.1-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:5eea5e025655a2b07c1573602e705013bd6d3ebdf432c39eb8eb29628c6df5a7", size = 31871294, upload-time = "2025-06-04T01:01:14.326Z" },
-    { url = "https://files.pythonhosted.org/packages/8d/e6/d5759fb20faceac183cfd45ede5fa0484f972eecfa2a851f8dd24aa4ace2/mlx-0.26.1-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:3b93e1e2375222a7b0f65b4be3c4010855db2ee8c40cc39f8bf2820ff1652b06", size = 31871854, upload-time = "2025-06-04T01:03:14.003Z" },
-    { url = "https://files.pythonhosted.org/packages/e0/98/a95ca1dfb3721d3dbdd197ab17e3bb79b695bf426e37939d71d24cca19ce/mlx-0.26.1-cp312-cp312-manylinux_2_31_x86_64.whl", hash = "sha256:07f2fdb11b908156a65823ed062682cad0b03325a4f11fb9a36ce5acfe7475e5", size = 10120549, upload-time = "2025-06-06T23:07:43.068Z" },
-    { url = "https://files.pythonhosted.org/packages/a2/a7/871c451fe81274d37022a62f825c1dcd22b30e1f8bd2241f91d9f508c9b9/mlx-0.26.1-cp313-cp313-macosx_13_0_arm64.whl", hash = "sha256:ccd8662abad0f1340326412d6051c116fcb5c923c4d2a25ba1277ae65ab140dd", size = 32396333, upload-time = "2025-06-04T01:02:29.963Z" },
-    { url = "https://files.pythonhosted.org/packages/82/77/720bea5a67934b50372dfd5043864458f103743edcc7c30049e788ea3762/mlx-0.26.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:0c113dd7c7ac13af6e39f0132d33a8dc78928e858ba8d18f8c89f8bfa694a358", size = 31871172, upload-time = "2025-06-04T01:03:05.075Z" },
-    { url = "https://files.pythonhosted.org/packages/15/4f/83f67bc4fe012dffffd2d96d2767b83fee9b2d7d185611d554ac659cfa4d/mlx-0.26.1-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:2ec37131dbb06c0be78ce56b1731ddab6e56183012e7b83bea79b5329ef7d695", size = 31871791, upload-time = "2025-06-04T01:03:15.384Z" },
-    { url = "https://files.pythonhosted.org/packages/4f/fb/4123952002fd91f096ba07ce797b6bb6a32cc7a89c988565e261559f77dd/mlx-0.26.1-cp313-cp313-manylinux_2_31_x86_64.whl", hash = "sha256:db96a53466d8efc6cf2a2918b2d4e29cbf9f25174c838fb3c380c8717a40752f", size = 10120515, upload-time = "2025-06-06T23:07:38.428Z" },
+    { url = "https://files.pythonhosted.org/packages/8b/77/89fa3327011f018638c9e943e1edc081ce9ca0ed296fe64d6daf93c6ff51/mlx_metal-0.27.1-py3-none-macosx_13_0_arm64.whl", hash = "sha256:c66d9b1adb3c0ea19492fba6493f672bc7542e65dd65f7e2995918815fbeb907", size = 33523035, upload-time = "2025-07-25T22:58:02.533Z" },
+    { url = "https://files.pythonhosted.org/packages/d7/a8/ac706ad6ce834834762d5146d791f77710efc896c13ef47fd7d672099056/mlx_metal-0.27.1-py3-none-macosx_14_0_arm64.whl", hash = "sha256:fe4415ddd242974d91c7ca0699cd01507d17da8a5ba304122ef137cdb5e7fff4", size = 32926383, upload-time = "2025-07-25T22:57:54.194Z" },
+    { url = "https://files.pythonhosted.org/packages/78/77/6963681fb54ecaa0ae5de4209c15504a803a0edd1a33fd074e6c558fd5e0/mlx_metal-0.27.1-py3-none-macosx_15_0_arm64.whl", hash = "sha256:d025dea30bda8baa32c928cfa333eac64a5adc8d07656f8fc55072d99403ebc9", size = 32897065, upload-time = "2025-07-25T22:58:26.362Z" },
 ]
 
 [[package]]