pipeline_ddim.py 4.19 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
# Copyright 2022 The HuggingFace Team. 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.


Sid Sahai's avatar
Sid Sahai committed
17
from typing import Optional, Tuple, Union
Pedro Cuenca's avatar
Pedro Cuenca committed
18

Patrick von Platen's avatar
Patrick von Platen committed
19
20
import torch

21
from ...pipeline_utils import DiffusionPipeline, ImagePipelineOutput
Patrick von Platen's avatar
Patrick von Platen committed
22
23


Patrick von Platen's avatar
Patrick von Platen committed
24
class DDIMPipeline(DiffusionPipeline):
Kashif Rasul's avatar
Kashif Rasul committed
25
26
27
28
29
30
31
32
33
34
35
    r"""
    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)

    Parameters:
        unet ([`UNet2DModel`]): U-Net architecture to denoise the encoded image.
        scheduler ([`SchedulerMixin`]):
            A scheduler to be used in combination with `unet` to denoise the encoded image. Can be one of
            [`DDPMScheduler`], or [`DDIMScheduler`].
    """

36
    def __init__(self, unet, scheduler):
Patrick von Platen's avatar
Patrick von Platen committed
37
        super().__init__()
38
        self.register_modules(unet=unet, scheduler=scheduler)
Patrick von Platen's avatar
Patrick von Platen committed
39

Patrick von Platen's avatar
Patrick von Platen committed
40
    @torch.no_grad()
41
42
    def __call__(
        self,
Sid Sahai's avatar
Sid Sahai committed
43
44
45
46
47
        batch_size: int = 1,
        generator: Optional[torch.Generator] = None,
        eta: float = 0.0,
        num_inference_steps: int = 50,
        output_type: Optional[str] = "pil",
48
49
50
        return_dict: bool = True,
        **kwargs,
    ) -> Union[ImagePipelineOutput, Tuple]:
Kashif Rasul's avatar
Kashif Rasul committed
51
52
        r"""
        Args:
53
            batch_size (`int`, *optional*, defaults to 1):
Kashif Rasul's avatar
Kashif Rasul committed
54
                The number of images to generate.
55
            generator (`torch.Generator`, *optional*):
Kashif Rasul's avatar
Kashif Rasul committed
56
57
                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation
                deterministic.
58
            eta (`float`, *optional*, defaults to 0.0):
Kashif Rasul's avatar
Kashif Rasul committed
59
                The eta parameter which controls the scale of the variance (0 is DDIM and 1 is one type of DDPM).
60
            num_inference_steps (`int`, *optional*, defaults to 50):
Kashif Rasul's avatar
Kashif Rasul committed
61
62
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
63
            output_type (`str`, *optional*, defaults to `"pil"`):
Kashif Rasul's avatar
Kashif Rasul committed
64
                The output format of the generate image. Choose between
65
                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
66
            return_dict (`bool`, *optional*, defaults to `True`):
Kashif Rasul's avatar
Kashif Rasul committed
67
                Whether or not to return a [`~pipeline_utils.ImagePipelineOutput`] instead of a plain tuple.
68
69
70
71
72

        Returns:
            [`~pipeline_utils.ImagePipelineOutput`] or `tuple`: [`~pipelines.utils.ImagePipelineOutput`] if
            `return_dict` is True, otherwise a `tuple. When returning a tuple, the first element is a list with the
            generated images.
Kashif Rasul's avatar
Kashif Rasul committed
73
        """
Pedro Cuenca's avatar
Pedro Cuenca committed
74

Patrick von Platen's avatar
Patrick von Platen committed
75
        # Sample gaussian noise to begin loop
Patrick von Platen's avatar
Patrick von Platen committed
76
        image = torch.randn(
Patrick von Platen's avatar
Patrick von Platen committed
77
            (batch_size, self.unet.in_channels, self.unet.sample_size, self.unet.sample_size),
Patrick von Platen's avatar
Patrick von Platen committed
78
79
            generator=generator,
        )
Pedro Cuenca's avatar
Pedro Cuenca committed
80
        image = image.to(self.device)
Patrick von Platen's avatar
Patrick von Platen committed
81

82
83
        # set step values
        self.scheduler.set_timesteps(num_inference_steps)
Patrick von Platen's avatar
Patrick von Platen committed
84

hysts's avatar
hysts committed
85
        for t in self.progress_bar(self.scheduler.timesteps):
Patrick von Platen's avatar
Patrick von Platen committed
86
            # 1. predict noise model_output
87
            model_output = self.unet(image, t).sample
Patrick von Platen's avatar
Patrick von Platen committed
88

89
            # 2. predict previous mean of image x_t-1 and add variance depending on eta
90
            # eta corresponds to η in paper and should be between [0, 1]
91
            # do x_t -> x_t-1
92
            image = self.scheduler.step(model_output, t, image, eta).prev_sample
Patrick von Platen's avatar
Patrick von Platen committed
93

94
95
        image = (image / 2 + 0.5).clamp(0, 1)
        image = image.cpu().permute(0, 2, 3, 1).numpy()
anton-l's avatar
anton-l committed
96
97
        if output_type == "pil":
            image = self.numpy_to_pil(image)
98

99
100
101
102
        if not return_dict:
            return (image,)

        return ImagePipelineOutput(images=image)