



import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()







class      __lowerCAmelCase       (     A							,		unittest.TestCase		):
			UpperCamelCase     =	KandinskyPipeline
			UpperCamelCase     =	[
			    '''prompt''',
			    '''image_embeds''',
			    '''negative_image_embeds''',
			]
			UpperCamelCase     =	['''prompt''', '''negative_prompt''', '''image_embeds''', '''negative_image_embeds''']
			UpperCamelCase     =	[
			    '''generator''',
			    '''height''',
			    '''width''',
			    '''latents''',
			    '''guidance_scale''',
			    '''negative_prompt''',
			    '''num_inference_steps''',
			    '''return_dict''',
			    '''guidance_scale''',
			    '''num_images_per_prompt''',
			    '''output_type''',
			    '''return_dict''',
			]
			UpperCamelCase     =	False



			@property
			def        _lowerCamelCase			(		self  :    Union[str, Any])  ->       Tuple:



						"""simple docstring"""






						return 32



			@property
			def        _lowerCamelCase			(		self  :    Optional[Any])  ->       Any:



						"""simple docstring"""






						return 32



			@property
			def        _lowerCamelCase			(		self  :    List[str])  ->       Dict:



						"""simple docstring"""






						return self.time_input_dim



			@property
			def        _lowerCamelCase			(		self  :    Dict)  ->       str:



						"""simple docstring"""






						return self.time_input_dim * 4



			@property
			def        _lowerCamelCase			(		self  :    Any)  ->       str:



						"""simple docstring"""






						return 1_00



			@property
			def        _lowerCamelCase			(		self  :    Dict)  ->       Optional[Any]:



						"""simple docstring"""






						_UpperCAmelCase									=  XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base')
						return tokenizer



			@property
			def        _lowerCamelCase			(		self  :    Union[str, Any])  ->       Optional[int]:



						"""simple docstring"""






						torch.manual_seed(0)
						_UpperCAmelCase									=  MCLIPConfig(
						    numDims=self.cross_attention_dim					,						transformerDimensions=self.text_embedder_hidden_size					,						hidden_size=self.text_embedder_hidden_size					,						intermediate_size=37					,						num_attention_heads=4					,						num_hidden_layers=5					,						vocab_size=10_05					,						)

						_UpperCAmelCase									=  MultilingualCLIP(A)
						_UpperCAmelCase									=  text_encoder.eval()

						return text_encoder



			@property
			def        _lowerCamelCase			(		self  :    List[str])  ->       Union[str, Any]:



						"""simple docstring"""






						torch.manual_seed(0)

						_UpperCAmelCase									=  {
						    'in_channels': 4,
						    # Out channels is double in channels because predicts mean and variance
						    'out_channels': 8,
						    'addition_embed_type': 'text_image',
						    'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
						    'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
						    'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
						    'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
						    'layers_per_block': 1,
						    'encoder_hid_dim': self.text_embedder_hidden_size,
						    'encoder_hid_dim_type': 'text_image_proj',
						    'cross_attention_dim': self.cross_attention_dim,
						    'attention_head_dim': 4,
						    'resnet_time_scale_shift': 'scale_shift',
						    'class_embed_type': None,
						}

						_UpperCAmelCase									=  UNetaDConditionModel(**A)
						return model



			@property
			def        _lowerCamelCase			(		self  :    Union[str, Any])  ->       Optional[Any]:



						"""simple docstring"""






						return {
						    "block_out_channels": [32, 64],
						    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
						    "in_channels": 3,
						    "latent_channels": 4,
						    "layers_per_block": 1,
						    "norm_num_groups": 8,
						    "norm_type": "spatial",
						    "num_vq_embeddings": 12,
						    "out_channels": 3,
						    "up_block_types": [
						        "AttnUpDecoderBlock2D",
						        "UpDecoderBlock2D",
						    ],
						    "vq_embed_dim": 4,
						}



			@property
			def        _lowerCamelCase			(		self  :    Any)  ->       Any:



						"""simple docstring"""






						torch.manual_seed(0)
						_UpperCAmelCase									=  VQModel(**self.dummy_movq_kwargs)
						return model



			def        _lowerCamelCase			(		self  :    int)  ->       Optional[int]:



						"""simple docstring"""






						_UpperCAmelCase									=  self.dummy_text_encoder
						_UpperCAmelCase									=  self.dummy_tokenizer
						_UpperCAmelCase									=  self.dummy_unet
						_UpperCAmelCase									=  self.dummy_movq

						_UpperCAmelCase									=  DDIMScheduler(
						    num_train_timesteps=10_00					,						beta_schedule='linear'					,						beta_start=0.0_0_0_8_5					,						beta_end=0.0_1_2					,						clip_sample=A					,						set_alpha_to_one=A					,						steps_offset=1					,						prediction_type='epsilon'					,						thresholding=A					,						)

						_UpperCAmelCase									=  {
						    'text_encoder': text_encoder,
						    'tokenizer': tokenizer,
						    'unet': unet,
						    'scheduler': scheduler,
						    'movq': movq,
						}
						return components



			def        _lowerCamelCase			(		self  :    List[str]					,						A  :    Tuple					,						A  :    Optional[Any]=0)  ->       List[Any]:



						"""simple docstring"""






						_UpperCAmelCase									=  floats_tensor((1, self.cross_attention_dim)					,						rng=random.Random(A)).to(A)
						_UpperCAmelCase									=  floats_tensor((1, self.cross_attention_dim)					,						rng=random.Random(seed + 1)).to(A)
						if str(A).startswith('mps'):
									_UpperCAmelCase									=  torch.manual_seed(A)
						else:
									_UpperCAmelCase									=  torch.Generator(device=A).manual_seed(A)
						_UpperCAmelCase									=  {
						    'prompt': 'horse',
						    'image_embeds': image_embeds,
						    'negative_image_embeds': negative_image_embeds,
						    'generator': generator,
						    'height': 64,
						    'width': 64,
						    'guidance_scale': 4.0,
						    'num_inference_steps': 2,
						    'output_type': 'np',
						}
						return inputs



			def        _lowerCamelCase			(		self  :    Optional[int])  ->       Optional[int]:



						"""simple docstring"""






						_UpperCAmelCase									=  'cpu'

						_UpperCAmelCase									=  self.get_dummy_components()

						_UpperCAmelCase									=  self.pipeline_class(**A)
						_UpperCAmelCase									=  pipe.to(A)

						pipe.set_progress_bar_config(disable=A)

						_UpperCAmelCase									=  pipe(**self.get_dummy_inputs(A))
						_UpperCAmelCase									=  output.images

						_UpperCAmelCase									=  pipe(
						    **self.get_dummy_inputs(A)					,						return_dict=A					,						)[0]

						_UpperCAmelCase									=  image[0, -3:, -3:, -1]
						_UpperCAmelCase									=  image_from_tuple[0, -3:, -3:, -1]

						assert image.shape == (1, 64, 64, 3)

						_UpperCAmelCase									=  np.array(
						    [0.3_2_8_6_6_3, 1.0, 0.2_3_2_1_6_8_7_3, 1.0, 0.9_2_7_1_7_5_6_4, 0.4_6_3_9_0_4_6, 0.9_6_8_9_4_7_7_7, 0.3_1_7_1_3_3_7_8, 0.6_2_9_3_9_5_3])

						assert (
						    np.abs(image_slice.flatten() - expected_slice).max() < 1E-2
						), F" expected_slice {expected_slice}, but got {image_slice.flatten()}"
						assert (
						    np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1E-2
						), F" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"







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



			def        _lowerCamelCase			(		self  :    Dict)  ->       List[Any]:



						"""simple docstring"""






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



			def        _lowerCamelCase			(		self  :    Any)  ->       int:



						"""simple docstring"""






						_UpperCAmelCase									=  load_numpy(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
						    '/kandinsky/kandinsky_text2img_cat_fp16.npy')

						_UpperCAmelCase									=  KandinskyPriorPipeline.from_pretrained(
						    'kandinsky-community/kandinsky-2-1-prior'					,						torch_dtype=torch.floataa)
						pipe_prior.to(A)

						_UpperCAmelCase									=  KandinskyPipeline.from_pretrained('kandinsky-community/kandinsky-2-1'					,						torch_dtype=torch.floataa)
						_UpperCAmelCase									=  pipeline.to(A)
						pipeline.set_progress_bar_config(disable=A)

						_UpperCAmelCase									=  'red cat, 4k photo'

						_UpperCAmelCase									=  torch.Generator(device='cuda').manual_seed(0)
						_UpperCAmelCase			,	_UpperCAmelCase									=  pipe_prior(
						    A					,						generator=A					,						num_inference_steps=5					,						negative_prompt=''					,						).to_tuple()

						_UpperCAmelCase									=  torch.Generator(device='cuda').manual_seed(0)
						_UpperCAmelCase									=  pipeline(
						    A					,						image_embeds=A					,						negative_image_embeds=A					,						generator=A					,						num_inference_steps=1_00					,						output_type='np'					,						)

						_UpperCAmelCase									=  output.images[0]

						assert image.shape == (5_12, 5_12, 3)

						assert_mean_pixel_difference(A					,						A)





