


'''simple docstring'''
import itertools
from dataclasses import dataclass
from typing import List, Optional

import pyarrow as pa
import pyarrow.parquet as pq

import datasets
from datasets.table import table_cast


UpperCAmelCase_ :     List[Any]								=      datasets.utils.logging.get_logger(__name__)

@dataclass
class 						_lowerCamelCase					(    datasets.BuilderConfig				):




 '''simple docstring'''
 __lowercase				:		str    							=						1_0_0_0_0
 __lowercase				:		Tuple    							=						None
 __lowercase				:		List[Any]    							=						None


class 						_lowerCamelCase					(    datasets.ArrowBasedBuilder				):




 '''simple docstring'''
 __lowercase				:		Any    							=						ParquetConfig


 def  snake_case__   (      self						):

  """simple docstring"""







  return datasets.DatasetInfo(features=self.config.features						)


 def  snake_case__   (      self      ,      __lowercase						):

  """simple docstring"""







  if not self.config.data_files:
   raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}"""						)
  __A  :		str       =			dl_manager.download_and_extract(self.config.data_files						)
  if isinstance(_lowercase      ,      (str, list, tuple)						):
   __A  :		Union[str, Any]       =			data_files
   if isinstance(_lowercase      ,      _lowercase						):
    __A  :		Any       =			[files]
   # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
   __A  :		str       =			[dl_manager.iter_files(_lowercase						) for file in files]
   return [datasets.SplitGenerator(name=datasets.Split.TRAIN      ,      gen_kwargs={'files': files}						)]
  __A  :		Union[str, Any]       =			[]
  for split_name, files in data_files.items():
   if isinstance(_lowercase      ,      _lowercase						):
    __A  :		Optional[int]       =			[files]
   # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
   __A  :		List[str]       =			[dl_manager.iter_files(_lowercase						) for file in files]
   # Infer features is they are stoed in the arrow schema
   if self.info.features is None:
    for file in itertools.chain.from_iterable(_lowercase						):
     with open(_lowercase      ,      'rb'						) as f:
      __A  :		List[str]       =			datasets.Features.from_arrow_schema(pq.read_schema(_lowercase						)						)
     break
   splits.append(datasets.SplitGenerator(name=_lowercase      ,      gen_kwargs={'files': files}						)						)
  return splits


 def  snake_case__   (      self      ,      __lowercase						):

  """simple docstring"""







  if self.info.features is not None:
   # more expensive cast to support nested features with keys in a different order
   # allows str <-> int/float or str to Audio for example
   __A  :		Dict       =			table_cast(_lowercase      ,      self.info.features.arrow_schema						)
  return pa_table







 def  snake_case__   (      self      ,      __lowercase						):

  """simple docstring"""







  __A  :		Optional[int]       =			self.info.features.arrow_schema if self.info.features is not None else None
  if self.info.features is not None and self.config.columns is not None:
   if sorted(field.name for field in schema						) != sorted(self.config.columns						):
    raise ValueError(
        F"""Tried to load parquet data with columns \'{self.config.columns}\' with mismatching features \'{self.info.features}\'"""						)
  for file_idx, file in enumerate(itertools.chain.from_iterable(_lowercase						)						):
   with open(_lowercase      ,      'rb'						) as f:
    __A  :		int       =			pq.ParquetFile(_lowercase						)
    try:
     for batch_idx, record_batch in enumerate(
         parquet_file.iter_batches(batch_size=self.config.batch_size      ,      columns=self.config.columns						)						):
      __A  :		Tuple       =			pa.Table.from_batches([record_batch]						)
      # Uncomment for debugging (will print the Arrow table size and elements)
      # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
      # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
      yield F"""{file_idx}_{batch_idx}""", self._cast_table(_lowercase						)
    except ValueError as e:
     logger.error(F"""Failed to read file \'{file}\' with error {type(_lowercase						)}: {e}"""						)
     raise










'''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  snake_case__   (      self						):

  """simple docstring"""







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


 @property
 def  snake_case__   (      self						):

  """simple docstring"""







  torch.manual_seed(0						)
  __A  :		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  snake_case__   (      self						):

  """simple docstring"""







  torch.manual_seed(0						)
  __A  :		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  snake_case__   (      self						):

  """simple docstring"""







  torch.manual_seed(0						)
  __A  :		Union[str, 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')      ,      )
  __A  :		Union[str, Any]       =			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  snake_case__   (      self						):

  """simple docstring"""







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

  __A  :		Tuple       =			DDPMScheduler()
  __A  :		Optional[int]       =			AudioDiffusionPipeline(vqvae=__lowercase      ,      unet=self.dummy_unet      ,      mel=__lowercase      ,      scheduler=__lowercase						)
  __A  :		Dict       =			pipe.to(__lowercase						)
  pipe.set_progress_bar_config(disable=__lowercase						)

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

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

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

  __A  :		Optional[Any]       =			DDIMScheduler()
  __A  :		List[Any]       =			self.dummy_vqvae_and_unet
  __A  :		int       =			AudioDiffusionPipeline(
      vqvae=self.dummy_vqvae_and_unet[0]      ,      unet=dummy_vqvae_and_unet[1]      ,      mel=__lowercase      ,      scheduler=__lowercase						)
  __A  :		Dict       =			pipe.to(__lowercase						)
  pipe.set_progress_bar_config(disable=__lowercase						)

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

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

  __A  :		Any       =			self.dummy_unet_condition
  __A  :		Optional[int]       =			AudioDiffusionPipeline(
      vqvae=self.dummy_vqvae_and_unet[0]      ,      unet=__lowercase      ,      mel=__lowercase      ,      scheduler=__lowercase						)
  __A  :		Union[str, Any]       =			pipe.to(__lowercase						)
  pipe.set_progress_bar_config(disable=__lowercase						)

  np.random.seed(0						)
  __A  :		Dict       =			torch.rand((1, 1, 10)						)
  __A  :		Any       =			pipe(generator=__lowercase      ,      encoding=__lowercase						)
  __A  :		Union[str, Any]       =			output.images[0]
  __A  :		List[Any]       =			np.frombuffer(image.tobytes()      ,      dtype='uint8'						)[:10]
  __A  :		int       =			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  snake_case__   (      self						):

  """simple docstring"""







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







 def  snake_case__   (      self						):

  """simple docstring"""







  __A  :		Union[str, Any]       =			torch_device

  __A  :		Union[str, Any]       =			DiffusionPipeline.from_pretrained('teticio/audio-diffusion-ddim-256'						)
  __A  :		Dict       =			pipe.to(__lowercase						)
  pipe.set_progress_bar_config(disable=__lowercase						)

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

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






