null on exported json config results in custom model not being loaded

[ysig]

Hi,

I have created my own CustomImage pipeline using the save_pretrained method using a CustomEncoder imported from a local file. The saved general config file looks like this:

{
  "_class_name": "CustomImagePipeline",
  "_diffusers_version": "0.14.0",
  "customencoder": [
    null,
    "CustomEncoder"
  ],
    "scheduler": [
    "diffusers",
    "DDPMScheduler"
  ],
  "unet": [
    "diffusers",
    "UNet2DConditionModel"
  ],
  "vqvae": [
    "diffusers",
    "AutoencoderKL"
  ]
}

When saved the library corresponding to the encoder becomes None and it is then removed through load_module and thus I get the error:

File ~/miniconda3/envs/diffusers/lib/python3.10/site-packages/diffusers/pipelines/pipeline_utils.py:1038, in DiffusionPipeline.from_pretrained(cls, pretrained_model_name_or_path, **kwargs)
   1036 elif len(missing_modules) > 0:
   1037     passed_modules = set(list(init_kwargs.keys()) + list(passed_class_obj.keys())) - optional_kwargs
-> 1038     raise ValueError(
   1039         f"Pipeline {pipeline_class} expected {expected_modules}, but only {passed_modules} were passed."
   1040     )
   1042 # 8. Instantiate the pipeline
   1043 model = pipeline_class(**init_kwargs)

ValueError: Pipeline <class 'model.CustomImagePipeline'> expected {'scheduler', 'vqvae', 'unet', 'customencoder'}, but only {'scheduler', 'vqvae', 'unet'} were passed.

If I replace null with the name of the file from which the encoder comes from, let's say encoderfname then the following line breaks:

File ~/miniconda3/envs/diffusers/lib/python3.10/site-packages/diffusers/pipelines/pipeline_utils.py:995, in DiffusionPipeline.from_pretrained(cls, pretrained_model_name_or_path, **kwargs)
    993 # 6.2 Define all importable classes
    994 is_pipeline_module = hasattr(pipelines, library_name)
--> 995 importable_classes = ALL_IMPORTABLE_CLASSES if is_pipeline_module else LOADABLE_CLASSES[library_name]
    996 loaded_sub_model = None
    998 # 6.3 Use passed sub model or load class_name from library_name

KeyError: 'encoderfname'

Could you propose a workaround for this problem (that ofc wouldn't involve making a library out of my code)?
(Future proofing libraries is always a good idea especially when the use case is that simple.)

[leopoldmaillard]

Hello @ysig,

I've recently encountered the same issue. I've created my own Pipeline that similarly had a CustomModule implemented in a local file.

First thing I tried was to create a minimal pip installable package to wrap my CustomModule so it can be imported anywhere in the environment and changed the config_dict accordingly, e.g.

"custommodule": [
    mypackage.modules,
    "CustomModule"
  ]

Unfortunately, it didn't work.

However, a quick solution is to make your CustomModule inherit from ModelMixin and ConfigMixin so you can instantiate and call from_pretrained on all the pipeline's components individually, including CustomModule, before creating it.

I might go deeper into the diffusers.pipeline code and will let you know here if a find a more elegant workaround. I agree that support for local, custom pipelines would be a plus in the lib !

[ysig]

I have a workaround that is also elegant, but this is not the point. The point is for the huggingface team to make their code more modular/hackable.

[leopoldmaillard]

Would you mind sharing your solution ?

As far as I know, Diffusers is a community-driven library and I'm sure the 🤗 Team would be happy to review any suggestions/PRs on how to make it more flexible !

[ysig]

First you rewrite your init so that it can be initialized with what only diffusers can import (those that are not None).
Then in your pipeline you add a from_pretrained method, that:

1. first loads everything from super
2. then manually loads the custom modules
3. register them in your pipeline

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        """Helper method to load the module `name` from `library_name` and `class_name`"""
        pipeline = super().from_pretrained(pretrained_model_name_or_path, **kwargs)

        # Imagining it's not online it should come from a full path
        path = os.path.join(pretrained_model_name_or_path, 'customencoder')

        # helper to find files
        def find_first(path, ext):
            return [file for file in os.listdir(path) if file.endswith(ext)][0]

        # open the json under dir and open
        with open(os.path.join(path, find_first(path, '.json')), 'r') as f:
            config = json.load(f)
        
        # load model and weights
        model = CustomEncoder(**config)
        model.load_state_dict(torch.load(os.path.join(path, find_first(path, '.bin'))))
        model = model.to(model.device)

        # add to the archi
        pipeline.register_modules(customencoder=model)

        return pipeline

But this should be instead addressed in a design way by diffusers' core developers (could be with post-init functions that perform the above logic - dunno - but in any case, the lack of anticipated modularity/hackability during design is a very central shortcoming of diffusers, something that is a common place problem of everyone I know that is using it and as such needs to be systematically addressed in the future).

[sayakpaul]

Ccing @patrickvonplaten

[patrickvonplaten]

I need to see a fully reproducible codesnippet here, it's hard to help otherwise :-)

[leopoldmaillard]

Hello @patrickvonplaten,
The following snippet should reproduce the error mentioned by @ysig

import torch.nn as nn
from diffusers import ModelMixin, ConfigMixin, DiffusionPipeline
from diffusers import UNet2DConditionModel, DDIMScheduler

# define a custom encoder
class CustomEncoder(ModelMixin, ConfigMixin):
    def __init__(self, input_dim=256, hidden_dim=128, num_layers=4, num_heads=8):
        super().__init__()
        self.embedding = nn.Embedding(input_dim, hidden_dim)
        self.transformer_encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(hidden_dim, num_heads), num_layers
        )
    def forward(self, x):
        x = self.embedding(x)
        x = self.transformer_encoder(x)
        return x

# define a custom pipeline
class CustomPipeline(DiffusionPipeline):
    def __init__(
        self,
        unet: UNet2DConditionModel,
        encoder: CustomEncoder,
        scheduler: DDIMScheduler,
    ):
        super().__init__()
        self.register_modules(unet=unet, encoder=encoder, scheduler=scheduler)

        # define a custom __call__ etc.

def main():
    # save a custom pipe
    pipe = CustomPipeline(
        unet = UNet2DConditionModel(), 
        encoder = CustomEncoder(), 
        scheduler = DDIMScheduler(),
    )
    pipe.save_pretrained("/path/to/pipe")

    # load back the pipe
    pipe = CustomPipeline.from_pretrained("/path/to/pipe")

if __name__ == "__main__":
    main()

This would create the following model_index.json:

{
  "_class_name": "CustomPipeline",
  "_diffusers_version": "0.16.1",
  "encoder": [
    null,
    "CustomEncoder"
  ],
  "scheduler": [
    "diffusers",
    "DDIMScheduler"
  ],
  "unet": [
    "diffusers",
    "UNet2DConditionModel"
  ]
}

However, loading back the pipeline with from_pretrained() raises the following error:

File "~/miniconda3/envs/diffusers/lib/python3.11/site-packages/diffusers/pipelines/pipeline_utils.py", line 1069, in from_pretrained
    raise ValueError(
ValueError: Pipeline <class '__main__.CustomPipeline'> expected {'scheduler', 'unet', 'encoder'}, but only {'scheduler', 'unet'} were passed.

As mentioned by @ysig, solutions involving replacing the null value in the config dict do not seem to work.

Hope it helps !

[patrickvonplaten]

Hey @leopoldmaillard,

Thanks for the snippet. This PR should solve it: #3504

[leopoldmaillard]

Hello @patrickvonplaten, thank you for the PR, it works perfectly now !

[github-actions]

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.