diffusers/community-pipelines-mirror

1import inspect2import time3from pathlib import Path4from typing import Callable, List, Optional, Union5 6import numpy as np7import torch8from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer9 10from diffusers.configuration_utils import FrozenDict11from diffusers.models import AutoencoderKL, UNet2DConditionModel12from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin13from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput14from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker15from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler16from diffusers.utils import deprecate, logging17 18 19logger = logging.get_logger(__name__)  # pylint: disable=invalid-name20 21 22def slerp(t, v0, v1, DOT_THRESHOLD=0.9995):23    """helper function to spherically interpolate two arrays v1 v2"""24 25    if not isinstance(v0, np.ndarray):26        inputs_are_torch = True27        input_device = v0.device28        v0 = v0.cpu().numpy()29        v1 = v1.cpu().numpy()30 31    dot = np.sum(v0 * v1 / (np.linalg.norm(v0) * np.linalg.norm(v1)))32    if np.abs(dot) > DOT_THRESHOLD:33        v2 = (1 - t) * v0 + t * v134    else:35        theta_0 = np.arccos(dot)36        sin_theta_0 = np.sin(theta_0)37        theta_t = theta_0 * t38        sin_theta_t = np.sin(theta_t)39        s0 = np.sin(theta_0 - theta_t) / sin_theta_040        s1 = sin_theta_t / sin_theta_041        v2 = s0 * v0 + s1 * v142 43    if inputs_are_torch:44        v2 = torch.from_numpy(v2).to(input_device)45 46    return v247 48 49class StableDiffusionWalkPipeline(DiffusionPipeline, StableDiffusionMixin):50    r"""51    Pipeline for text-to-image generation using Stable Diffusion.52 53    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the54    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)55 56    Args:57        vae ([`AutoencoderKL`]):58            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.59        text_encoder ([`CLIPTextModel`]):60            Frozen text-encoder. Stable Diffusion uses the text portion of61            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically62            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.63        tokenizer (`CLIPTokenizer`):64            Tokenizer of class65            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).66        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.67        scheduler ([`SchedulerMixin`]):68            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of69            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].70        safety_checker ([`StableDiffusionSafetyChecker`]):71            Classification module that estimates whether generated images could be considered offensive or harmful.72            Please, refer to the [model card](https://huggingface.co/CompVis/stable-diffusion-v1-4) for details.73        feature_extractor ([`CLIPImageProcessor`]):74            Model that extracts features from generated images to be used as inputs for the `safety_checker`.75    """76 77    def __init__(78        self,79        vae: AutoencoderKL,80        text_encoder: CLIPTextModel,81        tokenizer: CLIPTokenizer,82        unet: UNet2DConditionModel,83        scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],84        safety_checker: StableDiffusionSafetyChecker,85        feature_extractor: CLIPImageProcessor,86    ):87        super().__init__()88 89        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:90            deprecation_message = (91                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"92                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "93                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"94                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"95                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"96                " file"97            )98            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)99            new_config = dict(scheduler.config)100            new_config["steps_offset"] = 1101            scheduler._internal_dict = FrozenDict(new_config)102 103        if safety_checker is None:104            logger.warning(105                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"106                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"107                " results in services or applications open to the public. Both the diffusers team and Hugging Face"108                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"109                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"110                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."111            )112 113        self.register_modules(114            vae=vae,115            text_encoder=text_encoder,116            tokenizer=tokenizer,117            unet=unet,118            scheduler=scheduler,119            safety_checker=safety_checker,120            feature_extractor=feature_extractor,121        )122 123    @torch.no_grad()124    def __call__(125        self,126        prompt: Optional[Union[str, List[str]]] = None,127        height: int = 512,128        width: int = 512,129        num_inference_steps: int = 50,130        guidance_scale: float = 7.5,131        negative_prompt: Optional[Union[str, List[str]]] = None,132        num_images_per_prompt: Optional[int] = 1,133        eta: float = 0.0,134        generator: Optional[torch.Generator] = None,135        latents: Optional[torch.Tensor] = None,136        output_type: Optional[str] = "pil",137        return_dict: bool = True,138        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,139        callback_steps: int = 1,140        text_embeddings: Optional[torch.Tensor] = None,141        **kwargs,142    ):143        r"""144        Function invoked when calling the pipeline for generation.145 146        Args:147            prompt (`str` or `List[str]`, *optional*, defaults to `None`):148                The prompt or prompts to guide the image generation. If not provided, `text_embeddings` is required.149            height (`int`, *optional*, defaults to 512):150                The height in pixels of the generated image.151            width (`int`, *optional*, defaults to 512):152                The width in pixels of the generated image.153            num_inference_steps (`int`, *optional*, defaults to 50):154                The number of denoising steps. More denoising steps usually lead to a higher quality image at the155                expense of slower inference.156            guidance_scale (`float`, *optional*, defaults to 7.5):157                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).158                `guidance_scale` is defined as `w` of equation 2. of [Imagen159                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >160                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,161                usually at the expense of lower image quality.162            negative_prompt (`str` or `List[str]`, *optional*):163                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored164                if `guidance_scale` is less than `1`).165            num_images_per_prompt (`int`, *optional*, defaults to 1):166                The number of images to generate per prompt.167            eta (`float`, *optional*, defaults to 0.0):168                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to169                [`schedulers.DDIMScheduler`], will be ignored for others.170            generator (`torch.Generator`, *optional*):171                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation172                deterministic.173            latents (`torch.Tensor`, *optional*):174                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image175                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents176                tensor will ge generated by sampling using the supplied random `generator`.177            output_type (`str`, *optional*, defaults to `"pil"`):178                The output format of the generate image. Choose between179                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.180            return_dict (`bool`, *optional*, defaults to `True`):181                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a182                plain tuple.183            callback (`Callable`, *optional*):184                A function that will be called every `callback_steps` steps during inference. The function will be185                called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.186            callback_steps (`int`, *optional*, defaults to 1):187                The frequency at which the `callback` function will be called. If not specified, the callback will be188                called at every step.189            text_embeddings (`torch.Tensor`, *optional*, defaults to `None`):190                Pre-generated text embeddings to be used as inputs for image generation. Can be used in place of191                `prompt` to avoid re-computing the embeddings. If not provided, the embeddings will be generated from192                the supplied `prompt`.193 194        Returns:195            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:196            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.197            When returning a tuple, the first element is a list with the generated images, and the second element is a198            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"199            (nsfw) content, according to the `safety_checker`.200        """201 202        if height % 8 != 0 or width % 8 != 0:203            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")204 205        if (callback_steps is None) or (206            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)207        ):208            raise ValueError(209                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"210                f" {type(callback_steps)}."211            )212 213        if text_embeddings is None:214            if isinstance(prompt, str):215                batch_size = 1216            elif isinstance(prompt, list):217                batch_size = len(prompt)218            else:219                raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")220 221            # get prompt text embeddings222            text_inputs = self.tokenizer(223                prompt,224                padding="max_length",225                max_length=self.tokenizer.model_max_length,226                return_tensors="pt",227            )228            text_input_ids = text_inputs.input_ids229 230            if text_input_ids.shape[-1] > self.tokenizer.model_max_length:231                removed_text = self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :])232                print(233                    "The following part of your input was truncated because CLIP can only handle sequences up to"234                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"235                )236                text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length]237            text_embeddings = self.text_encoder(text_input_ids.to(self.device))[0]238        else:239            batch_size = text_embeddings.shape[0]240 241        # duplicate text embeddings for each generation per prompt, using mps friendly method242        bs_embed, seq_len, _ = text_embeddings.shape243        text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1)244        text_embeddings = text_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)245 246        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)247        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`248        # corresponds to doing no classifier free guidance.249        do_classifier_free_guidance = guidance_scale > 1.0250        # get unconditional embeddings for classifier free guidance251        if do_classifier_free_guidance:252            uncond_tokens: List[str]253            if negative_prompt is None:254                uncond_tokens = [""] * batch_size255            elif type(prompt) is not type(negative_prompt):256                raise TypeError(257                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="258                    f" {type(prompt)}."259                )260            elif isinstance(negative_prompt, str):261                uncond_tokens = [negative_prompt]262            elif batch_size != len(negative_prompt):263                raise ValueError(264                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"265                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"266                    " the batch size of `prompt`."267                )268            else:269                uncond_tokens = negative_prompt270 271            max_length = self.tokenizer.model_max_length272            uncond_input = self.tokenizer(273                uncond_tokens,274                padding="max_length",275                max_length=max_length,276                truncation=True,277                return_tensors="pt",278            )279            uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]280 281            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method282            seq_len = uncond_embeddings.shape[1]283            uncond_embeddings = uncond_embeddings.repeat(1, num_images_per_prompt, 1)284            uncond_embeddings = uncond_embeddings.view(batch_size * num_images_per_prompt, seq_len, -1)285 286            # For classifier free guidance, we need to do two forward passes.287            # Here we concatenate the unconditional and text embeddings into a single batch288            # to avoid doing two forward passes289            text_embeddings = torch.cat([uncond_embeddings, text_embeddings])290 291        # get the initial random noise unless the user supplied it292 293        # Unlike in other pipelines, latents need to be generated in the target device294        # for 1-to-1 results reproducibility with the CompVis implementation.295        # However this currently doesn't work in `mps`.296        latents_shape = (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)297        latents_dtype = text_embeddings.dtype298        if latents is None:299            if self.device.type == "mps":300                # randn does not work reproducibly on mps301                latents = torch.randn(latents_shape, generator=generator, device="cpu", dtype=latents_dtype).to(302                    self.device303                )304            else:305                latents = torch.randn(latents_shape, generator=generator, device=self.device, dtype=latents_dtype)306        else:307            if latents.shape != latents_shape:308                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}")309            latents = latents.to(self.device)310 311        # set timesteps312        self.scheduler.set_timesteps(num_inference_steps)313 314        # Some schedulers like PNDM have timesteps as arrays315        # It's more optimized to move all timesteps to correct device beforehand316        timesteps_tensor = self.scheduler.timesteps.to(self.device)317 318        # scale the initial noise by the standard deviation required by the scheduler319        latents = latents * self.scheduler.init_noise_sigma320 321        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature322        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.323        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502324        # and should be between [0, 1]325        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())326        extra_step_kwargs = {}327        if accepts_eta:328            extra_step_kwargs["eta"] = eta329 330        for i, t in enumerate(self.progress_bar(timesteps_tensor)):331            # expand the latents if we are doing classifier free guidance332            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents333            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)334 335            # predict the noise residual336            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample337 338            # perform guidance339            if do_classifier_free_guidance:340                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)341                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)342 343            # compute the previous noisy sample x_t -> x_t-1344            latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample345 346            # call the callback, if provided347            if callback is not None and i % callback_steps == 0:348                step_idx = i // getattr(self.scheduler, "order", 1)349                callback(step_idx, t, latents)350 351        latents = 1 / 0.18215 * latents352        image = self.vae.decode(latents).sample353 354        image = (image / 2 + 0.5).clamp(0, 1)355 356        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16357        image = image.cpu().permute(0, 2, 3, 1).float().numpy()358 359        if self.safety_checker is not None:360            safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(361                self.device362            )363            image, has_nsfw_concept = self.safety_checker(364                images=image, clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype)365            )366        else:367            has_nsfw_concept = None368 369        if output_type == "pil":370            image = self.numpy_to_pil(image)371 372        if not return_dict:373            return (image, has_nsfw_concept)374 375        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)376 377    def embed_text(self, text):378        """takes in text and turns it into text embeddings"""379        text_input = self.tokenizer(380            text,381            padding="max_length",382            max_length=self.tokenizer.model_max_length,383            truncation=True,384            return_tensors="pt",385        )386        with torch.no_grad():387            embed = self.text_encoder(text_input.input_ids.to(self.device))[0]388        return embed389 390    def get_noise(self, seed, dtype=torch.float32, height=512, width=512):391        """Takes in random seed and returns corresponding noise vector"""392        return torch.randn(393            (1, self.unet.config.in_channels, height // 8, width // 8),394            generator=torch.Generator(device=self.device).manual_seed(seed),395            device=self.device,396            dtype=dtype,397        )398 399    def walk(400        self,401        prompts: List[str],402        seeds: List[int],403        num_interpolation_steps: Optional[int] = 6,404        output_dir: Optional[str] = "./dreams",405        name: Optional[str] = None,406        batch_size: Optional[int] = 1,407        height: Optional[int] = 512,408        width: Optional[int] = 512,409        guidance_scale: Optional[float] = 7.5,410        num_inference_steps: Optional[int] = 50,411        eta: Optional[float] = 0.0,412    ) -> List[str]:413        """414        Walks through a series of prompts and seeds, interpolating between them and saving the results to disk.415 416        Args:417            prompts (`List[str]`):418                List of prompts to generate images for.419            seeds (`List[int]`):420                List of seeds corresponding to provided prompts. Must be the same length as prompts.421            num_interpolation_steps (`int`, *optional*, defaults to 6):422                Number of interpolation steps to take between prompts.423            output_dir (`str`, *optional*, defaults to `./dreams`):424                Directory to save the generated images to.425            name (`str`, *optional*, defaults to `None`):426                Subdirectory of `output_dir` to save the generated images to. If `None`, the name will427                be the current time.428            batch_size (`int`, *optional*, defaults to 1):429                Number of images to generate at once.430            height (`int`, *optional*, defaults to 512):431                Height of the generated images.432            width (`int`, *optional*, defaults to 512):433                Width of the generated images.434            guidance_scale (`float`, *optional*, defaults to 7.5):435                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).436                `guidance_scale` is defined as `w` of equation 2. of [Imagen437                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >438                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,439                usually at the expense of lower image quality.440            num_inference_steps (`int`, *optional*, defaults to 50):441                The number of denoising steps. More denoising steps usually lead to a higher quality image at the442                expense of slower inference.443            eta (`float`, *optional*, defaults to 0.0):444                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to445                [`schedulers.DDIMScheduler`], will be ignored for others.446 447        Returns:448            `List[str]`: List of paths to the generated images.449        """450        if not len(prompts) == len(seeds):451            raise ValueError(452                f"Number of prompts and seeds must be equalGot {len(prompts)} prompts and {len(seeds)} seeds"453            )454 455        name = name or time.strftime("%Y%m%d-%H%M%S")456        save_path = Path(output_dir) / name457        save_path.mkdir(exist_ok=True, parents=True)458 459        frame_idx = 0460        frame_filepaths = []461        for prompt_a, prompt_b, seed_a, seed_b in zip(prompts, prompts[1:], seeds, seeds[1:]):462            # Embed Text463            embed_a = self.embed_text(prompt_a)464            embed_b = self.embed_text(prompt_b)465 466            # Get Noise467            noise_dtype = embed_a.dtype468            noise_a = self.get_noise(seed_a, noise_dtype, height, width)469            noise_b = self.get_noise(seed_b, noise_dtype, height, width)470 471            noise_batch, embeds_batch = None, None472            T = np.linspace(0.0, 1.0, num_interpolation_steps)473            for i, t in enumerate(T):474                noise = slerp(float(t), noise_a, noise_b)475                embed = torch.lerp(embed_a, embed_b, t)476 477                noise_batch = noise if noise_batch is None else torch.cat([noise_batch, noise], dim=0)478                embeds_batch = embed if embeds_batch is None else torch.cat([embeds_batch, embed], dim=0)479 480                batch_is_ready = embeds_batch.shape[0] == batch_size or i + 1 == T.shape[0]481                if batch_is_ready:482                    outputs = self(483                        latents=noise_batch,484                        text_embeddings=embeds_batch,485                        height=height,486                        width=width,487                        guidance_scale=guidance_scale,488                        eta=eta,489                        num_inference_steps=num_inference_steps,490                    )491                    noise_batch, embeds_batch = None, None492 493                    for image in outputs["images"]:494                        frame_filepath = str(save_path / f"frame_{frame_idx:06d}.png")495                        image.save(frame_filepath)496                        frame_filepaths.append(frame_filepath)497                        frame_idx += 1498        return frame_filepaths499