

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 convert_to_rgb, 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


if is_vision_available():
    import PIL


lowerCAmelCase     :   Optional[int]           =  logging.get_logger(__name__)
class  _A		(  __magic_name__):
  SCREAMING_SNAKE_CASE    :      Tuple       =  ['''pixel_values''']

  def __init__(				self ,					_SCREAMING_SNAKE_CASE = True ,					_SCREAMING_SNAKE_CASE = None ,					_SCREAMING_SNAKE_CASE = PILImageResampling.BICUBIC ,					_SCREAMING_SNAKE_CASE = True ,					_SCREAMING_SNAKE_CASE = 1 / 255 ,					_SCREAMING_SNAKE_CASE = True ,					_SCREAMING_SNAKE_CASE = None ,					_SCREAMING_SNAKE_CASE = None ,					_SCREAMING_SNAKE_CASE = True ,					**_SCREAMING_SNAKE_CASE ,					):





        """simple docstring"""







        super().__init__(**_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     size if size is not None else {'height': 384, 'width': 384}
        SCREAMING_SNAKE_CASE_				:  Any           =     get_size_dict(_SCREAMING_SNAKE_CASE ,					default_to_square=_SCREAMING_SNAKE_CASE					)

        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     do_resize
        SCREAMING_SNAKE_CASE_				:  str           =     size
        SCREAMING_SNAKE_CASE_				:  Any           =     resample
        SCREAMING_SNAKE_CASE_				:  int           =     do_rescale
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     rescale_factor
        SCREAMING_SNAKE_CASE_				:  Tuple           =     do_normalize
        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     image_mean if image_mean is not None else OPENAI_CLIP_MEAN
        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     image_std if image_std is not None else OPENAI_CLIP_STD
        SCREAMING_SNAKE_CASE_				:  Any           =     do_convert_rgb

  def  UpperCAmelCase    (				self ,					_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE = PILImageResampling.BICUBIC ,					_SCREAMING_SNAKE_CASE = None ,					**_SCREAMING_SNAKE_CASE ,					):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_				:  List[Any]           =     get_size_dict(_SCREAMING_SNAKE_CASE ,					default_to_square=_SCREAMING_SNAKE_CASE					)
        if "height" not in size or "width" not in size:
              raise ValueError(f"The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}"					)
        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     (size['height'], size['width'])
        return resize(_SCREAMING_SNAKE_CASE ,					size=_SCREAMING_SNAKE_CASE ,					resample=_SCREAMING_SNAKE_CASE ,					data_format=_SCREAMING_SNAKE_CASE ,					**_SCREAMING_SNAKE_CASE					)

  def  UpperCAmelCase    (				self ,					_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE = None ,					**_SCREAMING_SNAKE_CASE ,					):





        """simple docstring"""







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

  def  UpperCAmelCase    (				self ,					_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE = None ,					**_SCREAMING_SNAKE_CASE ,					):





        """simple docstring"""







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



  def  UpperCAmelCase    (				self ,					_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE = None ,					_SCREAMING_SNAKE_CASE = None ,					_SCREAMING_SNAKE_CASE = None ,					_SCREAMING_SNAKE_CASE = None ,					_SCREAMING_SNAKE_CASE = None ,					_SCREAMING_SNAKE_CASE = None ,					_SCREAMING_SNAKE_CASE = None ,					_SCREAMING_SNAKE_CASE = None ,					_SCREAMING_SNAKE_CASE = None ,					_SCREAMING_SNAKE_CASE = None ,					_SCREAMING_SNAKE_CASE = ChannelDimension.FIRST ,					**_SCREAMING_SNAKE_CASE ,					):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_				:  int           =     do_resize if do_resize is not None else self.do_resize
        SCREAMING_SNAKE_CASE_				:  List[str]           =     resample if resample is not None else self.resample
        SCREAMING_SNAKE_CASE_				:  Dict           =     do_rescale if do_rescale is not None else self.do_rescale
        SCREAMING_SNAKE_CASE_				:  Dict           =     rescale_factor if rescale_factor is not None else self.rescale_factor
        SCREAMING_SNAKE_CASE_				:  List[str]           =     do_normalize if do_normalize is not None else self.do_normalize
        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     image_mean if image_mean is not None else self.image_mean
        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     image_std if image_std is not None else self.image_std
        SCREAMING_SNAKE_CASE_				:  Tuple           =     do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     size if size is not None else self.size
        SCREAMING_SNAKE_CASE_				:  List[str]           =     get_size_dict(_SCREAMING_SNAKE_CASE ,					default_to_square=_SCREAMING_SNAKE_CASE					)

        SCREAMING_SNAKE_CASE_				:  Tuple           =     make_list_of_images(_SCREAMING_SNAKE_CASE					)

        if not valid_images(_SCREAMING_SNAKE_CASE					):
              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 or resample is None:
              raise ValueError('Size and resample must be specified if do_resize 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:
              SCREAMING_SNAKE_CASE_				:  Optional[int]           =     [convert_to_rgb(_SCREAMING_SNAKE_CASE					) for image in images]

        # All transformations expect numpy arrays.
        SCREAMING_SNAKE_CASE_				:  Dict           =     [to_numpy_array(_SCREAMING_SNAKE_CASE					) for image in images]

        if do_resize:
              SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     [self.resize(image=_SCREAMING_SNAKE_CASE ,					size=_SCREAMING_SNAKE_CASE ,					resample=_SCREAMING_SNAKE_CASE					) for image in images]

        if do_rescale:
              SCREAMING_SNAKE_CASE_				:  Dict           =     [self.rescale(image=_SCREAMING_SNAKE_CASE ,					scale=_SCREAMING_SNAKE_CASE					) for image in images]

        if do_normalize:
              SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     [self.normalize(image=_SCREAMING_SNAKE_CASE ,					mean=_SCREAMING_SNAKE_CASE ,					std=_SCREAMING_SNAKE_CASE					) for image in images]

        SCREAMING_SNAKE_CASE_				:  int           =     [to_channel_dimension_format(_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE					) for image in images]

        SCREAMING_SNAKE_CASE_				:  Any           =     BatchFeature(data={'pixel_values': images} ,					tensor_type=_SCREAMING_SNAKE_CASE					)

        return encoded_outputs









import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()
class  _A		(  unittest.TestCase):

  def  UpperCAmelCase    (				self					):





        """simple docstring"""







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

  @property
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        torch.manual_seed(0					)
        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     UNetaDModel(
            sample_size=(32, 64) ,					in_channels=1 ,					out_channels=1 ,					layers_per_block=2 ,					block_out_channels=(128, 128) ,					down_block_types=('AttnDownBlock2D', 'DownBlock2D') ,					up_block_types=('UpBlock2D', 'AttnUpBlock2D') ,					)
        return model

  @property
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        torch.manual_seed(0					)
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     UNetaDConditionModel(
            sample_size=(64, 32) ,					in_channels=1 ,					out_channels=1 ,					layers_per_block=2 ,					block_out_channels=(128, 128) ,					down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D') ,					up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D') ,					cross_attention_dim=10 ,					)
        return model

  @property
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        torch.manual_seed(0					)
        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     AutoencoderKL(
            sample_size=(128, 64) ,					in_channels=1 ,					out_channels=1 ,					latent_channels=1 ,					layers_per_block=2 ,					block_out_channels=(128, 128) ,					down_block_types=('DownEncoderBlock2D', 'DownEncoderBlock2D') ,					up_block_types=('UpDecoderBlock2D', 'UpDecoderBlock2D') ,					)
        SCREAMING_SNAKE_CASE_				:  Dict           =     UNetaDModel(
            sample_size=(64, 32) ,					in_channels=1 ,					out_channels=1 ,					layers_per_block=2 ,					block_out_channels=(128, 128) ,					down_block_types=('AttnDownBlock2D', 'DownBlock2D') ,					up_block_types=('UpBlock2D', 'AttnUpBlock2D') ,					)
        return vqvae, unet



  @slow
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     'cpu'  # ensure determinism for the device-dependent torch.Generator
        SCREAMING_SNAKE_CASE_				:  Any           =     Mel(
            x_res=self.dummy_unet.config.sample_size[1] ,					y_res=self.dummy_unet.config.sample_size[0] ,					)

        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     DDPMScheduler()
        SCREAMING_SNAKE_CASE_				:  str           =     AudioDiffusionPipeline(vqvae=_SCREAMING_SNAKE_CASE ,					unet=self.dummy_unet ,					mel=_SCREAMING_SNAKE_CASE ,					scheduler=_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  Dict           =     pipe.to(_SCREAMING_SNAKE_CASE					)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

        SCREAMING_SNAKE_CASE_				:  Dict           =     torch.Generator(device=_SCREAMING_SNAKE_CASE					).manual_seed(42					)
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     pipe(generator=_SCREAMING_SNAKE_CASE ,					steps=4					)
        SCREAMING_SNAKE_CASE_				:  Any           =     output.audios[0]
        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     output.images[0]

        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     torch.Generator(device=_SCREAMING_SNAKE_CASE					).manual_seed(42					)
        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     pipe(generator=_SCREAMING_SNAKE_CASE ,					steps=4 ,					return_dict=_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     output[0][0]

        assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
        assert (
            image.height == self.dummy_unet.config.sample_size[0]
            and image.width == self.dummy_unet.config.sample_size[1]
        )
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     np.frombuffer(image.tobytes() ,					dtype='uint8'					)[:10]
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     np.frombuffer(image_from_tuple.tobytes() ,					dtype='uint8'					)[:10]
        SCREAMING_SNAKE_CASE_				:  Any           =     np.array([69, 255, 255, 255, 0, 0, 77, 181, 12, 127]					)

        assert np.abs(image_slice.flatten() - expected_slice					).max() == 0
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice					).max() == 0

        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     Mel(
            x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1] ,					y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0] ,					)

        SCREAMING_SNAKE_CASE_				:  str           =     DDIMScheduler()
        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     self.dummy_vqvae_and_unet
        SCREAMING_SNAKE_CASE_				:  int           =     AudioDiffusionPipeline(
            vqvae=self.dummy_vqvae_and_unet[0] ,					unet=dummy_vqvae_and_unet[1] ,					mel=_SCREAMING_SNAKE_CASE ,					scheduler=_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  str           =     pipe.to(_SCREAMING_SNAKE_CASE					)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

        np.random.seed(0					)
        SCREAMING_SNAKE_CASE_				:  Dict           =     np.random.uniform(-1 ,					1 ,					((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)					)
        SCREAMING_SNAKE_CASE_				:  str           =     torch.Generator(device=_SCREAMING_SNAKE_CASE					).manual_seed(42					)
        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     pipe(raw_audio=_SCREAMING_SNAKE_CASE ,					generator=_SCREAMING_SNAKE_CASE ,					start_step=5 ,					steps=10					)
        SCREAMING_SNAKE_CASE_				:  int           =     output.images[0]

        assert (
            image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
            and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
        )
        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     np.frombuffer(image.tobytes() ,					dtype='uint8'					)[:10]
        SCREAMING_SNAKE_CASE_				:  int           =     np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]					)

        assert np.abs(image_slice.flatten() - expected_slice					).max() == 0

        SCREAMING_SNAKE_CASE_				:  Dict           =     self.dummy_unet_condition
        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     AudioDiffusionPipeline(
            vqvae=self.dummy_vqvae_and_unet[0] ,					unet=_SCREAMING_SNAKE_CASE ,					mel=_SCREAMING_SNAKE_CASE ,					scheduler=_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  Dict           =     pipe.to(_SCREAMING_SNAKE_CASE					)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

        np.random.seed(0					)
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     torch.rand((1, 1, 10)					)
        SCREAMING_SNAKE_CASE_				:  Any           =     pipe(generator=_SCREAMING_SNAKE_CASE ,					encoding=_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     output.images[0]
        SCREAMING_SNAKE_CASE_				:  Dict           =     np.frombuffer(image.tobytes() ,					dtype='uint8'					)[:10]
        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111]					)

        assert np.abs(image_slice.flatten() - expected_slice					).max() == 0

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

  def  UpperCAmelCase    (				self					):





        """simple docstring"""







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



  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     torch_device

        SCREAMING_SNAKE_CASE_				:  str           =     DiffusionPipeline.from_pretrained('teticio/audio-diffusion-ddim-256'					)
        SCREAMING_SNAKE_CASE_				:  Dict           =     pipe.to(_SCREAMING_SNAKE_CASE					)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     torch.Generator(device=_SCREAMING_SNAKE_CASE					).manual_seed(42					)
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     pipe(generator=_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  Any           =     output.audios[0]
        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     output.images[0]

        assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
        assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     np.frombuffer(image.tobytes() ,					dtype='uint8'					)[:10]
        SCREAMING_SNAKE_CASE_				:  Any           =     np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]					)

        assert np.abs(image_slice.flatten() - expected_slice					).max() == 0






