Check correct model type is passed to from_pretrained

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
@@ -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)
@@ -832,6 +834,22 @@ def load_module(name, value):
             return True
 
         init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)}
+        scheduler_types = expected_types["scheduler"][0]
+        if isinstance(scheduler_types, enum.EnumType):
+            scheduler_types = list(scheduler_types)
+        else:
+            scheduler_types = [str(scheduler_types)]
+        scheduler_types = [str(scheduler).split(".")[-1].strip("'>") for scheduler in scheduler_types]
+
+        for key, (_, expected_class_name) in zip(init_dict.keys(), init_dict.values()):
+            if key not in passed_class_obj:
+                continue
+            class_name = passed_class_obj[key].__class__.__name__
+            class_name = class_name[4:] if class_name.startswith("Flax") else class_name
+            if key == "scheduler" and class_name not in scheduler_types:
+                raise ValueError(f"Expected {scheduler_types} for {key}, got {class_name}.")
+            elif key != "scheduler" and class_name != expected_class_name:
+                raise ValueError(f"Expected {expected_class_name} for {key}, got {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:

tests/pipelines/test_pipelines.py

@@ -1802,6 +1802,17 @@ 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 "Expected" in str(error_context.exception)
+        assert "text_encoder" in str(error_context.exception)
+        assert f"{tokenizer.__class__.__name}" in str(error_context.exception)
+
 
 @slow
 @require_torch_gpu