pipeline_utils.py 6.87 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

Patrick von Platen's avatar
Patrick von Platen committed
23
from .configuration_utils import ConfigMixin
patil-suraj's avatar
patil-suraj committed
24
from .dynamic_modules_utils import get_class_from_dynamic_module
Patrick von Platen's avatar
Patrick von Platen committed
25
from .utils import DIFFUSERS_CACHE, logging
Patrick von Platen's avatar
improve  
Patrick von Platen committed
26

Patrick von Platen's avatar
Patrick von Platen committed
27
28
29
30
31
32
33
34
35

INDEX_FILE = "diffusion_model.pt"


logger = logging.get_logger(__name__)


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

47
48
49
50
ALL_IMPORTABLE_CLASSES = {}
for library in LOADABLE_CLASSES:
    ALL_IMPORTABLE_CLASSES.update(LOADABLE_CLASSES[library])

Patrick von Platen's avatar
Patrick von Platen committed
51

Patrick von Platen's avatar
Patrick von Platen committed
52
class DiffusionPipeline(ConfigMixin):
Patrick von Platen's avatar
Patrick von Platen committed
53
54
55

    config_name = "model_index.json"

Patrick von Platen's avatar
up  
Patrick von Platen committed
56
    def register_modules(self, **kwargs):
Patrick von Platen's avatar
Patrick von Platen committed
57
58
59
        for name, module in kwargs.items():
            # retrive library
            library = module.__module__.split(".")[0]
patil-suraj's avatar
patil-suraj committed
60
61
62
63
            # 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
64
65
66
            # retrive class_name
            class_name = module.__class__.__name__

67
68
            register_dict = {name: (library, class_name)}

Patrick von Platen's avatar
Patrick von Platen committed
69
            # save model index config
70
            self.register(**register_dict)
Patrick von Platen's avatar
Patrick von Platen committed
71
72
73

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

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

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

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

86
        for name, (library_name, class_name) in model_index_dict.items():
Patrick von Platen's avatar
Patrick von Platen committed
87
88
            importable_classes = LOADABLE_CLASSES[library_name]

89
90
91
92
            # TODO: Suraj
            if library_name == self.__module__:
                library_name = self

Patrick von Platen's avatar
Patrick von Platen committed
93
94
95
96
97
98
99
100
101
102
103
104
105
106
            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):
107
        r"""
Patrick von Platen's avatar
Patrick von Platen committed
108
        Add docstrings
109
110
111
112
113
114
115
        """
        cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)
        resume_download = kwargs.pop("resume_download", False)
        proxies = kwargs.pop("proxies", None)
        local_files_only = kwargs.pop("local_files_only", False)
        use_auth_token = kwargs.pop("use_auth_token", None)

Patrick von Platen's avatar
Patrick von Platen committed
116
        # use snapshot download here to get it working from from_pretrained
Patrick von Platen's avatar
Patrick von Platen committed
117
        if not os.path.isdir(pretrained_model_name_or_path):
118
119
120
121
122
123
124
125
            cached_folder = snapshot_download(
                pretrained_model_name_or_path,
                cache_dir=cache_dir,
                resume_download=resume_download,
                proxies=proxies,
                local_files_only=local_files_only,
                use_auth_token=use_auth_token,
            )
Patrick von Platen's avatar
Patrick von Platen committed
126
127
        else:
            cached_folder = pretrained_model_name_or_path
128

patil-suraj's avatar
patil-suraj committed
129
        config_dict = cls.get_config_dict(cached_folder)
130

patil-suraj's avatar
patil-suraj committed
131
132
        module = config_dict["_module"]
        class_name_ = config_dict["_class_name"]
Patrick von Platen's avatar
fix  
Patrick von Platen committed
133
        module_candidate = config_dict["_module"]
patil-suraj's avatar
patil-suraj committed
134
        module_candidate_name = module_candidate.replace(".py", "")
Patrick von Platen's avatar
fix  
Patrick von Platen committed
135

136
137
        # if we load from explicit class, let's use it
        if cls != DiffusionPipeline:
138
139
            pipeline_class = cls
        else:
140
141
            # 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
142
            module = module_candidate
143
            pipeline_class = get_class_from_dynamic_module(cached_folder, module, class_name_, cached_folder)
144

145
        init_dict, _ = pipeline_class.extract_init_dict(config_dict, **kwargs)
Patrick von Platen's avatar
Patrick von Platen committed
146
147
148

        init_kwargs = {}

patil-suraj's avatar
patil-suraj committed
149
        for name, (library_name, class_name) in init_dict.items():
patil-suraj's avatar
patil-suraj committed
150
151
152
153
            # 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)
154
                # since it's not from a library, we need to check class candidates for all importable classes
155
156
                importable_classes = ALL_IMPORTABLE_CLASSES
                class_candidates = {c: class_obj for c in ALL_IMPORTABLE_CLASSES.keys()}
patil-suraj's avatar
patil-suraj committed
157
158
159
            else:
                library = importlib.import_module(library_name)
                class_obj = getattr(library, class_name)
160
                importable_classes = LOADABLE_CLASSES[library_name]
patil-suraj's avatar
patil-suraj committed
161
                class_candidates = {c: getattr(library, c) for c in importable_classes.keys()}
162

163
164
165
166
            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
167
168
169

            load_method = getattr(class_obj, load_method_name)

Patrick von Platen's avatar
Patrick von Platen committed
170
            if os.path.isdir(os.path.join(cached_folder, name)):
171
172
173
                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
174

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

177
        model = pipeline_class(**init_kwargs)
Patrick von Platen's avatar
Patrick von Platen committed
178
        return model