


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    (      _a       ,		_a		)					->							Tuple:
     '''simple docstring'''





     lowercase_       :Optional[int]     =       []
     for part_id in partition_order:
          lowercase_       :Any     =       df.where(f"SPARK_PARTITION_ID() = {part_id}"		).collect()
          for row_idx, row in enumerate(_a		):
               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    (      )					->							Dict:
     '''simple docstring'''





     lowercase_       :Tuple     =       pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
     lowercase_       :Dict     =       spark.range(1_0_0		).repartition(1		)
     lowercase_       :Optional[int]     =       Spark(_a		)
     # 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=1_6		)
     # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
     assert spark_builder.df.rdd.getNumPartitions() == 5_0




@require_not_windows
@require_dill_gt_0_3_2
def     UpperCamelCase    (      )					->							Optional[Any]:
     '''simple docstring'''





     lowercase_       :List[str]     =       pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
     lowercase_       :Tuple     =       spark.range(1_0		).repartition(2		)
     lowercase_       :int     =       [1, 0]
     lowercase_       :Tuple     =       _generate_iterable_examples(_a       ,		_a		)  # Reverse the partitions.
     lowercase_       :Optional[Any]     =       _get_expected_row_ids_and_row_dicts_for_partition_order(_a       ,		_a		)

     for i, (row_id, row_dict) in enumerate(generate_fn()		):
          lowercase_    ,      lowercase_       :Dict     =       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'''





     lowercase_       :int     =       pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
     lowercase_       :Optional[Any]     =       spark.range(1_0		).repartition(1		)
     lowercase_       :List[Any]     =       SparkExamplesIterable(_a		)
     assert it.n_shards == 1
     for i, (row_id, row_dict) in enumerate(_a		):
          assert row_id == f"0_{i}"
          assert row_dict == {"id": i}




@require_not_windows
@require_dill_gt_0_3_2
def     UpperCamelCase    (      )					->							Any:
     '''simple docstring'''





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

          lowercase_       :Optional[Any]     =       SparkExamplesIterable(_a		).shuffle_data_sources(_a		)
          assert shuffled_it.n_shards == 3
          for i, (row_id, row_dict) in enumerate(_a		):
               lowercase_    ,      lowercase_       :Dict     =       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    (      )					->							List[str]:
     '''simple docstring'''





     lowercase_       :int     =       pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
     lowercase_       :List[str]     =       spark.range(2_0		).repartition(4		)

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

     # Partitions 1 and 3
     lowercase_       :Tuple     =       SparkExamplesIterable(_a		).shard_data_sources(worker_id=1       ,		num_workers=2		)
     assert shard_it_a.n_shards == 2
     lowercase_       :List[Any]     =       _get_expected_row_ids_and_row_dicts_for_partition_order(_a       ,		[1, 3]		)
     for i, (row_id, row_dict) in enumerate(_a		):
          lowercase_    ,      lowercase_       :str     =       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    (      )					->							List[Any]:
     '''simple docstring'''





     lowercase_       :int     =       pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
     lowercase_       :str     =       spark.range(1_0_0		).repartition(1		)
     lowercase_       :Dict     =       Spark(_a		)
     # 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_0_0





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


SCREAMING_SNAKE_CASE       :		int         =				False






class      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   def 		UpperCamelCase  (     self		):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()
   @property
   def 		UpperCamelCase  (     self		):
        return 12
   @property
   def 		UpperCamelCase  (     self		):
        return 12
   @property
   def 		UpperCamelCase  (     self		):
        return 32
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)
        lowercase_       :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		):
        lowercase_       :Optional[Any]     =       CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)
        return tokenizer
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)
        lowercase_       :int     =       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(UpperCamelCase_		)
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)

        lowercase_       :int     =       12
        lowercase_       :List[Any]     =       12

        lowercase_       :Dict     =       {
            '''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''',
        }

        lowercase_       :int     =       TransformeraDModel(**UpperCamelCase_		)
        return model
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[str]     =       '''cpu'''

        lowercase_       :int     =       self.dummy_vqvae
        lowercase_       :int     =       self.dummy_text_encoder
        lowercase_       :Any     =       self.dummy_tokenizer
        lowercase_       :Optional[int]     =       self.dummy_transformer
        lowercase_       :List[str]     =       VQDiffusionScheduler(self.num_embed		)
        lowercase_       :int     =       LearnedClassifierFreeSamplingEmbeddings(learnable=UpperCamelCase_		)

        lowercase_       :List[Any]     =       VQDiffusionPipeline(
            vqvae=UpperCamelCase_   ,							text_encoder=UpperCamelCase_   ,							tokenizer=UpperCamelCase_   ,							transformer=UpperCamelCase_   ,							scheduler=UpperCamelCase_   ,							learned_classifier_free_sampling_embeddings=UpperCamelCase_   ,							)
        lowercase_       :Union[str, Any]     =       pipe.to(UpperCamelCase_		)
        pipe.set_progress_bar_config(disable=UpperCamelCase_		)

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

        lowercase_       :Dict     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :Optional[int]     =       pipe([prompt]   ,							generator=UpperCamelCase_   ,							num_inference_steps=2   ,							output_type='''np'''		)
        lowercase_       :Any     =       output.images

        lowercase_       :Tuple     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :str     =       pipe(
            [prompt]   ,							generator=UpperCamelCase_   ,							output_type='''np'''   ,							return_dict=UpperCamelCase_   ,							num_inference_steps=2		)[0]

        lowercase_       :Optional[Any]     =       image[0, -3:, -3:, -1]
        lowercase_       :Tuple     =       image_from_tuple[0, -3:, -3:, -1]

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

        lowercase_       :str     =       np.array([0.6551, 0.6168, 0.5008, 0.5676, 0.5659, 0.4295, 0.6073, 0.5599, 0.4992]		)

        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		):
        lowercase_       :int     =       '''cpu'''

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

        lowercase_       :Optional[int]     =       VQDiffusionPipeline(
            vqvae=UpperCamelCase_   ,							text_encoder=UpperCamelCase_   ,							tokenizer=UpperCamelCase_   ,							transformer=UpperCamelCase_   ,							scheduler=UpperCamelCase_   ,							learned_classifier_free_sampling_embeddings=UpperCamelCase_   ,							)
        lowercase_       :Dict     =       pipe.to(UpperCamelCase_		)
        pipe.set_progress_bar_config(disable=UpperCamelCase_		)

        lowercase_       :int     =       '''teddy bear playing in the pool'''

        lowercase_       :Tuple     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :Optional[int]     =       pipe([prompt]   ,							generator=UpperCamelCase_   ,							num_inference_steps=2   ,							output_type='''np'''		)
        lowercase_       :Optional[Any]     =       output.images

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

        lowercase_       :List[str]     =       image[0, -3:, -3:, -1]
        lowercase_       :Optional[int]     =       image_from_tuple[0, -3:, -3:, -1]

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

        lowercase_       :Dict     =       np.array([0.6693, 0.6075, 0.4959, 0.5701, 0.5583, 0.4333, 0.6171, 0.5684, 0.4988]		)

        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      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   def 		UpperCamelCase  (     self		):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()




   def 		UpperCamelCase  (     self		):
        lowercase_       :List[str]     =       load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''		)

        lowercase_       :int     =       VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''		)
        lowercase_       :Tuple     =       pipeline.to(UpperCamelCase_		)
        pipeline.set_progress_bar_config(disable=UpperCamelCase_		)

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

        lowercase_       :List[str]     =       output.images[0]

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

