


import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()







class 	__lowercase    (__SCREAMING_SNAKE_CASE	,       __SCREAMING_SNAKE_CASE	,       __SCREAMING_SNAKE_CASE	,       unittest.TestCase  ):





							"""simple docstring"""







							_UpperCAmelCase      =						StableUnCLIPImgaImgPipeline
							_UpperCAmelCase      =						TEXT_GUIDED_IMAGE_VARIATION_PARAMS
							_UpperCAmelCase      =						TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
							_UpperCAmelCase      =						frozenset(
							    []  )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
							_UpperCAmelCase      =						frozenset([]  )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[Any]					=					3_2
										SCREAMING_SNAKE_CASE_		:     Any					=					embedder_hidden_size

										# image encoding components

										SCREAMING_SNAKE_CASE_		:     List[Any]					=					CLIPImageProcessor(crop_size=3_2			,						size=3_2       )

										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     int					=					CLIPVisionModelWithProjection(
										    CLIPVisionConfig(
										        hidden_size=lowerCAmelCase__			,						projection_dim=lowerCAmelCase__			,						num_hidden_layers=5			,						num_attention_heads=4			,						image_size=3_2			,						intermediate_size=3_7			,						patch_size=1			,						)       )

										# regular denoising components

										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     str					=					StableUnCLIPImageNormalizer(embedding_dim=lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     int					=					DDPMScheduler(beta_schedule='squaredcos_cap_v2'       )

										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     str					=					CLIPTextModel(
										    CLIPTextConfig(
										        bos_token_id=0			,						eos_token_id=2			,						hidden_size=lowerCAmelCase__			,						projection_dim=3_2			,						intermediate_size=3_7			,						layer_norm_eps=1E-05			,						num_attention_heads=4			,						num_hidden_layers=5			,						pad_token_id=1			,						vocab_size=1_0_0_0			,						)       )

										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     int					=					UNetaDConditionModel(
										    sample_size=3_2			,						in_channels=4			,						out_channels=4			,						down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D')			,						up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D')			,						block_out_channels=(3_2, 6_4)			,						attention_head_dim=(2, 4)			,						class_embed_type='projection'			,						projection_class_embeddings_input_dim=embedder_projection_dim * 2			,						cross_attention_dim=lowerCAmelCase__			,						layers_per_block=1			,						upcast_attention=lowerCAmelCase__			,						use_linear_projection=lowerCAmelCase__			,						)

										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     List[str]					=					DDIMScheduler(
										    beta_schedule='scaled_linear'			,						beta_start=0.00_085			,						beta_end=0.012			,						prediction_type='v_prediction'			,						set_alpha_to_one=lowerCAmelCase__			,						steps_offset=1			,						)

										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     str					=					AutoencoderKL()

										SCREAMING_SNAKE_CASE_		:     str					=					{
										    # image encoding components
										    'feature_extractor': feature_extractor,
										    'image_encoder': image_encoder.eval(),
										    # image noising components
										    'image_normalizer': image_normalizer.eval(),
										    'image_noising_scheduler': image_noising_scheduler,
										    # regular denoising components
										    'tokenizer': tokenizer,
										    'text_encoder': text_encoder.eval(),
										    'unet': unet.eval(),
										    'scheduler': scheduler,
										    'vae': vae.eval(),
										}

										return components






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__=0			,						lowerCAmelCase__=True       ):





										"""simple docstring"""




										if str(lowerCAmelCase__       ).startswith('mps'       ):
													SCREAMING_SNAKE_CASE_		:     List[str]					=					torch.manual_seed(lowerCAmelCase__       )
										else:
													SCREAMING_SNAKE_CASE_		:     Any					=					torch.Generator(device=lowerCAmelCase__       ).manual_seed(lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     List[Any]					=					floats_tensor((1, 3, 3_2, 3_2)			,						rng=random.Random(lowerCAmelCase__       )       ).to(lowerCAmelCase__       )

										if pil_image:
													SCREAMING_SNAKE_CASE_		:     Any					=					input_image * 0.5 + 0.5
													SCREAMING_SNAKE_CASE_		:     Any					=					input_image.clamp(0			,						1       )
													SCREAMING_SNAKE_CASE_		:     List[str]					=					input_image.cpu().permute(0			,						2			,						3			,						1       ).float().numpy()
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					DiffusionPipeline.numpy_to_pil(lowerCAmelCase__       )[0]

										return {
										    "prompt": "An anime racoon running a marathon",
										    "image": input_image,
										    "generator": generator,
										    "num_inference_steps": 2,
										    "output_type": "np",
										}






							@skip_mps
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					'cpu'  # ensure determinism for the device-dependent torch.Generator
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					self.get_dummy_components()
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					StableUnCLIPImgaImgPipeline(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					sd_pipe.to(lowerCAmelCase__       )
										sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					self.get_dummy_inputs(lowerCAmelCase__       )
										inputs.update({'image_embeds': None}       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					sd_pipe(**lowerCAmelCase__       ).images
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					image[0, -3:, -3:, -1]

										assert image.shape == (1, 3_2, 3_2, 3)
										SCREAMING_SNAKE_CASE_		:     List[str]					=					np.array([0.3_872, 0.7_224, 0.5_601, 0.4_741, 0.6_872, 0.5_814, 0.4_636, 0.3_867, 0.5_078]       )

										assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-3






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					torch_device in ['cpu', 'mps']

										self._test_attention_slicing_forward_pass(test_max_difference=lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					torch_device in ['cpu', 'mps']

										self._test_inference_batch_single_identical(test_max_difference=lowerCAmelCase__       )






							@unittest.skipIf(
							    torch_device != 'cuda' or not is_xformers_available()			,						reason='XFormers attention is only available with CUDA and `xformers` installed'			,						)
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										self._test_xformers_attention_forwardGenerator_pass(test_max_difference=lowerCAmelCase__       )







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





							"""simple docstring"""







							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					load_image(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'       )

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					load_numpy(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy'       )

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					StableUnCLIPImgaImgPipeline.from_pretrained(
										    'fusing/stable-unclip-2-1-l-img2img'			,						torch_dtype=torch.floataa       )
										pipe.to(lowerCAmelCase__       )
										pipe.set_progress_bar_config(disable=lowerCAmelCase__       )
										# stable unclip will oom when integration tests are run on a V100,
										# so turn on memory savings
										pipe.enable_attention_slicing()
										pipe.enable_sequential_cpu_offload()

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					torch.Generator(device='cpu'       ).manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					pipe(lowerCAmelCase__			,						'anime turle'			,						generator=lowerCAmelCase__			,						output_type='np'       )

										SCREAMING_SNAKE_CASE_		:     Tuple					=					output.images[0]

										assert image.shape == (7_6_8, 7_6_8, 3)

										assert_mean_pixel_difference(lowerCAmelCase__			,						lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Any					=					load_image(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'       )

										SCREAMING_SNAKE_CASE_		:     Dict					=					load_numpy(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy'       )

										SCREAMING_SNAKE_CASE_		:     int					=					StableUnCLIPImgaImgPipeline.from_pretrained(
										    'fusing/stable-unclip-2-1-h-img2img'			,						torch_dtype=torch.floataa       )
										pipe.to(lowerCAmelCase__       )
										pipe.set_progress_bar_config(disable=lowerCAmelCase__       )
										# stable unclip will oom when integration tests are run on a V100,
										# so turn on memory savings
										pipe.enable_attention_slicing()
										pipe.enable_sequential_cpu_offload()

										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					torch.Generator(device='cpu'       ).manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     int					=					pipe(lowerCAmelCase__			,						'anime turle'			,						generator=lowerCAmelCase__			,						output_type='np'       )

										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					output.images[0]

										assert image.shape == (7_6_8, 7_6_8, 3)

										assert_mean_pixel_difference(lowerCAmelCase__			,						lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Tuple					=					load_image(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'       )

										torch.cuda.empty_cache()
										torch.cuda.reset_max_memory_allocated()
										torch.cuda.reset_peak_memory_stats()

										SCREAMING_SNAKE_CASE_		:     str					=					StableUnCLIPImgaImgPipeline.from_pretrained(
										    'fusing/stable-unclip-2-1-h-img2img'			,						torch_dtype=torch.floataa       )
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					pipe.to(lowerCAmelCase__       )
										pipe.set_progress_bar_config(disable=lowerCAmelCase__       )
										pipe.enable_attention_slicing()
										pipe.enable_sequential_cpu_offload()

										SCREAMING_SNAKE_CASE_		:     List[str]					=					pipe(
										    lowerCAmelCase__			,						'anime turtle'			,						num_inference_steps=2			,						output_type='np'			,						)

										SCREAMING_SNAKE_CASE_		:     List[Any]					=					torch.cuda.max_memory_allocated()
										# make sure that less than 7 GB is allocated
										assert mem_bytes < 7 * 1_0**9












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

import numpy as np

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


SCREAMING_SNAKE_CASE__							:      Optional[int]         =    logging.get_logger(__name__)







class 	__lowerCAmelCase      (       _UpperCamelCase	):
   _UpperCamelCase			:    Tuple  		=					["""input_values""", """attention_mask"""]





   def __init__(			self       ,      snake_case = 1       ,      snake_case = 16_000       ,      snake_case = 0.0       ,      snake_case = False       ,      snake_case = 80       ,      snake_case = 16       ,      snake_case = 64       ,      snake_case = "hann_window"       ,      snake_case = 1.0       ,      snake_case = 80       ,      snake_case = 7_600       ,      snake_case = 1E-10       ,      snake_case = 2       ,      snake_case = True       ,      **snake_case       ,      )					->	Dict:
      """simple docstring"""





      super().__init__(feature_size=snake_case       ,      sampling_rate=snake_case       ,      padding_value=snake_case       ,      **snake_case   )
      a__						:    Any   		=	do_normalize
      a__						:    List[str]   		=	return_attention_mask

      a__						:    List[Any]   		=	num_mel_bins
      a__						:    List[str]   		=	hop_length
      a__						:    int   		=	win_length
      a__						:    List[Any]   		=	win_function
      a__						:    List[str]   		=	frame_signal_scale
      a__						:    List[Any]   		=	fmin
      a__						:    Optional[Any]   		=	fmax
      a__						:    Union[str, Any]   		=	mel_floor
      a__						:    Union[str, Any]   		=	reduction_factor

      a__						:    List[str]   		=	win_length * sampling_rate // 1_000
      a__						:    List[Any]   		=	hop_length * sampling_rate // 1_000
      a__						:    List[Any]   		=	optimal_fft_length(self.sample_size   )
      a__						:    Dict   		=	(self.n_fft // 2) + 1

      a__						:    str   		=	window_function(window_length=self.sample_size       ,      name=self.win_function       ,      periodic=snake_case   )

      a__						:    Tuple   		=	mel_filter_bank(
          num_frequency_bins=self.n_freqs       ,      num_mel_filters=self.num_mel_bins       ,      min_frequency=self.fmin       ,      max_frequency=self.fmax       ,      sampling_rate=self.sampling_rate       ,      norm="slaney"       ,      mel_scale="slaney"       ,      )

      if frame_signal_scale != 1.0:
         warnings.warn(
             "The argument `frame_signal_scale` is deprecated and will be removed in version 4.30.0 of Transformers"       ,      snake_case       ,      )
      if reduction_factor != 2.0:
         warnings.warn(
             "The argument `reduction_factor` is deprecated and will be removed in version 4.30.0 of Transformers"       ,      snake_case       ,      )





   @staticmethod
   # Copied from transformers.models.wav2vec2.feature_extraction_wav2vec2.Wav2Vec2FeatureExtractor.zero_mean_unit_var_norm
   def 						_snake_case			(			snake_case       ,      snake_case       ,      snake_case = 0.0   )					->	List[np.ndarray]:
      """simple docstring"""





      if attention_mask is not None:
         a__						:    Tuple   		=	np.array(snake_case       ,      np.intaa   )
         a__						:    List[str]   		=	[]

         for vector, length in zip(snake_case       ,      attention_mask.sum(-1   )   ):
            a__						:    List[Any]   		=	(vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1E-7   )
            if length < normed_slice.shape[0]:
               a__						:    List[str]   		=	padding_value

            normed_input_values.append(snake_case   )
      else:
         a__						:    Any   		=	[(x - x.mean()) / np.sqrt(x.var() + 1E-7   ) for x in input_values]

      return normed_input_values





   def 						_snake_case			(			self       ,      snake_case       ,      )					->	np.ndarray:
      """simple docstring"""





      a__						:    str   		=	spectrogram(
          snake_case       ,      window=self.window       ,      frame_length=self.sample_size       ,      hop_length=self.sample_stride       ,      fft_length=self.n_fft       ,      mel_filters=self.mel_filters       ,      mel_floor=self.mel_floor       ,      log_mel="log10"       ,      )
      return log_mel_spec.T





   def __call__(			self       ,      snake_case = None       ,      snake_case = None       ,      snake_case = False       ,      snake_case = None       ,      snake_case = False       ,      snake_case = None       ,      snake_case = None       ,      snake_case = None       ,      snake_case = None       ,      **snake_case       ,      )					->	BatchFeature:
      """simple docstring"""





      if audio is None and audio_target is None:
         raise ValueError("You must provide either `audio` or `audio_target` values."   )

      if sampling_rate is not None:
         if sampling_rate != self.sampling_rate:
            raise ValueError(
                F"""The model corresponding to this feature extractor: {self} was trained using a sampling rate of"""
                F""" {self.sampling_rate}. Please make sure that the provided audio input was sampled with"""
                F""" {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."   )

      if audio is not None:
         a__						:    Dict   		=	self._process_audio(
             snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      **snake_case       ,      )
      else:
         a__						:    Optional[int]   		=	None

      if audio_target is not None:
         a__						:    List[Any]   		=	self._process_audio(
             snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      **snake_case       ,      )

         if inputs is None:
            return inputs_target
         else:
            a__						:    Tuple   		=	inputs_target["input_values"]
            a__						:    Tuple   		=	inputs_target.get("attention_mask"   )
            if decoder_attention_mask is not None:
               a__						:    Tuple   		=	decoder_attention_mask

      return inputs





   def 						_snake_case			(			self       ,      snake_case       ,      snake_case = False       ,      snake_case = False       ,      snake_case = None       ,      snake_case = False       ,      snake_case = None       ,      snake_case = None       ,      snake_case = None       ,      **snake_case       ,      )					->	BatchFeature:
      """simple docstring"""





      a__						:    Optional[int]   		=	isinstance(snake_case       ,      np.ndarray   ) and len(speech.shape   ) > 1
      if is_batched_numpy and len(speech.shape   ) > 2:
         raise ValueError(F"""Only mono-channel audio is supported for input to {self}"""   )
      a__						:    List[Any]   		=	is_batched_numpy or (
          isinstance(snake_case       ,      (list, tuple)   ) and (isinstance(speech[0]       ,      (np.ndarray, tuple, list)   ))
      )

      if is_batched:
         a__						:    int   		=	[np.asarray(snake_case       ,      dtype=np.floataa   ) for speech in speech]
      elif not is_batched and not isinstance(snake_case       ,      np.ndarray   ):
         a__						:    Any   		=	np.asarray(snake_case       ,      dtype=np.floataa   )
      elif isinstance(snake_case       ,      np.ndarray   ) and speech.dtype is np.dtype(np.floataa   ):
         a__						:    List[Any]   		=	speech.astype(np.floataa   )

      # always return batch
      if not is_batched:
         a__						:    Union[str, Any]   		=	[speech]

      # needed to make pad() work on spectrogram inputs
      a__						:    Optional[Any]   		=	self.feature_size

      # convert into correct format for padding
      if is_target:
         a__						:    List[str]   		=	[self._extract_mel_features(snake_case   ) for waveform in speech]
         a__						:    Optional[Any]   		=	BatchFeature({"input_values": features}   )
         a__						:    str   		=	self.num_mel_bins
      else:
         a__						:    int   		=	BatchFeature({"input_values": speech}   )

      a__						:    int   		=	self.pad(
          snake_case       ,      padding=snake_case       ,      max_length=snake_case       ,      truncation=snake_case       ,      pad_to_multiple_of=snake_case       ,      return_attention_mask=snake_case       ,      **snake_case       ,      )

      a__						:    Any   		=	feature_size_hack

      # convert input values to correct format
      a__						:    Tuple   		=	padded_inputs["input_values"]
      if not isinstance(input_values[0]       ,      np.ndarray   ):
         a__						:    int   		=	[np.asarray(snake_case       ,      dtype=np.floataa   ) for array in input_values]
      elif (
          not isinstance(snake_case       ,      np.ndarray   )
          and isinstance(input_values[0]       ,      np.ndarray   )
          and input_values[0].dtype is np.dtype(np.floataa   )
      ):
         a__						:    Union[str, Any]   		=	[array.astype(np.floataa   ) for array in input_values]
      elif isinstance(snake_case       ,      np.ndarray   ) and input_values.dtype is np.dtype(np.floataa   ):
         a__						:    Optional[int]   		=	input_values.astype(np.floataa   )

      # convert attention_mask to correct format
      a__						:    Optional[int]   		=	padded_inputs.get("attention_mask"   )
      if attention_mask is not None:
         a__						:    Tuple   		=	[np.asarray(snake_case       ,      dtype=np.intaa   ) for array in attention_mask]

      # zero-mean and unit-variance normalization
      if not is_target and self.do_normalize:
         a__						:    Any   		=	(
             attention_mask
             if self._get_padding_strategies(snake_case       ,      max_length=snake_case   ) is not PaddingStrategy.DO_NOT_PAD
             else None
         )
         a__						:    Optional[Any]   		=	self.zero_mean_unit_var_norm(
             padded_inputs["input_values"]       ,      attention_mask=snake_case       ,      padding_value=self.padding_value   )

      if return_tensors is not None:
         a__						:    int   		=	padded_inputs.convert_to_tensors(snake_case   )

      return padded_inputs





   def 						_snake_case			(			self   )					->	Dict[str, Any]:
      """simple docstring"""





      a__						:    int   		=	super().to_dict()

      # Don't serialize these as they are derived from the other properties.
      a__						:    str   		=	["window", "mel_filters", "sample_size", "sample_stride", "n_fft", "n_freqs"]
      for name in names:
         if name in output:
            del output[name]

      return output

