pipeline_repaint.py 9.74 KB
Newer Older
Patrick von Platen's avatar
Patrick von Platen committed
1
# Copyright 2023 ETH Zurich Computer Vision Lab and The HuggingFace Team. All rights reserved.
Revist's avatar
Revist 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
Revist's avatar
Revist committed
17
18
19

import numpy as np
import PIL
20
import torch
Revist's avatar
Revist committed
21
22
23

from ...models import UNet2DModel
from ...schedulers import RePaintScheduler
24
from ...utils import PIL_INTERPOLATION, deprecate, logging, randn_tensor
25
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
Revist's avatar
Revist committed
26
27


28
29
30
31
32
logger = logging.get_logger(__name__)  # pylint: disable=invalid-name


# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess
def _preprocess_image(image: Union[List, PIL.Image.Image, torch.Tensor]):
33
34
    deprecation_message = "The preprocess method is deprecated and will be removed in diffusers 1.0.0. Please use VaeImageProcessor.preprocess(...) instead"
    deprecate("preprocess", "1.0.0", deprecation_message, standard_warn=False)
35
36
37
38
39
40
41
    if isinstance(image, torch.Tensor):
        return image
    elif isinstance(image, PIL.Image.Image):
        image = [image]

    if isinstance(image[0], PIL.Image.Image):
        w, h = image[0].size
42
        w, h = (x - x % 8 for x in (w, h))  # resize to integer multiple of 8
