pipeline_ddpm.py 3.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
# 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 DDPMPipeline(DiffusionPipeline):
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
    def __call__(
Sid Sahai's avatar
Sid Sahai committed
42
43
44
45
46
47
        self,
        batch_size: int = 1,
        generator: Optional[torch.Generator] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        **kwargs,
48
    ) -> Union[ImagePipelineOutput, Tuple]:
49
50
        r"""
        Args:
51
            batch_size (`int`, *optional*, defaults to 1):
52
                The number of images to generate.
53
            generator (`torch.Generator`, *optional*):
54
55
                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation
                deterministic.
56
            output_type (`str`, *optional*, defaults to `"pil"`):
57
                The output format of the generate image. Choose between
58
                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
59
            return_dict (`bool`, *optional*, defaults to `True`):
60
                Whether or not to return a [`~pipeline_utils.ImagePipelineOutput`] instead of a plain tuple.
61
62
63
64
65

        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.
66
        """
Patrick von Platen's avatar
Patrick von Platen committed
67
68

        # Sample gaussian noise to begin loop
Patrick von Platen's avatar
Patrick von Platen committed
69
        image = torch.randn(
Patrick von Platen's avatar
Patrick von Platen committed
70
            (batch_size, self.unet.in_channels, self.unet.sample_size, self.unet.sample_size),
Patrick von Platen's avatar
Patrick von Platen committed
71
72
            generator=generator,
        )
Pedro Cuenca's avatar
Pedro Cuenca committed
73
        image = image.to(self.device)
Patrick von Platen's avatar
Patrick von Platen committed
74

75
76
77
        # set step values
        self.scheduler.set_timesteps(1000)

hysts's avatar
hysts committed
78
        for t in self.progress_bar(self.scheduler.timesteps):
Patrick von Platen's avatar
Patrick von Platen committed
79
            # 1. predict noise model_output
80
            model_output = self.unet(image, t).sample
Patrick von Platen's avatar
Patrick von Platen committed
81

Patrick von Platen's avatar
Patrick von Platen committed
82
            # 2. compute previous image: x_t -> t_t-1
83
            image = self.scheduler.step(model_output, t, image, generator=generator).prev_sample
Patrick von Platen's avatar
Patrick von Platen committed
84

85
86
        image = (image / 2 + 0.5).clamp(0, 1)
        image = image.cpu().permute(0, 2, 3, 1).numpy()
87
88
        if output_type == "pil":
            image = self.numpy_to_pil(image)
89

90
91
92
93
        if not return_dict:
            return (image,)

        return ImagePipelineOutput(images=image)