diffusers/community-pipelines-mirror

1# Copyright 2024 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15 16from math import pi17from typing import Callable, List, Optional, Tuple, Union18 19import numpy as np20import torch21from PIL import Image22 23from diffusers import DDPMScheduler, DiffusionPipeline, ImagePipelineOutput, UNet2DModel24from diffusers.utils.torch_utils import randn_tensor25 26 27class DPSPipeline(DiffusionPipeline):28    r"""29    Pipeline for Diffusion Posterior Sampling.30 31    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods32    implemented for all pipelines (downloading, saving, running on a particular device, etc.).33 34    Parameters:35        unet ([`UNet2DModel`]):36            A `UNet2DModel` to denoise the encoded image latents.37        scheduler ([`SchedulerMixin`]):38            A scheduler to be used in combination with `unet` to denoise the encoded image. Can be one of39            [`DDPMScheduler`], or [`DDIMScheduler`].40    """41 42    model_cpu_offload_seq = "unet"43 44    def __init__(self, unet, scheduler):45        super().__init__()46        self.register_modules(unet=unet, scheduler=scheduler)47 48    @torch.no_grad()49    def __call__(50        self,51        measurement: torch.Tensor,52        operator: torch.nn.Module,53        loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],54        batch_size: int = 1,55        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,56        num_inference_steps: int = 1000,57        output_type: Optional[str] = "pil",58        return_dict: bool = True,59        zeta: float = 0.3,60    ) -> Union[ImagePipelineOutput, Tuple]:61        r"""62        The call function to the pipeline for generation.63 64        Args:65            measurement (`torch.Tensor`, *required*):66                A 'torch.Tensor', the corrupted image67            operator (`torch.nn.Module`, *required*):68                A 'torch.nn.Module', the operator generating the corrupted image69            loss_fn (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *required*):70                A 'Callable[[torch.Tensor, torch.Tensor], torch.Tensor]', the loss function used71                between the measurements, for most of the cases using RMSE is fine.72            batch_size (`int`, *optional*, defaults to 1):73                The number of images to generate.74            generator (`torch.Generator`, *optional*):75                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make76                generation deterministic.77            num_inference_steps (`int`, *optional*, defaults to 1000):78                The number of denoising steps. More denoising steps usually lead to a higher quality image at the79                expense of slower inference.80            output_type (`str`, *optional*, defaults to `"pil"`):81                The output format of the generated image. Choose between `PIL.Image` or `np.array`.82            return_dict (`bool`, *optional*, defaults to `True`):83                Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple.84 85        Example:86 87        ```py88        >>> from diffusers import DDPMPipeline89 90        >>> # load model and scheduler91        >>> pipe = DDPMPipeline.from_pretrained("google/ddpm-cat-256")92 93        >>> # run pipeline in inference (sample random noise and denoise)94        >>> image = pipe().images[0]95 96        >>> # save image97        >>> image.save("ddpm_generated_image.png")98        ```99 100        Returns:101            [`~pipelines.ImagePipelineOutput`] or `tuple`:102                If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is103                returned where the first element is a list with the generated images104        """105        # Sample gaussian noise to begin loop106        if isinstance(self.unet.config.sample_size, int):107            image_shape = (108                batch_size,109                self.unet.config.in_channels,110                self.unet.config.sample_size,111                self.unet.config.sample_size,112            )113        else:114            image_shape = (batch_size, self.unet.config.in_channels, *self.unet.config.sample_size)115 116        if self.device.type == "mps":117            # randn does not work reproducibly on mps118            image = randn_tensor(image_shape, generator=generator)119            image = image.to(self.device)120        else:121            image = randn_tensor(image_shape, generator=generator, device=self.device)122 123        # set step values124        self.scheduler.set_timesteps(num_inference_steps)125 126        for t in self.progress_bar(self.scheduler.timesteps):127            with torch.enable_grad():128                # 1. predict noise model_output129                image = image.requires_grad_()130                model_output = self.unet(image, t).sample131 132                # 2. compute previous image x'_{t-1} and original prediction x0_{t}133                scheduler_out = self.scheduler.step(model_output, t, image, generator=generator)134                image_pred, origi_pred = scheduler_out.prev_sample, scheduler_out.pred_original_sample135 136                # 3. compute y'_t = f(x0_{t})137                measurement_pred = operator(origi_pred)138 139                # 4. compute loss = d(y, y'_t-1)140                loss = loss_fn(measurement, measurement_pred)141                loss.backward()142 143                print("distance: {0:.4f}".format(loss.item()))144 145                with torch.no_grad():146                    image_pred = image_pred - zeta * image.grad147                    image = image_pred.detach()148 149        image = (image / 2 + 0.5).clamp(0, 1)150        image = image.cpu().permute(0, 2, 3, 1).numpy()151        if output_type == "pil":152            image = self.numpy_to_pil(image)153 154        if not return_dict:155            return (image,)156 157        return ImagePipelineOutput(images=image)158 159 160if __name__ == "__main__":161    import scipy162    from torch import nn163    from torchvision.utils import save_image164 165    # defining the operators f(.) of y = f(x)166    # super-resolution operator167    class SuperResolutionOperator(nn.Module):168        def __init__(self, in_shape, scale_factor):169            super().__init__()170 171            # Resizer local class, do not use outiside the SR operator class172            class Resizer(nn.Module):173                def __init__(self, in_shape, scale_factor=None, output_shape=None, kernel=None, antialiasing=True):174                    super(Resizer, self).__init__()175 176                    # First standardize values and fill missing arguments (if needed) by deriving scale from output shape or vice versa177                    scale_factor, output_shape = self.fix_scale_and_size(in_shape, output_shape, scale_factor)178 179                    # Choose interpolation method, each method has the matching kernel size180                    def cubic(x):181                        absx = np.abs(x)182                        absx2 = absx**2183                        absx3 = absx**3184                        return (1.5 * absx3 - 2.5 * absx2 + 1) * (absx <= 1) + (185                            -0.5 * absx3 + 2.5 * absx2 - 4 * absx + 2186                        ) * ((1 < absx) & (absx <= 2))187 188                    def lanczos2(x):189                        return (190                            (np.sin(pi * x) * np.sin(pi * x / 2) + np.finfo(np.float32).eps)191                            / ((pi**2 * x**2 / 2) + np.finfo(np.float32).eps)192                        ) * (abs(x) < 2)193 194                    def box(x):195                        return ((-0.5 <= x) & (x < 0.5)) * 1.0196 197                    def lanczos3(x):198                        return (199                            (np.sin(pi * x) * np.sin(pi * x / 3) + np.finfo(np.float32).eps)200                            / ((pi**2 * x**2 / 3) + np.finfo(np.float32).eps)201                        ) * (abs(x) < 3)202 203                    def linear(x):204                        return (x + 1) * ((-1 <= x) & (x < 0)) + (1 - x) * ((0 <= x) & (x <= 1))205 206                    method, kernel_width = {207                        "cubic": (cubic, 4.0),208                        "lanczos2": (lanczos2, 4.0),209                        "lanczos3": (lanczos3, 6.0),210                        "box": (box, 1.0),211                        "linear": (linear, 2.0),212                        None: (cubic, 4.0),  # set default interpolation method as cubic213                    }.get(kernel)214 215                    # Antialiasing is only used when downscaling216                    antialiasing *= np.any(np.array(scale_factor) < 1)217 218                    # Sort indices of dimensions according to scale of each dimension. since we are going dim by dim this is efficient219                    sorted_dims = np.argsort(np.array(scale_factor))220                    self.sorted_dims = [int(dim) for dim in sorted_dims if scale_factor[dim] != 1]221 222                    # Iterate over dimensions to calculate local weights for resizing and resize each time in one direction223                    field_of_view_list = []224                    weights_list = []225                    for dim in self.sorted_dims:226                        # for each coordinate (along 1 dim), calculate which coordinates in the input image affect its result and the227                        # weights that multiply the values there to get its result.228                        weights, field_of_view = self.contributions(229                            in_shape[dim], output_shape[dim], scale_factor[dim], method, kernel_width, antialiasing230                        )231 232                        # convert to torch tensor233                        weights = torch.tensor(weights.T, dtype=torch.float32)234 235                        # We add singleton dimensions to the weight matrix so we can multiply it with the big tensor we get for236                        # tmp_im[field_of_view.T], (bsxfun style)237                        weights_list.append(238                            nn.Parameter(239                                torch.reshape(weights, list(weights.shape) + (len(scale_factor) - 1) * [1]),240                                requires_grad=False,241                            )242                        )243                        field_of_view_list.append(244                            nn.Parameter(245                                torch.tensor(field_of_view.T.astype(np.int32), dtype=torch.long), requires_grad=False246                            )247                        )248 249                    self.field_of_view = nn.ParameterList(field_of_view_list)250                    self.weights = nn.ParameterList(weights_list)251 252                def forward(self, in_tensor):253                    x = in_tensor254 255                    # Use the affecting position values and the set of weights to calculate the result of resizing along this 1 dim256                    for dim, fov, w in zip(self.sorted_dims, self.field_of_view, self.weights):257                        # To be able to act on each dim, we swap so that dim 0 is the wanted dim to resize258                        x = torch.transpose(x, dim, 0)259 260                        # This is a bit of a complicated multiplication: x[field_of_view.T] is a tensor of order image_dims+1.261                        # for each pixel in the output-image it matches the positions the influence it from the input image (along 1 dim262                        # only, this is why it only adds 1 dim to 5the shape). We then multiply, for each pixel, its set of positions with263                        # the matching set of weights. we do this by this big tensor element-wise multiplication (MATLAB bsxfun style:264                        # matching dims are multiplied element-wise while singletons mean that the matching dim is all multiplied by the265                        # same number266                        x = torch.sum(x[fov] * w, dim=0)267 268                        # Finally we swap back the axes to the original order269                        x = torch.transpose(x, dim, 0)270 271                    return x272 273                def fix_scale_and_size(self, input_shape, output_shape, scale_factor):274                    # First fixing the scale-factor (if given) to be standardized the function expects (a list of scale factors in the275                    # same size as the number of input dimensions)276                    if scale_factor is not None:277                        # By default, if scale-factor is a scalar we assume 2d resizing and duplicate it.278                        if np.isscalar(scale_factor) and len(input_shape) > 1:279                            scale_factor = [scale_factor, scale_factor]280 281                        # We extend the size of scale-factor list to the size of the input by assigning 1 to all the unspecified scales282                        scale_factor = list(scale_factor)283                        scale_factor = [1] * (len(input_shape) - len(scale_factor)) + scale_factor284 285                    # Fixing output-shape (if given): extending it to the size of the input-shape, by assigning the original input-size286                    # to all the unspecified dimensions287                    if output_shape is not None:288                        output_shape = list(input_shape[len(output_shape) :]) + list(np.uint(np.array(output_shape)))289 290                    # Dealing with the case of non-give scale-factor, calculating according to output-shape. note that this is291                    # sub-optimal, because there can be different scales to the same output-shape.292                    if scale_factor is None:293                        scale_factor = 1.0 * np.array(output_shape) / np.array(input_shape)294 295                    # Dealing with missing output-shape. calculating according to scale-factor296                    if output_shape is None:297                        output_shape = np.uint(np.ceil(np.array(input_shape) * np.array(scale_factor)))298 299                    return scale_factor, output_shape300 301                def contributions(self, in_length, out_length, scale, kernel, kernel_width, antialiasing):302                    # This function calculates a set of 'filters' and a set of field_of_view that will later on be applied303                    # such that each position from the field_of_view will be multiplied with a matching filter from the304                    # 'weights' based on the interpolation method and the distance of the sub-pixel location from the pixel centers305                    # around it. This is only done for one dimension of the image.306 307                    # When anti-aliasing is activated (default and only for downscaling) the receptive field is stretched to size of308                    # 1/sf. this means filtering is more 'low-pass filter'.309                    fixed_kernel = (lambda arg: scale * kernel(scale * arg)) if antialiasing else kernel310                    kernel_width *= 1.0 / scale if antialiasing else 1.0311 312                    # These are the coordinates of the output image313                    out_coordinates = np.arange(1, out_length + 1)314 315                    # since both scale-factor and output size can be provided simulatneously, perserving the center of the image requires shifting316                    # the output coordinates. the deviation is because out_length doesn't necesary equal in_length*scale.317                    # to keep the center we need to subtract half of this deivation so that we get equal margins for boths sides and center is preserved.318                    shifted_out_coordinates = out_coordinates - (out_length - in_length * scale) / 2319 320                    # These are the matching positions of the output-coordinates on the input image coordinates.321                    # Best explained by example: say we have 4 horizontal pixels for HR and we downscale by SF=2 and get 2 pixels:322                    # [1,2,3,4] -> [1,2]. Remember each pixel number is the middle of the pixel.323                    # The scaling is done between the distances and not pixel numbers (the right boundary of pixel 4 is transformed to324                    # the right boundary of pixel 2. pixel 1 in the small image matches the boundary between pixels 1 and 2 in the big325                    # one and not to pixel 2. This means the position is not just multiplication of the old pos by scale-factor).326                    # So if we measure distance from the left border, middle of pixel 1 is at distance d=0.5, border between 1 and 2 is327                    # at d=1, and so on (d = p - 0.5).  we calculate (d_new = d_old / sf) which means:328                    # (p_new-0.5 = (p_old-0.5) / sf)     ->          p_new = p_old/sf + 0.5 * (1-1/sf)329                    match_coordinates = shifted_out_coordinates / scale + 0.5 * (1 - 1 / scale)330 331                    # This is the left boundary to start multiplying the filter from, it depends on the size of the filter332                    left_boundary = np.floor(match_coordinates - kernel_width / 2)333 334                    # Kernel width needs to be enlarged because when covering has sub-pixel borders, it must 'see' the pixel centers335                    # of the pixels it only covered a part from. So we add one pixel at each side to consider (weights can zeroize them)336                    expanded_kernel_width = np.ceil(kernel_width) + 2337 338                    # Determine a set of field_of_view for each each output position, these are the pixels in the input image339                    # that the pixel in the output image 'sees'. We get a matrix whos horizontal dim is the output pixels (big) and the340                    # vertical dim is the pixels it 'sees' (kernel_size + 2)341                    field_of_view = np.squeeze(342                        np.int16(np.expand_dims(left_boundary, axis=1) + np.arange(expanded_kernel_width) - 1)343                    )344 345                    # Assign weight to each pixel in the field of view. A matrix whos horizontal dim is the output pixels and the346                    # vertical dim is a list of weights matching to the pixel in the field of view (that are specified in347                    # 'field_of_view')348                    weights = fixed_kernel(1.0 * np.expand_dims(match_coordinates, axis=1) - field_of_view - 1)349 350                    # Normalize weights to sum up to 1. be careful from dividing by 0351                    sum_weights = np.sum(weights, axis=1)352                    sum_weights[sum_weights == 0] = 1.0353                    weights = 1.0 * weights / np.expand_dims(sum_weights, axis=1)354 355                    # We use this mirror structure as a trick for reflection padding at the boundaries356                    mirror = np.uint(np.concatenate((np.arange(in_length), np.arange(in_length - 1, -1, step=-1))))357                    field_of_view = mirror[np.mod(field_of_view, mirror.shape[0])]358 359                    # Get rid of  weights and pixel positions that are of zero weight360                    non_zero_out_pixels = np.nonzero(np.any(weights, axis=0))361                    weights = np.squeeze(weights[:, non_zero_out_pixels])362                    field_of_view = np.squeeze(field_of_view[:, non_zero_out_pixels])363 364                    # Final products are the relative positions and the matching weights, both are output_size X fixed_kernel_size365                    return weights, field_of_view366 367            self.down_sample = Resizer(in_shape, 1 / scale_factor)368            for param in self.parameters():369                param.requires_grad = False370 371        def forward(self, data, **kwargs):372            return self.down_sample(data)373 374    # Gaussian blurring operator375    class GaussialBlurOperator(nn.Module):376        def __init__(self, kernel_size, intensity):377            super().__init__()378 379            class Blurkernel(nn.Module):380                def __init__(self, blur_type="gaussian", kernel_size=31, std=3.0):381                    super().__init__()382                    self.blur_type = blur_type383                    self.kernel_size = kernel_size384                    self.std = std385                    self.seq = nn.Sequential(386                        nn.ReflectionPad2d(self.kernel_size // 2),387                        nn.Conv2d(3, 3, self.kernel_size, stride=1, padding=0, bias=False, groups=3),388                    )389                    self.weights_init()390 391                def forward(self, x):392                    return self.seq(x)393 394                def weights_init(self):395                    if self.blur_type == "gaussian":396                        n = np.zeros((self.kernel_size, self.kernel_size))397                        n[self.kernel_size // 2, self.kernel_size // 2] = 1398                        k = scipy.ndimage.gaussian_filter(n, sigma=self.std)399                        k = torch.from_numpy(k)400                        self.k = k401                        for name, f in self.named_parameters():402                            f.data.copy_(k)403 404                def update_weights(self, k):405                    if not torch.is_tensor(k):406                        k = torch.from_numpy(k)407                    for name, f in self.named_parameters():408                        f.data.copy_(k)409 410                def get_kernel(self):411                    return self.k412 413            self.kernel_size = kernel_size414            self.conv = Blurkernel(blur_type="gaussian", kernel_size=kernel_size, std=intensity)415            self.kernel = self.conv.get_kernel()416            self.conv.update_weights(self.kernel.type(torch.float32))417 418            for param in self.parameters():419                param.requires_grad = False420 421        def forward(self, data, **kwargs):422            return self.conv(data)423 424        def transpose(self, data, **kwargs):425            return data426 427        def get_kernel(self):428            return self.kernel.view(1, 1, self.kernel_size, self.kernel_size)429 430    # assuming the forward process y = f(x) is polluted by Gaussian noise, use l2 norm431    def RMSELoss(yhat, y):432        return torch.sqrt(torch.sum((yhat - y) ** 2))433 434    # set up source image435    src = Image.open("sample.png")436    # read image into [1,3,H,W]437    src = torch.from_numpy(np.array(src, dtype=np.float32)).permute(2, 0, 1)[None]438    # normalize image to [-1,1]439    src = (src / 127.5) - 1.0440    src = src.to("cuda")441 442    # set up operator and measurement443    # operator = SuperResolutionOperator(in_shape=src.shape, scale_factor=4).to("cuda")444    operator = GaussialBlurOperator(kernel_size=61, intensity=3.0).to("cuda")445    measurement = operator(src)446 447    # set up scheduler448    scheduler = DDPMScheduler.from_pretrained("google/ddpm-celebahq-256")449    scheduler.set_timesteps(1000)450 451    # set up model452    model = UNet2DModel.from_pretrained("google/ddpm-celebahq-256").to("cuda")453 454    save_image((src + 1.0) / 2.0, "dps_src.png")455    save_image((measurement + 1.0) / 2.0, "dps_mea.png")456 457    # finally, the pipeline458    dpspipe = DPSPipeline(model, scheduler)459    image = dpspipe(460        measurement=measurement,461        operator=operator,462        loss_fn=RMSELoss,463        zeta=1.0,464    ).images[0]465 466    image.save("dps_generated_image.png")467