Fix: loading non-quantized saved models (#261)

Co-authored-by: Anthony Wu <pls-file-gh-issue@users.noreply.github.com>

src/mflux/weights/model_saver.py

@@ -34,8 +34,12 @@ class ModelSaver:
         weights = ModelSaver._split_weights(base_path, dict(tree_flatten(model.parameters())))
         for i, weight in enumerate(weights):
             mx.save_safetensors(
+                # usage: save_safetensors(file: str, arrays: dict[str, array], metadata: Optional[dict[str, str]] = None)
                 str(path / f"{i}.safetensors"),
+                # arrays (dict(str, array)): The dictionary of names to arrays to be saved.
                 weight,
+                # [save_safetensors] Metadata must be a dictionary with string keys and values
+                # i.e. 'None' and other special values are string-ified and need to be parsed by readers
                 {
                     "quantization_level": str(bits),
                     "mflux_version": VersionUtil.get_mflux_version(),

src/mflux/weights/quantization_util.py

@@ -19,6 +19,13 @@ class QuantizationUtil:
     ) -> None:
         q_level = weights.meta_data.quantization_level
 
+        # mx.save_tensors saves metadata dict kv 'quantization_level': 'None' as a str: str mapping
+        # we coerce both configs to NoneType to help users use non-quantized saved model files
+        if q_level == "None":
+            q_level = None
+        if quantize == "None":
+            quantize = None
+
         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(vae, bits=bits)