pipeline_utils.py 24.6 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
18
import inspect
Patrick von Platen's avatar
Patrick von Platen committed
19
import os
20
21
from dataclasses import dataclass
from typing import List, Optional, Union
anton-l's avatar
Style  
anton-l committed
22

23
import numpy as np
Pedro Cuenca's avatar
Pedro Cuenca committed
24
25
import torch

26
import diffusers
27
import PIL
Patrick von Platen's avatar
up  
Patrick von Platen committed
28
from huggingface_hub import snapshot_download
29
from PIL import Image
hysts's avatar
hysts committed
30
from tqdm.auto import tqdm
Patrick von Platen's avatar
Patrick von Platen committed
31

32
from . import __version__
Patrick von Platen's avatar
Patrick von Platen committed
33
from .configuration_utils import ConfigMixin
Patrick von Platen's avatar
Patrick von Platen committed
34
from .dynamic_modules_utils import get_class_from_dynamic_module
35
from .schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME
36
37
38
39
40
41
42
43
44
45
46
47
48
from .utils import (
    CONFIG_NAME,
    DIFFUSERS_CACHE,
    ONNX_WEIGHTS_NAME,
    WEIGHTS_NAME,
    BaseOutput,
    is_transformers_available,
    logging,
)


if is_transformers_available():
    from transformers import PreTrainedModel
Patrick von Platen's avatar
improve  
Patrick von Platen committed
49

Patrick von Platen's avatar
Patrick von Platen committed
50

Patrick von Platen's avatar
Patrick von Platen committed
51
INDEX_FILE = "diffusion_pytorch_model.bin"
Patrick von Platen's avatar
Patrick von Platen committed
52
CUSTOM_PIPELINE_FILE_NAME = "pipeline.py"
53
DUMMY_MODULES_FOLDER = "diffusers.utils"
Patrick von Platen's avatar
Patrick von Platen committed
54
55
56
57
58
59
60


logger = logging.get_logger(__name__)


LOADABLE_CLASSES = {
    "diffusers": {
Patrick von Platen's avatar
Patrick von Platen committed
61
        "ModelMixin": ["save_pretrained", "from_pretrained"],
Patrick von Platen's avatar
Patrick von Platen committed
62
        "SchedulerMixin": ["save_config", "from_config"],
Patrick von Platen's avatar
Patrick von Platen committed
63
        "DiffusionPipeline": ["save_pretrained", "from_pretrained"],
64
        "OnnxRuntimeModel": ["save_pretrained", "from_pretrained"],
Patrick von Platen's avatar
Patrick von Platen committed
65
66
    },
    "transformers": {
anton-l's avatar
anton-l committed
67
        "PreTrainedTokenizer": ["save_pretrained", "from_pretrained"],
68
        "PreTrainedTokenizerFast": ["save_pretrained", "from_pretrained"],
anton-l's avatar
anton-l committed
69
        "PreTrainedModel": ["save_pretrained", "from_pretrained"],
Suraj Patil's avatar
Suraj Patil committed
70
        "FeatureExtractionMixin": ["save_pretrained", "from_pretrained"],
Patrick von Platen's avatar
Patrick von Platen committed
71
72
73
    },
}

74
75
76
77
ALL_IMPORTABLE_CLASSES = {}
for library in LOADABLE_CLASSES:
    ALL_IMPORTABLE_CLASSES.update(LOADABLE_CLASSES[library])

Patrick von Platen's avatar
Patrick von Platen committed
78

79
80
81
82
83
84
85
86
87
88
89
90
91
92
@dataclass
class ImagePipelineOutput(BaseOutput):
    """
    Output class for image pipelines.

    Args:
        images (`List[PIL.Image.Image]` or `np.ndarray`)
            List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
            num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
    """

    images: Union[List[PIL.Image.Image], np.ndarray]


