pipeline_ddim.py 6.45 KB
Newer Older
Patrick von Platen's avatar
Patrick von Platen committed
1
# Copyright 2023 The HuggingFace Team. All rights reserved.
Patrick von Platen's avatar
Patrick von Platen committed
2
3
4
5
6
7
8
9
10
11
12
13
14
#
# 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.

15
from typing import List, Optional, Tuple, Union
Pedro Cuenca's avatar
Pedro Cuenca committed
16

Patrick von Platen's avatar
Patrick von Platen committed
17
18
import torch

19
from ...schedulers import DDIMScheduler
Dhruv Nair's avatar
Dhruv Nair committed
20
from ...utils.torch_utils import randn_tensor
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
    r"""
26
27
28
29
    Pipeline for image generation.

    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
Kashif Rasul's avatar
Kashif Rasul committed
30
31

    Parameters:
32
33
        unet ([`UNet2DModel`]):
            A `UNet2DModel` to denoise the encoded image latents.
Kashif Rasul's avatar
Kashif Rasul committed
34
35
36
37
        scheduler ([`SchedulerMixin`]):
            A scheduler to be used in combination with `unet` to denoise the encoded image. Can be one of
            [`DDPMScheduler`], or [`DDIMScheduler`].
    """
38

39
    model_cpu_offload_seq = "unet"
Kashif Rasul's avatar
Kashif Rasul committed
40

41
    def __init__(self, unet, scheduler):
Patrick von Platen's avatar
Patrick von Platen committed
42
        super().__init__()
43
44
45
46

        # make sure scheduler can always be converted to DDIM
        scheduler = DDIMScheduler.from_config(scheduler.config)

47
        self.register_modules(unet=unet, scheduler=scheduler)
Patrick von Platen's avatar
Patrick von Platen committed
48

Patrick von Platen's avatar
Patrick von Platen committed
49
    @torch.no_grad()
50
51
    def __call__(
        self,
Sid Sahai's avatar
Sid Sahai committed
52
        batch_size: int = 1,
53
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
Sid Sahai's avatar
Sid Sahai committed
54
55
        eta: float = 0.0,
        num_inference_steps: int = 50,
56
        use_clipped_model_output: Optional[bool] = None,
Sid Sahai's avatar
Sid Sahai committed
57
        output_type: Optional[str] = "pil",
58
59
        return_dict: bool = True,
    ) -> Union[ImagePipelineOutput, Tuple]:
Kashif Rasul's avatar
Kashif Rasul committed
60
        r"""
61
62
        The call function to the pipeline for generation.

Kashif Rasul's avatar
Kashif Rasul committed
63
        Args:
64
            batch_size (`int`, *optional*, defaults to 1):
Kashif Rasul's avatar
Kashif Rasul committed
65
                The number of images to generate.
66
            generator (`torch.Generator`, *optional*):
67
68
                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
                generation deterministic.
69
            eta (`float`, *optional*, defaults to 0.0):
70
71
72
                Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
                to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. A value of `0` corresponds to
                DDIM and `1` corresponds to DDPM.
73
            num_inference_steps (`int`, *optional*, defaults to 50):
Kashif Rasul's avatar
Kashif Rasul committed
74
75
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
76
            use_clipped_model_output (`bool`, *optional*, defaults to `None`):
77
78
                If `True` or `False`, see documentation for [`DDIMScheduler.step`]. If `None`, nothing is passed
                downstream to the scheduler (use `None` for schedulers which don't support this argument).
79
            output_type (`str`, *optional*, defaults to `"pil"`):
80
                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
81
            return_dict (`bool`, *optional*, defaults to `True`):
82
                Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple.
83

84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
        Example:

        ```py
        >>> from diffusers import DDIMPipeline
        >>> import PIL.Image
        >>> import numpy as np

        >>> # load model and scheduler
        >>> pipe = DDIMPipeline.from_pretrained("fusing/ddim-lsun-bedroom")

        >>> # run pipeline in inference (sample random noise and denoise)
        >>> image = pipe(eta=0.0, num_inference_steps=50)

        >>> # process image to PIL
        >>> image_processed = image.cpu().permute(0, 2, 3, 1)
        >>> image_processed = (image_processed + 1.0) * 127.5
        >>> image_processed = image_processed.numpy().astype(np.uint8)
        >>> image_pil = PIL.Image.fromarray(image_processed[0])

        >>> # save image
        >>> image_pil.save("test.png")
        ```

107
        Returns:
108
109
110
            [`~pipelines.ImagePipelineOutput`] or `tuple`:
                If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is
                returned where the first element is a list with the generated images
Kashif Rasul's avatar
Kashif Rasul committed
111
        """
Pedro Cuenca's avatar
Pedro Cuenca committed
112

Patrick von Platen's avatar
Patrick von Platen committed
113
        # Sample gaussian noise to begin loop
114
115
116
117
118
119
120
        if isinstance(self.unet.config.sample_size, int):
            image_shape = (
                batch_size,
                self.unet.config.in_channels,
                self.unet.config.sample_size,
                self.unet.config.sample_size,
            )
121
        else:
122
            image_shape = (batch_size, self.unet.config.in_channels, *self.unet.config.sample_size)
123

124
125
126
127
128
129
        if isinstance(generator, list) and len(generator) != batch_size:
            raise ValueError(
                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
            )

130
        image = randn_tensor(image_shape, generator=generator, device=self._execution_device, dtype=self.unet.dtype)
Patrick von Platen's avatar
Patrick von Platen committed
131

132
133
        # set step values
        self.scheduler.set_timesteps(num_inference_steps)
Patrick von Platen's avatar
Patrick von Platen committed
134

hysts's avatar
hysts committed
135
        for t in self.progress_bar(self.scheduler.timesteps):
Patrick von Platen's avatar
Patrick von Platen committed
136
            # 1. predict noise model_output
137
            model_output = self.unet(image, t).sample
Patrick von Platen's avatar
Patrick von Platen committed
138

139
            # 2. predict previous mean of image x_t-1 and add variance depending on eta
140
            # eta corresponds to η in paper and should be between [0, 1]
141
            # do x_t -> x_t-1
142
143
144
            image = self.scheduler.step(
                model_output, t, image, eta=eta, use_clipped_model_output=use_clipped_model_output, generator=generator
            ).prev_sample
Patrick von Platen's avatar
Patrick von Platen committed
145

146
147
        image = (image / 2 + 0.5).clamp(0, 1)
        image = image.cpu().permute(0, 2, 3, 1).numpy()
anton-l's avatar
anton-l committed
148
149
        if output_type == "pil":
            image = self.numpy_to_pil(image)
150

151
152
153
154
        if not return_dict:
            return (image,)

        return ImagePipelineOutput(images=image)