Check correct model type is passed to from_pretrained

src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py

@@ -17,7 +17,7 @@
 import torch
 
 from ...models import UNet2DConditionModel, VQModel
-from ...schedulers import DDPMScheduler
+from ...schedulers import DDPMScheduler, UnCLIPScheduler
 from ...utils import deprecate, logging, replace_example_docstring
 from ...utils.torch_utils import randn_tensor
 from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
@@ -83,7 +83,7 @@ class KandinskyV22Pipeline(DiffusionPipeline):
     def __init__(
         self,
         unet: UNet2DConditionModel,
-        scheduler: DDPMScheduler,
+        scheduler: Union[DDPMScheduler, UnCLIPScheduler],
         movq: VQModel,
     ):
         super().__init__()

src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py

@@ -140,13 +140,13 @@ class KandinskyV22CombinedPipeline(DiffusionPipeline):
     def __init__(
         self,
         unet: UNet2DConditionModel,
-        scheduler: DDPMScheduler,
+        scheduler: Union[DDPMScheduler, UnCLIPScheduler],
         movq: VQModel,
         prior_prior: PriorTransformer,
         prior_image_encoder: CLIPVisionModelWithProjection,
         prior_text_encoder: CLIPTextModelWithProjection,
         prior_tokenizer: CLIPTokenizer,
-        prior_scheduler: UnCLIPScheduler,
+        prior_scheduler: Union[DDPMScheduler, UnCLIPScheduler],
         prior_image_processor: CLIPImageProcessor,
     ):
         super().__init__()

src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py

@@ -5,7 +5,7 @@
 from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection
 
 from ...models import PriorTransformer
-from ...schedulers import UnCLIPScheduler
+from ...schedulers import DDPMScheduler, UnCLIPScheduler
 from ...utils import (
     logging,
     replace_example_docstring,
@@ -114,7 +114,7 @@ def __init__(
         image_encoder: CLIPVisionModelWithProjection,
         text_encoder: CLIPTextModelWithProjection,
         tokenizer: CLIPTokenizer,
-        scheduler: UnCLIPScheduler,
+        scheduler: Union[DDPMScheduler, UnCLIPScheduler],
         image_processor: CLIPImageProcessor,
     ):
         super().__init__()

src/diffusers/pipelines/pipeline_utils.py

@@ -13,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import enum
 import fnmatch
 import importlib
 import inspect
@@ -45,7 +46,7 @@
 from ..models.attention_processor import FusedAttnProcessor2_0
 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
+from ..schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME, SchedulerMixin
 from ..utils import (
     CONFIG_NAME,
     DEPRECATED_REVISION_ARGS,
@@ -811,6 +812,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
         # in this case they are already instantiated in `kwargs`
         # extract them here
         expected_modules, optional_kwargs = cls._get_signature_keys(pipeline_class)
+        expected_types = pipeline_class._get_signature_types()
         passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs}
         passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs}
         init_dict, unused_kwargs, _ = pipeline_class.extract_init_dict(config_dict, **kwargs)
@@ -833,6 +835,31 @@ def load_module(name, value):
 
         init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)}
 
+        for key in init_dict.keys():
+            if key not in passed_class_obj:
+                continue
+
+            class_obj = passed_class_obj[key]
+            _expected_class_types = []
+            for expected_type in expected_types[key]:
+                if isinstance(expected_type, enum.EnumMeta):
+                    _expected_class_types.extend(expected_type.__members__.keys())
+                else:
+                    _expected_class_types.append(expected_type.__name__)
+
+            _is_valid_type = class_obj.__class__.__name__ in _expected_class_types
+            if isinstance(class_obj, SchedulerMixin) 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:
+                    logger.warning(f"Expected FlowMatch scheduler, got {class_obj.__class__.__name__}.")
+                elif not _requires_flow_match and _is_flow_match:
+                    logger.warning(f"Expected non-FlowMatch scheduler, got {class_obj.__class__.__name__}.")
+            elif not _is_valid_type:
+                logger.warning(
+                    f"Expected types for {key}: {_expected_class_types}, got {class_obj.__class__.__name__}."
+                )
+
         # Special case: safety_checker must be loaded separately when using `from_flax`
         if from_flax and "safety_checker" in init_dict and "safety_checker" not in passed_class_obj:
             raise NotImplementedError(

src/diffusers/schedulers/scheduling_utils.py

@@ -46,6 +46,7 @@ class KarrasDiffusionSchedulers(Enum):
     UniPCMultistepScheduler = 13
     DPMSolverSDEScheduler = 14
     EDMEulerScheduler = 15
+    LCMScheduler = 16
 
 
 AysSchedules = {

tests/pipelines/stable_diffusion_2/test_stable_diffusion_latent_upscale.py

@@ -250,6 +250,7 @@ def test_karras_schedulers_shape(self):
             "KDPM2AncestralDiscreteScheduler",
             "DPMSolverSDEScheduler",
             "EDMEulerScheduler",
+            "LCMScheduler",
         ]
         components = self.get_dummy_components()
         pipe = self.pipeline_class(**components)

tests/pipelines/test_pipelines.py

@@ -47,6 +47,8 @@
     DPMSolverMultistepScheduler,
     EulerAncestralDiscreteScheduler,
     EulerDiscreteScheduler,
+    FlowMatchEulerDiscreteScheduler,
+    FluxPipeline,
     LMSDiscreteScheduler,
     ModelMixin,
     PNDMScheduler,
@@ -1802,6 +1804,42 @@ def test_pipe_same_device_id_offload(self):
         sd.maybe_free_model_hooks()
         assert sd._offload_gpu_id == 5
 
+    def test_wrong_model(self):
+        tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
+        with self.assertRaises(ValueError) as error_context:
+            _ = StableDiffusionPipeline.from_pretrained(
+                "hf-internal-testing/diffusers-stable-diffusion-tiny-all", text_encoder=tokenizer
+            )
+
+        assert "is of type" in str(error_context.exception)
+        assert "but should be" in str(error_context.exception)
+
+    def test_wrong_model_scheduler_type(self):
+        scheduler = EulerDiscreteScheduler.from_pretrained("hf-internal-testing/tiny-flux-pipe", subfolder="scheduler")
+        with self.assertLogs(
+            logging.get_logger("diffusers.pipelines.pipeline_utils"), level="WARNING"
+        ) as warning_context:
+            _ = FluxPipeline.from_pretrained("hf-internal-testing/tiny-flux-pipe", scheduler=scheduler)
+
+        assert any("Expected" in message for message in warning_context.output)
+        assert any("scheduler" in message for message in warning_context.output)
+        assert any("EulerDiscreteScheduler" in message for message in warning_context.output)
+
+    def test_wrong_model_scheduler_enum(self):
+        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            "hf-internal-testing/diffusers-stable-diffusion-tiny-all", subfolder="scheduler"
+        )
+        with self.assertLogs(
+            logging.get_logger("diffusers.pipelines.pipeline_utils"), level="WARNING"
+        ) as warning_context:
+            _ = StableDiffusionPipeline.from_pretrained(
+                "hf-internal-testing/diffusers-stable-diffusion-tiny-all", scheduler=scheduler
+            )
+
+        assert any("Expected" in message for message in warning_context.output)
+        assert any("scheduler" in message for message in warning_context.output)
+        assert any("FlowMatchEulerDiscreteScheduler" in message for message in warning_context.output)
+
 
 @slow
 @require_torch_gpu