"examples/vscode:/vscode.git/clone" did not exist on "81f56f6498108aea70efedb4b050f16a4085cc2e"
pipeline_utils.py 33.1 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
from dataclasses import dataclass
21
from typing import Any, Dict, 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 packaging import version
30
from PIL import Image
hysts's avatar
hysts committed
31
from tqdm.auto import tqdm
Patrick von Platen's avatar
Patrick von Platen committed
32

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 .hub_utils import http_user_agent
36
from .modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT
37
from .schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME
38
39
40
41
42
43
from .utils import (
    CONFIG_NAME,
    DIFFUSERS_CACHE,
    ONNX_WEIGHTS_NAME,
    WEIGHTS_NAME,
    BaseOutput,
44
    deprecate,
45
46
    is_accelerate_available,
    is_torch_version,
47
48
49
50
51
52
    is_transformers_available,
    logging,
)


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

Patrick von Platen's avatar
Patrick von Platen committed
56

Patrick von Platen's avatar
Patrick von Platen committed
57
INDEX_FILE = "diffusion_pytorch_model.bin"
Patrick von Platen's avatar
Patrick von Platen committed
58
CUSTOM_PIPELINE_FILE_NAME = "pipeline.py"
59
DUMMY_MODULES_FOLDER = "diffusers.utils"
60
TRANSFORMERS_DUMMY_MODULES_FOLDER = "transformers.utils"
Patrick von Platen's avatar
Patrick von Platen committed
61
62
63
64
65
66
67


logger = logging.get_logger(__name__)


LOADABLE_CLASSES = {
    "diffusers": {
Patrick von Platen's avatar
Patrick von Platen committed
68
        "ModelMixin": ["save_pretrained", "from_pretrained"],
69
        "SchedulerMixin": ["save_pretrained", "from_pretrained"],
Patrick von Platen's avatar
Patrick von Platen committed
70
        "DiffusionPipeline": ["save_pretrained", "from_pretrained"],
71
        "OnnxRuntimeModel": ["save_pretrained", "from_pretrained"],
Patrick von Platen's avatar
Patrick von Platen committed
72
73
    },
    "transformers": {
anton-l's avatar
anton-l committed
74
        "PreTrainedTokenizer": ["save_pretrained", "from_pretrained"],
75
        "PreTrainedTokenizerFast": ["save_pretrained", "from_pretrained"],
anton-l's avatar
anton-l committed
76
        "PreTrainedModel": ["save_pretrained", "from_pretrained"],
Suraj Patil's avatar
Suraj Patil committed
77
        "FeatureExtractionMixin": ["save_pretrained", "from_pretrained"],
78
79
        "ProcessorMixin": ["save_pretrained", "from_pretrained"],
        "ImageProcessingMixin": ["save_pretrained", "from_pretrained"],
Patrick von Platen's avatar
Patrick von Platen committed
80
81
82
    },
}

83
84
85
86
ALL_IMPORTABLE_CLASSES = {}
for library in LOADABLE_CLASSES:
    ALL_IMPORTABLE_CLASSES.update(LOADABLE_CLASSES[library])

Patrick von Platen's avatar
Patrick von Platen committed
87

88
89
90
91
92
93
94
95
96
97
98
99
100
101
@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]


102
103
104
105
106
107
108
109
110
111
112
113
114
115
@dataclass
class AudioPipelineOutput(BaseOutput):
    """
    Output class for audio pipelines.

    Args:
        audios (`np.ndarray`)
            List of denoised samples of shape `(batch_size, num_channels, sample_rate)`. Numpy array present the
            denoised audio samples of the diffusion pipeline.
    """

    audios: np.ndarray


Patrick von Platen's avatar
Patrick von Platen committed
116
class DiffusionPipeline(ConfigMixin):
117
118
119
120
121
122
123
124
125
126
127
128
    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
129
          components of the diffusion pipeline.
