
import gc
import unittest

import numpy as np
import torch

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS, UNCONDITIONAL_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class 					_lowerCamelCase			(   UpperCamelCase_						,				unittest.TestCase			):





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		 =	DanceDiffusionPipeline
     SCREAMING_SNAKE_CASE_		 =	UNCONDITIONAL_AUDIO_GENERATION_PARAMS
     SCREAMING_SNAKE_CASE_		 =	PipelineTesterMixin.required_optional_params - {
         '''callback''',
         '''latents''',
         '''callback_steps''',
         '''output_type''',
         '''num_images_per_prompt''',
     }
     SCREAMING_SNAKE_CASE_		 =	UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
     SCREAMING_SNAKE_CASE_		 =	False
     SCREAMING_SNAKE_CASE_		 =	False






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Any:




           """simple docstring"""







           torch.manual_seed(0	)
           UpperCamelCase__     : List[str]											=					UNetaDModel(
               block_out_channels=(3_2, 3_2, 6_4)    ,    extra_in_channels=1_6    ,    sample_size=5_1_2    ,    sample_rate=1_6_0_0_0    ,    in_channels=2    ,    out_channels=2    ,    flip_sin_to_cos=__SCREAMING_SNAKE_CASE    ,    use_timestep_embedding=__SCREAMING_SNAKE_CASE    ,    time_embedding_type='''fourier'''    ,    mid_block_type='''UNetMidBlock1D'''    ,    down_block_types=('''DownBlock1DNoSkip''', '''DownBlock1D''', '''AttnDownBlock1D''')    ,    up_block_types=('''AttnUpBlock1D''', '''UpBlock1D''', '''UpBlock1DNoSkip''')    ,    )
           UpperCamelCase__     : Union[str, Any]											=					IPNDMScheduler()

           UpperCamelCase__     : List[str]											=					{
               '''unet''': unet,
               '''scheduler''': scheduler,
           }
           return components






     def  __SCREAMING_SNAKE_CASE      (				self    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE=0	)	-> Any:




           """simple docstring"""







           if str(__SCREAMING_SNAKE_CASE	).startswith('''mps'''	):
                 UpperCamelCase__     : Optional[Any]											=					torch.manual_seed(__SCREAMING_SNAKE_CASE	)
           else:
                 UpperCamelCase__     : Optional[Any]											=					torch.Generator(device=__SCREAMING_SNAKE_CASE	).manual_seed(__SCREAMING_SNAKE_CASE	)
           UpperCamelCase__     : int											=					{
               '''batch_size''': 1,
               '''generator''': generator,
               '''num_inference_steps''': 4,
           }
           return inputs






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> List[Any]:




           """simple docstring"""







           UpperCamelCase__     : str											=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
           UpperCamelCase__     : Dict											=					self.get_dummy_components()
           UpperCamelCase__     : str											=					DanceDiffusionPipeline(**__SCREAMING_SNAKE_CASE	)
           UpperCamelCase__     : List[Any]											=					pipe.to(__SCREAMING_SNAKE_CASE	)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : int											=					self.get_dummy_inputs(__SCREAMING_SNAKE_CASE	)
           UpperCamelCase__     : Dict											=					pipe(**__SCREAMING_SNAKE_CASE	)
           UpperCamelCase__     : List[str]											=					output.audios

           UpperCamelCase__     : Dict											=					audio[0, -3:, -3:]

           assert audio.shape == (1, 2, components["unet"].sample_size)
           UpperCamelCase__     : List[Any]											=					np.array([-0.7265, 1.0000, -0.8388, 0.1175, 0.9498, -1.0000]	)
           assert np.abs(audio_slice.flatten() - expected_slice	).max() < 1e-2






     @skip_mps
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> List[Any]:




           """simple docstring"""







           return super().test_save_load_local()






     @skip_mps
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Tuple:




           """simple docstring"""







           return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3	)






     @skip_mps
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> str:




           """simple docstring"""







           return super().test_save_load_optional_components()






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




           """simple docstring"""







           return super().test_attention_slicing_forward_pass()






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Tuple:




           """simple docstring"""







           super().test_inference_batch_single_identical(expected_max_diff=3e-3	)




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





     """simple docstring"""







     def  __SCREAMING_SNAKE_CASE      (				self	)	-> int:




           """simple docstring"""







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






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




           """simple docstring"""







           UpperCamelCase__     : Any											=					torch_device

           UpperCamelCase__     : Any											=					DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''	)
           UpperCamelCase__     : int											=					pipe.to(__SCREAMING_SNAKE_CASE	)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : int											=					torch.manual_seed(0	)
           UpperCamelCase__     : Optional[Any]											=					pipe(generator=__SCREAMING_SNAKE_CASE    ,    num_inference_steps=1_0_0    ,    audio_length_in_s=4.096	)
           UpperCamelCase__     : str											=					output.audios

           UpperCamelCase__     : List[Any]											=					audio[0, -3:, -3:]

           assert audio.shape == (1, 2, pipe.unet.sample_size)
           UpperCamelCase__     : Tuple											=					np.array([-0.0192, -0.0231, -0.0318, -0.0059, 0.0002, -0.0020]	)

           assert np.abs(audio_slice.flatten() - expected_slice	).max() < 1e-2






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




           """simple docstring"""







           UpperCamelCase__     : Any											=					torch_device

           UpperCamelCase__     : Union[str, Any]											=					DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''    ,    torch_dtype=torch.floataa	)
           UpperCamelCase__     : Tuple											=					pipe.to(__SCREAMING_SNAKE_CASE	)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : int											=					torch.manual_seed(0	)
           UpperCamelCase__     : Union[str, Any]											=					pipe(generator=__SCREAMING_SNAKE_CASE    ,    num_inference_steps=1_0_0    ,    audio_length_in_s=4.096	)
           UpperCamelCase__     : List[Any]											=					output.audios

           UpperCamelCase__     : List[Any]											=					audio[0, -3:, -3:]

           assert audio.shape == (1, 2, pipe.unet.sample_size)
           UpperCamelCase__     : Optional[Any]											=					np.array([-0.0367, -0.0488, -0.0771, -0.0525, -0.0444, -0.0341]	)

           assert np.abs(audio_slice.flatten() - expected_slice	).max() < 1e-2




