
'''simple docstring'''
import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
		from transformers import WhisperFeatureExtractor

if is_torch_available():
		import torch

lowercase__											=random.Random()






def      UpperCamelCase_				(      A__  ,      A__=1.0  ,      A__=None  ,      A__=None						):





	if rng is None:
		a_      				=						global_rng

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

	return values







@require_torch
@require_torchaudio
class   a_ (      unittest.TestCase					):




			def __init__(       self					,   UpperCAmelCase					,   UpperCAmelCase=7					,   UpperCAmelCase=4_00					,   UpperCAmelCase=20_00					,   UpperCAmelCase=10					,   UpperCAmelCase=1_60					,   UpperCAmelCase=8					,   UpperCAmelCase=0.0					,   UpperCAmelCase=40_00					,   UpperCAmelCase=False					,   UpperCAmelCase=True					,   ):
				a_      				=						parent
				a_      				=						batch_size
				a_      				=						min_seq_length
				a_      				=						max_seq_length
				a_      				=						(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
				a_      				=						padding_value
				a_      				=						sampling_rate
				a_      				=						return_attention_mask
				a_      				=						do_normalize
				a_      				=						feature_size
				a_      				=						chunk_length
				a_      				=						hop_length




			def 					lowerCAmelCase__					(       self	):
				return {
				    "feature_size": self.feature_size,
				    "hop_length": self.hop_length,
				    "chunk_length": self.chunk_length,
				    "padding_value": self.padding_value,
				    "sampling_rate": self.sampling_rate,
				    "return_attention_mask": self.return_attention_mask,
				    "do_normalize": self.do_normalize,
				}






			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase=False					,   UpperCAmelCase=False	):




				def _flatten(UpperCAmelCase	):
					return list(itertools.chain(*_lowerCAmelCase	)	)

				if equal_length:
					a_      				=						[floats_list((self.max_seq_length, self.feature_size)	) for _ in range(self.batch_size	)]
				else:
					# make sure that inputs increase in size
					a_      				=						[
					    floats_list((x, self.feature_size)	)
					    for x in range(self.min_seq_length					,   self.max_seq_length					,   self.seq_length_diff	)
					]
				if numpify:
					a_      				=						[np.asarray(_lowerCAmelCase	) for x in speech_inputs]
				return speech_inputs







