feat: implement automodel for diffusers.

Author: sayakpaul

src/diffusers/models/auto_model.py

@@ -0,0 +1,55 @@
+import importlib
+import os
+from typing import Optional, Union
+
+from huggingface_hub.utils import validate_hf_hub_args
+
+from ..configuration_utils import ConfigMixin
+from ..utils import CONFIG_NAME
+from .modeling_utils import ModelMixin
+
+
+class AutoModel(ConfigMixin):
+    config_name = CONFIG_NAME
+
+    """TODO"""
+
+    def __init__(self, *args, **kwargs):
+        raise EnvironmentError(
+            f"{self.__class__.__name__} is designed to be instantiated "
+            f"using the `{self.__class__.__name__}.from_pretrained(pretrained_model_name_or_path)` or "
+            "from_config(config) methods."
+        )
+
+    @classmethod
+    @validate_hf_hub_args
+    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
+        """TODO"""
+        cache_dir = kwargs.pop("cache_dir", None)
+        force_download = kwargs.pop("force_download", False)
+        proxies = kwargs.pop("proxies", None)
+        token = kwargs.pop("token", None)
+        local_files_only = kwargs.pop("local_files_only", False)
+        revision = kwargs.pop("revision", None)
+        subfolder = kwargs.pop("subfolder", None)
+
+        load_config_kwargs = {
+            "cache_dir": cache_dir,
+            "force_download": force_download,
+            "proxies": proxies,
+            "token": token,
+            "local_files_only": local_files_only,
+            "revision": revision,
+            "subfolder": subfolder,
+        }
+        config = cls.load_config(pretrained_model_name_or_path, **load_config_kwargs)
+        class_name = config["_class_name"]
+        diffusers_module = importlib.import_module(__name__.split(".")[0])
+
+        try:
+            model_cls: ModelMixin = getattr(diffusers_module, class_name)
+        except Exception:
+            raise ValueError(f"Could not import the `{class_name}` class from diffusers.")
+
+        kwargs = {**load_config_kwargs, **kwargs}
+        return model_cls.from_pretrained(pretrained_model_name_or_path, **kwargs)

tests/models/test_modeling_common.py

@@ -44,6 +44,7 @@
     AttnProcessorNPU,
     XFormersAttnProcessor,
 )
+from diffusers.models.auto_model import AutoModel
 from diffusers.training_utils import EMAModel
 from diffusers.utils import (
     SAFE_WEIGHTS_INDEX_NAME,
@@ -1424,6 +1425,37 @@ def get_memory_usage(storage_dtype, compute_dtype):
             or abs(fp8_e4m3_fp32_max_memory - fp32_max_memory) < MB_TOLERANCE
         )
 
+    @parameterized.expand([None, "foo"])
+    def test_works_with_automodel(self, subfolder):
+        config, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+        model = self.model_class(**config).eval()
+        model_cls_name = model.__class__.__name__
+        model.to(torch_device)
+
+        torch.manual_seed(0)
+        output = model(**inputs_dict, return_dict=False)[0]
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = os.path.join(tmpdir, subfolder) if subfolder else tmpdir
+            model.save_pretrained(path)
+            automodel = AutoModel.from_pretrained(tmpdir, subfolder=subfolder).to(torch_device)
+
+        automodel_cls_name = automodel.__class__.__name__
+        self.assertTrue(model_cls_name == automodel_cls_name)
+        for p1, p2 in zip(model.parameters(), automodel.parameters()):
+            self.assertTrue(torch.equal(p1, p2))
+
+        torch.manual_seed(0)
+        output_automodel = model(**inputs_dict, return_dict=False)[0]
+
+        self.assertTrue(torch.allclose(output[0], output_automodel[0], atol=1e-5))
+
+    def test_automodel_raises_error_with_direct_init(self):
+        config, _ = self.prepare_init_args_and_inputs_for_common()
+        with self.assertRaises(EnvironmentError) as err_context:
+            _ = AutoModel(**config)
+        self.assertTrue("is designed to be instantiated" in str(err_context.exception))
+
 
 @is_staging_test
 class ModelPushToHubTester(unittest.TestCase):