


"""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     __lowerCamelCase  (					a__  ,					unittest.TestCase				):

     '''simple docstring'''
     A_       :	List[Any]    =		DanceDiffusionPipeline
     A_       :	List[str]    =		UNCONDITIONAL_AUDIO_GENERATION_PARAMS
     A_       :	Optional[int]    =		PipelineTesterMixin.required_optional_params - {
         'callback',
         'latents',
         'callback_steps',
         'output_type',
         'num_images_per_prompt',
     }
     A_       :	List[str]    =		UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
     A_       :	int    =		False
     A_       :	str    =		False


     def 							_UpperCAmelCase      (      self							)       ->							List[str]:
          torch.manual_seed(0							)
          _a    	=				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=__UpperCAmelCase    ,				use_timestep_embedding=__UpperCAmelCase    ,				time_embedding_type='''fourier'''    ,				mid_block_type='''UNetMidBlock1D'''    ,				down_block_types=('''DownBlock1DNoSkip''', '''DownBlock1D''', '''AttnDownBlock1D''')    ,				up_block_types=('''AttnUpBlock1D''', '''UpBlock1D''', '''UpBlock1DNoSkip''')    ,				)
          _a    	=				IPNDMScheduler()

          _a    	=				{
              '''unet''': unet,
              '''scheduler''': scheduler,
          }
          return components


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase=0							)       ->							List[Any]:
          if str(__UpperCAmelCase							).startswith('''mps'''							):
               _a    	=				torch.manual_seed(__UpperCAmelCase							)
          else:
               _a    	=				torch.Generator(device=__UpperCAmelCase							).manual_seed(__UpperCAmelCase							)
          _a    	=				{
              '''batch_size''': 1,
              '''generator''': generator,
              '''num_inference_steps''': 4,
          }
          return inputs


     def 							_UpperCAmelCase      (      self							)       ->							Any:
          _a    	=				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
          _a    	=				self.get_dummy_components()
          _a    	=				DanceDiffusionPipeline(**__UpperCAmelCase							)
          _a    	=				pipe.to(__UpperCAmelCase							)
          pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

          _a    	=				self.get_dummy_inputs(__UpperCAmelCase							)
          _a    	=				pipe(**__UpperCAmelCase							)
          _a    	=				output.audios

          _a    	=				audio[0, -3:, -3:]

          assert audio.shape == (1, 2, components["unet"].sample_size)
          _a    	=				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 							_UpperCAmelCase      (      self							)       ->							Tuple:
          return super().test_save_load_local()


     @skip_mps
     def 							_UpperCAmelCase      (      self							)       ->							Optional[int]:
          return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3							)


     @skip_mps
     def 							_UpperCAmelCase      (      self							)       ->							List[Any]:
          return super().test_save_load_optional_components()


     @skip_mps
     def 							_UpperCAmelCase      (      self							)       ->							int:
          return super().test_attention_slicing_forward_pass()





     def 							_UpperCAmelCase      (      self							)       ->							List[Any]:
          super().test_inference_batch_single_identical(expected_max_diff=3e-3							)







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

     '''simple docstring'''


     def 							_UpperCAmelCase      (      self							)       ->							Optional[int]:
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()


     def 							_UpperCAmelCase      (      self							)       ->							List[str]:
          _a    	=				torch_device

          _a    	=				DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''							)
          _a    	=				pipe.to(__UpperCAmelCase							)
          pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

          _a    	=				torch.manual_seed(0							)
          _a    	=				pipe(generator=__UpperCAmelCase    ,				num_inference_steps=100    ,				audio_length_in_s=4.096							)
          _a    	=				output.audios

          _a    	=				audio[0, -3:, -3:]

          assert audio.shape == (1, 2, pipe.unet.sample_size)
          _a    	=				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 							_UpperCAmelCase      (      self							)       ->							Optional[int]:
          _a    	=				torch_device

          _a    	=				DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''    ,				torch_dtype=torch.floataa							)
          _a    	=				pipe.to(__UpperCAmelCase							)
          pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

          _a    	=				torch.manual_seed(0							)
          _a    	=				pipe(generator=__UpperCAmelCase    ,				num_inference_steps=100    ,				audio_length_in_s=4.096							)
          _a    	=				output.audios

          _a    	=				audio[0, -3:, -3:]

          assert audio.shape == (1, 2, pipe.unet.sample_size)
          _a    	=				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



"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__snake_case		=       {
    '''configuration_lxmert''': ['''LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''LxmertConfig'''],
    '''tokenization_lxmert''': ['''LxmertTokenizer'''],
}

try:
   if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __snake_case		=       ['''LxmertTokenizerFast''']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __snake_case		=       [
       '''LxmertEncoder''',
       '''LxmertForPreTraining''',
       '''LxmertForQuestionAnswering''',
       '''LxmertModel''',
       '''LxmertPreTrainedModel''',
       '''LxmertVisualFeatureEncoder''',
       '''LxmertXLayer''',
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __snake_case		=       [
       '''TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''TFLxmertForPreTraining''',
       '''TFLxmertMainLayer''',
       '''TFLxmertModel''',
       '''TFLxmertPreTrainedModel''',
       '''TFLxmertVisualFeatureEncoder''',
   ]


if TYPE_CHECKING:
   from .configuration_lxmert import LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP, LxmertConfig
   from .tokenization_lxmert import LxmertTokenizer

   try:
      if not is_tokenizers_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_lxmert_fast import LxmertTokenizerFast

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_lxmert import (
          LxmertEncoder,
          LxmertForPreTraining,
          LxmertForQuestionAnswering,
          LxmertModel,
          LxmertPreTrainedModel,
          LxmertVisualFeatureEncoder,
          LxmertXLayer,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_lxmert import (
          TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFLxmertForPreTraining,
          TFLxmertMainLayer,
          TFLxmertModel,
          TFLxmertPreTrainedModel,
          TFLxmertVisualFeatureEncoder,
      )

else:
   import sys

   __snake_case		=       _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)