130
    """
Patrick von Platen's avatar
Patrick von Platen committed
131
132
    config_name = "model_index.json"

Patrick von Platen's avatar
up  
Patrick von Platen committed
133
    def register_modules(self, **kwargs):
134
135
        # import it here to avoid circular import
        from diffusers import pipelines
136

Patrick von Platen's avatar
Patrick von Platen committed
137
        for name, module in kwargs.items():
138
            # retrieve library
139
140
141
142
            if module is None:
                register_dict = {name: (None, None)}
            else:
                library = module.__module__.split(".")[0]
143

144
                # check if the module is a pipeline module
145
                pipeline_dir = module.__module__.split(".")[-2] if len(module.__module__.split(".")) > 2 else None
146
147
                path = module.__module__.split(".")
                is_pipeline_module = pipeline_dir in path and hasattr(pipelines, pipeline_dir)
148

149
150
151
152
153
                # 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
                # folder so we set the library to module name.
                if library not in LOADABLE_CLASSES or is_pipeline_module:
                    library = pipeline_dir
patil-suraj's avatar
patil-suraj committed
154

155
156
                # retrieve class_name
                class_name = module.__class__.__name__
Patrick von Platen's avatar
Patrick von Platen committed
157

158
                register_dict = {name: (library, class_name)}
159

Patrick von Platen's avatar
Patrick von Platen committed
160
            # save model index config
161
            self.register_to_config(**register_dict)
Patrick von Platen's avatar
Patrick von Platen committed
162
163
164

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

Patrick von Platen's avatar
Patrick von Platen committed
166
    def save_pretrained(self, save_directory: Union[str, os.PathLike]):
167
168
169
170
171
172
173
174
175
        """
        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
176
177
        self.save_config(save_directory)

Patrick von Platen's avatar
Patrick von Platen committed
178
        model_index_dict = dict(self.config)
Patrick von Platen's avatar
Patrick von Platen committed
179
        model_index_dict.pop("_class_name")
180
        model_index_dict.pop("_diffusers_version")
181
        model_index_dict.pop("_module", None)
Patrick von Platen's avatar
Patrick von Platen committed
182

anton-l's avatar
anton-l committed
183
184
        for pipeline_component_name in model_index_dict.keys():
            sub_model = getattr(self, pipeline_component_name)
185
186
187
188
            if sub_model is None:
                # edge case for saving a pipeline with safety_checker=None
                continue

anton-l's avatar
anton-l committed
189
            model_cls = sub_model.__class__
Patrick von Platen's avatar
Patrick von Platen committed
190
191

            save_method_name = None
anton-l's avatar
anton-l committed
192
193
194
195
            # 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():
196
197
                    class_candidate = getattr(library, base_class, None)
                    if class_candidate is not None and issubclass(model_cls, class_candidate):
anton-l's avatar
anton-l committed
198
199
200
201
202
203
204
205
                        # 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
206

Pedro Cuenca's avatar
Pedro Cuenca committed
207
208
209
210
    def to(self, torch_device: Optional[Union[str, torch.device]] = None):
        if torch_device is None:
            return self

211
        module_names, _, _ = self.extract_init_dict(dict(self.config))
Pedro Cuenca's avatar
Pedro Cuenca committed
212
213
214
        for name in module_names.keys():
            module = getattr(self, name)
            if isinstance(module, torch.nn.Module):
215
                if module.dtype == torch.float16 and str(torch_device) in ["cpu"]:
216
                    logger.warning(
217
218
219
220
221
                        "Pipelines loaded with `torch_dtype=torch.float16` cannot run with `cpu` device. It"
                        " is not recommended to move them to `cpu` as running them will fail. Please make"
                        " sure to use an accelerator to run the pipeline in inference, due to the lack of"
                        " support for`float16` operations on this device in PyTorch. Please, remove the"
                        " `torch_dtype=torch.float16` argument, or use another device for inference."
222
                    )
Pedro Cuenca's avatar
Pedro Cuenca committed
223
224
225
226
227
                module.to(torch_device)
        return self

    @property
    def device(self) -> torch.device:
228
229
230
231
        r"""
        Returns:
            `torch.device`: The torch device on which the pipeline is located.
        """
232
        module_names, _, _ = self.extract_init_dict(dict(self.config))
Pedro Cuenca's avatar
Pedro Cuenca committed
233
234
235
236
237
238
        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
