





"""simple docstring"""





import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


lowerCAmelCase__   	=	datasets.utils.logging.get_logger(__name__)





@dataclass
class 				_lowerCamelCase					(				datasets.BuilderConfig       ):
						UpperCAmelCase_       =	None
						UpperCAmelCase_       =	"utf-8"
						UpperCAmelCase_       =	None
						UpperCAmelCase_       =	None
						UpperCAmelCase_       =	True  # deprecated
						UpperCAmelCase_       =	None  # deprecated
						UpperCAmelCase_       =	10 << 20  # 10MB
						UpperCAmelCase_       =	None





class 				_lowerCamelCase					(				datasets.ArrowBasedBuilder       ):
						UpperCAmelCase_       =	JsonConfig


						def 			snake_case_      (self						)       ->      Union[str, Any]:







								if self.config.block_size is not None:
										logger.warning("The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead"						)
										UpperCamelCase      =		self.config.block_size
								if self.config.use_threads is not True:
										logger.warning(
										    "The JSON loader parameter `use_threads` is deprecated and doesn\'t have any effect anymore."						)
								if self.config.newlines_in_values is not None:
										raise ValueError("The JSON loader parameter `newlines_in_values` is no longer supported"						)
								return datasets.DatasetInfo(features=self.config.features						)


						def 			snake_case_      (self      ,							__a						)       ->      Optional[int]:







								if not self.config.data_files:
										raise ValueError(F"At least one data file must be specified, but got data_files={self.config.data_files}"						)
								UpperCamelCase      =		dl_manager.download_and_extract(self.config.data_files						)
								if isinstance(UpperCamelCase_      ,							(str, list, tuple)						):
										UpperCamelCase      =		data_files
										if isinstance(UpperCamelCase_      ,							UpperCamelCase_						):
												UpperCamelCase      =		[files]
										UpperCamelCase      =		[dl_manager.iter_files(UpperCamelCase_						) for file in files]
										return [datasets.SplitGenerator(name=datasets.Split.TRAIN      ,							gen_kwargs={"files": files}						)]
								UpperCamelCase      =		[]
								for split_name, files in data_files.items():
										if isinstance(UpperCamelCase_      ,							UpperCamelCase_						):
												UpperCamelCase      =		[files]
										UpperCamelCase      =		[dl_manager.iter_files(UpperCamelCase_						) for file in files]
										splits.append(datasets.SplitGenerator(name=UpperCamelCase_      ,							gen_kwargs={"files": files}						)						)
								return splits


						def 			snake_case_      (self      ,							__a						)       ->      pa.Table:







								if self.config.features is not None:
										# adding missing columns
										for column_name in set(self.config.features						) - set(pa_table.column_names						):
												UpperCamelCase      =		self.config.features.arrow_schema.field(UpperCamelCase_						).type
												UpperCamelCase      =		pa_table.append_column(UpperCamelCase_      ,							pa.array([None] * len(UpperCamelCase_						)      ,							type=UpperCamelCase_						)						)
										# more expensive cast to support nested structures with keys in a different order
										# allows str <-> int/float or str to Audio for example
										UpperCamelCase      =		table_cast(UpperCamelCase_      ,							self.config.features.arrow_schema						)
								return pa_table




						def 			snake_case_      (self      ,							__a						)       ->      str:







								for file_idx, file in enumerate(itertools.chain.from_iterable(UpperCamelCase_						)						):
										# If the file is one json object and if we need to look at the list of items in one specific field
										if self.config.field is not None:
												with open(UpperCamelCase_      ,							encoding=self.config.encoding      ,							errors=self.config.encoding_errors						) as f:
														UpperCamelCase      =		json.load(UpperCamelCase_						)

												# We keep only the field we are interested in
												UpperCamelCase      =		dataset[self.config.field]

												# We accept two format: a list of dicts or a dict of lists
												if isinstance(UpperCamelCase_      ,							(list, tuple)						):
														UpperCamelCase      =		set().union(*[row.keys() for row in dataset]						)
														UpperCamelCase      =		{col: [row.get(UpperCamelCase_						) for row in dataset] for col in keys}
												else:
														UpperCamelCase      =		dataset
												UpperCamelCase      =		pa.Table.from_pydict(UpperCamelCase_						)
												yield file_idx, self._cast_table(UpperCamelCase_						)

										# If the file has one json object per line
										else:
												with open(UpperCamelCase_      ,							"rb"						) as f:
														UpperCamelCase      =		0
														# Use block_size equal to the chunk size divided by 32 to leverage multithreading
														# Set a default minimum value of 16kB if the chunk size is really small
														UpperCamelCase      =		max(self.config.chunksize // 32      ,							16 << 10						)
														UpperCamelCase      =		(
														    self.config.encoding_errors if self.config.encoding_errors is not None else "strict"
														)
														while True:
																UpperCamelCase      =		f.read(self.config.chunksize						)
																if not batch:
																		break
																# Finish current line
																try:
																		batch += f.readline()
																except (AttributeError, io.UnsupportedOperation):
																		batch += readline(UpperCamelCase_						)
																# PyArrow only accepts utf-8 encoded bytes
																if self.config.encoding != "utf-8":
																		UpperCamelCase      =		batch.decode(self.config.encoding      ,							errors=UpperCamelCase_						).encode("utf-8"						)
																try:
																		while True:
																				try:
																						UpperCamelCase      =		paj.read_json(
																						    io.BytesIO(UpperCamelCase_						)      ,							read_options=paj.ReadOptions(block_size=UpperCamelCase_						)						)
																						break
																				except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
																						if (
																						    isinstance(UpperCamelCase_      ,							pa.ArrowInvalid						)
																						    and "straddling" not in str(UpperCamelCase_						)
																						    or block_size > len(UpperCamelCase_						)
																						):
																								raise
																						else:
																								# Increase the block size in case it was too small.
																								# The block size will be reset for the next file.
																								logger.debug(
																								    F"Batch of {len(UpperCamelCase_						)} bytes couldn\'t be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}."						)
																								block_size *= 2
																except pa.ArrowInvalid as e:
																		try:
																				with open(
																				    UpperCamelCase_      ,							encoding=self.config.encoding      ,							errors=self.config.encoding_errors						) as f:
																						UpperCamelCase      =		json.load(UpperCamelCase_						)
																		except json.JSONDecodeError:
																				logger.error(F"Failed to read file \'{file}\' with error {type(UpperCamelCase_						)}: {e}"						)
																				raise e
																		# If possible, parse the file as a list of json objects and exit the loop
																		if isinstance(UpperCamelCase_      ,							UpperCamelCase_						):  # list is the only sequence type supported in JSON
																				try:
																						UpperCamelCase      =		set().union(*[row.keys() for row in dataset]						)
																						UpperCamelCase      =		{col: [row.get(UpperCamelCase_						) for row in dataset] for col in keys}
																						UpperCamelCase      =		pa.Table.from_pydict(UpperCamelCase_						)
																				except (pa.ArrowInvalid, AttributeError) as e:
																						logger.error(F"Failed to read file \'{file}\' with error {type(UpperCamelCase_						)}: {e}"						)
																						raise ValueError(F"Not able to read records in the JSON file at {file}."						) from None
																				yield file_idx, self._cast_table(UpperCamelCase_						)
																				break
																		else:
																				logger.error(F"Failed to read file \'{file}\' with error {type(UpperCamelCase_						)}: {e}"						)
																				raise ValueError(
																				    F"Not able to read records in the JSON file at {file}. "
																				    F"You should probably indicate the field of the JSON file containing your records. "
																				    F"This JSON file contain the following fields: {str(list(dataset.keys()						)						)}. "
																				    F"Select the correct one and provide it as `field=\'XXX\'` to the dataset loading method. "						) from None
                        # Uncomment for debugging (will print the Arrow table size and elements)
                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
																yield (file_idx, batch_idx), self._cast_table(UpperCamelCase_						)
																batch_idx += 1










import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()


class        _a							(			UpperCamelCase__     ,	unittest.TestCase			):
							_lowercase	:			Union[str, Any]        = KandinskyVaaImgaImgPipeline
							_lowercase	:			Tuple        = ['''image_embeds''', '''negative_image_embeds''', '''image''']
							_lowercase	:			Any        = [
							    '''image_embeds''',
							    '''negative_image_embeds''',
							    '''image''',
							]
							_lowercase	:			Union[str, Any]        = [
							    '''generator''',
							    '''height''',
							    '''width''',
							    '''strength''',
							    '''guidance_scale''',
							    '''num_inference_steps''',
							    '''return_dict''',
							    '''guidance_scale''',
							    '''num_images_per_prompt''',
							    '''output_type''',
							    '''return_dict''',
							]
							_lowercase	:			Optional[Any]        = False

							@property
							def       lowerCamelCase_ (							self:	Union[str, Any]							)			->   Dict:





									"""simple docstring"""







									return 32

							@property
							def       lowerCamelCase_ (							self:	Optional[int]							)			->   Optional[Any]:





									"""simple docstring"""







									return 32

							@property
							def       lowerCamelCase_ (							self:	Any							)			->   Any:





									"""simple docstring"""







									return self.time_input_dim

							@property
							def       lowerCamelCase_ (							self:	Tuple							)			->   Any:





									"""simple docstring"""







									return self.time_input_dim * 4

							@property
							def       lowerCamelCase_ (							self:	List[Any]							)			->   Optional[Any]:





									"""simple docstring"""







									return 100

							@property
							def       lowerCamelCase_ (							self:	int							)			->   int:





									"""simple docstring"""







									torch.manual_seed(0							)

									lowercase__            =      {
									    '''in_channels''': 4,
									    # Out channels is double in channels because predicts mean and variance
									    '''out_channels''': 8,
									    '''addition_embed_type''': '''image''',
									    '''down_block_types''': ('''ResnetDownsampleBlock2D''', '''SimpleCrossAttnDownBlock2D'''),
									    '''up_block_types''': ('''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''),
									    '''mid_block_type''': '''UNetMidBlock2DSimpleCrossAttn''',
									    '''block_out_channels''': (self.block_out_channels_a, self.block_out_channels_a * 2),
									    '''layers_per_block''': 1,
									    '''encoder_hid_dim''': self.text_embedder_hidden_size,
									    '''encoder_hid_dim_type''': '''image_proj''',
									    '''cross_attention_dim''': self.cross_attention_dim,
									    '''attention_head_dim''': 4,
									    '''resnet_time_scale_shift''': '''scale_shift''',
									    '''class_embed_type''': None,
									}

									lowercase__            =      UNetaDConditionModel(**UpperCamelCase_							)
									return model

							@property
							def       lowerCamelCase_ (							self:	Optional[int]							)			->   Union[str, Any]:





									"""simple docstring"""







									return {
									    "block_out_channels": [32, 64],
									    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
									    "in_channels": 3,
									    "latent_channels": 4,
									    "layers_per_block": 1,
									    "norm_num_groups": 8,
									    "norm_type": "spatial",
									    "num_vq_embeddings": 12,
									    "out_channels": 3,
									    "up_block_types": [
									        "AttnUpDecoderBlock2D",
									        "UpDecoderBlock2D",
									    ],
									    "vq_embed_dim": 4,
									}

							@property
							def       lowerCamelCase_ (							self:	Optional[Any]							)			->   int:





									"""simple docstring"""







									torch.manual_seed(0							)
									lowercase__            =      VQModel(**self.dummy_movq_kwargs							)
									return model

							def       lowerCamelCase_ (							self:	Optional[int]							)			->   Optional[int]:





									"""simple docstring"""







									lowercase__            =      self.dummy_unet
									lowercase__            =      self.dummy_movq

									lowercase__            =      {
									    '''num_train_timesteps''': 1_000,
									    '''beta_schedule''': '''linear''',
									    '''beta_start''': 0.00085,
									    '''beta_end''': 0.012,
									    '''clip_sample''': False,
									    '''set_alpha_to_one''': False,
									    '''steps_offset''': 0,
									    '''prediction_type''': '''epsilon''',
									    '''thresholding''': False,
									}

									lowercase__            =      DDIMScheduler(**UpperCamelCase_							)

									lowercase__            =      {
									    '''unet''': unet,
									    '''scheduler''': scheduler,
									    '''movq''': movq,
									}

									return components

							def       lowerCamelCase_ (							self:	Optional[int]    ,						UpperCamelCase_:	Optional[Any]    ,						UpperCamelCase_:	Optional[int]=0							)			->   Optional[int]:





									"""simple docstring"""







									lowercase__            =      floats_tensor((1, self.text_embedder_hidden_size)    ,						rng=random.Random(UpperCamelCase_							)							).to(UpperCamelCase_							)
									lowercase__            =      floats_tensor((1, self.text_embedder_hidden_size)    ,						rng=random.Random(seed + 1							)							).to(
									    UpperCamelCase_							)
									# create init_image
									lowercase__            =      floats_tensor((1, 3, 64, 64)    ,						rng=random.Random(UpperCamelCase_							)							).to(UpperCamelCase_							)
									lowercase__            =      image.cpu().permute(0    ,						2    ,						3    ,						1							)[0]
									lowercase__            =      Image.fromarray(np.uinta(UpperCamelCase_							)							).convert('''RGB'''							).resize((256, 256)							)

									if str(UpperCamelCase_							).startswith('''mps'''							):
											lowercase__            =      torch.manual_seed(UpperCamelCase_							)
									else:
											lowercase__            =      torch.Generator(device=UpperCamelCase_							).manual_seed(UpperCamelCase_							)
									lowercase__            =      {
									    '''image''': init_image,
									    '''image_embeds''': image_embeds,
									    '''negative_image_embeds''': negative_image_embeds,
									    '''generator''': generator,
									    '''height''': 64,
									    '''width''': 64,
									    '''num_inference_steps''': 10,
									    '''guidance_scale''': 7.0,
									    '''strength''': 0.2,
									    '''output_type''': '''np''',
									}
									return inputs


							def       lowerCamelCase_ (							self:	Optional[int]							)			->   Dict:





									"""simple docstring"""







									lowercase__            =      '''cpu'''

									lowercase__            =      self.get_dummy_components()

									lowercase__            =      self.pipeline_class(**UpperCamelCase_							)
									lowercase__            =      pipe.to(UpperCamelCase_							)

									pipe.set_progress_bar_config(disable=UpperCamelCase_							)

									lowercase__            =      pipe(**self.get_dummy_inputs(UpperCamelCase_							)							)
									lowercase__            =      output.images

									lowercase__            =      pipe(
									    **self.get_dummy_inputs(UpperCamelCase_							)    ,						return_dict=UpperCamelCase_    ,						)[0]

									lowercase__            =      image[0, -3:, -3:, -1]
									lowercase__            =      image_from_tuple[0, -3:, -3:, -1]

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

									lowercase__            =      np.array(
									    [0.6199778, 0.63984406, 0.46145785, 0.62944984, 0.5622215, 0.47306132, 0.47441456, 0.4607606, 0.48719263]							)
									assert (
									    np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2
									), f' expected_slice {expected_slice}, but got {image_slice.flatten()}'
									assert (
									    np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() < 1E-2
									), f' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'


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

							def       lowerCamelCase_ (							self:	str							)			->   List[Any]:





									"""simple docstring"""







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


							def       lowerCamelCase_ (							self:	List[str]							)			->   Union[str, Any]:





									"""simple docstring"""







									lowercase__            =      load_numpy(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
									    '''/kandinskyv22/kandinskyv22_img2img_frog.npy'''							)

									lowercase__            =      load_image(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''							)
									lowercase__            =      '''A red cartoon frog, 4k'''

									lowercase__            =      KandinskyVaaPriorPipeline.from_pretrained(
									    '''kandinsky-community/kandinsky-2-2-prior'''    ,						torch_dtype=torch.floataa							)
									pipe_prior.to(UpperCamelCase_							)

									lowercase__            =      KandinskyVaaImgaImgPipeline.from_pretrained(
									    '''kandinsky-community/kandinsky-2-2-decoder'''    ,						torch_dtype=torch.floataa							)
									lowercase__            =      pipeline.to(UpperCamelCase_							)

									pipeline.set_progress_bar_config(disable=UpperCamelCase_							)

									lowercase__            =      torch.Generator(device='''cpu'''							).manual_seed(0							)
									lowercase__     ,		lowercase__            =      pipe_prior(
									    UpperCamelCase_    ,						generator=UpperCamelCase_    ,						num_inference_steps=5    ,						negative_prompt=''''''    ,						).to_tuple()

									lowercase__            =      pipeline(
									    image=UpperCamelCase_    ,						image_embeds=UpperCamelCase_    ,						negative_image_embeds=UpperCamelCase_    ,						generator=UpperCamelCase_    ,						num_inference_steps=100    ,						height=768    ,						width=768    ,						strength=0.2    ,						output_type='''np'''    ,						)

									lowercase__            =      output.images[0]

									assert image.shape == (768, 768, 3)

									assert_mean_pixel_difference(UpperCamelCase_    ,						UpperCamelCase_							)


