



import io
import json

import fsspec
import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.json import JsonDatasetReader, JsonDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases

def 			A_      (	__a     :				List[str] ,		__a     :				Optional[int]		):

				"""simple docstring"""
				assert isinstance(__a ,		__a		)
				assert dataset.num_rows == 4
				assert dataset.num_columns == 3
				assert dataset.column_names == ["col_1", "col_2", "col_3"]
				for feature, expected_dtype in expected_features.items():
								assert dataset.features[feature].dtype == expected_dtype

@pytest.mark.parametrize("""keep_in_memory""" ,		[False, True]		)
def 			A_      (	__a     :				Tuple ,		__a     :				Dict ,		__a     :				List[str]		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
								a__				      =		JsonDatasetReader(__a ,		cache_dir=__a ,		keep_in_memory=__a		).read()
				_check_json_dataset(__a ,		__a		)

@pytest.mark.parametrize(
    """features""" ,		[
        None,
        {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""},
        {"""col_1""": """string""", """col_2""": """string""", """col_3""": """string"""},
        {"""col_1""": """int32""", """col_2""": """int32""", """col_3""": """int32"""},
        {"""col_1""": """float32""", """col_2""": """float32""", """col_3""": """float32"""},
    ] ,		)
def 			A_      (	__a     :				int ,		__a     :				Dict ,		__a     :				List[Any]		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				a__				      =		features.copy() if features else default_expected_features
				a__				      =		(
				    Features({feature: Value(__a		) for feature, dtype in features.items()}		) if features is not None else None
				)
				a__				      =		JsonDatasetReader(__a ,		features=__a ,		cache_dir=__a		).read()
				_check_json_dataset(__a ,		__a		)

@pytest.mark.parametrize(
    """features""" ,		[
        None,
        {"""col_3""": """float64""", """col_1""": """string""", """col_2""": """int64"""},
    ] ,		)
def 			A_      (	__a     :				Tuple ,		__a     :				Tuple ,		__a     :				Tuple		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_3""": """float64""", """col_1""": """string""", """col_2""": """int64"""}
				a__				      =		features.copy() if features else default_expected_features
				a__				      =		(
				    Features({feature: Value(__a		) for feature, dtype in features.items()}		) if features is not None else None
				)
				a__				      =		JsonDatasetReader(__a ,		features=__a ,		cache_dir=__a		).read()
				assert isinstance(__a ,		__a		)
				assert dataset.num_rows == 2
				assert dataset.num_columns == 3
				assert dataset.column_names == ["col_3", "col_1", "col_2"]
				for feature, expected_dtype in expected_features.items():
								assert dataset.features[feature].dtype == expected_dtype

def 			A_      (	__a     :				int ,		__a     :				Optional[Any]		):

				"""simple docstring"""
				a__				      =		{"""col_2""": """int64""", """col_3""": """float64""", """col_1""": """string"""}
				a__				      =		features.copy()
				a__				      =		(
				    Features({feature: Value(__a		) for feature, dtype in features.items()}		) if features is not None else None
				)
				a__				      =		tmp_path / """cache"""
				a__				      =		JsonDatasetReader(__a ,		features=__a ,		cache_dir=__a		).read()
				assert isinstance(__a ,		__a		)
				assert dataset.num_rows == 2
				assert dataset.num_columns == 3
				assert dataset.column_names == ["col_2", "col_3", "col_1"]
				for feature, expected_dtype in expected_features.items():
								assert dataset.features[feature].dtype == expected_dtype

@pytest.mark.parametrize("""split""" ,		[None, NamedSplit("""train"""		), """train""", """test"""]		)
def 			A_      (	__a     :				List[Any] ,		__a     :				int ,		__a     :				List[str]		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				a__				      =		JsonDatasetReader(__a ,		cache_dir=__a ,		split=__a		).read()
				_check_json_dataset(__a ,		__a		)
				assert dataset.split == split if split else "train"

@pytest.mark.parametrize("""path_type""" ,		[str, list]		)
def 			A_      (	__a     :				Dict ,		__a     :				Tuple ,		__a     :				List[str]		):

				"""simple docstring"""
				if issubclass(__a ,		__a		):
								a__				      =		jsonl_path
				elif issubclass(__a ,		__a		):
								a__				      =		[jsonl_path]
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				a__				      =		JsonDatasetReader(__a ,		cache_dir=__a		).read()
				_check_json_dataset(__a ,		__a		)

def 			A_      (	__a     :				str ,		__a     :				List[Any] ,		__a     :				Dict=("train",)		):

				"""simple docstring"""
				assert isinstance(__a ,		__a		)
				for split in splits:
								a__				      =		dataset_dict[split]
								assert dataset.num_rows == 4
								assert dataset.num_columns == 3
								assert dataset.column_names == ["col_1", "col_2", "col_3"]
								for feature, expected_dtype in expected_features.items():
												assert dataset.features[feature].dtype == expected_dtype

@pytest.mark.parametrize("""keep_in_memory""" ,		[False, True]		)
def 			A_      (	__a     :				int ,		__a     :				int ,		__a     :				Dict		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
								a__				      =		JsonDatasetReader({"""train""": jsonl_path} ,		cache_dir=__a ,		keep_in_memory=__a		).read()
				_check_json_datasetdict(__a ,		__a		)

@pytest.mark.parametrize(
    """features""" ,		[
        None,
        {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""},
        {"""col_1""": """string""", """col_2""": """string""", """col_3""": """string"""},
        {"""col_1""": """int32""", """col_2""": """int32""", """col_3""": """int32"""},
        {"""col_1""": """float32""", """col_2""": """float32""", """col_3""": """float32"""},
    ] ,		)
def 			A_      (	__a     :				Tuple ,		__a     :				Dict ,		__a     :				Optional[Any]		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				a__				      =		features.copy() if features else default_expected_features
				a__				      =		(
				    Features({feature: Value(__a		) for feature, dtype in features.items()}		) if features is not None else None
				)
				a__				      =		JsonDatasetReader({"""train""": jsonl_path} ,		features=__a ,		cache_dir=__a		).read()
				_check_json_datasetdict(__a ,		__a		)

@pytest.mark.parametrize("""split""" ,		[None, NamedSplit("""train"""		), """train""", """test"""]		)
def 			A_      (	__a     :				Dict ,		__a     :				Any ,		__a     :				Tuple		):

				"""simple docstring"""
				if split:
								a__				      =		{split: jsonl_path}
				else:
								a__				      =		"""train"""
								a__				      =		{"""train""": jsonl_path, """test""": jsonl_path}
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				a__				      =		JsonDatasetReader(__a ,		cache_dir=__a		).read()
				_check_json_datasetdict(__a ,		__a ,		splits=list(path.keys()		)		)
				assert all(dataset[split].split == split for split in path.keys()		)

def 			A_      (	__a     :				Any		):

				"""simple docstring"""
				return json.load(__a		)

def 			A_      (	__a     :				Tuple		):

				"""simple docstring"""
				return [json.loads(__a		) for line in buffer]




class  __snake_case      :







							'''simple docstring'''






							@pytest.mark.parametrize("""lines, load_json_function"""  ,  [(True, load_json_lines), (False, load_json)]		)
							def _a      (							self  ,  a_  ,  a_  ,  a_		):
											with io.BytesIO() as buffer:
															JsonDatasetWriter(a_  ,  a_  ,  lines=a_		).write()
															buffer.seek(0		)
															a__				      =		load_json_function(a_		)
											assert isinstance(a_  ,  a_		)
											assert isinstance(exported_content[0]  ,  a_		)
											assert len(a_		) == 10


							@pytest.mark.parametrize(
							    """orient, container, keys, len_at"""  ,  [
							        ("""records""", list, {"""tokens""", """labels""", """answers""", """id"""}, None),
							        ("""split""", dict, {"""columns""", """data"""}, """data"""),
							        ("""index""", dict, set("""0123456789"""		), None),
							        ("""columns""", dict, {"""tokens""", """labels""", """answers""", """id"""}, """tokens"""),
							        ("""values""", list, None, None),
							        ("""table""", dict, {"""schema""", """data"""}, """data"""),
							    ]  ,  )
							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											with io.BytesIO() as buffer:
															JsonDatasetWriter(a_  ,  a_  ,  lines=a_  ,  orient=a_		).write()
															buffer.seek(0		)
															a__				      =		load_json(a_		)
											assert isinstance(a_  ,  a_		)
											if keys:
															if container is dict:
																			assert exported_content.keys() == keys
															else:
																			assert exported_content[0].keys() == keys
											else:
															assert not hasattr(a_  ,  """keys"""		) and not hasattr(exported_content[0]  ,  """keys"""		)
											if len_at:
															assert len(exported_content[len_at]		) == 10
											else:
															assert len(a_		) == 10


							@pytest.mark.parametrize("""lines, load_json_function"""  ,  [(True, load_json_lines), (False, load_json)]		)
							def _a      (							self  ,  a_  ,  a_  ,  a_		):
											with io.BytesIO() as buffer:
															JsonDatasetWriter(a_  ,  a_  ,  lines=a_  ,  num_proc=2		).write()
															buffer.seek(0		)
															a__				      =		load_json_function(a_		)
											assert isinstance(a_  ,  a_		)
											assert isinstance(exported_content[0]  ,  a_		)
											assert len(a_		) == 10


							@pytest.mark.parametrize(
							    """orient, container, keys, len_at"""  ,  [
							        ("""records""", list, {"""tokens""", """labels""", """answers""", """id"""}, None),
							        ("""split""", dict, {"""columns""", """data"""}, """data"""),
							        ("""index""", dict, set("""0123456789"""		), None),
							        ("""columns""", dict, {"""tokens""", """labels""", """answers""", """id"""}, """tokens"""),
							        ("""values""", list, None, None),
							        ("""table""", dict, {"""schema""", """data"""}, """data"""),
							    ]  ,  )
							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											with io.BytesIO() as buffer:
															JsonDatasetWriter(a_  ,  a_  ,  lines=a_  ,  orient=a_  ,  num_proc=2		).write()
															buffer.seek(0		)
															a__				      =		load_json(a_		)
											assert isinstance(a_  ,  a_		)
											if keys:
															if container is dict:
																			assert exported_content.keys() == keys
															else:
																			assert exported_content[0].keys() == keys
											else:
															assert not hasattr(a_  ,  """keys"""		) and not hasattr(exported_content[0]  ,  """keys"""		)
											if len_at:
															assert len(exported_content[len_at]		) == 10
											else:
															assert len(a_		) == 10


							def _a      (							self  ,  a_		):
											with pytest.raises(a_		):
															with io.BytesIO() as buffer:
																			JsonDatasetWriter(a_  ,  a_  ,  num_proc=0		)







							@pytest.mark.parametrize("""compression, extension"""  ,  [("""gzip""", """gz"""), ("""bz2""", """bz2"""), ("""xz""", """xz""")]		)
							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		tmp_path_factory.mktemp("""data"""		) / F'''test.json.{extension}'''
											a__				      =		str(shared_datadir / F'''test_file.json.{extension}'''		)
											JsonDatasetWriter(a_  ,  a_  ,  compression=a_		).write()

											with fsspec.open(a_  ,  """rb"""  ,  compression="""infer"""		) as f:
															a__				      =		f.read()
											with fsspec.open(a_  ,  """rb"""  ,  compression="""infer"""		) as f:
															a__				      =		f.read()
											assert exported_content == original_content






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 _a      (							self		):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()


							@property
							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		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 _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		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 _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		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__				      =		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 _a      (							self		):
											a__				      =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
											a__				      =		Mel(
											    x_res=self.dummy_unet.config.sample_size[1]  ,  y_res=self.dummy_unet.config.sample_size[0]  ,  )

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

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_  ,  steps=4		)
											a__				      =		output.audios[0]
											a__				      =		output.images[0]

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

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

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

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

											np.random.seed(0		)
											a__				      =		torch.rand((1, 1, 10)		)
											a__				      =		pipe(generator=a_  ,  encoding=a_		)
											a__				      =		output.images[0]
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		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 _a      (							self		):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()







							def _a      (							self		):
											a__				      =		torch_device

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

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_		)
											a__				      =		output.audios[0]
											a__				      =		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__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]		)

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

