Add LCM to DIFFUSERS_TASKS_TO_MODEL_LOADERS

l-bat · l-bat · commit 3058da7617d6 · 2024-03-21T12:23:12.000Z
diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -189,7 +189,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 = {

Original file line number	Diff line number	Diff line change
`@@ -189,7 +189,9 @@ class TasksManager:`
`189`	`189`
`190`	`190`	`_DIFFUSERS_TASKS_TO_MODEL_LOADERS = {`
`191`	`191`	`"stable-diffusion": "StableDiffusionPipeline",`
`192`		`- "stable-diffusion-xl": "StableDiffusionXLImg2ImgPipeline",`
	`192`	`+ "stable-diffusion-xl": "StableDiffusionXLPipeline",`
	`193`	`+ "stable-diffusion-xl-refiner": "StableDiffusionXLImg2ImgPipeline",`
	`194`	`+ "latent-consistency": "LatentConsistencyModelPipeline",`
`193`	`195`	`}`
`194`	`196`
`195`	`197`	`_TIMM_TASKS_TO_MODEL_LOADERS = {`