




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


__lowerCamelCase   :						List[str]				=			False


class    lowerCamelCase    (		unittest.TestCase						):






			'''simple docstring'''





			def 							UpperCAmelCase__       ( self			:   Optional[Any]  )						->   Any:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()

			@property
			def 							UpperCAmelCase__       ( self			:   List[str]  )						->   Any:
								return 12

			@property
			def 							UpperCAmelCase__       ( self			:   Dict  )						->   List[Any]:
								return 12

			@property
			def 							UpperCAmelCase__       ( self			:   List[str]  )						->   Optional[int]:
								return 32

			@property
			def 							UpperCAmelCase__       ( self			:   Union[str, Any]  )						->   str:
								torch.manual_seed(0  )
								__magic_name__							:	List[Any]    							=							VQModel(
								    block_out_channels=[32, 64]       ,						in_channels=3       ,						out_channels=3       ,						down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']       ,						up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']       ,						latent_channels=3       ,						num_vq_embeddings=self.num_embed       ,						vq_embed_dim=3       ,						)
								return model

			@property
			def 							UpperCAmelCase__       ( self			:   List[str]  )						->   Dict:
								__magic_name__							:	Tuple    							=							CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )
								return tokenizer

			@property
			def 							UpperCAmelCase__       ( self			:   Any  )						->   int:
								torch.manual_seed(0  )
								__magic_name__							:	List[str]    							=							CLIPTextConfig(
								    bos_token_id=0       ,						eos_token_id=2       ,						hidden_size=self.text_embedder_hidden_size       ,						intermediate_size=37       ,						layer_norm_eps=1E-05       ,						num_attention_heads=4       ,						num_hidden_layers=5       ,						pad_token_id=1       ,						vocab_size=1000       ,						)
								return CLIPTextModel(lowerCamelCase_  )

			@property
			def 							UpperCAmelCase__       ( self			:   int  )						->   int:
								torch.manual_seed(0  )

								__magic_name__							:	Optional[int]    							=							12
								__magic_name__							:	int    							=							12

								__magic_name__							:	str    							=							{
								    '''attention_bias''': True,
								    '''cross_attention_dim''': 32,
								    '''attention_head_dim''': height * width,
								    '''num_attention_heads''': 1,
								    '''num_vector_embeds''': self.num_embed,
								    '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
								    '''norm_num_groups''': 32,
								    '''sample_size''': width,
								    '''activation_fn''': '''geglu-approximate''',
								}

								__magic_name__							:	Tuple    							=							TransformeraDModel(**lowerCamelCase_  )
								return model

			def 							UpperCAmelCase__       ( self			:   List[Any]  )						->   str:
								__magic_name__							:	Tuple    							=							'''cpu'''

								__magic_name__							:	List[Any]    							=							self.dummy_vqvae
								__magic_name__							:	Tuple    							=							self.dummy_text_encoder
								__magic_name__							:	Optional[Any]    							=							self.dummy_tokenizer
								__magic_name__							:	Optional[int]    							=							self.dummy_transformer
								__magic_name__							:	List[Any]    							=							VQDiffusionScheduler(self.num_embed  )
								__magic_name__							:	Union[str, Any]    							=							LearnedClassifierFreeSamplingEmbeddings(learnable=lowerCamelCase_  )

								__magic_name__							:	Optional[Any]    							=							VQDiffusionPipeline(
								    vqvae=lowerCamelCase_       ,						text_encoder=lowerCamelCase_       ,						tokenizer=lowerCamelCase_       ,						transformer=lowerCamelCase_       ,						scheduler=lowerCamelCase_       ,						learned_classifier_free_sampling_embeddings=lowerCamelCase_       ,						)
								__magic_name__							:	Union[str, Any]    							=							pipe.to(lowerCamelCase_  )
								pipe.set_progress_bar_config(disable=lowerCamelCase_  )

								__magic_name__							:	Optional[int]    							=							'''teddy bear playing in the pool'''

								__magic_name__							:	Any    							=							torch.Generator(device=lowerCamelCase_  ).manual_seed(0  )
								__magic_name__							:	Union[str, Any]    							=							pipe([prompt]       ,						generator=lowerCamelCase_       ,						num_inference_steps=2       ,						output_type='''np'''  )
								__magic_name__							:	str    							=							output.images

								__magic_name__							:	int    							=							torch.Generator(device=lowerCamelCase_  ).manual_seed(0  )
								__magic_name__							:	Optional[Any]    							=							pipe(
								    [prompt]       ,						generator=lowerCamelCase_       ,						output_type='''np'''       ,						return_dict=lowerCamelCase_       ,						num_inference_steps=2  )[0]

								__magic_name__							:	Any    							=							image[0, -3:, -3:, -1]
								__magic_name__							:	Optional[int]    							=							image_from_tuple[0, -3:, -3:, -1]

								assert image.shape == (1, 24, 24, 3)

								__magic_name__							:	Optional[Any]    							=							np.array([0.6_5_5_1, 0.6_1_6_8, 0.5_0_0_8, 0.5_6_7_6, 0.5_6_5_9, 0.4_2_9_5, 0.6_0_7_3, 0.5_5_9_9, 0.4_9_9_2]  )

								assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2
								assert np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 1E-2


			def 							UpperCAmelCase__       ( self			:   Tuple  )						->   Optional[int]:
								__magic_name__							:	List[Any]    							=							'''cpu'''

								__magic_name__							:	Optional[Any]    							=							self.dummy_vqvae
								__magic_name__							:	int    							=							self.dummy_text_encoder
								__magic_name__							:	List[str]    							=							self.dummy_tokenizer
								__magic_name__							:	str    							=							self.dummy_transformer
								__magic_name__							:	Any    							=							VQDiffusionScheduler(self.num_embed  )
								__magic_name__							:	Any    							=							LearnedClassifierFreeSamplingEmbeddings(
								    learnable=lowerCamelCase_       ,						hidden_size=self.text_embedder_hidden_size       ,						length=tokenizer.model_max_length  )

								__magic_name__							:	Optional[Any]    							=							VQDiffusionPipeline(
								    vqvae=lowerCamelCase_       ,						text_encoder=lowerCamelCase_       ,						tokenizer=lowerCamelCase_       ,						transformer=lowerCamelCase_       ,						scheduler=lowerCamelCase_       ,						learned_classifier_free_sampling_embeddings=lowerCamelCase_       ,						)
								__magic_name__							:	Dict    							=							pipe.to(lowerCamelCase_  )
								pipe.set_progress_bar_config(disable=lowerCamelCase_  )

								__magic_name__							:	Dict    							=							'''teddy bear playing in the pool'''

								__magic_name__							:	List[str]    							=							torch.Generator(device=lowerCamelCase_  ).manual_seed(0  )
								__magic_name__							:	List[str]    							=							pipe([prompt]       ,						generator=lowerCamelCase_       ,						num_inference_steps=2       ,						output_type='''np'''  )
								__magic_name__							:	str    							=							output.images

								__magic_name__							:	Optional[int]    							=							torch.Generator(device=lowerCamelCase_  ).manual_seed(0  )
								__magic_name__							:	str    							=							pipe(
								    [prompt]       ,						generator=lowerCamelCase_       ,						output_type='''np'''       ,						return_dict=lowerCamelCase_       ,						num_inference_steps=2  )[0]

								__magic_name__							:	Optional[int]    							=							image[0, -3:, -3:, -1]
								__magic_name__							:	Tuple    							=							image_from_tuple[0, -3:, -3:, -1]

								assert image.shape == (1, 24, 24, 3)

								__magic_name__							:	List[str]    							=							np.array([0.6_6_9_3, 0.6_0_7_5, 0.4_9_5_9, 0.5_7_0_1, 0.5_5_8_3, 0.4_3_3_3, 0.6_1_7_1, 0.5_6_8_4, 0.4_9_8_8]  )

								assert np.abs(image_slice.flatten() - expected_slice  ).max() < 2.0
								assert np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 1E-2






