






import pyarrow.parquet as pq
import pytest

from datasets import Audio, Dataset, DatasetDict, Features, NamedSplit, Sequence, Value, config
from datasets.features.image import Image
from datasets.io.parquet import ParquetDatasetReader, ParquetDatasetWriter, get_writer_batch_size

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases




def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case )				->						Optional[Any]:
       assert isinstance(__snake_case   ,      __snake_case )
       assert dataset.num_rows == 4
       assert dataset.num_columns == 3
       assert dataset.column_names == ["col_1", "col_2", "col_3"]
       for feature, expected_dtype in expected_features.items():
              assert dataset.features[feature].dtype == expected_dtype




@pytest.mark.parametrize("""keep_in_memory"""   ,      [False, True] )
def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case   ,      __snake_case )				->						Optional[Any]:
       _UpperCAmelCase							=    tmp_path / """cache"""
       _UpperCAmelCase							=    {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
       with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
              _UpperCAmelCase							=    ParquetDatasetReader(__snake_case   ,      cache_dir=__snake_case   ,      keep_in_memory=__snake_case ).read()
       _check_parquet_dataset(__snake_case   ,      __snake_case )




@pytest.mark.parametrize(
    """features"""   ,      [
        None,
        {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""},
        {"""col_1""": """string""", """col_2""": """string""", """col_3""": """string"""},
        {"""col_1""": """int32""", """col_2""": """int32""", """col_3""": """int32"""},
        {"""col_1""": """float32""", """col_2""": """float32""", """col_3""": """float32"""},
    ]   ,      )
def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case   ,      __snake_case )				->						Optional[int]:
       _UpperCAmelCase							=    tmp_path / """cache"""
       _UpperCAmelCase							=    {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
       _UpperCAmelCase							=    features.copy() if features else default_expected_features
       _UpperCAmelCase							=    (
           Features({feature: Value(__snake_case ) for feature, dtype in features.items()} ) if features is not None else None
       )
       _UpperCAmelCase							=    ParquetDatasetReader(__snake_case   ,      features=__snake_case   ,      cache_dir=__snake_case ).read()
       _check_parquet_dataset(__snake_case   ,      __snake_case )




@pytest.mark.parametrize("""split"""   ,      [None, NamedSplit("""train""" ), """train""", """test"""] )
def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case   ,      __snake_case )				->						Any:
       _UpperCAmelCase							=    tmp_path / """cache"""
       _UpperCAmelCase							=    {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
       _UpperCAmelCase							=    ParquetDatasetReader(__snake_case   ,      cache_dir=__snake_case   ,      split=__snake_case ).read()
       _check_parquet_dataset(__snake_case   ,      __snake_case )
       assert dataset.split == split if split else "train"




@pytest.mark.parametrize("""path_type"""   ,      [str, list] )
def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case   ,      __snake_case )				->						Any:
       if issubclass(__snake_case   ,      __snake_case ):
              _UpperCAmelCase							=    parquet_path
       elif issubclass(__snake_case   ,      __snake_case ):
              _UpperCAmelCase							=    [parquet_path]
       _UpperCAmelCase							=    tmp_path / """cache"""
       _UpperCAmelCase							=    {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
       _UpperCAmelCase							=    ParquetDatasetReader(__snake_case   ,      cache_dir=__snake_case ).read()
       _check_parquet_dataset(__snake_case   ,      __snake_case )




def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case   ,      __snake_case=("train",) )				->						List[str]:
       assert isinstance(__snake_case   ,      __snake_case )
       for split in splits:
              _UpperCAmelCase							=    dataset_dict[split]
              assert dataset.num_rows == 4
              assert dataset.num_columns == 3
              assert dataset.column_names == ["col_1", "col_2", "col_3"]
              for feature, expected_dtype in expected_features.items():
                     assert dataset.features[feature].dtype == expected_dtype




@pytest.mark.parametrize("""keep_in_memory"""   ,      [False, True] )
def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case   ,      __snake_case )				->						Tuple:
       _UpperCAmelCase							=    tmp_path / """cache"""
       _UpperCAmelCase							=    {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
       with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
              _UpperCAmelCase							=    ParquetDatasetReader(
                  {"""train""": parquet_path}   ,      cache_dir=__snake_case   ,      keep_in_memory=__snake_case ).read()
       _check_parquet_datasetdict(__snake_case   ,      __snake_case )




@pytest.mark.parametrize(
    """features"""   ,      [
        None,
        {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""},
        {"""col_1""": """string""", """col_2""": """string""", """col_3""": """string"""},
        {"""col_1""": """int32""", """col_2""": """int32""", """col_3""": """int32"""},
        {"""col_1""": """float32""", """col_2""": """float32""", """col_3""": """float32"""},
    ]   ,      )
def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case   ,      __snake_case )				->						Optional[Any]:
       _UpperCAmelCase							=    tmp_path / """cache"""
       _UpperCAmelCase							=    {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
       _UpperCAmelCase							=    features.copy() if features else default_expected_features
       _UpperCAmelCase							=    (
           Features({feature: Value(__snake_case ) for feature, dtype in features.items()} ) if features is not None else None
       )
       _UpperCAmelCase							=    ParquetDatasetReader({"""train""": parquet_path}   ,      features=__snake_case   ,      cache_dir=__snake_case ).read()
       _check_parquet_datasetdict(__snake_case   ,      __snake_case )




@pytest.mark.parametrize("""split"""   ,      [None, NamedSplit("""train""" ), """train""", """test"""] )
def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case   ,      __snake_case )				->						List[Any]:
       if split:
              _UpperCAmelCase							=    {split: parquet_path}
       else:
              _UpperCAmelCase							=    """train"""
              _UpperCAmelCase							=    {"""train""": parquet_path, """test""": parquet_path}
       _UpperCAmelCase							=    tmp_path / """cache"""
       _UpperCAmelCase							=    {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
       _UpperCAmelCase							=    ParquetDatasetReader(__snake_case   ,      cache_dir=__snake_case ).read()
       _check_parquet_datasetdict(__snake_case   ,      __snake_case   ,      splits=list(path.keys() ) )
       assert all(dataset[split].split == split for split in path.keys() )




def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case )				->						Dict:
       _UpperCAmelCase							=    ParquetDatasetWriter(__snake_case   ,      tmp_path / """foo.parquet""" )
       assert writer.write() > 0
       _UpperCAmelCase							=    pq.ParquetFile(tmp_path / """foo.parquet""" )
       _UpperCAmelCase							=    pf.read()
       assert dataset.data.table == output_table




