


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__							:						int							     =   datasets.utils.logging.get_logger(__name__)
@dataclass
class 		__magic_name__						(datasets.BuilderConfig				):





     '''simple docstring'''



     __lowercase				:							Optional[datasets.Features]	   =       None
     __lowercase				:							str	   =       "utf-8"
     __lowercase				:							Optional[str]	   =       None
     __lowercase				:							Optional[str]	   =       None
     __lowercase				:							bool	   =       True  # deprecated
     __lowercase				:							Optional[int]	   =       None  # deprecated
     __lowercase				:							int	   =       10 << 20  # 10MB
     __lowercase				:							Optional[bool]	   =       None




class 		__magic_name__						(datasets.ArrowBasedBuilder				):





     '''simple docstring'''



     __lowercase				:							int	   =       JsonConfig

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

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







     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[Any]     ,		_a:List[Any]   ):
       for file_idx, file in enumerate(itertools.chain.from_iterable(_a   )   ):
         # 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(_a     ,		encoding=self.config.encoding     ,		errors=self.config.encoding_errors   ) as f:
             snake_case__						= json.load(_a   )

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

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

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




import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyImgaImgPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
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 SCREAMING_SNAKE_CASE	(		__snake_case				,    unittest.TestCase				):






			"""simple docstring"""




			__A    	=	KandinskyImgaImgPipeline
			__A    	=	["""prompt""", """image_embeds""", """negative_image_embeds""", """image"""]
			__A    	=	[
			    """prompt""",
			    """negative_prompt""",
			    """image_embeds""",
			    """negative_image_embeds""",
			    """image""",
			]
			__A    	=	[
			    """generator""",
			    """height""",
			    """width""",
			    """strength""",
			    """guidance_scale""",
			    """negative_prompt""",
			    """num_inference_steps""",
			    """return_dict""",
			    """guidance_scale""",
			    """num_images_per_prompt""",
			    """output_type""",
			    """return_dict""",
			]
			__A    	=	False


			@property
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				return 32


			@property
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				return 32


			@property
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				return self.time_input_dim


			@property
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				return self.time_input_dim * 4


			@property
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				return 1_00


			@property
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base'				)
				return tokenizer


			@property
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				torch.manual_seed(0				)
				snake_case_   		=			MCLIPConfig(
				    numDims=self.cross_attention_dim			,					transformerDimensions=self.text_embedder_hidden_size			,					hidden_size=self.text_embedder_hidden_size			,					intermediate_size=37			,					num_attention_heads=4			,					num_hidden_layers=5			,					vocab_size=10_05			,					)

				snake_case_   		=			MultilingualCLIP(__UpperCamelCase				)
				snake_case_   		=			text_encoder.eval()

				return text_encoder


			@property
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				torch.manual_seed(0				)

				snake_case_   		=			{
				    'in_channels': 4,
				    # Out channels is double in channels because predicts mean and variance
				    'out_channels': 8,
				    'addition_embed_type': 'text_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': 'text_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       __lowerCAmelCase			(     self				):

				"""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				):

				"""simple docstring"""


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


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self.dummy_text_encoder
				snake_case_   		=			self.dummy_tokenizer
				snake_case_   		=			self.dummy_unet
				snake_case_   		=			self.dummy_movq

				snake_case_   		=			{
				    'num_train_timesteps': 10_00,
				    'beta_schedule': 'linear',
				    'beta_start': 0.0_0085,
				    '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_   		=			{
				    'text_encoder': text_encoder,
				    'tokenizer': tokenizer,
				    'unet': unet,
				    'scheduler': scheduler,
				    'movq': movq,
				}

				return components


			def       __lowerCAmelCase			(     self			,					__UpperCamelCase			,					__UpperCamelCase=0				):

				"""simple docstring"""


				snake_case_   		=			floats_tensor((1, self.cross_attention_dim)			,					rng=random.Random(__UpperCamelCase				)				).to(__UpperCamelCase				)
				snake_case_   		=			floats_tensor((1, self.cross_attention_dim)			,					rng=random.Random(seed + 1				)				).to(__UpperCamelCase				)
				# create init_image
				snake_case_   		=			floats_tensor((1, 3, 64, 64)			,					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_56, 2_56)				)

				if str(__UpperCamelCase				).startswith('mps'				):
					snake_case_   		=			torch.manual_seed(__UpperCamelCase				)
				else:
					snake_case_   		=			torch.Generator(device=__UpperCamelCase				).manual_seed(__UpperCamelCase				)
				snake_case_   		=			{
				    'prompt': 'horse',
				    '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				):

				"""simple docstring"""


				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, 64, 64, 3)

				snake_case_   		=			np.array(
				    [0.6147_4943, 0.607_3539, 0.4330_8544, 0.592_8269, 0.4749_3595, 0.4675_5973, 0.461_3838, 0.4536_8797, 0.5011_9233]				)
				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 SCREAMING_SNAKE_CASE	(		unittest.TestCase				):






			"""simple docstring"""




			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


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







			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			load_numpy(
				    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
				    '/kandinsky/kandinsky_img2img_frog.npy'				)

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

				snake_case_   		=			KandinskyPriorPipeline.from_pretrained(
				    'kandinsky-community/kandinsky-2-1-prior'			,					torch_dtype=torch.floataa				)
				pipe_prior.to(__UpperCamelCase				)

				snake_case_   		=			KandinskyImgaImgPipeline.from_pretrained(
				    'kandinsky-community/kandinsky-2-1'			,					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			,					generator=__UpperCamelCase			,					num_inference_steps=5			,					negative_prompt=''			,					).to_tuple()

				snake_case_   		=			pipeline(
				    __UpperCamelCase			,					image=__UpperCamelCase			,					image_embeds=__UpperCamelCase			,					negative_image_embeds=__UpperCamelCase			,					generator=__UpperCamelCase			,					num_inference_steps=1_00			,					height=7_68			,					width=7_68			,					strength=0.2			,					output_type='np'			,					)

				snake_case_   		=			output.images[0]

				assert image.shape == (7_68, 7_68, 3)

				assert_mean_pixel_difference(__UpperCamelCase			,					__UpperCamelCase				)