@slow
@require_torch_gpu
class    lowerCamelCase    (		unittest.TestCase						):






			'''simple docstring'''





			def 							UpperCAmelCase__       ( self			:   Dict  )						->   Tuple:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()


			def 							UpperCAmelCase__       ( self			:   int  )						->   Optional[Any]:
								__magic_name__							:	int    							=							load_numpy(
								    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
								    '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''  )

								__magic_name__							:	Tuple    							=							VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''  )
								__magic_name__							:	str    							=							pipeline.to(lowerCamelCase_  )
								pipeline.set_progress_bar_config(disable=lowerCamelCase_  )

								# requires GPU generator for gumbel softmax
								# don't use GPU generator in tests though
								__magic_name__							:	int    							=							torch.Generator(device=lowerCamelCase_  ).manual_seed(0  )
								__magic_name__							:	List[Any]    							=							pipeline(
								    '''teddy bear playing in the pool'''       ,						num_images_per_prompt=1       ,						generator=lowerCamelCase_       ,						output_type='''np'''       ,						)

								__magic_name__							:	Optional[int]    							=							output.images[0]

								assert image.shape == (256, 256, 3)
								assert np.abs(expected_image - image  ).max() < 2.0









import multiprocessing
import os
from typing import BinaryIO, Optional, Union

