



'''simple docstring'''
import json
import os
import pickle
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
					import faiss






@require_faiss
class 							lowercase__     (							lowercase       ):





	def      UpperCamelCase_  (			self :    Any			):







							'''simple docstring'''
							_UpperCamelCase :		Union[str, Any]			    =       tempfile.mkdtemp()
							_UpperCamelCase :		List[str]			    =       8

							# DPR tok
							_UpperCamelCase :		int			    =       [
							    '[UNK]',
							    '[CLS]',
							    '[SEP]',
							    '[PAD]',
							    '[MASK]',
							    'want',
							    '##want',
							    '##ed',
							    'wa',
							    'un',
							    'runn',
							    '##ing',
							    ',',
							    'low',
							    'lowest',
							]
							_UpperCamelCase :		List[Any]			    =       os.path.join(self.tmpdirname			,'dpr_tokenizer'			)
							os.makedirs(lowerCamelCase__			,exist_ok=lowerCamelCase__			)
							_UpperCamelCase :		List[str]			    =       os.path.join(lowerCamelCase__			,DPR_VOCAB_FILES_NAMES['vocab_file']			)
							with open(self.vocab_file			,'w'			,encoding='utf-8'			) as vocab_writer:
													vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]			)			)

							# BART tok
							_UpperCamelCase :		Any			    =       [
							    'l',
							    'o',
							    'w',
							    'e',
							    'r',
							    's',
							    't',
							    'i',
							    'd',
							    'n',
							    '\u0120',
							    '\u0120l',
							    '\u0120n',
							    '\u0120lo',
							    '\u0120low',
							    'er',
							    '\u0120lowest',
							    '\u0120newer',
							    '\u0120wider',
							    '<unk>',
							]
							_UpperCamelCase :		Any			    =       dict(zip(lowerCamelCase__			,range(len(lowerCamelCase__			)			)			)			)
							_UpperCamelCase :		Optional[Any]			    =       ['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
							_UpperCamelCase :		List[Any]			    =       {'unk_token': '<unk>'}

							_UpperCamelCase :		Union[str, Any]			    =       os.path.join(self.tmpdirname			,'bart_tokenizer'			)
							os.makedirs(lowerCamelCase__			,exist_ok=lowerCamelCase__			)
							_UpperCamelCase :		int			    =       os.path.join(lowerCamelCase__			,BART_VOCAB_FILES_NAMES['vocab_file']			)
							_UpperCamelCase :		str			    =       os.path.join(lowerCamelCase__			,BART_VOCAB_FILES_NAMES['merges_file']			)
							with open(self.vocab_file			,'w'			,encoding='utf-8'			) as fp:
													fp.write(json.dumps(lowerCamelCase__			) + '\n'			)
							with open(self.merges_file			,'w'			,encoding='utf-8'			) as fp:
													fp.write('\n'.join(lowerCamelCase__			)			)





	def      UpperCamelCase_  (			self :    Tuple			):







							'''simple docstring'''
							return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname			,'dpr_tokenizer'			)			)





	def      UpperCamelCase_  (			self :    List[str]			):







							'''simple docstring'''
							return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname			,'dpr_tokenizer'			)			)





	def      UpperCamelCase_  (			self :    Optional[int]			):







							'''simple docstring'''
							return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname			,'bart_tokenizer'			)			)





	def      UpperCamelCase_  (			self :    List[Any]			):







							'''simple docstring'''
							shutil.rmtree(self.tmpdirname			)





	def      UpperCamelCase_  (			self :    int			):







							'''simple docstring'''
							_UpperCamelCase :		Union[str, Any]			    =       Dataset.from_dict(
							    {
							        'id': ['0', '1'],
							        'text': ['foo', 'bar'],
							        'title': ['Foo', 'Bar'],
							        'embeddings': [np.ones(self.retrieval_vector_size			), 2 * np.ones(self.retrieval_vector_size			)],
							    }			)
							dataset.add_faiss_index('embeddings'			,string_factory='Flat'			,metric_type=faiss.METRIC_INNER_PRODUCT			)
							return dataset





	def      UpperCamelCase_  (			self :    str			):







							'''simple docstring'''
							_UpperCamelCase :		List[str]			    =       self.get_dummy_dataset()
							_UpperCamelCase :		List[Any]			    =       RagConfig(
							    retrieval_vector_size=self.retrieval_vector_size			,question_encoder=DPRConfig().to_dict()			,generator=BartConfig().to_dict()			,)
							with patch('transformers.models.rag.retrieval_rag.load_dataset'			) as mock_load_dataset:
													_UpperCamelCase :		str			    =       dataset
													_UpperCamelCase :		int			    =       RagRetriever(
													    lowerCamelCase__			,question_encoder_tokenizer=self.get_dpr_tokenizer()			,generator_tokenizer=self.get_bart_tokenizer()			,)
							return retriever





	def      UpperCamelCase_  (			self :    Optional[int]			,lowerCamelCase__ :    bool			):







							'''simple docstring'''
							_UpperCamelCase :		str			    =       self.get_dummy_dataset()
							_UpperCamelCase :		Tuple			    =       RagConfig(
							    retrieval_vector_size=self.retrieval_vector_size			,question_encoder=DPRConfig().to_dict()			,generator=BartConfig().to_dict()			,index_name='custom'			,)
							if from_disk:
													_UpperCamelCase :		int			    =       os.path.join(self.tmpdirname			,'dataset'			)
													_UpperCamelCase :		Any			    =       os.path.join(self.tmpdirname			,'index.faiss'			)
													dataset.get_index('embeddings'			).save(os.path.join(self.tmpdirname			,'index.faiss'			)			)
													dataset.drop_index('embeddings'			)
													dataset.save_to_disk(os.path.join(self.tmpdirname			,'dataset'			)			)
													del dataset
													_UpperCamelCase :		Any			    =       RagRetriever(
													    lowerCamelCase__			,question_encoder_tokenizer=self.get_dpr_tokenizer()			,generator_tokenizer=self.get_bart_tokenizer()			,)
							else:
													_UpperCamelCase :		str			    =       RagRetriever(
													    lowerCamelCase__			,question_encoder_tokenizer=self.get_dpr_tokenizer()			,generator_tokenizer=self.get_bart_tokenizer()			,index=CustomHFIndex(config.retrieval_vector_size			,lowerCamelCase__			)			,)
							return retriever





	def      UpperCamelCase_  (			self :    str			):







							'''simple docstring'''
							_UpperCamelCase :		Optional[Any]			    =       Dataset.from_dict(
							    {
							        'id': ['0', '1'],
							        'text': ['foo', 'bar'],
							        'title': ['Foo', 'Bar'],
							        'embeddings': [np.ones(self.retrieval_vector_size + 1			), 2 * np.ones(self.retrieval_vector_size + 1			)],
							    }			)
							dataset.add_faiss_index('embeddings'			,string_factory='Flat'			,metric_type=faiss.METRIC_INNER_PRODUCT			)

							_UpperCamelCase :		Union[str, Any]			    =       os.path.join(self.tmpdirname			,'hf_bert_base.hnswSQ8_correct_phi_128.c_index'			)
							dataset.save_faiss_index('embeddings'			,index_file_name + '.index.dpr'			)
							pickle.dump(dataset['id']			,open(index_file_name + '.index_meta.dpr'			,'wb'			)			)

							_UpperCamelCase :		Dict			    =       os.path.join(self.tmpdirname			,'psgs_w100.tsv.pkl'			)
							_UpperCamelCase :		int			    =       {sample['id']: [sample['text'], sample['title']] for sample in dataset}
							pickle.dump(lowerCamelCase__			,open(lowerCamelCase__			,'wb'			)			)

							_UpperCamelCase :		Union[str, Any]			    =       RagConfig(
							    retrieval_vector_size=self.retrieval_vector_size			,question_encoder=DPRConfig().to_dict()			,generator=BartConfig().to_dict()			,index_name='legacy'			,index_path=self.tmpdirname			,)
							_UpperCamelCase :		List[Any]			    =       RagRetriever(
							    lowerCamelCase__			,question_encoder_tokenizer=self.get_dpr_tokenizer()			,generator_tokenizer=self.get_bart_tokenizer()			)
							return retriever





	def      UpperCamelCase_  (			self :    int			):







							'''simple docstring'''
							_UpperCamelCase :		Dict			    =       1
							_UpperCamelCase :		List[Any]			    =       self.get_dummy_canonical_hf_index_retriever()
							_UpperCamelCase :		Tuple			    =       np.array(
							    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,dtype=np.floataa			)
							_UpperCamelCase    ,      _UpperCamelCase    ,      _UpperCamelCase :		Optional[int]			    =       retriever.retrieve(lowerCamelCase__			,n_docs=lowerCamelCase__			)
							self.assertEqual(retrieved_doc_embeds.shape			,(2, n_docs, self.retrieval_vector_size)			)
							self.assertEqual(len(lowerCamelCase__			)			,2			)
							self.assertEqual(sorted(doc_dicts[0]			)			,['embeddings', 'id', 'text', 'title']			)
							self.assertEqual(len(doc_dicts[0]['id']			)			,lowerCamelCase__			)
							self.assertEqual(doc_dicts[0]['id'][0]			,'1'			)  # max inner product is reached with second doc
							self.assertEqual(doc_dicts[1]['id'][0]			,'0'			)  # max inner product is reached with first doc
							self.assertListEqual(doc_ids.tolist()			,[[1], [0]]			)





	def      UpperCamelCase_  (			self :    List[str]			):







							'''simple docstring'''
							_UpperCamelCase :		Dict			    =       self.get_dummy_canonical_hf_index_retriever()
							with tempfile.TemporaryDirectory() as tmp_dirname:
													with patch('transformers.models.rag.retrieval_rag.load_dataset'			) as mock_load_dataset:
																			_UpperCamelCase :		List[Any]			    =       self.get_dummy_dataset()
																			retriever.save_pretrained(lowerCamelCase__			)
																			_UpperCamelCase :		Optional[Any]			    =       RagRetriever.from_pretrained(lowerCamelCase__			)
																			self.assertIsInstance(lowerCamelCase__			,lowerCamelCase__			)
																			_UpperCamelCase :		Dict			    =       np.array(
																			    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,dtype=np.floataa			)
																			_UpperCamelCase :		int			    =       retriever.retrieve(lowerCamelCase__			,n_docs=1			)
																			self.assertTrue(out is not None			)





	def      UpperCamelCase_  (			self :    Union[str, Any]			):







							'''simple docstring'''
							_UpperCamelCase :		List[Any]			    =       1
							_UpperCamelCase :		List[Any]			    =       self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase__			)
							_UpperCamelCase :		Optional[Any]			    =       np.array(
							    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,dtype=np.floataa			)
							_UpperCamelCase    ,      _UpperCamelCase    ,      _UpperCamelCase :		List[str]			    =       retriever.retrieve(lowerCamelCase__			,n_docs=lowerCamelCase__			)
							self.assertEqual(retrieved_doc_embeds.shape			,(2, n_docs, self.retrieval_vector_size)			)
							self.assertEqual(len(lowerCamelCase__			)			,2			)
							self.assertEqual(sorted(doc_dicts[0]			)			,['embeddings', 'id', 'text', 'title']			)
							self.assertEqual(len(doc_dicts[0]['id']			)			,lowerCamelCase__			)
							self.assertEqual(doc_dicts[0]['id'][0]			,'1'			)  # max inner product is reached with second doc
							self.assertEqual(doc_dicts[1]['id'][0]			,'0'			)  # max inner product is reached with first doc
							self.assertListEqual(doc_ids.tolist()			,[[1], [0]]			)





	def      UpperCamelCase_  (			self :    Dict			):







							'''simple docstring'''
							_UpperCamelCase :		Optional[Any]			    =       self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase__			)
							with tempfile.TemporaryDirectory() as tmp_dirname:
													retriever.save_pretrained(lowerCamelCase__			)
													_UpperCamelCase :		Tuple			    =       RagRetriever.from_pretrained(lowerCamelCase__			)
													self.assertIsInstance(lowerCamelCase__			,lowerCamelCase__			)
													_UpperCamelCase :		Dict			    =       np.array(
													    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,dtype=np.floataa			)
													_UpperCamelCase :		List[Any]			    =       retriever.retrieve(lowerCamelCase__			,n_docs=1			)
													self.assertTrue(out is not None			)





	def      UpperCamelCase_  (			self :    Tuple			):







							'''simple docstring'''
							_UpperCamelCase :		Tuple			    =       1
							_UpperCamelCase :		Any			    =       self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase__			)
							_UpperCamelCase :		Dict			    =       np.array(
							    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,dtype=np.floataa			)
							_UpperCamelCase    ,      _UpperCamelCase    ,      _UpperCamelCase :		str			    =       retriever.retrieve(lowerCamelCase__			,n_docs=lowerCamelCase__			)
							self.assertEqual(retrieved_doc_embeds.shape			,(2, n_docs, self.retrieval_vector_size)			)
							self.assertEqual(len(lowerCamelCase__			)			,2			)
							self.assertEqual(sorted(doc_dicts[0]			)			,['embeddings', 'id', 'text', 'title']			)
							self.assertEqual(len(doc_dicts[0]['id']			)			,lowerCamelCase__			)
							self.assertEqual(doc_dicts[0]['id'][0]			,'1'			)  # max inner product is reached with second doc
							self.assertEqual(doc_dicts[1]['id'][0]			,'0'			)  # max inner product is reached with first doc
							self.assertListEqual(doc_ids.tolist()			,[[1], [0]]			)





	def      UpperCamelCase_  (			self :    Tuple			):







							'''simple docstring'''
							_UpperCamelCase :		Dict			    =       self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase__			)
							with tempfile.TemporaryDirectory() as tmp_dirname:
													retriever.save_pretrained(lowerCamelCase__			)
													_UpperCamelCase :		Union[str, Any]			    =       RagRetriever.from_pretrained(lowerCamelCase__			)
													self.assertIsInstance(lowerCamelCase__			,lowerCamelCase__			)
													_UpperCamelCase :		List[Any]			    =       np.array(
													    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,dtype=np.floataa			)
													_UpperCamelCase :		Optional[Any]			    =       retriever.retrieve(lowerCamelCase__			,n_docs=1			)
													self.assertTrue(out is not None			)





	def      UpperCamelCase_  (			self :    int			):







							'''simple docstring'''
							_UpperCamelCase :		int			    =       1
							_UpperCamelCase :		str			    =       self.get_dummy_legacy_index_retriever()
							_UpperCamelCase :		Union[str, Any]			    =       np.array(
							    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,dtype=np.floataa			)
							_UpperCamelCase    ,      _UpperCamelCase    ,      _UpperCamelCase :		List[str]			    =       retriever.retrieve(lowerCamelCase__			,n_docs=lowerCamelCase__			)
							self.assertEqual(retrieved_doc_embeds.shape			,(2, n_docs, self.retrieval_vector_size)			)
							self.assertEqual(len(lowerCamelCase__			)			,2			)
							self.assertEqual(sorted(doc_dicts[0]			)			,['text', 'title']			)
							self.assertEqual(len(doc_dicts[0]['text']			)			,lowerCamelCase__			)
							self.assertEqual(doc_dicts[0]['text'][0]			,'bar'			)  # max inner product is reached with second doc
							self.assertEqual(doc_dicts[1]['text'][0]			,'foo'			)  # max inner product is reached with first doc
							self.assertListEqual(doc_ids.tolist()			,[[1], [0]]			)





	def      UpperCamelCase_  (			self :    Optional[int]			):







							'''simple docstring'''
							_UpperCamelCase :		Optional[int]			    =       self.get_dummy_legacy_index_retriever()
							with tempfile.TemporaryDirectory() as tmp_dirname:
													retriever.save_pretrained(lowerCamelCase__			)
													_UpperCamelCase :		int			    =       RagRetriever.from_pretrained(lowerCamelCase__			)
													self.assertIsInstance(lowerCamelCase__			,lowerCamelCase__			)
													_UpperCamelCase :		List[Any]			    =       np.array(
													    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,dtype=np.floataa			)
													_UpperCamelCase :		List[Any]			    =       retriever.retrieve(lowerCamelCase__			,n_docs=1			)
													self.assertTrue(out is not None			)





	@require_torch
	@require_tokenizers
	@require_sentencepiece
	def      UpperCamelCase_  (			self :    str			):







							'''simple docstring'''
							import torch

							_UpperCamelCase :		int			    =       1
							_UpperCamelCase :		Any			    =       self.get_dummy_canonical_hf_index_retriever()
							_UpperCamelCase :		int			    =       [[5, 7], [10, 11]]
							_UpperCamelCase :		int			    =       np.array(
							    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,dtype=np.floataa			)
							_UpperCamelCase :		Tuple			    =       retriever(lowerCamelCase__			,lowerCamelCase__			,prefix=retriever.config.generator.prefix			,n_docs=lowerCamelCase__			)
							_UpperCamelCase    ,      _UpperCamelCase    ,      _UpperCamelCase :		int			    =       (
							    out['context_input_ids'],
							    out['context_attention_mask'],
							    out['retrieved_doc_embeds'],
							)
							self.assertEqual(retrieved_doc_embeds.shape			,(2, n_docs, self.retrieval_vector_size)			)
							self.assertIsInstance(lowerCamelCase__			,lowerCamelCase__			)
							self.assertIsInstance(lowerCamelCase__			,lowerCamelCase__			)
							self.assertIsInstance(lowerCamelCase__			,np.ndarray			)

							_UpperCamelCase :		Union[str, Any]			    =       retriever(
							    lowerCamelCase__			,lowerCamelCase__			,prefix=retriever.config.generator.prefix			,n_docs=lowerCamelCase__			,return_tensors='pt'			,)
							_UpperCamelCase    ,      _UpperCamelCase    ,      _UpperCamelCase    ,      _UpperCamelCase :		str			    =       (  # noqa: F841
							    out['context_input_ids'],
							    out['context_attention_mask'],
							    out['retrieved_doc_embeds'],
							    out['doc_ids'],
							)
							self.assertEqual(retrieved_doc_embeds.shape			,(2, n_docs, self.retrieval_vector_size)			)
							self.assertIsInstance(lowerCamelCase__			,torch.Tensor			)
							self.assertIsInstance(lowerCamelCase__			,torch.Tensor			)
							self.assertIsInstance(lowerCamelCase__			,torch.Tensor			)





	@require_torch
	@require_tokenizers
	@require_sentencepiece
	def      UpperCamelCase_  (			self :    Optional[int]			):







							'''simple docstring'''
							_UpperCamelCase :		Optional[Any]			    =       self.get_dpr_ctx_encoder_tokenizer()
							_UpperCamelCase :		Union[str, Any]			    =       1
							_UpperCamelCase :		Union[str, Any]			    =       self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase__			)
							retriever.set_ctx_encoder_tokenizer(lowerCamelCase__			)

							_UpperCamelCase :		int			    =       [[5, 7], [10, 11]]
							_UpperCamelCase :		Tuple			    =       np.array(
							    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,dtype=np.floataa			)
							_UpperCamelCase :		str			    =       retriever(lowerCamelCase__			,lowerCamelCase__			,prefix=retriever.config.generator.prefix			,n_docs=lowerCamelCase__			)

							self.assertEqual(
							    len(lowerCamelCase__			)			,6			)  # check whether the retriever output consist of 6 attributes including tokenized docs
							self.assertEqual(
							    all(k in out for k in ('tokenized_doc_ids', 'tokenized_doc_attention_mask')			)			,lowerCamelCase__			)  # check for doc token related keys in dictionary.







