
'''simple docstring'''
import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device

def   UpperCamelCase_       (		__SCREAMING_SNAKE_CASE    ):
    """simple docstring"""






    snake_case_	:      List[Any]            =  {}
    for name in unet.attn_processors.keys():
        snake_case_	:      int            =  None if name.endswith("attn1.processor"    ) else unet.config.cross_attention_dim
        if name.startswith("mid_block"    ):
            snake_case_	:      Optional[Any]            =  unet.config.block_out_channels[-1]
        elif name.startswith("up_blocks"    ):
            snake_case_	:      Any            =  int(name[len("up_blocks."    )]    )
            snake_case_	:      int            =  list(reversed(unet.config.block_out_channels    )    )[block_id]
        elif name.startswith("down_blocks"    ):
            snake_case_	:      Union[str, Any]            =  int(name[len("down_blocks."    )]    )
            snake_case_	:      Dict            =  unet.config.block_out_channels[block_id]
        snake_case_	:      List[str]            =  (
            LoRAAttnProcessor2_0 if hasattr(lowercase_,						"scaled_dot_product_attention"    ) else LoRAAttnProcessor
        )
        snake_case_	:      List[Any]            =  lora_attn_processor_class(
            hidden_size=lowercase_,						cross_attention_dim=lowercase_    )
    snake_case_	:      List[str]            =  AttnProcsLayers(lowercase_    )
    return lora_attn_procs, unet_lora_layers

def   UpperCamelCase_       (		__SCREAMING_SNAKE_CASE    ):
    """simple docstring"""






    snake_case_	:      Dict            =  {}
    snake_case_	:      Optional[Any]            =  (
        LoRAAttnProcessor2_0 if hasattr(lowercase_,						"scaled_dot_product_attention"    ) else LoRAAttnProcessor
    )
    for name, module in text_encoder.named_modules():
        if name.endswith(lowercase_    ):
            snake_case_	:      str            =  lora_attn_processor_class(
                hidden_size=module.out_proj.out_features,						cross_attention_dim=lowercase_    )
    return text_lora_attn_procs

def   UpperCamelCase_       (		__SCREAMING_SNAKE_CASE    ):
    """simple docstring"""






    snake_case_	:      Union[str, Any]            =  create_text_encoder_lora_attn_procs(lowercase_    )
    snake_case_	:      int            =  AttnProcsLayers(lowercase_    )
    return text_encoder_lora_layers

def   UpperCamelCase_       (		__SCREAMING_SNAKE_CASE,						__SCREAMING_SNAKE_CASE=False    ):
    """simple docstring"""






    for _, attn_proc in text_lora_attn_procs.items():
        # set up.weights
        for layer_name, layer_module in attn_proc.named_modules():
            if layer_name.endswith("_lora"    ):
                snake_case_	:      Optional[int]            =  (
                    torch.randn_like(layer_module.up.weight    )
                    if randn_weight
                    else torch.zeros_like(layer_module.up.weight    )
                )
                snake_case_	:      List[str]            =  torch.nn.Parameter(lowercase_    )






