Add getattr for fetching library_name

huggingface · Oct 17, 2023 · 5eb4120 · 5eb4120
1 parent e7bd60d
commit 5eb4120
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -1391,13 +1391,13 @@ def _infer_task_from_model_name_or_path(
                     "Cannot infer the task from a model repo with a subfolder yet, please specify the task manually."
                 )
             model_info = huggingface_hub.model_info(model_name_or_path, revision=revision)
-            if model_info.library_name == "diffusers":
+            if getattr(model_info, "library_name", None) == "diffusers":
                 # TODO : getattr(model_info, "model_index") defining auto_model_class_name currently set to None
                 for task in ("stable-diffusion-xl", "stable-diffusion"):
                     if task in model_info.tags:
                         inferred_task_name = task
                         break
-            elif model_info.library_name == "timm":
+            elif getattr(model_info, "library_name", None) == "timm":
                 inferred_task_name = "image-classification"
             else:
                 pipeline_tag = getattr(model_info, "pipeline_tag", None)