

import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetImgaImgPipeline,
    KandinskyVaaPriorEmbaEmbPipeline,
    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       snake_case_  (						__A      ,			unittest.TestCase   ):







			'''simple docstring'''

			SCREAMING_SNAKE_CASE	:  List[str]		     =							KandinskyVaaControlnetImgaImgPipeline
			SCREAMING_SNAKE_CASE	:  int		     =							["image_embeds", "negative_image_embeds", "image", "hint"]
			SCREAMING_SNAKE_CASE	:  List[str]		     =							["image_embeds", "negative_image_embeds", "image", "hint"]
			SCREAMING_SNAKE_CASE	:  Any		     =							[
			    "generator",
			    "height",
			    "width",
			    "strength",
			    "guidance_scale",
			    "num_inference_steps",
			    "return_dict",
			    "guidance_scale",
			    "num_images_per_prompt",
			    "output_type",
			    "return_dict",
			]
			SCREAMING_SNAKE_CASE	:  Tuple		     =							False







			@property
			def 							snake_case__(	self					:							Optional[Any]				) ->Tuple:
									return 3_2







			@property
			def 							snake_case__(	self					:							Dict				) ->Tuple:
									return 3_2







			@property
			def 							snake_case__(	self					:							List[str]				) ->Dict:
									return self.time_input_dim







			@property
			def 							snake_case__(	self					:							Optional[Any]				) ->List[str]:
									return self.time_input_dim * 4







			@property
			def 							snake_case__(	self					:							Any				) ->Optional[Any]:
									return 1_0_0







			@property
			def 							snake_case__(	self					:							Union[str, Any]				) ->str:
									torch.manual_seed(0				)

									snake_case_  =      {
									    '''in_channels''': 8,
									    # Out channels is double in channels because predicts mean and variance
									    '''out_channels''': 8,
									    '''addition_embed_type''': '''image_hint''',
									    '''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,
									}

									snake_case_  =      UNetaDConditionModel(**_UpperCamelCase				)
									return model







			@property
			def 							snake_case__(	self					:							Any				) ->Optional[int]:
									return {
									    "block_out_channels": [3_2, 3_2, 6_4, 6_4],
									    "down_block_types": [
									        "DownEncoderBlock2D",
									        "DownEncoderBlock2D",
									        "DownEncoderBlock2D",
									        "AttnDownEncoderBlock2D",
									    ],
									    "in_channels": 3,
									    "latent_channels": 4,
									    "layers_per_block": 1,
									    "norm_num_groups": 8,
									    "norm_type": "spatial",
									    "num_vq_embeddings": 1_2,
									    "out_channels": 3,
									    "up_block_types": ["AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"],
									    "vq_embed_dim": 4,
									}







			@property
			def 							snake_case__(	self					:							int				) ->int:
									torch.manual_seed(0				)
									snake_case_  =      VQModel(**self.dummy_movq_kwargs				)
									return model







			def 							snake_case__(	self					:							Optional[Any]				) ->Union[str, Any]:
									snake_case_  =      self.dummy_unet
									snake_case_  =      self.dummy_movq

									snake_case_  =      {
									    '''num_train_timesteps''': 1_0_0_0,
									    '''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,
									}

									snake_case_  =      DDIMScheduler(**_UpperCamelCase				)

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

									return components







			def 							snake_case__(	self					:							Dict				,  _UpperCamelCase					:							Union[str, Any]				,  _UpperCamelCase					:							Union[str, Any]=0				) ->List[str]:
									snake_case_  =      floats_tensor((1, self.text_embedder_hidden_size)				,  rng=random.Random(_UpperCamelCase				)				).to(_UpperCamelCase				)
									snake_case_  =      floats_tensor((1, self.text_embedder_hidden_size)				,  rng=random.Random(seed + 1				)				).to(
									    _UpperCamelCase				)
									# create init_image
									snake_case_  =      floats_tensor((1, 3, 6_4, 6_4)				,  rng=random.Random(_UpperCamelCase				)				).to(_UpperCamelCase				)
									snake_case_  =      image.cpu().permute(0				,  2				,  3				,  1				)[0]
									snake_case_  =      Image.fromarray(np.uinta(_UpperCamelCase				)				).convert('''RGB'''				).resize((2_5_6, 2_5_6)				)
									# create hint
									snake_case_  =      floats_tensor((1, 3, 6_4, 6_4)				,  rng=random.Random(_UpperCamelCase				)				).to(_UpperCamelCase				)

									if str(_UpperCamelCase				).startswith('''mps'''				):
															snake_case_  =      torch.manual_seed(_UpperCamelCase				)
									else:
															snake_case_  =      torch.Generator(device=_UpperCamelCase				).manual_seed(_UpperCamelCase				)
									snake_case_  =      {
									    '''image''': init_image,
									    '''image_embeds''': image_embeds,
									    '''negative_image_embeds''': negative_image_embeds,
									    '''hint''': hint,
									    '''generator''': generator,
									    '''height''': 6_4,
									    '''width''': 6_4,
									    '''num_inference_steps''': 1_0,
									    '''guidance_scale''': 7.0,
									    '''strength''': 0.2,
									    '''output_type''': '''np''',
									}
									return inputs







			def 							snake_case__(	self					:							Tuple				) ->Optional[Any]:
									snake_case_  =      '''cpu'''

									snake_case_  =      self.get_dummy_components()

									snake_case_  =      self.pipeline_class(**_UpperCamelCase				)
									snake_case_  =      pipe.to(_UpperCamelCase				)

									pipe.set_progress_bar_config(disable=_UpperCamelCase				)

									snake_case_  =      pipe(**self.get_dummy_inputs(_UpperCamelCase				)				)
									snake_case_  =      output.images

									snake_case_  =      pipe(
									    **self.get_dummy_inputs(_UpperCamelCase				)				,  return_dict=_UpperCamelCase				,  )[0]

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

									assert image.shape == (1, 6_4, 6_4, 3)

									snake_case_  =      np.array(
									    [0.54985034, 0.55509365, 0.52561504, 0.5570494, 0.5593818, 0.5263979, 0.50285643, 0.5069846, 0.51196736]				)
									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       snake_case_  (						unittest.TestCase   ):







			'''simple docstring'''

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







			def 							snake_case__(	self					:							Tuple				) ->Union[str, Any]:
									snake_case_  =      load_numpy(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
									    '''/kandinskyv22/kandinskyv22_controlnet_img2img_robotcat_fp16.npy'''				)

									snake_case_  =      load_image(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''				)
									snake_case_  =      init_image.resize((5_1_2, 5_1_2)				)

									snake_case_  =      load_image(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
									    '''/kandinskyv22/hint_image_cat.png'''				)
									snake_case_  =      torch.from_numpy(np.array(_UpperCamelCase				)				).float() / 255.0
									snake_case_  =      hint.permute(2				,  0				,  1				).unsqueeze(0				)

									snake_case_  =      '''A robot, 4k photo'''

									snake_case_  =      KandinskyVaaPriorEmbaEmbPipeline.from_pretrained(
									    '''kandinsky-community/kandinsky-2-2-prior'''				,  torch_dtype=torch.floataa				)
									pipe_prior.to(_UpperCamelCase				)

									snake_case_  =      KandinskyVaaControlnetImgaImgPipeline.from_pretrained(
									    '''kandinsky-community/kandinsky-2-2-controlnet-depth'''				,  torch_dtype=torch.floataa				)
									snake_case_  =      pipeline.to(_UpperCamelCase				)

									pipeline.set_progress_bar_config(disable=_UpperCamelCase				)

									snake_case_  =      torch.Generator(device='''cpu'''				).manual_seed(0				)

									snake_case_,       snake_case_  =      pipe_prior(
									    _UpperCamelCase				,  image=_UpperCamelCase				,  strength=0.85				,  generator=_UpperCamelCase				,  negative_prompt=''''''				,  ).to_tuple()

									snake_case_  =      pipeline(
									    image=_UpperCamelCase				,  image_embeds=_UpperCamelCase				,  negative_image_embeds=_UpperCamelCase				,  hint=_UpperCamelCase				,  generator=_UpperCamelCase				,  num_inference_steps=1_0_0				,  height=5_1_2				,  width=5_1_2				,  strength=0.5				,  output_type='''np'''				,  )

									snake_case_  =      output.images[0]

									assert image.shape == (5_1_2, 5_1_2, 3)

									assert_mean_pixel_difference(_UpperCamelCase				,  _UpperCamelCase				)

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


