
from typing import Callable, List, Optional, Tuple, Union

import torch
from transformers import CLIPTextModel, CLIPTokenizer

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin, TransformeraDModel, VQModel
from ...schedulers import VQDiffusionScheduler
from ...utils import logging
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


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






class   UpperCAmelCase		(						A_       ,A_ ):



   @register_to_config
   def __init__(self		:					Dict    ,						snake_case__		:					bool    ,						snake_case__		:					Optional[int] = None    ,						snake_case__		:					Optional[int] = None )		->	List[Any]:


          '''simple docstring'''


          super().__init__()

          snake_case       :   Optional[int]							=      learnable

          if self.learnable:
                 assert hidden_size is not None, "learnable=True requires `hidden_size` to be set"
                 assert length is not None, "learnable=True requires `length` to be set"

                 snake_case       :   Optional[int]							=      torch.zeros(snake_case__    ,						snake_case__ )
          else:
                 snake_case       :   Any							=      None

          snake_case       :   Optional[Any]							=      torch.nn.Parameter(snake_case__ )






class   UpperCAmelCase		(						A_ ):
   A__					:		VQModel
   A__					:		CLIPTextModel
   A__					:		CLIPTokenizer
   A__					:		TransformeraDModel
   A__					:		LearnedClassifierFreeSamplingEmbeddings
   A__					:		VQDiffusionScheduler



   def __init__(self		:					str    ,						snake_case__		:					VQModel    ,						snake_case__		:					CLIPTextModel    ,						snake_case__		:					CLIPTokenizer    ,						snake_case__		:					TransformeraDModel    ,						snake_case__		:					VQDiffusionScheduler    ,						snake_case__		:					LearnedClassifierFreeSamplingEmbeddings    ,						)		->	List[Any]:


          '''simple docstring'''


          super().__init__()

          self.register_modules(
              vqvae=snake_case__    ,						transformer=snake_case__    ,						text_encoder=snake_case__    ,						tokenizer=snake_case__    ,						scheduler=snake_case__    ,						learned_classifier_free_sampling_embeddings=snake_case__    ,						)



   def 				_SCREAMING_SNAKE_CASE   (self		:					str    ,						snake_case__		:					Any    ,						snake_case__		:					int    ,						snake_case__		:					Union[str, Any] )		->	Union[str, Any]:


          '''simple docstring'''


          snake_case       :   List[str]							=      len(snake_case__ ) if isinstance(snake_case__    ,						snake_case__ ) else 1

          # get prompt text embeddings
          snake_case       :   Optional[Any]							=      self.tokenizer(
              snake_case__    ,						padding="max_length"    ,						max_length=self.tokenizer.model_max_length    ,						return_tensors="pt"    ,						)
          snake_case       :   Dict							=      text_inputs.input_ids

          if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
                 snake_case       :   List[str]							=      self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :] )
                 logger.warning(
                     "The following part of your input was truncated because CLIP can only handle sequences up to"
                     f""" {self.tokenizer.model_max_length} tokens: {removed_text}""" )
                 snake_case       :   str							=      text_input_ids[:, : self.tokenizer.model_max_length]
          snake_case       :   str							=      self.text_encoder(text_input_ids.to(self.device ) )[0]

          # NOTE: This additional step of normalizing the text embeddings is from VQ-Diffusion.
          # While CLIP does normalize the pooled output of the text transformer when combining
          # the image and text embeddings, CLIP does not directly normalize the last hidden state.
          #
          # CLIP normalizing the pooled output.
          # https://github.com/huggingface/transformers/blob/d92e22d1f28324f513f3080e5c47c071a3916721/src/transformers/models/clip/modeling_clip.py#L1052-L1053
          snake_case       :   Dict							=      prompt_embeds / prompt_embeds.norm(dim=-1    ,						keepdim=snake_case__ )

          # duplicate text embeddings for each generation per prompt
          snake_case       :   Dict							=      prompt_embeds.repeat_interleave(snake_case__    ,						dim=0 )

          if do_classifier_free_guidance:
                 if self.learned_classifier_free_sampling_embeddings.learnable:
                        snake_case       :   Optional[int]							=      self.learned_classifier_free_sampling_embeddings.embeddings
                        snake_case       :   Any							=      negative_prompt_embeds.unsqueeze(0 ).repeat(snake_case__    ,						1    ,						1 )
                 else:
                        snake_case       :   Optional[int]							=      [""] * batch_size

                        snake_case       :   int							=      text_input_ids.shape[-1]
                        snake_case       :   Any							=      self.tokenizer(
                            snake_case__    ,						padding="max_length"    ,						max_length=snake_case__    ,						truncation=snake_case__    ,						return_tensors="pt"    ,						)
                        snake_case       :   Union[str, Any]							=      self.text_encoder(uncond_input.input_ids.to(self.device ) )[0]
                        # See comment for normalizing text embeddings
                        snake_case       :   Optional[Any]							=      negative_prompt_embeds / negative_prompt_embeds.norm(dim=-1    ,						keepdim=snake_case__ )

                 # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
                 snake_case       :   Dict							=      negative_prompt_embeds.shape[1]
                 snake_case       :   Any							=      negative_prompt_embeds.repeat(1    ,						snake_case__    ,						1 )
                 snake_case       :   List[Any]							=      negative_prompt_embeds.view(batch_size * num_images_per_prompt    ,						snake_case__    ,						-1 )

                 # For classifier free guidance, we need to do two forward passes.
                 # Here we concatenate the unconditional and text embeddings into a single batch
                 # to avoid doing two forward passes
                 snake_case       :   Tuple							=      torch.cat([negative_prompt_embeds, prompt_embeds] )

          return prompt_embeds



   @torch.no_grad()
   def __call__(self		:					Any    ,						snake_case__		:					Union[str, List[str]]    ,						snake_case__		:					int = 1_00    ,						snake_case__		:					float = 5.0    ,						snake_case__		:					float = 1.0    ,						snake_case__		:					int = 1    ,						snake_case__		:					Optional[Union[torch.Generator, List[torch.Generator]]] = None    ,						snake_case__		:					Optional[torch.FloatTensor] = None    ,						snake_case__		:					Optional[str] = "pil"    ,						snake_case__		:					bool = True    ,						snake_case__		:					Optional[Callable[[int, int, torch.FloatTensor], None]] = None    ,						snake_case__		:					int = 1    ,						)		->	Union[ImagePipelineOutput, Tuple]:


          '''simple docstring'''


          if isinstance(snake_case__    ,						snake_case__ ):
                 snake_case       :   Optional[Any]							=      1
          elif isinstance(snake_case__    ,						snake_case__ ):
                 snake_case       :   int							=      len(snake_case__ )
          else:
                 raise ValueError(f"""`prompt` has to be of type `str` or `list` but is {type(snake_case__ )}""" )

          snake_case       :   List[str]							=      batch_size * num_images_per_prompt

          snake_case       :   Union[str, Any]							=      guidance_scale > 1.0

          snake_case       :   Optional[int]							=      self._encode_prompt(snake_case__    ,						snake_case__    ,						snake_case__ )

          if (callback_steps is None) or (
              callback_steps is not None and (not isinstance(snake_case__    ,						snake_case__ ) or callback_steps <= 0)
          ):
                 raise ValueError(
                     f"""`callback_steps` has to be a positive integer but is {callback_steps} of type"""
                     f""" {type(snake_case__ )}.""" )

          # get the initial completely masked latents unless the user supplied it

          snake_case       :   Tuple							=      (batch_size, self.transformer.num_latent_pixels)
          if latents is None:
                 snake_case       :   Optional[int]							=      self.transformer.num_vector_embeds - 1
                 snake_case       :   Optional[int]							=      torch.full(snake_case__    ,						snake_case__ ).to(self.device )
          else:
                 if latents.shape != latents_shape:
                        raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {latents_shape}""" )
                 if (latents < 0).any() or (latents >= self.transformer.num_vector_embeds).any():
                        raise ValueError(
                            "Unexpected latents value(s). All latents be valid embedding indices i.e. in the range 0,"
                            f""" {self.transformer.num_vector_embeds - 1} (inclusive).""" )
                 snake_case       :   List[str]							=      latents.to(self.device )

          # set timesteps
          self.scheduler.set_timesteps(snake_case__    ,						device=self.device )

          snake_case       :   Optional[int]							=      self.scheduler.timesteps.to(self.device )

          snake_case       :   List[str]							=      latents

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

                 # predict the un-noised image
                 # model_output == `log_p_x_0`
                 snake_case       :   Tuple							=      self.transformer(snake_case__    ,						encoder_hidden_states=snake_case__    ,						timestep=snake_case__ ).sample

                 if do_classifier_free_guidance:
                        snake_case				,					snake_case       :   Optional[Any]							=      model_output.chunk(2 )
                        snake_case       :   List[Any]							=      model_output_uncond + guidance_scale * (model_output_text - model_output_uncond)
                        model_output -= torch.logsumexp(snake_case__    ,						dim=1    ,						keepdim=snake_case__ )

                 snake_case       :   Tuple							=      self.truncate(snake_case__    ,						snake_case__ )

                 # remove `log(0)`'s (`-inf`s)
                 snake_case       :   List[str]							=      model_output.clamp(-70 )

                 # compute the previous noisy sample x_t -> x_t-1
                 snake_case       :   int							=      self.scheduler.step(snake_case__    ,						timestep=snake_case__    ,						sample=snake_case__    ,						generator=snake_case__ ).prev_sample

                 # call the callback, if provided
                 if callback is not None and i % callback_steps == 0:
                        callback(snake_case__    ,						snake_case__    ,						snake_case__ )

          snake_case       :   List[Any]							=      self.vqvae.config.vq_embed_dim
          snake_case       :   List[str]							=      (batch_size, self.transformer.height, self.transformer.width, embedding_channels)
          snake_case       :   List[str]							=      self.vqvae.quantize.get_codebook_entry(snake_case__    ,						shape=snake_case__ )
          snake_case       :   Optional[Any]							=      self.vqvae.decode(snake_case__    ,						force_not_quantize=snake_case__ ).sample

          snake_case       :   Tuple							=      (image / 2 + 0.5).clamp(0    ,						1 )
          snake_case       :   Tuple							=      image.cpu().permute(0    ,						2    ,						3    ,						1 ).numpy()

          if output_type == "pil":
                 snake_case       :   int							=      self.numpy_to_pil(snake_case__ )

          if not return_dict:
                 return (image,)

          return ImagePipelineOutput(images=snake_case__ )



   def 				_SCREAMING_SNAKE_CASE   (self		:					List[Any]    ,						snake_case__		:					torch.FloatTensor    ,						snake_case__		:					float )		->	torch.FloatTensor:


          '''simple docstring'''


          snake_case				,					snake_case       :   Optional[int]							=      torch.sort(snake_case__    ,						1    ,						descending=snake_case__ )
          snake_case       :   List[Any]							=      torch.exp(snake_case__ )
          snake_case       :   List[Any]							=      sorted_p_x_0.cumsum(dim=1 ) < truncation_rate

          # Ensure that at least the largest probability is not zeroed out
          snake_case       :   Optional[Any]							=      torch.full_like(keep_mask[:, 0:1, :]    ,						snake_case__ )
          snake_case       :   List[Any]							=      torch.cat((all_true, keep_mask)    ,						dim=1 )
          snake_case       :   str							=      keep_mask[:, :-1, :]

          snake_case       :   Tuple							=      keep_mask.gather(1    ,						indices.argsort(1 ) )

          snake_case       :   Any							=      log_p_x_0.clone()

          snake_case       :   List[Any]							=      -torch.inf  # -inf = log(0)

          return rv








import itertools
import random
import unittest

import numpy as np

from transformers import WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST, WavaVecaConfig, WavaVecaFeatureExtractor
from transformers.testing_utils import require_torch, slow

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


_UpperCamelCase 			=		random.Random()







def 		_lowercase (     lowercase__   ,	lowercase__=1.0   ,	lowercase__=None   ,	lowercase__=None						):
      if rng is None:
            __lowerCAmelCase					:						Optional[Any]					=					global_rng

      __lowerCAmelCase					:						Tuple					=					[]
      for batch_idx in range(shape[0]						):
            values.append([]						)
            for _ in range(shape[1]						):
                  values[-1].append(rng.random() * scale						)

      return values







class   __lowercase      (unittest.TestCase ):
 def __init__(							self      ,							A_      ,							A_=7      ,							A_=400      ,							A_=2000      ,							A_=1      ,							A_=0.0      ,							A_=1_6000      ,							A_=True      ,							A_=True      ,							)					->List[Any]:



       '''simple docstring'''





       __lowerCAmelCase					:						List[Any]					=					parent
       __lowerCAmelCase					:						Optional[int]					=					batch_size
       __lowerCAmelCase					:						Any					=					min_seq_length
       __lowerCAmelCase					:						Tuple					=					max_seq_length
       __lowerCAmelCase					:						Tuple					=					(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
       __lowerCAmelCase					:						Dict					=					feature_size
       __lowerCAmelCase					:						Optional[int]					=					padding_value
       __lowerCAmelCase					:						Tuple					=					sampling_rate
       __lowerCAmelCase					:						Union[str, Any]					=					return_attention_mask
       __lowerCAmelCase					:						Dict					=					do_normalize
 def 	UpperCamelCase__  (							self					)					->Optional[Any]:



       '''simple docstring'''





       return {
           "feature_size": self.feature_size,
           "padding_value": self.padding_value,
           "sampling_rate": self.sampling_rate,
           "return_attention_mask": self.return_attention_mask,
           "do_normalize": self.do_normalize,
       }





 def 	UpperCamelCase__  (							self      ,							A_=False      ,							A_=False					)					->Union[str, Any]:



       '''simple docstring'''





       def _flatten(A_					):
             return list(itertools.chain(*A_					)					)

       if equal_length:
             __lowerCAmelCase					:						Dict					=					floats_list((self.batch_size, self.max_seq_length)					)
       else:
             # make sure that inputs increase in size
             __lowerCAmelCase					:						Tuple					=					[
                 _flatten(floats_list((x, self.feature_size)					)					)
                 for x in range(self.min_seq_length      ,							self.max_seq_length      ,							self.seq_length_diff					)
             ]

       if numpify:
             __lowerCAmelCase					:						Tuple					=					[np.asarray(A_					) for x in speech_inputs]

       return speech_inputs







class   __lowercase      (_UpperCAmelCase		,   unittest.TestCase ):
 _UpperCamelCase			  =  WavaVecaFeatureExtractor
 def 	UpperCamelCase__  (							self					)					->Dict:



       '''simple docstring'''





       __lowerCAmelCase					:						List[Any]					=					WavaVecaFeatureExtractionTester(self					)
 def 	UpperCamelCase__  (							self      ,							A_					)					->Optional[Any]:



       '''simple docstring'''





       self.assertTrue(np.all(np.mean(A_      ,							axis=0					) < 1e-3					)					)
       self.assertTrue(np.all(np.abs(np.var(A_      ,							axis=0					) - 1					) < 1e-3					)					)
 def 	UpperCamelCase__  (							self					)					->Tuple:



       '''simple docstring'''





       __lowerCAmelCase					:						List[str]					=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()					)
       # create three inputs of length 800, 1000, and 1200
       __lowerCAmelCase					:						Optional[Any]					=					[floats_list((1, x)					)[0] for x in range(800      ,							1400      ,							200					)]
       __lowerCAmelCase					:						Any					=					[np.asarray(A_					) for speech_input in speech_inputs]

       # Test not batched input
       __lowerCAmelCase					:						Optional[Any]					=					feat_extract(speech_inputs[0]      ,							return_tensors='''np'''					).input_values
       __lowerCAmelCase					:						Dict					=					feat_extract(np_speech_inputs[0]      ,							return_tensors='''np'''					).input_values
       self.assertTrue(np.allclose(A_      ,							A_      ,							atol=1e-3					)					)

       # Test batched
       __lowerCAmelCase					:						Dict					=					feat_extract(A_      ,							return_tensors='''np'''					).input_values
       __lowerCAmelCase					:						Dict					=					feat_extract(A_      ,							return_tensors='''np'''					).input_values
       for enc_seq_a, enc_seq_a in zip(A_      ,							A_					):
             self.assertTrue(np.allclose(A_      ,							A_      ,							atol=1e-3					)					)

       # Test 2-D numpy arrays are batched.
       __lowerCAmelCase					:						int					=					[floats_list((1, x)					)[0] for x in (800, 800, 800)]
       __lowerCAmelCase					:						List[Any]					=					np.asarray(A_					)
       __lowerCAmelCase					:						Any					=					feat_extract(A_      ,							return_tensors='''np'''					).input_values
       __lowerCAmelCase					:						Union[str, Any]					=					feat_extract(A_      ,							return_tensors='''np'''					).input_values
       for enc_seq_a, enc_seq_a in zip(A_      ,							A_					):
             self.assertTrue(np.allclose(A_      ,							A_      ,							atol=1e-3					)					)
 def 	UpperCamelCase__  (							self					)					->Union[str, Any]:



       '''simple docstring'''





       __lowerCAmelCase					:						Optional[int]					=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()					)
       __lowerCAmelCase					:						str					=					[floats_list((1, x)					)[0] for x in range(800      ,							1400      ,							200					)]

       __lowerCAmelCase					:						str					=					['''longest''', '''max_length''', '''do_not_pad''']
       __lowerCAmelCase					:						str					=					[None, 1600, None]
       for max_length, padding in zip(A_      ,							A_					):
             __lowerCAmelCase					:						Optional[int]					=					feat_extract(A_      ,							padding=A_      ,							max_length=A_      ,							return_tensors='''np'''					)
             __lowerCAmelCase					:						Optional[Any]					=					processed.input_values

             self._check_zero_mean_unit_variance(input_values[0][:800]					)
             self.assertTrue(input_values[0][800:].sum() < 1e-6					)
             self._check_zero_mean_unit_variance(input_values[1][:1000]					)
             self.assertTrue(input_values[0][1000:].sum() < 1e-6					)
             self._check_zero_mean_unit_variance(input_values[2][:1200]					)
 def 	UpperCamelCase__  (							self					)					->Dict:



       '''simple docstring'''





       __lowerCAmelCase					:						List[str]					=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()					)
       __lowerCAmelCase					:						Optional[int]					=					range(800      ,							1400      ,							200					)
       __lowerCAmelCase					:						Union[str, Any]					=					[floats_list((1, x)					)[0] for x in lengths]

       __lowerCAmelCase					:						int					=					['''longest''', '''max_length''', '''do_not_pad''']
       __lowerCAmelCase					:						List[str]					=					[None, 1600, None]

       for max_length, padding in zip(A_      ,							A_					):
             __lowerCAmelCase					:						Union[str, Any]					=					feat_extract(A_      ,							max_length=A_      ,							padding=A_					)
             __lowerCAmelCase					:						Union[str, Any]					=					processed.input_values

             self._check_zero_mean_unit_variance(input_values[0][:800]					)
             self._check_zero_mean_unit_variance(input_values[1][:1000]					)
             self._check_zero_mean_unit_variance(input_values[2][:1200]					)
 def 	UpperCamelCase__  (							self					)					->Tuple:



       '''simple docstring'''





       __lowerCAmelCase					:						Dict					=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()					)
       __lowerCAmelCase					:						Optional[Any]					=					[floats_list((1, x)					)[0] for x in range(800      ,							1400      ,							200					)]
       __lowerCAmelCase					:						List[str]					=					feat_extract(
           A_      ,							truncation=A_      ,							max_length=1000      ,							padding='''max_length'''      ,							return_tensors='''np'''					)
       __lowerCAmelCase					:						int					=					processed.input_values

       self._check_zero_mean_unit_variance(input_values[0, :800]					)
       self._check_zero_mean_unit_variance(input_values[1]					)
       self._check_zero_mean_unit_variance(input_values[2]					)
 def 	UpperCamelCase__  (							self					)					->List[str]:



       '''simple docstring'''





       __lowerCAmelCase					:						Dict					=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()					)
       __lowerCAmelCase					:						int					=					[floats_list((1, x)					)[0] for x in range(800      ,							1400      ,							200					)]
       __lowerCAmelCase					:						int					=					feat_extract(
           A_      ,							truncation=A_      ,							max_length=1000      ,							padding='''longest'''      ,							return_tensors='''np'''					)
       __lowerCAmelCase					:						Optional[Any]					=					processed.input_values

       self._check_zero_mean_unit_variance(input_values[0, :800]					)
       self._check_zero_mean_unit_variance(input_values[1, :1000]					)
       self._check_zero_mean_unit_variance(input_values[2]					)

       # make sure that if max_length < longest -> then pad to max_length
       self.assertTrue(input_values.shape == (3, 1000)					)

       __lowerCAmelCase					:						Any					=					[floats_list((1, x)					)[0] for x in range(800      ,							1400      ,							200					)]
       __lowerCAmelCase					:						Optional[int]					=					feat_extract(
           A_      ,							truncation=A_      ,							max_length=2000      ,							padding='''longest'''      ,							return_tensors='''np'''					)
       __lowerCAmelCase					:						List[str]					=					processed.input_values

       self._check_zero_mean_unit_variance(input_values[0, :800]					)
       self._check_zero_mean_unit_variance(input_values[1, :1000]					)
       self._check_zero_mean_unit_variance(input_values[2]					)

       # make sure that if max_length > longest -> then pad to longest
       self.assertTrue(input_values.shape == (3, 1200)					)
 @require_torch
 def 	UpperCamelCase__  (							self					)					->Any:



       '''simple docstring'''





       import torch

       __lowerCAmelCase					:						Union[str, Any]					=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()					)
       __lowerCAmelCase					:						Any					=					np.random.rand(100					).astype(np.floataa					)
       __lowerCAmelCase					:						List[Any]					=					np_speech_inputs.tolist()

       for inputs in [py_speech_inputs, np_speech_inputs]:
             __lowerCAmelCase					:						Any					=					feature_extractor.pad([{'''input_values''': inputs}]      ,							return_tensors='''np'''					)
             self.assertTrue(np_processed.input_values.dtype == np.floataa					)
             __lowerCAmelCase					:						List[str]					=					feature_extractor.pad([{'''input_values''': inputs}]      ,							return_tensors='''pt'''					)
             self.assertTrue(pt_processed.input_values.dtype == torch.floataa					)





 @slow
 @require_torch
 def 	UpperCamelCase__  (							self					)					->int:



       '''simple docstring'''





       for model_id in WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST:
             __lowerCAmelCase					:						Any					=					WavaVecaConfig.from_pretrained(A_					)
             __lowerCAmelCase					:						Tuple					=					WavaVecaFeatureExtractor.from_pretrained(A_					)

             # only "layer" feature extraction norm should make use of
             # attention_mask
             self.assertEqual(feat_extract.return_attention_mask      ,							config.feat_extract_norm == '''layer'''					)