Patrick von Platen's avatar
Patrick von Platen committed
93
class DiffusionPipeline(ConfigMixin):
94
95
96
97
98
99
100
101
102
103
104
105
    r"""
    Base class for all models.

    [`DiffusionPipeline`] takes care of storing all components (models, schedulers, processors) for diffusion pipelines
    and handles methods for loading, downloading and saving models as well as a few methods common to all pipelines to:

        - move all PyTorch modules to the device of your choice
        - enabling/disabling the progress bar for the denoising iteration

    Class attributes:

        - **config_name** ([`str`]) -- name of the config file that will store the class and module names of all
106
          components of the diffusion pipeline.
107
    """
Patrick von Platen's avatar
Patrick von Platen committed
108
109
    config_name = "model_index.json"

Patrick von Platen's avatar
up  
Patrick von Platen committed
110
    def register_modules(self, **kwargs):
111
112
        # import it here to avoid circular import
        from diffusers import pipelines
113

Patrick von Platen's avatar
Patrick von Platen committed
114
        for name, module in kwargs.items():
115
            # retrieve library
Patrick von Platen's avatar
Patrick von Platen committed
116
            library = module.__module__.split(".")[0]
117

118
119
            # check if the module is a pipeline module
            pipeline_dir = module.__module__.split(".")[-2]
Suraj Patil's avatar
Suraj Patil committed
120
121
            path = module.__module__.split(".")
            is_pipeline_module = pipeline_dir in path and hasattr(pipelines, pipeline_dir)
122

123
124
            # if library is not in LOADABLE_CLASSES, then it is a custom module.
            # Or if it's a pipeline module, then the module is inside the pipeline
125
            # folder so we set the library to module name.
126
            if library not in LOADABLE_CLASSES or is_pipeline_module:
127
                library = pipeline_dir
patil-suraj's avatar
patil-suraj committed
128

129
            # retrieve class_name
Patrick von Platen's avatar
Patrick von Platen committed
130
131
            class_name = module.__class__.__name__

132
133
            register_dict = {name: (library, class_name)}

Patrick von Platen's avatar
Patrick von Platen committed
134
            # save model index config
135
            self.register_to_config(**register_dict)
Patrick von Platen's avatar
Patrick von Platen committed
136
137
138

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

Patrick von Platen's avatar
Patrick von Platen committed
140
    def save_pretrained(self, save_directory: Union[str, os.PathLike]):
141
142
143
144
145
146
147
148
149
        """
        Save all variables of the pipeline that can be saved and loaded as well as the pipelines configuration file to
        a directory. A pipeline variable can be saved and loaded if its class implements both a save and loading
        method. The pipeline can easily be re-loaded using the `[`~DiffusionPipeline.from_pretrained`]` class method.

        Arguments:
            save_directory (`str` or `os.PathLike`):
                Directory to which to save. Will be created if it doesn't exist.
        """
Patrick von Platen's avatar
Patrick von Platen committed
150
151
        self.save_config(save_directory)

Patrick von Platen's avatar
Patrick von Platen committed
152
        model_index_dict = dict(self.config)
Patrick von Platen's avatar
Patrick von Platen committed
153
        model_index_dict.pop("_class_name")
154
        model_index_dict.pop("_diffusers_version")
155
        model_index_dict.pop("_module", None)
Patrick von Platen's avatar
Patrick von Platen committed
156

anton-l's avatar
anton-l committed
157
158
159
        for pipeline_component_name in model_index_dict.keys():
            sub_model = getattr(self, pipeline_component_name)
            model_cls = sub_model.__class__
Patrick von Platen's avatar
Patrick von Platen committed
160
161

            save_method_name = None
anton-l's avatar
anton-l committed
162
163
164
165
166
167
168
169
170
171
172
173
174
175
            # search for the model's base class in LOADABLE_CLASSES
            for library_name, library_classes in LOADABLE_CLASSES.items():
                library = importlib.import_module(library_name)
                for base_class, save_load_methods in library_classes.items():
                    class_candidate = getattr(library, base_class)
                    if issubclass(model_cls, class_candidate):
                        # if we found a suitable base class in LOADABLE_CLASSES then grab its save method
                        save_method_name = save_load_methods[0]
                        break
                if save_method_name is not None:
                    break

            save_method = getattr(sub_model, save_method_name)
            save_method(os.path.join(save_directory, pipeline_component_name))
