Release 0.9.6 (#238)

CHANGELOG.md

@@ -5,6 +5,14 @@ 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.9.6] - 2025-07-20
+
+# MFLUX v.0.9.6 Release Notes
+
+### 🔧 Technical Details
+
+- Cap the upper MLX dependency to a known working version (0.26.1) to avoid compatibility issues with newer MLX releases that enforce stricter weight validation (see [#238](https://github.com/filipstrand/mflux/pull/238))
+
 ## [0.9.5] - 2025-07-17
 
 # MFLUX v.0.9.5 Release Notes

pyproject.toml

@@ -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.27.0",
+    "mlx>=0.22.0,<=0.26.1",
     "numpy>=2.0.1,<3.0",
     "opencv-python>=4.10.0,<5.0",
     "piexif>=1.1.3,<2.0",

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:

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(