

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__										=      "./tests/fixtures/elise_format0.mid"







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


  """simple docstring"""
  SCREAMING_SNAKE_CASE__					:			List[Any]					    =   SpectrogramDiffusionPipeline
  SCREAMING_SNAKE_CASE__					:			Dict					    =   PipelineTesterMixin.required_optional_params - {
      """callback""",
      """latents""",
      """callback_steps""",
      """output_type""",
      """num_images_per_prompt""",
  }
  SCREAMING_SNAKE_CASE__					:			List[Any]					    =   False
  SCREAMING_SNAKE_CASE__					:			int					    =   False
  SCREAMING_SNAKE_CASE__					:			List[Any]					    =   TOKENS_TO_AUDIO_GENERATION_PARAMS
  SCREAMING_SNAKE_CASE__					:			int					    =   TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS



  def   _UpperCamelCase     (						self     :					List[Any]						)       ->			int:
      torch.manual_seed(0						)
      SCREAMING_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"  ,	)

      SCREAMING_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"  ,	)

      SCREAMING_SNAKE_CASE__         =    TaFilmDecoder(
          input_dims=128  ,	targets_length=256  ,	max_decoder_noise_time=20_000.0  ,	d_model=768  ,	num_layers=1  ,	num_heads=1  ,	d_kv=4  ,	d_ff=2048  ,	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   _UpperCamelCase     (						self     :					int  ,	A_     :					str  ,	A_     :					Dict=0						)       ->			Optional[int]:
      if str(A_						).startswith("mps"						):
          SCREAMING_SNAKE_CASE__         =    torch.manual_seed(A_						)
      else:
          SCREAMING_SNAKE_CASE__         =    torch.Generator(device=A_						).manual_seed(A_						)
      SCREAMING_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   _UpperCamelCase     (						self     :					Dict						)       ->			Tuple:
      SCREAMING_SNAKE_CASE__         =    "cpu"  # ensure determinism for the device-dependent torch.Generator
      SCREAMING_SNAKE_CASE__         =    self.get_dummy_components()
      SCREAMING_SNAKE_CASE__         =    SpectrogramDiffusionPipeline(**A_						)
      SCREAMING_SNAKE_CASE__         =    pipe.to(A_						)
      pipe.set_progress_bar_config(disable=A_						)

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

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

      assert mel_slice.shape == (3, 3)
      SCREAMING_SNAKE_CASE__         =    np.array(
          [-11.512_925, -4.788_215, -0.46_172_905, -2.051_715, -10.539_147, -10.970_963, -9.091_634, 4.0, 4.0]						)
      assert np.abs(mel_slice.flatten() - expected_slice						).max() < 1E-2



  @skip_mps
  def   _UpperCamelCase     (						self     :					Tuple						)       ->			Union[str, Any]:
      return super().test_save_load_local()



  @skip_mps
  def   _UpperCamelCase     (						self     :					List[str]						)       ->			str:
      return super().test_dict_tuple_outputs_equivalent()



  @skip_mps
  def   _UpperCamelCase     (						self     :					Tuple						)       ->			Optional[int]:
      return super().test_save_load_optional_components()



  @skip_mps
  def   _UpperCamelCase     (						self     :					Optional[int]						)       ->			Tuple:
      return super().test_attention_slicing_forward_pass()



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



  def   _UpperCamelCase     (						self     :					List[Any]						)       ->			Union[str, Any]:
      pass







  @skip_mps
  def   _UpperCamelCase     (						self     :					int						)       ->			Dict:
      return super().test_progress_bar()







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


  """simple docstring"""



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



  def   _UpperCamelCase     (						self     :					Tuple						)       ->			Tuple:
      # TODO - test that pipeline can decode tokens in a callback
      # so that music can be played live
      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(A_						)
      pipe.set_progress_bar_config(disable=A_						)

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

      SCREAMING_SNAKE_CASE__         =    MidiProcessor()
      SCREAMING_SNAKE_CASE__         =    processor(A_						)

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



  def   _UpperCamelCase     (						self     :					Tuple						)       ->			int:
      SCREAMING_SNAKE_CASE__         =    torch_device

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

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

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

      SCREAMING_SNAKE_CASE__         =    output.audios[0]

      assert abs(np.abs(A_						).sum() - 3_612.841						) < 1E-1







  def   _UpperCamelCase     (						self     :					str						)       ->			Optional[Any]:
      SCREAMING_SNAKE_CASE__         =    torch_device

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

      SCREAMING_SNAKE_CASE__         =    MidiProcessor()

      SCREAMING_SNAKE_CASE__         =    processor(A_						)

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

      SCREAMING_SNAKE_CASE__         =    torch.manual_seed(0						)
      SCREAMING_SNAKE_CASE__         =    pipe(A_  ,	num_inference_steps=100  ,	generator=A_						)

      SCREAMING_SNAKE_CASE__         =    output.audios[0]
      assert abs(np.abs(A_						).sum() - 9_389.1_111						) < 5E-2





from typing import SupportsIndex

import numpy as np
from scipy.ndimage import convolve





def 				__a							(       __A			,							__A			,							__A						):
    SCREAMING_SNAKE_CASE__         =    np.stack((horizontal_flow, vertical_flow)			,							2						)

    # Create a grid of all pixel coordinates and subtract the flow to get the
    # target pixels coordinates
    SCREAMING_SNAKE_CASE__         =    np.stack(
        np.meshgrid(np.arange(0			,							image.shape[1]						)			,							np.arange(0			,							image.shape[0]						)						)			,							2						)
    SCREAMING_SNAKE_CASE__         =    np.round(grid - flow						).astype(np.intaa						)

    # Find the locations outside of the original image
    SCREAMING_SNAKE_CASE__         =    (grid < 0) | (grid >= np.array([image.shape[1], image.shape[0]]						))
    SCREAMING_SNAKE_CASE__         =    0

    SCREAMING_SNAKE_CASE__         =    image[grid[:, :, 1], grid[:, :, 0]]

    # Set pixels at invalid locations to 0
    SCREAMING_SNAKE_CASE__         =    0

    return warped





def 				__a							(       __A			,							__A			,							__A			,							__A = None			,							):
    if alpha is None:
        SCREAMING_SNAKE_CASE__         =    0.1

    # Initialize flow
    SCREAMING_SNAKE_CASE__         =    np.zeros_like(__A						)
    SCREAMING_SNAKE_CASE__         =    np.zeros_like(__A						)

    # Prepare kernels for the calculation of the derivatives and the average velocity
    SCREAMING_SNAKE_CASE__         =    np.array([[-1, 1], [-1, 1]]						) * 0.25
    SCREAMING_SNAKE_CASE__         =    np.array([[-1, -1], [1, 1]]						) * 0.25
    SCREAMING_SNAKE_CASE__         =    np.array([[1, 1], [1, 1]]						) * 0.25
    SCREAMING_SNAKE_CASE__         =    np.array(
        [[1 / 12, 1 / 6, 1 / 12], [1 / 6, 0, 1 / 6], [1 / 12, 1 / 6, 1 / 12]]						)

    # Iteratively refine the flow
    for _ in range(__A						):
        SCREAMING_SNAKE_CASE__         =    warp(__A			,							__A			,							__A						)
        SCREAMING_SNAKE_CASE__         =    convolve(__A			,							__A						) + convolve(__A			,							__A						)
        SCREAMING_SNAKE_CASE__         =    convolve(__A			,							__A						) + convolve(__A			,							__A						)
        SCREAMING_SNAKE_CASE__         =    convolve(__A			,							__A						) + convolve(__A			,							-kernel_t						)

        SCREAMING_SNAKE_CASE__         =    convolve(__A			,							__A						)
        SCREAMING_SNAKE_CASE__         =    convolve(__A			,							__A						)

        # This updates the flow as proposed in the paper (Step 12)
        SCREAMING_SNAKE_CASE__         =    (
            derivative_x * avg_horizontal_velocity
            + derivative_y * avg_vertical_velocity
            + derivative_t
        )
        SCREAMING_SNAKE_CASE__         =    update / (alpha**2 + derivative_x**2 + derivative_y**2)

        SCREAMING_SNAKE_CASE__         =    avg_horizontal_velocity - derivative_x * update
        SCREAMING_SNAKE_CASE__         =    avg_vertical_velocity - derivative_y * update

    return horizontal_flow, vertical_flow


if __name__ == "__main__":
  import doctest

  doctest.testmod()


