Handle strict argument in load_weights for better model compatibility (#94)

Author: kamillobinski

mlx_audio/tts/models/sesame/model.py

@@ -344,8 +344,8 @@ def sanitize(self, weights):
 
         return sanitized_weights
 
-    def load_weights(self, weights):
-        self.model.load_weights(weights)
+    def load_weights(self, weights, strict: bool = True):
+        self.model.load_weights(weights, strict=strict)
 
     def generate(
         self,

mlx_audio/tts/utils.py

@@ -1,5 +1,6 @@
 import glob
 import importlib
+import inspect
 import logging
 import shutil
 from pathlib import Path
@@ -191,7 +192,13 @@ def get_class_predicate(p, m):
             class_predicate=get_class_predicate,
         )
 
-    model.load_weights(list(weights.items()), strict=strict)
+    load_weights_sig = inspect.signature(model.load_weights)
+
+    kwargs = {"weights": list(weights.items())}
+    if "strict" in load_weights_sig.parameters:
+        kwargs["strict"] = strict
+
+    model.load_weights(**kwargs)
 
     if not lazy:
         mx.eval(model.parameters())