From e4da154d1853c107e241067d19720ef4a5c5f423 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Fri, 12 Jul 2024 15:11:21 +0200 Subject: [PATCH] fix --- optimum_benchmark/task_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/optimum_benchmark/task_utils.py b/optimum_benchmark/task_utils.py index 35c533d1..08532195 100644 --- a/optimum_benchmark/task_utils.py +++ b/optimum_benchmark/task_utils.py @@ -152,7 +152,9 @@ def infer_task_from_model_name_or_path(model_name_or_path: str, revision: Option AUTO_TEXT2IMAGE_PIPELINES_MAPPING, ) - config, _ = DiffusionPipeline.load_config(model_name_or_path) + config = DiffusionPipeline.load_config(model_name_or_path) + config, _ = config if isinstance(config, tuple) else (config, None) + model_class_name = config["_class_name"] for task_name, model_mapping in (