239
240
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
241
        r"""
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
        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
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
            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
307
308
                Adding Custom
                Pipelines](https://huggingface.co/docs/diffusers/using-diffusers/custom_pipeline_overview)
Patrick von Platen's avatar
Patrick von Platen committed
309
310

            torch_dtype (`str` or `torch.dtype`, *optional*):
311
312
313
314
315
316
317
318
319
320
            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`):
321
                Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
322
323
324
325
326
327
328
329
330
331
332
333
334
            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.
335
336
337
338
339
340
341
342
343
344
345
346
347
            device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
                A map that specifies where each submodule should go. It doesn't need to be refined to each
                parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the
                same device.

                To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For
                more information about each option see [designing a device
                map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map).
            low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
                Speed up model loading by not initializing the weights and only loading the pre-trained weights. This
                also tries to not use more than 1x model size in CPU memory (including peak memory) while loading the
                model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch,
                setting this argument to `True` will raise an error.
348
349
350

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

        <Tip>

356
         It is required to be logged in (`huggingface-cli login`) when you want to use private or [gated
apolinario's avatar
apolinario committed
357
         models](https://huggingface.co/docs/hub/models-gated#gated-models), *e.g.* `"runwayml/stable-diffusion-v1-5"`
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378

        </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)
apolinario's avatar
apolinario committed
379
        >>> pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
380

381
        >>> # Use a different scheduler
382
383
        >>> from diffusers import LMSDiscreteScheduler

384
385
        >>> scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config)
        >>> pipeline.scheduler = scheduler
386
        ```
387
388
389
        """
        cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)
        resume_download = kwargs.pop("resume_download", False)
390
        force_download = kwargs.pop("force_download", False)
391
392
393
        proxies = kwargs.pop("proxies", None)
        local_files_only = kwargs.pop("local_files_only", False)
        use_auth_token = kwargs.pop("use_auth_token", None)
394
        revision = kwargs.pop("revision", None)
395
        torch_dtype = kwargs.pop("torch_dtype", None)
Patrick von Platen's avatar
Patrick von Platen committed
396
        custom_pipeline = kwargs.pop("custom_pipeline", None)
397
        provider = kwargs.pop("provider", None)
398
        sess_options = kwargs.pop("sess_options", None)
399
        device_map = kwargs.pop("device_map", None)
400
401
        low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT)

402
403
404
405
406
407
408
409
410
        if low_cpu_mem_usage and not is_accelerate_available():
            low_cpu_mem_usage = False
            logger.warn(
                "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the"
                " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install"
                " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip"
                " install accelerate\n```\n."
            )

411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
        if device_map is not None and not is_torch_version(">=", "1.9.0"):
            raise NotImplementedError(
                "Loading and dispatching requires torch >= 1.9.0. Please either update your PyTorch version or set"
                " `device_map=None`."
            )

        if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"):
            raise NotImplementedError(
                "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set"
                " `low_cpu_mem_usage=False`."
            )

        if low_cpu_mem_usage is False and device_map is not None:
            raise ValueError(
                f"You cannot set `low_cpu_mem_usage` to False while using device_map={device_map} for loading and"
                " dispatching. Please make sure to set `low_cpu_mem_usage=True`."
            )
Patrick von Platen's avatar
Patrick von Platen committed
428

patil-suraj's avatar
patil-suraj committed
429
        # 1. Download the checkpoints and configs
Patrick von Platen's avatar
Patrick von Platen committed
430
        # use snapshot download here to get it working from from_pretrained
Patrick von Platen's avatar
Patrick von Platen committed
431
        if not os.path.isdir(pretrained_model_name_or_path):
432
            config_dict = cls.load_config(
433
434
435
                pretrained_model_name_or_path,
                cache_dir=cache_dir,
                resume_download=resume_download,
436
                force_download=force_download,
437
438
439
440
441
442
443
444
445
446
                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]

447
448
449
            # make sure we don't download flax weights
            ignore_patterns = "*.msgpack"

Patrick von Platen's avatar
Patrick von Platen committed
450
451
452
            if custom_pipeline is not None:
                allow_patterns += [CUSTOM_PIPELINE_FILE_NAME]

453
454
455
456
457
            if cls != DiffusionPipeline:
                requested_pipeline_class = cls.__name__
            else:
                requested_pipeline_class = config_dict.get("_class_name", cls.__name__)
            user_agent = {"pipeline_class": requested_pipeline_class}
458
459
            if custom_pipeline is not None:
                user_agent["custom_pipeline"] = custom_pipeline
460
            user_agent = http_user_agent(user_agent)
461

462
            # download all allow_patterns
463
464
465
466
467
468
469
            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,
470
                revision=revision,
471
                allow_patterns=allow_patterns,
472
                ignore_patterns=ignore_patterns,
473
                user_agent=user_agent,
474
            )
Patrick von Platen's avatar
Patrick von Platen committed
475
476
        else:
            cached_folder = pretrained_model_name_or_path
477

478
        config_dict = cls.load_config(cached_folder)
479

Patrick von Platen's avatar
Patrick von Platen committed
480
        # 2. Load the pipeline class, if using custom module then load it from the hub
481
        # if we load from explicit class, let's use it
Patrick von Platen's avatar
Patrick von Platen committed
482
483
484
485
486
        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:
487
488
            pipeline_class = cls
        else:
