






import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.text import TextDatasetReader

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases


def   lowerCamelCase__			(     lowercase						,     lowercase				):

  """simple docstring"""






  assert isinstance(lowercase						,     lowercase				)
  assert dataset.num_rows == 4
  assert dataset.num_columns == 1
  assert dataset.column_names == ["text"]
  for feature, expected_dtype in expected_features.items():
    assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("keep_in_memory"						,     [False, True]				)
def   lowerCamelCase__			(     lowercase						,     lowercase						,     lowercase				):

  """simple docstring"""






  SCREAMING_SNAKE_CASE       :      List[Any]							=    tmp_path / "cache"
  SCREAMING_SNAKE_CASE       :      Tuple							=    {"text": "string"}
  with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
    SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    TextDatasetReader(lowercase						,     cache_dir=lowercase						,     keep_in_memory=lowercase				).read()
  _check_text_dataset(lowercase						,     lowercase				)


@pytest.mark.parametrize(
    "features"						,     [
        None,
        {"text": "string"},
        {"text": "int32"},
        {"text": "float32"},
    ]						,     )
def   lowerCamelCase__			(     lowercase						,     lowercase						,     lowercase				):

  """simple docstring"""






  SCREAMING_SNAKE_CASE       :      Optional[Any]							=    tmp_path / "cache"
  SCREAMING_SNAKE_CASE       :      int							=    {"text": "string"}
  SCREAMING_SNAKE_CASE       :      List[Any]							=    features.copy() if features else default_expected_features
  SCREAMING_SNAKE_CASE       :      Any							=    (
      Features({feature: Value(lowercase				) for feature, dtype in features.items()}				) if features is not None else None
  )
  SCREAMING_SNAKE_CASE       :      List[Any]							=    TextDatasetReader(lowercase						,     features=lowercase						,     cache_dir=lowercase				).read()
  _check_text_dataset(lowercase						,     lowercase				)


@pytest.mark.parametrize("split"						,     [None, NamedSplit("train"				), "train", "test"]				)
def   lowerCamelCase__			(     lowercase						,     lowercase						,     lowercase				):

  """simple docstring"""






  SCREAMING_SNAKE_CASE       :      List[Any]							=    tmp_path / "cache"
  SCREAMING_SNAKE_CASE       :      List[Any]							=    {"text": "string"}
  SCREAMING_SNAKE_CASE       :      Optional[int]							=    TextDatasetReader(lowercase						,     cache_dir=lowercase						,     split=lowercase				).read()
  _check_text_dataset(lowercase						,     lowercase				)
  assert dataset.split == split if split else "train"


@pytest.mark.parametrize("path_type"						,     [str, list]				)
def   lowerCamelCase__			(     lowercase						,     lowercase						,     lowercase				):

  """simple docstring"""






  if issubclass(lowercase						,     lowercase				):
    SCREAMING_SNAKE_CASE       :      List[Any]							=    text_path
  elif issubclass(lowercase						,     lowercase				):
    SCREAMING_SNAKE_CASE       :      Any							=    [text_path]
  SCREAMING_SNAKE_CASE       :      str							=    tmp_path / "cache"
  SCREAMING_SNAKE_CASE       :      Optional[int]							=    {"text": "string"}
  SCREAMING_SNAKE_CASE       :      Dict							=    TextDatasetReader(lowercase						,     cache_dir=lowercase				).read()
  _check_text_dataset(lowercase						,     lowercase				)


def   lowerCamelCase__			(     lowercase						,     lowercase						,     lowercase=("train",)				):

  """simple docstring"""






  assert isinstance(lowercase						,     lowercase				)
  for split in splits:
    SCREAMING_SNAKE_CASE       :      str							=    dataset_dict[split]
    assert dataset.num_rows == 4
    assert dataset.num_columns == 1
    assert dataset.column_names == ["text"]
    for feature, expected_dtype in expected_features.items():
      assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("keep_in_memory"						,     [False, True]				)
def   lowerCamelCase__			(     lowercase						,     lowercase						,     lowercase				):

  """simple docstring"""






  SCREAMING_SNAKE_CASE       :      Optional[Any]							=    tmp_path / "cache"
  SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    {"text": "string"}
  with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
    SCREAMING_SNAKE_CASE       :      Tuple							=    TextDatasetReader({"train": text_path}						,     cache_dir=lowercase						,     keep_in_memory=lowercase				).read()
  _check_text_datasetdict(lowercase						,     lowercase				)


