onnx-diagnostic 0.7.10__py3-none-any.whl → 0.7.12__py3-none-any.whl

onnx_diagnostic/torch_models/hghub/hub_api.py

@@ -177,6 +177,51 @@ def task_from_arch(
     return data[arch]
 
 
+def _trygetattr(config, attname):
+    try:
+        return getattr(config, attname)
+    except AttributeError:
+        return None
+
+
+def architecture_from_config(config) -> Optional[str]:
+    """Guesses the architecture (class) of the model described by this config."""
+    if isinstance(config, dict):
+        if "_class_name" in config:
+            return config["_class_name"]
+        if "architecture" in config:
+            return config["architecture"]
+        if config.get("architectures", []):
+            return config["architectures"][0]
+    if hasattr(config, "_class_name"):
+        return config._class_name
+    if hasattr(config, "architecture"):
+        return config.architecture
+    if hasattr(config, "architectures") and config.architectures:
+        return config.architectures[0]
+    if hasattr(config, "__dict__"):
+        if "_class_name" in config.__dict__:
+            return config.__dict__["_class_name"]
+        if "architecture" in config.__dict__:
+            return config.__dict__["architecture"]
+        if config.__dict__.get("architectures", []):
+            return config.__dict__["architectures"][0]
+    return None
+
+
+def find_package_source(config) -> Optional[str]:
+    """Guesses the package the class models from."""
+    if isinstance(config, dict):
+        if "_diffusers_version" in config:
+            return "diffusers"
+    if hasattr(config, "_diffusers_version"):
+        return "diffusers"
+    if hasattr(config, "__dict__"):
+        if "_diffusers_version" in config.__dict__:
+            return "diffusers"
+    return "transformers"
+
+
 def task_from_id(
     model_id: str,
     default_value: Optional[str] = None,
@@ -202,28 +247,30 @@ def task_from_id(
             if not fall_back_to_pretrained:
                 raise
     config = get_pretrained_config(model_id, subfolder=subfolder)
-    try:
-        return config.pipeline_tag
-    except AttributeError:
-        guess = _guess_task_from_config(config)
-        if guess is not None:
-            return guess
-        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_"):
-                return "fill-mask"
-        assert config.architectures is not None and len(config.architectures) == 1, (
-            f"Cannot return the task of {model_id!r}, pipeline_tag is not setup, "
-            f"architectures={config.architectures} in config={config}. "
-            f"The task can be added in "
-            f"``onnx_diagnostic.torch_models.hghub.hub_data.__data_arch__``."
-        )
-        return task_from_arch(config.architectures[0], default_value=default_value)
+    tag = _trygetattr(config, "pipeline_tag")
+    if tag is not None:
+        return tag
+
+    guess = _guess_task_from_config(config)
+    if guess is not None:
+        return guess
+    data = load_architecture_task()
+    if subfolder:
+        full_id = f"{model_id}//{subfolder}"
+        if full_id in data:
+            return data[full_id]
+    if model_id in data:
+        return data[model_id]
+    arch = architecture_from_config(config)
+    if arch is None:
+        if model_id.startswith("google/bert_"):
+            return "fill-mask"
+    assert arch is not None, (
+        f"Cannot return the task of {model_id!r}, pipeline_tag is not setup, "
+        f"config={config}. The task can be added in "
+        f"``onnx_diagnostic.torch_models.hghub.hub_data.__data_arch__``."
+    )
+    return task_from_arch(arch, default_value=default_value)
 
 
 def task_from_tags(tags: Union[str, List[str]]) -> str:
@@ -242,21 +289,17 @@ def task_from_tags(tags: Union[str, List[str]]) -> str:
 
 def enumerate_model_list(
     n: int = 50,
-    task: Optional[str] = None,
-    library: Optional[str] = None,
-    tags: Optional[Union[str, List[str]]] = None,
+    pipeline_tag: Optional[str] = None,
     search: Optional[str] = None,
     dump: Optional[str] = None,
-    filter: Optional[str] = None,
+    filter: Optional[Union[str, List[str]]] = None,
     verbose: int = 0,
 ):
     """
     Enumerates models coming from :epkg:`huggingface_hub`.
 
     :param n: number of models to retrieve (-1 for all)
-    :param task: see :meth:`huggingface_hub.HfApi.list_models`
-    :param tags: see :meth:`huggingface_hub.HfApi.list_models`
-    :param library: see :meth:`huggingface_hub.HfApi.list_models`
+    :param pipeline_tag: see :meth:`huggingface_hub.HfApi.list_models`
     :param search: see :meth:`huggingface_hub.HfApi.list_models`
     :param filter: see :meth:`huggingface_hub.HfApi.list_models`
     :param dump: dumps the result in this csv file
@@ -264,9 +307,7 @@ def enumerate_model_list(
     """
     api = HfApi()
     models = api.list_models(
-        task=task,
-        library=library,
-        tags=tags,
+        pipeline_tag=pipeline_tag,
         search=search,
         full=True,
         filter=filter,

onnx_diagnostic/torch_models/hghub/hub_data.py

@@ -30,6 +30,7 @@ __data_arch__ = textwrap.dedent(
     ConvBertModel,feature-extraction
     ConvNextForImageClassification,image-classification
     ConvNextV2Model,image-feature-extraction
+    CosmosTransformer3DModel,image-to-video
     CvtModel,feature-extraction
     DPTModel,image-feature-extraction
     Data2VecAudioModel,feature-extraction
@@ -156,7 +157,8 @@ __data_arch__ = textwrap.dedent(
     YolosForObjectDetection,object-detection
     YolosModel,image-feature-extraction
     Alibaba-NLP/gte-large-en-v1.5,sentence-similarity
-    emilyalsentzer/Bio_ClinicalBERT,fill-mask"""
+    emilyalsentzer/Bio_ClinicalBERT,fill-mask
+    nvidia/Cosmos-Predict2-2B-Video2World//transformer,image-to-video"""
 )
 
 __data_tasks__ = [

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -2,13 +2,21 @@ import copy
 import inspect
 import os
 import pprint
+import time
 from typing import Any, Dict, Optional, Tuple
 import torch
 import transformers
 from ...helpers.config_helper import update_config, build_diff_config
 from ...tasks import reduce_model_config, random_input_kwargs
-from .hub_api import task_from_arch, task_from_id, get_pretrained_config, download_code_modelid
-from .model_specific import HANDLED_MODELS, load_specific_model
+from .hub_api import (
+    task_from_arch,
+    task_from_id,
+    get_pretrained_config,
+    download_code_modelid,
+    architecture_from_config,
+    find_package_source,
+)
+from .model_specific import HANDLED_MODELS, load_specific_model, instantiate_specific_model
 
 
 def _code_needing_rewriting(model: Any) -> Any:
@@ -96,27 +104,18 @@ def get_untrained_model_with_inputs(
             model, task, config = load_specific_model(model_id, verbose=verbose)
 
     if model is None:
-        if hasattr(config, "architecture") and config.architecture:
-            archs = [config.architecture]
-        if type(config) is dict:
-            assert (
-                "_class_name" in config
-            ), f"Unable to get the architecture from config={config}"
-            archs = [config["_class_name"]]
-        else:
-            archs = config.architectures  # type: ignore
-        task = None
-        if archs is None:
-            task = task_from_id(model_id)
-        assert task is not None or (archs is not None and len(archs) == 1), (
+        arch = architecture_from_config(config)
+        if arch is None:
+            task = task_from_id(model_id, subfolder=subfolder)
+        assert task is not None or arch is not None, (
             f"Unable to determine the architecture for model {model_id!r}, "
-            f"architectures={archs!r}, conf={config}"
+            f"archs={arch!r}, conf={config}"
         )
         if verbose:
-            print(f"[get_untrained_model_with_inputs] architectures={archs!r}")
+            print(f"[get_untrained_model_with_inputs] architecture={arch!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, subfolder=subfolder)
+            task = task_from_arch(arch, model_id=model_id, subfolder=subfolder)
         if verbose:
             print(f"[get_untrained_model_with_inputs] task={task!r}")
 
@@ -170,36 +169,58 @@ def get_untrained_model_with_inputs(
                 f"{getattr(config, '_attn_implementation', '?')!r}"  # type: ignore[union-attr]
             )
 
-        if type(config) is dict and "_diffusers_version" in config:
+        if find_package_source(config) == "diffusers":
             import diffusers
 
             package_source = diffusers
         else:
             package_source = transformers
 
+        if verbose:
+            print(
+                f"[get_untrained_model_with_inputs] package_source={package_source.__name__} é"
+                f"from {package_source.__file__}"
+            )
         if use_pretrained:
+            begin = time.perf_counter()
+            if verbose:
+                print(
+                    f"[get_untrained_model_with_inputs] pretrained model_id {model_id!r}, "
+                    f"subfolder={subfolder!r}"
+                )
             model = transformers.AutoModel.from_pretrained(
-                model_id, trust_remote_code=True, **mkwargs
+                model_id, subfolder=subfolder or "", trust_remote_code=True, **mkwargs
             )
+            if verbose:
+                print(
+                    f"[get_untrained_model_with_inputs] -- done in "
+                    f"{time.perf_counter() - begin}s"
+                )
         else:
-            if archs is not None:
+            begin = time.perf_counter()
+            if verbose:
+                print(
+                    f"[get_untrained_model_with_inputs] instantiate model_id {model_id!r}, "
+                    f"subfolder={subfolder!r}"
+                )
+            if arch is not None:
                 try:
-                    cls_model = getattr(package_source, archs[0])
+                    cls_model = getattr(package_source, arch)
                 except AttributeError as e:
                     # The code of the models is not in transformers but in the
                     # repository of the model. We need to download it.
                     pyfiles = download_code_modelid(model_id, verbose=verbose)
                     if pyfiles:
-                        if "." in archs[0]:
-                            cls_name = archs[0]
+                        if "." in arch:
+                            cls_name = arch
                         else:
                             modeling = [_ for _ in pyfiles if "/modeling_" in _]
                             assert len(modeling) == 1, (
                                 f"Unable to guess the main file implemented class "
-                                f"{archs[0]!r} from {pyfiles}, found={modeling}."
+                                f"{arch!r} from {pyfiles}, found={modeling}."
                             )
                             last_name = os.path.splitext(os.path.split(modeling[0])[-1])[0]
-                            cls_name = f"{last_name}.{archs[0]}"
+                            cls_name = f"{last_name}.{arch}"
                         if verbose:
                             print(
                                 f"[get_untrained_model_with_inputs] "
@@ -217,7 +238,7 @@ def get_untrained_model_with_inputs(
                         )
                     else:
                         raise AttributeError(
-                            f"Unable to find class 'tranformers.{archs[0]}'. "
+                            f"Unable to find class 'tranformers.{arch}'. "
                             f"The code needs to be downloaded, config="
                             f"\n{pprint.pformat(config)}."
                         ) from e
@@ -225,20 +246,31 @@ def get_untrained_model_with_inputs(
                 assert same_as_pretrained and use_pretrained, (
                     f"Model {model_id!r} cannot be built, the model cannot be built. "
                     f"It must be downloaded. Use same_as_pretrained=True "
-                    f"and use_pretrained=True."
+                    f"and use_pretrained=True, arch={arch!r}, config={config}"
+                )
+            if verbose:
+                print(
+                    f"[get_untrained_model_with_inputs] -- done in "
+                    f"{time.perf_counter() - begin}s"
                 )
 
-            try:
-                if type(config) is dict:
-                    model = cls_model(**config)
-                else:
-                    model = cls_model(config)
-            except RuntimeError as e:
-                raise RuntimeError(
-                    f"Unable to instantiate class {cls_model.__name__} with\n{config}"
-                ) from e
+            seed = int(os.environ.get("SEED", "17"))
+            torch.manual_seed(seed)
+            model = instantiate_specific_model(cls_model, config)
+            if model is None:
+                try:
+                    if type(config) is dict:
+                        model = cls_model(**config)
+                    else:
+                        model = cls_model(config)
+                except RuntimeError as e:
+                    raise RuntimeError(
+                        f"Unable to instantiate class {cls_model.__name__} with\n{config}"
+                    ) from e
 
     # input kwargs
+    seed = int(os.environ.get("SEED", "17")) + 1
+    torch.manual_seed(seed)
     kwargs, fct = random_input_kwargs(config, task)  # type: ignore[arg-type]
     if verbose:
         print(f"[get_untrained_model_with_inputs] use fct={fct}")
@@ -250,7 +282,7 @@ def get_untrained_model_with_inputs(
 
     # This line is important. Some models may produce different
     # outputs even with the same inputs in training mode.
-    model.eval()
+    model.eval()  # type: ignore[union-attr]
     res = fct(model, config, add_second_input=add_second_input, **kwargs)
 
     res["input_kwargs"] = kwargs

onnx_diagnostic/torch_models/hghub/model_specific.py

@@ -1,6 +1,33 @@
 from typing import Any, Dict, Tuple
 
 
+def instantiate_specific_model(cls_model: type, config: Any) -> object:
+    """
+    Instantiates some model requiring some specific code.
+    """
+    if cls_model.__name__ == "CosmosTransformer3DModel":
+        return instantiate_CosmosTransformer3DModel(cls_model, config)
+    return None
+
+
+def instantiate_CosmosTransformer3DModel(cls_model: type, config: Any) -> object:
+    kwargs = dict(
+        in_channels=config.in_channels,
+        out_channels=config.out_channels,
+        attention_head_dim=config.attention_head_dim,
+        mlp_ratio=config.mlp_ratio,
+        num_layers=config.num_layers,
+        text_embed_dim=config.text_embed_dim,
+        adaln_lora_dim=config.adaln_lora_dim,
+        max_size=config.max_size,
+        patch_size=config.patch_size,
+        rope_scale=config.rope_scale,
+        concat_padding_mask=config.concat_padding_mask,
+        extra_pos_embed_type=config.extra_pos_embed_type,
+    )
+    return cls_model(**kwargs)
+
+
 class SpecificConfig:
     """Creates a specific configuration for the loaded model."""