


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        __lowercase       (							__lowerCAmelCase			,							unittest.TestCase  ):



		"""simple docstring"""

		UpperCamelCase     :					List[Any]			       =		KandinskyImgaImgPipeline
		UpperCamelCase     :					int			       =		['''prompt''', '''image_embeds''', '''negative_image_embeds''', '''image''']
		UpperCamelCase     :					List[Any]			       =		[
		    '''prompt''',
		    '''negative_prompt''',
		    '''image_embeds''',
		    '''negative_image_embeds''',
		    '''image''',
		]
		UpperCamelCase     :					Optional[int]			       =		[
		    '''generator''',
		    '''height''',
		    '''width''',
		    '''strength''',
		    '''guidance_scale''',
		    '''negative_prompt''',
		    '''num_inference_steps''',
		    '''return_dict''',
		    '''guidance_scale''',
		    '''num_images_per_prompt''',
		    '''output_type''',
		    '''return_dict''',
		]
		UpperCamelCase     :					List[Any]			       =		False







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







					'''simple docstring'''







					return 32







		@property
		def 					__A   ( self	)  -> str:







					'''simple docstring'''







					return 32







		@property
		def 					__A   ( self	)  -> Any:







					'''simple docstring'''







					return self.time_input_dim







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







					'''simple docstring'''







					return self.time_input_dim * 4







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







					'''simple docstring'''







					return 1_00







		@property
		def 					__A   ( self	)  -> Tuple:







					'''simple docstring'''







					lowerCamelCase    		=  XLMRobertaTokenizerFast.from_pretrained("""YiYiXu/tiny-random-mclip-base"""	)
					return tokenizer







		@property
		def 					__A   ( self	)  -> List[str]:







					'''simple docstring'''







					torch.manual_seed(0	)
					lowerCamelCase    		=  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		,      )

					lowerCamelCase    		=  MultilingualCLIP(__SCREAMING_SNAKE_CASE	)
					lowerCamelCase    		=  text_encoder.eval()

					return text_encoder







		@property
		def 					__A   ( self	)  -> Tuple:







					'''simple docstring'''







					torch.manual_seed(0	)

					lowerCamelCase    		=  {
					    """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,
					}

					lowerCamelCase    		=  UNetaDConditionModel(**__SCREAMING_SNAKE_CASE	)
					return model







		@property
		def 					__A   ( self	)  -> Dict:







					'''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 					__A   ( self	)  -> Dict:







					'''simple docstring'''







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







		def 					__A   ( self	)  -> str:







					'''simple docstring'''







					lowerCamelCase    		=  self.dummy_text_encoder
					lowerCamelCase    		=  self.dummy_tokenizer
					lowerCamelCase    		=  self.dummy_unet
					lowerCamelCase    		=  self.dummy_movq

					lowerCamelCase    		=  {
					    """num_train_timesteps""": 10_00,
					    """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,
					}

					lowerCamelCase    		=  DDIMScheduler(**__SCREAMING_SNAKE_CASE	)

					lowerCamelCase    		=  {
					    """text_encoder""": text_encoder,
					    """tokenizer""": tokenizer,
					    """unet""": unet,
					    """scheduler""": scheduler,
					    """movq""": movq,
					}

					return components







		def 					__A   ( self		,      A		,      A=0	)  -> List[str]:







					'''simple docstring'''







					lowerCamelCase    		=  floats_tensor((1, self.cross_attention_dim)		,      rng=random.Random(__SCREAMING_SNAKE_CASE	)	).to(__SCREAMING_SNAKE_CASE	)
					lowerCamelCase    		=  floats_tensor((1, self.cross_attention_dim)		,      rng=random.Random(seed + 1	)	).to(__SCREAMING_SNAKE_CASE	)
					# create init_image
					lowerCamelCase    		=  floats_tensor((1, 3, 64, 64)		,      rng=random.Random(__SCREAMING_SNAKE_CASE	)	).to(__SCREAMING_SNAKE_CASE	)
					lowerCamelCase    		=  image.cpu().permute(0		,      2		,      3		,      1	)[0]
					lowerCamelCase    		=  Image.fromarray(np.uinta(__SCREAMING_SNAKE_CASE	)	).convert("""RGB"""	).resize((2_56, 2_56)	)

					if str(__SCREAMING_SNAKE_CASE	).startswith("""mps"""	):
								lowerCamelCase    		=  torch.manual_seed(__SCREAMING_SNAKE_CASE	)
					else:
								lowerCamelCase    		=  torch.Generator(device=__SCREAMING_SNAKE_CASE	).manual_seed(__SCREAMING_SNAKE_CASE	)
					lowerCamelCase    		=  {
					    """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 					__A   ( self	)  -> int:







					'''simple docstring'''







					lowerCamelCase    		=  """cpu"""

					lowerCamelCase    		=  self.get_dummy_components()

					lowerCamelCase    		=  self.pipeline_class(**__SCREAMING_SNAKE_CASE	)
					lowerCamelCase    		=  pipe.to(__SCREAMING_SNAKE_CASE	)

					pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE	)

					lowerCamelCase    		=  pipe(**self.get_dummy_inputs(__SCREAMING_SNAKE_CASE	)	)
					lowerCamelCase    		=  output.images

					lowerCamelCase    		=  pipe(
					    **self.get_dummy_inputs(__SCREAMING_SNAKE_CASE	)		,      return_dict=__SCREAMING_SNAKE_CASE		,      )[0]

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

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

					lowerCamelCase    		=  np.array(
					    [0.61474943, 0.6073539, 0.43308544, 0.5928269, 0.47493595, 0.46755973, 0.4613838, 0.45368797, 0.50119233]	)
					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        __lowercase       (							unittest.TestCase  ):



		"""simple docstring"""

		def 					__A   ( self	)  -> List[str]:







					'''simple docstring'''







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







		def 					__A   ( self	)  -> Any:







					'''simple docstring'''







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

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

					lowerCamelCase    		=  KandinskyPriorPipeline.from_pretrained(
					    """kandinsky-community/kandinsky-2-1-prior"""		,      torch_dtype=torch.floataa	)
					pipe_prior.to(__SCREAMING_SNAKE_CASE	)

					lowerCamelCase    		=  KandinskyImgaImgPipeline.from_pretrained(
					    """kandinsky-community/kandinsky-2-1"""		,      torch_dtype=torch.floataa	)
					lowerCamelCase    		=  pipeline.to(__SCREAMING_SNAKE_CASE	)

					pipeline.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE	)

					lowerCamelCase    		=  torch.Generator(device="""cpu"""	).manual_seed(0	)
					lowerCamelCase		,		lowerCamelCase    		=  pipe_prior(
					    __SCREAMING_SNAKE_CASE		,      generator=__SCREAMING_SNAKE_CASE		,      num_inference_steps=5		,      negative_prompt=""""""		,      ).to_tuple()

					lowerCamelCase    		=  pipeline(
					    __SCREAMING_SNAKE_CASE		,      image=__SCREAMING_SNAKE_CASE		,      image_embeds=__SCREAMING_SNAKE_CASE		,      negative_image_embeds=__SCREAMING_SNAKE_CASE		,      generator=__SCREAMING_SNAKE_CASE		,      num_inference_steps=1_00		,      height=7_68		,      width=7_68		,      strength=0.2		,      output_type="""np"""		,      )

					lowerCamelCase    		=  output.images[0]

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

					assert_mean_pixel_difference(__SCREAMING_SNAKE_CASE		,      __SCREAMING_SNAKE_CASE	)







'''simple docstring'''



from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)







def       _UpperCamelCase			(_lowerCamelCase				:  Union[str, Any] ,						_lowerCamelCase				:  List[Any]						)-> Optional[Any]:




     '''simple docstring'''





     __snake_case            =	[]
     for part_id in partition_order:
          __snake_case            =	df.where(f'''SPARK_PARTITION_ID() = {part_id}'''						).collect()
          for row_idx, row in enumerate(_lowerCamelCase						):
               expected_row_ids_and_row_dicts.append((f'''{part_id}_{row_idx}''', row.asDict())						)
     return expected_row_ids_and_row_dicts







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Any:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(1_00						).repartition(1						)
     __snake_case            =	Spark(_lowerCamelCase						)
     # The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
     # that each partition can hold 2 rows.
     spark_builder._repartition_df_if_needed(max_shard_size=16						)
     # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
     assert spark_builder.df.rdd.getNumPartitions() == 50







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Tuple:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(10						).repartition(2						)
     __snake_case            =	[1, 0]
     __snake_case            =	_generate_iterable_examples(_lowerCamelCase ,						_lowerCamelCase						)  # Reverse the partitions.
     __snake_case            =	_get_expected_row_ids_and_row_dicts_for_partition_order(_lowerCamelCase ,						_lowerCamelCase						)

     for i, (row_id, row_dict) in enumerate(generate_fn()						):
          __snake_case       ,     __snake_case            =	expected_row_ids_and_row_dicts[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> int:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(10						).repartition(1						)
     __snake_case            =	SparkExamplesIterable(_lowerCamelCase						)
     assert it.n_shards == 1
     for i, (row_id, row_dict) in enumerate(_lowerCamelCase						):
          assert row_id == f'''0_{i}'''
          assert row_dict == {"id": i}







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Union[str, Any]:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(30						).repartition(3						)
     # Mock the generator so that shuffle reverses the partition indices.
     with patch('''numpy.random.Generator'''						) as generator_mock:
          __snake_case            =	lambda _lowerCamelCase						: x.reverse()
          __snake_case            =	_get_expected_row_ids_and_row_dicts_for_partition_order(_lowerCamelCase ,						[2, 1, 0]						)

          __snake_case            =	SparkExamplesIterable(_lowerCamelCase						).shuffle_data_sources(_lowerCamelCase						)
          assert shuffled_it.n_shards == 3
          for i, (row_id, row_dict) in enumerate(_lowerCamelCase						):
               __snake_case       ,     __snake_case            =	expected_row_ids_and_row_dicts[i]
               assert row_id == expected_row_id
               assert row_dict == expected_row_dict







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Tuple:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(20						).repartition(4						)

     # Partitions 0 and 2
     __snake_case            =	SparkExamplesIterable(_lowerCamelCase						).shard_data_sources(worker_id=0 ,						num_workers=2						)
     assert shard_it_a.n_shards == 2
     __snake_case            =	_get_expected_row_ids_and_row_dicts_for_partition_order(_lowerCamelCase ,						[0, 2]						)
     for i, (row_id, row_dict) in enumerate(_lowerCamelCase						):
          __snake_case       ,     __snake_case            =	expected_row_ids_and_row_dicts_a[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict

     # Partitions 1 and 3
     __snake_case            =	SparkExamplesIterable(_lowerCamelCase						).shard_data_sources(worker_id=1 ,						num_workers=2						)
     assert shard_it_a.n_shards == 2
     __snake_case            =	_get_expected_row_ids_and_row_dicts_for_partition_order(_lowerCamelCase ,						[1, 3]						)
     for i, (row_id, row_dict) in enumerate(_lowerCamelCase						):
          __snake_case       ,     __snake_case            =	expected_row_ids_and_row_dicts_a[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Optional[int]:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(1_00						).repartition(1						)
     __snake_case            =	Spark(_lowerCamelCase						)
     # Choose a small max_shard_size for maximum partitioning.
     spark_builder._repartition_df_if_needed(max_shard_size=1						)
     # The new number of partitions should not be greater than the number of rows.
     assert spark_builder.df.rdd.getNumPartitions() == 1_00