Patrick von Platen's avatar
Patrick von Platen committed
176

Pedro Cuenca's avatar
Pedro Cuenca committed
177
178
179
180
181
182
183
184
    def to(self, torch_device: Optional[Union[str, torch.device]] = None):
        if torch_device is None:
            return self

        module_names, _ = self.extract_init_dict(dict(self.config))
        for name in module_names.keys():
            module = getattr(self, name)
            if isinstance(module, torch.nn.Module):
185
                if module.dtype == torch.float16 and str(torch_device) in ["cpu", "mps"]:
186
187
188
189
190
191
                    logger.warning(
                        "Pipelines loaded with `torch_dtype=torch.float16` cannot run with `cpu` or `mps` device. It"
                        " is not recommended to move them to `cpu` or `mps` as running them will fail. Please make"
                        " sure to use a `cuda` device to run the pipeline in inference. due to the lack of support for"
                        " `float16` operations on those devices in PyTorch. Please remove the"
                        " `torch_dtype=torch.float16` argument, or use a `cuda` device to run inference."
192
                    )
Pedro Cuenca's avatar
Pedro Cuenca committed
193
194
195
196
197
                module.to(torch_device)
        return self

    @property
    def device(self) -> torch.device:
198
199
200
201
        r"""
        Returns:
            `torch.device`: The torch device on which the pipeline is located.
        """
Pedro Cuenca's avatar
Pedro Cuenca committed
202
203
204
205
206
207
208
        module_names, _ = self.extract_init_dict(dict(self.config))
        for name in module_names.keys():
            module = getattr(self, name)
            if isinstance(module, torch.nn.Module):
                return module.device
        return torch.device("cpu")

