


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


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[Any] ):
    torch.manual_seed(0 )
    UpperCAmelCase__	:Union[str, Any]      =      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=__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__	:Union[str, Any]      =      {
        '''unet''': unet,
        '''scheduler''': scheduler,
    }
    return components


 def        __SCREAMING_SNAKE_CASE					(      self							:	str			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	Dict=0 ):
    if str(__lowerCamelCase ).startswith('''mps''' ):
       UpperCAmelCase__	:Dict      =      torch.manual_seed(__lowerCamelCase )
    else:
       UpperCAmelCase__	:List[str]      =      torch.Generator(device=__lowerCamelCase ).manual_seed(__lowerCamelCase )
    UpperCAmelCase__	:str      =      {
        '''batch_size''': 1,
        '''generator''': generator,
        '''num_inference_steps''': 4,
    }
    return inputs


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

    UpperCAmelCase__	:List[str]      =      self.get_dummy_inputs(__lowerCamelCase )
    UpperCAmelCase__	:List[Any]      =      pipe(**__lowerCamelCase )
    UpperCAmelCase__	:Any      =      output.audios

    UpperCAmelCase__	:Any      =      audio[0, -3:, -3:]

    assert audio.shape == (1, 2, components["unet"].sample_size)
    UpperCAmelCase__	:List[Any]      =      np.array([-0.72_65, 1.00_00, -0.83_88, 0.11_75, 0.94_98, -1.00_00] )
    assert np.abs(audio_slice.flatten() - expected_slice ).max() < 1e-2


 @skip_mps
 def        __SCREAMING_SNAKE_CASE					(      self							:	Tuple ):
    return super().test_save_load_local()


 @skip_mps
 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3 )


 @skip_mps
 def        __SCREAMING_SNAKE_CASE					(      self							:	Union[str, Any] ):
    return super().test_save_load_optional_components()


 @skip_mps
 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int] ):
    return super().test_attention_slicing_forward_pass()





 def        __SCREAMING_SNAKE_CASE					(      self							:	Union[str, Any] ):
    super().test_inference_batch_single_identical(expected_max_diff=3e-3 )




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


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any] ):
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    UpperCAmelCase__	:List[Any]      =      torch_device

    UpperCAmelCase__	:Union[str, Any]      =      DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k''' )
    UpperCAmelCase__	:int      =      pipe.to(__lowerCamelCase )
    pipe.set_progress_bar_config(disable=__lowerCamelCase )

    UpperCAmelCase__	:Optional[int]      =      torch.manual_seed(0 )
    UpperCAmelCase__	:Optional[Any]      =      pipe(generator=__lowerCamelCase			,       num_inference_steps=1_0_0			,       audio_length_in_s=4.0_96 )
    UpperCAmelCase__	:str      =      output.audios

    UpperCAmelCase__	:Union[str, Any]      =      audio[0, -3:, -3:]

    assert audio.shape == (1, 2, pipe.unet.sample_size)
    UpperCAmelCase__	:Optional[int]      =      np.array([-0.01_92, -0.02_31, -0.03_18, -0.00_59, 0.00_02, -0.00_20] )

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





 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    UpperCAmelCase__	:Optional[Any]      =      torch_device

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

    UpperCAmelCase__	:Tuple      =      torch.manual_seed(0 )
    UpperCAmelCase__	:Tuple      =      pipe(generator=__lowerCamelCase			,       num_inference_steps=1_0_0			,       audio_length_in_s=4.0_96 )
    UpperCAmelCase__	:Union[str, Any]      =      output.audios

    UpperCAmelCase__	:Union[str, Any]      =      audio[0, -3:, -3:]

    assert audio.shape == (1, 2, pipe.unet.sample_size)
    UpperCAmelCase__	:Union[str, Any]      =      np.array([-0.03_67, -0.04_88, -0.07_71, -0.05_25, -0.04_44, -0.03_41] )

    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


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

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


try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCamelCase												=						[
        '''MASKFORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
        '''MaskFormerForInstanceSegmentation''',
        '''MaskFormerModel''',
        '''MaskFormerPreTrainedModel''',
    ]
    __lowerCamelCase												=						[
        '''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

    __lowerCamelCase												=						_LazyModule(__name__, globals()['''__file__'''], _import_structure)



