

import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()

class     __UpperCamelCase     (       SCREAMING_SNAKE_CASE__  ,  SCREAMING_SNAKE_CASE__  ,  SCREAMING_SNAKE_CASE__  ,  unittest.TestCase    ):
	SCREAMING_SNAKE_CASE__     					=  StableUnCLIPPipeline
	SCREAMING_SNAKE_CASE__     					=  TEXT_TO_IMAGE_PARAMS
	SCREAMING_SNAKE_CASE__     					=  TEXT_TO_IMAGE_BATCH_PARAMS
	SCREAMING_SNAKE_CASE__     					=  TEXT_TO_IMAGE_IMAGE_PARAMS
	SCREAMING_SNAKE_CASE__     					=  TEXT_TO_IMAGE_IMAGE_PARAMS

	# TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
	SCREAMING_SNAKE_CASE__     					=  False




	def    __A       (		self						:    Optional[int]				):


		'''simple docstring'''







		UpperCAmelCase_												=      32
		UpperCAmelCase_												=      embedder_hidden_size

		# prior components

		torch.manual_seed(0				)
		UpperCAmelCase_												=      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				)

		torch.manual_seed(0				)
		UpperCAmelCase_												=      CLIPTextModelWithProjection(
		    CLIPTextConfig(
		        bos_token_id=0		,						eos_token_id=2		,						hidden_size=__lowerCAmelCase		,						projection_dim=__lowerCAmelCase		,						intermediate_size=37		,						layer_norm_eps=1e-05		,						num_attention_heads=4		,						num_hidden_layers=5		,						pad_token_id=1		,						vocab_size=1_000		,						)				)

		torch.manual_seed(0				)
		UpperCAmelCase_												=      PriorTransformer(
		    num_attention_heads=2		,						attention_head_dim=12		,						embedding_dim=__lowerCAmelCase		,						num_layers=1		,						)

		torch.manual_seed(0				)
		UpperCAmelCase_												=      DDPMScheduler(
		    variance_type="fixed_small_log"		,						prediction_type="sample"		,						num_train_timesteps=1_000		,						clip_sample=__lowerCAmelCase		,						clip_sample_range=5.0		,						beta_schedule="squaredcos_cap_v2"		,						)

		# regular denoising components

		torch.manual_seed(0				)
		UpperCAmelCase_												=      StableUnCLIPImageNormalizer(embedding_dim=__lowerCAmelCase				)
		UpperCAmelCase_												=      DDPMScheduler(beta_schedule="squaredcos_cap_v2"				)

		torch.manual_seed(0				)
		UpperCAmelCase_												=      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				)

		torch.manual_seed(0				)
		UpperCAmelCase_												=      CLIPTextModel(
		    CLIPTextConfig(
		        bos_token_id=0		,						eos_token_id=2		,						hidden_size=__lowerCAmelCase		,						projection_dim=32		,						intermediate_size=37		,						layer_norm_eps=1e-05		,						num_attention_heads=4		,						num_hidden_layers=5		,						pad_token_id=1		,						vocab_size=1_000		,						)				)

		torch.manual_seed(0				)
		UpperCAmelCase_												=      UNetaDConditionModel(
		    sample_size=32		,						in_channels=4		,						out_channels=4		,						down_block_types=("CrossAttnDownBlock2D", "DownBlock2D")		,						up_block_types=("UpBlock2D", "CrossAttnUpBlock2D")		,						block_out_channels=(32, 64)		,						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				)
		UpperCAmelCase_												=      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				)
		UpperCAmelCase_												=      AutoencoderKL()

		UpperCAmelCase_												=      {
		    # prior components
		    """prior_tokenizer""": prior_tokenizer,
		    """prior_text_encoder""": prior_text_encoder,
		    """prior""": prior,
		    """prior_scheduler""": prior_scheduler,
		    # image noising components
		    """image_normalizer""": image_normalizer,
		    """image_noising_scheduler""": image_noising_scheduler,
		    # regular denoising components
		    """tokenizer""": tokenizer,
		    """text_encoder""": text_encoder,
		    """unet""": unet,
		    """scheduler""": scheduler,
		    """vae""": vae,
		}

		return components




	def    __A       (		self						:    int		,						lowerCAmelCase						:    Dict		,						lowerCAmelCase						:    int=0				):


		'''simple docstring'''







		if str(__lowerCAmelCase				).startswith("mps"				):
			UpperCAmelCase_												=      torch.manual_seed(__lowerCAmelCase				)
		else:
			UpperCAmelCase_												=      torch.Generator(device=__lowerCAmelCase				).manual_seed(__lowerCAmelCase				)
		UpperCAmelCase_												=      {
		    """prompt""": """A painting of a squirrel eating a burger""",
		    """generator""": generator,
		    """num_inference_steps""": 2,
		    """prior_num_inference_steps""": 2,
		    """output_type""": """numpy""",
		}
		return inputs




	def    __A       (		self						:    Tuple				):


		'''simple docstring'''







		UpperCAmelCase_												=      torch_device == """cpu"""

		self._test_attention_slicing_forward_pass(test_max_difference=__lowerCAmelCase				)




	def    __A       (		self						:    Any				):


		'''simple docstring'''







		UpperCAmelCase_												=      torch_device in ["""cpu""", """mps"""]

		self._test_inference_batch_single_identical(test_max_difference=__lowerCAmelCase				)


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




	def    __A       (		self						:    int				):


		'''simple docstring'''







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




	def    __A       (		self						:    Any				):


		'''simple docstring'''







		UpperCAmelCase_												=      load_numpy(
		    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"				)

		UpperCAmelCase_												=      StableUnCLIPPipeline.from_pretrained("fusing/stable-unclip-2-1-l"		,						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()

		UpperCAmelCase_												=      torch.Generator(device="cpu"				).manual_seed(0				)
		UpperCAmelCase_												=      pipe("anime turle"		,						generator=__lowerCAmelCase		,						output_type="np"				)

		UpperCAmelCase_												=      output.images[0]

		assert image.shape == (768, 768, 3)

		assert_mean_pixel_difference(__lowerCAmelCase		,						__lowerCAmelCase				)




	def    __A       (		self						:    Optional[Any]				):


		'''simple docstring'''







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

		UpperCAmelCase_												=      StableUnCLIPPipeline.from_pretrained("fusing/stable-unclip-2-1-l"		,						torch_dtype=torch.floataa				)
		UpperCAmelCase_												=      pipe.to(__lowerCAmelCase				)
		pipe.set_progress_bar_config(disable=__lowerCAmelCase				)
		pipe.enable_attention_slicing()
		pipe.enable_sequential_cpu_offload()

		UpperCAmelCase_												=      pipe(
		    "anime turtle"		,						prior_num_inference_steps=2		,						num_inference_steps=2		,						output_type="np"		,						)

		UpperCAmelCase_												=      torch.cuda.max_memory_allocated()
		# make sure that less than 7 GB is allocated
		assert mem_bytes < 7 * 10**9







'''simple docstring'''

from typing import List, Optional, Union

import numpy as np
import torch
import torchaudio.compliance.kaldi as ta_kaldi

from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import PaddingStrategy, TensorType, logging


UpperCamelCase      					=logging.get_logger(__name__)




class      A     (   SCREAMING_SNAKE_CASE__    ):






 """simple docstring"""





 __a      :    List[Any]			  =    ['''input_features''', '''attention_mask''']
 def __init__(				self      ,  __lowerCAmelCase=80      ,  __lowerCAmelCase=1_60_00      ,  __lowerCAmelCase=80      ,  __lowerCAmelCase=0.0      ,  __lowerCAmelCase=True      ,  __lowerCAmelCase=True      ,  __lowerCAmelCase=True      ,  **__lowerCAmelCase      ,  ):
       super().__init__(feature_size=__lowerCAmelCase      ,  sampling_rate=__lowerCAmelCase      ,  padding_value=__lowerCAmelCase      ,  **__lowerCAmelCase			)
       UpperCamelCase_					: int            =				num_mel_bins
       UpperCamelCase_					: Union[str, Any]            =				do_ceptral_normalize
       UpperCamelCase_					: Any            =				normalize_means
       UpperCamelCase_					: int            =				normalize_vars
       UpperCamelCase_					: int            =				True
 def 	_UpperCAmelCase							(				self      ,  __lowerCAmelCase      ,  ):
       UpperCamelCase_					: Optional[Any]            =				waveform * (2**15)  # Kaldi compliance: 16-bit signed integers
       UpperCamelCase_					: Optional[int]            =				torch.from_numpy(__lowerCAmelCase			).unsqueeze(0			)
       UpperCamelCase_					: Tuple            =				ta_kaldi.fbank(__lowerCAmelCase      ,  num_mel_bins=self.num_mel_bins      ,  sample_frequency=self.sampling_rate			)
       return features.numpy()
 @staticmethod
 def 	_UpperCAmelCase							(				__lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase = True      ,  __lowerCAmelCase = True      ,  __lowerCAmelCase = 0.0      ,  ):
       # make sure we normalize float32 arrays
       if normalize_means:
             UpperCamelCase_					: int            =				x[:input_length].mean(axis=0			)
             UpperCamelCase_					: str            =				np.subtract(__lowerCAmelCase      ,  __lowerCAmelCase			)
       if normalize_vars:
             UpperCamelCase_					: Optional[int]            =				x[:input_length].std(axis=0			)
             UpperCamelCase_					: Optional[Any]            =				np.divide(__lowerCAmelCase      ,  __lowerCAmelCase			)

       if input_length < x.shape[0]:
             UpperCamelCase_					: Any            =				padding_value

       # make sure array is in float32
       UpperCamelCase_					: str            =				x.astype(np.floataa			)

       return x
 def 	_UpperCAmelCase							(				self      ,  __lowerCAmelCase      ,  __lowerCAmelCase = None			):
       UpperCamelCase_					: List[str]            =				attention_mask.sum(-1			) if attention_mask is not None else [x.shape[0] for x in input_features]
       return [
           self.utterance_cmvn(__lowerCAmelCase      ,  __lowerCAmelCase      ,  self.normalize_means      ,  self.normalize_vars      ,  self.padding_value			)
           for x, n in zip(__lowerCAmelCase      ,  __lowerCAmelCase			)
       ]





 def __call__(				self      ,  __lowerCAmelCase      ,  __lowerCAmelCase = False      ,  __lowerCAmelCase = None      ,  __lowerCAmelCase = False      ,  __lowerCAmelCase = None      ,  __lowerCAmelCase = None      ,  __lowerCAmelCase = None      ,  __lowerCAmelCase = None      ,  **__lowerCAmelCase      ,  ):

       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 `raw_speech` 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."""			)

       UpperCamelCase_					: Optional[int]            =				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}"			)
       UpperCamelCase_					: List[str]            =				is_batched_numpy or (
           isinstance(__lowerCAmelCase      ,  (list, tuple)			) and (isinstance(raw_speech[0]      ,  (np.ndarray, tuple, list)			))
       )

       if is_batched:
             UpperCamelCase_					: str            =				[np.asarray(__lowerCAmelCase      ,  dtype=np.floataa			) for speech in raw_speech]
       elif not is_batched and not isinstance(__lowerCAmelCase      ,  np.ndarray			):
             UpperCamelCase_					: Optional[Any]            =				np.asarray(__lowerCAmelCase      ,  dtype=np.floataa			)
       elif isinstance(__lowerCAmelCase      ,  np.ndarray			) and raw_speech.dtype is np.dtype(np.floataa			):
             UpperCamelCase_					: int            =				raw_speech.astype(np.floataa			)

       # always return batch
       if not is_batched:
             UpperCamelCase_					: Tuple            =				[raw_speech]

       # extract fbank features
       UpperCamelCase_					: Any            =				[self._extract_fbank_features(__lowerCAmelCase			) for waveform in raw_speech]

       # convert into correct format for padding
       UpperCamelCase_					: Tuple            =				BatchFeature({"""input_features""": features}			)

       UpperCamelCase_					: Tuple            =				self.pad(
           __lowerCAmelCase      ,  padding=__lowerCAmelCase      ,  max_length=__lowerCAmelCase      ,  truncation=__lowerCAmelCase      ,  pad_to_multiple_of=__lowerCAmelCase      ,  return_attention_mask=__lowerCAmelCase      ,  **__lowerCAmelCase      ,  )

       # make sure list is in array format
       UpperCamelCase_					: str            =				padded_inputs.get("""input_features"""			)
       if isinstance(input_features[0]      ,  __lowerCAmelCase			):
             UpperCamelCase_					: Any            =				[np.asarray(__lowerCAmelCase      ,  dtype=np.floataa			) for feature in input_features]

       UpperCamelCase_					: Union[str, Any]            =				padded_inputs.get("""attention_mask"""			)
       if attention_mask is not None:
             UpperCamelCase_					: Union[str, Any]            =				[np.asarray(__lowerCAmelCase      ,  dtype=np.intaa			) for array in attention_mask]

       # Utterance-level cepstral mean and variance normalization
       if self.do_ceptral_normalize:
             UpperCamelCase_					: List[str]            =				(
                 np.array(__lowerCAmelCase      ,  dtype=np.intaa			)
                 if self._get_padding_strategies(__lowerCAmelCase      ,  max_length=__lowerCAmelCase			) is not PaddingStrategy.DO_NOT_PAD
                 else None
             )
             UpperCamelCase_					: Optional[int]            =				self.normalize(
                 padded_inputs["""input_features"""]      ,  attention_mask=__lowerCAmelCase			)

       if return_tensors is not None:
             UpperCamelCase_					: Optional[Any]            =				padded_inputs.convert_to_tensors(__lowerCAmelCase			)

       return padded_inputs






