'''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_       (		_UpperCamelCase    ,    unittest.TestCase			):







   """simple docstring"""

   __lowerCAmelCase         =   DanceDiffusionPipeline
   __lowerCAmelCase         =   UNCONDITIONAL_AUDIO_GENERATION_PARAMS
   __lowerCAmelCase         =   PipelineTesterMixin.required_optional_params - {
       "callback",
       "latents",
       "callback_steps",
       "output_type",
       "num_images_per_prompt",
   }
   __lowerCAmelCase         =   UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
   __lowerCAmelCase         =   False
   __lowerCAmelCase         =   False
   def 	__UpperCAmelCase   (      self      :			int     )  ->			Optional[Any]:
        torch.manual_seed(0     )
        _A    		=  UNetaDModel(
            block_out_channels=(32, 32, 64), extra_in_channels=16, sample_size=5_12, sample_rate=1_60_00, 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      :			Dict, UpperCamelCase__      :			Optional[Any], UpperCamelCase__      :			str=0     )  ->			List[str]:
        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      :			Optional[int]     )  ->			Union[str, 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.7_265, 1.0_000, -0.8_388, 0.1_175, 0.9_498, -1.0_000]     )
        assert np.abs(audio_slice.flatten() - expected_slice     ).max() < 1e-2
   @skip_mps
   def 	__UpperCAmelCase   (      self      :			Tuple     )  ->			Tuple:
        return super().test_save_load_local()
   @skip_mps
   def 	__UpperCAmelCase   (      self      :			Dict     )  ->			Union[str, Any]:
        return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3     )
   @skip_mps
   def 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			List[Any]:
        return super().test_save_load_optional_components()
   @skip_mps
   def 	__UpperCAmelCase   (      self      :			int     )  ->			Optional[int]:
        return super().test_attention_slicing_forward_pass()

   def 	__UpperCAmelCase   (      self      :			Any     )  ->			Any:
        super().test_inference_batch_single_identical(expected_max_diff=3e-3     )







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







   """simple docstring"""

   def 	__UpperCAmelCase   (      self      :			Optional[Any]     )  ->			Dict:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()
   def 	__UpperCAmelCase   (      self      :			int     )  ->			Tuple:
        _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=1_00, 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.0_192, -0.0_231, -0.0_318, -0.0_059, 0.0_002, -0.0_020]     )

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

   def 	__UpperCAmelCase   (      self      :			str     )  ->			Dict:
        _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=1_00, 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.0_367, -0.0_488, -0.0_771, -0.0_525, -0.0_444, -0.0_341]     )

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




'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


_UpperCAmelCase    :     Any         =       {
    '''configuration_maskformer''': ['''MASKFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''MaskFormerConfig'''],
    '''configuration_maskformer_swin''': ['''MaskFormerSwinConfig'''],
}

try:
   if not is_vision_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCAmelCase    :     str         =       ['''MaskFormerFeatureExtractor''']
   _UpperCAmelCase    :     str         =       ['''MaskFormerImageProcessor''']


try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCAmelCase    :     Tuple         =       [
       '''MASKFORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''MaskFormerForInstanceSegmentation''',
       '''MaskFormerModel''',
       '''MaskFormerPreTrainedModel''',
   ]
   _UpperCAmelCase    :     int         =       [
       '''MaskFormerSwinBackbone''',
       '''MaskFormerSwinModel''',
       '''MaskFormerSwinPreTrainedModel''',
   ]

if TYPE_CHECKING:
   from .configuration_maskformer import MASKFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, MaskFormerConfig
   from .configuration_maskformer_swin import MaskFormerSwinConfig

   try:
      if not is_vision_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .feature_extraction_maskformer import MaskFormerFeatureExtractor
      from .image_processing_maskformer import MaskFormerImageProcessor
   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_maskformer import (
          MASKFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
          MaskFormerForInstanceSegmentation,
          MaskFormerModel,
          MaskFormerPreTrainedModel,
      )
      from .modeling_maskformer_swin import (
          MaskFormerSwinBackbone,
          MaskFormerSwinModel,
          MaskFormerSwinPreTrainedModel,
      )


else:
   import sys

   _UpperCAmelCase    :     str         =       _LazyModule(__name__, globals()['''__file__'''], _import_structure)



