llava-next wip

huggingface · Aug 30, 2024 · c3fc608 · c3fc608
1 parent d0a72ed
commit c3fc608
Show file tree

Hide file tree

Showing 3 changed files with 254 additions and 27 deletions.
diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py
@@ -850,7 +850,7 @@ def _get_submodels_and_export_configs(
     exporter: str = "openvino",
     stateful:bool = False
 ):
-    if not custom_architecture and library_name == "transformers" and model.config.model_type == "llava":
+    if not custom_architecture and library_name == "transformers" and model.config.model_type in ["llava", "llava_next"]:
         return _get_llava_submodels_and_export_configs(model, task, library_name, int_dtype, float_dtype, preprocessors, model_kwargs, stateful)
 
     export_config, models_for_export =  _default_get_submodels_and_export_configs(model, task, monolith, custom_export_configs, custom_architecture, _variant, library_name, int_dtype, float_dtype, fn_get_submodels, preprocessors, legacy, model_kwargs, exporter)

diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py
@@ -76,10 +76,14 @@
 
 
 def init_model_configs():
-    TasksManager._SUPPORTED_MODEL_TYPE["llava"] = {"openvino": {"image-to-text": "LlamaOpenVINOConfig", "image-text-to-text": "LlamaOpenVINOConfig", "vision-question-answering": "LlamaOpenVINOConfig"}}
+    TasksManager._SUPPORTED_MODEL_TYPE["llava"] = {"openvino": {"image-to-text": "", "image-text-to-text": "", "vision-question-answering": ""}}
+    TasksManager._SUPPORTED_MODEL_TYPE["llava-next"] = {"openvino": {"image-to-text": "", "image-text-to-text": "", "vision-question-answering": ""}}
     TasksManager._CUSTOM_CLASSES[("pt", "llava", "image-to-text")] = ("transformers", "LlavaForConditionalGeneration")
     TasksManager._CUSTOM_CLASSES[("pt", "llava", "image-text-to-text")] = ("transformers", "LlavaForConditionalGeneration")
     TasksManager._CUSTOM_CLASSES[("pt", "llava", "visual-question-answering")] = ("transformers", "LlavaForConditionalGeneration")
+    TasksManager._CUSTOM_CLASSES[("pt", "llava-next", "image-to-text")] = ("transformers", "LlavaNextForConditionalGeneration")
+    TasksManager._CUSTOM_CLASSES[("pt", "llava-next", "image-text-to-text")] = ("transformers", "LlavaNextForConditionalGeneration")
+    TasksManager._CUSTOM_CLASSES[("pt", "llava-next", "visual-question-answering")] = ("transformers", "LlavaNextForConditionalGeneration")
     TasksManager._TRANSFORMERS_TASKS_TO_MODEL_LOADERS["image-text-to-text"] = TasksManager._TRANSFORMERS_TASKS_TO_MODEL_LOADERS["image-to-text"]
 
     supported_model_types = [