Refactor diffusers tasks #1947
Refactor diffusers tasks #1947
Conversation
… into "text-to-image", "image-to-image" and "inpainting"
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| if task in ["stable-diffusion", "stable-diffusion-xl"]: | ||
| logger.warning( | ||
| f"The task `{task}` is deprecated and will be removed in a future release of Optimum. " | ||
| "Please use one of the following tasks instead: `text-to-image`, `image-to-image`, `inpainting`." | ||
| ) | ||
|
|
There was a problem hiding this comment.
We warn and then it gets replaced with text-to-image by default. This is the same as the previous behavior for stable-diffusion but not stable-diffusion-xl as the model loader in that case used to be an image-to-image pipeline.
There was a problem hiding this comment.
_optional_components was introduced in diffusers v0.22.0 https://github.com/huggingface/diffusers/blob/v0.22.0/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py#L144 for StableDiffusionXLPipeline (not the case for StableDiffusionXLImg2ImgPipeline). I think we might need to upgrade DIFFUSERS_MINIMUM_VERSION as well to avoid any potential issue when loading a SDXL model with missing optional components
| @@ -260,6 +280,9 @@ class TasksManager: | |||
| "vision2seq-lm": "image-to-text", | |||
| "zero-shot-classification": "text-classification", | |||
| "image-feature-extraction": "feature-extraction", | |||
| # for backward compatibility | |||
| "stable-diffusion": "text-to-image", | |||
| "stable-diffusion-xl": "text-to-image", | |||
There was a problem hiding this comment.
should this be image-to-image for stable-diffusion-xl ? I don't think it matters
optimum/exporters/tasks.py
Outdated
| # Torch model mappings | ||
| _TRANSFORMERS_TASKS_TO_MODEL_MAPPINGS = {} | ||
| _DIFFUSERS_TASKS_TO_PIPELINE_MAPPINGS = {} | ||
|
|
||
| # TF model mappings | ||
| _TRANSFORMERS_TASKS_TO_TF_MODEL_MAPPINGS = {} |
There was a problem hiding this comment.
uses the mappings from diffusers and transformers to construct a task->model_type->model_class
optimum/exporters/tasks.py
Outdated
| for task_name in _DIFFUSERS_TASKS_TO_PIPELINE_MAPPINGS: | ||
| pipeline_mapping = _DIFFUSERS_TASKS_TO_PIPELINE_MAPPINGS[task_name] | ||
| _DIFFUSERS_TASKS_TO_PIPELINE_MAPPINGS[task_name] = { | ||
| pipeline_name: pipeline_class.__name__ for pipeline_name, pipeline_class in pipeline_mapping.items() | ||
| } |
There was a problem hiding this comment.
transforming classes into class names to conform with transformers mappings
| is_stable_diffusion = isinstance( | ||
| pipeline, (StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, StableDiffusionInpaintPipeline) | ||
| ) | ||
| is_stable_diffusion_xl = isinstance( | ||
| pipeline, (StableDiffusionXLImg2ImgPipeline, StableDiffusionXLInpaintPipeline, StableDiffusionXLPipeline) | ||
| ) | ||
| is_latent_consistency_model = isinstance( | ||
| pipeline, (LatentConsistencyModelPipeline, LatentConsistencyModelImg2ImgPipeline) | ||
| ) |
There was a problem hiding this comment.
could be simplified using an _infer_model_type_from_model_or_model_class method and _DIFFUSERS_TASKS_TO_PIPELINE_MAPPINGS
| if task in ["stable-diffusion", "stable-diffusion-xl"]: | ||
| logger.warning( | ||
| f"The task `{task}` is deprecated and will be removed in a future release of Optimum. " | ||
| "Please use one of the following tasks instead: `text-to-image`, `image-to-image`, `inpainting`." | ||
| ) | ||
|
|
There was a problem hiding this comment.
_optional_components was introduced in diffusers v0.22.0 https://github.com/huggingface/diffusers/blob/v0.22.0/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py#L144 for StableDiffusionXLPipeline (not the case for StableDiffusionXLImg2ImgPipeline). I think we might need to upgrade DIFFUSERS_MINIMUM_VERSION as well to avoid any potential issue when loading a SDXL model with missing optional components
What does this PR do?
This PR refactors "stable-diffusion" and "stable-diffusion-xl" into "text-to-image", "image-to-image" and "inpainting" tasks.
"stable-diffusion" and "stable-diffusion-xl" become model types in this case, with extensibility to other model types like "lcm" (latent consistency, control net, etc).
This PR also adds new class variables to the
TasksManagermostly to create a mapping from tasks to model types to model classes. Which is used to infer things like task and model type from model or model class (e.g. text-to-image and stable-diffusion fromStableDiffusionPipeline).Before submitting
Who can review?