diffusers/community-pipelines-mirror

1from typing import Any, Callable, Dict, List, Optional, Union2 3import numpy as np4import PIL.Image5import torch6 7from diffusers import StableDiffusionImg2ImgPipeline8from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput9 10 11class MaskedStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):12    debug_save = False13 14    @torch.no_grad()15    def __call__(16        self,17        prompt: Union[str, List[str]] = None,18        image: Union[19            torch.Tensor,20            PIL.Image.Image,21            np.ndarray,22            List[torch.Tensor],23            List[PIL.Image.Image],24            List[np.ndarray],25        ] = None,26        strength: float = 0.8,27        num_inference_steps: Optional[int] = 50,28        guidance_scale: Optional[float] = 7.5,29        negative_prompt: Optional[Union[str, List[str]]] = None,30        num_images_per_prompt: Optional[int] = 1,31        eta: Optional[float] = 0.0,32        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,33        prompt_embeds: Optional[torch.Tensor] = None,34        negative_prompt_embeds: Optional[torch.Tensor] = None,35        output_type: Optional[str] = "pil",36        return_dict: bool = True,37        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,38        callback_steps: int = 1,39        cross_attention_kwargs: Optional[Dict[str, Any]] = None,40        mask: Union[41            torch.Tensor,42            PIL.Image.Image,43            np.ndarray,44            List[torch.Tensor],45            List[PIL.Image.Image],46            List[np.ndarray],47        ] = None,48    ):49        r"""50        The call function to the pipeline for generation.51 52        Args:53            prompt (`str` or `List[str]`, *optional*):54                The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.55            image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):56                `Image` or tensor representing an image batch to be used as the starting point. Can also accept image57                latents as `image`, but if passing latents directly it is not encoded again.58            strength (`float`, *optional*, defaults to 0.8):59                Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a60                starting point and more noise is added the higher the `strength`. The number of denoising steps depends61                on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising62                process runs for the full number of iterations specified in `num_inference_steps`. A value of 163                essentially ignores `image`.64            num_inference_steps (`int`, *optional*, defaults to 50):65                The number of denoising steps. More denoising steps usually lead to a higher quality image at the66                expense of slower inference. This parameter is modulated by `strength`.67            guidance_scale (`float`, *optional*, defaults to 7.5):68                A higher guidance scale value encourages the model to generate images closely linked to the text69                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.70            negative_prompt (`str` or `List[str]`, *optional*):71                The prompt or prompts to guide what to not include in image generation. If not defined, you need to72                pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).73            num_images_per_prompt (`int`, *optional*, defaults to 1):74                The number of images to generate per prompt.75            eta (`float`, *optional*, defaults to 0.0):76                Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies77                to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.78            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):79                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make80                generation deterministic.81            prompt_embeds (`torch.Tensor`, *optional*):82                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not83                provided, text embeddings are generated from the `prompt` input argument.84            negative_prompt_embeds (`torch.Tensor`, *optional*):85                Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If86                not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.87            output_type (`str`, *optional*, defaults to `"pil"`):88                The output format of the generated image. Choose between `PIL.Image` or `np.array`.89            return_dict (`bool`, *optional*, defaults to `True`):90                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a91                plain tuple.92            callback (`Callable`, *optional*):93                A function that calls every `callback_steps` steps during inference. The function is called with the94                following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.95            callback_steps (`int`, *optional*, defaults to 1):96                The frequency at which the `callback` function is called. If not specified, the callback is called at97                every step.98            cross_attention_kwargs (`dict`, *optional*):99                A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in100                [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).101            mask (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`, *optional*):102                A mask with non-zero elements for the area to be inpainted. If not specified, no mask is applied.103        Examples:104 105        Returns:106            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:107                If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,108                otherwise a `tuple` is returned where the first element is a list with the generated images and the109                second element is a list of `bool`s indicating whether the corresponding generated image contains110                "not-safe-for-work" (nsfw) content.111        """112        # code adapted from parent class StableDiffusionImg2ImgPipeline113 114        # 0. Check inputs. Raise error if not correct115        self.check_inputs(prompt, strength, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds)116 117        # 1. Define call parameters118        if prompt is not None and isinstance(prompt, str):119            batch_size = 1120        elif prompt is not None and isinstance(prompt, list):121            batch_size = len(prompt)122        else:123            batch_size = prompt_embeds.shape[0]124        device = self._execution_device125        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)126        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`127        # corresponds to doing no classifier free guidance.128        do_classifier_free_guidance = guidance_scale > 1.0129 130        # 2. Encode input prompt131        text_encoder_lora_scale = (132            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None133        )134        prompt_embeds = self._encode_prompt(135            prompt,136            device,137            num_images_per_prompt,138            do_classifier_free_guidance,139            negative_prompt,140            prompt_embeds=prompt_embeds,141            negative_prompt_embeds=negative_prompt_embeds,142            lora_scale=text_encoder_lora_scale,143        )144 145        # 3. Preprocess image146        image = self.image_processor.preprocess(image)147 148        # 4. set timesteps149        self.scheduler.set_timesteps(num_inference_steps, device=device)150        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)151        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)152 153        # 5. Prepare latent variables154        # it is sampled from the latent distribution of the VAE155        latents = self.prepare_latents(156            image, latent_timestep, batch_size, num_images_per_prompt, prompt_embeds.dtype, device, generator157        )158 159        # mean of the latent distribution160        init_latents = [161            self.vae.encode(image.to(device=device, dtype=prompt_embeds.dtype)[i : i + 1]).latent_dist.mean162            for i in range(batch_size)163        ]164        init_latents = torch.cat(init_latents, dim=0)165 166        # 6. create latent mask167        latent_mask = self._make_latent_mask(latents, mask)168 169        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline170        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)171 172        # 8. Denoising loop173        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order174        with self.progress_bar(total=num_inference_steps) as progress_bar:175            for i, t in enumerate(timesteps):176                # expand the latents if we are doing classifier free guidance177                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents178                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)179 180                # predict the noise residual181                noise_pred = self.unet(182                    latent_model_input,183                    t,184                    encoder_hidden_states=prompt_embeds,185                    cross_attention_kwargs=cross_attention_kwargs,186                    return_dict=False,187                )[0]188 189                # perform guidance190                if do_classifier_free_guidance:191                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)192                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)193 194                if latent_mask is not None:195                    latents = torch.lerp(init_latents * self.vae.config.scaling_factor, latents, latent_mask)196                    noise_pred = torch.lerp(torch.zeros_like(noise_pred), noise_pred, latent_mask)197 198                # compute the previous noisy sample x_t -> x_t-1199                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]200 201                # call the callback, if provided202                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):203                    progress_bar.update()204                    if callback is not None and i % callback_steps == 0:205                        step_idx = i // getattr(self.scheduler, "order", 1)206                        callback(step_idx, t, latents)207 208        if not output_type == "latent":209            scaled = latents / self.vae.config.scaling_factor210            if latent_mask is not None:211                # scaled = latents / self.vae.config.scaling_factor * latent_mask + init_latents * (1 - latent_mask)212                scaled = torch.lerp(init_latents, scaled, latent_mask)213            image = self.vae.decode(scaled, return_dict=False)[0]214            if self.debug_save:215                image_gen = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]216                image_gen = self.image_processor.postprocess(image_gen, output_type=output_type, do_denormalize=[True])217                image_gen[0].save("from_latent.png")218            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)219        else:220            image = latents221            has_nsfw_concept = None222 223        if has_nsfw_concept is None:224            do_denormalize = [True] * image.shape[0]225        else:226            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]227 228        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)229 230        # Offload last model to CPU231        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:232            self.final_offload_hook.offload()233 234        if not return_dict:235            return (image, has_nsfw_concept)236 237        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)238 239    def _make_latent_mask(self, latents, mask):240        if mask is not None:241            latent_mask = []242            if not isinstance(mask, list):243                tmp_mask = [mask]244            else:245                tmp_mask = mask246            _, l_channels, l_height, l_width = latents.shape247            for m in tmp_mask:248                if not isinstance(m, PIL.Image.Image):249                    if len(m.shape) == 2:250                        m = m[..., np.newaxis]251                    if m.max() > 1:252                        m = m / 255.0253                    m = self.image_processor.numpy_to_pil(m)[0]254                if m.mode != "L":255                    m = m.convert("L")256                resized = self.image_processor.resize(m, l_height, l_width)257                if self.debug_save:258                    resized.save("latent_mask.png")259                latent_mask.append(np.repeat(np.array(resized)[np.newaxis, :, :], l_channels, axis=0))260            latent_mask = torch.as_tensor(np.stack(latent_mask)).to(latents)261            latent_mask = latent_mask / latent_mask.max()262        return latent_mask263