


'''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  ):





							'''simple docstring'''






							lowercase__      :       List[str]     						=				DanceDiffusionPipeline
							lowercase__      :       List[Any]     						=				UNCONDITIONAL_AUDIO_GENERATION_PARAMS
							lowercase__      :       Dict     						=				PipelineTesterMixin.required_optional_params - {
							    "callback",
							    "latents",
							    "callback_steps",
							    "output_type",
							    "num_images_per_prompt",
							}
							lowercase__      :       str     						=				UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
							lowercase__      :       Union[str, Any]     						=				False
							lowercase__      :       Optional[Any]     						=				False







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[Any]:
														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=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''') ,   )
														lowerCAmelCase__	       =     IPNDMScheduler()

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







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_ ,   lowerCamelCase_=0       )	->							Dict:
														if str(lowerCamelCase_       ).startswith('''mps'''       ):
																					lowerCAmelCase__	       =     torch.manual_seed(lowerCamelCase_       )
														else:
																					lowerCAmelCase__	       =     torch.Generator(device=lowerCamelCase_       ).manual_seed(lowerCamelCase_       )
														lowerCAmelCase__	       =     {
														    '''batch_size''': 1,
														    '''generator''': generator,
														    '''num_inference_steps''': 4,
														}
														return inputs







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[Any]:
														lowerCAmelCase__	       =     '''cpu'''  # ensure determinism for the device-dependent torch.Generator
														lowerCAmelCase__	       =     self.get_dummy_components()
														lowerCAmelCase__	       =     DanceDiffusionPipeline(**lowerCamelCase_       )
														lowerCAmelCase__	       =     pipe.to(lowerCamelCase_       )
														pipe.set_progress_bar_config(disable=lowerCamelCase_       )

														lowerCAmelCase__	       =     self.get_dummy_inputs(lowerCamelCase_       )
														lowerCAmelCase__	       =     pipe(**lowerCamelCase_       )
														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   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[Any]:
														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       )	->							int:
														return super().test_save_load_optional_components()







							@skip_mps
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							str:
														return super().test_attention_slicing_forward_pass()







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[int]:
														super().test_inference_batch_single_identical(expected_max_diff=3e-3       )




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





							'''simple docstring'''






							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       )	->							Optional[Any]:
														lowerCAmelCase__	       =     torch_device

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

														lowerCAmelCase__	       =     torch.manual_seed(0       )
														lowerCAmelCase__	       =     pipe(generator=lowerCamelCase_ ,   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   __SCREAMING_SNAKE_CASE							(		self       )	->							Any:
														lowerCAmelCase__	       =     torch_device

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

														lowerCAmelCase__	       =     torch.manual_seed(0       )
														lowerCAmelCase__	       =     pipe(generator=lowerCamelCase_ ,   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



'''simple docstring'''






from __future__ import annotations

import numpy as np
from numpy import floataa
from numpy.typing import NDArray


def  _snake_case       (					A			,   A			,   A			,   A			,   )						->   list[float]:
							lowerCAmelCase__       , lowerCAmelCase__	       =     coefficient_matrix.shape
							lowerCAmelCase__       , lowerCAmelCase__	       =     constant_matrix.shape

							if rowsa != colsa:
														lowerCAmelCase__	       =     F"""Coefficient matrix dimensions must be nxn but received {rowsa}x{colsa}"""
														raise ValueError(A	)

							if colsa != 1:
														lowerCAmelCase__	       =     F"""Constant matrix must be nx1 but received {rowsa}x{colsa}"""
														raise ValueError(A	)

							if rowsa != rowsa:
														lowerCAmelCase__	       =     (
														    '''Coefficient and constant matrices dimensions must be nxn and nx1 but '''
														    F"""received {rowsa}x{colsa} and {rowsa}x{colsa}"""
														)
														raise ValueError(A	)

							if len(A	) != rowsa:
														lowerCAmelCase__	       =     (
														    '''Number of initial values must be equal to number of rows in coefficient '''
														    F"""matrix but received {len(A	)} and {rowsa}"""
														)
														raise ValueError(A	)

							if iterations <= 0:
														raise ValueError('''Iterations must be at least 1'''	)

							lowerCAmelCase__	       =     np.concatenate(
							    (coefficient_matrix, constant_matrix)			,   axis=1	)

							lowerCAmelCase__       , lowerCAmelCase__	       =     table.shape

							strictly_diagonally_dominant(A	)

							# Iterates the whole matrix for given number of times
							for _ in range(A	):
														lowerCAmelCase__	       =     []
														for row in range(A	):
																					lowerCAmelCase__	       =     0
																					for col in range(A	):
																												if col == row:
																																			lowerCAmelCase__	       =     table[row][col]
																												elif col == cols - 1:
																																			lowerCAmelCase__	       =     table[row][col]
																												else:
																																			temp += (-1) * table[row][col] * init_val[col]
																					lowerCAmelCase__	       =     (temp + val) / denom
																					new_val.append(A	)
														lowerCAmelCase__	       =     new_val

							return [float(A	) for i in new_val]


def  _snake_case       (					A	)						->   bool:
							lowerCAmelCase__       , lowerCAmelCase__	       =     table.shape

							lowerCAmelCase__	       =     True

							for i in range(0			,   A	):
														lowerCAmelCase__	       =     0
														for j in range(0			,   cols - 1	):
																					if i == j:
																												continue
																					else:
																												total += table[i][j]

														if table[i][i] <= total:
																					raise ValueError('''Coefficient matrix is not strictly diagonally dominant'''	)

							return is_diagonally_dominant


# Test Cases
if __name__ == "__main__":
		import doctest

		doctest.testmod()