add pick

Author: xadupre

onnx_diagnostic/helpers/config_helper.py

@@ -69,6 +69,18 @@ def _pick(config, *atts, exceptions: Optional[Dict[str, Callable]] = None):
     raise AssertionError(f"Unable to find any of these {atts!r} in {config}")
 
 
+def pick(config, name: str, default_value: Any) -> Any:
+    """
+    Returns the vlaue of a attribute if config has it
+    otherwise the default value.
+    """
+    if not config:
+        return default_value
+    if type(config) is dict:
+        return config.get(name, default_value)
+    return getattr(config, name, default_value)
+
+
 @functools.cache
 def config_class_from_architecture(arch: str, exc: bool = False) -> Optional[type]:
     """

onnx_diagnostic/tasks/text_to_image.py

@@ -1,6 +1,6 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import update_config, check_hasattr, pick
 
 __TASK__ = "text-to-image"
 
@@ -82,10 +82,10 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
         check_hasattr(config, "sample_size", "cross_attention_dim", "in_channels")
     kwargs = dict(
         batch_size=2,
-        sequence_length=config["in_channels"],
+        sequence_length=pick(config, "in_channels", 4),
         cache_length=77,
-        in_channels=config["in_channels"],
-        sample_size=config["sample_size"],
-        cross_attention_dim=config["cross_attention_dim"],
+        in_channels=pick(config, "in_channels", 4),
+        sample_size=pick(config, "sample_size", 32),
+        cross_attention_dim=pick(config, "cross_attention_dim", 64),
     )
     return kwargs, get_inputs

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -200,10 +200,15 @@ def get_untrained_model_with_inputs(
                 f"and use_pretrained=True."
             )
 
-        if type(config) is dict:
-            model = cls_model(**config)
-        else:
-            model = cls_model(config)
+        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
     kwargs, fct = random_input_kwargs(config, task)