update

src/diffusers/modular_pipelines/modular_pipeline_utils.py

@@ -309,16 +309,16 @@ class ComponentSpec:
                 f"`type_hint` is required when loading a single file model but is missing for component: {self.name}"
             )
 
+        from diffusers import AutoModel
+
         # `torch_dtype` is not an accepted parameter for tokenizers and processors.
         # As a result, it gets stored in `init_kwargs`, which are written to the config
         # during save. This causes JSON serialization to fail when saving the component.
-        if self.type_hint is not None and not issubclass(self.type_hint, torch.nn.Module):
+        if self.type_hint is not None and not issubclass(self.type_hint, (torch.nn.Module, AutoModel)):
             kwargs.pop("torch_dtype", None)
 
         if self.type_hint is None:
             try:
-                from diffusers import AutoModel
-
                 component = AutoModel.from_pretrained(pretrained_model_name_or_path, **load_kwargs, **kwargs)
             except Exception as e:
                 raise ValueError(f"Unable to load {self.name} without `type_hint`: {e}")
@@ -332,12 +332,6 @@ class ComponentSpec:
                 else getattr(self.type_hint, "from_pretrained")
             )
 
-            # `torch_dtype` is not an accepted parameter for tokenizers and processors.
-            # As a result, it gets stored in `init_kwargs`, which are written to the config
-            # during save. This causes JSON serialization to fail when saving the component.
-            if not issubclass(self.type_hint, torch.nn.Module):
-                kwargs.pop("torch_dtype", None)
-
             try:
                 component = load_method(pretrained_model_name_or_path, **load_kwargs, **kwargs)
             except Exception as e: