pipeline_utils.py 5.28 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
Patrick von Platen's avatar
up  
Patrick von Platen committed
20
from huggingface_hub import snapshot_download
Patrick von Platen's avatar
Patrick von Platen committed
21
22
23
24

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

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

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

INDEX_FILE = "diffusion_model.pt"


logger = logging.get_logger(__name__)


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


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

    config_name = "model_index.json"

Patrick von Platen's avatar
up  
Patrick von Platen committed
52
    def register_modules(self, **kwargs):
Patrick von Platen's avatar
Patrick von Platen committed
53
54
55
56
57
58
        for name, module in kwargs.items():
            # retrive library
            library = module.__module__.split(".")[0]
            # retrive class_name
            class_name = module.__class__.__name__

59
60
            register_dict = {name: (library, class_name)}

Patrick von Platen's avatar
Patrick von Platen committed
61
            # save model index config
62
            self.register(**register_dict)
Patrick von Platen's avatar
Patrick von Platen committed
63
64
65

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

67
68
        register_dict = {"_module" : self.__module__.split(".")[-1] + ".py"}
        self.register(**register_dict)
Patrick von Platen's avatar
Patrick von Platen committed
69
70
71
72
73
74

    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")
75
        model_index_dict.pop("_module")
Patrick von Platen's avatar
Patrick von Platen committed
76
77
78
79

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

80
81
82
83
            # TODO: Suraj
            if library_name == self.__module__:
                library_name = self

Patrick von Platen's avatar
Patrick von Platen committed
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
            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
99
100
101
102
        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
103

patil-suraj's avatar
patil-suraj committed
104
        config_dict = cls.get_config_dict(cached_folder)
105

Patrick von Platen's avatar
fix  
Patrick von Platen committed
106
107
        module_candidate = config_dict["_module"]

108
109
        # if we load from explicit class, let's use it
        if cls != DiffusionPipeline:
110
111
            pipeline_class = cls
        else:
112
113
            # 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
114
            module = module_candidate
115
            pipeline_class = get_class_from_dynamic_module(cached_folder, module, class_name_, cached_folder)
116

117
        init_dict, _ = pipeline_class.extract_init_dict(config_dict, **kwargs)
Patrick von Platen's avatar
Patrick von Platen committed
118
119
120

        init_kwargs = {}

patil-suraj's avatar
patil-suraj committed
121
        for name, (library_name, class_name) in init_dict.items():
Patrick von Platen's avatar
Patrick von Platen committed
122
123
            importable_classes = LOADABLE_CLASSES[library_name]

Patrick von Platen's avatar
fix  
Patrick von Platen committed
124
            if library_name == module_candidate:
125
                # TODO(Suraj)
126
                # for vq
127
128
                pass

Patrick von Platen's avatar
Patrick von Platen committed
129
130
131
132
133
134
135
136
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()}

            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]

            load_method = getattr(class_obj, load_method_name)

Patrick von Platen's avatar
Patrick von Platen committed
140
            if os.path.isdir(os.path.join(cached_folder, name)):
141
142
143
                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
144

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

147
        model = pipeline_class(**init_kwargs)
Patrick von Platen's avatar
Patrick von Platen committed
148
        return model