

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()


lowerCamelCase_					  =   '''./tests/fixtures/elise_format0.mid'''







@unittest.skip('The note-seq package currently throws an error on import'      )
class 					__lowerCamelCase							(      __snake_case    ,				unittest.TestCase      ):
      lowerCamelCase_      : List[str]  			=   SpectrogramDiffusionPipeline
      lowerCamelCase_      : Optional[int]  			=   PipelineTesterMixin.required_optional_params - {
          'callback',
          'latents',
          'callback_steps',
          'output_type',
          'num_images_per_prompt',
      }
      lowerCamelCase_      : Any  			=   False
      lowerCamelCase_      : int  			=   False
      lowerCamelCase_      : Union[str, Any]  			=   TOKENS_TO_AUDIO_GENERATION_PARAMS
      lowerCamelCase_      : Union[str, Any]  			=   TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS







      def       lowerCAmelCase_       (	self  )		->			List[Any]:
        torch.manual_seed(0  )
        snake_case_					= SpectrogramNotesEncoder(
            max_length=2048				,    vocab_size=1536				,    d_model=768				,    dropout_rate=0.1				,    num_layers=1				,    num_heads=1				,    d_kv=4				,    d_ff=2048				,    feed_forward_proj="""gated-gelu"""				,    )

        snake_case_					= SpectrogramContEncoder(
            input_dims=128				,    targets_context_length=256				,    d_model=768				,    dropout_rate=0.1				,    num_layers=1				,    num_heads=1				,    d_kv=4				,    d_ff=2048				,    feed_forward_proj="""gated-gelu"""				,    )

        snake_case_					= TaFilmDecoder(
            input_dims=128				,    targets_length=256				,    max_decoder_noise_time=2_0000.0				,    d_model=768				,    num_layers=1				,    num_heads=1				,    d_kv=4				,    d_ff=2048				,    dropout_rate=0.1				,    )

        snake_case_					= DDPMScheduler()

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







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase=0  )		->			Optional[Any]:
        if str(lowerCamelCase  ).startswith("""mps"""  ):
          snake_case_					= torch.manual_seed(lowerCamelCase  )
        else:
          snake_case_					= torch.Generator(device=lowerCamelCase  ).manual_seed(lowerCamelCase  )
        snake_case_					= {
            """input_tokens""": [
                [1134, 90, 1135, 1133, 1080, 112, 1132, 1080, 1133, 1079, 133, 1132, 1079, 1133, 1] + [0] * 2033
            ],
            """generator""": generator,
            """num_inference_steps""": 4,
            """output_type""": """mel""",
        }
        return inputs







      def       lowerCAmelCase_       (	self  )		->			Optional[int]:
        snake_case_					= """cpu"""  # ensure determinism for the device-dependent torch.Generator
        snake_case_					= self.get_dummy_components()
        snake_case_					= SpectrogramDiffusionPipeline(**lowerCamelCase  )
        snake_case_					= pipe.to(lowerCamelCase  )
        pipe.set_progress_bar_config(disable=lowerCamelCase  )

        snake_case_					= self.get_dummy_inputs(lowerCamelCase  )
        snake_case_					= pipe(**lowerCamelCase  )
        snake_case_					= output.audios

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

        assert mel_slice.shape == (3, 3)
        snake_case_					= np.array(
            [-11.51_2925, -4.78_8215, -0.4617_2905, -2.05_1715, -10.53_9147, -10.97_0963, -9.09_1634, 4.0, 4.0]  )
        assert np.abs(mel_slice.flatten() - expected_slice  ).max() < 1e-2







      @skip_mps
      def       lowerCAmelCase_       (	self  )		->			Dict:
        return super().test_save_load_local()







      @skip_mps
      def       lowerCAmelCase_       (	self  )		->			str:
        return super().test_dict_tuple_outputs_equivalent()







      @skip_mps
      def       lowerCAmelCase_       (	self  )		->			List[Any]:
        return super().test_save_load_optional_components()







      @skip_mps
      def       lowerCAmelCase_       (	self  )		->			List[Any]:
        return super().test_attention_slicing_forward_pass()







      def       lowerCAmelCase_       (	self  )		->			List[Any]:
        pass







      def       lowerCAmelCase_       (	self  )		->			Tuple:
        pass







      @skip_mps
      def       lowerCAmelCase_       (	self  )		->			Any:
        return super().test_progress_bar()







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







      def       lowerCAmelCase_       (	self  )		->			Optional[Any]:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







      def       lowerCAmelCase_       (	self  )		->			List[Any]:
        # TODO - test that pipeline can decode tokens in a callback
        # so that music can be played live
        snake_case_					= torch_device

        snake_case_					= SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""  )
        snake_case_					= pipe.melgan
        snake_case_					= None

        snake_case_					= pipe.to(lowerCamelCase  )
        pipe.set_progress_bar_config(disable=lowerCamelCase  )

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

        snake_case_					= MidiProcessor()
        snake_case_					= processor(lowerCamelCase  )

        snake_case_					= input_tokens[:3]
        snake_case_					= torch.manual_seed(0  )
        pipe(lowerCamelCase				,    num_inference_steps=5				,    generator=lowerCamelCase				,    callback=lowerCamelCase				,    output_type="""mel"""  )







      def       lowerCAmelCase_       (	self  )		->			Optional[Any]:
        snake_case_					= torch_device

        snake_case_					= SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""  )
        snake_case_					= pipe.to(lowerCamelCase  )
        pipe.set_progress_bar_config(disable=lowerCamelCase  )
        snake_case_					= MidiProcessor()

        snake_case_					= processor(lowerCamelCase  )
        # just run two denoising loops
        snake_case_					= input_tokens[:2]

        snake_case_					= torch.manual_seed(0  )
        snake_case_					= pipe(lowerCamelCase				,    num_inference_steps=2				,    generator=lowerCamelCase  )

        snake_case_					= output.audios[0]

        assert abs(np.abs(lowerCamelCase  ).sum() - 3612.841  ) < 1e-1







      def       lowerCAmelCase_       (	self  )		->			int:
        snake_case_					= torch_device

        snake_case_					= SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""  )
        snake_case_					= pipe.to(lowerCamelCase  )
        pipe.set_progress_bar_config(disable=lowerCamelCase  )

        snake_case_					= MidiProcessor()

        snake_case_					= processor(lowerCamelCase  )

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

        snake_case_					= torch.manual_seed(0  )
        snake_case_					= pipe(lowerCamelCase				,    num_inference_steps=100				,    generator=lowerCamelCase  )

        snake_case_					= output.audios[0]
        assert abs(np.abs(lowerCamelCase  ).sum() - 9389.1111  ) < 5e-2