@pytest.mark.parametrize(
    "features"						,     [
        None,
        {"text": "string"},
        {"text": "int32"},
        {"text": "float32"},
    ]						,     )
def   lowerCamelCase__			(     lowercase						,     lowercase						,     lowercase				):

  """simple docstring"""






  SCREAMING_SNAKE_CASE       :      List[str]							=    tmp_path / "cache"
  # CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
  SCREAMING_SNAKE_CASE       :      List[Any]							=    {"text": "string"}
  SCREAMING_SNAKE_CASE       :      List[str]							=    features.copy() if features else default_expected_features
  SCREAMING_SNAKE_CASE       :      List[str]							=    (
      Features({feature: Value(lowercase				) for feature, dtype in features.items()}				) if features is not None else None
  )
  SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    TextDatasetReader({"train": text_path}						,     features=lowercase						,     cache_dir=lowercase				).read()
  _check_text_datasetdict(lowercase						,     lowercase				)


@pytest.mark.parametrize("split"						,     [None, NamedSplit("train"				), "train", "test"]				)
def   lowerCamelCase__			(     lowercase						,     lowercase						,     lowercase				):

  """simple docstring"""






  if split:
    SCREAMING_SNAKE_CASE       :      List[str]							=    {split: text_path}
  else:
    SCREAMING_SNAKE_CASE       :      Any							=    "train"
    SCREAMING_SNAKE_CASE       :      Optional[Any]							=    {"train": text_path, "test": text_path}
  SCREAMING_SNAKE_CASE       :      Optional[int]							=    tmp_path / "cache"
  SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    {"text": "string"}
  SCREAMING_SNAKE_CASE       :      Optional[int]							=    TextDatasetReader(lowercase						,     cache_dir=lowercase				).read()
  _check_text_datasetdict(lowercase						,     lowercase						,     splits=list(path.keys()				)				)
  assert all(dataset[split].split == split for split in path.keys()				)














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  SCREAMING_SNAKE_CASE			(			lowerCAmelCase      ,		lowerCAmelCase      ,		unittest.TestCase					):
     '''simple docstring'''


     UpperCamelCase_							:		Optional[Any]			=	IFImgaImgSuperResolutionPipeline
     UpperCamelCase_							:		Any			=	TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''width''', '''height'''}
     UpperCamelCase_							:		Optional[Any]			=	TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'''original_image'''}					)
     UpperCamelCase_							:		str			=	PipelineTesterMixin.required_optional_params - {'''latents'''}


     def 				_A       (				self       :	List[Any]      ):
       return self._get_superresolution_dummy_components()


     def 				_A       (				self       :	Any					, UpperCAmelCase_       :	int					, UpperCAmelCase_       :	List[str]=0      ):
       if str(UpperCAmelCase_      ).startswith("mps"      ):
         SCREAMING_SNAKE_CASE       :      Any							=    torch.manual_seed(UpperCAmelCase_      )
       else:
         SCREAMING_SNAKE_CASE       :      List[str]							=    torch.Generator(device=UpperCAmelCase_      ).manual_seed(UpperCAmelCase_      )

       SCREAMING_SNAKE_CASE       :      Any							=    floats_tensor((1, 3, 32, 32)					, rng=random.Random(UpperCAmelCase_      )      ).to(UpperCAmelCase_      )
       SCREAMING_SNAKE_CASE       :      Any							=    floats_tensor((1, 3, 16, 16)					, rng=random.Random(UpperCAmelCase_      )      ).to(UpperCAmelCase_      )

       SCREAMING_SNAKE_CASE       :      int							=    {
           "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 				_A       (				self       :	Tuple      ):
       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3      )


     def 				_A       (				self       :	Optional[Any]      ):
       self._test_save_load_optional_components()


     @unittest.skipIf(torch_device != "cuda"					, reason="float16 requires CUDA"      )
     def 				_A       (				self       :	List[Any]      ):
       # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
       super().test_save_load_floataa(expected_max_diff=1E-1      )


     def 				_A       (				self       :	Optional[int]      ):
       self._test_attention_slicing_forward_pass(expected_max_diff=1E-2      )


     def 				_A       (				self       :	Tuple      ):
       self._test_save_load_local()






     def 				_A       (				self       :	List[str]      ):
       self._test_inference_batch_single_identical(
           expected_max_diff=1E-2					, )