class 				UpperCAmelCase_			(		unittest.TestCase   ):





 def     snake_case__						(      self):
     torch.manual_seed(0)
     snake_case_	:      str            =  UNetaDConditionModel(
         block_out_channels=(32, 64)	,    layers_per_block=2	,    sample_size=32	,    in_channels=4	,    out_channels=4	,    down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")	,    up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")	,    cross_attention_dim=32	,    )
     snake_case_	:      Union[str, Any]            =  DDIMScheduler(
         beta_start=0.00_085	,    beta_end=0.012	,    beta_schedule="scaled_linear"	,    clip_sample=lowerCAmelCase__	,    set_alpha_to_one=lowerCAmelCase__	,    steps_offset=1	,    )
     torch.manual_seed(0)
     snake_case_	:      Any            =  AutoencoderKL(
         block_out_channels=[32, 64]	,    in_channels=3	,    out_channels=3	,    down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]	,    up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]	,    latent_channels=4	,    )
     snake_case_	:      Tuple            =  CLIPTextConfig(
         bos_token_id=0	,    eos_token_id=2	,    hidden_size=32	,    intermediate_size=37	,    layer_norm_eps=1E-05	,    num_attention_heads=4	,    num_hidden_layers=5	,    pad_token_id=1	,    vocab_size=10_00	,    )
     snake_case_	:      int            =  CLIPTextModel(lowerCAmelCase__)
     snake_case_	:      int            =  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

     snake_case_     ,	snake_case_	:      str            =  create_unet_lora_layers(lowerCAmelCase__)
     snake_case_	:      List[Any]            =  create_text_encoder_lora_layers(lowerCAmelCase__)

     snake_case_	:      Union[str, Any]            =  {
         "unet": unet,
         "scheduler": scheduler,
         "vae": vae,
         "text_encoder": text_encoder,
         "tokenizer": tokenizer,
         "safety_checker": None,
         "feature_extractor": None,
     }
     snake_case_	:      Optional[int]            =  {
         "unet_lora_layers": unet_lora_layers,
         "text_encoder_lora_layers": text_encoder_lora_layers,
         "unet_lora_attn_procs": unet_lora_attn_procs,
     }
     return pipeline_components, lora_components





 def     snake_case__						(      self):
     snake_case_	:      str            =  1
     snake_case_	:      List[str]            =  10
     snake_case_	:      Tuple            =  4
     snake_case_	:      List[Any]            =  (32, 32)

     snake_case_	:      Union[str, Any]            =  torch.manual_seed(0)
     snake_case_	:      Optional[Any]            =  floats_tensor((batch_size, num_channels) + sizes)
     snake_case_	:      Tuple            =  torch.randint(1	,    lowerCAmelCase__	,    size=(batch_size, sequence_length)	,    generator=lowerCAmelCase__)

     snake_case_	:      Any            =  {
         "prompt": "A painting of a squirrel eating a burger",
         "generator": generator,
         "num_inference_steps": 2,
         "guidance_scale": 6.0,
         "output_type": "numpy",
     }

     return noise, input_ids, pipeline_inputs

     # copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb





 def     snake_case__						(      self):
     snake_case_	:      Any            =  77

     snake_case_	:      Dict            =  torch.randint(2	,    56	,    size=(1, max_seq_length)	,    generator=torch.manual_seed(0))

     snake_case_	:      int            =  {}
     snake_case_	:      Any            =  inputs
     return prepared_inputs





 def     snake_case__						(      self	,    lowercase_):
     snake_case_     ,	snake_case_	:      str            =  self.get_dummy_components()
     LoraLoaderMixin.save_lora_weights(
         save_directory=lowerCAmelCase__	,    unet_lora_layers=lora_components["unet_lora_layers"]	,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]	,    )
     self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__	,    "pytorch_lora_weights.bin")))





 def     snake_case__						(      self):
     snake_case_     ,	snake_case_	:      Optional[Any]            =  self.get_dummy_components()
     snake_case_	:      List[str]            =  StableDiffusionPipeline(**lowerCAmelCase__)
     snake_case_	:      Tuple            =  sd_pipe.to(lowerCAmelCase__)
     sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

     snake_case_     ,	snake_case_     ,	snake_case_	:      int            =  self.get_dummy_inputs()

     snake_case_	:      Optional[Any]            =  sd_pipe(**lowerCAmelCase__).images
     snake_case_	:      Any            =  original_images[0, -3:, -3:, -1]

     with tempfile.TemporaryDirectory() as tmpdirname:
         LoraLoaderMixin.save_lora_weights(
             save_directory=lowerCAmelCase__	,    unet_lora_layers=lora_components["unet_lora_layers"]	,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]	,    )
         self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__	,    "pytorch_lora_weights.bin")))
         sd_pipe.load_lora_weights(lowerCAmelCase__)

     snake_case_	:      int            =  sd_pipe(**lowerCAmelCase__).images
     snake_case_	:      Dict            =  lora_images[0, -3:, -3:, -1]

     # Outputs shouldn't match.
     self.assertFalse(torch.allclose(torch.from_numpy(lowerCAmelCase__)	,    torch.from_numpy(lowerCAmelCase__)))





 def     snake_case__						(      self):
     snake_case_     ,	snake_case_	:      Optional[int]            =  self.get_dummy_components()
     snake_case_	:      int            =  StableDiffusionPipeline(**lowerCAmelCase__)
     snake_case_	:      Dict            =  sd_pipe.to(lowerCAmelCase__)
     sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

     snake_case_     ,	snake_case_     ,	snake_case_	:      Tuple            =  self.get_dummy_inputs()

     snake_case_	:      str            =  sd_pipe(**lowerCAmelCase__).images
     snake_case_	:      Dict            =  original_images[0, -3:, -3:, -1]

     with tempfile.TemporaryDirectory() as tmpdirname:
         LoraLoaderMixin.save_lora_weights(
             save_directory=lowerCAmelCase__	,    unet_lora_layers=lora_components["unet_lora_layers"]	,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]	,    safe_serialization=lowerCAmelCase__	,    )
         self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__	,    "pytorch_lora_weights.safetensors")))
         sd_pipe.load_lora_weights(lowerCAmelCase__)

     snake_case_	:      Optional[Any]            =  sd_pipe(**lowerCAmelCase__).images
     snake_case_	:      Optional[Any]            =  lora_images[0, -3:, -3:, -1]

     # Outputs shouldn't match.
     self.assertFalse(torch.allclose(torch.from_numpy(lowerCAmelCase__)	,    torch.from_numpy(lowerCAmelCase__)))





 def     snake_case__						(      self):
     snake_case_     ,	snake_case_	:      Optional[int]            =  self.get_dummy_components()
     snake_case_	:      int            =  lora_components["unet_lora_attn_procs"]
     snake_case_	:      Any            =  StableDiffusionPipeline(**lowerCAmelCase__)
     snake_case_	:      Union[str, Any]            =  sd_pipe.to(lowerCAmelCase__)
     sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

     snake_case_     ,	snake_case_     ,	snake_case_	:      List[Any]            =  self.get_dummy_inputs()

     snake_case_	:      int            =  sd_pipe(**lowerCAmelCase__).images
     snake_case_	:      List[Any]            =  original_images[0, -3:, -3:, -1]

     with tempfile.TemporaryDirectory() as tmpdirname:
         snake_case_	:      Optional[int]            =  sd_pipe.unet
         unet.set_attn_processor(lowerCAmelCase__)
         unet.save_attn_procs(lowerCAmelCase__)
         self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__	,    "pytorch_lora_weights.bin")))
         sd_pipe.load_lora_weights(lowerCAmelCase__)

     snake_case_	:      Optional[Any]            =  sd_pipe(**lowerCAmelCase__).images
     snake_case_	:      Dict            =  lora_images[0, -3:, -3:, -1]

     # Outputs shouldn't match.
     self.assertFalse(torch.allclose(torch.from_numpy(lowerCAmelCase__)	,    torch.from_numpy(lowerCAmelCase__)))





 def     snake_case__						(      self):
     snake_case_     ,	snake_case_	:      Any            =  self.get_dummy_components()
     snake_case_	:      Dict            =  StableDiffusionPipeline(**lowerCAmelCase__)

     snake_case_	:      Optional[int]            =  self.get_dummy_tokens()

     # inference without lora
     snake_case_	:      Optional[Any]            =  pipe.text_encoder(**lowerCAmelCase__)[0]
     assert outputs_without_lora.shape == (1, 77, 32)

     # create lora_attn_procs with zeroed out up.weights
     snake_case_	:      Union[str, Any]            =  create_text_encoder_lora_attn_procs(pipe.text_encoder)
     set_lora_up_weights(lowerCAmelCase__	,    randn_weight=lowerCAmelCase__)

     # monkey patch
     pipe._modify_text_encoder(lowerCAmelCase__)

     # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
     del text_attn_procs
     gc.collect()

     # inference with lora
     snake_case_	:      Any            =  pipe.text_encoder(**lowerCAmelCase__)[0]
     assert outputs_with_lora.shape == (1, 77, 32)

     assert torch.allclose(
         lowerCAmelCase__	,    lowerCAmelCase__), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

     # create lora_attn_procs with randn up.weights
     snake_case_	:      Tuple            =  create_text_encoder_lora_attn_procs(pipe.text_encoder)
     set_lora_up_weights(lowerCAmelCase__	,    randn_weight=lowerCAmelCase__)

     # monkey patch
     pipe._modify_text_encoder(lowerCAmelCase__)

     # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
     del text_attn_procs
     gc.collect()

     # inference with lora
     snake_case_	:      str            =  pipe.text_encoder(**lowerCAmelCase__)[0]
     assert outputs_with_lora.shape == (1, 77, 32)

     assert not torch.allclose(
         lowerCAmelCase__	,    lowerCAmelCase__), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"





 def     snake_case__						(      self):
     snake_case_     ,	snake_case_	:      List[str]            =  self.get_dummy_components()
     snake_case_	:      Dict            =  StableDiffusionPipeline(**lowerCAmelCase__)

     snake_case_	:      List[str]            =  self.get_dummy_tokens()

     # inference without lora
     snake_case_	:      int            =  pipe.text_encoder(**lowerCAmelCase__)[0]
     assert outputs_without_lora.shape == (1, 77, 32)

     # create lora_attn_procs with randn up.weights
     snake_case_	:      List[Any]            =  create_text_encoder_lora_attn_procs(pipe.text_encoder)
     set_lora_up_weights(lowerCAmelCase__	,    randn_weight=lowerCAmelCase__)

     # monkey patch
     pipe._modify_text_encoder(lowerCAmelCase__)

     # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
     del text_attn_procs
     gc.collect()

     # inference with lora
     snake_case_	:      List[Any]            =  pipe.text_encoder(**lowerCAmelCase__)[0]
     assert outputs_with_lora.shape == (1, 77, 32)

     assert not torch.allclose(
         lowerCAmelCase__	,    lowerCAmelCase__), "lora outputs should be different to without lora outputs"

     # remove monkey patch
     pipe._remove_text_encoder_monkey_patch()

     # inference with removed lora
     snake_case_	:      Union[str, Any]            =  pipe.text_encoder(**lowerCAmelCase__)[0]
     assert outputs_without_lora_removed.shape == (1, 77, 32)

     assert torch.allclose(
         lowerCAmelCase__	,    lowerCAmelCase__), "remove lora monkey patch should restore the original outputs"





 def     snake_case__						(      self):
     snake_case_     ,	snake_case_	:      str            =  self.get_dummy_components()
     snake_case_	:      List[str]            =  StableDiffusionPipeline(**lowerCAmelCase__)
     snake_case_	:      Tuple            =  sd_pipe.to(lowerCAmelCase__)
     sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

     snake_case_     ,	snake_case_     ,	snake_case_	:      Optional[Any]            =  self.get_dummy_inputs()

     with tempfile.TemporaryDirectory() as tmpdirname:
         LoraLoaderMixin.save_lora_weights(
             save_directory=lowerCAmelCase__	,    unet_lora_layers=lora_components["unet_lora_layers"]	,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]	,    )
         self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__	,    "pytorch_lora_weights.bin")))
         sd_pipe.load_lora_weights(lowerCAmelCase__)

     snake_case_	:      Any            =  sd_pipe(**lowerCAmelCase__).images
     snake_case_	:      Optional[int]            =  lora_images[0, -3:, -3:, -1]

     snake_case_	:      List[Any]            =  sd_pipe(**lowerCAmelCase__	,    cross_attention_kwargs={"scale": 0.5}).images
     snake_case_	:      Tuple            =  lora_images_with_scale[0, -3:, -3:, -1]

     # Outputs shouldn't match.
     self.assertFalse(
         torch.allclose(torch.from_numpy(lowerCAmelCase__)	,    torch.from_numpy(lowerCAmelCase__)))





 def     snake_case__						(      self):
     with tempfile.TemporaryDirectory() as tmpdirname:
         self.create_lora_weight_file(lowerCAmelCase__)

         snake_case_     ,	snake_case_	:      int            =  self.get_dummy_components()
         snake_case_	:      Dict            =  StableDiffusionPipeline(**lowerCAmelCase__)
         snake_case_	:      int            =  sd_pipe.to(lowerCAmelCase__)
         sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

         # check if vanilla attention processors are used
         for _, module in sd_pipe.unet.named_modules():
             if isinstance(lowerCAmelCase__	,    lowerCAmelCase__):
                 self.assertIsInstance(module.processor	,    (AttnProcessor, AttnProcessor2_0))

                    # load LoRA weight file
         sd_pipe.load_lora_weights(lowerCAmelCase__)

         # check if lora attention processors are used
         for _, module in sd_pipe.unet.named_modules():
             if isinstance(lowerCAmelCase__	,    lowerCAmelCase__):
                 snake_case_	:      Any            =  (
                     LoRAAttnProcessor2_0 if hasattr(lowerCAmelCase__	,    "scaled_dot_product_attention") else LoRAAttnProcessor
                 )
                 self.assertIsInstance(module.processor	,    lowerCAmelCase__)





 @unittest.skipIf(torch_device != "cuda"	,    "This test is supposed to run on GPU")
 def     snake_case__						(      self):
     with tempfile.TemporaryDirectory() as tmpdirname:
         self.create_lora_weight_file(lowerCAmelCase__)

         snake_case_     ,	snake_case_	:      List[str]            =  self.get_dummy_components()
         snake_case_	:      Dict            =  StableDiffusionPipeline(**lowerCAmelCase__)
         snake_case_	:      int            =  sd_pipe.to(lowerCAmelCase__)
         sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

         # enable XFormers
         sd_pipe.enable_xformers_memory_efficient_attention()

         # check if xFormers attention processors are used
         for _, module in sd_pipe.unet.named_modules():
             if isinstance(lowerCAmelCase__	,    lowerCAmelCase__):
                 self.assertIsInstance(module.processor	,    lowerCAmelCase__)

                    # load LoRA weight file
         sd_pipe.load_lora_weights(lowerCAmelCase__)

         # check if lora attention processors are used
         for _, module in sd_pipe.unet.named_modules():
             if isinstance(lowerCAmelCase__	,    lowerCAmelCase__):
                 self.assertIsInstance(module.processor	,    lowerCAmelCase__)





 @unittest.skipIf(torch_device != "cuda"	,    "This test is supposed to run on GPU")
 def     snake_case__						(      self):
     snake_case_     ,	snake_case_	:      Dict            =  self.get_dummy_components()
     snake_case_	:      List[Any]            =  StableDiffusionPipeline(**lowerCAmelCase__)
     snake_case_	:      str            =  sd_pipe.to(lowerCAmelCase__)
     sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

     snake_case_     ,	snake_case_     ,	snake_case_	:      int            =  self.get_dummy_inputs()

     # enable XFormers
     sd_pipe.enable_xformers_memory_efficient_attention()

     snake_case_	:      List[Any]            =  sd_pipe(**lowerCAmelCase__).images
     snake_case_	:      Any            =  original_images[0, -3:, -3:, -1]

     with tempfile.TemporaryDirectory() as tmpdirname:
         LoraLoaderMixin.save_lora_weights(
             save_directory=lowerCAmelCase__	,    unet_lora_layers=lora_components["unet_lora_layers"]	,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]	,    )
         self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__	,    "pytorch_lora_weights.bin")))
         sd_pipe.load_lora_weights(lowerCAmelCase__)

     snake_case_	:      Union[str, Any]            =  sd_pipe(**lowerCAmelCase__).images
     snake_case_	:      Optional[int]            =  lora_images[0, -3:, -3:, -1]

     # Outputs shouldn't match.
     self.assertFalse(torch.allclose(torch.from_numpy(lowerCAmelCase__)	,    torch.from_numpy(lowerCAmelCase__)))






