pipeline_pndm.py 4.53 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
15
#
# 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.


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

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

Partho's avatar
Partho committed
20
21
from ...models import UNet2DModel
from ...schedulers import PNDMScheduler
22
from ...utils import randn_tensor
23
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
Patrick von Platen's avatar
Patrick von Platen committed
24
25


Patrick von Platen's avatar
Patrick von Platen committed
26
class PNDMPipeline(DiffusionPipeline):
Kashif Rasul's avatar
Kashif Rasul committed
27
    r"""
28
29
30
31
    Pipeline for unconditional 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
32
33

    Parameters:
34
35
36
37
        unet ([`UNet2DModel`]):
            A `UNet2DModel` to denoise the encoded image latents.
        scheduler ([`PNDMScheduler`]):
            A `PNDMScheduler` to be used in combination with `unet` to denoise the encoded image.
Kashif Rasul's avatar
Kashif Rasul committed
38
39
    """

Partho's avatar
Partho committed
40
41
42
43
    unet: UNet2DModel
    scheduler: PNDMScheduler

    def __init__(self, unet: UNet2DModel, scheduler: PNDMScheduler):
Patrick von Platen's avatar
Patrick von Platen committed
44
        super().__init__()
45
46
47

        scheduler = PNDMScheduler.from_config(scheduler.config)

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

Patrick von Platen's avatar
Patrick von Platen committed
50
    @torch.no_grad()
Partho's avatar
Partho committed
51
52
53
54
    def __call__(
        self,
        batch_size: int = 1,
        num_inference_steps: int = 50,
55
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
Partho's avatar
Partho committed
56
        output_type: Optional[str] = "pil",
57
        return_dict: bool = True,
Partho's avatar
Partho committed
58
        **kwargs,
59
    ) -> 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
65
            batch_size (`int`, `optional`, defaults to 1):
                The number of images to generate.
66
67
68
            num_inference_steps (`int`, `optional`, defaults to 50):
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
            generator (`torch.Generator`, `optional`):
                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
                generation deterministic.
            output_type (`str`, `optional`, defaults to `"pil"`):
                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`ImagePipelineOutput`] instead of a plain tuple.

        Example:

        ```py
        >>> from diffusers import PNDMPipeline

        >>> # load model and scheduler
        >>> pndm = PNDMPipeline.from_pretrained("google/ddpm-cifar10-32")

        >>> # run pipeline in inference (sample random noise and denoise)
        >>> image = pndm().images[0]

        >>> # save image
        >>> image.save("pndm_generated_image.png")
        ```
91
92

        Returns:
93
94
95
            [`~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
96
        """
Patrick von Platen's avatar
Patrick von Platen committed
97
98
        # For more information on the sampling method you can take a look at Algorithm 2 of
        # the official paper: https://arxiv.org/pdf/2202.09778.pdf
Patrick von Platen's avatar
Patrick von Platen committed
99
100

        # Sample gaussian noise to begin loop
101
        image = randn_tensor(
102
            (batch_size, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size),
Patrick von Platen's avatar
Patrick von Platen committed
103
            generator=generator,
104
            device=self.device,
Patrick von Platen's avatar
Patrick von Platen committed
105
106
        )

107
        self.scheduler.set_timesteps(num_inference_steps)
hysts's avatar
hysts committed
108
        for t in self.progress_bar(self.scheduler.timesteps):
109
            model_output = self.unet(image, t).sample
Patrick von Platen's avatar
Patrick von Platen committed
110

111
            image = self.scheduler.step(model_output, t, image).prev_sample
Patrick von Platen's avatar
Patrick von Platen committed
112

113
114
        image = (image / 2 + 0.5).clamp(0, 1)
        image = image.cpu().permute(0, 2, 3, 1).numpy()
115
116
        if output_type == "pil":
            image = self.numpy_to_pil(image)
117

118
119
120
121
        if not return_dict:
            return (image,)

        return ImagePipelineOutput(images=image)