Patrick von Platen's avatar
Patrick von Platen committed
489
490
491
            diffusers_module = importlib.import_module(cls.__module__.split(".")[0])
            pipeline_class = getattr(diffusers_module, config_dict["_class_name"])

492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
        # To be removed in 1.0.0
        if pipeline_class.__name__ == "StableDiffusionInpaintPipeline" and version.parse(
            version.parse(config_dict["_diffusers_version"]).base_version
        ) <= version.parse("0.5.1"):
            from diffusers import StableDiffusionInpaintPipeline, StableDiffusionInpaintPipelineLegacy

            pipeline_class = StableDiffusionInpaintPipelineLegacy

            deprecation_message = (
                "You are using a legacy checkpoint for inpainting with Stable Diffusion, therefore we are loading the"
                f" {StableDiffusionInpaintPipelineLegacy} class instead of {StableDiffusionInpaintPipeline}. For"
                " better inpainting results, we strongly suggest using Stable Diffusion's official inpainting"
                " checkpoint: https://huggingface.co/runwayml/stable-diffusion-inpainting instead or adapting your"
                f" checkpoint {pretrained_model_name_or_path} to the format of"
                " https://huggingface.co/runwayml/stable-diffusion-inpainting. Note that we do not actively maintain"
                " the {StableDiffusionInpaintPipelineLegacy} class and will likely remove it in version 1.0.0."
            )
            deprecate("StableDiffusionInpaintPipelineLegacy", "1.0.0", deprecation_message, standard_warn=False)

511
512
513
        # 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
514
        expected_modules = set(inspect.signature(pipeline_class.__init__).parameters.keys()) - set(["self"])
515
516
        passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs}

517
        init_dict, unused_kwargs, _ = pipeline_class.extract_init_dict(config_dict, **kwargs)
518
519
520

        if len(unused_kwargs) > 0:
            logger.warning(f"Keyword arguments {unused_kwargs} not recognized.")
Patrick von Platen's avatar
Patrick von Platen committed
521
522

        init_kwargs = {}
523

524
525
        # import it here to avoid circular import
        from diffusers import pipelines
526

Patrick von Platen's avatar
Patrick von Platen committed
527
        # 3. Load each module in the pipeline
patil-suraj's avatar
patil-suraj committed
528
        for name, (library_name, class_name) in init_dict.items():
529
530
531
532
533
            if class_name is None:
                # edge case for when the pipeline was saved with safety_checker=None
                init_kwargs[name] = None
                continue

534
535
536
537
            # 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:]

538
            is_pipeline_module = hasattr(pipelines, library_name)
539
            loaded_sub_model = None
540
            sub_model_should_be_defined = True
541

542
            # if the model is in a pipeline module, then we load it from the pipeline
543
544
            if name in passed_class_obj:
                # 1. check that passed_class_obj has correct parent class
545
                if not is_pipeline_module and passed_class_obj[name] is not None:
546
547
548
                    library = importlib.import_module(library_name)
                    class_obj = getattr(library, class_name)
                    importable_classes = LOADABLE_CLASSES[library_name]
549
                    class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()}
550
551
552

                    expected_class_obj = None
                    for class_name, class_candidate in class_candidates.items():
553
                        if class_candidate is not None and issubclass(class_obj, class_candidate):
554
555
556
557
558
559
560
                            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}"
                        )
561
562
563
564
565
566
                elif passed_class_obj[name] is None:
                    logger.warn(
                        f"You have passed `None` for {name} to disable its functionality in {pipeline_class}. Note"
                        f" that this might lead to problems when using {pipeline_class} and is not recommended."
                    )
                    sub_model_should_be_defined = False
567
568
569
570
571
572
573
574
575
                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:
576
577
578
                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
579
                class_candidates = {c: class_obj for c in importable_classes.keys()}
patil-suraj's avatar
patil-suraj committed
580
            else:
patil-suraj's avatar
patil-suraj committed
581
                # else we just import it from the library.
patil-suraj's avatar
patil-suraj committed
582
                library = importlib.import_module(library_name)
583

patil-suraj's avatar
patil-suraj committed
584
                class_obj = getattr(library, class_name)
585
                importable_classes = LOADABLE_CLASSES[library_name]
586
                class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()}
587

588
            if loaded_sub_model is None and sub_model_should_be_defined:
589
590
                load_method_name = None
                for class_name, class_candidate in class_candidates.items():
591
                    if class_candidate is not None and issubclass(class_obj, class_candidate):
592
                        load_method_name = importable_classes[class_name][1]
