diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py index 7ccb0d9c7b..b74ccf8b20 100644 --- a/optimum/exporters/tasks.py +++ b/optimum/exporters/tasks.py @@ -188,7 +188,9 @@ class TasksManager: _DIFFUSERS_TASKS_TO_MODEL_LOADERS = { "stable-diffusion": "StableDiffusionPipeline", - "stable-diffusion-xl": "StableDiffusionXLImg2ImgPipeline", + "stable-diffusion-xl": "StableDiffusionXLPipeline", + "stable-diffusion-xl-refiner": "StableDiffusionXLImg2ImgPipeline", + "latent-consistency": "LatentConsistencyModelPipeline", } _TIMM_TASKS_TO_MODEL_LOADERS = { @@ -1561,7 +1563,12 @@ def _infer_task_from_model_name_or_path( if library_name == "diffusers": class_name = model_info.config["diffusers"]["class_name"] - inferred_task_name = "stable-diffusion-xl" if "StableDiffusionXL" in class_name else "stable-diffusion" + if "LatentConsistency" in class_name: + inferred_task_name = "latent-consistency" + elif "StableDiffusionXL" in class_name: + inferred_task_name = f"stable-diffusion-xl{'-refiner' if 'Img2Img' in class_name else ''}" + else: + inferred_task_name = "stable-diffusion" elif library_name == "timm": inferred_task_name = "image-classification" else: diff --git a/optimum/utils/import_utils.py b/optimum/utils/import_utils.py index a5df9e2624..4a57fda79c 100644 --- a/optimum/utils/import_utils.py +++ b/optimum/utils/import_utils.py @@ -50,7 +50,7 @@ def _is_package_available(pkg_name: str, return_version: bool = False) -> Union[ TORCH_MINIMUM_VERSION = version.parse("1.11.0") TRANSFORMERS_MINIMUM_VERSION = version.parse("4.25.0") -DIFFUSERS_MINIMUM_VERSION = version.parse("0.18.0") +DIFFUSERS_MINIMUM_VERSION = version.parse("0.22.0") AUTOGPTQ_MINIMUM_VERSION = version.parse("0.4.99") # Allows 0.5.0.dev0