import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers.testing_utils import require_vision
from transformers.utils import is_vision_available


if is_vision_available():
    from PIL import Image

    from transformers import AutoProcessor, BlipaProcessor, BlipImageProcessor, GPTaTokenizer, PreTrainedTokenizerFast




@require_vision
class 					_lowerCamelCase			(   unittest.TestCase			):





     """simple docstring"""







     def  __SCREAMING_SNAKE_CASE      (				self	)	-> int:




           """simple docstring"""







           UpperCamelCase__     : int											=					tempfile.mkdtemp()

           UpperCamelCase__     : str											=					BlipImageProcessor()
           UpperCamelCase__     : List[Any]											=					GPTaTokenizer.from_pretrained('''hf-internal-testing/tiny-random-GPT2Model'''	)

           UpperCamelCase__     : List[Any]											=					BlipaProcessor(__SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE	)

           processor.save_pretrained(self.tmpdirname	)






     def  __SCREAMING_SNAKE_CASE      (				self    ,    **__SCREAMING_SNAKE_CASE	)	-> int:




           """simple docstring"""







           return AutoProcessor.from_pretrained(self.tmpdirname    ,    **__SCREAMING_SNAKE_CASE	).tokenizer






     def  __SCREAMING_SNAKE_CASE      (				self    ,    **__SCREAMING_SNAKE_CASE	)	-> int:




           """simple docstring"""







           return AutoProcessor.from_pretrained(self.tmpdirname    ,    **__SCREAMING_SNAKE_CASE	).image_processor






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Tuple:




           """simple docstring"""







           shutil.rmtree(self.tmpdirname	)






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> List[Any]:




           """simple docstring"""







           UpperCamelCase__     : List[Any]											=					[np.random.randint(2_5_5    ,    size=(3, 3_0, 4_0_0)    ,    dtype=np.uinta	)]

           UpperCamelCase__     : Dict											=					[Image.fromarray(np.moveaxis(__SCREAMING_SNAKE_CASE    ,    0    ,    -1	)	) for x in image_inputs]

           return image_inputs






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Optional[Any]:




           """simple docstring"""







           UpperCamelCase__     : str											=					BlipaProcessor(tokenizer=self.get_tokenizer()    ,    image_processor=self.get_image_processor()	)
           processor.save_pretrained(self.tmpdirname	)

           UpperCamelCase__     : Any											=					self.get_tokenizer(bos_token='''(BOS)'''    ,    eos_token='''(EOS)'''	)
           UpperCamelCase__     : str											=					self.get_image_processor(do_normalize=__SCREAMING_SNAKE_CASE    ,    padding_value=1.0	)

           UpperCamelCase__     : Optional[Any]											=					BlipaProcessor.from_pretrained(
               self.tmpdirname    ,    bos_token='''(BOS)'''    ,    eos_token='''(EOS)'''    ,    do_normalize=__SCREAMING_SNAKE_CASE    ,    padding_value=1.0	)

           self.assertEqual(processor.tokenizer.get_vocab()    ,    tokenizer_add_kwargs.get_vocab()	)
           self.assertIsInstance(processor.tokenizer    ,    __SCREAMING_SNAKE_CASE	)

           self.assertEqual(processor.image_processor.to_json_string()    ,    image_processor_add_kwargs.to_json_string()	)
           self.assertIsInstance(processor.image_processor    ,    __SCREAMING_SNAKE_CASE	)






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> List[Any]:




           """simple docstring"""







           UpperCamelCase__     : Dict											=					self.get_image_processor()
           UpperCamelCase__     : Tuple											=					self.get_tokenizer()

           UpperCamelCase__     : int											=					BlipaProcessor(tokenizer=__SCREAMING_SNAKE_CASE    ,    image_processor=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : Union[str, Any]											=					self.prepare_image_inputs()

           UpperCamelCase__     : Optional[Any]											=					image_processor(__SCREAMING_SNAKE_CASE    ,    return_tensors='''np'''	)
           UpperCamelCase__     : List[Any]											=					processor(images=__SCREAMING_SNAKE_CASE    ,    return_tensors='''np'''	)

           for key in input_feat_extract.keys():
                 self.assertAlmostEqual(input_feat_extract[key].sum()    ,    input_processor[key].sum()    ,    delta=1e-2	)






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Dict:




           """simple docstring"""







           UpperCamelCase__     : str											=					self.get_image_processor()
           UpperCamelCase__     : Any											=					self.get_tokenizer()

           UpperCamelCase__     : Tuple											=					BlipaProcessor(tokenizer=__SCREAMING_SNAKE_CASE    ,    image_processor=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : List[str]											=					'''lower newer'''

           UpperCamelCase__     : Tuple											=					processor(text=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : int											=					tokenizer(__SCREAMING_SNAKE_CASE    ,    return_token_type_ids=__SCREAMING_SNAKE_CASE	)

           for key in encoded_tok.keys():
                 self.assertListEqual(encoded_tok[key]    ,    encoded_processor[key]	)






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Tuple:




           """simple docstring"""







           UpperCamelCase__     : Any											=					self.get_image_processor()
           UpperCamelCase__     : Optional[int]											=					self.get_tokenizer()

           UpperCamelCase__     : Dict											=					BlipaProcessor(tokenizer=__SCREAMING_SNAKE_CASE    ,    image_processor=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : Tuple											=					'''lower newer'''
           UpperCamelCase__     : Any											=					self.prepare_image_inputs()

           UpperCamelCase__     : Any											=					processor(text=__SCREAMING_SNAKE_CASE    ,    images=__SCREAMING_SNAKE_CASE	)

           self.assertListEqual(list(inputs.keys()	)    ,    ['''pixel_values''', '''input_ids''', '''attention_mask''']	)

           # test if it raises when no input is passed
           with pytest.raises(__SCREAMING_SNAKE_CASE	):
                 processor()






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> int:




           """simple docstring"""







           UpperCamelCase__     : List[str]											=					self.get_image_processor()
           UpperCamelCase__     : Optional[int]											=					self.get_tokenizer()

           UpperCamelCase__     : Tuple											=					BlipaProcessor(tokenizer=__SCREAMING_SNAKE_CASE    ,    image_processor=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : List[Any]											=					[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

           UpperCamelCase__     : Tuple											=					processor.batch_decode(__SCREAMING_SNAKE_CASE	)
           UpperCamelCase__     : List[str]											=					tokenizer.batch_decode(__SCREAMING_SNAKE_CASE	)

           self.assertListEqual(__SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE	)






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> str:




           """simple docstring"""







           UpperCamelCase__     : Optional[Any]											=					self.get_image_processor()
           UpperCamelCase__     : Dict											=					self.get_tokenizer()

           UpperCamelCase__     : int											=					BlipaProcessor(tokenizer=__SCREAMING_SNAKE_CASE    ,    image_processor=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : List[Any]											=					'''lower newer'''
           UpperCamelCase__     : List[str]											=					self.prepare_image_inputs()

           UpperCamelCase__     : List[Any]											=					processor(text=__SCREAMING_SNAKE_CASE    ,    images=__SCREAMING_SNAKE_CASE	)

           # For now the processor supports only ['pixel_values', 'input_ids', 'attention_mask']
           self.assertListEqual(list(inputs.keys()	)    ,    ['''pixel_values''', '''input_ids''', '''attention_mask''']	)


