





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







					'''simple docstring'''

					def        SCREAMING_SNAKE_CASE							(						self    :  Any		):






						'''simple docstring'''







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





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






						'''simple docstring'''







						torch.manual_seed(0		)
						UpperCAmelCase_								=     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        SCREAMING_SNAKE_CASE							(						self    :  int		):






						'''simple docstring'''







						torch.manual_seed(0		)
						UpperCAmelCase_								=     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        SCREAMING_SNAKE_CASE							(						self    :  int		):






						'''simple docstring'''







						torch.manual_seed(0		)
						UpperCAmelCase_								=     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_								=     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        SCREAMING_SNAKE_CASE							(						self    :  int		):






						'''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(42		)
						UpperCAmelCase_								=     pipe(generator=__lowercase ,					steps=4		)
						UpperCAmelCase_								=     output.audios[0]
						UpperCAmelCase_								=     output.images[0]

						UpperCAmelCase_								=     torch.Generator(device=__lowercase		).manual_seed(42		)
						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"""		)[:10]
						UpperCAmelCase_								=     np.frombuffer(image_from_tuple.tobytes() ,					dtype="""uint8"""		)[:10]
						UpperCAmelCase_								=     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_								=     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(42		)
						UpperCAmelCase_								=     pipe(raw_audio=__lowercase ,					generator=__lowercase ,					start_step=5 ,					steps=10		)
						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"""		)[:10]
						UpperCAmelCase_								=     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_								=     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, 10)		)
						UpperCAmelCase_								=     pipe(generator=__lowercase ,					encoding=__lowercase		)
						UpperCAmelCase_								=     output.images[0]
						UpperCAmelCase_								=     np.frombuffer(image.tobytes() ,					dtype="""uint8"""		)[:10]
						UpperCAmelCase_								=     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   _UpperCamelCase  (       unittest.TestCase  ):







					'''simple docstring'''

					def        SCREAMING_SNAKE_CASE							(						self    :  int		):






						'''simple docstring'''







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





					def        SCREAMING_SNAKE_CASE							(						self    :  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(42		)
						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"""		)[:10]
						UpperCAmelCase_								=     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









import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


UpperCamelCase__   : Tuple					  = 5_00_00
UpperCamelCase__   : Any					  = 50_00

UpperCamelCase__  ,			UpperCamelCase__   : Tuple					  = os.path.split(__file__)
UpperCamelCase__   : Union[str, Any]					  = os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))




@get_duration
def    A_(			A				,      A   ):
	for i in range(A   ):
		UpperCAmelCase_								=     dataset[i]




@get_duration
def    A_(			A				,      A				,      A   ):
	for i in range(0				,      len(A   )				,      A   ):
		UpperCAmelCase_								=     dataset[i : i + batch_size]




@get_duration
def    A_(			A				,      A				,      A   ):
	with dataset.formatted_as(type=A   ):
		for i in range(A   ):
			UpperCAmelCase_								=     dataset[i]




@get_duration
def    A_(			A				,      A				,      A				,      A   ):
	with dataset.formatted_as(type=A   ):
		for i in range(0				,      A				,      A   ):
			UpperCAmelCase_								=     dataset[i : i + batch_size]




def    A_(			):
	UpperCAmelCase_								=     {"""num examples""": SPEED_TEST_N_EXAMPLES}
	UpperCAmelCase_								=     [
	    (read, {"""length""": SMALL_TEST}),
	    (read, {"""length""": SPEED_TEST_N_EXAMPLES}),
	    (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 10}),
	    (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 100}),
	    (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 1000}),
	    (read_formatted, {"""type""": """numpy""", """length""": SMALL_TEST}),
	    (read_formatted, {"""type""": """pandas""", """length""": SMALL_TEST}),
	    (read_formatted, {"""type""": """torch""", """length""": SMALL_TEST}),
	    (read_formatted, {"""type""": """tensorflow""", """length""": SMALL_TEST}),
	    (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 10}),
	    (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 1000}),
	]

	UpperCAmelCase_								=     [
	    (read, {"""length""": SMALL_TEST}),
	    (read, {"""length""": SPEED_TEST_N_EXAMPLES}),
	    (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 10}),
	    (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 100}),
	    (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 1000}),
	    (read_formatted, {"""type""": """numpy""", """length""": SMALL_TEST}),
	    (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 10}),
	    (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 1000}),
	]
	with tempfile.TemporaryDirectory() as tmp_dir:
		print("""generating dataset"""   )
		UpperCAmelCase_								=     datasets.Features(
		    {"""list""": datasets.Sequence(datasets.Value("""float32"""   )   ), """numbers""": datasets.Value("""float32"""   )}   )
		UpperCAmelCase_								=     generate_example_dataset(
		    os.path.join(A				,      """dataset.arrow"""   )				,      A				,      num_examples=A				,      seq_shapes={"""list""": (100,)}				,      )
		print("""first set of iterations"""   )
		for func, kwargs in functions:
			print(func.__name__				,      str(A   )   )
			UpperCAmelCase_								=     func(A				,      **A   )

		print("""shuffling dataset"""   )
		UpperCAmelCase_								=     dataset.shuffle()
		print("""Second set of iterations (after shuffling"""   )
		for func, kwargs in functions_shuffled:
			print("""shuffled """				,      func.__name__				,      str(A   )   )
			UpperCAmelCase_								=     func(
			    A				,      **A   )

	with open(A				,      """wb"""   ) as f:
		f.write(json.dumps(A   ).encode("""utf-8"""   )   )


if __name__ == "__main__":  # useful to run the profiler
				benchmark_iterating()


