






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 					__lowercase					(		self					:   Any						):


      '''simple docstring'''






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







 @property
 def 					__lowercase					(		self					:   Optional[Any]						):


      '''simple docstring'''






      torch.manual_seed(0						)
      UpperCamelCase__			:   Dict       				=  UNetaDModel(
          sample_size=(32, 64)							,   in_channels=1							,   out_channels=1							,   layers_per_block=2							,   block_out_channels=(1_28, 1_28)							,   down_block_types=("AttnDownBlock2D", "DownBlock2D")							,   up_block_types=("UpBlock2D", "AttnUpBlock2D")							,   )
      return model







 @property
 def 					__lowercase					(		self					:   str						):


      '''simple docstring'''






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







 @property
 def 					__lowercase					(		self					:   Union[str, Any]						):


      '''simple docstring'''






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







 @slow
 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






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

      UpperCamelCase__			:   Optional[Any]       				=  DDPMScheduler()
      UpperCamelCase__			:   Optional[Any]       				=  AudioDiffusionPipeline(vqvae=SCREAMING_SNAKE_CASE							,   unet=self.dummy_unet							,   mel=SCREAMING_SNAKE_CASE							,   scheduler=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Optional[int]       				=  pipe.to(SCREAMING_SNAKE_CASE						)
      pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

      UpperCamelCase__			:   Tuple       				=  torch.Generator(device=SCREAMING_SNAKE_CASE						).manual_seed(42						)
      UpperCamelCase__			:   Tuple       				=  pipe(generator=SCREAMING_SNAKE_CASE							,   steps=4						)
      UpperCamelCase__			:   List[Any]       				=  output.audios[0]
      UpperCamelCase__			:   str       				=  output.images[0]

      UpperCamelCase__			:   Any       				=  torch.Generator(device=SCREAMING_SNAKE_CASE						).manual_seed(42						)
      UpperCamelCase__			:   Dict       				=  pipe(generator=SCREAMING_SNAKE_CASE							,   steps=4							,   return_dict=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   int       				=  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__			:   Optional[Any]       				=  np.frombuffer(image.tobytes()							,   dtype="uint8"						)[:10]
      UpperCamelCase__			:   Tuple       				=  np.frombuffer(image_from_tuple.tobytes()							,   dtype="uint8"						)[:10]
      UpperCamelCase__			:   Dict       				=  np.array([69, 2_55, 2_55, 2_55, 0, 0, 77, 1_81, 12, 1_27]						)

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

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

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

      np.random.seed(0						)
      UpperCamelCase__			:   str       				=  np.random.uniform(-1							,   1							,   ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)						)
      UpperCamelCase__			:   Optional[int]       				=  torch.Generator(device=SCREAMING_SNAKE_CASE						).manual_seed(42						)
      UpperCamelCase__			:   Any       				=  pipe(raw_audio=SCREAMING_SNAKE_CASE							,   generator=SCREAMING_SNAKE_CASE							,   start_step=5							,   steps=10						)
      UpperCamelCase__			:   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]
      )
      UpperCamelCase__			:   str       				=  np.frombuffer(image.tobytes()							,   dtype="uint8"						)[:10]
      UpperCamelCase__			:   str       				=  np.array([1_20, 1_17, 1_10, 1_09, 1_38, 1_67, 1_38, 1_48, 1_32, 1_21]						)

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

      UpperCamelCase__			:   List[Any]       				=  self.dummy_unet_condition
      UpperCamelCase__			:   Any       				=  AudioDiffusionPipeline(
          vqvae=self.dummy_vqvae_and_unet[0]							,   unet=SCREAMING_SNAKE_CASE							,   mel=SCREAMING_SNAKE_CASE							,   scheduler=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Tuple       				=  pipe.to(SCREAMING_SNAKE_CASE						)
      pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

      np.random.seed(0						)
      UpperCamelCase__			:   List[Any]       				=  torch.rand((1, 1, 10)						)
      UpperCamelCase__			:   Any       				=  pipe(generator=SCREAMING_SNAKE_CASE							,   encoding=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Tuple       				=  output.images[0]
      UpperCamelCase__			:   Dict       				=  np.frombuffer(image.tobytes()							,   dtype="uint8"						)[:10]
      UpperCamelCase__			:   Optional[int]       				=  np.array([1_07, 1_03, 1_20, 1_27, 1_42, 1_22, 1_13, 1_22, 97, 1_11]						)

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







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







 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






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







 def 					__lowercase					(		self					:   str						):


      '''simple docstring'''






      UpperCamelCase__			:   Union[str, Any]       				=  torch_device

      UpperCamelCase__			:   Tuple       				=  DiffusionPipeline.from_pretrained("teticio/audio-diffusion-ddim-256"						)
      UpperCamelCase__			:   List[Any]       				=  pipe.to(SCREAMING_SNAKE_CASE						)
      pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

      UpperCamelCase__			:   Union[str, Any]       				=  torch.Generator(device=SCREAMING_SNAKE_CASE						).manual_seed(42						)
      UpperCamelCase__			:   Optional[Any]       				=  pipe(generator=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Any       				=  output.audios[0]
      UpperCamelCase__			:   List[str]       				=  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__			:   List[str]       				=  np.frombuffer(image.tobytes()							,   dtype="uint8"						)[:10]
      UpperCamelCase__			:   Union[str, Any]       				=  np.array([1_51, 1_67, 1_54, 1_44, 1_22, 1_34, 1_21, 1_05, 70, 26]						)

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







from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    flip_channel_order,
    get_resize_output_image_size,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, is_vision_available, logging


if is_vision_available():
 import PIL

if is_torch_available():
 import torch


lowerCamelCase	:							Any				=logging.get_logger(__name__)





class 			__a  (    A__       ):
 _lowerCAmelCase       :       List[str]        = ['''pixel_values''']







 def __init__(		self					:   List[str]							,   SCREAMING_SNAKE_CASE					:   bool = True							,   SCREAMING_SNAKE_CASE					:   Dict[str, int] = None							,   SCREAMING_SNAKE_CASE					:   PILImageResampling = PILImageResampling.BILINEAR							,   SCREAMING_SNAKE_CASE					:   bool = True							,   SCREAMING_SNAKE_CASE					:   Union[int, float] = 1 / 2_55							,   SCREAMING_SNAKE_CASE					:   bool = True							,   SCREAMING_SNAKE_CASE					:   Dict[str, int] = None							,   SCREAMING_SNAKE_CASE					:   bool = True							,   **SCREAMING_SNAKE_CASE					:   Tuple							,   ):


      '''simple docstring'''






      super().__init__(**SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   List[str]       				=  size if size is not None else {"shortest_edge": 2_24}
      UpperCamelCase__			:   Any       				=  get_size_dict(SCREAMING_SNAKE_CASE							,   default_to_square=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   List[Any]       				=  crop_size if crop_size is not None else {"height": 2_56, "width": 2_56}
      UpperCamelCase__			:   int       				=  get_size_dict(SCREAMING_SNAKE_CASE							,   param_name="crop_size"						)

      UpperCamelCase__			:   Dict       				=  do_resize
      UpperCamelCase__			:   List[str]       				=  size
      UpperCamelCase__			:   int       				=  resample
      UpperCamelCase__			:   Optional[int]       				=  do_rescale
      UpperCamelCase__			:   List[Any]       				=  rescale_factor
      UpperCamelCase__			:   Union[str, Any]       				=  do_center_crop
      UpperCamelCase__			:   int       				=  crop_size
      UpperCamelCase__			:   Optional[int]       				=  do_flip_channel_order







 def 					__lowercase					(		self					:   Optional[Any]							,   SCREAMING_SNAKE_CASE					:   np.ndarray							,   SCREAMING_SNAKE_CASE					:   Dict[str, int]							,   SCREAMING_SNAKE_CASE					:   PILImageResampling = PIL.Image.BILINEAR							,   SCREAMING_SNAKE_CASE					:   Optional[Union[str, ChannelDimension]] = None							,   **SCREAMING_SNAKE_CASE					:   Optional[Any]							,   ):


      '''simple docstring'''






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







 def 					__lowercase					(		self					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   np.ndarray							,   SCREAMING_SNAKE_CASE					:   Dict[str, int]							,   SCREAMING_SNAKE_CASE					:   Optional[Union[str, ChannelDimension]] = None							,   **SCREAMING_SNAKE_CASE					:   Optional[Any]							,   ):


      '''simple docstring'''






      UpperCamelCase__			:   List[Any]       				=  get_size_dict(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()}'						)
      return center_crop(SCREAMING_SNAKE_CASE							,   size=(size["height"], size["width"])							,   data_format=SCREAMING_SNAKE_CASE							,   **SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   int							,   SCREAMING_SNAKE_CASE					:   np.ndarray							,   SCREAMING_SNAKE_CASE					:   Union[int, float]							,   SCREAMING_SNAKE_CASE					:   Optional[Union[str, ChannelDimension]] = None							,   **SCREAMING_SNAKE_CASE					:   Union[str, Any]							,   ):


      '''simple docstring'''






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







 def 					__lowercase					(		self					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   np.ndarray							,   SCREAMING_SNAKE_CASE					:   Optional[Union[str, ChannelDimension]] = None						):


      '''simple docstring'''






      return flip_channel_order(SCREAMING_SNAKE_CASE							,   data_format=SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   Optional[Any]							,   SCREAMING_SNAKE_CASE					:   ImageInput							,   SCREAMING_SNAKE_CASE					:   bool = None							,   SCREAMING_SNAKE_CASE					:   Dict[str, int] = None							,   SCREAMING_SNAKE_CASE					:   PILImageResampling = None							,   SCREAMING_SNAKE_CASE					:   bool = None							,   SCREAMING_SNAKE_CASE					:   float = None							,   SCREAMING_SNAKE_CASE					:   bool = None							,   SCREAMING_SNAKE_CASE					:   Dict[str, int] = None							,   SCREAMING_SNAKE_CASE					:   bool = None							,   SCREAMING_SNAKE_CASE					:   Optional[Union[str, TensorType]] = None							,   SCREAMING_SNAKE_CASE					:   ChannelDimension = ChannelDimension.FIRST							,   **SCREAMING_SNAKE_CASE					:   Dict							,   ):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  do_resize if do_resize is not None else self.do_resize
      UpperCamelCase__			:   List[str]       				=  resample if resample is not None else self.resample
      UpperCamelCase__			:   Any       				=  do_rescale if do_rescale is not None else self.do_rescale
      UpperCamelCase__			:   Tuple       				=  rescale_factor if rescale_factor is not None else self.rescale_factor
      UpperCamelCase__			:   Dict       				=  do_center_crop if do_center_crop is not None else self.do_center_crop
      UpperCamelCase__			:   Any       				=  (
          do_flip_channel_order if do_flip_channel_order is not None else self.do_flip_channel_order
      )

      UpperCamelCase__			:   Optional[int]       				=  size if size is not None else self.size
      UpperCamelCase__			:   List[str]       				=  get_size_dict(SCREAMING_SNAKE_CASE							,   default_to_square=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Any       				=  crop_size if crop_size is not None else self.crop_size
      UpperCamelCase__			:   Union[str, Any]       				=  get_size_dict(SCREAMING_SNAKE_CASE							,   param_name="crop_size"						)

      UpperCamelCase__			:   int       				=  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:
           raise ValueError("Size 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_center_crop and crop_size is None:
           raise ValueError("Crop size must be specified if do_center_crop is True."						)

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

      if do_resize:
           UpperCamelCase__			:   Optional[int]       				=  [self.resize(image=SCREAMING_SNAKE_CASE							,   size=SCREAMING_SNAKE_CASE							,   resample=SCREAMING_SNAKE_CASE						) for image in images]

      if do_center_crop:
           UpperCamelCase__			:   Any       				=  [self.center_crop(image=SCREAMING_SNAKE_CASE							,   size=SCREAMING_SNAKE_CASE						) for image in images]

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

      # the pretrained checkpoints assume images are BGR, not RGB
      if do_flip_channel_order:
           UpperCamelCase__			:   Any       				=  [self.flip_channel_order(image=SCREAMING_SNAKE_CASE						) for image in images]

      UpperCamelCase__			:   Optional[int]       				=  [to_channel_dimension_format(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						) for image in images]

      UpperCamelCase__			:   Optional[int]       				=  {"pixel_values": images}
      return BatchFeature(data=SCREAMING_SNAKE_CASE							,   tensor_type=SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   Dict							,   SCREAMING_SNAKE_CASE					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   List[Tuple] = None						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[int]       				=  outputs.logits

      # Resize logits and compute semantic segmentation maps
      if target_sizes is not None:
           if len(SCREAMING_SNAKE_CASE						) != len(SCREAMING_SNAKE_CASE						):
                raise ValueError(
                    "Make sure that you pass in as many target sizes as the batch dimension of the logits"						)

           if is_torch_tensor(SCREAMING_SNAKE_CASE						):
                UpperCamelCase__			:   Optional[Any]       				=  target_sizes.numpy()

           UpperCamelCase__			:   Any       				=  []

           for idx in range(len(SCREAMING_SNAKE_CASE						)						):
                UpperCamelCase__			:   Optional[Any]       				=  torch.nn.functional.interpolate(
                    logits[idx].unsqueeze(dim=0						)							,   size=target_sizes[idx]							,   mode="bilinear"							,   align_corners=SCREAMING_SNAKE_CASE						)
                UpperCamelCase__			:   List[str]       				=  resized_logits[0].argmax(dim=0						)
                semantic_segmentation.append(SCREAMING_SNAKE_CASE						)
      else:
           UpperCamelCase__			:   List[str]       				=  logits.argmax(dim=1						)
           UpperCamelCase__			:   Dict       				=  [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]						)]

      return semantic_segmentation