

import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

import diffusers
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
def     UpperCAmelCase_       (						_A    ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__   =   [tensor.shape for tensor in tensor_list]
     return all(shape == shapes[0] for shape in shapes[1:]    )






class 	UpperCAmelCase__		(	A__						,					A__						,					A__						,					unittest.TestCase	):






  """simple docstring"""


  a									= StableDiffusionLatentUpscalePipeline
  a									= TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
      "height",
      "width",
      "cross_attention_kwargs",
      "negative_prompt_embeds",
      "prompt_embeds",
  }
  a									= PipelineTesterMixin.required_optional_params - {"num_images_per_prompt"}
  a									= TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  a									= frozenset(
      []	)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
  a									= frozenset([]	)

  a									= True






  @property
  def 				lowercase_		(							self						:		Optional[Any] )		->							Union[str, Any]:
       SCREAMING_SNAKE_CASE__   =   1
       SCREAMING_SNAKE_CASE__   =   4
       SCREAMING_SNAKE_CASE__   =   (16, 16)

       SCREAMING_SNAKE_CASE__   =   floats_tensor((batch_size, num_channels) + sizes     ,       rng=random.Random(0 ) ).to(__lowerCamelCase )
       return image






  def 				lowercase_		(							self						:		Optional[int] )		->							List[Any]:
       torch.manual_seed(0 )
       SCREAMING_SNAKE_CASE__   =   UNetaDConditionModel(
           act_fn='''gelu'''     ,       attention_head_dim=8     ,       norm_num_groups=__lowerCamelCase     ,       block_out_channels=[32, 32, 64, 64]     ,       time_cond_proj_dim=160     ,       conv_in_kernel=1     ,       conv_out_kernel=1     ,       cross_attention_dim=32     ,       down_block_types=(
               '''KDownBlock2D''',
               '''KCrossAttnDownBlock2D''',
               '''KCrossAttnDownBlock2D''',
               '''KCrossAttnDownBlock2D''',
           )     ,       in_channels=8     ,       mid_block_type=__lowerCamelCase     ,       only_cross_attention=__lowerCamelCase     ,       out_channels=5     ,       resnet_time_scale_shift='''scale_shift'''     ,       time_embedding_type='''fourier'''     ,       timestep_post_act='''gelu'''     ,       up_block_types=('''KCrossAttnUpBlock2D''', '''KCrossAttnUpBlock2D''', '''KCrossAttnUpBlock2D''', '''KUpBlock2D''')     ,       )
       SCREAMING_SNAKE_CASE__   =   AutoencoderKL(
           block_out_channels=[32, 32, 64, 64]     ,       in_channels=3     ,       out_channels=3     ,       down_block_types=[
               '''DownEncoderBlock2D''',
               '''DownEncoderBlock2D''',
               '''DownEncoderBlock2D''',
               '''DownEncoderBlock2D''',
           ]     ,       up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''', '''UpDecoderBlock2D''', '''UpDecoderBlock2D''']     ,       latent_channels=4     ,       )
       SCREAMING_SNAKE_CASE__   =   EulerDiscreteScheduler(prediction_type='''sample''' )
       SCREAMING_SNAKE_CASE__   =   CLIPTextConfig(
           bos_token_id=0     ,       eos_token_id=2     ,       hidden_size=32     ,       intermediate_size=37     ,       layer_norm_eps=1e-05     ,       num_attention_heads=4     ,       num_hidden_layers=5     ,       pad_token_id=1     ,       vocab_size=1000     ,       hidden_act='''quick_gelu'''     ,       projection_dim=512     ,       )
       SCREAMING_SNAKE_CASE__   =   CLIPTextModel(__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )

       SCREAMING_SNAKE_CASE__   =   {
           '''unet''': model.eval(),
           '''vae''': vae.eval(),
           '''scheduler''': scheduler,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
       }

       return components






  def 				lowercase_		(							self						:		List[Any]     ,       __lowerCamelCase						:		List[str]     ,       __lowerCamelCase						:		Union[str, Any]=0 )		->							int:
       if str(__lowerCamelCase ).startswith('''mps''' ):
            SCREAMING_SNAKE_CASE__   =   torch.manual_seed(__lowerCamelCase )
       else:
            SCREAMING_SNAKE_CASE__   =   torch.Generator(device=__lowerCamelCase ).manual_seed(__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   {
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''image''': self.dummy_image.cpu(),
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''output_type''': '''numpy''',
       }
       return inputs






  def 				lowercase_		(							self						:		Any )		->							Dict:
       SCREAMING_SNAKE_CASE__   =   '''cpu'''

       SCREAMING_SNAKE_CASE__   =   self.get_dummy_components()
       SCREAMING_SNAKE_CASE__   =   self.pipeline_class(**__lowerCamelCase )
       pipe.to(__lowerCamelCase )
       pipe.set_progress_bar_config(disable=__lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   self.get_dummy_inputs(__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   pipe(**__lowerCamelCase ).images
       SCREAMING_SNAKE_CASE__   =   image[0, -3:, -3:, -1]

       self.assertEqual(image.shape     ,       (1, 256, 256, 3) )
       SCREAMING_SNAKE_CASE__   =   np.array(
           [0.47222412, 0.41921633, 0.44717434, 0.46874192, 0.42588258, 0.46150726, 0.4677534, 0.45583832, 0.48579055] )
       SCREAMING_SNAKE_CASE__   =   np.abs(image_slice.flatten() - expected_slice ).max()
       self.assertLessEqual(__lowerCamelCase     ,       1e-3 )






  def 				lowercase_		(							self						:		Optional[Any] )		->							Any:
       super().test_attention_slicing_forward_pass(expected_max_diff=7e-3 )






  def 				lowercase_		(							self						:		int )		->							str:
       super().test_cpu_offload_forward_pass(expected_max_diff=3e-3 )






  def 				lowercase_		(							self						:		int )		->							Optional[int]:
       super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3 )






  def 				lowercase_		(							self						:		Any )		->							Optional[Any]:
       super().test_inference_batch_single_identical(expected_max_diff=7e-3 )






  def 				lowercase_		(							self						:		List[Any] )		->							str:
       super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3e-3 )






  def 				lowercase_		(							self						:		Any )		->							int:
       super().test_save_load_local(expected_max_difference=3e-3 )






  def 				lowercase_		(							self						:		List[Any] )		->							Tuple:
       super().test_save_load_optional_components(expected_max_difference=3e-3 )







  def 				lowercase_		(							self						:		Optional[Any] )		->							Tuple:
       SCREAMING_SNAKE_CASE__   =   [
           '''DDIMScheduler''',
           '''DDPMScheduler''',
           '''PNDMScheduler''',
           '''HeunDiscreteScheduler''',
           '''EulerAncestralDiscreteScheduler''',
           '''KDPM2DiscreteScheduler''',
           '''KDPM2AncestralDiscreteScheduler''',
           '''DPMSolverSDEScheduler''',
       ]
       SCREAMING_SNAKE_CASE__   =   self.get_dummy_components()
       SCREAMING_SNAKE_CASE__   =   self.pipeline_class(**__lowerCamelCase )

       # make sure that PNDM does not need warm-up
       pipe.scheduler.register_to_config(skip_prk_steps=__lowerCamelCase )

       pipe.to(__lowerCamelCase )
       pipe.set_progress_bar_config(disable=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   self.get_dummy_inputs(__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   2

       SCREAMING_SNAKE_CASE__   =   []
       for scheduler_enum in KarrasDiffusionSchedulers:
            if scheduler_enum.name in skip_schedulers:
                 # no sigma schedulers are not supported
                 # no schedulers
                 continue

            SCREAMING_SNAKE_CASE__   =   getattr(__lowerCamelCase     ,       scheduler_enum.name )
            SCREAMING_SNAKE_CASE__   =   scheduler_cls.from_config(pipe.scheduler.config )
            SCREAMING_SNAKE_CASE__   =   pipe(**__lowerCamelCase )[0]
            outputs.append(__lowerCamelCase )

       assert check_same_shape(__lowerCamelCase )






@require_torch_gpu
@slow
class 	UpperCAmelCase__		(	unittest.TestCase	):






  """simple docstring"""


  def 				lowercase_		(							self						:		Any )		->							List[Any]:
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






  def 				lowercase_		(							self						:		List[Any] )		->							Tuple:
       SCREAMING_SNAKE_CASE__   =   torch.manual_seed(33 )

       SCREAMING_SNAKE_CASE__   =   StableDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''     ,       torch_dtype=torch.floataa )
       pipe.to('''cuda''' )

       SCREAMING_SNAKE_CASE__   =   StableDiffusionLatentUpscalePipeline.from_pretrained(
           '''stabilityai/sd-x2-latent-upscaler'''     ,       torch_dtype=torch.floataa )
       upscaler.to('''cuda''' )

       SCREAMING_SNAKE_CASE__   =   '''a photo of an astronaut high resolution, unreal engine, ultra realistic'''

       SCREAMING_SNAKE_CASE__   =   pipe(__lowerCamelCase     ,       generator=__lowerCamelCase     ,       output_type='''latent''' ).images

       SCREAMING_SNAKE_CASE__   =   upscaler(
           prompt=__lowerCamelCase     ,       image=__lowerCamelCase     ,       num_inference_steps=20     ,       guidance_scale=0     ,       generator=__lowerCamelCase     ,       output_type='''np'''     ,       ).images[0]

       SCREAMING_SNAKE_CASE__   =   load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/astronaut_1024.npy''' )
       assert np.abs((expected_image - image).mean() ) < 5e-2







  def 				lowercase_		(							self						:		Union[str, Any] )		->							Dict:
       SCREAMING_SNAKE_CASE__   =   torch.manual_seed(33 )

       SCREAMING_SNAKE_CASE__   =   StableDiffusionLatentUpscalePipeline.from_pretrained(
           '''stabilityai/sd-x2-latent-upscaler'''     ,       torch_dtype=torch.floataa )
       upscaler.to('''cuda''' )

       SCREAMING_SNAKE_CASE__   =   '''the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas'''

       SCREAMING_SNAKE_CASE__   =   load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_512.png''' )

       SCREAMING_SNAKE_CASE__   =   upscaler(
           prompt=__lowerCamelCase     ,       image=__lowerCamelCase     ,       num_inference_steps=20     ,       guidance_scale=0     ,       generator=__lowerCamelCase     ,       output_type='''np'''     ,       ).images[0]

       SCREAMING_SNAKE_CASE__   =   load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_1024.npy''' )
       assert np.abs((expected_image - image).max() ) < 5e-2





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

