`DiffusionPipeline` cannot serialize or load custom huggingface `transformers` model

[ChenchaoZhao]

Is your feature request related to a problem? Please describe. Currently DiffusionPipeline class can only serialize official transformers models correctly and load them using from_pretrained. However, it does not work for custom huggingface models which could be serialized using transformers and loaded using transformers.AutoModel.from_pretrained

Describe the solution you'd like. Adding support for custom huggingface models.

Describe alternatives you've considered. There is a hacky way to do it but I don't want to keep it long term.

Additional context. The error message I got was "the model class is not one of the loadable classes ModelMixin, ...."

[sayakpaul]

Can you post a fully reproducible code snippet?

[ChenchaoZhao]

class CustomConfig(transformers.PretrainedConfig)
  model_type = "custom"

  def __init__(self, in_features=1, out_features=1)
    
    self.in_features = in_features
    self.out_features = out_features


class CustomModel(transformers.PretrainedModel)

 config_type = CustomConfig

  def __init__(self, config):
    super().__init__(config)
    self.f = nn.Linear(config.in_features, config.out_features)



# in a different file pipeline.py
class MyPipeline(diffusers.DiffusionPipeline)

  def __init__(self, model):

    self.register_modules(model)

Instantiate a pipeline and save it as pretrained then load it using DiffusionPipeline.from_pretrained('localdir', custom_pipeline='localdir/mypipe')

[sayakpaul]

Have you followed https://huggingface.co/docs/diffusers/en/using-diffusers/custom_pipeline_overview#community-components ?

The implementation of the CustomModel also needs to follow https://huggingface.co/docs/transformers/en/custom_models. Make sure you follow all of this.

If none of these work, please share a minimally reproducible code snippet.

[ChenchaoZhao]

Yes each of them works individually but they do not work together. Updated the minimum code above.

[sayakpaul]

I will just repeat myself saying that we then need a minimally reproducible snippet to help you. Until and unless that's provided we won't be able to take any further actions.

[ChenchaoZhao]

The model can be whatever custom model

class TestConfig(transformers.PretrainedConfig):
    
    model_type = "test_model"
    
    def __init__(self, d_model: int=None, n_class: int=None, bias: bool=True, **kwargs):
        super().__init__(**kwargs)
        self.d_model = d_model
        self.n_class = n_class
        self.bias = bias


class TestModel(transformers.PreTrainedModel):
    
    config_class = TestConfig
    
    def __init__(self, config: TestConfig):
        super().__init__(config)
        
        self.linear = torch.nn.Linear(
            in_features=config.d_model, 
            out_features=config.n_class,
            bias=config.bias
            )
        
    def forward(self, inputs):
        return dict(logits=self.linear(inputs))

transformers.AutoConfig.register("test_model", TestConfig)
transformers.AutoModel.register(TestConfig, TestModel)

TestConfig.register_for_auto_class()
TestModel.register_for_auto_class(auto_class="AutoModel")

Define the pipeline

class MyPipe(diffusers.DiffusionPipeline):
    
    def __init__(self, model):
        self.register_modules(model=model)

Then do pipe = MyPipe(model=test_model) and then pipe.save_pretrained('my-pipe') Then do MyPipe.from_pretrained('my-pipe')

