pipeline_utils.py 6.77 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"],
Patrick von Platen's avatar
improve  
Patrick von Platen committed
37
        "GaussianDDPMScheduler": ["save_config", "from_config"],
38
        "ClassifierFreeGuidanceScheduler": ["save_config", "from_config"],
39
        "GlideDDIMScheduler": ["save_config", "from_config"],
Patrick von Platen's avatar
Patrick von Platen committed
40
41
    },
    "transformers": {
anton-l's avatar
anton-l committed
42
        "PreTrainedTokenizer": ["save_pretrained", "from_pretrained"],
Patrick von Platen's avatar
Patrick von Platen committed
43
44
45
    },
}

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

Patrick von Platen's avatar
Patrick von Platen committed
50

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

    config_name = "model_index.json"

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

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

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

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

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

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

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

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

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

Patrick von Platen's avatar
Patrick von Platen committed
92
93
94
95
96
97
98
99
100
101
102
103
104
105
            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):
106
        r"""
Patrick von Platen's avatar
Patrick von Platen committed
107
        Add docstrings
108
109
110
111
112
113
114
        """
        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
115
        # use snapshot download here to get it working from from_pretrained
Patrick von Platen's avatar
Patrick von Platen committed
116
        if not os.path.isdir(pretrained_model_name_or_path):
117
118
119
120
121
122
123
124
            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
125
126
        else:
            cached_folder = pretrained_model_name_or_path
127

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

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

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

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

        init_kwargs = {}

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

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

            load_method = getattr(class_obj, load_method_name)

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

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

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