





"""simple docstring"""



# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
       import jax
       import jaxlib

a_					=     get_logger()

a_					=     None






class 						__snake_case   (						TensorFormatter[Mapping, """jax.Array""", Mapping]	):





 """simple docstring"""


 def __init__(       self   ,		__lowerCamelCase=None   ,		__lowerCamelCase=None   ,		**__lowerCamelCase    ):



     '''simple docstring'''




     super().__init__(features=__lowerCamelCase    )
     import jax
     from jaxlib.xla_client import Device

     if isinstance(__lowerCamelCase   ,		__lowerCamelCase    ):
         raise ValueError(
             F"""Expected {device} to be a `str` not {type(__lowerCamelCase    )}, as `jaxlib.xla_extension.Device` """
             '''is not serializable neither with `pickle` nor with `dill`. Instead you can surround '''
             '''the device with `str()` to get its string identifier that will be internally mapped '''
             '''to the actual `jaxlib.xla_extension.Device`.'''    )
     __A       :		List[Any]							=	device if isinstance(__lowerCamelCase   ,		__lowerCamelCase    ) else str(jax.devices()[0]    )
     # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
     # with `pickle` nor with `dill`, so we need to use a global variable instead
     global DEVICE_MAPPING
     if DEVICE_MAPPING is None:
         __A       :		Tuple							=	self._map_devices_to_str()
     if self.device not in list(DEVICE_MAPPING.keys()    ):
         logger.warning(
             F"""Device with string identifier {self.device} not listed among the available """
             F"""devices: {list(DEVICE_MAPPING.keys()    )}, so falling back to the default """
             F"""device: {str(jax.devices()[0]    )}."""    )
         __A       :		Dict							=	str(jax.devices()[0]    )
     __A       :		Union[str, Any]							=	jnp_array_kwargs

 @staticmethod
 def    UpperCamelCase__(       ):



     '''simple docstring'''




     import jax

     return {str(__lowerCamelCase    ): device for device in jax.devices()}

 def    UpperCamelCase__(       self   ,		__lowerCamelCase    ):



     '''simple docstring'''




     import jax
     import jax.numpy as jnp

     if isinstance(__lowerCamelCase   ,		__lowerCamelCase    ) and column:
         if all(
             isinstance(__lowerCamelCase   ,		jax.Array    ) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column    ):
             return jnp.stack(__lowerCamelCase   ,		axis=0    )
     return column

 def    UpperCamelCase__(       self   ,		__lowerCamelCase    ):



     '''simple docstring'''




     import jax
     import jax.numpy as jnp

     if isinstance(__lowerCamelCase   ,		(str, bytes, type(__lowerCamelCase    ))    ):
         return value
     elif isinstance(__lowerCamelCase   ,		(np.character, np.ndarray)    ) and np.issubdtype(value.dtype   ,		np.character    ):
         return value.tolist()

     __A       :		Union[str, Any]							=	{}

     if isinstance(__lowerCamelCase   ,		(np.number, np.ndarray)    ) and np.issubdtype(value.dtype   ,		np.integer    ):
         # the default int precision depends on the jax config
         # see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
         if jax.config.jax_enable_xaa:
             __A       :		Union[str, Any]							=	{'''dtype''': jnp.intaa}
         else:
             __A       :		int							=	{'''dtype''': jnp.intaa}
     elif isinstance(__lowerCamelCase   ,		(np.number, np.ndarray)    ) and np.issubdtype(value.dtype   ,		np.floating    ):
         __A       :		Any							=	{'''dtype''': jnp.floataa}
     elif config.PIL_AVAILABLE and "PIL" in sys.modules:
         import PIL.Image

         if isinstance(__lowerCamelCase   ,		PIL.Image.Image    ):
             __A       :		Optional[int]							=	np.asarray(__lowerCamelCase    )

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
     global DEVICE_MAPPING
     if DEVICE_MAPPING is None:
         __A       :		Dict							=	self._map_devices_to_str()

     with jax.default_device(DEVICE_MAPPING[self.device]    ):
         # calling jnp.array on a np.ndarray does copy the data
         # see https://github.com/google/jax/issues/4486
         return jnp.array(__lowerCamelCase   ,		**{**default_dtype, **self.jnp_array_kwargs}    )

 def    UpperCamelCase__(       self   ,		__lowerCamelCase    ):



     '''simple docstring'''




     import jax

     # support for torch, tf, jax etc.
     if config.TORCH_AVAILABLE and "torch" in sys.modules:
         import torch

         if isinstance(__lowerCamelCase   ,		torch.Tensor    ):
             return self._tensorize(data_struct.detach().cpu().numpy()[()]    )
     if hasattr(__lowerCamelCase   ,		'''__array__'''    ) and not isinstance(__lowerCamelCase   ,		jax.Array    ):
         __A       :		Dict							=	data_struct.__array__()
     # support for nested types like struct of list of struct
     if isinstance(__lowerCamelCase   ,		np.ndarray    ):
         if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
             return self._consolidate([self.recursive_tensorize(__lowerCamelCase    ) for substruct in data_struct]    )
     elif isinstance(__lowerCamelCase   ,		(list, tuple)    ):
         return self._consolidate([self.recursive_tensorize(__lowerCamelCase    ) for substruct in data_struct]    )
     return self._tensorize(__lowerCamelCase    )

 def    UpperCamelCase__(       self   ,		__lowerCamelCase    ):



     '''simple docstring'''




     return map_nested(self._recursive_tensorize   ,		__lowerCamelCase   ,		map_list=__lowerCamelCase    )

 def    UpperCamelCase__(       self   ,		__lowerCamelCase    ):



     '''simple docstring'''




     __A       :		Dict							=	self.numpy_arrow_extractor().extract_row(__lowerCamelCase    )
     __A       :		Optional[Any]							=	self.python_features_decoder.decode_row(__lowerCamelCase    )
     return self.recursive_tensorize(__lowerCamelCase    )

 def    UpperCamelCase__(       self   ,		__lowerCamelCase    ):



     '''simple docstring'''




     __A       :		str							=	self.numpy_arrow_extractor().extract_column(__lowerCamelCase    )
     __A       :		List[str]							=	self.python_features_decoder.decode_column(__lowerCamelCase   ,		pa_table.column_names[0]    )
     __A       :		Optional[int]							=	self.recursive_tensorize(__lowerCamelCase    )
     __A       :		Any							=	self._consolidate(__lowerCamelCase    )
     return column







 def    UpperCamelCase__(       self   ,		__lowerCamelCase    ):



     '''simple docstring'''




     __A       :		List[str]							=	self.numpy_arrow_extractor().extract_batch(__lowerCamelCase    )
     __A       :		Optional[int]							=	self.python_features_decoder.decode_batch(__lowerCamelCase    )
     __A       :		Tuple							=	self.recursive_tensorize(__lowerCamelCase    )
     for column_name in batch:
         __A       :		Any							=	self._consolidate(batch[column_name]    )
     return batch











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





 """simple docstring"""


 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




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

 @property
 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




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



     '''simple docstring'''




     torch.manual_seed(0    )
     __A       :		str							=	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    )
     __A       :		int							=	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       :		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    UpperCamelCase__(       self    ):



     '''simple docstring'''




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

     __A       :		List[Any]							=	DDPMScheduler()
     __A       :		Optional[Any]							=	AudioDiffusionPipeline(vqvae=__lowerCamelCase   ,		unet=self.dummy_unet   ,		mel=__lowerCamelCase   ,		scheduler=__lowerCamelCase    )
     __A       :		Tuple							=	pipe.to(__lowerCamelCase    )
     pipe.set_progress_bar_config(disable=__lowerCamelCase    )

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

     __A       :		List[Any]							=	torch.Generator(device=__lowerCamelCase    ).manual_seed(42    )
     __A       :		Any							=	pipe(generator=__lowerCamelCase   ,		steps=4   ,		return_dict=__lowerCamelCase    )
     __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       :		Tuple							=	np.frombuffer(image.tobytes()   ,		dtype='''uint8'''    )[:10]
     __A       :		str							=	np.frombuffer(image_from_tuple.tobytes()   ,		dtype='''uint8'''    )[:10]
     __A       :		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       :		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       :		Any							=	DDIMScheduler()
     __A       :		Optional[Any]							=	self.dummy_vqvae_and_unet
     __A       :		int							=	AudioDiffusionPipeline(
         vqvae=self.dummy_vqvae_and_unet[0]   ,		unet=dummy_vqvae_and_unet[1]   ,		mel=__lowerCamelCase   ,		scheduler=__lowerCamelCase    )
     __A       :		Any							=	pipe.to(__lowerCamelCase    )
     pipe.set_progress_bar_config(disable=__lowerCamelCase    )

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

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

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

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





 """simple docstring"""


 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




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







 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		int							=	torch_device

     __A       :		List[Any]							=	DiffusionPipeline.from_pretrained('''teticio/audio-diffusion-ddim-256'''    )
     __A       :		List[Any]							=	pipe.to(__lowerCamelCase    )
     pipe.set_progress_bar_config(disable=__lowerCamelCase    )

     __A       :		Dict							=	torch.Generator(device=__lowerCamelCase    ).manual_seed(42    )
     __A       :		str							=	pipe(generator=__lowerCamelCase    )
     __A       :		Dict							=	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       :		Optional[int]							=	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