import fsspec

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.json.json import Json
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader


class    lowerCamelCase    (		_lowerCamelCase						):






			'''simple docstring'''





			def __init__( self			:   List[Any]       ,						lowerCamelCase_			:   NestedDataStructureLike[PathLike]       ,						lowerCamelCase_			:   Optional[NamedSplit] = None       ,						lowerCamelCase_			:   Optional[Features] = None       ,						lowerCamelCase_			:   str = None       ,						lowerCamelCase_			:   bool = False       ,						lowerCamelCase_			:   bool = False       ,						lowerCamelCase_			:   Optional[str] = None       ,						lowerCamelCase_			:   Optional[int] = None       ,						**lowerCamelCase_			:   str       ,						)						->   Dict:
								super().__init__(
								    lowerCamelCase_       ,						split=lowerCamelCase_       ,						features=lowerCamelCase_       ,						cache_dir=lowerCamelCase_       ,						keep_in_memory=lowerCamelCase_       ,						streaming=lowerCamelCase_       ,						num_proc=lowerCamelCase_       ,						**lowerCamelCase_       ,						)
								__magic_name__							:	Optional[int]    							=							field
								__magic_name__							:	Optional[Any]    							=							path_or_paths if isinstance(lowerCamelCase_       ,						lowerCamelCase_  ) else {self.split: path_or_paths}
								__magic_name__							:	Optional[Any]    							=							Json(
								    cache_dir=lowerCamelCase_       ,						data_files=lowerCamelCase_       ,						features=lowerCamelCase_       ,						field=lowerCamelCase_       ,						**lowerCamelCase_       ,						)


			def 							UpperCAmelCase__       ( self			:   Optional[Any]  )						->   List[str]:
								# Build iterable dataset
								if self.streaming:
													__magic_name__							:	Optional[Any]    							=							self.builder.as_streaming_dataset(split=self.split  )
								# Build regular (map-style) dataset
								else:
													__magic_name__							:	int    							=							None
													__magic_name__							:	Tuple    							=							None
													__magic_name__							:	Tuple    							=							None
													__magic_name__							:	Optional[int]    							=							None

													self.builder.download_and_prepare(
													    download_config=lowerCamelCase_       ,						download_mode=lowerCamelCase_       ,						verification_mode=lowerCamelCase_       ,						base_path=lowerCamelCase_       ,						num_proc=self.num_proc       ,						)
													__magic_name__							:	Any    							=							self.builder.as_dataset(
													    split=self.split       ,						verification_mode=lowerCamelCase_       ,						in_memory=self.keep_in_memory  )
								return dataset






