Support unet model

Author: titaiwangms

onnx_diagnostic/export/dynamic_shapes.py

@@ -341,7 +341,7 @@ def _generic_walker_step(
                     if any(v is not None for v in value)
                     else None
                 )
-            assert type(inputs) is dict, f"Unexpected type for inputs {type(inputs)}"
+            assert isinstance(inputs, dict), f"Unexpected type for inputs {type(inputs)}"
             assert set(inputs) == set(ds), (
                 f"Keys mismatch between inputs {set(inputs)} and ds={set(ds)}, "
                 f"inputs={string_type(inputs, with_shape=True)}, ds={ds}"

onnx_diagnostic/helpers/config_helper.py

@@ -38,12 +38,12 @@ def update_config(config: Any, mkwargs: Dict[str, Any]):
                 setattr(config, k, v)
                 continue
             existing = getattr(config, k)
-            if type(existing) is dict:
+            if isinstance(existing, dict):
                 existing.update(v)
             else:
                 update_config(getattr(config, k), v)
             continue
-        if type(config) is dict:
+        if isinstance(config, dict):
             config[k] = v
         else:
             setattr(config, k, v)
@@ -76,7 +76,7 @@ def pick(config, name: str, default_value: Any) -> Any:
     """
     if not config:
         return default_value
-    if type(config) is dict:
+    if isinstance(config, dict):
         return config.get(name, default_value)
     return getattr(config, name, default_value)
 

onnx_diagnostic/helpers/helper.py

@@ -280,7 +280,7 @@ def string_type(
             print(f"[string_type] L:{type(obj)}")
         return f"{{...}}#{len(obj)}" if with_shape else "{...}"
     # dict
-    if isinstance(obj, dict) and type(obj) is dict:
+    if isinstance(obj, dict):
         if len(obj) == 0:
             if verbose:
                 print(f"[string_type] M:{type(obj)}")

onnx_diagnostic/helpers/mini_onnx_builder.py

@@ -498,7 +498,7 @@ def _make(ty: type, res: Any) -> Any:
             for k, v in res:
                 setattr(r, k, v)
             return r
-        if ty is dict:
+        if isinstance(res, dict):
             d = {}
             for k, v in res:
                 if k.startswith("((") and k.endswith("))"):

onnx_diagnostic/helpers/torch_helper.py

@@ -717,7 +717,7 @@ def to_any(value: Any, to_value: Union[torch.dtype, torch.device, str]) -> Any:
         return tuple(to_any(t, to_value) for t in value)
     if isinstance(value, set):
         return {to_any(t, to_value) for t in value}
-    if type(value) is dict:
+    if isinstance(value, dict):
         return {k: to_any(t, to_value) for k, t in value.items()}
     if value.__class__.__name__ == "DynamicCache":
         return make_dynamic_cache(

onnx_diagnostic/torch_models/hghub/hub_api.py

@@ -10,6 +10,7 @@
 from ...helpers.config_helper import update_config
 from . import hub_data_cached_configs
 from .hub_data import __date__, __data_tasks__, load_architecture_task, __data_arch_values__
+import diffusers
 
 
 @functools.cache
@@ -119,9 +120,21 @@ def get_pretrained_config(
                 # Diffusers uses a dictionayr.
                 with open(config, "r") as f:
                     return json.load(f)
-    return transformers.AutoConfig.from_pretrained(
-        model_id, trust_remote_code=trust_remote_code, **kwargs
-    )
+    try:
+        config = transformers.AutoConfig.from_pretrained(
+            model_id, trust_remote_code=trust_remote_code, **kwargs
+        )
+    except ValueError:
+        # The model might be from diffusers, not transformers.
+        try:
+            pipe = diffusers.DiffusionPipeline.from_pretrained(
+                model_id, trust_remote_code=trust_remote_code, **kwargs
+            )
+            config = pipe.unet.config
+        except Exception as exc:
+            raise ValueError(f"Unable to retrieve the configuration for {model_id!r}") from exc
+
+    return config
 
 
 def get_model_info(model_id) -> Any:
@@ -211,7 +224,7 @@ 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:
+        if isinstance(config, 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.
@@ -362,7 +375,7 @@ def download_code_modelid(
     paths = set()
     for i, name in enumerate(pyfiles):
         if verbose:
-            print(f"[download_code_modelid] download file {i+1}/{len(pyfiles)}: {name!r}")
+            print(f"[download_code_modelid] download file {i + 1}/{len(pyfiles)}: {name!r}")
         r = hf_hub_download(repo_id=model_id, filename=name)
         p = os.path.split(r)[0]
         paths.add(p)

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -7,6 +7,7 @@
 from ...helpers.config_helper import update_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
+import diffusers
 
 
 def _code_needing_rewriting(model: Any) -> Any:
@@ -18,7 +19,7 @@ def _code_needing_rewriting(model: Any) -> Any:
 def get_untrained_model_with_inputs(
     model_id: str,
     config: Optional[Any] = None,
-    task: Optional[str] = "",
+    task: Optional[str] = None,
     inputs_kwargs: Optional[Dict[str, Any]] = None,
     model_kwargs: Optional[Dict[str, Any]] = None,
     verbose: int = 0,
@@ -88,14 +89,20 @@ def get_untrained_model_with_inputs(
             **(model_kwargs or {}),
         )
 
-    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"]]
+    # Extract architecture information from config
+    archs = None
+    if isinstance(config, dict):
+        if "_class_name" in config:
+            archs = [config["_class_name"]]
+        else:
+            raise ValueError(f"Unable to get the architecture from config={config}")
     else:
-        archs = config.architectures  # type: ignore
-    task = None
+        # Config is an object (e.g., transformers config)
+        if hasattr(config, "architecture") and config.architecture:
+            archs = [config.architecture]
+        elif hasattr(config, "architectures") and config.architectures:
+            archs = config.architectures
+
     if archs is None:
         task = task_from_id(model_id)
     assert task is not None or (archs is not None and len(archs) == 1), (
@@ -112,9 +119,9 @@ def get_untrained_model_with_inputs(
 
     # model kwagrs
     if dynamic_rope is not None:
-        assert (
-            type(config) is not dict
-        ), f"Unable to set dynamic_rope if the configuration is a dictionary\n{config}"
+        assert type(config) is not dict, (
+            f"Unable to set dynamic_rope if the configuration is a dictionary\n{config}"
+        )
         assert hasattr(config, "rope_scaling"), f"Missing 'rope_scaling' in\n{config}"
         config.rope_scaling = (
             {"rope_type": "dynamic", "factor": 10.0} if dynamic_rope else None
@@ -150,9 +157,7 @@ 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:
-        import diffusers
-
+    if isinstance(config, dict) and "_diffusers_version" in config:
         package_source = diffusers
     else:
         package_source = transformers
@@ -206,7 +211,7 @@ def get_untrained_model_with_inputs(
             )
 
         try:
-            if type(config) is dict:
+            if isinstance(config, dict):
                 model = cls_model(**config)
             else:
                 model = cls_model(config)

onnx_diagnostic/torch_models/validate.py

@@ -576,7 +576,7 @@ def validate_model(
     summary["model_config"] = str(
         shrink_config(
             data["configuration"]
-            if type(data["configuration"]) is dict
+            if isinstance(data["configuration"], dict)
             else data["configuration"].to_dict()
         )
     ).replace(" ", "")