





# XXX: we want transformers master here - in the absense of conftest manipulating sys.path:
# hack it in for now:
import sys
from pathlib import Path


lowerCAmelCase_        =							Path(__file__).resolve().parents[3] / '''src'''
sys.path.insert(1, str(git_repo_path))

import dataclasses  # noqa
import io  # noqa
import itertools  # noqa
import json  # noqa
import os  # noqa
import unittest  # noqa
from copy import deepcopy  # noqa

from parameterized import parameterized  # noqa
from transformers import TrainingArguments, is_torch_available  # noqa
from transformers.deepspeed import is_deepspeed_available  # noqa
from transformers.file_utils import WEIGHTS_NAME  # noqa
from transformers.testing_utils import (  # noqa
    CaptureLogger,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    mockenv_context,
    require_deepspeed,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
from transformers.trainer_utils import set_seed  # noqa


set_seed(4_2)

lowerCAmelCase_        =							{'''base''': '''patrickvonplaten/wav2vec2_tiny_random''', '''robust''': '''patrickvonplaten/wav2vec2_tiny_random_robust'''}

lowerCAmelCase_        =							'''zero2'''
lowerCAmelCase_        =							'''zero3'''
lowerCAmelCase_        =							[ZEROa, ZEROa]






def 	lowerCamelCase_		(							_UpperCamelCase					,							_UpperCamelCase					,							_UpperCamelCase					)						->		Union[str, Any]:






				"""simple docstring"""


				snake_case_		:       Tuple           =    parameterized.to_safe_name('''_'''.join(str(_UpperCamelCase					) for x in param.args					)					)
				return f'''{func.__name__}_{param_based_name}'''


# Cartesian-product of zero stages with models to test
lowerCAmelCase_        =							list(itertools.product(stages, models.keys()))





@slow
@require_deepspeed
@require_torch_gpu
class    __lowerCAmelCase			(     _a	):






			@parameterized.expand(__magic_name__   ,				name_func=__magic_name__       )
			def      lowerCamelCase    (self   ,				__magic_name__   ,				__magic_name__       )	->  Dict:


							'''simple docstring'''



							self.run_and_check(
							    stage=__magic_name__   ,				model=__magic_name__   ,				distributed=__magic_name__   ,				fpaa=__magic_name__   ,				)






			@require_torch_multi_gpu
			@parameterized.expand(__magic_name__   ,				name_func=__magic_name__       )
			def      lowerCamelCase    (self   ,				__magic_name__   ,				__magic_name__       )	->  int:


							'''simple docstring'''



							self.run_and_check(
							    stage=__magic_name__   ,				model=__magic_name__   ,				distributed=__magic_name__   ,				fpaa=__magic_name__   ,				)






			@parameterized.expand(__magic_name__   ,				name_func=__magic_name__       )
			def      lowerCamelCase    (self   ,				__magic_name__   ,				__magic_name__       )	->  List[Any]:


							'''simple docstring'''



							self.run_and_check(
							    stage=__magic_name__   ,				model=__magic_name__   ,				distributed=__magic_name__   ,				fpaa=__magic_name__   ,				)






			@require_torch_multi_gpu
			@parameterized.expand(__magic_name__   ,				name_func=__magic_name__       )
			def      lowerCamelCase    (self   ,				__magic_name__   ,				__magic_name__       )	->  Optional[Any]:


							'''simple docstring'''



							self.run_and_check(
							    stage=__magic_name__   ,				model=__magic_name__   ,				distributed=__magic_name__   ,				fpaa=__magic_name__   ,				)






			def      lowerCamelCase    (self   ,				__magic_name__       )	->  Dict:


							'''simple docstring'''



							pass






			def      lowerCamelCase    (self   ,				__magic_name__   ,				__magic_name__   ,				__magic_name__ = 10   ,				__magic_name__ = True   ,				__magic_name__ = True   ,				__magic_name__ = True   ,				)	->  Any:


							'''simple docstring'''



							snake_case_		:       List[Any]           =    models[model]

							snake_case_		:       Optional[int]           =    self.run_trainer(
							    stage=__magic_name__   ,				model_name=__magic_name__   ,				eval_steps=__magic_name__   ,				num_train_epochs=1   ,				distributed=__magic_name__   ,				fpaa=__magic_name__   ,				)

							self.do_checks(__magic_name__       )

							return output_dir






			def      lowerCamelCase    (self   ,				__magic_name__   ,				__magic_name__   ,				__magic_name__ = 10   ,				__magic_name__ = 1   ,				__magic_name__ = True   ,				__magic_name__ = True   ,				)	->  Union[str, Any]:


							'''simple docstring'''



							snake_case_		:       Union[str, Any]           =    self.get_auto_remove_tmp_dir('''./xxx'''   ,				after=__magic_name__       )
							snake_case_		:       Optional[Any]           =    F'''
            --model_name_or_path {model_name}
            --dataset_name hf-internal-testing/librispeech_asr_dummy
            --dataset_config_name clean
            --train_split_name validation
            --validation_split_name validation
            --output_dir {output_dir}
            --num_train_epochs {str(__magic_name__       )}
            --per_device_train_batch_size 2
            --per_device_eval_batch_size 2
            --evaluation_strategy steps
            --learning_rate 5e-4
            --warmup_steps 8
            --orthography timit
            --preprocessing_num_workers 1
            --group_by_length
            --freeze_feature_extractor
            --report_to none
            --save_steps 0
            --eval_steps {eval_steps}
            --report_to none
        '''.split()

							if fpaa:
											args.extend(['''--fp16''']       )

							# currently ds_config_wav2vec2_zero.json requires "zero_optimization.find_unused_parameters": true,
							# hence the separate config files
							snake_case_		:       Tuple           =    F'''--deepspeed {self.test_file_dir_str}/ds_config_wav2vec2_{stage}.json'''.split()
							snake_case_		:       Dict           =    [F'''{self.examples_dir_str}/research_projects/wav2vec2/run_asr.py''']
							snake_case_		:       Optional[int]           =    self.get_launcher(__magic_name__       )

							snake_case_		:       List[Any]           =    launcher + script + args + ds_args
							# keep for quick debug
							# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
							execute_subprocess_async(__magic_name__   ,				env=self.get_env()       )

							return output_dir






			def      lowerCamelCase    (self   ,				__magic_name__=False       )	->  str:


							'''simple docstring'''



							snake_case_		:       Optional[Any]           =    min(2   ,				get_gpu_count()       ) if distributed else 1
							return F'''deepspeed --num_nodes 1 --num_gpus {num_gpus}'''.split()







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



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  torch.manual_seed(0     )
  lowerCAmelCase    				= 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=_snake_case				,       use_timestep_embedding=_snake_case				,       time_embedding_type='fourier'				,       mid_block_type='UNetMidBlock1D'				,       down_block_types=('DownBlock1DNoSkip', 'DownBlock1D', 'AttnDownBlock1D')				,       up_block_types=('AttnUpBlock1D', 'UpBlock1D', 'UpBlock1DNoSkip')				,       )
  lowerCAmelCase    				= IPNDMScheduler()

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



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case=0     ):







  """simple docstring"""






  if str(_snake_case     ).startswith('mps'     ):
   lowerCAmelCase    				= torch.manual_seed(_snake_case     )
  else:
   lowerCAmelCase    				= torch.Generator(device=_snake_case     ).manual_seed(_snake_case     )
  lowerCAmelCase    				= {
      'batch_size': 1,
      'generator': generator,
      'num_inference_steps': 4,
  }
  return inputs



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= 'cpu'  # ensure determinism for the device-dependent torch.Generator
  lowerCAmelCase    				= self.get_dummy_components()
  lowerCAmelCase    				= DanceDiffusionPipeline(**_snake_case     )
  lowerCAmelCase    				= pipe.to(_snake_case     )
  pipe.set_progress_bar_config(disable=_snake_case     )

  lowerCAmelCase    				= self.get_dummy_inputs(_snake_case     )
  lowerCAmelCase    				= pipe(**_snake_case     )
  lowerCAmelCase    				= output.audios

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

  assert audio.shape == (1, 2, components["unet"].sample_size)
  lowerCAmelCase    				= 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     ):







  """simple docstring"""






  return super().test_save_load_local()



 @skip_mps
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






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



 @skip_mps
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  return super().test_save_load_optional_components()



 @skip_mps
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  return super().test_attention_slicing_forward_pass()



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






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







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



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






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



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= torch_device

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

  lowerCAmelCase    				= torch.manual_seed(0     )
  lowerCAmelCase    				= pipe(generator=_snake_case				,       num_inference_steps=1_00				,       audio_length_in_s=4.096     )
  lowerCAmelCase    				= output.audios

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

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







  """simple docstring"""






  lowerCAmelCase    				= torch_device

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

  lowerCAmelCase    				= torch.manual_seed(0     )
  lowerCAmelCase    				= pipe(generator=_snake_case				,       num_inference_steps=1_00				,       audio_length_in_s=4.096     )
  lowerCAmelCase    				= output.audios

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

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





