Ensure checks see if vision model is in any of visible models list

h2oai · Jun 13, 2024 · 7340826 · 7340826
1 parent 7c703cb
commit 7340826
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/src/enums.py b/src/enums.py
@@ -374,10 +374,10 @@ def is_gradio_vision_model(base_model):
         base_model.startswith('Qwen/Qwen-VL')
 
 
-def is_vision_model(base_model, visible_models=[], visible_vision_models=[]):
+def is_vision_model(base_model, all_visible_models=[], visible_vision_models=[]):
     if not base_model:
         return False
-    if visible_vision_models and visible_models and visible_vision_models[0] in visible_models:
+    if visible_vision_models and all_visible_models and visible_vision_models[0] in all_visible_models:
         # all models are vision models by proxy
         return True
     return is_gradio_vision_model(base_model) or \

diff --git a/src/gen.py b/src/gen.py
@@ -2423,14 +2423,15 @@ def main(
             model_visible_vision_models = visible_vision_models
         if isinstance(model_visible_vision_models, str):
             model_visible_vision_models = [model_visible_vision_models]
+        all_visible_models = [x.get('visible_models') or x.get('base_model') for x in model_lock]
         model_state_trial['is_vision_model'] = is_vision_model(model_state_trial['base_model'],
-                                                               visible_models=visible_models,
+                                                               all_visible_models=all_visible_models,
                                                                visible_vision_models=model_visible_vision_models)
         if model_state_trial['is_actually_vision_model']:
             model_state_trial['images_num_max'] = images_num_max_dict.get(model_state_trial['base_model'],
                                                                           images_num_max or 1) or 1
         elif model_state_trial['is_vision_model'] and model_visible_vision_models and len(model_visible_vision_models) > 0:
-            model_state_trial['images_num_max'] = images_num_max_dict.get(model_visible_vision_models,
+            model_state_trial['images_num_max'] = images_num_max_dict.get(model_visible_vision_models[0],
                                                                           images_num_max or 1) or 1
         else:
             model_state_trial['images_num_max'] = 0

diff --git a/src/version.py b/src/version.py
@@ -1 +1 @@
-__version__ = "3e0ef9c00816f985e5ccdf7f742a5b4cb86682b0"
+__version__ = "7c703cbc281ef9ef384406334c244be83273a91d"