fix withut dropping stable diffusion

huggingface · May 17, 2024 · 5323c75 · 5323c75
1 parent b814085
commit 5323c75
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/optimum_benchmark/backends/pytorch/backend.py b/optimum_benchmark/backends/pytorch/backend.py
@@ -154,7 +154,8 @@ def load_model_from_pretrained(self) -> None:
             self.logger.info("\t+ Loading Diffusion Pipeline")
             self.pretrained_model = self.automodel_class.from_pretrained(
                 # pretrained_model_name_or_path=self.config.model,
-                pretrained_model_or_path=self.config.model,
+                # pretrained_model_or_path=self.config.model,
+                self.config.model,
                 device_map=self.config.device_map,
                 **self.config.model_kwargs,
                 **self.automodel_kwargs,

diff --git a/optimum_benchmark/task_utils.py b/optimum_benchmark/task_utils.py
@@ -38,6 +38,8 @@
 }
 
 _DIFFUSERS_TASKS_TO_MODEL_LOADERS = {
+    "stable-diffusion": "StableDiffusion", # should be deprecated
+    "stable-diffusion-xl": "StableDiffusionXL", # should be deprecated
     "inpainting": "AutoPipelineForInpainting",
     "text-to-image": "AutoPipelineForText2Image",
     "image-to-image": "AutoPipelineForImage2Image",
@@ -147,7 +149,7 @@ def infer_task_from_model_name_or_path(model_name_or_path: str, revision: Option
         elif "inpainting" in model_info.tags:
             inferred_task_name = "inpainting"
         else:
-            inferred_task_name = None
+            inferred_task_name = "StableDiffusionXL" if "stable-diffusion-xl" in model_info.tags else "StableDiffusion"
 
     elif library_name == "transformers":
         if model_info.pipeline_tag is not None: