




"""simple docstring"""






from typing import Dict, List, Optional, Union

import numpy as np

from transformers.utils import is_vision_available
from transformers.utils.generic import TensorType

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    is_valid_image,
    to_numpy_array,
    valid_images,
)
from ...utils import logging


if is_vision_available():
			import PIL

lowercase__			=		logging.get_logger(__name__)







def 	_snake_case     (							lowercase__						):
				if isinstance(lowercase__				,       (list, tuple)						) and isinstance(videos[0]				,       (list, tuple)						) and is_valid_image(videos[0][0]						):
								return videos

				elif isinstance(lowercase__				,       (list, tuple)						) and is_valid_image(videos[0]						):
								return [videos]

				elif is_valid_image(lowercase__						):
								return [[videos]]

				raise ValueError(f'''Could not make batched video from {videos}'''						)






class lowerCAmelCase__      (	lowercase		):






			'''simple docstring'''



			lowerCamelCase__ =     ["""pixel_values"""]
			def __init__(    self    , lowercase = True    , lowercase = None    , lowercase = PILImageResampling.BILINEAR    , lowercase = True    , lowercase = None    , lowercase = True    , lowercase = 1 / 255    , lowercase = True    , lowercase = True    , lowercase = None    , lowercase = None    , **lowercase    , ):
							super().__init__(**lowercase    )
							_lowerCamelCase	:     Union[str, Any]       =       size if size is not None else {'shortest_edge': 256}
							_lowerCamelCase	:     str       =       get_size_dict(lowercase    , default_to_square=lowercase    )
							_lowerCamelCase	:     Any       =       crop_size if crop_size is not None else {'height': 224, 'width': 224}
							_lowerCamelCase	:     Dict       =       get_size_dict(lowercase    , param_name='crop_size'    )

							_lowerCamelCase	:     int       =       do_resize
							_lowerCamelCase	:     str       =       size
							_lowerCamelCase	:     Optional[Any]       =       do_center_crop
							_lowerCamelCase	:     List[Any]       =       crop_size
							_lowerCamelCase	:     str       =       resample
							_lowerCamelCase	:     str       =       do_rescale
							_lowerCamelCase	:     Optional[int]       =       rescale_factor
							_lowerCamelCase	:     Optional[int]       =       offset
							_lowerCamelCase	:     Optional[Any]       =       do_normalize
							_lowerCamelCase	:     Union[str, Any]       =       image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
							_lowerCamelCase	:     str       =       image_std if image_std is not None else IMAGENET_STANDARD_STD
			def    A_       (    self    , lowercase    , lowercase    , lowercase = PILImageResampling.BILINEAR    , lowercase = None    , **lowercase    , ):
							_lowerCamelCase	:     List[str]       =       get_size_dict(lowercase    , default_to_square=lowercase    )
							if "shortest_edge" in size:
											_lowerCamelCase	:     Union[str, Any]       =       get_resize_output_image_size(lowercase    , size['shortest_edge']    , default_to_square=lowercase    )
							elif "height" in size and "width" in size:
											_lowerCamelCase	:     Union[str, Any]       =       (size['height'], size['width'])
							else:
											raise ValueError(F'''Size must have \'height\' and \'width\' or \'shortest_edge\' as keys. Got {size.keys()}'''    )
							return resize(lowercase    , size=lowercase    , resample=lowercase    , data_format=lowercase    , **lowercase    )
			def    A_       (    self    , lowercase    , lowercase    , lowercase = None    , **lowercase    , ):
							_lowerCamelCase	:     Dict       =       get_size_dict(lowercase    )
							if "height" not in size or "width" not in size:
											raise ValueError(F'''Size must have \'height\' and \'width\' as keys. Got {size.keys()}'''    )
							return center_crop(lowercase    , size=(size['height'], size['width'])    , data_format=lowercase    , **lowercase    )
			def    A_       (    self    , lowercase    , lowercase    , lowercase = True    , lowercase = None    , **lowercase    , ):
							_lowerCamelCase	:     Tuple       =       image.astype(np.floataa    )
							if offset:
											_lowerCamelCase	:     Optional[Any]       =       image - (scale / 2)
							return rescale(lowercase    , scale=lowercase    , data_format=lowercase    , **lowercase    )
			def    A_       (    self    , lowercase    , lowercase    , lowercase    , lowercase = None    , **lowercase    , ):
							return normalize(lowercase    , mean=lowercase    , std=lowercase    , data_format=lowercase    , **lowercase    )
			def    A_       (    self    , lowercase    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = ChannelDimension.FIRST    , ):
							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_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.'    )

							if offset and not do_rescale:
											raise ValueError('For offset, do_rescale must also be set to True.'    )

							# All transformations expect numpy arrays.
							_lowerCamelCase	:     List[Any]       =       to_numpy_array(lowercase    )

							if do_resize:
											_lowerCamelCase	:     int       =       self.resize(image=lowercase    , size=lowercase    , resample=lowercase    )

							if do_center_crop:
											_lowerCamelCase	:     List[str]       =       self.center_crop(lowercase    , size=lowercase    )

							if do_rescale:
											_lowerCamelCase	:     int       =       self.rescale(image=lowercase    , scale=lowercase    , offset=lowercase    )

							if do_normalize:
											_lowerCamelCase	:     Union[str, Any]       =       self.normalize(image=lowercase    , mean=lowercase    , std=lowercase    )

							_lowerCamelCase	:     Optional[int]       =       to_channel_dimension_format(lowercase    , lowercase    )
							return image






			def    A_       (    self    , lowercase    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = ChannelDimension.FIRST    , **lowercase    , ):
							_lowerCamelCase	:     Tuple       =       do_resize if do_resize is not None else self.do_resize
							_lowerCamelCase	:     Any       =       resample if resample is not None else self.resample
							_lowerCamelCase	:     Union[str, Any]       =       do_center_crop if do_center_crop is not None else self.do_center_crop
							_lowerCamelCase	:     List[Any]       =       do_rescale if do_rescale is not None else self.do_rescale
							_lowerCamelCase	:     Optional[Any]       =       rescale_factor if rescale_factor is not None else self.rescale_factor
							_lowerCamelCase	:     List[str]       =       offset if offset is not None else self.offset
							_lowerCamelCase	:     Union[str, Any]       =       do_normalize if do_normalize is not None else self.do_normalize
							_lowerCamelCase	:     Optional[Any]       =       image_mean if image_mean is not None else self.image_mean
							_lowerCamelCase	:     str       =       image_std if image_std is not None else self.image_std

							_lowerCamelCase	:     str       =       size if size is not None else self.size
							_lowerCamelCase	:     Dict       =       get_size_dict(lowercase    , default_to_square=lowercase    )
							_lowerCamelCase	:     Dict       =       crop_size if crop_size is not None else self.crop_size
							_lowerCamelCase	:     str       =       get_size_dict(lowercase    , param_name='crop_size'    )

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

							_lowerCamelCase	:     Optional[Any]       =       make_batched(lowercase    )

							_lowerCamelCase	:     List[str]       =       [
							    [
							        self._preprocess_image(
							            image=lowercase    , do_resize=lowercase    , size=lowercase    , resample=lowercase    , do_center_crop=lowercase    , crop_size=lowercase    , do_rescale=lowercase    , rescale_factor=lowercase    , offset=lowercase    , do_normalize=lowercase    , image_mean=lowercase    , image_std=lowercase    , data_format=lowercase    , )
							        for img in video
							    ]
							    for video in videos
							]

							_lowerCamelCase	:     Union[str, Any]       =       {'pixel_values': videos}
							return BatchFeature(data=lowercase    , tensor_type=lowercase    )





