pipeline_ddim.py 6.74 KB
Newer Older
1
# Copyright 2024 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 ...models import UNet2DModel
20
from ...schedulers import DDIMScheduler
hlky's avatar
hlky committed
21
from ...utils import is_torch_xla_available
Dhruv Nair's avatar
Dhruv Nair committed
22
from ...utils.torch_utils import randn_tensor
23
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
Patrick von Platen's avatar
Patrick von Platen committed
24
25


hlky's avatar
hlky committed
26
27
28
29
30
31
32
33
if is_torch_xla_available():
    import torch_xla.core.xla_model as xm

    XLA_AVAILABLE = True
else:
    XLA_AVAILABLE = False


Patrick von Platen's avatar
Patrick von Platen committed
34
class DDIMPipeline(DiffusionPipeline):
Kashif Rasul's avatar
Kashif Rasul committed
35
    r"""
36
37
38
39
    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
40
41

    Parameters:
42
43
        unet ([`UNet2DModel`]):
            A `UNet2DModel` to denoise the encoded image latents.
Kashif Rasul's avatar
Kashif Rasul committed
44
45
46
47
        scheduler ([`SchedulerMixin`]):
            A scheduler to be used in combination with `unet` to denoise the encoded image. Can be one of
            [`DDPMScheduler`], or [`DDIMScheduler`].
    """
48

49
    model_cpu_offload_seq = "unet"
Kashif Rasul's avatar
Kashif Rasul committed
50

51
    def __init__(self, unet: UNet2DModel, scheduler: DDIMScheduler):
Patrick von Platen's avatar
Patrick von Platen committed
52
        super().__init__()
53
54
55
56

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

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

Patrick von Platen's avatar
Patrick von Platen committed
59
    @torch.no_grad()
60
61
    def __call__(
        self,
Sid Sahai's avatar
Sid Sahai committed
62
        batch_size: int = 1,
63
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
Sid Sahai's avatar
Sid Sahai committed
64
65
        eta: float = 0.0,
        num_inference_steps: int = 50,
66
        use_clipped_model_output: Optional[bool] = None,
Sid Sahai's avatar
Sid Sahai committed
67
        output_type: Optional[str] = "pil",
68
69
        return_dict: bool = True,
    ) -> Union[ImagePipelineOutput, Tuple]:
Kashif Rasul's avatar
Kashif Rasul committed
70
        r"""
71
72
        The call function to the pipeline for generation.

Kashif Rasul's avatar
Kashif Rasul committed
73
        Args:
74
            batch_size (`int`, *optional*, defaults to 1):
Kashif Rasul's avatar
Kashif Rasul committed
75
                The number of images to generate.
76
            generator (`torch.Generator`, *optional*):
77
78
                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
                generation deterministic.
79
            eta (`float`, *optional*, defaults to 0.0):
80
81
82
                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.
83
            num_inference_steps (`int`, *optional*, defaults to 50):
Kashif Rasul's avatar
Kashif Rasul committed
84
85
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
86
            use_clipped_model_output (`bool`, *optional*, defaults to `None`):
87
88
                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).
89
            output_type (`str`, *optional*, defaults to `"pil"`):
90
                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
91
            return_dict (`bool`, *optional*, defaults to `True`):
92
                Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple.
93

94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
        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")
        ```

117
        Returns:
118
119
120
            [`~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
121
        """
Pedro Cuenca's avatar
Pedro Cuenca committed
122

Patrick von Platen's avatar
Patrick von Platen committed
123
        # Sample gaussian noise to begin loop
124
125
126
127
128
129
130
        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,
            )
131
        else:
132
            image_shape = (batch_size, self.unet.config.in_channels, *self.unet.config.sample_size)
133

134
135
136
137
138
139
        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."
            )

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

142
143
        # set step values
        self.scheduler.set_timesteps(num_inference_steps)
Patrick von Platen's avatar
Patrick von Platen committed
144

hysts's avatar
hysts committed
145
        for t in self.progress_bar(self.scheduler.timesteps):
Patrick von Platen's avatar
Patrick von Platen committed
146
            # 1. predict noise model_output
147
            model_output = self.unet(image, t).sample
Patrick von Platen's avatar
Patrick von Platen committed
148

149
            # 2. predict previous mean of image x_t-1 and add variance depending on eta
150
            # eta corresponds to η in paper and should be between [0, 1]
151
            # do x_t -> x_t-1
152
153
154
            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
155

hlky's avatar
hlky committed
156
157
158
            if XLA_AVAILABLE:
                xm.mark_step()

159
160
        image = (image / 2 + 0.5).clamp(0, 1)
        image = image.cpu().permute(0, 2, 3, 1).numpy()
anton-l's avatar
anton-l committed
161
162
        if output_type == "pil":
            image = self.numpy_to_pil(image)
163

164
165
166
167
        if not return_dict:
            return (image,)

        return ImagePipelineOutput(images=image)