Patrick von Platen's avatar
Patrick von Platen committed
209
210
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
211
        r"""
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
        Instantiate a PyTorch diffusion pipeline from pre-trained pipeline weights.

        The pipeline is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated).

        The warning *Weights from XXX not initialized from pretrained model* means that the weights of XXX do not come
        pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning
        task.

        The warning *Weights from XXX not used in YYY* means that the layer XXX is not used by YYY, therefore those
        weights are discarded.

        Parameters:
            pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
                Can be either:

                    - A string, the *repo id* of a pretrained pipeline hosted inside a model repo on
                      https://huggingface.co/ Valid repo ids have to be located under a user or organization name, like
                      `CompVis/ldm-text2im-large-256`.
                    - A path to a *directory* containing pipeline weights saved using
                      [`~DiffusionPipeline.save_pretrained`], e.g., `./my_pipeline_directory/`.
            torch_dtype (`str` or `torch.dtype`, *optional*):
                Override the default `torch.dtype` and load the model under this dtype. If `"auto"` is passed the dtype
                will be automatically derived from the model's weights.
Patrick von Platen's avatar
Patrick von Platen committed
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
            custom_pipeline (`str`, *optional*):

                <Tip warning={true}>

                    This is an experimental feature and is likely to change in the future.

                </Tip>

                Can be either:

                    - A string, the *repo id* of a custom pipeline hosted inside a model repo on
                      https://huggingface.co/. Valid repo ids have to be located under a user or organization name,
                      like `hf-internal-testing/diffusers-dummy-pipeline`.

                        <Tip>

                         It is required that the model repo has a file, called `pipeline.py` that defines the custom
                         pipeline.

                        </Tip>

                    - A string, the *file name* of a community pipeline hosted on GitHub under
                      https://github.com/huggingface/diffusers/tree/main/examples/community. Valid file names have to
                      match exactly the file name without `.py` located under the above link, *e.g.*
                      `clip_guided_stable_diffusion`.

                        <Tip>

                         Community pipelines are always loaded from the current `main` branch of GitHub.

                        </Tip>

                    - A path to a *directory* containing a custom pipeline, e.g., `./my_pipeline_directory/`.

                        <Tip>

                         It is required that the directory has a file, called `pipeline.py` that defines the custom
                         pipeline.

                        </Tip>

                For more information on how to load and create custom pipelines, please have a look at [Loading and
                Creating Custom
                Pipelines](https://huggingface.co/docs/diffusers/main/en/using-diffusers/custom_pipelines)

            torch_dtype (`str` or `torch.dtype`, *optional*):
281
282
283
284
285
286
287
288
289
290
            force_download (`bool`, *optional*, defaults to `False`):
                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
                cached versions if they exist.
            resume_download (`bool`, *optional*, defaults to `False`):
                Whether or not to delete incompletely received files. Will attempt to resume the download if such a
                file exists.
            proxies (`Dict[str, str]`, *optional*):
                A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
            output_loading_info(`bool`, *optional*, defaults to `False`):
291
                Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
            local_files_only(`bool`, *optional*, defaults to `False`):
                Whether or not to only look at local files (i.e., do not try to download the model).
            use_auth_token (`str` or *bool*, *optional*):
                The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
                when running `huggingface-cli login` (stored in `~/.huggingface`).
            revision (`str`, *optional*, defaults to `"main"`):
                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
                git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
                identifier allowed by git.
            mirror (`str`, *optional*):
                Mirror source to accelerate downloads in China. If you are from China and have an accessibility
                problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety.
                Please refer to the mirror site for more information. specify the folder name here.

            kwargs (remaining dictionary of keyword arguments, *optional*):
                Can be used to overwrite load - and saveable variables - *i.e.* the pipeline components - of the
308
309
                specific pipeline class. The overwritten components are then directly passed to the pipelines
                `__init__` method. See example below for more information.
310
311
312

        <Tip>

313
314
         It is required to be logged in (`huggingface-cli login`) when you want to use private or [gated
         models](https://huggingface.co/docs/hub/models-gated#gated-models), *e.g.* `"CompVis/stable-diffusion-v1-4"`
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335

        </Tip>

        <Tip>

        Activate the special ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use
        this method in a firewalled environment.

        </Tip>

        Examples:

        ```py
        >>> from diffusers import DiffusionPipeline

        >>> # Download pipeline from huggingface.co and cache.
        >>> pipeline = DiffusionPipeline.from_pretrained("CompVis/ldm-text2im-large-256")

        >>> # Download pipeline that requires an authorization token
        >>> # For more information on access tokens, please refer to this section
        >>> # of the documentation](https://huggingface.co/docs/hub/security-tokens)
336
        >>> pipeline = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
337
338
339
340
341

        >>> # Download pipeline, but overwrite scheduler
        >>> from diffusers import LMSDiscreteScheduler

        >>> scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear")
342
        >>> pipeline = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", scheduler=scheduler)
343
        ```
344
345
346
347
348
349
        """
        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)
350
        revision = kwargs.pop("revision", None)
351
        torch_dtype = kwargs.pop("torch_dtype", None)
Patrick von Platen's avatar
Patrick von Platen committed
352
        custom_pipeline = kwargs.pop("custom_pipeline", None)
353
        provider = kwargs.pop("provider", None)
354
        sess_options = kwargs.pop("sess_options", None)
355
        device_map = kwargs.pop("device_map", None)
Patrick von Platen's avatar
Patrick von Platen committed
356

patil-suraj's avatar
patil-suraj committed
357
        # 1. Download the checkpoints and configs
Patrick von Platen's avatar
Patrick von Platen committed
358
        # use snapshot download here to get it working from from_pretrained
Patrick von Platen's avatar
Patrick von Platen committed
359
        if not os.path.isdir(pretrained_model_name_or_path):
360
361
362
363
364
365
366
367
368
369
370
371
372
373
            config_dict = cls.get_config_dict(
                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,
                revision=revision,
            )
            # make sure we only download sub-folders and `diffusers` filenames
            folder_names = [k for k in config_dict.keys() if not k.startswith("_")]
            allow_patterns = [os.path.join(k, "*") for k in folder_names]
            allow_patterns += [WEIGHTS_NAME, SCHEDULER_CONFIG_NAME, CONFIG_NAME, ONNX_WEIGHTS_NAME, cls.config_name]

Patrick von Platen's avatar
Patrick von Platen committed
374
375
376
            if custom_pipeline is not None:
                allow_patterns += [CUSTOM_PIPELINE_FILE_NAME]