"""simple docstring"""






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 lowerCAmelCase__      (	unittest.TestCase		):






			'''simple docstring'''



			def    A_       (    self    ):
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()
			@property
			def    A_       (    self    ):
							torch.manual_seed(0    )
							_lowerCamelCase	:     Union[str, Any]       =       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    A_       (    self    ):
							torch.manual_seed(0    )
							_lowerCamelCase	:     Dict       =       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    A_       (    self    ):
							torch.manual_seed(0    )
							_lowerCamelCase	:     List[str]       =       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')    , )
							_lowerCamelCase	:     int       =       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    A_       (    self    ):
							_lowerCamelCase	:     str       =       'cpu'  # ensure determinism for the device-dependent torch.Generator
							_lowerCamelCase	:     Optional[int]       =       Mel(
							    x_res=self.dummy_unet.config.sample_size[1]    , y_res=self.dummy_unet.config.sample_size[0]    , )

							_lowerCamelCase	:     List[str]       =       DDPMScheduler()
							_lowerCamelCase	:     List[Any]       =       AudioDiffusionPipeline(vqvae=lowercase    , unet=self.dummy_unet    , mel=lowercase    , scheduler=lowercase    )
							_lowerCamelCase	:     List[Any]       =       pipe.to(lowercase    )
							pipe.set_progress_bar_config(disable=lowercase    )

							_lowerCamelCase	:     Union[str, Any]       =       torch.Generator(device=lowercase    ).manual_seed(42    )
							_lowerCamelCase	:     Union[str, Any]       =       pipe(generator=lowercase    , steps=4    )
							_lowerCamelCase	:     Optional[Any]       =       output.audios[0]
							_lowerCamelCase	:     int       =       output.images[0]

							_lowerCamelCase	:     Dict       =       torch.Generator(device=lowercase    ).manual_seed(42    )
							_lowerCamelCase	:     Dict       =       pipe(generator=lowercase    , steps=4    , return_dict=lowercase    )
							_lowerCamelCase	:     List[str]       =       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]
							)
							_lowerCamelCase	:     Optional[Any]       =       np.frombuffer(image.tobytes()    , dtype='uint8'    )[:10]
							_lowerCamelCase	:     Dict       =       np.frombuffer(image_from_tuple.tobytes()    , dtype='uint8'    )[:10]
							_lowerCamelCase	:     List[str]       =       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

							_lowerCamelCase	:     List[Any]       =       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]    , )

							_lowerCamelCase	:     Dict       =       DDIMScheduler()
							_lowerCamelCase	:     Tuple       =       self.dummy_vqvae_and_unet
							_lowerCamelCase	:     List[str]       =       AudioDiffusionPipeline(
							    vqvae=self.dummy_vqvae_and_unet[0]    , unet=dummy_vqvae_and_unet[1]    , mel=lowercase    , scheduler=lowercase    )
							_lowerCamelCase	:     Tuple       =       pipe.to(lowercase    )
							pipe.set_progress_bar_config(disable=lowercase    )

							np.random.seed(0    )
							_lowerCamelCase	:     Any       =       np.random.uniform(-1    , 1    , ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)    )
							_lowerCamelCase	:     Optional[int]       =       torch.Generator(device=lowercase    ).manual_seed(42    )
							_lowerCamelCase	:     Dict       =       pipe(raw_audio=lowercase    , generator=lowercase    , start_step=5    , steps=10    )
							_lowerCamelCase	:     str       =       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]
							)
							_lowerCamelCase	:     Dict       =       np.frombuffer(image.tobytes()    , dtype='uint8'    )[:10]
							_lowerCamelCase	:     Union[str, Any]       =       np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]    )

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

							_lowerCamelCase	:     Dict       =       self.dummy_unet_condition
							_lowerCamelCase	:     Optional[int]       =       AudioDiffusionPipeline(
							    vqvae=self.dummy_vqvae_and_unet[0]    , unet=lowercase    , mel=lowercase    , scheduler=lowercase    )
							_lowerCamelCase	:     Dict       =       pipe.to(lowercase    )
							pipe.set_progress_bar_config(disable=lowercase    )

							np.random.seed(0    )
							_lowerCamelCase	:     Optional[int]       =       torch.rand((1, 1, 10)    )
							_lowerCamelCase	:     Optional[Any]       =       pipe(generator=lowercase    , encoding=lowercase    )
							_lowerCamelCase	:     Dict       =       output.images[0]
							_lowerCamelCase	:     Tuple       =       np.frombuffer(image.tobytes()    , dtype='uint8'    )[:10]
							_lowerCamelCase	:     List[str]       =       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 lowerCAmelCase__      (	unittest.TestCase		):






			'''simple docstring'''



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






			def    A_       (    self    ):
							_lowerCamelCase	:     Optional[Any]       =       torch_device

							_lowerCamelCase	:     Dict       =       DiffusionPipeline.from_pretrained('teticio/audio-diffusion-ddim-256'    )
							_lowerCamelCase	:     List[str]       =       pipe.to(lowercase    )
							pipe.set_progress_bar_config(disable=lowercase    )

							_lowerCamelCase	:     int       =       torch.Generator(device=lowercase    ).manual_seed(42    )
							_lowerCamelCase	:     Tuple       =       pipe(generator=lowercase    )
							_lowerCamelCase	:     Dict       =       output.audios[0]
							_lowerCamelCase	:     Dict       =       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]
							_lowerCamelCase	:     Optional[int]       =       np.frombuffer(image.tobytes()    , dtype='uint8'    )[:10]
							_lowerCamelCase	:     int       =       np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]    )

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