'''simple docstring'''
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   ):
 UpperCAmelCase_						      =  SpectrogramDiffusionPipeline
 UpperCAmelCase_						      =  PipelineTesterMixin.required_optional_params - {
     """callback""",
     """latents""",
     """callback_steps""",
     """output_type""",
     """num_images_per_prompt""",
 }
 UpperCAmelCase_						      =  False
 UpperCAmelCase_						      =  False
 UpperCAmelCase_						      =  TOKENS_TO_AUDIO_GENERATION_PARAMS
 UpperCAmelCase_						      =  TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS





 def     snake_case__						(      self):
     torch.manual_seed(0)
     snake_case_	:      List[str]            =  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"	,    )

     snake_case_	:      int            =  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"	,    )

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

     snake_case_	:      Union[str, Any]            =  DDPMScheduler()

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





 def     snake_case__						(      self	,    lowercase_	,    lowercase_=0):
     if str(lowercase_).startswith("mps"):
         snake_case_	:      Dict            =  torch.manual_seed(lowercase_)
     else:
         snake_case_	:      Dict            =  torch.Generator(device=lowercase_).manual_seed(lowercase_)
     snake_case_	:      int            =  {
         "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     snake_case__						(      self):
     snake_case_	:      List[Any]            =  "cpu"  # ensure determinism for the device-dependent torch.Generator
     snake_case_	:      List[Any]            =  self.get_dummy_components()
     snake_case_	:      Tuple            =  SpectrogramDiffusionPipeline(**lowercase_)
     snake_case_	:      Tuple            =  pipe.to(lowercase_)
     pipe.set_progress_bar_config(disable=lowercase_)

     snake_case_	:      List[str]            =  self.get_dummy_inputs(lowercase_)
     snake_case_	:      Any            =  pipe(**lowercase_)
     snake_case_	:      str            =  output.audios

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

     assert mel_slice.shape == (3, 3)
     snake_case_	:      Any            =  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     snake_case__						(      self):
     return super().test_save_load_local()





 @skip_mps
 def     snake_case__						(      self):
     return super().test_dict_tuple_outputs_equivalent()





 @skip_mps
 def     snake_case__						(      self):
     return super().test_save_load_optional_components()





 @skip_mps
 def     snake_case__						(      self):
     return super().test_attention_slicing_forward_pass()





 def     snake_case__						(      self):
     pass





 def     snake_case__						(      self):
     pass





 @skip_mps
 def     snake_case__						(      self):
     return super().test_progress_bar()






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





 def     snake_case__						(      self):
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()





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

     snake_case_	:      str            =  SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion")
     snake_case_	:      Tuple            =  pipe.melgan
     snake_case_	:      List[str]            =  None

     snake_case_	:      int            =  pipe.to(lowercase_)
     pipe.set_progress_bar_config(disable=lowercase_)

     def callback(lowercase_	,    lowercase_):
         # decode mel to audio
         snake_case_	:      Dict            =  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

     snake_case_	:      Tuple            =  MidiProcessor()
     snake_case_	:      int            =  processor(lowercase_)

     snake_case_	:      int            =  input_tokens[:3]
     snake_case_	:      Union[str, Any]            =  torch.manual_seed(0)
     pipe(lowercase_	,    num_inference_steps=5	,    generator=lowercase_	,    callback=lowercase_	,    output_type="mel")





 def     snake_case__						(      self):
     snake_case_	:      Any            =  torch_device

     snake_case_	:      Optional[Any]            =  SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion")
     snake_case_	:      Optional[int]            =  pipe.to(lowercase_)
     pipe.set_progress_bar_config(disable=lowercase_)
     snake_case_	:      Union[str, Any]            =  MidiProcessor()

     snake_case_	:      List[str]            =  processor(lowercase_)
     # just run two denoising loops
     snake_case_	:      int            =  input_tokens[:2]

     snake_case_	:      Optional[Any]            =  torch.manual_seed(0)
     snake_case_	:      Tuple            =  pipe(lowercase_	,    num_inference_steps=2	,    generator=lowercase_)

     snake_case_	:      Optional[Any]            =  output.audios[0]

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





 def     snake_case__						(      self):
     snake_case_	:      Union[str, Any]            =  torch_device

     snake_case_	:      Tuple            =  SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion")
     snake_case_	:      Optional[Any]            =  pipe.to(lowercase_)
     pipe.set_progress_bar_config(disable=lowercase_)

     snake_case_	:      Tuple            =  MidiProcessor()

     snake_case_	:      int            =  processor(lowercase_)

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

     snake_case_	:      Optional[Any]            =  torch.manual_seed(0)
     snake_case_	:      Tuple            =  pipe(lowercase_	,    num_inference_steps=1_00	,    generator=lowercase_)

     snake_case_	:      Tuple            =  output.audios[0]
     assert abs(np.abs(lowercase_).sum() - 9_389.1_111) < 5E-2