@require_torch
@require_torchaudio
class   a_ (      __lowercase		,  unittest.TestCase					):
			lowerCamelCase__	:			Any       =    WhisperFeatureExtractor if is_speech_available() else None




			def 					lowerCAmelCase__					(       self	):
				a_      				=						WhisperFeatureExtractionTester(self	)




			def 					lowerCAmelCase__					(       self	):
				a_      				=						self.feature_extraction_class(**self.feat_extract_dict	)

				with tempfile.TemporaryDirectory() as tmpdirname:
					a_      				=						feat_extract_first.save_pretrained(_lowerCAmelCase	)[0]
					check_json_file_has_correct_format(_lowerCAmelCase	)
					a_      				=						self.feature_extraction_class.from_pretrained(_lowerCAmelCase	)

				a_      				=						feat_extract_first.to_dict()
				a_      				=						feat_extract_second.to_dict()
				a_      				=						feat_extract_first.mel_filters
				a_      				=						feat_extract_second.mel_filters
				self.assertTrue(np.allclose(_lowerCAmelCase					,   _lowerCAmelCase	)	)
				self.assertEqual(_lowerCAmelCase					,   _lowerCAmelCase	)




			def 					lowerCAmelCase__					(       self	):
				a_      				=						self.feature_extraction_class(**self.feat_extract_dict	)

				with tempfile.TemporaryDirectory() as tmpdirname:
					a_      				=						os.path.join(_lowerCAmelCase					,   """feat_extract.json"""	)
					feat_extract_first.to_json_file(_lowerCAmelCase	)
					a_      				=						self.feature_extraction_class.from_json_file(_lowerCAmelCase	)

				a_      				=						feat_extract_first.to_dict()
				a_      				=						feat_extract_second.to_dict()
				a_      				=						feat_extract_first.mel_filters
				a_      				=						feat_extract_second.mel_filters
				self.assertTrue(np.allclose(_lowerCAmelCase					,   _lowerCAmelCase	)	)
				self.assertEqual(_lowerCAmelCase					,   _lowerCAmelCase	)




			def 					lowerCAmelCase__					(       self	):
				a_      				=						self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()	)
				# create three inputs of length 800, 1000, and 1200
				a_      				=						[floats_list((1, x)	)[0] for x in range(8_00					,   14_00					,   2_00	)]
				a_      				=						[np.asarray(_lowerCAmelCase	) for speech_input in speech_inputs]

				# Test feature size
				a_      				=						feature_extractor(_lowerCAmelCase					,   padding="""max_length"""					,   return_tensors="""np"""	).input_features
				self.assertTrue(input_features.ndim == 3	)
				self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames	)
				self.assertTrue(input_features.shape[-2] == feature_extractor.feature_size	)

				# Test not batched input
				a_      				=						feature_extractor(speech_inputs[0]					,   return_tensors="""np"""	).input_features
				a_      				=						feature_extractor(np_speech_inputs[0]					,   return_tensors="""np"""	).input_features
				self.assertTrue(np.allclose(_lowerCAmelCase					,   _lowerCAmelCase					,   atol=1e-3	)	)

				# Test batched
				a_      				=						feature_extractor(_lowerCAmelCase					,   return_tensors="""np"""	).input_features
				a_      				=						feature_extractor(_lowerCAmelCase					,   return_tensors="""np"""	).input_features
				for enc_seq_a, enc_seq_a in zip(_lowerCAmelCase					,   _lowerCAmelCase	):
					self.assertTrue(np.allclose(_lowerCAmelCase					,   _lowerCAmelCase					,   atol=1e-3	)	)

				# Test 2-D numpy arrays are batched.
				a_      				=						[floats_list((1, x)	)[0] for x in (8_00, 8_00, 8_00)]
				a_      				=						np.asarray(_lowerCAmelCase	)
				a_      				=						feature_extractor(_lowerCAmelCase					,   return_tensors="""np"""	).input_features
				a_      				=						feature_extractor(_lowerCAmelCase					,   return_tensors="""np"""	).input_features
				for enc_seq_a, enc_seq_a in zip(_lowerCAmelCase					,   _lowerCAmelCase	):
					self.assertTrue(np.allclose(_lowerCAmelCase					,   _lowerCAmelCase					,   atol=1e-3	)	)

				# Test truncation required
				a_      				=						[floats_list((1, x)	)[0] for x in range(2_00					,   (feature_extractor.n_samples + 5_00)					,   2_00	)]
				a_      				=						[np.asarray(_lowerCAmelCase	) for speech_input in speech_inputs]

				a_      				=						[x[: feature_extractor.n_samples] for x in speech_inputs]
				a_      				=						[np.asarray(_lowerCAmelCase	) for speech_input in speech_inputs_truncated]

				a_      				=						feature_extractor(_lowerCAmelCase					,   return_tensors="""np"""	).input_features
				a_      				=						feature_extractor(_lowerCAmelCase					,   return_tensors="""np"""	).input_features
				for enc_seq_a, enc_seq_a in zip(_lowerCAmelCase					,   _lowerCAmelCase	):
					self.assertTrue(np.allclose(_lowerCAmelCase					,   _lowerCAmelCase					,   atol=1e-3	)	)




			def 					lowerCAmelCase__					(       self	):
				import torch

				a_      				=						self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()	)
				a_      				=						np.random.rand(1_00					,   32	).astype(np.floataa	)
				a_      				=						np_speech_inputs.tolist()

				for inputs in [py_speech_inputs, np_speech_inputs]:
					a_      				=						feature_extractor.pad([{"""input_features""": inputs}]					,   return_tensors="""np"""	)
					self.assertTrue(np_processed.input_features.dtype == np.floataa	)
					a_      				=						feature_extractor.pad([{"""input_features""": inputs}]					,   return_tensors="""pt"""	)
					self.assertTrue(pt_processed.input_features.dtype == torch.floataa	)




			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				a_      				=						load_dataset("""hf-internal-testing/librispeech_asr_dummy"""					,   """clean"""					,   split="""validation"""	)
				# automatic decoding with librispeech
				a_      				=						ds.sort("""id"""	).select(range(_lowerCAmelCase	)	)[:num_samples]['''audio''']

				return [x["array"] for x in speech_samples]




			def 					lowerCAmelCase__					(       self	):
				a_      				=						torch.tensor(
				    [
				        0.11_93, -0.09_46, -0.10_98, -0.01_96, 0.02_25, -0.06_90, -0.17_36, 0.09_51,
				        0.09_71, -0.08_17, -0.07_02, 0.01_62, 0.02_60, 0.00_17, -0.01_92, -0.16_78,
				        0.07_09, -0.18_67, -0.06_55, -0.02_74, -0.02_34, -0.18_84, -0.05_16, -0.05_54,
				        -0.02_74, -0.14_25, -0.14_23, 0.08_37, 0.03_77, -0.08_54
				    ]	)
				# fmt: on

				a_      				=						self._load_datasamples(1	)
				a_      				=						WhisperFeatureExtractor()
				a_      				=						feature_extractor(_lowerCAmelCase					,   return_tensors="""pt"""	).input_features
				self.assertEqual(input_features.shape					,   (1, 80, 30_00)	)
				self.assertTrue(torch.allclose(input_features[0, 0, :30]					,   _lowerCAmelCase					,   atol=1e-4	)	)






			def 					lowerCAmelCase__					(       self	):
				a_      				=						self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()	)
				a_      				=						self._load_datasamples(1	)[0]
				a_      				=						((audio - audio.min()) / (audio.max() - audio.min())) * 6_55_35  # Rescale to [0, 65535] to show issue
				a_      				=						feat_extract.zero_mean_unit_var_norm([audio]					,   attention_mask=_lowerCAmelCase	)[0]

				self.assertTrue(np.all(np.mean(_lowerCAmelCase	) < 1e-3	)	)
				self.assertTrue(np.all(np.abs(np.var(_lowerCAmelCase	) - 1	) < 1e-3	)	)







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 				a			(      __lowercase   ,__lowercase   ,__lowercase   ,unittest.TestCase			):
				SCREAMING_SNAKE_CASE__     :   str		      =     StableUnCLIPPipeline
				SCREAMING_SNAKE_CASE__     :   str		      =     TEXT_TO_IMAGE_PARAMS
				SCREAMING_SNAKE_CASE__     :   int		      =     TEXT_TO_IMAGE_BATCH_PARAMS
				SCREAMING_SNAKE_CASE__     :   Dict		      =     TEXT_TO_IMAGE_IMAGE_PARAMS
				SCREAMING_SNAKE_CASE__     :   List[Any]		      =     TEXT_TO_IMAGE_IMAGE_PARAMS

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






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Tuple           =   32
										__SCREAMING_SNAKE_CASE:				List[Any]           =   embedder_hidden_size

										# prior components

										torch.manual_seed(0					)
										__SCREAMING_SNAKE_CASE:				Dict           =   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

										torch.manual_seed(0					)
										__SCREAMING_SNAKE_CASE:				int           =   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=1000			, )					)

										torch.manual_seed(0					)
										__SCREAMING_SNAKE_CASE:				str           =   PriorTransformer(
										    num_attention_heads=2			, attention_head_dim=12			, embedding_dim=_lowerCAmelCase			, num_layers=1			, )

										torch.manual_seed(0					)
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   DDPMScheduler(
										    variance_type='''fixed_small_log'''			, prediction_type='''sample'''			, num_train_timesteps=1000			, clip_sample=_lowerCAmelCase			, clip_sample_range=5.0			, beta_schedule='''squaredcos_cap_v2'''			, )

										# regular denoising components

										torch.manual_seed(0					)
										__SCREAMING_SNAKE_CASE:				List[Any]           =   StableUnCLIPImageNormalizer(embedding_dim=_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Tuple           =   DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''					)

										torch.manual_seed(0					)
										__SCREAMING_SNAKE_CASE:				str           =   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

										torch.manual_seed(0					)
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   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=1000			, )					)

										torch.manual_seed(0					)
										__SCREAMING_SNAKE_CASE:				str           =   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					)
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   DDIMScheduler(
										    beta_schedule='''scaled_linear'''			, beta_start=0.00085			, beta_end=0.012			, prediction_type='''v_prediction'''			, set_alpha_to_one=_lowerCAmelCase			, steps_offset=1			, )

										torch.manual_seed(0					)
										__SCREAMING_SNAKE_CASE:				int           =   AutoencoderKL()

										__SCREAMING_SNAKE_CASE:				Dict           =   {
										    # 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 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase=0					):

										"""simple docstring"""




										if str(_lowerCAmelCase					).startswith('''mps'''					):
																__SCREAMING_SNAKE_CASE:				List[str]           =   torch.manual_seed(_lowerCAmelCase					)
										else:
																__SCREAMING_SNAKE_CASE:				List[str]           =   torch.Generator(device=_lowerCAmelCase					).manual_seed(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				List[str]           =   {
										    '''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 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				str           =   torch_device == '''cpu'''

										self._test_attention_slicing_forward_pass(test_max_difference=_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[str]           =   torch_device in ['''cpu''', '''mps''']

										self._test_inference_batch_single_identical(test_max_difference=_lowerCAmelCase					)





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






				def 			snake_case_    (  self					):

										"""simple docstring"""




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






				def 			snake_case_    (  self					):

										"""simple docstring"""




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

										__SCREAMING_SNAKE_CASE:				Optional[int]           =   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()

										__SCREAMING_SNAKE_CASE:				Tuple           =   torch.Generator(device='''cpu'''					).manual_seed(0					)
										__SCREAMING_SNAKE_CASE:				List[str]           =   pipe('''anime turle'''			, generator=_lowerCAmelCase			, output_type='''np'''					)

										__SCREAMING_SNAKE_CASE:				Optional[int]           =   output.images[0]

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

										assert_mean_pixel_difference(_lowerCAmelCase			, _lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




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

										__SCREAMING_SNAKE_CASE:				int           =   StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''			, 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:				Dict           =   pipe(
										    '''anime turtle'''			, prior_num_inference_steps=2			, num_inference_steps=2			, output_type='''np'''			, )

										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   torch.cuda.max_memory_allocated()
										# make sure that less than 7 GB is allocated
										assert mem_bytes < 7 * 10**9



