


'''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 	_A				(lowerCAmelCase__  :List[str]				,						lowerCAmelCase__  :int		)   ->  Tuple:




     '''simple docstring'''



     _a			      =							[]
     for part_id in partition_order:
          _a			      =							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 	_A				()   ->  int:




     '''simple docstring'''



     _a			      =							pyspark.sql.SparkSession.builder.master('local[*]'		).appName('pyspark'		).getOrCreate()
     _a			      =							spark.range(1_00		).repartition(1		)
     _a			      =							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 	_A				()   ->  Any:




     '''simple docstring'''



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

     for i, (row_id, row_dict) in enumerate(generate_fn()		):
          _a      ,   _a			      =							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 	_A				()   ->  Optional[int]:




     '''simple docstring'''



     _a			      =							pyspark.sql.SparkSession.builder.master('local[*]'		).appName('pyspark'		).getOrCreate()
     _a			      =							spark.range(10		).repartition(1		)
     _a			      =							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 	_A				()   ->  int:




     '''simple docstring'''



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

          _a			      =							SparkExamplesIterable(lowerCAmelCase__		).shuffle_data_sources(lowerCAmelCase__		)
          assert shuffled_it.n_shards == 3
          for i, (row_id, row_dict) in enumerate(lowerCAmelCase__		):
               _a      ,   _a			      =							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 	_A				()   ->  Tuple:




     '''simple docstring'''



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

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

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




     '''simple docstring'''



     _a			      =							pyspark.sql.SparkSession.builder.master('local[*]'		).appName('pyspark'		).getOrCreate()
     _a			      =							spark.range(1_00		).repartition(1		)
     _a			      =							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





'''simple docstring'''






import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class  a     (     _SCREAMING_SNAKE_CASE  ,				unittest.TestCase       ):
 _lowerCAmelCase 					=	DiTPipeline
 _lowerCAmelCase 					=	CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
 _lowerCAmelCase 					=	PipelineTesterMixin.required_optional_params - {
     """latents""",
     """num_images_per_prompt""",
     """callback""",
     """callback_steps""",
 }
 _lowerCAmelCase 					=	CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
 _lowerCAmelCase 					=	False
 def 							__UpperCAmelCase		( self						)					->							Dict:
      torch.manual_seed(0						)
      _a			      =							TransformeraDModel(
          sample_size=16      ,				num_layers=2      ,				patch_size=4      ,				attention_head_dim=8      ,				num_attention_heads=2      ,				in_channels=4      ,				out_channels=8      ,				attention_bias=__magic_name__      ,				activation_fn='gelu-approximate'      ,				num_embeds_ada_norm=10_00      ,				norm_type='ada_norm_zero'      ,				norm_elementwise_affine=__magic_name__      ,				)
      _a			      =							AutoencoderKL()
      _a			      =							DDIMScheduler()
      _a			      =							{'transformer': transformer.eval(), 'vae': vae.eval(), 'scheduler': scheduler}
      return components
 def 							__UpperCAmelCase		( self      ,				__magic_name__      ,				__magic_name__=0						)					->							Union[str, Any]:
      if str(__magic_name__						).startswith('mps'						):
           _a			      =							torch.manual_seed(__magic_name__						)
      else:
           _a			      =							torch.Generator(device=__magic_name__						).manual_seed(__magic_name__						)
      _a			      =							{
          'class_labels': [1],
          'generator': generator,
          'num_inference_steps': 2,
          'output_type': 'numpy',
      }
      return inputs
 def 							__UpperCAmelCase		( self						)					->							Tuple:
      _a			      =							'cpu'

      _a			      =							self.get_dummy_components()
      _a			      =							self.pipeline_class(**__magic_name__						)
      pipe.to(__magic_name__						)
      pipe.set_progress_bar_config(disable=__magic_name__						)

      _a			      =							self.get_dummy_inputs(__magic_name__						)
      _a			      =							pipe(**__magic_name__						).images
      _a			      =							image[0, -3:, -3:, -1]

      self.assertEqual(image.shape      ,				(1, 16, 16, 3)						)
      _a			      =							np.array([0.2_9_4_6, 0.6_6_0_1, 0.4_3_2_9, 0.3_2_9_6, 0.4_1_4_4, 0.5_3_1_9, 0.7_2_7_3, 0.5_0_1_3, 0.4_4_5_7]						)
      _a			      =							np.abs(image_slice.flatten() - expected_slice						).max()
      self.assertLessEqual(__magic_name__      ,				1e-3						)
 def 							__UpperCAmelCase		( self						)					->							Any:
      self._test_inference_batch_single_identical(relax_max_difference=__magic_name__      ,				expected_max_diff=1e-3						)




 @unittest.skipIf(
     torch_device != 'cuda' or not is_xformers_available()      ,				reason='XFormers attention is only available with CUDA and `xformers` installed'      ,				)
 def 							__UpperCAmelCase		( self						)					->							Any:
      self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3						)





@require_torch_gpu
@slow
class  a     (     unittest.TestCase       ):
 def 							__UpperCAmelCase		( self						)					->							List[str]:
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()
 def 							__UpperCAmelCase		( self						)					->							Union[str, Any]:
      _a			      =							torch.manual_seed(0						)

      _a			      =							DiTPipeline.from_pretrained('facebook/DiT-XL-2-256'						)
      pipe.to('cuda'						)

      _a			      =							['vase', 'umbrella', 'white shark', 'white wolf']
      _a			      =							pipe.get_label_ids(__magic_name__						)

      _a			      =							pipe(__magic_name__      ,				generator=__magic_name__      ,				num_inference_steps=40      ,				output_type='np'						).images

      for word, image in zip(__magic_name__      ,				__magic_name__						):
           _a			      =							load_numpy(
               f'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy'						)
           assert np.abs((expected_image - image).max()						) < 1e-2




 def 							__UpperCAmelCase		( self						)					->							Any:
      _a			      =							DiTPipeline.from_pretrained('facebook/DiT-XL-2-512'						)
      _a			      =							DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
      pipe.to('cuda'						)

      _a			      =							['vase', 'umbrella']
      _a			      =							pipe.get_label_ids(__magic_name__						)

      _a			      =							torch.manual_seed(0						)
      _a			      =							pipe(__magic_name__      ,				generator=__magic_name__      ,				num_inference_steps=25      ,				output_type='np'						).images

      for word, image in zip(__magic_name__      ,				__magic_name__						):
           _a			      =							load_numpy(
               'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
               f'/dit/{word}_512.npy'						)

           assert np.abs((expected_image - image).max()						) < 1e-1

