






'''simple docstring'''


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,
    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,
    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[Any]    	=							logging.get_logger(__name__)
class    _SCREAMING_SNAKE_CASE							(		SCREAMING_SNAKE_CASE	):


			'''simple docstring'''





			__a							:							Any					     =					["pixel_values"]






			def __init__(   self							:						Dict ,      lowercase							:						bool = True ,      lowercase							:						Dict[str, int] = None ,      lowercase							:						float = None ,      lowercase							:						PILImageResampling = PILImageResampling.BILINEAR ,      lowercase							:						bool = True ,      lowercase							:						Union[int, float] = 1 / 2_5_5 ,      lowercase							:						bool = True ,      lowercase							:						Optional[Union[float, List[float]]] = None ,      lowercase							:						Optional[Union[float, List[float]]] = None ,      **lowercase							:						Union[str, Any] ,      )							->    None:

									'''simple docstring'''







									super().__init__(**lowercase			)
									UpperCamelCase__         = size if size is not None else {"""shortest_edge""": 3_8_4}
									UpperCamelCase__         = get_size_dict(lowercase ,      default_to_square=lowercase			)

									UpperCamelCase__         = do_resize
									UpperCamelCase__         = size
									# Default value set here for backwards compatibility where the value in config is None
									UpperCamelCase__         = crop_pct if crop_pct is not None else 2_2_4 / 2_5_6
									UpperCamelCase__         = resample
									UpperCamelCase__         = do_rescale
									UpperCamelCase__         = rescale_factor
									UpperCamelCase__         = do_normalize
									UpperCamelCase__         = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
									UpperCamelCase__         = image_std if image_std is not None else IMAGENET_STANDARD_STD






			def 				A      (   self							:						Dict ,      lowercase							:						np.ndarray ,      lowercase							:						Dict[str, int] ,      lowercase							:						float ,      lowercase							:						PILImageResampling = PILImageResampling.BICUBIC ,      lowercase							:						Optional[Union[str, ChannelDimension]] = None ,      **lowercase							:						Any ,      )							->    np.ndarray:

									'''simple docstring'''







									UpperCamelCase__         = get_size_dict(lowercase ,      default_to_square=lowercase			)
									if "shortest_edge" not in size:
															raise ValueError(f"Size dictionary must contain 'shortest_edge' key. Got {size.keys()}"			)
									UpperCamelCase__         = size["""shortest_edge"""]

									if shortest_edge < 3_8_4:
															# maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
															UpperCamelCase__         = int(shortest_edge / crop_pct			)
															UpperCamelCase__         = get_resize_output_image_size(lowercase ,      size=lowercase ,      default_to_square=lowercase			)
															UpperCamelCase__         = resize(image=lowercase ,      size=lowercase ,      resample=lowercase ,      data_format=lowercase ,      **lowercase			)
															# then crop to (shortest_edge, shortest_edge)
															return center_crop(image=lowercase ,      size=(shortest_edge, shortest_edge) ,      data_format=lowercase ,      **lowercase			)
									else:
															# warping (no cropping) when evaluated at 384 or larger
															return resize(
															    lowercase ,      size=(shortest_edge, shortest_edge) ,      resample=lowercase ,      data_format=lowercase ,      **lowercase			)






			def 				A      (   self							:						Tuple ,      lowercase							:						np.ndarray ,      lowercase							:						Union[int, float] ,      lowercase							:						Optional[Union[str, ChannelDimension]] = None ,      **lowercase							:						Optional[int] ,      )							->    List[Any]:

									'''simple docstring'''







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






			def 				A      (   self							:						Optional[int] ,      lowercase							:						np.ndarray ,      lowercase							:						Union[float, List[float]] ,      lowercase							:						Union[float, List[float]] ,      lowercase							:						Optional[Union[str, ChannelDimension]] = None ,      **lowercase							:						str ,      )							->    np.ndarray:

									'''simple docstring'''







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







			def 				A      (   self							:						Dict ,      lowercase							:						ImageInput ,      lowercase							:						bool = None ,      lowercase							:						Dict[str, int] = None ,      lowercase							:						float = None ,      lowercase							:						PILImageResampling = None ,      lowercase							:						bool = None ,      lowercase							:						float = None ,      lowercase							:						bool = None ,      lowercase							:						Optional[Union[float, List[float]]] = None ,      lowercase							:						Optional[Union[float, List[float]]] = None ,      lowercase							:						Optional[Union[str, TensorType]] = None ,      lowercase							:						ChannelDimension = ChannelDimension.FIRST ,      **lowercase							:						int ,      )							->    PIL.Image.Image:

									'''simple docstring'''







									UpperCamelCase__         = do_resize if do_resize is not None else self.do_resize
									UpperCamelCase__         = crop_pct if crop_pct is not None else self.crop_pct
									UpperCamelCase__         = resample if resample is not None else self.resample
									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__         = size if size is not None else self.size
									UpperCamelCase__         = get_size_dict(lowercase ,      default_to_square=lowercase			)

									UpperCamelCase__         = make_list_of_images(lowercase			)

									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."""			)

									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_resize and size["shortest_edge"] < 3_8_4 and crop_pct is None:
															raise ValueError("""crop_pct must be specified if size < 384."""			)

									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."""			)

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

									if do_resize:
															UpperCamelCase__         = [self.resize(image=lowercase ,      size=lowercase ,      crop_pct=lowercase ,      resample=lowercase			) for image in images]

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

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

									UpperCamelCase__         = [to_channel_dimension_format(lowercase ,      lowercase			) for image in images]

									UpperCamelCase__         = {"""pixel_values""": images}
									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    _SCREAMING_SNAKE_CASE							(		unittest.TestCase	):


			'''simple docstring'''





			def 				A      (   self							:						Any			)							->    Any:

									'''simple docstring'''







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






			@property
			def 				A      (   self							:						Optional[Any]			)							->    Union[str, Any]:

									'''simple docstring'''







									torch.manual_seed(0			)
									UpperCamelCase__         = UNetaDModel(
									    sample_size=(3_2, 6_4) ,      in_channels=1 ,      out_channels=1 ,      layers_per_block=2 ,      block_out_channels=(1_2_8, 1_2_8) ,      down_block_types=("""AttnDownBlock2D""", """DownBlock2D""") ,      up_block_types=("""UpBlock2D""", """AttnUpBlock2D""") ,      )
									return model






			@property
			def 				A      (   self							:						List[Any]			)							->    Union[str, Any]:

									'''simple docstring'''







									torch.manual_seed(0			)
									UpperCamelCase__         = UNetaDConditionModel(
									    sample_size=(6_4, 3_2) ,      in_channels=1 ,      out_channels=1 ,      layers_per_block=2 ,      block_out_channels=(1_2_8, 1_2_8) ,      down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""") ,      up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""") ,      cross_attention_dim=1_0 ,      )
									return model






			@property
			def 				A      (   self							:						int			)							->    Optional[Any]:

									'''simple docstring'''







									torch.manual_seed(0			)
									UpperCamelCase__         = AutoencoderKL(
									    sample_size=(1_2_8, 6_4) ,      in_channels=1 ,      out_channels=1 ,      latent_channels=1 ,      layers_per_block=2 ,      block_out_channels=(1_2_8, 1_2_8) ,      down_block_types=("""DownEncoderBlock2D""", """DownEncoderBlock2D""") ,      up_block_types=("""UpDecoderBlock2D""", """UpDecoderBlock2D""") ,      )
									UpperCamelCase__         = UNetaDModel(
									    sample_size=(6_4, 3_2) ,      in_channels=1 ,      out_channels=1 ,      layers_per_block=2 ,      block_out_channels=(1_2_8, 1_2_8) ,      down_block_types=("""AttnDownBlock2D""", """DownBlock2D""") ,      up_block_types=("""UpBlock2D""", """AttnUpBlock2D""") ,      )
									return vqvae, unet







			@slow
			def 				A      (   self							:						Tuple			)							->    List[str]:

									'''simple docstring'''







									UpperCamelCase__         = """cpu"""  # ensure determinism for the device-dependent torch.Generator
									UpperCamelCase__         = Mel(
									    x_res=self.dummy_unet.config.sample_size[1] ,      y_res=self.dummy_unet.config.sample_size[0] ,      )

									UpperCamelCase__         = DDPMScheduler()
									UpperCamelCase__         = AudioDiffusionPipeline(vqvae=lowercase ,      unet=self.dummy_unet ,      mel=lowercase ,      scheduler=lowercase			)
									UpperCamelCase__         = pipe.to(lowercase			)
									pipe.set_progress_bar_config(disable=lowercase			)

									UpperCamelCase__         = torch.Generator(device=lowercase			).manual_seed(4_2			)
									UpperCamelCase__         = pipe(generator=lowercase ,      steps=4			)
									UpperCamelCase__         = output.audios[0]
									UpperCamelCase__         = output.images[0]

									UpperCamelCase__         = torch.Generator(device=lowercase			).manual_seed(4_2			)
									UpperCamelCase__         = pipe(generator=lowercase ,      steps=4 ,      return_dict=lowercase			)
									UpperCamelCase__         = 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]
									)
									UpperCamelCase__         = np.frombuffer(image.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.frombuffer(image_from_tuple.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.array([6_9, 2_5_5, 2_5_5, 2_5_5, 0, 0, 7_7, 1_8_1, 1_2, 1_2_7]			)

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

									UpperCamelCase__         = 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] ,      )

									UpperCamelCase__         = DDIMScheduler()
									UpperCamelCase__         = self.dummy_vqvae_and_unet
									UpperCamelCase__         = AudioDiffusionPipeline(
									    vqvae=self.dummy_vqvae_and_unet[0] ,      unet=dummy_vqvae_and_unet[1] ,      mel=lowercase ,      scheduler=lowercase			)
									UpperCamelCase__         = pipe.to(lowercase			)
									pipe.set_progress_bar_config(disable=lowercase			)

									np.random.seed(0			)
									UpperCamelCase__         = np.random.uniform(-1 ,      1 ,      ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)			)
									UpperCamelCase__         = torch.Generator(device=lowercase			).manual_seed(4_2			)
									UpperCamelCase__         = pipe(raw_audio=lowercase ,      generator=lowercase ,      start_step=5 ,      steps=1_0			)
									UpperCamelCase__         = 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]
									)
									UpperCamelCase__         = np.frombuffer(image.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.array([1_2_0, 1_1_7, 1_1_0, 1_0_9, 1_3_8, 1_6_7, 1_3_8, 1_4_8, 1_3_2, 1_2_1]			)

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

									UpperCamelCase__         = self.dummy_unet_condition
									UpperCamelCase__         = AudioDiffusionPipeline(
									    vqvae=self.dummy_vqvae_and_unet[0] ,      unet=lowercase ,      mel=lowercase ,      scheduler=lowercase			)
									UpperCamelCase__         = pipe.to(lowercase			)
									pipe.set_progress_bar_config(disable=lowercase			)

									np.random.seed(0			)
									UpperCamelCase__         = torch.rand((1, 1, 1_0)			)
									UpperCamelCase__         = pipe(generator=lowercase ,      encoding=lowercase			)
									UpperCamelCase__         = output.images[0]
									UpperCamelCase__         = np.frombuffer(image.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.array([1_0_7, 1_0_3, 1_2_0, 1_2_7, 1_4_2, 1_2_2, 1_1_3, 1_2_2, 9_7, 1_1_1]			)

									assert np.abs(image_slice.flatten() - expected_slice			).max() == 0
@slow
@require_torch_gpu
class    _SCREAMING_SNAKE_CASE							(		unittest.TestCase	):


			'''simple docstring'''





			def 				A      (   self							:						int			)							->    str:

									'''simple docstring'''







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







			def 				A      (   self							:						Any			)							->    Optional[Any]:

									'''simple docstring'''







									UpperCamelCase__         = torch_device

									UpperCamelCase__         = DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""			)
									UpperCamelCase__         = pipe.to(lowercase			)
									pipe.set_progress_bar_config(disable=lowercase			)

									UpperCamelCase__         = torch.Generator(device=lowercase			).manual_seed(4_2			)
									UpperCamelCase__         = pipe(generator=lowercase			)
									UpperCamelCase__         = output.audios[0]
									UpperCamelCase__         = 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]
									UpperCamelCase__         = np.frombuffer(image.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.array([1_5_1, 1_6_7, 1_5_4, 1_4_4, 1_2_2, 1_3_4, 1_2_1, 1_0_5, 7_0, 2_6]			)

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