Patrick von Platen's avatar
Patrick von Platen committed
593

594
595
                if load_method_name is None:
                    none_module = class_obj.__module__
596
597
598
599
                    is_dummy_path = none_module.startswith(DUMMY_MODULES_FOLDER) or none_module.startswith(
                        TRANSFORMERS_DUMMY_MODULES_FOLDER
                    )
                    if is_dummy_path and "dummy" in none_module:
600
601
602
603
604
605
606
                        # 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
607

608
                load_method = getattr(class_obj, load_method_name)
609
                loading_kwargs = {}
610

611
612
                if issubclass(class_obj, torch.nn.Module):
                    loading_kwargs["torch_dtype"] = torch_dtype
613
614
                if issubclass(class_obj, diffusers.OnnxRuntimeModel):
                    loading_kwargs["provider"] = provider
615
                    loading_kwargs["sess_options"] = sess_options
616

617
618
619
                is_diffusers_model = issubclass(class_obj, diffusers.ModelMixin)
                is_transformers_model = (
                    is_transformers_available()
620
                    and issubclass(class_obj, PreTrainedModel)
621
622
623
                    and version.parse(version.parse(transformers.__version__).base_version) >= version.parse("4.20.0")
                )

624
                # When loading a transformers model, if the device_map is None, the weights will be initialized as opposed to diffusers.
625
                # To make default loading faster we set the `low_cpu_mem_usage=low_cpu_mem_usage` flag which is `True` by default.
626
                # This makes sure that the weights won't be initialized which significantly speeds up loading.
627
                if is_diffusers_model or is_transformers_model:
628
                    loading_kwargs["device_map"] = device_map
629
                    loading_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage
630

631
632
                # check if the module is in a subdirectory
                if os.path.isdir(os.path.join(cached_folder, name)):
633
                    loaded_sub_model = load_method(os.path.join(cached_folder, name), **loading_kwargs)
634
635
                else:
                    # else load from the root directory
636
                    loaded_sub_model = load_method(cached_folder, **loading_kwargs)
Patrick von Platen's avatar
Patrick von Platen committed
637

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

Patrick von Platen's avatar
Patrick von Platen committed
640
641
642
643
644
645
        # 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:
646
            passed_modules = set(list(init_kwargs.keys()) + list(passed_class_obj.keys()))
Patrick von Platen's avatar
Patrick von Platen committed
647
648
649
650
651
            raise ValueError(
                f"Pipeline {pipeline_class} expected {expected_modules}, but only {passed_modules} were passed."
            )

        # 5. Instantiate the pipeline
652
        model = pipeline_class(**init_kwargs)
Patrick von Platen's avatar
Patrick von Platen committed
653
        return model
654

655
656
657
658
    @property
    def components(self) -> Dict[str, Any]:
        r"""

Yuta Hayashibe's avatar
Yuta Hayashibe committed
659
        The `self.components` property can be useful to run different pipelines with the same weights and
660
661
662
663
664
665
666
667
668
669
670
        configurations to not have to re-allocate memory.

        Examples:

        ```py
        >>> from diffusers import (
        ...     StableDiffusionPipeline,
        ...     StableDiffusionImg2ImgPipeline,
        ...     StableDiffusionInpaintPipeline,
        ... )

671
        >>> img2text = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
672
673
674
675
676
        >>> img2img = StableDiffusionImg2ImgPipeline(**img2text.components)
        >>> inpaint = StableDiffusionInpaintPipeline(**img2text.components)
        ```

        Returns:
Yuta Hayashibe's avatar
Yuta Hayashibe committed
677
            A dictionaly containing all the modules needed to initialize the pipeline.
678
679
680
681
682
683
684
685
686
687
688
689
        """
        components = {k: getattr(self, k) for k in self.config.keys() if not k.startswith("_")}
        expected_modules = set(inspect.signature(self.__init__).parameters.keys()) - set(["self"])

        if set(components.keys()) != expected_modules:
            raise ValueError(
                f"{self} has been incorrectly initialized or {self.__class__} is incorrectly implemented. Expected"
                f" {expected_modules} to be defined, but {components} are defined."
            )

        return components

690
691
692
693
694
695
696
697
    @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")
698
699
700
701
702
        if images.shape[-1] == 1:
            # special case for grayscale (single channel) images
            pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
        else:
            pil_images = [Image.fromarray(image) for image in images]
703
704

        return pil_images
hysts's avatar
hysts committed
705
706
707
708
709
710
711
712
713
714
715
716
717

    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