377
378
379
380
381
            requested_pipeline_class = config_dict.get("_class_name", cls.__name__)
            user_agent = {"diffusers": __version__, "pipeline_class": requested_pipeline_class}
            if custom_pipeline is not None:
                user_agent["custom_pipeline"] = custom_pipeline

382
            # download all allow_patterns
383
384
385
386
387
388
389
            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,
390
                revision=revision,
391
                allow_patterns=allow_patterns,
392
                user_agent=user_agent,
393
            )
Patrick von Platen's avatar
Patrick von Platen committed
394
395
        else:
            cached_folder = pretrained_model_name_or_path
396

patil-suraj's avatar
patil-suraj committed
397
        config_dict = cls.get_config_dict(cached_folder)
398

Patrick von Platen's avatar
Patrick von Platen committed
399
        # 2. Load the pipeline class, if using custom module then load it from the hub
400
        # if we load from explicit class, let's use it
Patrick von Platen's avatar
Patrick von Platen committed
401
402
403
404
405
        if custom_pipeline is not None:
            pipeline_class = get_class_from_dynamic_module(
                custom_pipeline, module_file=CUSTOM_PIPELINE_FILE_NAME, cache_dir=custom_pipeline
            )
        elif cls != DiffusionPipeline:
406
407
            pipeline_class = cls
        else:
Patrick von Platen's avatar
Patrick von Platen committed
408
409
410
            diffusers_module = importlib.import_module(cls.__module__.split(".")[0])
            pipeline_class = getattr(diffusers_module, config_dict["_class_name"])

411
412
413
        # some modules can be passed directly to the init
        # in this case they are already instantiated in `kwargs`
        # extract them here
Patrick von Platen's avatar
Patrick von Platen committed
414
        expected_modules = set(inspect.signature(pipeline_class.__init__).parameters.keys()) - set(["self"])
415
416
        passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs}

417
        init_dict, _ = pipeline_class.extract_init_dict(config_dict, **kwargs)
Patrick von Platen's avatar
Patrick von Platen committed
418
419

        init_kwargs = {}
420

421
422
        # import it here to avoid circular import
        from diffusers import pipelines
423

Patrick von Platen's avatar
Patrick von Platen committed
424
        # 3. Load each module in the pipeline
patil-suraj's avatar
patil-suraj committed
425
        for name, (library_name, class_name) in init_dict.items():
426
427
428
429
            # 3.1 - now that JAX/Flax is an official framework of the library, we might load from Flax names
            if class_name.startswith("Flax"):
                class_name = class_name[4:]

430
            is_pipeline_module = hasattr(pipelines, library_name)
431
432
            loaded_sub_model = None

433
            # if the model is in a pipeline module, then we load it from the pipeline
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
            if name in passed_class_obj:
                # 1. check that passed_class_obj has correct parent class
                if not is_pipeline_module:
                    library = importlib.import_module(library_name)
                    class_obj = getattr(library, class_name)
                    importable_classes = LOADABLE_CLASSES[library_name]
                    class_candidates = {c: getattr(library, c) for c in importable_classes.keys()}

                    expected_class_obj = None
                    for class_name, class_candidate in class_candidates.items():
                        if issubclass(class_obj, class_candidate):
                            expected_class_obj = class_candidate

                    if not issubclass(passed_class_obj[name].__class__, expected_class_obj):
                        raise ValueError(
                            f"{passed_class_obj[name]} is of type: {type(passed_class_obj[name])}, but should be"
                            f" {expected_class_obj}"
                        )
                else:
                    logger.warn(
                        f"You have passed a non-standard module {passed_class_obj[name]}. We cannot verify whether it"
                        " has the correct type"
                    )

                # set passed class object
                loaded_sub_model = passed_class_obj[name]
            elif is_pipeline_module:
461
462
463
                pipeline_module = getattr(pipelines, library_name)
                class_obj = getattr(pipeline_module, class_name)
                importable_classes = ALL_IMPORTABLE_CLASSES