UpperCamelCase_						:     Union[str, Any]	       =							datasets.utils.logging.get_logger(__name__)




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




class   __lowercase						(						datasets.ArrowBasedBuilder       ):
	_A       =   JsonConfig






	def       _a(self				: List[Any]    ) ->			Dict:
							if self.config.block_size is not None:
													logger.warning("The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead"    )
													_lowercase       :	Union[str, Any]           =   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       _a(self				: Any      ,	snake_case				: Union[str, Any]    ) ->			List[str]:
							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}"""    )
							_lowercase       :	str           =   dl_manager.download_and_extract(self.config.data_files    )
							if isinstance(snake_case      ,	(str, list, tuple)    ):
													_lowercase       :	int           =   data_files
													if isinstance(snake_case      ,	snake_case    ):
																			_lowercase       :	Union[str, Any]           =   [files]
													_lowercase       :	Union[str, Any]           =   [dl_manager.iter_files(snake_case    ) for file in files]
													return [datasets.SplitGenerator(name=datasets.Split.TRAIN      ,	gen_kwargs={"files": files}    )]
							_lowercase       :	Union[str, Any]           =   []
							for split_name, files in data_files.items():
													if isinstance(snake_case      ,	snake_case    ):
																			_lowercase       :	List[Any]           =   [files]
													_lowercase       :	Optional[Any]           =   [dl_manager.iter_files(snake_case    ) for file in files]
													splits.append(datasets.SplitGenerator(name=snake_case      ,	gen_kwargs={"files": files}    )    )
							return splits






	def       _a(self				: List[Any]      ,	snake_case				: pa.Table    ) ->			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    ):
																			_lowercase       :	Any           =   self.config.features.arrow_schema.field(snake_case    ).type
																			_lowercase       :	Optional[Any]           =   pa_table.append_column(snake_case      ,	pa.array([None] * len(snake_case    )      ,	type=snake_case    )    )
													# more expensive cast to support nested structures with keys in a different order
													# allows str <-> int/float or str to Audio for example
													_lowercase       :	str           =   table_cast(snake_case      ,	self.config.features.arrow_schema    )
							return pa_table






	def       _a(self				: str      ,	snake_case				: Tuple    ) ->			Tuple:
							for file_idx, file in enumerate(itertools.chain.from_iterable(snake_case    )    ):
													# 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(snake_case      ,	encoding=self.config.encoding      ,	errors=self.config.encoding_errors    ) as f:
																									_lowercase       :	Any           =   json.load(snake_case    )

																			# We keep only the field we are interested in
																			_lowercase       :	List[Any]           =   dataset[self.config.field]

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

													# If the file has one json object per line
													else:
																			with open(snake_case      ,	"rb"    ) as f:
																									_lowercase       :	Union[str, Any]           =   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
																									_lowercase       :	Dict           =   max(self.config.chunksize // 32      ,	16 << 10    )
																									_lowercase       :	int           =   (
																									    self.config.encoding_errors if self.config.encoding_errors is not None else "strict"
																									)
																									while True:
																															_lowercase       :	Tuple           =   f.read(self.config.chunksize    )
																															if not batch:
																																					break
																															# Finish current line
																															try:
																																					batch += f.readline()
																															except (AttributeError, io.UnsupportedOperation):
																																					batch += readline(snake_case    )
																															# PyArrow only accepts utf-8 encoded bytes
																															if self.config.encoding != "utf-8":
																																					_lowercase       :	Union[str, Any]           =   batch.decode(self.config.encoding      ,	errors=snake_case    ).encode("utf-8"    )
																															try:
																																					while True:
																																											try:
																																																	_lowercase       :	Optional[Any]           =   paj.read_json(
																																																	    io.BytesIO(snake_case    )      ,	read_options=paj.ReadOptions(block_size=snake_case    )    )
																																																	break
																																											except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
																																																	if (
																																																	    isinstance(snake_case      ,	pa.ArrowInvalid    )
																																																	    and "straddling" not in str(snake_case    )
																																																	    or block_size > len(snake_case    )
																																																	):
																																																							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(snake_case    )} 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(
																																											    snake_case      ,	encoding=self.config.encoding      ,	errors=self.config.encoding_errors    ) as f:
																																																	_lowercase       :	Any           =   json.load(snake_case    )
																																					except json.JSONDecodeError:
																																											logger.error(F"""Failed to read file '{file}' with error {type(snake_case    )}: {e}"""    )
																																											raise e
																																					# If possible, parse the file as a list of json objects and exit the loop
																																					if isinstance(snake_case      ,	snake_case    ):  # list is the only sequence type supported in JSON
																																											try:
																																																	_lowercase       :	int           =   set().union(*[row.keys() for row in dataset]    )
																																																	_lowercase       :	Dict           =   {col: [row.get(snake_case    ) for row in dataset] for col in keys}
																																																	_lowercase       :	str           =   pa.Table.from_pydict(snake_case    )
																																											except (pa.ArrowInvalid, AttributeError) as e:
																																																	logger.error(F"""Failed to read file '{file}' with error {type(snake_case    )}: {e}"""    )
																																																	raise ValueError(F"""Not able to read records in the JSON file at {file}."""    ) from None
																																											yield file_idx, self._cast_table(snake_case    )
																																											break
																																					else:
																																											logger.error(F"""Failed to read file '{file}' with error {type(snake_case    )}: {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(snake_case    )
																															batch_idx += 1