import gc
import unittest

import numpy as np
import torch

from diffusers import DDPMScheduler, MidiProcessor, SpectrogramDiffusionPipeline
from diffusers.pipelines.spectrogram_diffusion import SpectrogramContEncoder, SpectrogramNotesEncoder, TaFilmDecoder
from diffusers.utils import require_torch_gpu, skip_mps, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_note_seq, require_onnxruntime

from ..pipeline_params import TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS, TOKENS_TO_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


UpperCAmelCase__  =			"./tests/fixtures/elise_format0.mid"







@unittest.skip('''The note-seq package currently throws an error on import'''		)
class      __lowerCAmelCase       (     A							,		unittest.TestCase		):
			UpperCamelCase     =	SpectrogramDiffusionPipeline
			UpperCamelCase     =	PipelineTesterMixin.required_optional_params - {
			    '''callback''',
			    '''latents''',
			    '''callback_steps''',
			    '''output_type''',
			    '''num_images_per_prompt''',
			}
			UpperCamelCase     =	False
			UpperCamelCase     =	False
			UpperCamelCase     =	TOKENS_TO_AUDIO_GENERATION_PARAMS
			UpperCamelCase     =	TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS



			def        _lowerCamelCase			(		self  :    str)  ->       Any:



						"""simple docstring"""






						torch.manual_seed(0)
						_UpperCAmelCase									=  SpectrogramNotesEncoder(
						    max_length=20_48					,						vocab_size=15_36					,						d_model=7_68					,						dropout_rate=0.1					,						num_layers=1					,						num_heads=1					,						d_kv=4					,						d_ff=20_48					,						feed_forward_proj='gated-gelu'					,						)

						_UpperCAmelCase									=  SpectrogramContEncoder(
						    input_dims=1_28					,						targets_context_length=2_56					,						d_model=7_68					,						dropout_rate=0.1					,						num_layers=1					,						num_heads=1					,						d_kv=4					,						d_ff=20_48					,						feed_forward_proj='gated-gelu'					,						)

						_UpperCAmelCase									=  TaFilmDecoder(
						    input_dims=1_28					,						targets_length=2_56					,						max_decoder_noise_time=2_0_0_0_0.0					,						d_model=7_68					,						num_layers=1					,						num_heads=1					,						d_kv=4					,						d_ff=20_48					,						dropout_rate=0.1					,						)

						_UpperCAmelCase									=  DDPMScheduler()

						_UpperCAmelCase									=  {
						    'notes_encoder': notes_encoder.eval(),
						    'continuous_encoder': continuous_encoder.eval(),
						    'decoder': decoder.eval(),
						    'scheduler': scheduler,
						    'melgan': None,
						}
						return components



			def        _lowerCamelCase			(		self  :    Optional[Any]					,						A  :    str					,						A  :    Union[str, Any]=0)  ->       Optional[Any]:



						"""simple docstring"""






						if str(A).startswith('mps'):
									_UpperCAmelCase									=  torch.manual_seed(A)
						else:
									_UpperCAmelCase									=  torch.Generator(device=A).manual_seed(A)
						_UpperCAmelCase									=  {
						    'input_tokens': [
						        [11_34, 90, 11_35, 11_33, 10_80, 1_12, 11_32, 10_80, 11_33, 10_79, 1_33, 11_32, 10_79, 11_33, 1] + [0] * 20_33
						    ],
						    'generator': generator,
						    'num_inference_steps': 4,
						    'output_type': 'mel',
						}
						return inputs



			def        _lowerCamelCase			(		self  :    str)  ->       Optional[int]:



						"""simple docstring"""






						_UpperCAmelCase									=  'cpu'  # ensure determinism for the device-dependent torch.Generator
						_UpperCAmelCase									=  self.get_dummy_components()
						_UpperCAmelCase									=  SpectrogramDiffusionPipeline(**A)
						_UpperCAmelCase									=  pipe.to(A)
						pipe.set_progress_bar_config(disable=A)

						_UpperCAmelCase									=  self.get_dummy_inputs(A)
						_UpperCAmelCase									=  pipe(**A)
						_UpperCAmelCase									=  output.audios

						_UpperCAmelCase									=  mel[0, -3:, -3:]

						assert mel_slice.shape == (3, 3)
						_UpperCAmelCase									=  np.array(
						    [-1_1.5_1_2_9_2_5, -4.7_8_8_2_1_5, -0.4_6_1_7_2_9_0_5, -2.0_5_1_7_1_5, -1_0.5_3_9_1_4_7, -1_0.9_7_0_9_6_3, -9.0_9_1_6_3_4, 4.0, 4.0])
						assert np.abs(mel_slice.flatten() - expected_slice).max() < 1E-2



			@skip_mps
			def        _lowerCamelCase			(		self  :    Optional[Any])  ->       Optional[Any]:



						"""simple docstring"""






						return super().test_save_load_local()



			@skip_mps
			def        _lowerCamelCase			(		self  :    Optional[Any])  ->       List[str]:



						"""simple docstring"""






						return super().test_dict_tuple_outputs_equivalent()



			@skip_mps
			def        _lowerCamelCase			(		self  :    Optional[Any])  ->       Dict:



						"""simple docstring"""






						return super().test_save_load_optional_components()



			@skip_mps
			def        _lowerCamelCase			(		self  :    Union[str, Any])  ->       List[str]:



						"""simple docstring"""






						return super().test_attention_slicing_forward_pass()



			def        _lowerCamelCase			(		self  :    List[Any])  ->       Dict:



						"""simple docstring"""






						pass



			def        _lowerCamelCase			(		self  :    str)  ->       Any:



						"""simple docstring"""






						pass



			@skip_mps
			def        _lowerCamelCase			(		self  :    Optional[int])  ->       int:



						"""simple docstring"""






						return super().test_progress_bar()