from dataclasses import dataclass
from typing import Optional, Tuple, Union

import flax
import jax.numpy as jnp
from scipy import integrate

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils_flax import (
    CommonSchedulerState,
    FlaxKarrasDiffusionSchedulers,
    FlaxSchedulerMixin,
    FlaxSchedulerOutput,
    broadcast_to_shape_from_left,
)







@flax.struct.dataclass
class 					__lowerCamelCase							:
      lowerCamelCase_      : CommonSchedulerState

      # setable values
      lowerCamelCase_      : jnp.ndarray
      lowerCamelCase_      : jnp.ndarray
      lowerCamelCase_      : jnp.ndarray
      lowerCamelCase_      : Optional[int]  			=   None

      # running values
      lowerCamelCase_      : Optional[jnp.ndarray]  			=   None







      @classmethod
      def       lowerCAmelCase_       (	cls				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )		->			Union[str, Any]:
        return cls(common=lowerCamelCase				,    init_noise_sigma=lowerCamelCase				,    timesteps=lowerCamelCase				,    sigmas=lowerCamelCase  )







@dataclass
class 					__lowerCamelCase							(      __snake_case      ):
      lowerCamelCase_      : LMSDiscreteSchedulerState







class 					__lowerCamelCase							(      __snake_case    ,				__snake_case      ):
      lowerCamelCase_      : str  			=   [e.name for e in FlaxKarrasDiffusionSchedulers]

      lowerCamelCase_      : jnp.dtype







      @property
      def       lowerCAmelCase_       (	self  )		->			List[str]:
        return True







      @register_to_config
      def __init__(	self				,    lowerCamelCase = 1000				,    lowerCamelCase = 0.0001				,    lowerCamelCase = 0.02				,    lowerCamelCase = "linear"				,    lowerCamelCase = None				,    lowerCamelCase = "epsilon"				,    lowerCamelCase = jnp.floataa				,    )		->			Optional[Any]:
        snake_case_					= dtype







      def       lowerCAmelCase_       (	self				,    lowerCamelCase = None  )		->			LMSDiscreteSchedulerState:
        if common is None:
          snake_case_					= CommonSchedulerState.create(self  )

        snake_case_					= jnp.arange(0				,    self.config.num_train_timesteps  ).round()[::-1]
        snake_case_					= ((1 - common.alphas_cumprod) / common.alphas_cumprod) ** 0.5

        # standard deviation of the initial noise distribution
        snake_case_					= sigmas.max()

        return LMSDiscreteSchedulerState.create(
            common=lowerCamelCase				,    init_noise_sigma=lowerCamelCase				,    timesteps=lowerCamelCase				,    sigmas=lowerCamelCase				,    )







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )		->			jnp.ndarray:
        ((snake_case_)		,							)					= jnp.where(state.timesteps == timestep				,    size=1  )
        snake_case_					= step_index[0]

        snake_case_					= state.sigmas[step_index]
        snake_case_					= sample / ((sigma**2 + 1) ** 0.5)
        return sample







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )		->			Any:

        def lms_derivative(lowerCamelCase  ):
          snake_case_					= 1.0
          for k in range(lowerCamelCase  ):
            if current_order == k:
              continue
            prod *= (tau - state.sigmas[t - k]) / (state.sigmas[t - current_order] - state.sigmas[t - k])
          return prod

        snake_case_					= integrate.quad(lowerCamelCase				,    state.sigmas[t]				,    state.sigmas[t + 1]				,    epsrel=1e-4  )[0]

        return integrated_coeff







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase = ()  )		->			LMSDiscreteSchedulerState:
        snake_case_					= jnp.linspace(self.config.num_train_timesteps - 1				,    0				,    lowerCamelCase				,    dtype=self.dtype  )

        snake_case_					= jnp.floor(lowerCamelCase  ).astype(jnp.intaa  )
        snake_case_					= jnp.ceil(lowerCamelCase  ).astype(jnp.intaa  )

        snake_case_					= jnp.mod(lowerCamelCase				,    1.0  )

        snake_case_					= ((1 - state.common.alphas_cumprod) / state.common.alphas_cumprod) ** 0.5
        snake_case_					= (1 - frac) * sigmas[low_idx] + frac * sigmas[high_idx]
        snake_case_					= jnp.concatenate([sigmas, jnp.array([0.0]				,    dtype=self.dtype  )]  )

        snake_case_					= timesteps.astype(jnp.intaa  )

        # initial running values
        snake_case_					= jnp.zeros((0,) + shape				,    dtype=self.dtype  )

        return state.replace(
            timesteps=lowerCamelCase				,    sigmas=lowerCamelCase				,    num_inference_steps=lowerCamelCase				,    derivatives=lowerCamelCase				,    )







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase = 4				,    lowerCamelCase = True				,    )		->			Union[FlaxLMSSchedulerOutput, Tuple]:
        if state.num_inference_steps is None:
          raise ValueError(
              """Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"""  )

        snake_case_					= state.sigmas[timestep]

        # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
        if self.config.prediction_type == "epsilon":
          snake_case_					= sample - sigma * model_output
        elif self.config.prediction_type == "v_prediction":
          # * c_out + input * c_skip
          snake_case_					= model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))
        else:
          raise ValueError(
              f'''prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`'''  )

        # 2. Convert to an ODE derivative
        snake_case_					= (sample - pred_original_sample) / sigma
        snake_case_					= state.replace(derivatives=jnp.append(state.derivatives				,    lowerCamelCase  )  )
        if len(state.derivatives  ) > order:
          snake_case_					= state.replace(derivatives=jnp.delete(state.derivatives				,    0  )  )

        # 3. Compute linear multistep coefficients
        snake_case_					= min(timestep + 1				,    lowerCamelCase  )
        snake_case_					= [self.get_lms_coefficient(lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  ) for curr_order in range(lowerCamelCase  )]

        # 4. Compute previous sample based on the derivatives path
        snake_case_					= sample + sum(
            coeff * derivative for coeff, derivative in zip(lowerCamelCase				,    reversed(state.derivatives  )  )  )

        if not return_dict:
          return (prev_sample, state)

        return FlaxLMSSchedulerOutput(prev_sample=lowerCamelCase				,    state=lowerCamelCase  )







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    )		->			jnp.ndarray:
        snake_case_					= state.sigmas[timesteps].flatten()
        snake_case_					= broadcast_to_shape_from_left(lowerCamelCase				,    noise.shape  )

        snake_case_					= original_samples + noise * sigma

        return noisy_samples







      def __len__(	self  )		->			Optional[Any]:
        return self.config.num_train_timesteps