



import gc
import unittest

import numpy as np
import torch

from diffusers import DDPMScheduler, MidiProcessor, SpectrogramDiffusionPipeline
from diffusers.pipelines.spectrogram_diffusion import SpectrogramContEncoder, SpectrogramNotesEncoder, TaFilmDecoder
from diffusers.utils import require_torch_gpu, skip_mps, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_note_seq, require_onnxruntime

from ..pipeline_params import TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS, TOKENS_TO_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


_A				:							str							= """./tests/fixtures/elise_format0.mid"""



@unittest.skip("""The note-seq package currently throws an error on import"""						)
class 				__snake_case  (     lowerCAmelCase_ ,    unittest.TestCase						):


    '''simple docstring'''







    lowerCamelCase__						:		Tuple							=  SpectrogramDiffusionPipeline
    lowerCamelCase__						:		int							=  PipelineTesterMixin.required_optional_params - {
        """callback""",
        """latents""",
        """callback_steps""",
        """output_type""",
        """num_images_per_prompt""",
    }
    lowerCamelCase__						:		Any							=  False
    lowerCamelCase__						:		str							=  False
    lowerCamelCase__						:		Tuple							=  TOKENS_TO_AUDIO_GENERATION_PARAMS
    lowerCamelCase__						:		int							=  TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS



    def    lowercase_				(					self    ):





        '''simple docstring'''

        torch.manual_seed(0    )
        SCREAMING_SNAKE_CASE__       							= SpectrogramNotesEncoder(
            max_length=20_48      ,							vocab_size=15_36      ,							d_model=7_68      ,							dropout_rate=0.1      ,							num_layers=1      ,							num_heads=1      ,							d_kv=4      ,							d_ff=20_48      ,							feed_forward_proj='''gated-gelu'''      ,							)

        SCREAMING_SNAKE_CASE__       							= SpectrogramContEncoder(
            input_dims=1_28      ,							targets_context_length=2_56      ,							d_model=7_68      ,							dropout_rate=0.1      ,							num_layers=1      ,							num_heads=1      ,							d_kv=4      ,							d_ff=20_48      ,							feed_forward_proj='''gated-gelu'''      ,							)

        SCREAMING_SNAKE_CASE__       							= TaFilmDecoder(
            input_dims=1_28      ,							targets_length=2_56      ,							max_decoder_noise_time=2_00_00.0      ,							d_model=7_68      ,							num_layers=1      ,							num_heads=1      ,							d_kv=4      ,							d_ff=20_48      ,							dropout_rate=0.1      ,							)

        SCREAMING_SNAKE_CASE__       							= DDPMScheduler()

        SCREAMING_SNAKE_CASE__       							= {
            '''notes_encoder''': notes_encoder.eval(),
            '''continuous_encoder''': continuous_encoder.eval(),
            '''decoder''': decoder.eval(),
            '''scheduler''': scheduler,
            '''melgan''': None,
        }
        return components



    def    lowercase_				(					self      ,							A_      ,							A_=0    ):





        '''simple docstring'''

        if str(__SCREAMING_SNAKE_CASE    ).startswith('''mps'''    ):
            SCREAMING_SNAKE_CASE__       							= torch.manual_seed(__SCREAMING_SNAKE_CASE    )
        else:
            SCREAMING_SNAKE_CASE__       							= torch.Generator(device=__SCREAMING_SNAKE_CASE    ).manual_seed(__SCREAMING_SNAKE_CASE    )
        SCREAMING_SNAKE_CASE__       							= {
            '''input_tokens''': [
                [11_34, 90, 11_35, 11_33, 10_80, 1_12, 11_32, 10_80, 11_33, 10_79, 1_33, 11_32, 10_79, 11_33, 1] + [0] * 20_33
            ],
            '''generator''': generator,
            '''num_inference_steps''': 4,
            '''output_type''': '''mel''',
        }
        return inputs



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        SCREAMING_SNAKE_CASE__       							= self.get_dummy_components()
        SCREAMING_SNAKE_CASE__       							= SpectrogramDiffusionPipeline(**__SCREAMING_SNAKE_CASE    )
        SCREAMING_SNAKE_CASE__       							= pipe.to(__SCREAMING_SNAKE_CASE    )
        pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE    )

        SCREAMING_SNAKE_CASE__       							= self.get_dummy_inputs(__SCREAMING_SNAKE_CASE    )
        SCREAMING_SNAKE_CASE__       							= pipe(**__SCREAMING_SNAKE_CASE    )
        SCREAMING_SNAKE_CASE__       							= output.audios

        SCREAMING_SNAKE_CASE__       							= mel[0, -3:, -3:]

        assert mel_slice.shape == (3, 3)
        SCREAMING_SNAKE_CASE__       							= np.array(
            [-11.512925, -4.788215, -0.46172905, -2.051715, -10.539147, -10.970963, -9.091634, 4.0, 4.0]    )
        assert np.abs(mel_slice.flatten() - expected_slice    ).max() < 1E-2



    @skip_mps
    def    lowercase_				(					self    ):





        '''simple docstring'''

        return super().test_save_load_local()



    @skip_mps
    def    lowercase_				(					self    ):





        '''simple docstring'''

        return super().test_dict_tuple_outputs_equivalent()



    @skip_mps
    def    lowercase_				(					self    ):





        '''simple docstring'''

        return super().test_save_load_optional_components()



    @skip_mps
    def    lowercase_				(					self    ):





        '''simple docstring'''

        return super().test_attention_slicing_forward_pass()



    def    lowercase_				(					self    ):





        '''simple docstring'''

        pass



    def    lowercase_				(					self    ):





        '''simple docstring'''

        pass




    @skip_mps
    def    lowercase_				(					self    ):





        '''simple docstring'''

        return super().test_progress_bar()