43
44
45
46
47
48
49
50
51

        image = [np.array(i.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image]
        image = np.concatenate(image, axis=0)
        image = np.array(image).astype(np.float32) / 255.0
        image = image.transpose(0, 3, 1, 2)
        image = 2.0 * image - 1.0
        image = torch.from_numpy(image)
    elif isinstance(image[0], torch.Tensor):
        image = torch.cat(image, dim=0)
Revist's avatar
Revist committed
52
53
54
    return image


55
56
57
58
59
60
61
62
def _preprocess_mask(mask: Union[List, PIL.Image.Image, torch.Tensor]):
    if isinstance(mask, torch.Tensor):
        return mask
    elif isinstance(mask, PIL.Image.Image):
        mask = [mask]

    if isinstance(mask[0], PIL.Image.Image):
        w, h = mask[0].size
63
        w, h = (x - x % 32 for x in (w, h))  # resize to integer multiple of 32
64
65
66
67
68
69
70
71
        mask = [np.array(m.convert("L").resize((w, h), resample=PIL_INTERPOLATION["nearest"]))[None, :] for m in mask]
        mask = np.concatenate(mask, axis=0)
        mask = mask.astype(np.float32) / 255.0
        mask[mask < 0.5] = 0
        mask[mask >= 0.5] = 1
        mask = torch.from_numpy(mask)
    elif isinstance(mask[0], torch.Tensor):
        mask = torch.cat(mask, dim=0)
Revist's avatar
Revist committed
72
73
74
75
    return mask


class RePaintPipeline(DiffusionPipeline):
76
77
78
79
80
81
82
83
84
85
86
87
88
    r"""
    Pipeline for image inpainting using RePaint.

    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.).

    Parameters:
        unet ([`UNet2DModel`]):
            A `UNet2DModel` to denoise the encoded image latents.
        scheduler ([`RePaintScheduler`]):
            A `RePaintScheduler` to be used in combination with `unet` to denoise the encoded image.
    """

Revist's avatar
Revist committed
89
90
91
92
93
94
95
96
97
98
    unet: UNet2DModel
    scheduler: RePaintScheduler

    def __init__(self, unet, scheduler):
        super().__init__()
        self.register_modules(unet=unet, scheduler=scheduler)

    @torch.no_grad()
    def __call__(
        self,
99
100
        image: Union[torch.Tensor, PIL.Image.Image],
        mask_image: Union[torch.Tensor, PIL.Image.Image],
Revist's avatar
Revist committed
101
102
103
104
        num_inference_steps: int = 250,
        eta: float = 0.0,
        jump_length: int = 10,
        jump_n_sample: int = 10,
105
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
Revist's avatar
Revist committed
106
107
108
109
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
    ) -> Union[ImagePipelineOutput, Tuple]:
        r"""
110
111
        The call function to the pipeline for generation.

Revist's avatar
Revist committed
112
        Args:
113
            image (`torch.FloatTensor` or `PIL.Image.Image`):
Revist's avatar
Revist committed
114
115
                The original image to inpaint on.
            mask_image (`torch.FloatTensor` or `PIL.Image.Image`):
116
                The mask_image where 0.0 define which part of the original image to inpaint.
Revist's avatar
Revist committed
117
118
119
120
            num_inference_steps (`int`, *optional*, defaults to 1000):
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
            eta (`float`):
121
122
                The weight of the added noise in a diffusion step. Its value is between 0.0 and 1.0; 0.0 corresponds to
                DDIM and 1.0 is the DDPM scheduler.
Revist's avatar
Revist committed
123
124
            jump_length (`int`, *optional*, defaults to 10):
                The number of steps taken forward in time before going backward in time for a single jump ("j" in
125
                RePaint paper). Take a look at Figure 9 and 10 in the [paper](https://arxiv.org/pdf/2201.09865.pdf).
Revist's avatar
Revist committed
126
            jump_n_sample (`int`, *optional*, defaults to 10):
127
128
                The number of times to make a forward time jump for a given chosen time sample. Take a look at Figure 9
                and 10 in the [paper](https://arxiv.org/pdf/2201.09865.pdf).
Revist's avatar
Revist committed
129
            generator (`torch.Generator`, *optional*):
130
131
132
133
                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`.
Revist's avatar
Revist committed
134
            return_dict (`bool`, *optional*, defaults to `True`):
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
                Whether or not to return a [`ImagePipelineOutput`] instead of a plain tuple.

        Example:

        ```py
        >>> from io import BytesIO
        >>> import torch
        >>> import PIL
        >>> import requests
        >>> from diffusers import RePaintPipeline, RePaintScheduler


        >>> def download_image(url):
        ...     response = requests.get(url)
        ...     return PIL.Image.open(BytesIO(response.content)).convert("RGB")


        >>> img_url = "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/celeba_hq_256.png"
        >>> mask_url = "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/mask_256.png"

        >>> # Load the original image and the mask as PIL images
        >>> original_image = download_image(img_url).resize((256, 256))
        >>> mask_image = download_image(mask_url).resize((256, 256))

        >>> # Load the RePaint scheduler and pipeline based on a pretrained DDPM model
        >>> scheduler = RePaintScheduler.from_pretrained("google/ddpm-ema-celebahq-256")
        >>> pipe = RePaintPipeline.from_pretrained("google/ddpm-ema-celebahq-256", scheduler=scheduler)
        >>> pipe = pipe.to("cuda")

        >>> generator = torch.Generator(device="cuda").manual_seed(0)
        >>> output = pipe(
        ...     image=original_image,
        ...     mask_image=mask_image,
        ...     num_inference_steps=250,
        ...     eta=0.0,
        ...     jump_length=10,
        ...     jump_n_sample=10,
        ...     generator=generator,
        ... )
        >>> inpainted_image = output.images[0]
        ```
Revist's avatar
Revist committed
176
177

        Returns:
178
179
180
            [`~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.
Revist's avatar
Revist committed
181
182
        """

183
        original_image = image
184
185

        original_image = _preprocess_image(original_image)
186
        original_image = original_image.to(device=self._execution_device, dtype=self.unet.dtype)
187
        mask_image = _preprocess_mask(mask_image)
188
        mask_image = mask_image.to(device=self._execution_device, dtype=self.unet.dtype)
Revist's avatar
Revist committed
189

190
191
        batch_size = original_image.shape[0]

Revist's avatar
Revist committed
192
        # sample gaussian noise to begin the loop
193
194
195
196
197
198
199
        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."
            )

        image_shape = original_image.shape
200
        image = randn_tensor(image_shape, generator=generator, device=self._execution_device, dtype=self.unet.dtype)
Revist's avatar
Revist committed
201
202

        # set step values
203
        self.scheduler.set_timesteps(num_inference_steps, jump_length, jump_n_sample, self._execution_device)
Revist's avatar
Revist committed
204
205
206
        self.scheduler.eta = eta

        t_last = self.scheduler.timesteps[0] + 1
207
        generator = generator[0] if isinstance(generator, list) else generator
208
        for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)):
Revist's avatar
Revist committed
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
            if t < t_last:
                # predict the noise residual
                model_output = self.unet(image, t).sample
                # compute previous image: x_t -> x_t-1
                image = self.scheduler.step(model_output, t, image, original_image, mask_image, generator).prev_sample

            else:
                # compute the reverse: x_t-1 -> x_t
                image = self.scheduler.undo_step(image, t_last, generator)
            t_last = t

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

        if not return_dict:
            return (image,)

        return ImagePipelineOutput(images=image)