@slow
@require_torch_gpu
@require_onnxruntime
@require_note_seq
class      __lowerCAmelCase       (     unittest.TestCase		):



			def        _lowerCamelCase			(		self  :    List[str])  ->       Union[str, Any]:



						"""simple docstring"""






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



			def        _lowerCamelCase			(		self  :    int)  ->       Tuple:



						"""simple docstring"""






						_UpperCAmelCase									=  torch_device

						_UpperCAmelCase									=  SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion')
						_UpperCAmelCase									=  pipe.melgan
						_UpperCAmelCase									=  None

						_UpperCAmelCase									=  pipe.to(A)
						pipe.set_progress_bar_config(disable=A)

						def callback(A  :    Dict					,						A  :    Optional[int]):
									# decode mel to audio
									_UpperCAmelCase									=  melgan(input_features=mel_output.astype(np.floataa))[0]
									assert len(audio[0]) == 8_19_20 * (step + 1)
									# simulate that audio is played
									return audio

						_UpperCAmelCase									=  MidiProcessor()
						_UpperCAmelCase									=  processor(A)

						_UpperCAmelCase									=  input_tokens[:3]
						_UpperCAmelCase									=  torch.manual_seed(0)
						pipe(A					,						num_inference_steps=5					,						generator=A					,						callback=A					,						output_type='mel')



			def        _lowerCamelCase			(		self  :    str)  ->       Optional[int]:



						"""simple docstring"""






						_UpperCAmelCase									=  torch_device

						_UpperCAmelCase									=  SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion')
						_UpperCAmelCase									=  pipe.to(A)
						pipe.set_progress_bar_config(disable=A)
						_UpperCAmelCase									=  MidiProcessor()

						_UpperCAmelCase									=  processor(A)
						# just run two denoising loops
						_UpperCAmelCase									=  input_tokens[:2]

						_UpperCAmelCase									=  torch.manual_seed(0)
						_UpperCAmelCase									=  pipe(A					,						num_inference_steps=2					,						generator=A)

						_UpperCAmelCase									=  output.audios[0]

						assert abs(np.abs(A).sum() - 3_6_1_2.8_4_1) < 1E-1



			def        _lowerCamelCase			(		self  :    List[str])  ->       Tuple:



						"""simple docstring"""






						_UpperCAmelCase									=  torch_device

						_UpperCAmelCase									=  SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion')
						_UpperCAmelCase									=  pipe.to(A)
						pipe.set_progress_bar_config(disable=A)

						_UpperCAmelCase									=  MidiProcessor()

						_UpperCAmelCase									=  processor(A)

						# just run 4 denoising loops
						_UpperCAmelCase									=  input_tokens[:4]

						_UpperCAmelCase									=  torch.manual_seed(0)
						_UpperCAmelCase									=  pipe(A					,						num_inference_steps=1_00					,						generator=A)

						_UpperCAmelCase									=  output.audios[0]
						assert abs(np.abs(A).sum() - 9_3_8_9.1_1_1_1) < 5E-2
