diffusers/community-pipelines-mirror

1"""2    modified based on diffusion library from Huggingface: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py3"""4import inspect5import warnings6from typing import List, Optional, Union7 8import torch9 10from diffusers.models import AutoencoderKL, UNet2DConditionModel11from diffusers.pipeline_utils import DiffusionPipeline12from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput13from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker14from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler15from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer16 17 18class ComposableStableDiffusionPipeline(DiffusionPipeline):19    r"""20    Pipeline for text-to-image generation using Stable Diffusion.21    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the22    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)23    Args:24        vae ([`AutoencoderKL`]):25            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.26        text_encoder ([`CLIPTextModel`]):27            Frozen text-encoder. Stable Diffusion uses the text portion of28            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically29            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.30        tokenizer (`CLIPTokenizer`):31            Tokenizer of class32            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).33        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.34        scheduler ([`SchedulerMixin`]):35            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of36            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].37        safety_checker ([`StableDiffusionSafetyChecker`]):38            Classification module that estimates whether generated images could be considered offsensive or harmful.39            Please, refer to the [model card](https://huggingface.co/CompVis/stable-diffusion-v1-4) for details.40        feature_extractor ([`CLIPFeatureExtractor`]):41            Model that extracts features from generated images to be used as inputs for the `safety_checker`.42    """43 44    def __init__(45        self,46        vae: AutoencoderKL,47        text_encoder: CLIPTextModel,48        tokenizer: CLIPTokenizer,49        unet: UNet2DConditionModel,50        scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],51        safety_checker: StableDiffusionSafetyChecker,52        feature_extractor: CLIPFeatureExtractor,53    ):54        super().__init__()55        self.register_modules(56            vae=vae,57            text_encoder=text_encoder,58            tokenizer=tokenizer,59            unet=unet,60            scheduler=scheduler,61            safety_checker=safety_checker,62            feature_extractor=feature_extractor,63        )64 65    def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"):66        r"""67        Enable sliced attention computation.68        When this option is enabled, the attention module will split the input tensor in slices, to compute attention69        in several steps. This is useful to save some memory in exchange for a small speed decrease.70        Args:71            slice_size (`str` or `int`, *optional*, defaults to `"auto"`):72                When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If73                a number is provided, uses as many slices as `attention_head_dim // slice_size`. In this case,74                `attention_head_dim` must be a multiple of `slice_size`.75        """76        if slice_size == "auto":77            # half the attention head size is usually a good trade-off between78            # speed and memory79            slice_size = self.unet.config.attention_head_dim // 280        self.unet.set_attention_slice(slice_size)81 82    def disable_attention_slicing(self):83        r"""84        Disable sliced attention computation. If `enable_attention_slicing` was previously invoked, this method will go85        back to computing attention in one step.86        """87        # set slice_size = `None` to disable `attention slicing`88        self.enable_attention_slicing(None)89 90    @torch.no_grad()91    def __call__(92        self,93        prompt: Union[str, List[str]],94        height: Optional[int] = 512,95        width: Optional[int] = 512,96        num_inference_steps: Optional[int] = 50,97        guidance_scale: Optional[float] = 7.5,98        eta: Optional[float] = 0.0,99        generator: Optional[torch.Generator] = None,100        latents: Optional[torch.FloatTensor] = None,101        output_type: Optional[str] = "pil",102        return_dict: bool = True,103        weights: Optional[str] = "",104        **kwargs,105    ):106        r"""107        Function invoked when calling the pipeline for generation.108        Args:109            prompt (`str` or `List[str]`):110                The prompt or prompts to guide the image generation.111            height (`int`, *optional*, defaults to 512):112                The height in pixels of the generated image.113            width (`int`, *optional*, defaults to 512):114                The width in pixels of the generated image.115            num_inference_steps (`int`, *optional*, defaults to 50):116                The number of denoising steps. More denoising steps usually lead to a higher quality image at the117                expense of slower inference.118            guidance_scale (`float`, *optional*, defaults to 7.5):119                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).120                `guidance_scale` is defined as `w` of equation 2. of [Imagen121                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >122                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,123                usually at the expense of lower image quality.124            eta (`float`, *optional*, defaults to 0.0):125                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to126                [`schedulers.DDIMScheduler`], will be ignored for others.127            generator (`torch.Generator`, *optional*):128                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation129                deterministic.130            latents (`torch.FloatTensor`, *optional*):131                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image132                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents133                tensor will ge generated by sampling using the supplied random `generator`.134            output_type (`str`, *optional*, defaults to `"pil"`):135                The output format of the generate image. Choose between136                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.137            return_dict (`bool`, *optional*, defaults to `True`):138                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a139                plain tuple.140        Returns:141            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:142            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.143            When returning a tuple, the first element is a list with the generated images, and the second element is a144            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"145            (nsfw) content, according to the `safety_checker`.146        """147 148        if "torch_device" in kwargs:149            device = kwargs.pop("torch_device")150            warnings.warn(151                "`torch_device` is deprecated as an input argument to `__call__` and will be removed in v0.3.0."152                " Consider using `pipe.to(torch_device)` instead."153            )154 155            # Set device as before (to be removed in 0.3.0)156            if device is None:157                device = "cuda" if torch.cuda.is_available() else "cpu"158            self.to(device)159 160        if isinstance(prompt, str):161            batch_size = 1162        elif isinstance(prompt, list):163            batch_size = len(prompt)164        else:165            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")166 167        if height % 8 != 0 or width % 8 != 0:168            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")169 170        if "|" in prompt:171            prompt = [x.strip() for x in prompt.split("|")]172            print(f"composing {prompt}...")173 174        # get prompt text embeddings175        text_input = self.tokenizer(176            prompt,177            padding="max_length",178            max_length=self.tokenizer.model_max_length,179            truncation=True,180            return_tensors="pt",181        )182        text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0]183 184        if not weights:185            # specify weights for prompts (excluding the unconditional score)186            print("using equal weights for all prompts...")187            pos_weights = torch.tensor(188                [1 / (text_embeddings.shape[0] - 1)] * (text_embeddings.shape[0] - 1), device=self.device189            ).reshape(-1, 1, 1, 1)190            neg_weights = torch.tensor([1.0], device=self.device).reshape(-1, 1, 1, 1)191            mask = torch.tensor([False] + [True] * pos_weights.shape[0], dtype=torch.bool)192        else:193            # set prompt weight for each194            num_prompts = len(prompt) if isinstance(prompt, list) else 1195            weights = [float(w.strip()) for w in weights.split("|")]196            if len(weights) < num_prompts:197                weights.append(1.0)198            weights = torch.tensor(weights, device=self.device)199            assert len(weights) == text_embeddings.shape[0], "weights specified are not equal to the number of prompts"200            pos_weights = []201            neg_weights = []202            mask = []  # first one is unconditional score203            for w in weights:204                if w > 0:205                    pos_weights.append(w)206                    mask.append(True)207                else:208                    neg_weights.append(abs(w))209                    mask.append(False)210            # normalize the weights211            pos_weights = torch.tensor(pos_weights, device=self.device).reshape(-1, 1, 1, 1)212            pos_weights = pos_weights / pos_weights.sum()213            neg_weights = torch.tensor(neg_weights, device=self.device).reshape(-1, 1, 1, 1)214            neg_weights = neg_weights / neg_weights.sum()215            mask = torch.tensor(mask, device=self.device, dtype=torch.bool)216 217        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)218        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`219        # corresponds to doing no classifier free guidance.220        do_classifier_free_guidance = guidance_scale > 1.0221        # get unconditional embeddings for classifier free guidance222        if do_classifier_free_guidance:223            max_length = text_input.input_ids.shape[-1]224 225            if torch.all(mask):226                # no negative prompts, so we use empty string as the negative prompt227                uncond_input = self.tokenizer(228                    [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"229                )230                uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]231 232                # For classifier free guidance, we need to do two forward passes.233                # Here we concatenate the unconditional and text embeddings into a single batch234                # to avoid doing two forward passes235                text_embeddings = torch.cat([uncond_embeddings, text_embeddings])236 237                # update negative weights238                neg_weights = torch.tensor([1.0], device=self.device)239                mask = torch.tensor([False] + mask.detach().tolist(), device=self.device, dtype=torch.bool)240 241        # get the initial random noise unless the user supplied it242 243        # Unlike in other pipelines, latents need to be generated in the target device244        # for 1-to-1 results reproducibility with the CompVis implementation.245        # However this currently doesn't work in `mps`.246        latents_device = "cpu" if self.device.type == "mps" else self.device247        latents_shape = (batch_size, self.unet.in_channels, height // 8, width // 8)248        if latents is None:249            latents = torch.randn(250                latents_shape,251                generator=generator,252                device=latents_device,253            )254        else:255            if latents.shape != latents_shape:256                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}")257        latents = latents.to(self.device)258 259        # set timesteps260        accepts_offset = "offset" in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys())261        extra_set_kwargs = {}262        if accepts_offset:263            extra_set_kwargs["offset"] = 1264 265        self.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)266 267        # if we use LMSDiscreteScheduler, let's make sure latents are multiplied by sigmas268        if isinstance(self.scheduler, LMSDiscreteScheduler):269            latents = latents * self.scheduler.sigmas[0]270 271        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature272        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.273        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502274        # and should be between [0, 1]275        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())276        extra_step_kwargs = {}277        if accepts_eta:278            extra_step_kwargs["eta"] = eta279 280        for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)):281            # expand the latents if we are doing classifier free guidance282            latent_model_input = (283                torch.cat([latents] * text_embeddings.shape[0]) if do_classifier_free_guidance else latents284            )285            if isinstance(self.scheduler, LMSDiscreteScheduler):286                sigma = self.scheduler.sigmas[i]287                # the model input needs to be scaled to match the continuous ODE formulation in K-LMS288                latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5)289 290            # reduce memory by predicting each score sequentially291            noise_preds = []292            # predict the noise residual293            for latent_in, text_embedding_in in zip(294                torch.chunk(latent_model_input, chunks=latent_model_input.shape[0], dim=0),295                torch.chunk(text_embeddings, chunks=text_embeddings.shape[0], dim=0),296            ):297                noise_preds.append(self.unet(latent_in, t, encoder_hidden_states=text_embedding_in).sample)298            noise_preds = torch.cat(noise_preds, dim=0)299 300            # perform guidance301            if do_classifier_free_guidance:302                noise_pred_uncond = (noise_preds[~mask] * neg_weights).sum(dim=0, keepdims=True)303                noise_pred_text = (noise_preds[mask] * pos_weights).sum(dim=0, keepdims=True)304                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)305 306            # compute the previous noisy sample x_t -> x_t-1307            if isinstance(self.scheduler, LMSDiscreteScheduler):308                latents = self.scheduler.step(noise_pred, i, latents, **extra_step_kwargs).prev_sample309            else:310                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample311 312        # scale and decode the image latents with vae313        latents = 1 / 0.18215 * latents314        image = self.vae.decode(latents).sample315 316        image = (image / 2 + 0.5).clamp(0, 1)317        image = image.cpu().permute(0, 2, 3, 1).numpy()318 319        # run safety checker320        safety_cheker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(self.device)321        image, has_nsfw_concept = self.safety_checker(images=image, clip_input=safety_cheker_input.pixel_values)322 323        if output_type == "pil":324            image = self.numpy_to_pil(image)325 326        if not return_dict:327            return (image, has_nsfw_concept)328 329        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)330