"""simple docstring"""







from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


__SCREAMING_SNAKE_CASE								=logging.get_logger(__name__)  # pylint: disable=invalid-name

__SCREAMING_SNAKE_CASE								='''
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior")
        >>> pipe_prior.to("cuda")
        >>> prompt = "red cat, 4k photo"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> zero_image_emb = out.negative_image_embeds
        >>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder")
        >>> pipe.to("cuda")
        >>> image = pipe(
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=50,
        ... ).images
        >>> image[0].save("cat.png")
        ```
'''



def     lowercase__( __SCREAMING_SNAKE_CASE							:  List[Any]  ,    __SCREAMING_SNAKE_CASE							:  List[Any]  ,    __SCREAMING_SNAKE_CASE							:  Optional[int]=8   ):
     lowercase_			:       str									=			height // scale_factor**2
     if height % scale_factor**2 != 0:
          new_height += 1
     lowercase_			:       int									=			width // scale_factor**2
     if width % scale_factor**2 != 0:
          new_width += 1
     return new_height * scale_factor, new_width * scale_factor







class  UpperCamelCase	(    a__						):
   def __init__(    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,)			->   Tuple:


        '''simple docstring'''


        super().__init__()

        self.register_modules(
            unet=lowerCAmelCase__  ,scheduler=lowerCAmelCase__  ,movq=lowerCAmelCase__  ,)
        lowercase_			:       str									=			2 ** (len(self.movq.config.block_out_channels    ) - 1)
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )			->   Optional[int]:


        '''simple docstring'''


        if latents is None:
             lowercase_			:       List[Any]									=			randn_tensor(lowerCAmelCase__  ,generator=lowerCAmelCase__  ,device=lowerCAmelCase__  ,dtype=lowerCAmelCase__    )
        else:
             if latents.shape != shape:
                  raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {shape}'''    )
             lowercase_			:       str									=			latents.to(lowerCAmelCase__    )

        lowercase_			:       List[Any]									=			latents * scheduler.init_noise_sigma
        return latents
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase=0    )			->   List[str]:


        '''simple docstring'''


        if is_accelerate_available():
             from accelerate import cpu_offload
        else:
             raise ImportError('Please install accelerate via `pip install accelerate`'    )

        lowercase_			:       Optional[int]									=			torch.device(f'''cuda:{gpu_id}'''    )

        lowercase_			:       str									=			[
            self.unet,
            self.movq,
        ]
        for cpu_offloaded_model in models:
             if cpu_offloaded_model is not None:
                  cpu_offload(lowerCAmelCase__  ,lowerCAmelCase__    )
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase=0    )			->   Tuple:


        '''simple docstring'''


        if is_accelerate_available() and is_accelerate_version('>='  ,'0.17.0.dev0'    ):
             from accelerate import cpu_offload_with_hook
        else:
             raise ImportError('`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'    )

        lowercase_			:       Union[str, Any]									=			torch.device(f'''cuda:{gpu_id}'''    )

        if self.device.type != "cpu":
             self.to('cpu'  ,silence_dtype_warnings=lowerCAmelCase__    )
             torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

        lowercase_			:       List[Any]									=			None
        for cpu_offloaded_model in [self.unet, self.movq]:
             lowercase_			:       Any									=			cpu_offload_with_hook(lowerCAmelCase__  ,lowerCAmelCase__  ,prev_module_hook=lowerCAmelCase__    )

        # We'll offload the last model manually.
        lowercase_			:       Union[str, Any]									=			hook
   @property
   # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
   def     _UpperCAmelCase     (    self    )			->   Optional[int]:


        '''simple docstring'''


        if not hasattr(self.unet  ,'_hf_hook'    ):
             return self.device
        for module in self.unet.modules():
             if (
                 hasattr(lowerCAmelCase__  ,'_hf_hook'    )
                 and hasattr(module._hf_hook  ,'execution_device'    )
                 and module._hf_hook.execution_device is not None
             ):
                  return torch.device(module._hf_hook.execution_device    )
        return self.device







   @torch.no_grad()
   @replace_example_docstring(lowerCAmelCase__    )
   def __call__(    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase = 512  ,__UpperCamelCase = 512  ,__UpperCamelCase = 100  ,__UpperCamelCase = 4.0  ,__UpperCamelCase = 1  ,__UpperCamelCase = None  ,__UpperCamelCase = None  ,__UpperCamelCase = "pil"  ,__UpperCamelCase = True  ,)			->   List[str]:


        '''simple docstring'''


        lowercase_			:       Any									=			self._execution_device

        lowercase_			:       List[str]									=			guidance_scale > 1.0

        if isinstance(lowerCAmelCase__  ,lowerCAmelCase__    ):
             lowercase_			:       str									=			torch.cat(lowerCAmelCase__  ,dim=0    )
        lowercase_			:       List[Any]									=			image_embeds.shape[0] * num_images_per_prompt
        if isinstance(lowerCAmelCase__  ,lowerCAmelCase__    ):
             lowercase_			:       List[str]									=			torch.cat(lowerCAmelCase__  ,dim=0    )

        if do_classifier_free_guidance:
             lowercase_			:       Optional[int]									=			image_embeds.repeat_interleave(lowerCAmelCase__  ,dim=0    )
             lowercase_			:       Any									=			negative_image_embeds.repeat_interleave(lowerCAmelCase__  ,dim=0    )

        lowercase_			:       int									=			torch.cat([negative_image_embeds, image_embeds]  ,dim=0    ).to(dtype=self.unet.dtype  ,device=lowerCAmelCase__    )

        self.scheduler.set_timesteps(lowerCAmelCase__  ,device=lowerCAmelCase__    )
        lowercase_			:       Any									=			self.scheduler.timesteps

        lowercase_			:       Any									=			self.unet.config.in_channels

        lowercase_			:       Dict									=			downscale_height_and_width(lowerCAmelCase__  ,lowerCAmelCase__  ,self.movq_scale_factor    )

        # create initial latent
        lowercase_			:       Union[str, Any]									=			self.prepare_latents(
            (batch_size, num_channels_latents, height, width)  ,image_embeds.dtype  ,lowerCAmelCase__  ,lowerCAmelCase__  ,lowerCAmelCase__  ,self.scheduler  ,)

        for i, t in enumerate(self.progress_bar(lowerCAmelCase__    )    ):
             # expand the latents if we are doing classifier free guidance
             lowercase_			:       Optional[Any]									=			torch.cat([latents] * 2    ) if do_classifier_free_guidance else latents

             lowercase_			:       Any									=			{"image_embeds": image_embeds}
             lowercase_			:       List[Any]									=			self.unet(
                 sample=lowerCAmelCase__  ,timestep=lowerCAmelCase__  ,encoder_hidden_states=lowerCAmelCase__  ,added_cond_kwargs=lowerCAmelCase__  ,return_dict=lowerCAmelCase__  ,)[0]

             if do_classifier_free_guidance:
                  lowercase_			:       Tuple									=			noise_pred.split(latents.shape[1]  ,dim=1    )
                  lowercase_			:       str									=			noise_pred.chunk(2    )
                  lowercase_			:       str									=			variance_pred.chunk(2    )
                  lowercase_			:       List[str]									=			noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                  lowercase_			:       Dict									=			torch.cat([noise_pred, variance_pred_text]  ,dim=1    )

             if not (
                 hasattr(self.scheduler.config  ,'variance_type'    )
                 and self.scheduler.config.variance_type in ["learned", "learned_range"]
             ):
                  lowercase_			:       Any									=			noise_pred.split(latents.shape[1]  ,dim=1    )

             # compute the previous noisy sample x_t -> x_t-1
             lowercase_			:       Tuple									=			self.scheduler.step(
                 lowerCAmelCase__  ,lowerCAmelCase__  ,lowerCAmelCase__  ,generator=lowerCAmelCase__  ,)[0]
        # post-processing
        lowercase_			:       Union[str, Any]									=			self.movq.decode(lowerCAmelCase__  ,force_not_quantize=lowerCAmelCase__    )["sample"]

        if output_type not in ["pt", "np", "pil"]:
             raise ValueError(f'''Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}'''    )

        if output_type in ["np", "pil"]:
             lowercase_			:       Tuple									=			image * 0.5 + 0.5
             lowercase_			:       int									=			image.clamp(0  ,1    )
             lowercase_			:       Tuple									=			image.cpu().permute(0  ,2  ,3  ,1    ).float().numpy()

        if output_type == "pil":
             lowercase_			:       Optional[int]									=			self.numpy_to_pil(lowerCAmelCase__    )

        if not return_dict:
             return (image,)

        return ImagePipelineOutput(images=lowerCAmelCase__    )


"""simple docstring"""







import copy
from typing import Any, Dict, List, Optional, Union

import numpy as np
import torch

from ...audio_utils import mel_filter_bank, spectrogram, window_function
from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import TensorType, logging


__SCREAMING_SNAKE_CASE								=logging.get_logger(__name__)







class  UpperCamelCase	(    lowercase_						):
   lowercase							=						['input_features', 'is_longer']
   def __init__(    self  ,__UpperCamelCase=64  ,__UpperCamelCase=4_8000  ,__UpperCamelCase=480  ,__UpperCamelCase=10  ,__UpperCamelCase=1024  ,__UpperCamelCase=0.0  ,__UpperCamelCase=False  ,__UpperCamelCase = 0  ,__UpperCamelCase = 1_4000  ,__UpperCamelCase = None  ,__UpperCamelCase = "fusion"  ,__UpperCamelCase = "repeatpad"  ,**__UpperCamelCase  ,)			->   Optional[int]:


        '''simple docstring'''


        super().__init__(
            feature_size=__UpperCamelCase  ,sampling_rate=__UpperCamelCase  ,padding_value=__UpperCamelCase  ,return_attention_mask=__UpperCamelCase  ,**__UpperCamelCase  ,)
        lowercase_			:       Union[str, Any]									=			top_db
        lowercase_			:       Any									=			truncation
        lowercase_			:       str									=			padding
        lowercase_			:       Optional[Any]									=			fft_window_size
        lowercase_			:       List[Any]									=			(fft_window_size >> 1) + 1
        lowercase_			:       Any									=			hop_length
        lowercase_			:       List[Any]									=			max_length_s
        lowercase_			:       Any									=			max_length_s * sampling_rate
        lowercase_			:       Optional[int]									=			sampling_rate
        lowercase_			:       List[Any]									=			frequency_min
        lowercase_			:       str									=			frequency_max
        lowercase_			:       Any									=			mel_filter_bank(
            num_frequency_bins=self.nb_frequency_bins  ,num_mel_filters=__UpperCamelCase  ,min_frequency=__UpperCamelCase  ,max_frequency=__UpperCamelCase  ,sampling_rate=__UpperCamelCase  ,norm=__UpperCamelCase  ,mel_scale='htk'  ,)
        lowercase_			:       Optional[Any]									=			mel_filter_bank(
            num_frequency_bins=self.nb_frequency_bins  ,num_mel_filters=__UpperCamelCase  ,min_frequency=__UpperCamelCase  ,max_frequency=__UpperCamelCase  ,sampling_rate=__UpperCamelCase  ,norm='slaney'  ,mel_scale='slaney'  ,)
   def     _UpperCAmelCase     (    self    )			->   Dict[str, Any]:


        '''simple docstring'''


        lowercase_			:       int									=			copy.deepcopy(self.__dict__    )
        lowercase_			:       Any									=			self.__class__.__name__
        if "mel_filters" in output:
             del output["mel_filters"]
        if "mel_filters_slaney" in output:
             del output["mel_filters_slaney"]
        return output
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase = None    )			->   np.ndarray:


        '''simple docstring'''


        lowercase_			:       Union[str, Any]									=			spectrogram(
            __UpperCamelCase  ,window_function(self.fft_window_size  ,'hann'    )  ,frame_length=self.fft_window_size  ,hop_length=self.hop_length  ,power=2.0  ,mel_filters=__UpperCamelCase  ,log_mel='dB'  ,)
        return log_mel_spectrogram.T
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )			->   Any:


        '''simple docstring'''


        lowercase_			:       Dict									=			np.array_split(list(range(0  ,total_frames - chunk_frames + 1    )    )  ,3    )
        if len(ranges[1]    ) == 0:
             # if the audio is too short, we just use the first chunk
             lowercase_			:       Optional[int]									=			[0]
        if len(ranges[2]    ) == 0:
             # if the audio is too short, we just use the first chunk
             lowercase_			:       List[str]									=			[0]
        # randomly choose index for each part
        lowercase_			:       List[str]									=			np.random.choice(ranges[0]    )
        lowercase_			:       Optional[Any]									=			np.random.choice(ranges[1]    )
        lowercase_			:       Union[str, Any]									=			np.random.choice(ranges[2]    )

        lowercase_			:       Tuple									=			mel[idx_front : idx_front + chunk_frames, :]
        lowercase_			:       str									=			mel[idx_middle : idx_middle + chunk_frames, :]
        lowercase_			:       Optional[Any]									=			mel[idx_back : idx_back + chunk_frames, :]

        lowercase_			:       Tuple									=			torch.tensor(mel[None, None, :]    )
        lowercase_			:       Optional[Any]									=			torch.nn.functional.interpolate(
            __UpperCamelCase  ,size=[chunk_frames, 64]  ,mode='bilinear'  ,align_corners=__UpperCamelCase    )
        lowercase_			:       str									=			mel_shrink[0][0].numpy()
        lowercase_			:       Optional[int]									=			np.stack([mel_shrink, mel_chunk_front, mel_chunk_middle, mel_chunk_back]  ,axis=0    )
        return mel_fusion
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )			->   np.array:


        '''simple docstring'''


        if waveform.shape[0] > max_length:
             if truncation == "rand_trunc":
                  lowercase_			:       Any									=			True
                  # random crop to max_length (for compatibility) -> this should be handled by self.pad
                  lowercase_			:       Any									=			len(__UpperCamelCase    ) - max_length
                  lowercase_			:       List[str]									=			np.random.randint(0  ,overflow + 1    )
                  lowercase_			:       Any									=			waveform[idx : idx + max_length]
                  lowercase_			:       Optional[int]									=			self._np_extract_fbank_features(__UpperCamelCase  ,self.mel_filters_slaney    )[None, :]
             elif truncation == "fusion":
                  lowercase_			:       str									=			self._np_extract_fbank_features(__UpperCamelCase  ,self.mel_filters    )
                  lowercase_			:       List[str]									=			max_length // self.hop_length + 1  # the +1 related to how the spectrogram is computed
                  lowercase_			:       Dict									=			mel.shape[0]
                  if chunk_frames == total_frames:
                       # there is a corner case where the audio length is larger than max_length but smaller than max_length+hop_length.
                       # In this case, we just use the whole audio.
                       lowercase_			:       int									=			np.stack([mel, mel, mel, mel]  ,axis=0    )
                       lowercase_			:       Optional[Any]									=			False
                  else:
                       lowercase_			:       int									=			self._random_mel_fusion(__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )
                       lowercase_			:       Tuple									=			True
             else:
                  raise NotImplementedError(f'''data_truncating {truncation} not implemented'''    )

        else:
             lowercase_			:       Dict									=			False
             # only use repeat as a new possible value for padding. you repeat the audio before applying the usual max_length padding
             if waveform.shape[0] < max_length:
                  if padding == "repeat":
                       lowercase_			:       Optional[int]									=			int(max_length / len(__UpperCamelCase    )    )
                       lowercase_			:       Any									=			np.stack(np.tile(__UpperCamelCase  ,n_repeat + 1    )    )[:max_length]
                  if padding == "repeatpad":
                       lowercase_			:       Union[str, Any]									=			int(max_length / len(__UpperCamelCase    )    )
                       lowercase_			:       int									=			np.stack(np.tile(__UpperCamelCase  ,__UpperCamelCase    )    )
                  lowercase_			:       str									=			np.pad(__UpperCamelCase  ,(0, max_length - waveform.shape[0])  ,mode='constant'  ,constant_values=0    )

             if truncation == "fusion":
                  lowercase_			:       List[str]									=			self._np_extract_fbank_features(__UpperCamelCase  ,self.mel_filters    )
                  lowercase_			:       List[str]									=			np.stack([input_mel, input_mel, input_mel, input_mel]  ,axis=0    )
             else:
                  lowercase_			:       Any									=			self._np_extract_fbank_features(__UpperCamelCase  ,self.mel_filters_slaney    )[None, :]

        return input_mel, longer







   def __call__(    self  ,__UpperCamelCase  ,__UpperCamelCase = None  ,__UpperCamelCase = None  ,__UpperCamelCase = None  ,__UpperCamelCase = None  ,__UpperCamelCase = None  ,**__UpperCamelCase  ,)			->   BatchFeature:


        '''simple docstring'''


        lowercase_			:       Union[str, Any]									=			truncation if truncation is not None else self.truncation
        lowercase_			:       List[Any]									=			padding if padding else self.padding

        if sampling_rate is not None:
             if sampling_rate != self.sampling_rate:
                  raise ValueError(
                      f'''The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a'''
                      f''' sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input'''
                      f''' was sampled with {self.sampling_rate} and not {sampling_rate}.'''    )
        else:
             logger.warning(
                 'It is strongly recommended to pass the `sampling_rate` argument to this function. '
                 'Failing to do so can result in silent errors that might be hard to debug.'    )

        lowercase_			:       Tuple									=			isinstance(__UpperCamelCase  ,np.ndarray    ) and len(raw_speech.shape    ) > 1
        if is_batched_numpy and len(raw_speech.shape    ) > 2:
             raise ValueError(f'''Only mono-channel audio is supported for input to {self}'''    )
        lowercase_			:       List[str]									=			is_batched_numpy or (
            isinstance(__UpperCamelCase  ,(list, tuple)    ) and (isinstance(raw_speech[0]  ,(np.ndarray, tuple, list)    ))
        )

        if is_batched:
             lowercase_			:       str									=			[np.asarray(__UpperCamelCase  ,dtype=np.floataa    ) for speech in raw_speech]
        elif not is_batched and not isinstance(__UpperCamelCase  ,np.ndarray    ):
             lowercase_			:       Any									=			np.asarray(__UpperCamelCase  ,dtype=np.floataa    )
        elif isinstance(__UpperCamelCase  ,np.ndarray    ) and raw_speech.dtype is np.dtype(np.floataa    ):
             lowercase_			:       Dict									=			raw_speech.astype(np.floataa    )

        # always return batch
        if not is_batched:
             lowercase_			:       int									=			[np.asarray(__UpperCamelCase    )]

        # convert to mel spectrogram, truncate and pad if needed.
        lowercase_			:       List[str]									=			[
            self._get_input_mel(__UpperCamelCase  ,max_length if max_length else self.nb_max_samples  ,__UpperCamelCase  ,__UpperCamelCase    )
            for waveform in raw_speech
        ]

        lowercase_			:       int									=			[]
        lowercase_			:       int									=			[]
        for mel, longer in padded_inputs:
             input_mel.append(__UpperCamelCase    )
             is_longer.append(__UpperCamelCase    )

        if truncation == "fusion" and sum(__UpperCamelCase    ) == 0:
             # if no audio is longer than 10s, then randomly select one audio to be longer
             lowercase_			:       int									=			np.random.randint(0  ,len(__UpperCamelCase    )    )
             lowercase_			:       Tuple									=			True

        if isinstance(input_mel[0]  ,__UpperCamelCase    ):
             lowercase_			:       Union[str, Any]									=			[np.asarray(__UpperCamelCase  ,dtype=np.floataa    ) for feature in input_mel]

        # is_longer is a list of bool
        lowercase_			:       Any									=			[[longer] for longer in is_longer]

        lowercase_			:       List[Any]									=			{'input_features': input_mel, 'is_longer': is_longer}
        lowercase_			:       Union[str, Any]									=			BatchFeature(__UpperCamelCase    )

        if return_tensors is not None:
             lowercase_			:       Any									=			input_features.convert_to_tensors(__UpperCamelCase    )

        return input_features

