Extend dynamic loader to discover any diffusers class type, not just DiffusionPipeline

Co-authored-by: mudler <[email protected]>

Author: Copilot

backend/python/diffusers/README.md

@@ -33,6 +33,8 @@ from diffusers_dynamic_loader import (
     get_available_pipelines,
     get_available_tasks,
     resolve_pipeline_class,
+    discover_diffusers_classes,
+    get_available_classes,
 )
 
 # List all available pipelines
@@ -63,6 +65,28 @@ pipe = load_diffusers_pipeline(
     from_single_file=True,
     torch_dtype=torch.float16
 )
+
+# Discover other diffusers classes (schedulers, models, etc.)
+schedulers = discover_diffusers_classes("SchedulerMixin")
+print(f"Available schedulers: {list(schedulers.keys())[:5]}...")
+
+# Get list of available scheduler classes
+scheduler_list = get_available_classes("SchedulerMixin")
+```
+
+### Generic Class Discovery
+
+The dynamic loader can discover not just pipelines but any class type from diffusers:
+
+```python
+# Discover all scheduler classes
+schedulers = discover_diffusers_classes("SchedulerMixin")
+
+# Discover all model classes
+models = discover_diffusers_classes("ModelMixin")
+
+# Get a sorted list of available classes
+scheduler_names = get_available_classes("SchedulerMixin")
 ```
 
 ### Special Pipeline Handling

backend/python/diffusers/diffusers_dynamic_loader.py

@@ -5,6 +5,9 @@
 eliminating the need for per-pipeline conditional statements. New pipelines added to
 diffusers become available automatically without code changes.
 
+The module also supports discovering other diffusers classes like schedulers, models,
+and other components, making it a generic solution for dynamic class loading.
+
 Usage:
     from diffusers_dynamic_loader import load_diffusers_pipeline, get_available_pipelines
 
@@ -19,6 +22,10 @@
 
     # Get list of available pipelines
     available = get_available_pipelines()
+
+    # Discover other diffusers classes (schedulers, models, etc.)
+    schedulers = discover_diffusers_classes("SchedulerMixin")
+    models = discover_diffusers_classes("ModelMixin")
 """
 
 import importlib
@@ -31,6 +38,9 @@
 _pipeline_registry: Optional[Dict[str, Type]] = None
 _task_aliases: Optional[Dict[str, List[str]]] = None
 
+# Global cache for other discovered class types
+_class_registries: Dict[str, Dict[str, Type]] = {}
+
 
 def _camel_to_kebab(name: str) -> str:
     """
@@ -111,50 +121,126 @@ def _extract_task_keywords(class_name: str) -> List[str]:
     return list(set(aliases))  # Remove duplicates
 
 
-def _discover_pipelines() -> Tuple[Dict[str, Type], Dict[str, List[str]]]:
+def discover_diffusers_classes(
+    base_class_name: str,
+    include_base: bool = True
+) -> Dict[str, Type]:
     """
-    Discover all subclasses of DiffusionPipeline from diffusers.pipelines.
+    Discover all subclasses of a given base class from diffusers.
+
+    This function provides a generic way to discover any type of diffusers class,
+    not just pipelines. It can be used to discover schedulers, models, processors,
+    and other components.
 
-    This function imports diffusers.pipelines modules and collects all classes
-    that are subclasses of DiffusionPipeline.
+    Args:
+        base_class_name: Name of the base class to search for subclasses
+                        (e.g., "DiffusionPipeline", "SchedulerMixin", "ModelMixin")
+        include_base: Whether to include the base class itself in results
 
     Returns:
-        A tuple of (pipeline_registry, task_aliases) where:
-        - pipeline_registry: Dict mapping class names to class objects
-        - task_aliases: Dict mapping task aliases to lists of class names
-    """
-    from diffusers import DiffusionPipeline
+        Dict mapping class names to class objects
 
-    pipeline_registry: Dict[str, Type] = {}
-    task_aliases: Dict[str, List[str]] = {}
+    Examples:
+        # Discover all pipeline classes
+        pipelines = discover_diffusers_classes("DiffusionPipeline")
+
+        # Discover all scheduler classes
+        schedulers = discover_diffusers_classes("SchedulerMixin")
+
+        # Discover all model classes
+        models = discover_diffusers_classes("ModelMixin")
+
+        # Discover AutoPipeline classes
+        auto_pipelines = discover_diffusers_classes("AutoPipelineForText2Image")
+    """
+    global _class_registries
 
-    # Also add DiffusionPipeline itself as it's a valid generic pipeline
-    pipeline_registry['DiffusionPipeline'] = DiffusionPipeline
+    # Check cache first
+    if base_class_name in _class_registries:
+        return _class_registries[base_class_name]
 
-    # Get all pipeline classes that are exposed in diffusers
     import diffusers
+
+    # Try to get the base class from diffusers
+    base_class = None
+    try:
+        base_class = getattr(diffusers, base_class_name)
+    except AttributeError:
+        # Try to find in submodules
+        for submodule in ['schedulers', 'models', 'pipelines']:
+            try:
+                module = importlib.import_module(f'diffusers.{submodule}')
+                if hasattr(module, base_class_name):
+                    base_class = getattr(module, base_class_name)
+                    break
+            except (ImportError, ModuleNotFoundError):
+                continue
+
+    if base_class is None:
+        raise ValueError(f"Could not find base class '{base_class_name}' in diffusers")
+
+    registry: Dict[str, Type] = {}
+
+    # Include base class if requested
+    if include_base:
+        registry[base_class_name] = base_class
+
+    # Scan diffusers module for subclasses
     for attr_name in dir(diffusers):
         try:
             attr = getattr(diffusers, attr_name)
-            # Check if it's a class and a subclass of DiffusionPipeline
             if (isinstance(attr, type) and
-                issubclass(attr, DiffusionPipeline) and
-                attr is not DiffusionPipeline):
+                issubclass(attr, base_class) and
+                (include_base or attr is not base_class)):
+                registry[attr_name] = attr
+        except (ImportError, AttributeError, TypeError, RuntimeError, ModuleNotFoundError):
+            continue
 
-                pipeline_registry[attr_name] = attr
+    # Cache the results
+    _class_registries[base_class_name] = registry
+    return registry
 
-                # Generate task aliases for this pipeline
-                aliases = _extract_task_keywords(attr_name)
-                for alias in aliases:
-                    if alias not in task_aliases:
-                        task_aliases[alias] = []
-                    if attr_name not in task_aliases[alias]:
-                        task_aliases[alias].append(attr_name)
 
-        except (ImportError, AttributeError, TypeError, RuntimeError, ModuleNotFoundError):
-            # Skip any problematic attributes - some pipelines may have
-            # missing optional dependencies (e.g., ftfy, sentencepiece)
-            continue
+def get_available_classes(base_class_name: str) -> List[str]:
+    """
+    Get a sorted list of all discovered class names for a given base class.
+
+    Args:
+        base_class_name: Name of the base class (e.g., "SchedulerMixin")
+
+    Returns:
+        Sorted list of discovered class names
+    """
+    return sorted(discover_diffusers_classes(base_class_name).keys())
+
+
+def _discover_pipelines() -> Tuple[Dict[str, Type], Dict[str, List[str]]]:
+    """
+    Discover all subclasses of DiffusionPipeline from diffusers.
+
+    This function uses the generic discover_diffusers_classes() internally
+    and adds pipeline-specific task alias generation.
+
+    Returns:
+        A tuple of (pipeline_registry, task_aliases) where:
+        - pipeline_registry: Dict mapping class names to class objects
+        - task_aliases: Dict mapping task aliases to lists of class names
+    """
+    # Use the generic discovery function
+    pipeline_registry = discover_diffusers_classes("DiffusionPipeline", include_base=True)
+
+    # Generate task aliases for pipelines
+    task_aliases: Dict[str, List[str]] = {}
+    for attr_name in pipeline_registry:
+        if attr_name == "DiffusionPipeline":
+            continue  # Skip base class for alias generation
+
+        aliases = _extract_task_keywords(attr_name)
+        for alias in aliases:
+            if alias not in task_aliases:
+                task_aliases[alias] = []
+            if attr_name not in task_aliases[alias]:
+                task_aliases[alias].append(attr_name)
 
     return pipeline_registry, task_aliases
 

backend/python/diffusers/test.py

@@ -256,6 +256,43 @@ def test_get_pipeline_info_unknown_raises(self):
 
         self.assertIn("Unknown pipeline", str(ctx.exception))
 
+    def test_discover_diffusers_classes_pipelines(self):
+        """Test generic class discovery for DiffusionPipeline."""
+        classes = loader.discover_diffusers_classes("DiffusionPipeline")
+
+        # Should return a dict
+        self.assertIsInstance(classes, dict)
+
+        # Should contain known pipeline classes
+        self.assertIn("DiffusionPipeline", classes)
+        self.assertIn("StableDiffusionPipeline", classes)
+
+    def test_discover_diffusers_classes_caches_results(self):
+        """Test that class discovery results are cached."""
+        classes1 = loader.discover_diffusers_classes("DiffusionPipeline")
+        classes2 = loader.discover_diffusers_classes("DiffusionPipeline")
+
+        # Should be the same object (cached)
+        self.assertIs(classes1, classes2)
+
+    def test_discover_diffusers_classes_exclude_base(self):
+        """Test discovering classes without base class."""
+        classes = loader.discover_diffusers_classes("DiffusionPipeline", include_base=False)
+
+        # Should still contain subclasses
+        self.assertIn("StableDiffusionPipeline", classes)
+
+    def test_get_available_classes(self):
+        """Test getting list of available classes for a base class."""
+        classes = loader.get_available_classes("DiffusionPipeline")
+
+        # Should return a sorted list
+        self.assertIsInstance(classes, list)
+        self.assertEqual(classes, sorted(classes))
+
+        # Should contain known classes
+        self.assertIn("StableDiffusionPipeline", classes)
+
 
 class TestDiffusersDynamicLoaderWithMocks(unittest.TestCase):
     """Test cases using mocks to test edge cases."""