






import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import datasets
import datasets.config

from .utils import require_beam


class    _lowerCAmelCase			(		datasets.BeamBasedBuilder      ):



				"""simple docstring"""



				def 					__magic_name__							(      self    :    Optional[Any]):






											'''simple docstring'''







											return datasets.DatasetInfo(
											    features=datasets.Features({"""content""": datasets.Value("""string""")})			,		supervised_keys=UpperCamelCase__			,		)



				def 					__magic_name__							(      self    :    Dict			,		UpperCamelCase__    :    Optional[int]			,		UpperCamelCase__    :    Optional[Any]):






											'''simple docstring'''







											return [datasets.SplitGenerator(name=datasets.Split.TRAIN			,		gen_kwargs={"""examples""": get_test_dummy_examples()})]



				def 					__magic_name__							(      self    :    Optional[Any]			,		UpperCamelCase__    :    int			,		UpperCamelCase__    :    Dict):






											'''simple docstring'''







											import apache_beam as beam

											return pipeline | "Load Examples" >> beam.Create(UpperCamelCase__)


class    _lowerCAmelCase			(		datasets.BeamBasedBuilder      ):



				"""simple docstring"""



				def 					__magic_name__							(      self    :    Dict):






											'''simple docstring'''







											return datasets.DatasetInfo(
											    features=datasets.Features({"""a""": datasets.Sequence({"""b""": datasets.Value("""string""")})})			,		supervised_keys=UpperCamelCase__			,		)



				def 					__magic_name__							(      self    :    str			,		UpperCamelCase__    :    List[Any]			,		UpperCamelCase__    :    Dict):






											'''simple docstring'''







											return [
											    datasets.SplitGenerator(name=datasets.Split.TRAIN			,		gen_kwargs={"""examples""": get_test_nested_examples()})
											]



				def 					__magic_name__							(      self    :    Union[str, Any]			,		UpperCamelCase__    :    Tuple			,		UpperCamelCase__    :    str):






											'''simple docstring'''







											import apache_beam as beam

											return pipeline | "Load Examples" >> beam.Create(UpperCamelCase__)


def _UpperCAmelCase   (	):
							return [(i, {"content": content}) for i, content in enumerate(["""foo""", """bar""", """foobar"""] )]

def _UpperCAmelCase   (	):
							return [(i, {"a": {"b": [content]}}) for i, content in enumerate(["""foo""", """bar""", """foobar"""] )]



class    _lowerCAmelCase			(		lowercase_      ):



				"""simple docstring"""



				@require_beam
				def 					__magic_name__							(      self    :    Any):






											'''simple docstring'''







											snake_case__	  =       len(get_test_dummy_examples())
											with tempfile.TemporaryDirectory() as tmp_cache_dir:
																		snake_case__	  =       DummyBeamDataset(cache_dir=UpperCamelCase__			,		beam_runner="""DirectRunner""")
																		builder.download_and_prepare()
																		self.assertTrue(
																		    os.path.exists(
																		        os.path.join(UpperCamelCase__			,		builder.name			,		"""default"""			,		"""0.0.0"""			,		F'''{builder.name}-train.arrow''')))
																		self.assertDictEqual(builder.info.features			,		datasets.Features({"""content""": datasets.Value("""string""")}))
																		snake_case__	  =       builder.as_dataset()
																		self.assertEqual(dset["""train"""].num_rows			,		UpperCamelCase__)
																		self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples			,		UpperCamelCase__)
																		self.assertDictEqual(dset["""train"""][0]			,		get_test_dummy_examples()[0][1])
																		self.assertDictEqual(
																		    dset["""train"""][expected_num_examples - 1]			,		get_test_dummy_examples()[expected_num_examples - 1][1])
																		self.assertTrue(
																		    os.path.exists(os.path.join(UpperCamelCase__			,		builder.name			,		"""default"""			,		"""0.0.0"""			,		"""dataset_info.json""")))
																		del dset



				@require_beam
				def 					__magic_name__							(      self    :    Dict):






											'''simple docstring'''







											import apache_beam as beam

											snake_case__	  =       beam.io.parquetio.WriteToParquet

											snake_case__	  =       len(get_test_dummy_examples())
											with tempfile.TemporaryDirectory() as tmp_cache_dir:
																		snake_case__	  =       DummyBeamDataset(cache_dir=UpperCamelCase__			,		beam_runner="""DirectRunner""")
																		with patch("""apache_beam.io.parquetio.WriteToParquet""") as write_parquet_mock:
																									snake_case__	  =       partial(UpperCamelCase__			,		num_shards=2)
																									builder.download_and_prepare()
																		self.assertTrue(
																		    os.path.exists(
																		        os.path.join(
																		            UpperCamelCase__			,		builder.name			,		"""default"""			,		"""0.0.0"""			,		F'''{builder.name}-train-00000-of-00002.arrow''')))
																		self.assertTrue(
																		    os.path.exists(
																		        os.path.join(
																		            UpperCamelCase__			,		builder.name			,		"""default"""			,		"""0.0.0"""			,		F'''{builder.name}-train-00000-of-00002.arrow''')))
																		self.assertDictEqual(builder.info.features			,		datasets.Features({"""content""": datasets.Value("""string""")}))
																		snake_case__	  =       builder.as_dataset()
																		self.assertEqual(dset["""train"""].num_rows			,		UpperCamelCase__)
																		self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples			,		UpperCamelCase__)
																		# Order is not preserved when sharding, so we just check that all the elements are there
																		self.assertListEqual(sorted(dset["""train"""]["""content"""])			,		sorted(["""foo""", """bar""", """foobar"""]))
																		self.assertTrue(
																		    os.path.exists(os.path.join(UpperCamelCase__			,		builder.name			,		"""default"""			,		"""0.0.0"""			,		"""dataset_info.json""")))
																		del dset



				@require_beam
				def 					__magic_name__							(      self    :    Any):






											'''simple docstring'''







											with tempfile.TemporaryDirectory() as tmp_cache_dir:
																		snake_case__	  =       DummyBeamDataset(cache_dir=UpperCamelCase__)
																		self.assertRaises(datasets.builder.MissingBeamOptions			,		builder.download_and_prepare)



				@require_beam
				def 					__magic_name__							(      self    :    int):






											'''simple docstring'''







											snake_case__	  =       len(get_test_nested_examples())
											with tempfile.TemporaryDirectory() as tmp_cache_dir:
																		snake_case__	  =       NestedBeamDataset(cache_dir=UpperCamelCase__			,		beam_runner="""DirectRunner""")
																		builder.download_and_prepare()
																		self.assertTrue(
																		    os.path.exists(
																		        os.path.join(UpperCamelCase__			,		builder.name			,		"""default"""			,		"""0.0.0"""			,		F'''{builder.name}-train.arrow''')))
																		self.assertDictEqual(
																		    builder.info.features			,		datasets.Features({"""a""": datasets.Sequence({"""b""": datasets.Value("""string""")})}))
																		snake_case__	  =       builder.as_dataset()
																		self.assertEqual(dset["""train"""].num_rows			,		UpperCamelCase__)
																		self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples			,		UpperCamelCase__)
																		self.assertDictEqual(dset["""train"""][0]			,		get_test_nested_examples()[0][1])
																		self.assertDictEqual(
																		    dset["""train"""][expected_num_examples - 1]			,		get_test_nested_examples()[expected_num_examples - 1][1])
																		self.assertTrue(
																		    os.path.exists(os.path.join(UpperCamelCase__			,		builder.name			,		"""default"""			,		"""0.0.0"""			,		"""dataset_info.json""")))
																		del dset