'''simple docstring'''
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 							lowercase__     (							lowercase						, unittest.TestCase       ):
	lowercase__      =      DanceDiffusionPipeline
	lowercase__      =      UNCONDITIONAL_AUDIO_GENERATION_PARAMS
	lowercase__      =      PipelineTesterMixin.required_optional_params - {
	    """callback""",
	    """latents""",
	    """callback_steps""",
	    """output_type""",
	    """num_images_per_prompt""",
	}
	lowercase__      =      UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
	lowercase__      =      False
	lowercase__      =      False





	def      UpperCamelCase_  (			self :    Union[str, Any]			):







							'''simple docstring'''
							torch.manual_seed(0			)
							_UpperCamelCase :		str			    =       UNetaDModel(
							    block_out_channels=(32, 32, 64)			,extra_in_channels=16			,sample_size=512			,sample_rate=16000			,in_channels=2			,out_channels=2			,flip_sin_to_cos=lowerCamelCase__			,use_timestep_embedding=lowerCamelCase__			,time_embedding_type='fourier'			,mid_block_type='UNetMidBlock1D'			,down_block_types=('DownBlock1DNoSkip', 'DownBlock1D', 'AttnDownBlock1D')			,up_block_types=('AttnUpBlock1D', 'UpBlock1D', 'UpBlock1DNoSkip')			,)
							_UpperCamelCase :		int			    =       IPNDMScheduler()

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





	def      UpperCamelCase_  (			self :    Tuple			,lowerCamelCase__ :    Tuple			,lowerCamelCase__ :    int=0			):







							'''simple docstring'''
							if str(lowerCamelCase__			).startswith('mps'			):
													_UpperCamelCase :		Union[str, Any]			    =       torch.manual_seed(lowerCamelCase__			)
							else:
													_UpperCamelCase :		str			    =       torch.Generator(device=lowerCamelCase__			).manual_seed(lowerCamelCase__			)
							_UpperCamelCase :		str			    =       {
							    'batch_size': 1,
							    'generator': generator,
							    'num_inference_steps': 4,
							}
							return inputs





	def      UpperCamelCase_  (			self :    int			):







							'''simple docstring'''
							_UpperCamelCase :		Union[str, Any]			    =       'cpu'  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase :		List[str]			    =       self.get_dummy_components()
							_UpperCamelCase :		int			    =       DanceDiffusionPipeline(**lowerCamelCase__			)
							_UpperCamelCase :		List[str]			    =       pipe.to(lowerCamelCase__			)
							pipe.set_progress_bar_config(disable=lowerCamelCase__			)

							_UpperCamelCase :		Union[str, Any]			    =       self.get_dummy_inputs(lowerCamelCase__			)
							_UpperCamelCase :		List[Any]			    =       pipe(**lowerCamelCase__			)
							_UpperCamelCase :		Union[str, Any]			    =       output.audios

							_UpperCamelCase :		List[str]			    =       audio[0, -3:, -3:]

							assert audio.shape == (1, 2, components["unet"].sample_size)
							_UpperCamelCase :		Dict			    =       np.array([-0.7_2_6_5, 1.0_0_0_0, -0.8_3_8_8, 0.1_1_7_5, 0.9_4_9_8, -1.0_0_0_0]			)
							assert np.abs(audio_slice.flatten() - expected_slice			).max() < 1E-2





	@skip_mps
	def      UpperCamelCase_  (			self :    str			):







							'''simple docstring'''
							return super().test_save_load_local()





	@skip_mps
	def      UpperCamelCase_  (			self :    Dict			):







							'''simple docstring'''
							return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3			)





	@skip_mps
	def      UpperCamelCase_  (			self :    str			):







							'''simple docstring'''
							return super().test_save_load_optional_components()





	@skip_mps
	def      UpperCamelCase_  (			self :    List[Any]			):







							'''simple docstring'''
							return super().test_attention_slicing_forward_pass()





	def      UpperCamelCase_  (			self :    List[str]			):







							'''simple docstring'''
							super().test_inference_batch_single_identical(expected_max_diff=3E-3			)






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





	def      UpperCamelCase_  (			self :    List[str]			):







							'''simple docstring'''
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()





	def      UpperCamelCase_  (			self :    int			):







							'''simple docstring'''
							_UpperCamelCase :		Union[str, Any]			    =       torch_device

							_UpperCamelCase :		Dict			    =       DanceDiffusionPipeline.from_pretrained('harmonai/maestro-150k'			)
							_UpperCamelCase :		str			    =       pipe.to(lowerCamelCase__			)
							pipe.set_progress_bar_config(disable=lowerCamelCase__			)

							_UpperCamelCase :		int			    =       torch.manual_seed(0			)
							_UpperCamelCase :		Optional[Any]			    =       pipe(generator=lowerCamelCase__			,num_inference_steps=100			,audio_length_in_s=4.0_9_6			)
							_UpperCamelCase :		Optional[int]			    =       output.audios

							_UpperCamelCase :		int			    =       audio[0, -3:, -3:]

							assert audio.shape == (1, 2, pipe.unet.sample_size)
							_UpperCamelCase :		Tuple			    =       np.array([-0.0_1_9_2, -0.0_2_3_1, -0.0_3_1_8, -0.0_0_5_9, 0.0_0_0_2, -0.0_0_2_0]			)

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





	def      UpperCamelCase_  (			self :    Union[str, Any]			):







							'''simple docstring'''
							_UpperCamelCase :		List[Any]			    =       torch_device

							_UpperCamelCase :		int			    =       DanceDiffusionPipeline.from_pretrained('harmonai/maestro-150k'			,torch_dtype=torch.floataa			)
							_UpperCamelCase :		Any			    =       pipe.to(lowerCamelCase__			)
							pipe.set_progress_bar_config(disable=lowerCamelCase__			)

							_UpperCamelCase :		Union[str, Any]			    =       torch.manual_seed(0			)
							_UpperCamelCase :		Union[str, Any]			    =       pipe(generator=lowerCamelCase__			,num_inference_steps=100			,audio_length_in_s=4.0_9_6			)
							_UpperCamelCase :		Any			    =       output.audios

							_UpperCamelCase :		str			    =       audio[0, -3:, -3:]

							assert audio.shape == (1, 2, pipe.unet.sample_size)
							_UpperCamelCase :		Any			    =       np.array([-0.0_3_6_7, -0.0_4_8_8, -0.0_7_7_1, -0.0_5_2_5, -0.0_4_4_4, -0.0_3_4_1]			)

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


