pipeline_ddpm.py 3.65 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
39
        scheduler = scheduler.set_format("pt")
        self.register_modules(unet=unet, scheduler=scheduler)
Patrick von Platen's avatar
Patrick von Platen committed
40

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

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

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

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

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

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

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

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

        return ImagePipelineOutput(images=image)