class    lowerCamelCase    :






			'''simple docstring'''





			def __init__( self			:   int       ,						lowerCamelCase_			:   Dataset       ,						lowerCamelCase_			:   Union[PathLike, BinaryIO]       ,						lowerCamelCase_			:   Optional[int] = None       ,						lowerCamelCase_			:   Optional[int] = None       ,						**lowerCamelCase_			:   List[Any]       ,						)						->   List[Any]:
								if num_proc is not None and num_proc <= 0:
													raise ValueError(F'''num_proc {num_proc} must be an integer > 0.'''  )

								__magic_name__							:	Optional[Any]    							=							dataset
								__magic_name__							:	Tuple    							=							path_or_buf
								__magic_name__							:	Dict    							=							batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
								__magic_name__							:	Any    							=							num_proc
								__magic_name__							:	Tuple    							=							'''utf-8'''
								__magic_name__							:	Tuple    							=							to_json_kwargs

			def 							UpperCAmelCase__       ( self			:   Dict  )						->   int:
								__magic_name__							:	int    							=							self.to_json_kwargs.pop('''path_or_buf'''       ,						lowerCamelCase_  )
								__magic_name__							:	Any    							=							self.to_json_kwargs.pop('''orient'''       ,						'''records'''  )
								__magic_name__							:	List[Any]    							=							self.to_json_kwargs.pop('''lines'''       ,						True if orient == '''records''' else False  )
								__magic_name__							:	Optional[int]    							=							self.to_json_kwargs.pop('''index'''       ,						False if orient in ['''split''', '''table'''] else True  )
								__magic_name__							:	List[str]    							=							self.to_json_kwargs.pop('''compression'''       ,						lowerCamelCase_  )

								if compression not in [None, "infer", "gzip", "bz2", "xz"]:
													raise NotImplementedError(F'''`datasets` currently does not support {compression} compression'''  )

								if isinstance(self.path_or_buf       ,						(str, bytes, os.PathLike)  ):
													with fsspec.open(self.path_or_buf       ,						'''wb'''       ,						compression=lowerCamelCase_  ) as buffer:
																		__magic_name__							:	str    							=							self._write(file_obj=lowerCamelCase_       ,						orient=lowerCamelCase_       ,						lines=lowerCamelCase_       ,						index=lowerCamelCase_       ,						**self.to_json_kwargs  )
								else:
													if compression:
																		raise NotImplementedError(
																		    F'''The compression parameter is not supported when writing to a buffer, but compression={compression}'''
																		    ''' was passed. Please provide a local path instead.'''  )
													__magic_name__							:	List[str]    							=							self._write(
													    file_obj=self.path_or_buf       ,						orient=lowerCamelCase_       ,						lines=lowerCamelCase_       ,						index=lowerCamelCase_       ,						**self.to_json_kwargs  )
								return written

			def 							UpperCAmelCase__       ( self			:   Optional[Any]       ,						lowerCamelCase_			:   Dict  )						->   Optional[int]:
								__magic_name__							,    __magic_name__							,    __magic_name__							,    __magic_name__							,    __magic_name__							:	Any    							=							args

								__magic_name__							:	List[Any]    							=							query_table(
								    table=self.dataset.data       ,						key=slice(lowerCamelCase_       ,						offset + self.batch_size  )       ,						indices=self.dataset._indices       ,						)
								__magic_name__							:	Optional[Any]    							=							batch.to_pandas().to_json(
								    path_or_buf=lowerCamelCase_       ,						orient=lowerCamelCase_       ,						lines=lowerCamelCase_       ,						index=lowerCamelCase_       ,						**lowerCamelCase_  )
								if not json_str.endswith('''\n'''  ):
													json_str += "\n"
								return json_str.encode(self.encoding  )


			def 							UpperCAmelCase__       ( self			:   int       ,						lowerCamelCase_			:   BinaryIO       ,						lowerCamelCase_			:   List[str]       ,						lowerCamelCase_			:   Union[str, Any]       ,						lowerCamelCase_			:   Any       ,						**lowerCamelCase_			:   Optional[Any]       ,						)						->   int:
								__magic_name__							:	Union[str, Any]    							=							0

								if self.num_proc is None or self.num_proc == 1:
													for offset in logging.tqdm(
													    range(0       ,						len(self.dataset  )       ,						self.batch_size  )       ,						unit='''ba'''       ,						disable=not logging.is_progress_bar_enabled()       ,						desc='''Creating json from Arrow format'''       ,						):
																		__magic_name__							:	List[str]    							=							self._batch_json((offset, orient, lines, index, to_json_kwargs)  )
																		written += file_obj.write(lowerCamelCase_  )
								else:
													__magic_name__							,    __magic_name__							:	Optional[Any]    							=							len(self.dataset  ), self.batch_size
													with multiprocessing.Pool(self.num_proc  ) as pool:
																		for json_str in logging.tqdm(
																		    pool.imap(
																		        self._batch_json       ,						[(offset, orient, lines, index, to_json_kwargs) for offset in range(0       ,						lowerCamelCase_       ,						lowerCamelCase_  )]       ,						)       ,						total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size       ,						unit='''ba'''       ,						disable=not logging.is_progress_bar_enabled()       ,						desc='''Creating json from Arrow format'''       ,						):
																							written += file_obj.write(lowerCamelCase_  )

								return written