@slow
@require_torch_gpu
@require_onnxruntime
@require_note_seq
class 				__snake_case  (     unittest.TestCase						):


    '''simple docstring'''







    def    lowercase_				(					self    ):





        '''simple docstring'''

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



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= torch_device

        SCREAMING_SNAKE_CASE__       							= SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''    )
        SCREAMING_SNAKE_CASE__       							= pipe.melgan
        SCREAMING_SNAKE_CASE__       							= None

        SCREAMING_SNAKE_CASE__       							= pipe.to(__SCREAMING_SNAKE_CASE    )
        pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE    )

        def callback(A_      ,							A_    ):
            # decode mel to audio
            SCREAMING_SNAKE_CASE__       							= melgan(input_features=mel_output.astype(np.floataa    )    )[0]
            assert len(audio[0]    ) == 8_19_20 * (step + 1)
            # simulate that audio is played
            return audio

        SCREAMING_SNAKE_CASE__       							= MidiProcessor()
        SCREAMING_SNAKE_CASE__       							= processor(__SCREAMING_SNAKE_CASE    )

        SCREAMING_SNAKE_CASE__       							= input_tokens[:3]
        SCREAMING_SNAKE_CASE__       							= torch.manual_seed(0    )
        pipe(__SCREAMING_SNAKE_CASE      ,							num_inference_steps=5      ,							generator=__SCREAMING_SNAKE_CASE      ,							callback=__SCREAMING_SNAKE_CASE      ,							output_type='''mel'''    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= torch_device

        SCREAMING_SNAKE_CASE__       							= SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''    )
        SCREAMING_SNAKE_CASE__       							= pipe.to(__SCREAMING_SNAKE_CASE    )
        pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE    )
        SCREAMING_SNAKE_CASE__       							= MidiProcessor()

        SCREAMING_SNAKE_CASE__       							= processor(__SCREAMING_SNAKE_CASE    )
        # just run two denoising loops
        SCREAMING_SNAKE_CASE__       							= input_tokens[:2]

        SCREAMING_SNAKE_CASE__       							= torch.manual_seed(0    )
        SCREAMING_SNAKE_CASE__       							= pipe(__SCREAMING_SNAKE_CASE      ,							num_inference_steps=2      ,							generator=__SCREAMING_SNAKE_CASE    )

        SCREAMING_SNAKE_CASE__       							= output.audios[0]

        assert abs(np.abs(__SCREAMING_SNAKE_CASE    ).sum() - 3612.841    ) < 1E-1




    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= torch_device

        SCREAMING_SNAKE_CASE__       							= SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''    )
        SCREAMING_SNAKE_CASE__       							= pipe.to(__SCREAMING_SNAKE_CASE    )
        pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE    )

        SCREAMING_SNAKE_CASE__       							= MidiProcessor()

        SCREAMING_SNAKE_CASE__       							= processor(__SCREAMING_SNAKE_CASE    )

        # just run 4 denoising loops
        SCREAMING_SNAKE_CASE__       							= input_tokens[:4]

        SCREAMING_SNAKE_CASE__       							= torch.manual_seed(0    )
        SCREAMING_SNAKE_CASE__       							= pipe(__SCREAMING_SNAKE_CASE      ,							num_inference_steps=1_00      ,							generator=__SCREAMING_SNAKE_CASE    )

        SCREAMING_SNAKE_CASE__       							= output.audios[0]
        assert abs(np.abs(__SCREAMING_SNAKE_CASE    ).sum() - 9389.1111    ) < 5E-2







from __future__ import annotations

from collections.abc import Generator


def        __snake_case       (						lowerCAmelCase_  )							->	Generator[int, None, None]:
    if not isinstance(lowerCAmelCase_   ,       lowerCAmelCase_  ) or n < 1:
        raise Exception('''Sequence only defined for positive integers'''  )

    yield n
    while n != 1:
        if n % 2 == 0:
            n //= 2
        else:
            SCREAMING_SNAKE_CASE__       							= 3 * n + 1
        yield n


def        __snake_case       (						)							->	str:
    SCREAMING_SNAKE_CASE__       							= 4_3
    SCREAMING_SNAKE_CASE__       							= tuple(collatz_sequence(lowerCAmelCase_  )  )
    print(lowerCAmelCase_  )
    print(f'''Collatz sequence from {n} took {len(lowerCAmelCase_  )} steps.'''  )


if __name__ == "__main__":
  main()
