Fix edge case error when converting vlm models: "Received parameters not in model"

app.py

@@ -1,8 +1,14 @@
+from pathlib import Path
+import traceback
+
 import os
 import tempfile
 import importlib.util
 from enum import Enum
 
+from contextlib import contextmanager, AbstractContextManager
+from functools import wraps
+
 os.environ["HF_HUB_CACHE"] = "cache"
 os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
 import gradio as gr
@@ -18,9 +24,12 @@ from apscheduler.schedulers.background import BackgroundScheduler
 
 from textwrap import dedent
 from typing import (
+    Any,
     Callable,
     Dict,
     Optional,
+    Tuple,
+    Type,
     Union,
     NamedTuple,
 )
@@ -172,6 +181,47 @@ def upload_to_hub(path, upload_repo, hf_path, oauth_token, runtime: Runtime):
         
     print(f"Upload successful, go to https://huggingface.co/{upload_repo} for details.")    
 
+
+@contextmanager
+def patch_strict_default_methods_ctx() -> AbstractContextManager[Callable[[Any, str], None]]:
+    """
+    Context manager to temporarily set the default value of the 'strict' arg to `False` 
+    for specified class methods.
+    Does not affect explict `strict=True`.
+
+    (e.g. `def update(self, parameters: dict, strict: bool = True)` 
+    becomes `def update(self, parameters: dict, strict: bool = False)`)
+
+    Typical usage:
+
+        with patch_strict_default_methods_ctx() as patch:
+            patch(Foo, "bar")
+            patch(Foo, "baz")
+            patch(Bar, "foo")
+            # Patched methods active here
+        # Originals restored here
+    """
+
+    originals: Dict[Tuple[Type[Any], str], Callable] = {}
+
+    def patch(cls: Any, method_name: str):
+        method = getattr(cls, method_name)
+        originals[(cls, method_name)] = method
+
+        @wraps(method)
+        def wrapper(self, *args, strict=False, **kwargs):
+            return method(self, *args, strict=strict, **kwargs)
+
+        setattr(cls, method_name, wrapper)
+
+    try:
+        yield patch
+    finally:
+        # Restore all patched methods
+        for (cls, method_name), original in originals.items():
+            setattr(cls, method_name, original)
+        originals.clear()
+
 def convert(
     hf_path: str,
     mlx_path: str = "mlx_model",
@@ -188,6 +238,8 @@ def convert(
     skip_vision: bool = False, # mlx-vlm
     trust_remote_code: bool = True, # mlx-vlm
 ) -> Runtime :
+    model_path = get_model_path(hf_path, revision=revision)
+
     def mlx_lm_convert():
         mlx_lm.convert(
             hf_path=hf_path,
@@ -203,21 +255,44 @@ def convert(
         )
 
     def mlx_vlm_convert():
-        mlx_vlm.convert(
-            hf_path=hf_path,
-            mlx_path=mlx_path,
-            quantize=quantize,
-            q_group_size=q_group_size,
-            q_bits=q_bits,
-            dtype=dtype,
-            upload_repo=upload_repo,
-            revision=revision,
-            dequantize=dequantize,
-            skip_vision=skip_vision,
-            trust_remote_code=trust_remote_code,
-        )
-        
-    model_path = get_model_path(hf_path, revision=revision)
+        # try:
+        #     new_model_path = remove_extra_parameters_from_weights(model_path=model_path)
+        #     print(f"{new_model_path} exists: {Path(new_model_path).exists()}")
+        # except Exception as e:
+        #     new_model_path = model_path
+        #     print(f"Unexpected error while trying to fix model weights: {e}")
+        #     traceback.print_exc()
+        #     raise e
+
+        def _mlx_vlm_convert():
+            mlx_vlm.convert(
+                #hf_path=new_model_path,
+                hf_path=hf_path,
+                mlx_path=mlx_path,
+                quantize=quantize,
+                q_group_size=q_group_size,
+                q_bits=q_bits,
+                dtype=dtype,
+                upload_repo=upload_repo,
+                revision=revision,
+                dequantize=dequantize,
+                skip_vision=skip_vision,
+                trust_remote_code=trust_remote_code,
+            )
+
+        try:
+            _mlx_vlm_convert()
+        except ValueError as e:
+            print(e)
+            print(f"Error converting, try again with strict = False")
+            with patch_strict_default_methods_ctx() as patch:
+                import mlx.nn as n
+                patch(nn.Module, "load_weights")
+                patch(nn.Module, "update")
+                patch(nn.Module, "update_modules")
+                # patched strict=False by default, try again
+                _mlx_vlm_convert()
+
     config = load_config(model_path)
     model_type = config["model_type"]
     model_type = MODEL_REMAPPING.get(model_type, model_type)
@@ -273,6 +348,7 @@ def process_model(model_id, q_method, oauth_token: gr.OAuthToken | None):
             "llama.png",
         )
     except Exception as e:
+        traceback.print_exc()
         return (f"Error: {e}", "error.png")
     finally:
         clear_hf_cache_space()