finalize janus support

Author: eaidova

optimum/exporters/openvino/__main__.py

@@ -13,7 +13,6 @@
 #  limitations under the License.
 
 import gc
-import importlib
 import logging
 import operator
 import warnings
@@ -40,13 +39,13 @@
     _infer_library_from_model_name_or_path,
     _OpenClipForZeroShotImageClassification,
 )
-from optimum.utils.save_utils import maybe_load_preprocessors
 
 from .utils import (
     _MAX_UNCOMPRESSED_SIZE,
     MULTI_MODAL_TEXT_GENERATION_MODELS,
     clear_class_registry,
     deduce_diffusers_dtype,
+    load_preprocessors,
 )
 
 
@@ -193,6 +192,7 @@ def main_export(
     ```
     """
     from optimum.exporters.openvino.convert import export_from_model
+
     if use_auth_token is not None:
         warnings.warn(
             "The `use_auth_token` argument is deprecated and will be removed soon. Please use the `token` argument instead.",
@@ -214,7 +214,7 @@ def main_export(
             revision=revision,
             cache_dir=cache_dir,
             token=token,
-            library_name=library_name
+            library_name=library_name,
         )
         if library_name == "sentence_transformers":
             logger.warning(
@@ -434,7 +434,7 @@ class StoreAttr(object):
                 possible_synonyms = ""
             logger.info(f"Automatic task detection to {task}{possible_synonyms}.")
 
-        preprocessors = maybe_load_preprocessors(
+        preprocessors = load_preprocessors(
             model_name_or_path, subfolder=subfolder, trust_remote_code=trust_remote_code
         )
 

optimum/exporters/openvino/convert.py

@@ -757,7 +757,7 @@ def export_from_model(
                 logger.warning(
                     f"The generation config will not be saved, saving failed with following error:\n{exception}"
                 )
-
+        logger.warn(preprocessors)
         save_preprocessors(preprocessors, model.config, output, trust_remote_code)
 
         files_subpaths = ["openvino_" + model_name + ".xml" for model_name in models_and_export_configs.keys()]

optimum/exporters/openvino/model_configs.py

@@ -145,11 +145,10 @@ def init_model_configs():
     # for model registration in auto transformers classses
     if importlib.util.find_spec("janus") is not None:
         try:
-            from janus.models import MultiModalityCausalLM
+            from janus.models import MultiModalityCausalLM, VLChatProcessor
         except ImportError:
             pass
 
-
     if is_diffusers_available() and "fill" not in TasksManager._DIFFUSERS_TASKS_TO_MODEL_LOADERS:
         TasksManager._DIFFUSERS_TASKS_TO_MODEL_LOADERS["fill"] = "FluxFillPipeline"
         TasksManager._DIFFUSERS_TASKS_TO_MODEL_MAPPINGS["fill"] = {"flux": "FluxFillPipeline"}
@@ -1353,7 +1352,9 @@ def patch_model_for_export(
 
 
 class LMInputEmbedsConfigHelper(TextDecoderWithPositionIdsOnnxConfig):
-    def __init__(self, export_config, patcher_cls=None, dummy_input_generator=None, inputs_update=None, remove_lm_head=False):
+    def __init__(
+        self, export_config, patcher_cls=None, dummy_input_generator=None, inputs_update=None, remove_lm_head=False
+    ):
         self.orig_export_config = export_config
         if dummy_input_generator is not None:
             export_config.DUMMY_INPUT_GENERATOR_CLASSES = (
@@ -1372,16 +1373,15 @@ def __init__(self, export_config, patcher_cls=None, dummy_input_generator=None,
     def patch_model_for_export(
         self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
     ) -> "ModelPatcher":
-        
         if self.patcher_cls is not None:
             patcher = self.patcher_cls(self, model, model_kwargs=model_kwargs)
         # Refer to DecoderModelPatcher.
-        else: 
+        else:
             patcher = self.orig_export_config.patch_model_for_export(model, model_kwargs=model_kwargs)
-        
+
         if self.remove_lm_head:
             patcher = RemoveLMHeadPatcherHelper(self, model, model_kwargs, patcher)
-        
+
         return patcher
 
     @property
@@ -1390,7 +1390,7 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
         if self.remove_lm_head:
             logits_info = outputs.pop("logits")
             updated_outputs = {"last_hidden_state": logits_info}
-            return {**updated_outputs, **outputs} 
+            return {**updated_outputs, **outputs}
         return outputs
 
     @property
@@ -1479,15 +1479,15 @@ def get_vlm_text_generation_config(
     model_patcher=None,
     dummy_input_generator=None,
     inputs_update=None,
-    remove_lm_head=False
+    remove_lm_head=False,
 ):
     internal_export_config = get_vlm_internal_text_generation_config(model_type, model_config, int_dtype, float_dtype)
     export_config = LMInputEmbedsConfigHelper(
         internal_export_config,
         patcher_cls=model_patcher,
         dummy_input_generator=dummy_input_generator,
         inputs_update=inputs_update,
-        remove_lm_head=remove_lm_head
+        remove_lm_head=remove_lm_head,
     )
     export_config._normalized_config = internal_export_config._normalized_config
     return export_config
@@ -2812,45 +2812,60 @@ class JanusConfigBehavior(str, enum.Enum):
 
 
 class JanusDummyVisionGenInputGenerator(DummyInputGenerator):
-    SUPPORTED_INPUT_NAMES = (
-        "pixel_values",
-        "image_ids",
-        "code_b",
-        "image_shape",
-        "lm_hidden_state",
-        "hidden_state"
-    )
+    SUPPORTED_INPUT_NAMES = ("pixel_values", "image_ids", "code_b", "image_shape", "lm_hidden_state", "hidden_state")
 
     def __init__(
-            self,
-            task: str,
-            normalized_config: NormalizedConfig,
-            batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"],
-            sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"],
-            **kwargs,
-        ):
-            self.task = task
-            self.batch_size = batch_size
-            self.sequence_length = sequence_length
-            self.normalized_config = normalized_config
-    
+        self,
+        task: str,
+        normalized_config: NormalizedConfig,
+        batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"],
+        sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"],
+        **kwargs,
+    ):
+        self.task = task
+        self.batch_size = batch_size
+        self.sequence_length = sequence_length
+        self.normalized_config = normalized_config
+
     def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
         if input_name == "pixel_values":
-            return self.random_float_tensor([self.batch_size, 1, 3, self.normalized_config.config.params.image_size, self.normalized_config.config.params.image_size])
-        
+            return self.random_float_tensor(
+                [
+                    self.batch_size,
+                    1,
+                    3,
+                    self.normalized_config.config.params.image_size,
+                    self.normalized_config.config.params.image_size,
+                ]
+            )
+
         if input_name == "image_ids":
-            return self.random_int_tensor([self.sequence_length], max_value=self.normalized_config.config.params.image_token_size, framework=framework, dtype=int_dtype)
+            return self.random_int_tensor(
+                [self.sequence_length],
+                max_value=self.normalized_config.config.params.image_token_size,
+                framework=framework,
+                dtype=int_dtype,
+            )
         if input_name == "code_b":
-            return self.random_int_tensor([self.batch_size, 576], max_value=self.normalized_config.config.params.image_token_size, framework=framework, dtype=int_dtype)
+            return self.random_int_tensor(
+                [self.batch_size, 576],
+                max_value=self.normalized_config.config.params.image_token_size,
+                framework=framework,
+                dtype=int_dtype,
+            )
         if input_name == "image_shape":
             import torch
-            return torch.tensor([self.batch_size, self.normalized_config.config.params.n_embed, 24, 24], dtype=torch.int64)
+
+            return torch.tensor(
+                [self.batch_size, self.normalized_config.config.params.n_embed, 24, 24], dtype=torch.int64
+            )
         if input_name == "hidden_state":
-            return self.random_float_tensor([self.batch_size, self.sequence_length, self.normalized_config.hidden_size])
+            return self.random_float_tensor(
+                [self.batch_size, self.sequence_length, self.normalized_config.hidden_size]
+            )
         if input_name == "lm_hidden_state":
             return self.random_float_tensor([self.sequence_length, self.normalized_config.hidden_size])
         return super().generate(input_name, framework, int_dtype, float_dtype)
-    
 
 
 @register_in_tasks_manager("multi-modality", *["image-text-to-text", "any-to-any"], library_name="transformers")
@@ -2868,7 +2883,7 @@ def __init__(
         float_dtype: str = "fp32",
         behavior: JanusConfigBehavior = JanusConfigBehavior.VISION_EMBEDDINGS,
         preprocessors: Optional[List[Any]] = None,
-        **kwargs
+        **kwargs,
     ):
         super().__init__(
             config=config,
@@ -2882,7 +2897,9 @@ def __init__(
         if self._behavior == JanusConfigBehavior.VISION_EMBEDDINGS and hasattr(config, "vision_config"):
             self._config = config.vision_config
             self._normalized_config = NormalizedVisionConfig(self._config)
-        if self._behavior in [JanusConfigBehavior.LM_HEAD, JanusConfigBehavior.VISION_GEN_HEAD] and hasattr(config, "language_config"):
+        if self._behavior in [JanusConfigBehavior.LM_HEAD, JanusConfigBehavior.VISION_GEN_HEAD] and hasattr(
+            config, "language_config"
+        ):
             self._config = config.language_config
             self._normalized_config = NormalizedTextConfig(self._config)
         if self._behavior == JanusConfigBehavior.VISION_GEN_EMBEDDINGS and hasattr(config, "gen_head_config"):
@@ -2912,7 +2929,7 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
             return {"last_hidden_state": {0: "batch_size"}}
         if self._behavior == JanusConfigBehavior.VISION_GEN_EMBEDDINGS:
             return {"last_hidden_state": {0: "num_tokens"}}
-        
+
         if self._behavior == JanusConfigBehavior.LM_HEAD:
             return {"logits": {0: "batch_size", 1: "sequence_length"}}
 
@@ -2979,7 +2996,6 @@ def with_behavior(
                 preprocessors=self._preprocessors,
             )
 
-
         if behavior == JanusConfigBehavior.VISION_EMBEDDINGS:
             return self.__class__(
                 self._orig_config,
@@ -2989,7 +3005,7 @@ def with_behavior(
                 behavior=behavior,
                 preprocessors=self._preprocessors,
             )
-        
+
         if behavior == JanusConfigBehavior.VISION_GEN_DECODER:
             return self.__class__(
                 self._orig_config,
@@ -3000,7 +3016,6 @@ def with_behavior(
                 preprocessors=self._preprocessors,
             )
 
-
     def get_model_for_behavior(self, model, behavior: Union[str, JanusConfigBehavior]):
         if isinstance(behavior, str) and not isinstance(behavior, JanusConfigBehavior):
             behavior = JanusConfigBehavior(behavior)
@@ -3023,7 +3038,7 @@ def get_model_for_behavior(self, model, behavior: Union[str, JanusConfigBehavior
 
         if behavior == JanusConfigBehavior.VISION_GEN_EMBEDDINGS:
             return model
-    
+
         if behavior == JanusConfigBehavior.VISION_GEN_HEAD:
             gen_head = model.gen_head
             gen_head.config = model.language_model.config
@@ -3032,7 +3047,6 @@ def get_model_for_behavior(self, model, behavior: Union[str, JanusConfigBehavior
         if behavior == JanusConfigBehavior.VISION_GEN_DECODER:
             return model.gen_vision_model
 
-
     def patch_model_for_export(
         self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
     ):
@@ -3045,7 +3059,6 @@ def patch_model_for_export(
             return JanusVisionGenDecoderModelPatcher(self, model, model_kwargs)
         return super().patch_model_for_export(model, model_kwargs)
 
-
     def rename_ambiguous_inputs(self, inputs):
         if self._behavior == JanusConfigBehavior.VISION_GEN_HEAD:
             data = inputs.pop("lm_hidden_state")
@@ -3056,4 +3069,4 @@ def rename_ambiguous_inputs(self, inputs):
         if self._behavior == JanusConfigBehavior.VISION_GEN_DECODER:
             data = inputs.pop("image_shape")
             inputs["shape"] = data
-        return inputs
+        return inputs

optimum/exporters/openvino/model_patcher.py

@@ -3905,7 +3905,7 @@ def __exit__(self, exc_type, exc_value, traceback):
 
 def janus_vision_embed_forward(self, pixel_values):
     from einops import rearrange
-    
+
     bs, n = pixel_values.shape[0:2]
     images = rearrange(pixel_values, "b n c h w -> (b n) c h w")
     # [b x n, T2, D]
@@ -3968,23 +3968,25 @@ def __exit__(self, exc_type, exc_value, traceback):
 
 
 class RemoveLMHeadPatcherHelper(DecoderModelPatcher):
-    def __init__(self,
+    def __init__(
+        self,
         config: "OnnxConfig",
         model: Union["PreTrainedModel", "TFPreTrainedModel"],
         model_kwargs: Dict[str, Any],
-        internal_patcher = None
+        internal_patcher=None,
     ):
         model.__orig_forward = model.forward
+
         @functools.wraps(model.__orig_forward)
         def patched_forward(*args, **kwargs):
             return model.model.forward(*args, **kwargs)
+
         model.forward = patched_forward
         self._internal_patcher = internal_patcher
         if self._internal_patcher is not None:
             self._patched_forward = self._internal_patcher.patched_forward
         super().__init__(config, model, model_kwargs)
 
-
     def __enter__(self):
         if self._internal_patcher is not None:
             return self._internal_patcher.__enter__()
@@ -4007,4 +4009,4 @@ def patched_forward(self):
     def patched_forward(self, fn):
         self._patched_forward = fn
         if self._internal_patcher is not None:
-            self._internal_patcher.patched_forward = fn
+            self._internal_patcher.patched_forward = fn

optimum/exporters/openvino/utils.py

@@ -12,6 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import importlib
 import inspect
 import logging
 from collections import namedtuple
@@ -28,7 +29,7 @@
 from optimum.intel.utils import is_transformers_version
 from optimum.intel.utils.import_utils import is_safetensors_available
 from optimum.utils import is_diffusers_available
-from optimum.utils.save_utils import maybe_save_preprocessors
+from optimum.utils.save_utils import maybe_load_preprocessors, maybe_save_preprocessors
 
 
 logger = logging.getLogger(__name__)
@@ -225,7 +226,7 @@ def get_submodels(model):
     "minicpmv",
     "phi3-v",
     "qwen2-vl",
-    "multi-modality"
+    "multi-modality",
 ]
 
 
@@ -303,3 +304,20 @@ def save_preprocessors(
                 logger.error(f"Saving {type(processor)} failed with {ex}")
     else:
         maybe_save_preprocessors(model_name_or_path, output, trust_remote_code=trust_remote_code)
+
+
+def load_preprocessors(src_name_or_path: Union[str, Path], subfolder: str = "", trust_remote_code: bool = False):
+    preprocessors = maybe_load_preprocessors(
+        src_name_or_path, subfolder=subfolder, trust_remote_code=trust_remote_code
+    )
+    if importlib.util.find_spec("janus") is not None:
+        from janus.models import VLChatProcessor
+
+        try:
+            processor = VLChatProcessor.from_pretrained(
+                src_name_or_path, subfolder=subfolder, trust_remote_code=trust_remote_code
+            )
+            preprocessors.append(processor)
+        except Exception:
+            pass
+    return preprocessors