




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

    '''simple docstring'''






    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				:     Tuple)						->List[Any]:





           '''simple docstring'''

           torch.manual_seed(0)
           A__							=   UNetaDModel(
               block_out_channels=(32, 32, 64) ,			extra_in_channels=16 ,			sample_size=512 ,			sample_rate=16_000 ,			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 	SCREAMING_SNAKE_CASE							(      self				:     Optional[Any] ,			UpperCAmelCase__				:     List[Any] ,			UpperCAmelCase__				:     Tuple=0)						->Union[str, Any]:





           '''simple docstring'''

           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 	SCREAMING_SNAKE_CASE							(      self				:     Union[str, Any])						->Any:





           '''simple docstring'''

           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 	SCREAMING_SNAKE_CASE							(      self				:     Union[str, Any])						->Union[str, Any]:





           '''simple docstring'''

           return super().test_save_load_local()


    @skip_mps
    def 	SCREAMING_SNAKE_CASE							(      self				:     Any)						->Any:





           '''simple docstring'''

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


    @skip_mps
    def 	SCREAMING_SNAKE_CASE							(      self				:     Optional[Any])						->str:





           '''simple docstring'''

           return super().test_save_load_optional_components()


    @skip_mps
    def 	SCREAMING_SNAKE_CASE							(      self				:     List[Any])						->List[str]:





           '''simple docstring'''

           return super().test_attention_slicing_forward_pass()





    def 	SCREAMING_SNAKE_CASE							(      self				:     Optional[Any])						->Union[str, Any]:





           '''simple docstring'''

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


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

    '''simple docstring'''






    def 	SCREAMING_SNAKE_CASE							(      self				:     int)						->Optional[int]:





           '''simple docstring'''

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


    def 	SCREAMING_SNAKE_CASE							(      self				:     Optional[Any])						->int:





           '''simple docstring'''

           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 	SCREAMING_SNAKE_CASE							(      self				:     Dict)						->Any:





           '''simple docstring'''

           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








from __future__ import annotations

from typing import Any


class     UpperCamelCase_   :

    '''simple docstring'''






    def __init__(      self				:     str ,			UpperCAmelCase__				:     int ,			UpperCAmelCase__				:     int ,			UpperCAmelCase__				:     float = 0)						->None:





           '''simple docstring'''

           A__			,							A__							=   row, column
           A__							=   [[default_value for c in range(UpperCAmelCase__)] for r in range(UpperCAmelCase__)]


    def __str__(      self				:     List[str])						->str:





           '''simple docstring'''

           A__							=   f"""Matrix consist of {self.row} rows and {self.column} columns\n"""

           # Make string identifier
           A__							=   0
           for row_vector in self.array:
                  for obj in row_vector:
                         A__							=   max(UpperCAmelCase__ ,			len(str(UpperCAmelCase__)))
           A__							=   f"""%{max_element_length}s"""

           # Make string and return
           def single_line(UpperCAmelCase__				:     list[float]) -> str:
                  nonlocal string_format_identifier
                  A__							=   '''['''
                  line += ", ".join(string_format_identifier % (obj,) for obj in row_vector)
                  line += "]"
                  return line

           s += "\n".join(single_line(UpperCAmelCase__) for row_vector in self.array)
           return s


    def __repr__(      self				:     Tuple)						->str:





           '''simple docstring'''

           return str(self)


    def 	SCREAMING_SNAKE_CASE							(      self				:     List[str] ,			UpperCAmelCase__				:     tuple[int, int])						->bool:





           '''simple docstring'''

           if not (isinstance(UpperCAmelCase__ ,			(list, tuple)) and len(UpperCAmelCase__) == 2):
                  return False
           elif not (0 <= loc[0] < self.row and 0 <= loc[1] < self.column):
                  return False
           else:
                  return True


    def __getitem__(      self				:     List[Any] ,			UpperCAmelCase__				:     tuple[int, int])						->Any:





           '''simple docstring'''

           assert self.validate_indicies(UpperCAmelCase__)
           return self.array[loc[0]][loc[1]]


    def __setitem__(      self				:     List[Any] ,			UpperCAmelCase__				:     tuple[int, int] ,			UpperCAmelCase__				:     float)						->None:





           '''simple docstring'''

           assert self.validate_indicies(UpperCAmelCase__)
           A__							=   value


    def __add__(      self				:     Optional[int] ,			UpperCAmelCase__				:     Matrix)						->Matrix:





           '''simple docstring'''

           assert isinstance(UpperCAmelCase__ ,			UpperCAmelCase__)
           assert self.row == another.row and self.column == another.column

           # Add
           A__							=   Matrix(self.row ,			self.column)
           for r in range(self.row):
                  for c in range(self.column):
                         A__							=   self[r, c] + another[r, c]
           return result


    def __neg__(      self				:     str)						->Matrix:





           '''simple docstring'''

           A__							=   Matrix(self.row ,			self.column)
           for r in range(self.row):
                  for c in range(self.column):
                         A__							=   -self[r, c]
           return result


    def __sub__(      self				:     str ,			UpperCAmelCase__				:     Matrix)						->Matrix:





           '''simple docstring'''

           return self + (-another)


    def __mul__(      self				:     Union[str, Any] ,			UpperCAmelCase__				:     int | float | Matrix)						->Matrix:





           '''simple docstring'''

           if isinstance(UpperCAmelCase__ ,			(int, float)):  # Scalar multiplication
                  A__							=   Matrix(self.row ,			self.column)
                  for r in range(self.row):
                         for c in range(self.column):
                                A__							=   self[r, c] * another
                  return result
           elif isinstance(UpperCAmelCase__ ,			UpperCAmelCase__):  # Matrix multiplication
                  assert self.column == another.row
                  A__							=   Matrix(self.row ,			another.column)
                  for r in range(self.row):
                         for c in range(another.column):
                                for i in range(self.column):
                                       result[r, c] += self[r, i] * another[i, c]
                  return result
           else:
                  A__							=   f"""Unsupported type given for another ({type(UpperCAmelCase__)})"""
                  raise TypeError(UpperCAmelCase__)


    def 	SCREAMING_SNAKE_CASE							(      self				:     List[str])						->Matrix:





           '''simple docstring'''

           A__							=   Matrix(self.column ,			self.row)
           for r in range(self.row):
                  for c in range(self.column):
                         A__							=   self[r, c]
           return result





    def 	SCREAMING_SNAKE_CASE							(      self				:     Tuple ,			UpperCAmelCase__				:     Matrix ,			UpperCAmelCase__				:     Matrix)						->Any:





           '''simple docstring'''

           assert isinstance(UpperCAmelCase__ ,			UpperCAmelCase__) and isinstance(UpperCAmelCase__ ,			UpperCAmelCase__)
           assert self.row == self.column == u.row == v.row  # u, v should be column vector
           assert u.column == v.column == 1  # u, v should be column vector

           # Calculate
           A__							=   v.transpose()
           A__							=   (v_t * self * u)[0, 0] + 1
           if numerator_factor == 0:
                  return None  # It's not invertable
           return self - ((self * u) * (v_t * self) * (1.0 / numerator_factor))


# Testing
if __name__ == "__main__":







  def SCREAMING_SNAKE_CASE    (	)       ->	None:



         """simple docstring"""

         A__							=   Matrix(3			,    3			,    0						)
         for i in range(3						):
                A__							=   1
         print(f"""a^(-1) is {ainv}"""						)
         # u, v
         A__							=   Matrix(3			,    1			,    0						)
         A__			,							A__			,							A__							=   1, 2, -3
         A__							=   Matrix(3			,    1			,    0						)
         A__			,							A__			,							A__							=   4, -2, 5
         print(f"""u is {u}"""						)
         print(f"""v is {v}"""						)
         print(f"""uv^T is {u * v.transpose()}"""						)
         # Sherman Morrison
         print(f"""(a + uv^T)^(-1) is {ainv.sherman_morrison(lowercase_			,    lowercase_						)}"""						)







  def SCREAMING_SNAKE_CASE    (	)       ->	None:



         """simple docstring"""

         import doctest

         doctest.testmod()
  testa()


