import FlaxSchedulerMixin

Author: hlky

src/diffusers/pipelines/pipeline_utils.py

@@ -47,6 +47,7 @@
 from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, ModelMixin
 from ..quantizers.bitsandbytes.utils import _check_bnb_status
 from ..schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME, SchedulerMixin
+from ..schedulers.scheduling_utils_flax import FlaxSchedulerMixin
 from ..utils import (
     CONFIG_NAME,
     DEPRECATED_REVISION_ARGS,
@@ -848,7 +849,9 @@ def load_module(name, value):
                     _expected_class_types.append(expected_type.__name__)
 
             _is_valid_type = class_obj.__class__.__name__ in _expected_class_types
-            if (isinstance(class_obj, SchedulerMixin) or isinstance(class_obj, FlaxSchedulerMixin)) and not _is_valid_type:
+            if (
+                isinstance(class_obj, SchedulerMixin) or isinstance(class_obj, FlaxSchedulerMixin)
+            ) and not _is_valid_type:
                 _requires_flow_match = any("FlowMatch" in class_type for class_type in _expected_class_types)
                 _is_flow_match = "FlowMatch" in class_obj.__class__.__name__
                 if _requires_flow_match and not _is_flow_match: