pipeline_utils.py 5.83 KB
Newer Older
Patrick von Platen's avatar
Patrick von Platen committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# coding=utf-8
# Copyright 2022 The HuggingFace Inc. team.
# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.

Patrick von Platen's avatar
improve  
Patrick von Platen committed
17
import importlib
Patrick von Platen's avatar
Patrick von Platen committed
18
19
import os
from typing import Optional, Union
anton-l's avatar
Style  
anton-l committed
20

Patrick von Platen's avatar
up  
Patrick von Platen committed
21
from huggingface_hub import snapshot_download
Patrick von Platen's avatar
Patrick von Platen committed
22
23
24
25

# CHANGE to diffusers.utils
from transformers.utils import logging

Patrick von Platen's avatar
Patrick von Platen committed
26
from .configuration_utils import ConfigMixin
patil-suraj's avatar
patil-suraj committed
27
from .dynamic_modules_utils import get_class_from_dynamic_module
Patrick von Platen's avatar
improve  
Patrick von Platen committed
28

Patrick von Platen's avatar
Patrick von Platen committed
29
30
31
32
33
34
35
36
37

INDEX_FILE = "diffusion_model.pt"


logger = logging.get_logger(__name__)


LOADABLE_CLASSES = {
    "diffusers": {
Patrick von Platen's avatar
Patrick von Platen committed
38
        "ModelMixin": ["save_pretrained", "from_pretrained"],
39
        "CLIPTextModel": ["save_pretrained", "from_pretrained"],  # TODO (Anton): move to transformers
Patrick von Platen's avatar
improve  
Patrick von Platen committed
40
        "GaussianDDPMScheduler": ["save_config", "from_config"],
41
        "ClassifierFreeGuidanceScheduler": ["save_config", "from_config"],
Patrick von Platen's avatar
Patrick von Platen committed
42
43
    },
    "transformers": {
44
        "GPT2Tokenizer": ["save_pretrained", "from_pretrained"],
Patrick von Platen's avatar
Patrick von Platen committed
45
46
47
48
    },
}


Patrick von Platen's avatar
Patrick von Platen committed
49
class DiffusionPipeline(ConfigMixin):
Patrick von Platen's avatar
Patrick von Platen committed
50
51
52

    config_name = "model_index.json"

Patrick von Platen's avatar
up  
Patrick von Platen committed
53
    def register_modules(self, **kwargs):
Patrick von Platen's avatar
Patrick von Platen committed
54
55
56
        for name, module in kwargs.items():
            # retrive library
            library = module.__module__.split(".")[0]
patil-suraj's avatar
patil-suraj committed
57
58
59
60
            # if library is not in LOADABLE_CLASSES, then it is a custom module
            if library not in LOADABLE_CLASSES:
                library = module.__module__.split(".")[-1]

Patrick von Platen's avatar
Patrick von Platen committed
61
62
63
            # retrive class_name
            class_name = module.__class__.__name__

64
65
            register_dict = {name: (library, class_name)}

Patrick von Platen's avatar
Patrick von Platen committed
66
            # save model index config
67
            self.register(**register_dict)
Patrick von Platen's avatar
Patrick von Platen committed
68
69
70

            # set models
            setattr(self, name, module)
71

anton-l's avatar
Style  
anton-l committed
72
        register_dict = {"_module": self.__module__.split(".")[-1] + ".py"}
73
        self.register(**register_dict)
Patrick von Platen's avatar
Patrick von Platen committed
74
75
76
77
78
79

    def save_pretrained(self, save_directory: Union[str, os.PathLike]):
        self.save_config(save_directory)

        model_index_dict = self._dict_to_save
        model_index_dict.pop("_class_name")
80
        model_index_dict.pop("_module")
Patrick von Platen's avatar
Patrick von Platen committed
81
82
83
84

        for name, (library_name, class_name) in self._dict_to_save.items():
            importable_classes = LOADABLE_CLASSES[library_name]

85
86
87
88
            # TODO: Suraj
            if library_name == self.__module__:
                library_name = self

Patrick von Platen's avatar
Patrick von Platen committed
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
            library = importlib.import_module(library_name)
            class_obj = getattr(library, class_name)
            class_candidates = {c: getattr(library, c) for c in importable_classes.keys()}

            save_method_name = None
            for class_name, class_candidate in class_candidates.items():
                if issubclass(class_obj, class_candidate):
                    save_method_name = importable_classes[class_name][0]

            save_method = getattr(getattr(self, name), save_method_name)
            save_method(os.path.join(save_directory, name))

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
        # use snapshot download here to get it working from from_pretrained
Patrick von Platen's avatar
Patrick von Platen committed
104
105
106
107
        if not os.path.isdir(pretrained_model_name_or_path):
            cached_folder = snapshot_download(pretrained_model_name_or_path)
        else:
            cached_folder = pretrained_model_name_or_path
108

patil-suraj's avatar
patil-suraj committed
109
        config_dict = cls.get_config_dict(cached_folder)
110

Patrick von Platen's avatar
fix  
Patrick von Platen committed
111
        module_candidate = config_dict["_module"]
patil-suraj's avatar
patil-suraj committed
112
        module_candidate_name = module_candidate.replace(".py", "")
Patrick von Platen's avatar
fix  
Patrick von Platen committed
113

114
115
        # if we load from explicit class, let's use it
        if cls != DiffusionPipeline:
116
117
            pipeline_class = cls
        else:
118
119
            # else we need to load the correct module from the Hub
            class_name_ = config_dict["_class_name"]
Patrick von Platen's avatar
fix  
Patrick von Platen committed
120
            module = module_candidate
121
            pipeline_class = get_class_from_dynamic_module(cached_folder, module, class_name_, cached_folder)
122

123
        init_dict, _ = pipeline_class.extract_init_dict(config_dict, **kwargs)
Patrick von Platen's avatar
Patrick von Platen committed
124
125
126

        init_kwargs = {}

patil-suraj's avatar
patil-suraj committed
127
        for name, (library_name, class_name) in init_dict.items():
patil-suraj's avatar
patil-suraj committed
128
129
130
131
132
133
134
135
            
            # if the model is not in diffusers or transformers, we need to load it from the hub
            # assumes that it's a subclass of ModelMixin
            if library_name == module_candidate_name:
                class_obj = get_class_from_dynamic_module(cached_folder, module, class_name, cached_folder)
                load_method_name = "from_pretrained"
            else:
                importable_classes = LOADABLE_CLASSES[library_name]
Patrick von Platen's avatar
Patrick von Platen committed
136

patil-suraj's avatar
patil-suraj committed
137
138
139
                library = importlib.import_module(library_name)
                class_obj = getattr(library, class_name)
                class_candidates = {c: getattr(library, c) for c in importable_classes.keys()}
140

patil-suraj's avatar
patil-suraj committed
141
142
143
144
                load_method_name = None
                for class_name, class_candidate in class_candidates.items():
                    if issubclass(class_obj, class_candidate):
                        load_method_name = importable_classes[class_name][1]
Patrick von Platen's avatar
Patrick von Platen committed
145
146
147

            load_method = getattr(class_obj, load_method_name)

Patrick von Platen's avatar
Patrick von Platen committed
148
            if os.path.isdir(os.path.join(cached_folder, name)):
149
150
151
                loaded_sub_model = load_method(os.path.join(cached_folder, name))
            else:
                loaded_sub_model = load_method(cached_folder)
Patrick von Platen's avatar
Patrick von Platen committed
152

153
            init_kwargs[name] = loaded_sub_model  # UNet(...), # DiffusionSchedule(...)
Patrick von Platen's avatar
Patrick von Platen committed
154

155
        model = pipeline_class(**init_kwargs)
Patrick von Platen's avatar
Patrick von Platen committed
156
        return model