def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case )				->						List[Any]:
       _UpperCAmelCase							=    str(shared_datadir / """test_image_rgb.jpg""" )
       _UpperCAmelCase							=    {"""image""": [image_path]}
       _UpperCAmelCase							=    Features({"""image""": Image()} )
       _UpperCAmelCase							=    Dataset.from_dict(__snake_case   ,      features=__snake_case )
       _UpperCAmelCase							=    ParquetDatasetWriter(__snake_case   ,      tmp_path / """foo.parquet""" )
       assert writer.write() > 0

       _UpperCAmelCase							=    Dataset.from_parquet(str(tmp_path / """foo.parquet""" ) )
       assert dataset.features == reloaded_dataset.features

       _UpperCAmelCase							=    ParquetDatasetReader(str(tmp_path / """foo.parquet""" )   ,      streaming=__snake_case ).read()
       assert dataset.features == reloaded_iterable_dataset.features





@pytest.mark.parametrize(
    """feature, expected"""   ,      [
        (Features({"""foo""": Value("""int32""" )} ), None),
        (Features({"""image""": Image(), """foo""": Value("""int32""" )} ), config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS),
        (Features({"""nested""": Sequence(Audio() )} ), config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS),
    ]   ,      )
def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case )				->						List[Any]:
       assert get_writer_batch_size(__snake_case ) == expected




import random
import unittest

import torch

from diffusers import IFImgaImgSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin



@skip_mps
class    UpperCAmelCase__(				lowerCamelCase  ,				lowerCamelCase  ,				unittest.TestCase						):

   '''simple docstring'''




   A     :	str												=     IFImgaImgSuperResolutionPipeline
   A     :	Dict												=     TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"width", "height"}
   A     :	Any												=     TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({"original_image"}						)
   A     :	str												=     PipelineTesterMixin.required_optional_params - {"latents"}


   def       UpperCAmelCase		(    self	:      Dict)						->					Dict:




         """simple docstring"""

         return self._get_superresolution_dummy_components()


   def       UpperCAmelCase		(    self	:      Optional[Any]      ,					lowerCAmelCase	:      Tuple      ,					lowerCAmelCase	:      Union[str, Any]=0)						->					str:




         """simple docstring"""

         if str(lowerCAmelCase).startswith('mps'):
               lowercase__  =	torch.manual_seed(lowerCAmelCase)
         else:
               lowercase__  =	torch.Generator(device=lowerCAmelCase).manual_seed(lowerCAmelCase)

         lowercase__  =	floats_tensor((1, 3, 32, 32)      ,					rng=random.Random(lowerCAmelCase)).to(lowerCAmelCase)
         lowercase__  =	floats_tensor((1, 3, 16, 16)      ,					rng=random.Random(lowerCAmelCase)).to(lowerCAmelCase)

         lowercase__  =	{
             'prompt': 'A painting of a squirrel eating a burger',
             'image': image,
             'original_image': original_image,
             'generator': generator,
             'num_inference_steps': 2,
             'output_type': 'numpy',
         }

         return inputs


   @unittest.skipIf(
       torch_device != 'cuda' or not is_xformers_available()      ,					reason='XFormers attention is only available with CUDA and `xformers` installed'      ,					)
   def       UpperCAmelCase		(    self	:      str)						->					Dict:




         """simple docstring"""

         self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3)


   def       UpperCAmelCase		(    self	:      List[Any])						->					Dict:




         """simple docstring"""

         self._test_save_load_optional_components()


   @unittest.skipIf(torch_device != 'cuda'      ,					reason='float16 requires CUDA')
   def       UpperCAmelCase		(    self	:      Optional[Any])						->					List[Any]:




         """simple docstring"""

         super().test_save_load_floataa(expected_max_diff=1E-1)


   def       UpperCAmelCase		(    self	:      str)						->					Union[str, Any]:




         """simple docstring"""

         self._test_attention_slicing_forward_pass(expected_max_diff=1E-2)


   def       UpperCAmelCase		(    self	:      Tuple)						->					Tuple:




         """simple docstring"""

         self._test_save_load_local()



   def       UpperCAmelCase		(    self	:      List[str])						->					Tuple:




         """simple docstring"""

         self._test_inference_batch_single_identical(
             expected_max_diff=1E-2      ,					)