import numpy as np

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					:  List[str]             =      logging.get_logger(__name__)






class 	UpperCAmelCase__		(	A__	):






  """simple docstring"""


  a									= ["input_features"]






  def __init__(							self						:		Dict     ,       __lowerCamelCase						:		Tuple=80     ,       __lowerCamelCase						:		List[Any]=1_6000     ,       __lowerCamelCase						:		Optional[int]=160     ,       __lowerCamelCase						:		List[str]=30     ,       __lowerCamelCase						:		List[Any]=400     ,       __lowerCamelCase						:		Union[str, Any]=0.0     ,       __lowerCamelCase						:		str=False     ,       **__lowerCamelCase						:		List[str]     ,       )		->							Any:
       super().__init__(
           feature_size=__lowerCamelCase     ,       sampling_rate=__lowerCamelCase     ,       padding_value=__lowerCamelCase     ,       return_attention_mask=__lowerCamelCase     ,       **__lowerCamelCase     ,       )
       SCREAMING_SNAKE_CASE__   =   n_fft
       SCREAMING_SNAKE_CASE__   =   hop_length
       SCREAMING_SNAKE_CASE__   =   chunk_length
       SCREAMING_SNAKE_CASE__   =   chunk_length * sampling_rate
       SCREAMING_SNAKE_CASE__   =   self.n_samples // hop_length
       SCREAMING_SNAKE_CASE__   =   sampling_rate
       SCREAMING_SNAKE_CASE__   =   mel_filter_bank(
           num_frequency_bins=1 + n_fft // 2     ,       num_mel_filters=__lowerCamelCase     ,       min_frequency=0.0     ,       max_frequency=8000.0     ,       sampling_rate=__lowerCamelCase     ,       norm='''slaney'''     ,       mel_scale='''slaney'''     ,       )






  def 				lowercase_		(							self						:		int     ,       __lowerCamelCase						:		np.array )		->							np.ndarray:
       SCREAMING_SNAKE_CASE__   =   spectrogram(
           __lowerCamelCase     ,       window_function(self.n_fft     ,       '''hann''' )     ,       frame_length=self.n_fft     ,       hop_length=self.hop_length     ,       power=2.0     ,       mel_filters=self.mel_filters     ,       log_mel='''log10'''     ,       )
       SCREAMING_SNAKE_CASE__   =   log_spec[:, :-1]
       SCREAMING_SNAKE_CASE__   =   np.maximum(__lowerCamelCase     ,       log_spec.max() - 8.0 )
       SCREAMING_SNAKE_CASE__   =   (log_spec + 4.0) / 4.0
       return log_spec






  @staticmethod
  # Copied from transformers.models.wav2vec2.feature_extraction_wav2vec2.Wav2Vec2FeatureExtractor.zero_mean_unit_var_norm
  def 				lowercase_		(							__lowerCamelCase						:		List[np.ndarray]     ,       __lowerCamelCase						:		List[np.ndarray]     ,       __lowerCamelCase						:		float = 0.0 )		->							List[np.ndarray]:
       if attention_mask is not None:
            SCREAMING_SNAKE_CASE__   =   np.array(__lowerCamelCase     ,       np.intaa )
            SCREAMING_SNAKE_CASE__   =   []

            for vector, length in zip(__lowerCamelCase     ,       attention_mask.sum(-1 ) ):
                 SCREAMING_SNAKE_CASE__   =   (vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1e-7 )
                 if length < normed_slice.shape[0]:
                      SCREAMING_SNAKE_CASE__   =   padding_value

                 normed_input_values.append(__lowerCamelCase )
       else:
            SCREAMING_SNAKE_CASE__   =   [(x - x.mean()) / np.sqrt(x.var() + 1e-7 ) for x in input_values]

       return normed_input_values






  def __call__(							self						:		List[str]     ,       __lowerCamelCase						:		Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]     ,       __lowerCamelCase						:		bool = True     ,       __lowerCamelCase						:		Optional[int] = None     ,       __lowerCamelCase						:		Optional[Union[str, TensorType]] = None     ,       __lowerCamelCase						:		Optional[bool] = None     ,       __lowerCamelCase						:		Optional[str] = "max_length"     ,       __lowerCamelCase						:		Optional[int] = None     ,       __lowerCamelCase						:		Optional[int] = None     ,       __lowerCamelCase						:		Optional[bool] = None     ,       **__lowerCamelCase						:		List[str]     ,       )		->							BatchFeature:

       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.''' )

       SCREAMING_SNAKE_CASE__   =   isinstance(__lowerCamelCase     ,       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}''' )
       SCREAMING_SNAKE_CASE__   =   is_batched_numpy or (
           isinstance(__lowerCamelCase     ,       (list, tuple) ) and (isinstance(raw_speech[0]     ,       (np.ndarray, tuple, list) ))
       )

       if is_batched:
            SCREAMING_SNAKE_CASE__   =   [np.asarray([speech]     ,       dtype=np.floataa ).T for speech in raw_speech]
       elif not is_batched and not isinstance(__lowerCamelCase     ,       np.ndarray ):
            SCREAMING_SNAKE_CASE__   =   np.asarray(__lowerCamelCase     ,       dtype=np.floataa )
       elif isinstance(__lowerCamelCase     ,       np.ndarray ) and raw_speech.dtype is np.dtype(np.floataa ):
            SCREAMING_SNAKE_CASE__   =   raw_speech.astype(np.floataa )

       # always return batch
       if not is_batched:
            SCREAMING_SNAKE_CASE__   =   [np.asarray([raw_speech] ).T]

       SCREAMING_SNAKE_CASE__   =   BatchFeature({'''input_features''': raw_speech} )

       # convert into correct format for padding

       SCREAMING_SNAKE_CASE__   =   self.pad(
           __lowerCamelCase     ,       padding=__lowerCamelCase     ,       max_length=max_length if max_length else self.n_samples     ,       truncation=__lowerCamelCase     ,       pad_to_multiple_of=__lowerCamelCase     ,       return_attention_mask=return_attention_mask or do_normalize     ,       )

       # zero-mean and unit-variance normalization
       if do_normalize:
            SCREAMING_SNAKE_CASE__   =   self.zero_mean_unit_var_norm(
                padded_inputs['''input_features''']     ,       attention_mask=padded_inputs['''attention_mask''']     ,       padding_value=self.padding_value     ,       )
            SCREAMING_SNAKE_CASE__   =   np.stack(padded_inputs['''input_features''']     ,       axis=0 )

       # make sure list is in array format
       SCREAMING_SNAKE_CASE__   =   padded_inputs.get('''input_features''' ).transpose(2     ,       0     ,       1 )

       SCREAMING_SNAKE_CASE__   =   [self._np_extract_fbank_features(__lowerCamelCase ) for waveform in input_features[0]]

       if isinstance(input_features[0]     ,       __lowerCamelCase ):
            SCREAMING_SNAKE_CASE__   =   [np.asarray(__lowerCamelCase     ,       dtype=np.floataa ) for feature in input_features]
       else:
            SCREAMING_SNAKE_CASE__   =   input_features

       if return_attention_mask:
            # rescale from sample (48000) to feature (3000)
            SCREAMING_SNAKE_CASE__   =   padded_inputs['''attention_mask'''][:, :: self.hop_length]

       if return_tensors is not None:
            SCREAMING_SNAKE_CASE__   =   padded_inputs.convert_to_tensors(__lowerCamelCase )

       return padded_inputs







  def 				lowercase_		(							self						:		str )		->							Dict[str, Any]:
       SCREAMING_SNAKE_CASE__   =   copy.deepcopy(self.__dict__ )
       SCREAMING_SNAKE_CASE__   =   self.__class__.__name__
       if "mel_filters" in output:
            del output["mel_filters"]
       return output