{
	"name": "ValueError",
	"message": "The component <class 'transformers_modules.model.define.TestModel'> of <class '__main__.MyPipe'> cannot be loaded as it does not seem to have any of the loading methods defined in {'ModelMixin': ['save_pretrained', 'from_pretrained'], 'SchedulerMixin': ['save_pretrained', 'from_pretrained'], 'DiffusionPipeline': ['save_pretrained', 'from_pretrained'], 'OnnxRuntimeModel': ['save_pretrained', 'from_pretrained'], 'PreTrainedTokenizer': ['save_pretrained', 'from_pretrained'], 'PreTrainedTokenizerFast': ['save_pretrained', 'from_pretrained'], 'PreTrainedModel': ['save_pretrained', 'from_pretrained'], 'FeatureExtractionMixin': ['save_pretrained', 'from_pretrained'], 'ProcessorMixin': ['save_pretrained', 'from_pretrained'], 'ImageProcessingMixin': ['save_pretrained', 'from_pretrained'], 'ORTModule': ['save_pretrained', 'from_pretrained']}.",
	"stack": "---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/Users/cczhao/work/notebooks/test-diffusion-pipeline/export.ipynb Cell 6 line 1
----> <a href='vscode-notebook-cell:/Users/cczhao/work/notebooks/test-diffusion-pipeline/export.ipynb#W5sZmlsZQ%3D%3D?line=0'>1</a> MyPipe.from_pretrained('pipe-model')

File ~/venv/dev/lib/python3.10/site-packages/huggingface_hub/utils/_validators.py:118, in validate_hf_hub_args.<locals>._inner_fn(*args, **kwargs)
    115 if check_use_auth_token:
    116     kwargs = smoothly_deprecate_use_auth_token(fn_name=fn.__name__, has_token=has_token, kwargs=kwargs)
--> 118 return fn(*args, **kwargs)

File ~/venv/dev/lib/python3.10/site-packages/diffusers/pipelines/pipeline_utils.py:1286, in DiffusionPipeline.from_pretrained(cls, pretrained_model_name_or_path, **kwargs)
   1283     loaded_sub_model = passed_class_obj[name]
   1284 else:
   1285     # load sub model
-> 1286     loaded_sub_model = load_sub_model(
   1287         library_name=library_name,
   1288         class_name=class_name,
   1289         importable_classes=importable_classes,
   1290         pipelines=pipelines,
   1291         is_pipeline_module=is_pipeline_module,
   1292         pipeline_class=pipeline_class,
   1293         torch_dtype=torch_dtype,
   1294         provider=provider,
   1295         sess_options=sess_options,
   1296         device_map=device_map,
   1297         max_memory=max_memory,
   1298         offload_folder=offload_folder,
   1299         offload_state_dict=offload_state_dict,
   1300         model_variants=model_variants,
   1301         name=name,
   1302         from_flax=from_flax,
   1303         variant=variant,
   1304         low_cpu_mem_usage=low_cpu_mem_usage,
   1305         cached_folder=cached_folder,
   1306         revision=revision,
   1307     )
   1308     logger.info(
   1309         f\"Loaded {name} as {class_name} from `{name}` subfolder of {pretrained_model_name_or_path}.\"
   1310     )
   1312 init_kwargs[name] = loaded_sub_model  # UNet(...), # DiffusionSchedule(...)

File ~/venv/dev/lib/python3.10/site-packages/diffusers/pipelines/pipeline_utils.py:469, in load_sub_model(library_name, class_name, importable_classes, pipelines, is_pipeline_module, pipeline_class, torch_dtype, provider, sess_options, device_map, max_memory, offload_folder, offload_state_dict, model_variants, name, from_flax, variant, low_cpu_mem_usage, cached_folder, revision)
    465     if is_dummy_path and \"dummy\" in none_module:
    466         # call class_obj for nice error message of missing requirements
    467         class_obj()
--> 469     raise ValueError(
    470         f\"The component {class_obj} of {pipeline_class} cannot be loaded as it does not seem to have\"
    471         f\" any of the loading methods defined in {ALL_IMPORTABLE_CLASSES}.\"
    472     )
    474 load_method = getattr(class_obj, load_method_name)
    476 # add kwargs to loading method

ValueError: The component <class 'transformers_modules.model.define.TestModel'> of <class '__main__.MyPipe'> cannot be loaded as it does not seem to have any of the loading methods defined in {'ModelMixin': ['save_pretrained', 'from_pretrained'], 'SchedulerMixin': ['save_pretrained', 'from_pretrained'], 'DiffusionPipeline': ['save_pretrained', 'from_pretrained'], 'OnnxRuntimeModel': ['save_pretrained', 'from_pretrained'], 'PreTrainedTokenizer': ['save_pretrained', 'from_pretrained'], 'PreTrainedTokenizerFast': ['save_pretrained', 'from_pretrained'], 'PreTrainedModel': ['save_pretrained', 'from_pretrained'], 'FeatureExtractionMixin': ['save_pretrained', 'from_pretrained'], 'ProcessorMixin': ['save_pretrained', 'from_pretrained'], 'ImageProcessingMixin': ['save_pretrained', 'from_pretrained'], 'ORTModule': ['save_pretrained', 'from_pretrained']}."
}

The custom model is a subclass of PreTrainedModel. The ValueError is clearly lying.

[sayakpaul]

Would something like this work for you?

[ChenchaoZhao]

Right that's what I was saying the tranformers custom model and diffusers custom pipeline can work independently but not together even though custom model is a subclass of pretrained model -- a serializable object according to diffusers.

Actually the transformers custom model can write the definition files to pretrained dir automatically. Each py file is a module in the directory.

Ideally, it would be great the pipeline can serialize the custom model just like official models and load them back without any additional dependencies

[sayakpaul]

I showed you a simple way to use a custom model from transformers. I agree it’s not fully complete but this feature is also not a highly requested one. So, if this can be solved with a few user side tweaks, we will encourage it. Because a custom pipeline with all custom components is quite of an advanced use case and I think it should be reserved for power users.

Hope that makes sense.

[ChenchaoZhao]

Ok that makes sense. I do have a hack to make it work for some cases but sometimes it fails for mysterious reasons. I will dig into the code further. Thanks!

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.