import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()


class    _lowerCAmelCase			(		lowercase_     ,      lowercase_     ,      unittest.TestCase      ):



				"""simple docstring"""



				_lowercase				:   Optional[int]      =  VQModel
				_lowercase				:   str      =  '''sample'''



				@property
				def 					__magic_name__							(      self    :    Union[str, Any]			,		UpperCamelCase__    :    Tuple=(3_2, 3_2)):






											'''simple docstring'''







											snake_case__	  =       4
											snake_case__	  =       3

											snake_case__	  =       floats_tensor((batch_size, num_channels) + sizes).to(UpperCamelCase__)

											return {"sample": image}



				@property
				def 					__magic_name__							(      self    :    str):






											'''simple docstring'''







											return (3, 3_2, 3_2)



				@property
				def 					__magic_name__							(      self    :    List[str]):






											'''simple docstring'''







											return (3, 3_2, 3_2)



				def 					__magic_name__							(      self    :    int):






											'''simple docstring'''







											snake_case__	  =       {
											    """block_out_channels""": [3_2, 6_4],
											    """in_channels""": 3,
											    """out_channels""": 3,
											    """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
											    """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
											    """latent_channels""": 3,
											}
											snake_case__	  =       self.dummy_input
											return init_dict, inputs_dict



				def 					__magic_name__							(      self    :    Optional[int]):






											'''simple docstring'''







											pass



				def 					__magic_name__							(      self    :    Tuple):






											'''simple docstring'''







											pass



				def 					__magic_name__							(      self    :    str):






											'''simple docstring'''







											snake_case__  , snake_case__	  =       VQModel.from_pretrained("""fusing/vqgan-dummy"""			,		output_loading_info=UpperCamelCase__)
											self.assertIsNotNone(UpperCamelCase__)
											self.assertEqual(len(loading_info["""missing_keys"""])			,		0)

											model.to(UpperCamelCase__)
											snake_case__	  =       model(**self.dummy_input)

											assert image is not None, "Make sure output is not None"



				def 					__magic_name__							(      self    :    Dict):






											'''simple docstring'''







											snake_case__	  =       VQModel.from_pretrained("""fusing/vqgan-dummy""")
											model.to(UpperCamelCase__).eval()

											torch.manual_seed(0)
											if torch.cuda.is_available():
																		torch.cuda.manual_seed_all(0)

											snake_case__	  =       torch.randn(1			,		model.config.in_channels			,		model.config.sample_size			,		model.config.sample_size)
											snake_case__	  =       image.to(UpperCamelCase__)
											with torch.no_grad():
																		snake_case__	  =       model(UpperCamelCase__).sample

											snake_case__	  =       output[0, -1, -3:, -3:].flatten().cpu()
											# fmt: off
											snake_case__	  =       torch.tensor([-0.01_53, -0.40_44, -0.18_80, -0.51_61, -0.24_18, -0.40_72, -0.16_12, -0.06_33, -0.01_43])
											# fmt: on
											self.assertTrue(torch.allclose(UpperCamelCase__			,		UpperCamelCase__			,		atol=1E-3))
