diffusers/community-pipelines-mirror

1import inspect2from typing import List, Optional, Tuple, Union3 4import torch5from torch.nn import functional as F6from transformers import CLIPTextModelWithProjection, CLIPTokenizer7from transformers.models.clip.modeling_clip import CLIPTextModelOutput8 9from diffusers import (10    DiffusionPipeline,11    ImagePipelineOutput,12    PriorTransformer,13    UnCLIPScheduler,14    UNet2DConditionModel,15    UNet2DModel,16)17from diffusers.pipelines.unclip import UnCLIPTextProjModel18from diffusers.utils import logging19from diffusers.utils.torch_utils import randn_tensor20 21 22logger = logging.get_logger(__name__)  # pylint: disable=invalid-name23 24 25def slerp(val, low, high):26    """27    Find the interpolation point between the 'low' and 'high' values for the given 'val'. See https://en.wikipedia.org/wiki/Slerp for more details on the topic.28    """29    low_norm = low / torch.norm(low)30    high_norm = high / torch.norm(high)31    omega = torch.acos((low_norm * high_norm))32    so = torch.sin(omega)33    res = (torch.sin((1.0 - val) * omega) / so) * low + (torch.sin(val * omega) / so) * high34    return res35 36 37class UnCLIPTextInterpolationPipeline(DiffusionPipeline):38    """39    Pipeline for prompt-to-prompt interpolation on CLIP text embeddings and using the UnCLIP / Dall-E to decode them to images.40 41    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the42    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)43 44    Args:45        text_encoder ([`CLIPTextModelWithProjection`]):46            Frozen text-encoder.47        tokenizer (`CLIPTokenizer`):48            Tokenizer of class49            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).50        prior ([`PriorTransformer`]):51            The canonical unCLIP prior to approximate the image embedding from the text embedding.52        text_proj ([`UnCLIPTextProjModel`]):53            Utility class to prepare and combine the embeddings before they are passed to the decoder.54        decoder ([`UNet2DConditionModel`]):55            The decoder to invert the image embedding into an image.56        super_res_first ([`UNet2DModel`]):57            Super resolution unet. Used in all but the last step of the super resolution diffusion process.58        super_res_last ([`UNet2DModel`]):59            Super resolution unet. Used in the last step of the super resolution diffusion process.60        prior_scheduler ([`UnCLIPScheduler`]):61            Scheduler used in the prior denoising process. Just a modified DDPMScheduler.62        decoder_scheduler ([`UnCLIPScheduler`]):63            Scheduler used in the decoder denoising process. Just a modified DDPMScheduler.64        super_res_scheduler ([`UnCLIPScheduler`]):65            Scheduler used in the super resolution denoising process. Just a modified DDPMScheduler.66 67    """68 69    prior: PriorTransformer70    decoder: UNet2DConditionModel71    text_proj: UnCLIPTextProjModel72    text_encoder: CLIPTextModelWithProjection73    tokenizer: CLIPTokenizer74    super_res_first: UNet2DModel75    super_res_last: UNet2DModel76 77    prior_scheduler: UnCLIPScheduler78    decoder_scheduler: UnCLIPScheduler79    super_res_scheduler: UnCLIPScheduler80 81    # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.__init__82    def __init__(83        self,84        prior: PriorTransformer,85        decoder: UNet2DConditionModel,86        text_encoder: CLIPTextModelWithProjection,87        tokenizer: CLIPTokenizer,88        text_proj: UnCLIPTextProjModel,89        super_res_first: UNet2DModel,90        super_res_last: UNet2DModel,91        prior_scheduler: UnCLIPScheduler,92        decoder_scheduler: UnCLIPScheduler,93        super_res_scheduler: UnCLIPScheduler,94    ):95        super().__init__()96 97        self.register_modules(98            prior=prior,99            decoder=decoder,100            text_encoder=text_encoder,101            tokenizer=tokenizer,102            text_proj=text_proj,103            super_res_first=super_res_first,104            super_res_last=super_res_last,105            prior_scheduler=prior_scheduler,106            decoder_scheduler=decoder_scheduler,107            super_res_scheduler=super_res_scheduler,108        )109 110    # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents111    def prepare_latents(self, shape, dtype, device, generator, latents, scheduler):112        if latents is None:113            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)114        else:115            if latents.shape != shape:116                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")117            latents = latents.to(device)118 119        latents = latents * scheduler.init_noise_sigma120        return latents121 122    # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline._encode_prompt123    def _encode_prompt(124        self,125        prompt,126        device,127        num_images_per_prompt,128        do_classifier_free_guidance,129        text_model_output: Optional[Union[CLIPTextModelOutput, Tuple]] = None,130        text_attention_mask: Optional[torch.Tensor] = None,131    ):132        if text_model_output is None:133            batch_size = len(prompt) if isinstance(prompt, list) else 1134            # get prompt text embeddings135            text_inputs = self.tokenizer(136                prompt,137                padding="max_length",138                max_length=self.tokenizer.model_max_length,139                truncation=True,140                return_tensors="pt",141            )142            text_input_ids = text_inputs.input_ids143            text_mask = text_inputs.attention_mask.bool().to(device)144 145            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids146 147            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(148                text_input_ids, untruncated_ids149            ):150                removed_text = self.tokenizer.batch_decode(151                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]152                )153                logger.warning(154                    "The following part of your input was truncated because CLIP can only handle sequences up to"155                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"156                )157                text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length]158 159            text_encoder_output = self.text_encoder(text_input_ids.to(device))160 161            prompt_embeds = text_encoder_output.text_embeds162            text_encoder_hidden_states = text_encoder_output.last_hidden_state163 164        else:165            batch_size = text_model_output[0].shape[0]166            prompt_embeds, text_encoder_hidden_states = text_model_output[0], text_model_output[1]167            text_mask = text_attention_mask168 169        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)170        text_encoder_hidden_states = text_encoder_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)171        text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0)172 173        if do_classifier_free_guidance:174            uncond_tokens = [""] * batch_size175 176            uncond_input = self.tokenizer(177                uncond_tokens,178                padding="max_length",179                max_length=self.tokenizer.model_max_length,180                truncation=True,181                return_tensors="pt",182            )183            uncond_text_mask = uncond_input.attention_mask.bool().to(device)184            negative_prompt_embeds_text_encoder_output = self.text_encoder(uncond_input.input_ids.to(device))185 186            negative_prompt_embeds = negative_prompt_embeds_text_encoder_output.text_embeds187            uncond_text_encoder_hidden_states = negative_prompt_embeds_text_encoder_output.last_hidden_state188 189            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method190 191            seq_len = negative_prompt_embeds.shape[1]192            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt)193            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len)194 195            seq_len = uncond_text_encoder_hidden_states.shape[1]196            uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.repeat(1, num_images_per_prompt, 1)197            uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.view(198                batch_size * num_images_per_prompt, seq_len, -1199            )200            uncond_text_mask = uncond_text_mask.repeat_interleave(num_images_per_prompt, dim=0)201 202            # done duplicates203 204            # For classifier free guidance, we need to do two forward passes.205            # Here we concatenate the unconditional and text embeddings into a single batch206            # to avoid doing two forward passes207            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])208            text_encoder_hidden_states = torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states])209 210            text_mask = torch.cat([uncond_text_mask, text_mask])211 212        return prompt_embeds, text_encoder_hidden_states, text_mask213 214    @torch.no_grad()215    def __call__(216        self,217        start_prompt: str,218        end_prompt: str,219        steps: int = 5,220        prior_num_inference_steps: int = 25,221        decoder_num_inference_steps: int = 25,222        super_res_num_inference_steps: int = 7,223        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,224        prior_guidance_scale: float = 4.0,225        decoder_guidance_scale: float = 8.0,226        enable_sequential_cpu_offload=True,227        gpu_id=0,228        output_type: Optional[str] = "pil",229        return_dict: bool = True,230    ):231        """232        Function invoked when calling the pipeline for generation.233 234        Args:235            start_prompt (`str`):236                The prompt to start the image generation interpolation from.237            end_prompt (`str`):238                The prompt to end the image generation interpolation at.239            steps (`int`, *optional*, defaults to 5):240                The number of steps over which to interpolate from start_prompt to end_prompt. The pipeline returns241                the same number of images as this value.242            prior_num_inference_steps (`int`, *optional*, defaults to 25):243                The number of denoising steps for the prior. More denoising steps usually lead to a higher quality244                image at the expense of slower inference.245            decoder_num_inference_steps (`int`, *optional*, defaults to 25):246                The number of denoising steps for the decoder. More denoising steps usually lead to a higher quality247                image at the expense of slower inference.248            super_res_num_inference_steps (`int`, *optional*, defaults to 7):249                The number of denoising steps for super resolution. More denoising steps usually lead to a higher250                quality image at the expense of slower inference.251            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):252                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)253                to make generation deterministic.254            prior_guidance_scale (`float`, *optional*, defaults to 4.0):255                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).256                `guidance_scale` is defined as `w` of equation 2. of [Imagen257                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >258                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,259                usually at the expense of lower image quality.260            decoder_guidance_scale (`float`, *optional*, defaults to 4.0):261                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).262                `guidance_scale` is defined as `w` of equation 2. of [Imagen263                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >264                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,265                usually at the expense of lower image quality.266            output_type (`str`, *optional*, defaults to `"pil"`):267                The output format of the generated image. Choose between268                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.269            enable_sequential_cpu_offload (`bool`, *optional*, defaults to `True`):270                If True, offloads all models to CPU using accelerate, significantly reducing memory usage. When called, the pipeline's271                models have their state dicts saved to CPU and then are moved to a `torch.device('meta') and loaded to GPU only272                when their specific submodule has its `forward` method called.273            gpu_id (`int`, *optional*, defaults to `0`):274                The gpu_id to be passed to enable_sequential_cpu_offload. Only works when enable_sequential_cpu_offload is set to True.275            return_dict (`bool`, *optional*, defaults to `True`):276                Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple.277        """278 279        if not isinstance(start_prompt, str) or not isinstance(end_prompt, str):280            raise ValueError(281                f"`start_prompt` and `end_prompt` should be of type `str` but got {type(start_prompt)} and"282                f" {type(end_prompt)} instead"283            )284 285        if enable_sequential_cpu_offload:286            self.enable_sequential_cpu_offload(gpu_id=gpu_id)287 288        device = self._execution_device289 290        # Turn the prompts into embeddings.291        inputs = self.tokenizer(292            [start_prompt, end_prompt],293            padding="max_length",294            truncation=True,295            max_length=self.tokenizer.model_max_length,296            return_tensors="pt",297        )298        inputs.to(device)299        text_model_output = self.text_encoder(**inputs)300 301        text_attention_mask = torch.max(inputs.attention_mask[0], inputs.attention_mask[1])302        text_attention_mask = torch.cat([text_attention_mask.unsqueeze(0)] * steps).to(device)303 304        # Interpolate from the start to end prompt using slerp and add the generated images to an image output pipeline305        batch_text_embeds = []306        batch_last_hidden_state = []307 308        for interp_val in torch.linspace(0, 1, steps):309            text_embeds = slerp(interp_val, text_model_output.text_embeds[0], text_model_output.text_embeds[1])310            last_hidden_state = slerp(311                interp_val, text_model_output.last_hidden_state[0], text_model_output.last_hidden_state[1]312            )313            batch_text_embeds.append(text_embeds.unsqueeze(0))314            batch_last_hidden_state.append(last_hidden_state.unsqueeze(0))315 316        batch_text_embeds = torch.cat(batch_text_embeds)317        batch_last_hidden_state = torch.cat(batch_last_hidden_state)318 319        text_model_output = CLIPTextModelOutput(320            text_embeds=batch_text_embeds, last_hidden_state=batch_last_hidden_state321        )322 323        batch_size = text_model_output[0].shape[0]324 325        do_classifier_free_guidance = prior_guidance_scale > 1.0 or decoder_guidance_scale > 1.0326 327        prompt_embeds, text_encoder_hidden_states, text_mask = self._encode_prompt(328            prompt=None,329            device=device,330            num_images_per_prompt=1,331            do_classifier_free_guidance=do_classifier_free_guidance,332            text_model_output=text_model_output,333            text_attention_mask=text_attention_mask,334        )335 336        # prior337 338        self.prior_scheduler.set_timesteps(prior_num_inference_steps, device=device)339        prior_timesteps_tensor = self.prior_scheduler.timesteps340 341        embedding_dim = self.prior.config.embedding_dim342 343        prior_latents = self.prepare_latents(344            (batch_size, embedding_dim),345            prompt_embeds.dtype,346            device,347            generator,348            None,349            self.prior_scheduler,350        )351 352        for i, t in enumerate(self.progress_bar(prior_timesteps_tensor)):353            # expand the latents if we are doing classifier free guidance354            latent_model_input = torch.cat([prior_latents] * 2) if do_classifier_free_guidance else prior_latents355 356            predicted_image_embedding = self.prior(357                latent_model_input,358                timestep=t,359                proj_embedding=prompt_embeds,360                encoder_hidden_states=text_encoder_hidden_states,361                attention_mask=text_mask,362            ).predicted_image_embedding363 364            if do_classifier_free_guidance:365                predicted_image_embedding_uncond, predicted_image_embedding_text = predicted_image_embedding.chunk(2)366                predicted_image_embedding = predicted_image_embedding_uncond + prior_guidance_scale * (367                    predicted_image_embedding_text - predicted_image_embedding_uncond368                )369 370            if i + 1 == prior_timesteps_tensor.shape[0]:371                prev_timestep = None372            else:373                prev_timestep = prior_timesteps_tensor[i + 1]374 375            prior_latents = self.prior_scheduler.step(376                predicted_image_embedding,377                timestep=t,378                sample=prior_latents,379                generator=generator,380                prev_timestep=prev_timestep,381            ).prev_sample382 383        prior_latents = self.prior.post_process_latents(prior_latents)384 385        image_embeddings = prior_latents386 387        # done prior388 389        # decoder390 391        text_encoder_hidden_states, additive_clip_time_embeddings = self.text_proj(392            image_embeddings=image_embeddings,393            prompt_embeds=prompt_embeds,394            text_encoder_hidden_states=text_encoder_hidden_states,395            do_classifier_free_guidance=do_classifier_free_guidance,396        )397 398        if device.type == "mps":399            # HACK: MPS: There is a panic when padding bool tensors,400            # so cast to int tensor for the pad and back to bool afterwards401            text_mask = text_mask.type(torch.int)402            decoder_text_mask = F.pad(text_mask, (self.text_proj.clip_extra_context_tokens, 0), value=1)403            decoder_text_mask = decoder_text_mask.type(torch.bool)404        else:405            decoder_text_mask = F.pad(text_mask, (self.text_proj.clip_extra_context_tokens, 0), value=True)406 407        self.decoder_scheduler.set_timesteps(decoder_num_inference_steps, device=device)408        decoder_timesteps_tensor = self.decoder_scheduler.timesteps409 410        num_channels_latents = self.decoder.config.in_channels411        height = self.decoder.config.sample_size412        width = self.decoder.config.sample_size413 414        decoder_latents = self.prepare_latents(415            (batch_size, num_channels_latents, height, width),416            text_encoder_hidden_states.dtype,417            device,418            generator,419            None,420            self.decoder_scheduler,421        )422 423        for i, t in enumerate(self.progress_bar(decoder_timesteps_tensor)):424            # expand the latents if we are doing classifier free guidance425            latent_model_input = torch.cat([decoder_latents] * 2) if do_classifier_free_guidance else decoder_latents426 427            noise_pred = self.decoder(428                sample=latent_model_input,429                timestep=t,430                encoder_hidden_states=text_encoder_hidden_states,431                class_labels=additive_clip_time_embeddings,432                attention_mask=decoder_text_mask,433            ).sample434 435            if do_classifier_free_guidance:436                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)437                noise_pred_uncond, _ = noise_pred_uncond.split(latent_model_input.shape[1], dim=1)438                noise_pred_text, predicted_variance = noise_pred_text.split(latent_model_input.shape[1], dim=1)439                noise_pred = noise_pred_uncond + decoder_guidance_scale * (noise_pred_text - noise_pred_uncond)440                noise_pred = torch.cat([noise_pred, predicted_variance], dim=1)441 442            if i + 1 == decoder_timesteps_tensor.shape[0]:443                prev_timestep = None444            else:445                prev_timestep = decoder_timesteps_tensor[i + 1]446 447            # compute the previous noisy sample x_t -> x_t-1448            decoder_latents = self.decoder_scheduler.step(449                noise_pred, t, decoder_latents, prev_timestep=prev_timestep, generator=generator450            ).prev_sample451 452        decoder_latents = decoder_latents.clamp(-1, 1)453 454        image_small = decoder_latents455 456        # done decoder457 458        # super res459 460        self.super_res_scheduler.set_timesteps(super_res_num_inference_steps, device=device)461        super_res_timesteps_tensor = self.super_res_scheduler.timesteps462 463        channels = self.super_res_first.config.in_channels // 2464        height = self.super_res_first.config.sample_size465        width = self.super_res_first.config.sample_size466 467        super_res_latents = self.prepare_latents(468            (batch_size, channels, height, width),469            image_small.dtype,470            device,471            generator,472            None,473            self.super_res_scheduler,474        )475 476        if device.type == "mps":477            # MPS does not support many interpolations478            image_upscaled = F.interpolate(image_small, size=[height, width])479        else:480            interpolate_antialias = {}481            if "antialias" in inspect.signature(F.interpolate).parameters:482                interpolate_antialias["antialias"] = True483 484            image_upscaled = F.interpolate(485                image_small, size=[height, width], mode="bicubic", align_corners=False, **interpolate_antialias486            )487 488        for i, t in enumerate(self.progress_bar(super_res_timesteps_tensor)):489            # no classifier free guidance490 491            if i == super_res_timesteps_tensor.shape[0] - 1:492                unet = self.super_res_last493            else:494                unet = self.super_res_first495 496            latent_model_input = torch.cat([super_res_latents, image_upscaled], dim=1)497 498            noise_pred = unet(499                sample=latent_model_input,500                timestep=t,501            ).sample502 503            if i + 1 == super_res_timesteps_tensor.shape[0]:504                prev_timestep = None505            else:506                prev_timestep = super_res_timesteps_tensor[i + 1]507 508            # compute the previous noisy sample x_t -> x_t-1509            super_res_latents = self.super_res_scheduler.step(510                noise_pred, t, super_res_latents, prev_timestep=prev_timestep, generator=generator511            ).prev_sample512 513        image = super_res_latents514        # done super res515 516        # post processing517 518        image = image * 0.5 + 0.5519        image = image.clamp(0, 1)520        image = image.cpu().permute(0, 2, 3, 1).float().numpy()521 522        if output_type == "pil":523            image = self.numpy_to_pil(image)524 525        if not return_dict:526            return (image,)527 528        return ImagePipelineOutput(images=image)529