





"""simple docstring"""
import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only
from utils_rag import save_json

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    int			):
 lowerCAmelCase    				= filter(lambda _UpperCAmelCase			: p.requires_grad  ,     model.parameters()			)
 lowerCAmelCase    				= sum([np.prod(p.size()			) for p in model_parameters]			)
 return params


__UpperCamelCase     :						Union[str, Any]						       =  logging.getLogger(__name__)

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Optional[int]  ,     _UpperCAmelCase					:    str			):
 if metric == "rouge2":
  lowerCAmelCase    				= '{val_avg_rouge2:.4f}-{step_count}'
 elif metric == "bleu":
  lowerCAmelCase    				= '{val_avg_bleu:.4f}-{step_count}'
 elif metric == "em":
  lowerCAmelCase    				= '{val_avg_em:.4f}-{step_count}'
 elif metric == "loss":
  lowerCAmelCase    				= '{val_avg_loss:.4f}-{step_count}'
 else:
  raise NotImplementedError(
      F'seq2seq callbacks only support rouge2 and bleu, got {metric}, You can make your own by adding to this'
      ' function.'			)

 lowerCAmelCase    				= ModelCheckpoint(
     dirpath=_UpperCAmelCase  ,     filename=_UpperCAmelCase  ,     monitor=F'val_{metric}'  ,     mode='max'  ,     save_top_k=1  ,     every_n_epochs=1  ,     )

 return checkpoint_callback

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Union[str, Any]  ,     _UpperCAmelCase					:    Dict			):
 return EarlyStopping(
     monitor=F'val_{metric}'  ,     mode='min' if 'loss' in metric else 'max'  ,     patience=_UpperCAmelCase  ,     verbose=_UpperCAmelCase  ,     )







class       a							(   pl.Callback  ):



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= {F'lr_group_{i}': param['lr'] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups     )}
  pl_module.logger.log_metrics(_snake_case     )



 @rank_zero_only
 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case=True     ):







  """simple docstring"""






  logger.info(F'***** {type_path} results at step {trainer.global_step:05d} *****'     )
  lowerCAmelCase    				= trainer.callback_metrics
  trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ['log', 'progress_bar', 'preds']}     )
  # Log results
  lowerCAmelCase    				= Path(pl_module.hparams.output_dir     )
  if type_path == "test":
   lowerCAmelCase    				= od / 'test_results.txt'
   lowerCAmelCase    				= od / 'test_generations.txt'
  else:
   # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
   # If people want this it will be easy enough to add back.
   lowerCAmelCase    				= od / F'{type_path}_results/{trainer.global_step:05d}.txt'
   lowerCAmelCase    				= od / F'{type_path}_generations/{trainer.global_step:05d}.txt'
   results_file.parent.mkdir(exist_ok=_snake_case     )
   generations_file.parent.mkdir(exist_ok=_snake_case     )
  with open(_snake_case				,       'a+'     ) as writer:
   for key in sorted(_snake_case     ):
    if key in ["log", "progress_bar", "preds"]:
     continue
    lowerCAmelCase    				= metrics[key]
    if isinstance(_snake_case				,       torch.Tensor     ):
     lowerCAmelCase    				= val.item()
    lowerCAmelCase    				= F'{key}: {val:.6f}\n'
    writer.write(_snake_case     )

  if not save_generations:
   return

  if "preds" in metrics:
   lowerCAmelCase    				= '\n'.join(metrics['preds']     )
   generations_file.open('w+'     ).write(_snake_case     )



 @rank_zero_only
 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  try:
   lowerCAmelCase    				= pl_module.model.model.num_parameters()
  except AttributeError:
   lowerCAmelCase    				= pl_module.model.num_parameters()

  lowerCAmelCase    				= count_trainable_parameters(_snake_case     )
  # mp stands for million parameters
  trainer.logger.log_metrics({'n_params': npars, 'mp': npars / 1E6, 'grad_mp': n_trainable_pars / 1E6}     )



 @rank_zero_only
 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  save_json(pl_module.metrics				,       pl_module.metrics_save_path     )
  return self._write_logs(_snake_case				,       _snake_case				,       'test'     )



 @rank_zero_only
 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  save_json(pl_module.metrics				,       pl_module.metrics_save_path     )
  # Uncommenting this will save val generations
  # return self._write_logs(trainer, pl_module, "valid")












"""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





