include patch_diffusers=True when using --patch

Author: xadupre

onnx_diagnostic/torch_models/hghub/hub_api.py

@@ -140,14 +140,18 @@ def _guess_task_from_config(config: Any) -> Optional[str]:
 
 @functools.cache
 def task_from_arch(
-    arch: str, default_value: Optional[str] = None, model_id: Optional[str] = None
+    arch: str,
+    default_value: Optional[str] = None,
+    model_id: Optional[str] = None,
+    subfolder: Optional[str] = None,
 ) -> str:
     """
     This function relies on stored information. That information needs to be refresh.
 
     :param arch: architecture name
     :param default_value: default value in case the task cannot be determined
     :param model_id: unused unless the architecture does not help.
+    :param subfolder: subfolder
     :return: task
 
     .. runpython::
@@ -162,7 +166,7 @@ def task_from_arch(
     data = load_architecture_task()
     if arch not in data and model_id:
         # Let's try with the model id.
-        return task_from_id(model_id)
+        return task_from_id(model_id, subfolder=subfolder)
     if default_value is not None:
         return data.get(arch, default_value)
     assert arch in data, (
@@ -178,6 +182,7 @@ def task_from_id(
     default_value: Optional[str] = None,
     pretrained: bool = False,
     fall_back_to_pretrained: bool = True,
+    subfolder: Optional[str] = None,
 ) -> str:
     """
     Returns the task attached to a model id.
@@ -187,7 +192,7 @@ def task_from_id(
         if the task cannot be determined
     :param pretrained: uses the config
     :param fall_back_to_pretrained: falls back to pretrained config
-    :param exc: raises an exception if True
+    :param subfolder: subfolder
     :return: task
     """
     if not pretrained:
@@ -196,7 +201,7 @@ def task_from_id(
         except RuntimeError:
             if not fall_back_to_pretrained:
                 raise
-    config = get_pretrained_config(model_id)
+    config = get_pretrained_config(model_id, subfolder=subfolder)
     try:
         return config.pipeline_tag
     except AttributeError:
@@ -206,6 +211,8 @@ def task_from_id(
         data = load_architecture_task()
         if model_id in data:
             return data[model_id]
+        if type(config) is dict and "_class_name" in config:
+            return task_from_arch(config["_class_name"], default_value=default_value)
         if not config.architectures or not config.architectures:
             # Some hardcoded values until a better solution is found.
             if model_id.startswith("google/bert_"):

onnx_diagnostic/torch_models/hghub/hub_data.py

@@ -22,6 +22,7 @@
     BlenderbotModel,feature-extraction
     BloomModel,feature-extraction
     CLIPModel,zero-shot-image-classification
+    CLIPTextModel,feature-extraction
     CLIPVisionModel,feature-extraction
     CamembertModel,feature-extraction
     CodeGenModel,feature-extraction

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -106,7 +106,7 @@ def get_untrained_model_with_inputs(
         print(f"[get_untrained_model_with_inputs] architectures={archs!r}")
         print(f"[get_untrained_model_with_inputs] cls={config.__class__.__name__!r}")
     if task is None:
-        task = task_from_arch(archs[0], model_id=model_id)
+        task = task_from_arch(archs[0], model_id=model_id, subfolder=subfolder)
     if verbose:
         print(f"[get_untrained_model_with_inputs] task={task!r}")
 

onnx_diagnostic/torch_models/validate.py

@@ -263,7 +263,7 @@ def validate_model(
     use_pretrained: bool = False,
     optimization: Optional[str] = None,
     quiet: bool = False,
-    patch: bool = False,
+    patch: Union[bool, str, Dict[str, bool]] = False,
     rewrite: bool = False,
     stop_if_static: int = 1,
     dump_folder: Optional[str] = None,
@@ -301,8 +301,10 @@ def validate_model(
     :param optimization: optimization to apply to the exported model,
         depend on the the exporter
     :param quiet: if quiet, catches exception if any issue
-    :param patch: applies patches (``patch_transformers=True``) before exporting,
-        see :func:`onnx_diagnostic.torch_export_patches.torch_export_patches`
+    :param patch: applies patches (``patch_transformers=True, path_diffusers=True``)
+        if True before exporting
+        see :func:`onnx_diagnostic.torch_export_patches.torch_export_patches`,
+        a string can be used to specify only one of them
     :param rewrite: applies known rewriting (``patch_transformers=True``) before exporting,
         see :func:`onnx_diagnostic.torch_export_patches.torch_export_patches`
     :param stop_if_static: stops if a dynamic dimension becomes static,
@@ -346,6 +348,13 @@ def validate_model(
     exported model returns the same outputs as the original one, otherwise,
     :class:`onnx_diagnostic.reference.TorchOnnxEvaluator` is used.
     """
+    if isinstance(patch, bool):
+        patch = dict(patch_transformers=True, patch_diffusers=True) if patch else {}
+    elif isinstance(patch, str):
+        patch = {p: True for p in patch.split(",")}  # noqa: C420
+    else:
+        assert isinstance(patch, dict), f"Unable to interpret patch={patch!r}"
+
     assert not rewrite or patch, (
         f"rewrite={rewrite}, patch={patch}, patch must be True to enable rewriting, "
         f"if --no-patch was specified on the command line, --no-rewrite must be added."
@@ -580,11 +589,11 @@ def validate_model(
                     f"stop_if_static={stop_if_static}"
                 )
             with torch_export_patches(  # type: ignore
-                patch_transformers=True,
                 stop_if_static=stop_if_static,
                 verbose=max(0, verbose - 1),
                 rewrite=data.get("rewrite", None),
                 dump_rewriting=(os.path.join(dump_folder, "rewrite") if dump_folder else None),
+                **patch,
             ) as modificator:
                 data["inputs_export"] = modificator(data["inputs"])  # type: ignore