Wauplin/diffusers-community-pipelines-mirror

1import inspect2from typing import Callable, List, Optional, Union3 4import torch5from transformers import (6    CLIPImageProcessor,7    CLIPTextModel,8    CLIPTokenizer,9    WhisperForConditionalGeneration,10    WhisperProcessor,11)12 13from diffusers import (14    AutoencoderKL,15    DDIMScheduler,16    DiffusionPipeline,17    LMSDiscreteScheduler,18    PNDMScheduler,19    UNet2DConditionModel,20)21from diffusers.pipelines.pipeline_utils import StableDiffusionMixin22from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput23from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker24from diffusers.utils import logging25 26 27logger = logging.get_logger(__name__)  # pylint: disable=invalid-name28 29 30class SpeechToImagePipeline(DiffusionPipeline, StableDiffusionMixin):31    def __init__(32        self,33        speech_model: WhisperForConditionalGeneration,34        speech_processor: WhisperProcessor,35        vae: AutoencoderKL,36        text_encoder: CLIPTextModel,37        tokenizer: CLIPTokenizer,38        unet: UNet2DConditionModel,39        scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],40        safety_checker: StableDiffusionSafetyChecker,41        feature_extractor: CLIPImageProcessor,42    ):43        super().__init__()44 45        if safety_checker is None:46            logger.warning(47                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"48                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"49                " results in services or applications open to the public. Both the diffusers team and Hugging Face"50                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"51                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"52                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."53            )54 55        self.register_modules(56            speech_model=speech_model,57            speech_processor=speech_processor,58            vae=vae,59            text_encoder=text_encoder,60            tokenizer=tokenizer,61            unet=unet,62            scheduler=scheduler,63            feature_extractor=feature_extractor,64        )65 66    @torch.no_grad()67    def __call__(68        self,69        audio,70        sampling_rate=16_000,71        height: int = 512,72        width: int = 512,73        num_inference_steps: int = 50,74        guidance_scale: float = 7.5,75        negative_prompt: Optional[Union[str, List[str]]] = None,76        num_images_per_prompt: Optional[int] = 1,77        eta: float = 0.0,78        generator: Optional[torch.Generator] = None,79        latents: Optional[torch.Tensor] = None,80        output_type: Optional[str] = "pil",81        return_dict: bool = True,82        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,83        callback_steps: int = 1,84        **kwargs,85    ):86        inputs = self.speech_processor.feature_extractor(87            audio, return_tensors="pt", sampling_rate=sampling_rate88        ).input_features.to(self.device)89        predicted_ids = self.speech_model.generate(inputs, max_length=480_000)90 91        prompt = self.speech_processor.tokenizer.batch_decode(predicted_ids, skip_special_tokens=True, normalize=True)[92            093        ]94 95        if isinstance(prompt, str):96            batch_size = 197        elif isinstance(prompt, list):98            batch_size = len(prompt)99        else:100            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")101 102        if height % 8 != 0 or width % 8 != 0:103            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")104 105        if (callback_steps is None) or (106            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)107        ):108            raise ValueError(109                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"110                f" {type(callback_steps)}."111            )112 113        # get prompt text embeddings114        text_inputs = self.tokenizer(115            prompt,116            padding="max_length",117            max_length=self.tokenizer.model_max_length,118            return_tensors="pt",119        )120        text_input_ids = text_inputs.input_ids121 122        if text_input_ids.shape[-1] > self.tokenizer.model_max_length:123            removed_text = self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :])124            logger.warning(125                "The following part of your input was truncated because CLIP can only handle sequences up to"126                f" {self.tokenizer.model_max_length} tokens: {removed_text}"127            )128            text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length]129        text_embeddings = self.text_encoder(text_input_ids.to(self.device))[0]130 131        # duplicate text embeddings for each generation per prompt, using mps friendly method132        bs_embed, seq_len, _ = text_embeddings.shape133        text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1)134        text_embeddings = text_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)135 136        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)137        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`138        # corresponds to doing no classifier free guidance.139        do_classifier_free_guidance = guidance_scale > 1.0140        # get unconditional embeddings for classifier free guidance141        if do_classifier_free_guidance:142            uncond_tokens: List[str]143            if negative_prompt is None:144                uncond_tokens = [""] * batch_size145            elif type(prompt) is not type(negative_prompt):146                raise TypeError(147                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="148                    f" {type(prompt)}."149                )150            elif isinstance(negative_prompt, str):151                uncond_tokens = [negative_prompt]152            elif batch_size != len(negative_prompt):153                raise ValueError(154                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"155                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"156                    " the batch size of `prompt`."157                )158            else:159                uncond_tokens = negative_prompt160 161            max_length = text_input_ids.shape[-1]162            uncond_input = self.tokenizer(163                uncond_tokens,164                padding="max_length",165                max_length=max_length,166                truncation=True,167                return_tensors="pt",168            )169            uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]170 171            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method172            seq_len = uncond_embeddings.shape[1]173            uncond_embeddings = uncond_embeddings.repeat(1, num_images_per_prompt, 1)174            uncond_embeddings = uncond_embeddings.view(batch_size * num_images_per_prompt, seq_len, -1)175 176            # For classifier free guidance, we need to do two forward passes.177            # Here we concatenate the unconditional and text embeddings into a single batch178            # to avoid doing two forward passes179            text_embeddings = torch.cat([uncond_embeddings, text_embeddings])180 181        # get the initial random noise unless the user supplied it182 183        # Unlike in other pipelines, latents need to be generated in the target device184        # for 1-to-1 results reproducibility with the CompVis implementation.185        # However this currently doesn't work in `mps`.186        latents_shape = (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)187        latents_dtype = text_embeddings.dtype188        if latents is None:189            if self.device.type == "mps":190                # randn does not exist on mps191                latents = torch.randn(latents_shape, generator=generator, device="cpu", dtype=latents_dtype).to(192                    self.device193                )194            else:195                latents = torch.randn(latents_shape, generator=generator, device=self.device, dtype=latents_dtype)196        else:197            if latents.shape != latents_shape:198                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}")199            latents = latents.to(self.device)200 201        # set timesteps202        self.scheduler.set_timesteps(num_inference_steps)203 204        # Some schedulers like PNDM have timesteps as arrays205        # It's more optimized to move all timesteps to correct device beforehand206        timesteps_tensor = self.scheduler.timesteps.to(self.device)207 208        # scale the initial noise by the standard deviation required by the scheduler209        latents = latents * self.scheduler.init_noise_sigma210 211        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature212        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.213        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502214        # and should be between [0, 1]215        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())216        extra_step_kwargs = {}217        if accepts_eta:218            extra_step_kwargs["eta"] = eta219 220        for i, t in enumerate(self.progress_bar(timesteps_tensor)):221            # expand the latents if we are doing classifier free guidance222            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents223            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)224 225            # predict the noise residual226            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample227 228            # perform guidance229            if do_classifier_free_guidance:230                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)231                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)232 233            # compute the previous noisy sample x_t -> x_t-1234            latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample235 236            # call the callback, if provided237            if callback is not None and i % callback_steps == 0:238                step_idx = i // getattr(self.scheduler, "order", 1)239                callback(step_idx, t, latents)240 241        latents = 1 / 0.18215 * latents242        image = self.vae.decode(latents).sample243 244        image = (image / 2 + 0.5).clamp(0, 1)245 246        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16247        image = image.cpu().permute(0, 2, 3, 1).float().numpy()248 249        if output_type == "pil":250            image = self.numpy_to_pil(image)251 252        if not return_dict:253            return image254 255        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=None)256