





from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


__lowercase  :   List[str]						      =logging.get_logger(__name__)


if is_vision_available():
    import PIL

class      A       (	__lowercase						):
      _snake_case							     =['''pixel_values''']



      def __init__(	self:						List[str]		,  _lowerCAmelCase:						bool = True		,  _lowerCAmelCase:						Dict[str, int] = None		,  _lowerCAmelCase:						PILImageResampling = PILImageResampling.BICUBIC		,  _lowerCAmelCase:						bool = True		,  _lowerCAmelCase:						Dict[str, int] = None		,  _lowerCAmelCase:						bool = True		,  _lowerCAmelCase:						Union[int, float] = 1 / 255		,  _lowerCAmelCase:						bool = True		,  _lowerCAmelCase:						Optional[Union[float, List[float]]] = None		,  _lowerCAmelCase:						Optional[Union[float, List[float]]] = None		,  _lowerCAmelCase:						bool = True		,  **_lowerCAmelCase:						Dict		,  )						->    None:
         '''simple docstring'''

         super().__init__(**UpperCamelCase_       )
         UpperCAmelCase_											=size if size is not None else {'shortest_edge': 224}
         UpperCAmelCase_											=get_size_dict(UpperCamelCase_		,  default_to_square=UpperCamelCase_       )
         UpperCAmelCase_											=crop_size if crop_size is not None else {'height': 224, 'width': 224}
         UpperCAmelCase_											=get_size_dict(UpperCamelCase_		,  default_to_square=UpperCamelCase_		,  param_name="crop_size"       )

         UpperCAmelCase_											=do_resize
         UpperCAmelCase_											=size
         UpperCAmelCase_											=resample
         UpperCAmelCase_											=do_center_crop
         UpperCAmelCase_											=crop_size
         UpperCAmelCase_											=do_rescale
         UpperCAmelCase_											=rescale_factor
         UpperCAmelCase_											=do_normalize
         UpperCAmelCase_											=image_mean if image_mean is not None else OPENAI_CLIP_MEAN
         UpperCAmelCase_											=image_std if image_std is not None else OPENAI_CLIP_STD
         UpperCAmelCase_											=do_convert_rgb



      def 	lowerCAmelCase__    (	self:						Optional[int]		,  _lowerCAmelCase:						np.ndarray		,  _lowerCAmelCase:						Dict[str, int]		,  _lowerCAmelCase:						PILImageResampling = PILImageResampling.BICUBIC		,  _lowerCAmelCase:						Optional[Union[str, ChannelDimension]] = None		,  **_lowerCAmelCase:						Dict		,  )						->    np.ndarray:
         '''simple docstring'''

         UpperCAmelCase_											=get_size_dict(UpperCamelCase_		,  default_to_square=UpperCamelCase_       )
         if "shortest_edge" not in size:
            raise ValueError(F'The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}'       )
         UpperCAmelCase_											=get_resize_output_image_size(UpperCamelCase_		,  size=size["shortest_edge"]		,  default_to_square=UpperCamelCase_       )
         return resize(UpperCamelCase_		,  size=UpperCamelCase_		,  resample=UpperCamelCase_		,  data_format=UpperCamelCase_		,  **UpperCamelCase_       )



      def 	lowerCAmelCase__    (	self:						Tuple		,  _lowerCAmelCase:						np.ndarray		,  _lowerCAmelCase:						Dict[str, int]		,  _lowerCAmelCase:						Optional[Union[str, ChannelDimension]] = None		,  **_lowerCAmelCase:						Tuple		,  )						->    np.ndarray:
         '''simple docstring'''

         UpperCAmelCase_											=get_size_dict(UpperCamelCase_       )
         if "height" not in size or "width" not in size:
            raise ValueError(F'The `size` parameter must contain the keys (height, width). Got {size.keys()}'       )
         return center_crop(UpperCamelCase_		,  size=(size["height"], size["width"])		,  data_format=UpperCamelCase_		,  **UpperCamelCase_       )



      def 	lowerCAmelCase__    (	self:						Tuple		,  _lowerCAmelCase:						np.ndarray		,  _lowerCAmelCase:						Union[int, float]		,  _lowerCAmelCase:						Optional[Union[str, ChannelDimension]] = None		,  **_lowerCAmelCase:						List[str]		,  )						->    int:
         '''simple docstring'''

         return rescale(UpperCamelCase_		,  scale=UpperCamelCase_		,  data_format=UpperCamelCase_		,  **UpperCamelCase_       )



      def 	lowerCAmelCase__    (	self:						Tuple		,  _lowerCAmelCase:						np.ndarray		,  _lowerCAmelCase:						Union[float, List[float]]		,  _lowerCAmelCase:						Union[float, List[float]]		,  _lowerCAmelCase:						Optional[Union[str, ChannelDimension]] = None		,  **_lowerCAmelCase:						Tuple		,  )						->    np.ndarray:
         '''simple docstring'''

         return normalize(UpperCamelCase_		,  mean=UpperCamelCase_		,  std=UpperCamelCase_		,  data_format=UpperCamelCase_		,  **UpperCamelCase_       )






      def 	lowerCAmelCase__    (	self:						List[Any]		,  _lowerCAmelCase:						ImageInput		,  _lowerCAmelCase:						bool = None		,  _lowerCAmelCase:						Dict[str, int] = None		,  _lowerCAmelCase:						PILImageResampling = None		,  _lowerCAmelCase:						bool = None		,  _lowerCAmelCase:						int = None		,  _lowerCAmelCase:						bool = None		,  _lowerCAmelCase:						float = None		,  _lowerCAmelCase:						bool = None		,  _lowerCAmelCase:						Optional[Union[float, List[float]]] = None		,  _lowerCAmelCase:						Optional[Union[float, List[float]]] = None		,  _lowerCAmelCase:						bool = None		,  _lowerCAmelCase:						Optional[Union[str, TensorType]] = None		,  _lowerCAmelCase:						Optional[ChannelDimension] = ChannelDimension.FIRST		,  **_lowerCAmelCase:						Optional[Any]		,  )						->    PIL.Image.Image:
         '''simple docstring'''

         UpperCAmelCase_											=do_resize if do_resize is not None else self.do_resize
         UpperCAmelCase_											=size if size is not None else self.size
         UpperCAmelCase_											=get_size_dict(UpperCamelCase_		,  param_name="size"		,  default_to_square=UpperCamelCase_       )
         UpperCAmelCase_											=resample if resample is not None else self.resample
         UpperCAmelCase_											=do_center_crop if do_center_crop is not None else self.do_center_crop
         UpperCAmelCase_											=crop_size if crop_size is not None else self.crop_size
         UpperCAmelCase_											=get_size_dict(UpperCamelCase_		,  param_name="crop_size"		,  default_to_square=UpperCamelCase_       )
         UpperCAmelCase_											=do_rescale if do_rescale is not None else self.do_rescale
         UpperCAmelCase_											=rescale_factor if rescale_factor is not None else self.rescale_factor
         UpperCAmelCase_											=do_normalize if do_normalize is not None else self.do_normalize
         UpperCAmelCase_											=image_mean if image_mean is not None else self.image_mean
         UpperCAmelCase_											=image_std if image_std is not None else self.image_std
         UpperCAmelCase_											=do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

         UpperCAmelCase_											=make_list_of_images(UpperCamelCase_       )

         if not valid_images(UpperCamelCase_       ):
            raise ValueError(
                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
                "torch.Tensor, tf.Tensor or jax.ndarray."       )

         if do_resize and size is None:
            raise ValueError("Size must be specified if do_resize is True."       )

         if do_center_crop and crop_size is None:
            raise ValueError("Crop size must be specified if do_center_crop is True."       )

         if do_rescale and rescale_factor is None:
            raise ValueError("Rescale factor must be specified if do_rescale is True."       )

         if do_normalize and (image_mean is None or image_std is None):
            raise ValueError("Image mean and std must be specified if do_normalize is True."       )

         # PIL RGBA images are converted to RGB
         if do_convert_rgb:
            UpperCAmelCase_											=[convert_to_rgb(UpperCamelCase_       ) for image in images]

         # All transformations expect numpy arrays.
         UpperCAmelCase_											=[to_numpy_array(UpperCamelCase_       ) for image in images]

         if do_resize:
            UpperCAmelCase_											=[self.resize(image=UpperCamelCase_		,  size=UpperCamelCase_		,  resample=UpperCamelCase_       ) for image in images]

         if do_center_crop:
            UpperCAmelCase_											=[self.center_crop(image=UpperCamelCase_		,  size=UpperCamelCase_       ) for image in images]

         if do_rescale:
            UpperCAmelCase_											=[self.rescale(image=UpperCamelCase_		,  scale=UpperCamelCase_       ) for image in images]

         if do_normalize:
            UpperCAmelCase_											=[self.normalize(image=UpperCamelCase_		,  mean=UpperCamelCase_		,  std=UpperCamelCase_       ) for image in images]

         UpperCAmelCase_											=[to_channel_dimension_format(UpperCamelCase_		,  UpperCamelCase_       ) for image in images]

         UpperCAmelCase_											={'pixel_values': images}
         return BatchFeature(data=UpperCamelCase_		,  tensor_type=UpperCamelCase_       )











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       (	__lowercase		,      unittest.TestCase						):
      _snake_case							     =DanceDiffusionPipeline
      _snake_case							     =UNCONDITIONAL_AUDIO_GENERATION_PARAMS
      _snake_case							     =PipelineTesterMixin.required_optional_params - {
          '''callback''',
          '''latents''',
          '''callback_steps''',
          '''output_type''',
          '''num_images_per_prompt''',
      }
      _snake_case							     =UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
      _snake_case							     =False
      _snake_case							     =False



      def 	lowerCAmelCase__    (	self:						Tuple       )						->    Tuple:
         '''simple docstring'''

         torch.manual_seed(0       )
         UpperCAmelCase_											=UNetaDModel(
             block_out_channels=(32, 32, 64)		,  extra_in_channels=16		,  sample_size=512		,  sample_rate=1_6000		,  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")		,  )
         UpperCAmelCase_											=IPNDMScheduler()

         UpperCAmelCase_											={
             "unet": unet,
             "scheduler": scheduler,
         }
         return components



      def 	lowerCAmelCase__    (	self:						str		,  _lowerCAmelCase:						Tuple		,  _lowerCAmelCase:						List[str]=0       )						->    str:
         '''simple docstring'''

         if str(_lowerCAmelCase       ).startswith("mps"       ):
            UpperCAmelCase_											=torch.manual_seed(_lowerCAmelCase       )
         else:
            UpperCAmelCase_											=torch.Generator(device=_lowerCAmelCase       ).manual_seed(_lowerCAmelCase       )
         UpperCAmelCase_											={
             "batch_size": 1,
             "generator": generator,
             "num_inference_steps": 4,
         }
         return inputs



      def 	lowerCAmelCase__    (	self:						Dict       )						->    Tuple:
         '''simple docstring'''

         UpperCAmelCase_											="cpu"  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase_											=self.get_dummy_components()
         UpperCAmelCase_											=DanceDiffusionPipeline(**_lowerCAmelCase       )
         UpperCAmelCase_											=pipe.to(_lowerCAmelCase       )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         UpperCAmelCase_											=self.get_dummy_inputs(_lowerCAmelCase       )
         UpperCAmelCase_											=pipe(**_lowerCAmelCase       )
         UpperCAmelCase_											=output.audios

         UpperCAmelCase_											=audio[0, -3:, -3:]

         assert audio.shape == (1, 2, components["unet"].sample_size)
         UpperCAmelCase_											=np.array([-0.72_65, 1.00_00, -0.83_88, 0.11_75, 0.94_98, -1.00_00]       )
         assert np.abs(audio_slice.flatten() - expected_slice       ).max() < 1e-2



      @skip_mps
      def 	lowerCAmelCase__    (	self:						Optional[int]       )						->    Any:
         '''simple docstring'''

         return super().test_save_load_local()



      @skip_mps
      def 	lowerCAmelCase__    (	self:						Union[str, Any]       )						->    int:
         '''simple docstring'''

         return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3       )



      @skip_mps
      def 	lowerCAmelCase__    (	self:						List[str]       )						->    Optional[int]:
         '''simple docstring'''

         return super().test_save_load_optional_components()



      @skip_mps
      def 	lowerCAmelCase__    (	self:						Optional[Any]       )						->    Union[str, Any]:
         '''simple docstring'''

         return super().test_attention_slicing_forward_pass()






      def 	lowerCAmelCase__    (	self:						Optional[int]       )						->    Optional[Any]:
         '''simple docstring'''

         super().test_inference_batch_single_identical(expected_max_diff=3e-3       )

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



      def 	lowerCAmelCase__    (	self:						Any       )						->    Optional[Any]:
         '''simple docstring'''

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



      def 	lowerCAmelCase__    (	self:						Union[str, Any]       )						->    Tuple:
         '''simple docstring'''

         UpperCAmelCase_											=torch_device

         UpperCAmelCase_											=DanceDiffusionPipeline.from_pretrained("harmonai/maestro-150k"       )
         UpperCAmelCase_											=pipe.to(_lowerCAmelCase       )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         UpperCAmelCase_											=torch.manual_seed(0       )
         UpperCAmelCase_											=pipe(generator=_lowerCAmelCase		,  num_inference_steps=100		,  audio_length_in_s=4.0_96       )
         UpperCAmelCase_											=output.audios

         UpperCAmelCase_											=audio[0, -3:, -3:]

         assert audio.shape == (1, 2, pipe.unet.sample_size)
         UpperCAmelCase_											=np.array([-0.01_92, -0.02_31, -0.03_18, -0.00_59, 0.00_02, -0.00_20]       )

         assert np.abs(audio_slice.flatten() - expected_slice       ).max() < 1e-2






      def 	lowerCAmelCase__    (	self:						Tuple       )						->    Any:
         '''simple docstring'''

         UpperCAmelCase_											=torch_device

         UpperCAmelCase_											=DanceDiffusionPipeline.from_pretrained("harmonai/maestro-150k"		,  torch_dtype=torch.floataa       )
         UpperCAmelCase_											=pipe.to(_lowerCAmelCase       )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         UpperCAmelCase_											=torch.manual_seed(0       )
         UpperCAmelCase_											=pipe(generator=_lowerCAmelCase		,  num_inference_steps=100		,  audio_length_in_s=4.0_96       )
         UpperCAmelCase_											=output.audios

         UpperCAmelCase_											=audio[0, -3:, -3:]

         assert audio.shape == (1, 2, pipe.unet.sample_size)
         UpperCAmelCase_											=np.array([-0.03_67, -0.04_88, -0.07_71, -0.05_25, -0.04_44, -0.03_41]       )

         assert np.abs(audio_slice.flatten() - expected_slice       ).max() < 1e-2




