Skip to content

Commit c25b305

Browse files
Copilotmudler
andcommitted
Add AutoPipeline classes to pipeline registry for default model loading
Co-authored-by: mudler <[email protected]>
1 parent efcd070 commit c25b305

File tree

1 file changed

+20
-1
lines changed

1 file changed

+20
-1
lines changed

backend/python/diffusers/diffusers_dynamic_loader.py

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -219,7 +219,9 @@ def _discover_pipelines() -> Tuple[Dict[str, Type], Dict[str, List[str]]]:
219219
Discover all subclasses of DiffusionPipeline from diffusers.
220220
221221
This function uses the generic discover_diffusers_classes() internally
222-
and adds pipeline-specific task alias generation.
222+
and adds pipeline-specific task alias generation. It also includes
223+
AutoPipeline classes which are special utility classes for automatic
224+
pipeline selection.
223225
224226
Returns:
225227
A tuple of (pipeline_registry, task_aliases) where:
@@ -229,6 +231,23 @@ def _discover_pipelines() -> Tuple[Dict[str, Type], Dict[str, List[str]]]:
229231
# Use the generic discovery function
230232
pipeline_registry = discover_diffusers_classes("DiffusionPipeline", include_base=True)
231233

234+
# Also add AutoPipeline classes - these are special utility classes that are
235+
# NOT subclasses of DiffusionPipeline but are commonly used
236+
import diffusers
237+
auto_pipeline_classes = [
238+
"AutoPipelineForText2Image",
239+
"AutoPipelineForImage2Image",
240+
"AutoPipelineForInpainting",
241+
]
242+
for cls_name in auto_pipeline_classes:
243+
try:
244+
cls = getattr(diffusers, cls_name)
245+
if cls is not None:
246+
pipeline_registry[cls_name] = cls
247+
except AttributeError:
248+
# Class not available in this version of diffusers
249+
pass
250+
232251
# Generate task aliases for pipelines
233252
task_aliases: Dict[str, List[str]] = {}
234253
for attr_name in pipeline_registry:

0 commit comments

Comments
 (0)