Patrick von Platen's avatar
Patrick von Platen committed
464
                class_candidates = {c: class_obj for c in importable_classes.keys()}
patil-suraj's avatar
patil-suraj committed
465
            else:
patil-suraj's avatar
patil-suraj committed
466
                # else we just import it from the library.
patil-suraj's avatar
patil-suraj committed
467
468
                library = importlib.import_module(library_name)
                class_obj = getattr(library, class_name)
469
                importable_classes = LOADABLE_CLASSES[library_name]
patil-suraj's avatar
patil-suraj committed
470
                class_candidates = {c: getattr(library, c) for c in importable_classes.keys()}
471

472
473
474
475
476
            if loaded_sub_model is None:
                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
477

478
479
480
481
482
483
484
485
486
487
                if load_method_name is None:
                    none_module = class_obj.__module__
                    if none_module.startswith(DUMMY_MODULES_FOLDER) and "dummy" in none_module:
                        # call class_obj for nice error message of missing requirements
                        class_obj()

                    raise ValueError(
                        f"The component {class_obj} of {pipeline_class} cannot be loaded as it does not seem to have"
                        f" any of the loading methods defined in {ALL_IMPORTABLE_CLASSES}."
                    )
Patrick von Platen's avatar
Patrick von Platen committed
488

489
                load_method = getattr(class_obj, load_method_name)
490
                loading_kwargs = {}
491

492
493
                if issubclass(class_obj, torch.nn.Module):
                    loading_kwargs["torch_dtype"] = torch_dtype
494
495
                if issubclass(class_obj, diffusers.OnnxRuntimeModel):
                    loading_kwargs["provider"] = provider
496
                    loading_kwargs["sess_options"] = sess_options
497

498
499
500
501
502
503
504
                if (
                    issubclass(class_obj, diffusers.ModelMixin)
                    or is_transformers_available()
                    and issubclass(class_obj, PreTrainedModel)
                ):
                    loading_kwargs["device_map"] = device_map

505
506
                # check if the module is in a subdirectory
                if os.path.isdir(os.path.join(cached_folder, name)):
507
                    loaded_sub_model = load_method(os.path.join(cached_folder, name), **loading_kwargs)
508
509
                else:
                    # else load from the root directory
510
                    loaded_sub_model = load_method(cached_folder, **loading_kwargs)
Patrick von Platen's avatar
Patrick von Platen committed
511

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

Patrick von Platen's avatar
Patrick von Platen committed
514
515
516
517
518
519
        # 4. Potentially add passed objects if expected
        missing_modules = set(expected_modules) - set(init_kwargs.keys())
        if len(missing_modules) > 0 and missing_modules <= set(passed_class_obj.keys()):
            for module in missing_modules:
                init_kwargs[module] = passed_class_obj[module]
        elif len(missing_modules) > 0:
520
            passed_modules = set(list(init_kwargs.keys()) + list(passed_class_obj.keys()))
Patrick von Platen's avatar
Patrick von Platen committed
521
522
523
524
525
            raise ValueError(
                f"Pipeline {pipeline_class} expected {expected_modules}, but only {passed_modules} were passed."
            )

        # 5. Instantiate the pipeline
526
        model = pipeline_class(**init_kwargs)
Patrick von Platen's avatar
Patrick von Platen committed
527
        return model
528
529
530
531
532
533
534
535
536
537
538
539

    @staticmethod
    def numpy_to_pil(images):
        """
        Convert a numpy image or a batch of images to a PIL image.
        """
        if images.ndim == 3:
            images = images[None, ...]
        images = (images * 255).round().astype("uint8")
        pil_images = [Image.fromarray(image) for image in images]

        return pil_images
hysts's avatar
hysts committed
540
541
542
543
544
545
546
547
548
549
550
551
552

    def progress_bar(self, iterable):
        if not hasattr(self, "_progress_bar_config"):
            self._progress_bar_config = {}
        elif not isinstance(self._progress_bar_config, dict):
            raise ValueError(
                f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}."
            )

        return tqdm(iterable, **self._progress_bar_config)

    def set_progress_bar_config(self, **kwargs):
        self._progress_bar_config = kwargs