





"""simple docstring"""


import random
import unittest

import numpy as np
import torch
from PIL import Image
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorEmbaEmbPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class    __SCREAMING_SNAKE_CASE			(  _UpperCamelCase						,     unittest.TestCase      ):







  '''simple docstring'''


  SCREAMING_SNAKE_CASE__							:List[Any]		   =    KandinskyVaaPriorEmbaEmbPipeline
  SCREAMING_SNAKE_CASE__							:List[Any]		   =    ["prompt", "image"]
  SCREAMING_SNAKE_CASE__							:Optional[int]		   =    ["prompt", "image"]
  SCREAMING_SNAKE_CASE__							:List[Any]		   =    [
      "num_images_per_prompt",
      "strength",
      "generator",
      "num_inference_steps",
      "latents",
      "negative_prompt",
      "guidance_scale",
      "output_type",
      "return_dict",
  ]
  SCREAMING_SNAKE_CASE__							:Optional[int]		   =    False

  @property
  def       __SCREAMING_SNAKE_CASE (						self      :				Optional[Any]		)						-> Optional[int]:
        return 32

  @property
  def       __SCREAMING_SNAKE_CASE (						self      :				Optional[Any]		)						-> List[Any]:
        return 32

  @property
  def       __SCREAMING_SNAKE_CASE (						self      :				Optional[Any]		)						-> Dict:
        return self.time_input_dim

  @property
  def       __SCREAMING_SNAKE_CASE (						self      :				Any		)						-> Optional[Any]:
        return self.time_input_dim * 4

  @property
  def       __SCREAMING_SNAKE_CASE (						self      :				List[str]		)						-> int:
        return 100

  @property
  def       __SCREAMING_SNAKE_CASE (						self      :				str		)						-> Any:
        _UpperCamelCase					:		str  =				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)
        return tokenizer

  @property
  def       __SCREAMING_SNAKE_CASE (						self      :				str		)						-> Tuple:
        torch.manual_seed(0		)
        _UpperCamelCase					:		Dict  =				CLIPTextConfig(
            bos_token_id=0    ,      eos_token_id=2    ,      hidden_size=self.text_embedder_hidden_size    ,      projection_dim=self.text_embedder_hidden_size    ,      intermediate_size=37    ,      layer_norm_eps=1e-0_5    ,      num_attention_heads=4    ,      num_hidden_layers=5    ,      pad_token_id=1    ,      vocab_size=1000    ,      )
        return CLIPTextModelWithProjection(__a		)

  @property
  def       __SCREAMING_SNAKE_CASE (						self      :				int		)						-> Any:
        torch.manual_seed(0		)

        _UpperCamelCase					:		Union[str, Any]  =				{
            'num_attention_heads': 2,
            'attention_head_dim': 12,
            'embedding_dim': self.text_embedder_hidden_size,
            'num_layers': 1,
        }

        _UpperCamelCase					:		Optional[int]  =				PriorTransformer(**__a		)
        # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
        _UpperCamelCase					:		Optional[Any]  =				nn.Parameter(torch.ones(model.clip_std.shape		)		)
        return model

  @property
  def       __SCREAMING_SNAKE_CASE (						self      :				Optional[int]		)						-> List[Any]:
        torch.manual_seed(0		)
        _UpperCamelCase					:		Any  =				CLIPVisionConfig(
            hidden_size=self.text_embedder_hidden_size    ,      image_size=224    ,      projection_dim=self.text_embedder_hidden_size    ,      intermediate_size=37    ,      num_attention_heads=4    ,      num_channels=3    ,      num_hidden_layers=5    ,      patch_size=14    ,      )

        _UpperCamelCase					:		int  =				CLIPVisionModelWithProjection(__a		)
        return model

  @property
  def       __SCREAMING_SNAKE_CASE (						self      :				Optional[Any]		)						-> Dict:
        _UpperCamelCase					:		str  =				CLIPImageProcessor(
            crop_size=224    ,      do_center_crop=__a    ,      do_normalize=__a    ,      do_resize=__a    ,      image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]    ,      image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]    ,      resample=3    ,      size=224    ,      )

        return image_processor

  def       __SCREAMING_SNAKE_CASE (						self      :				List[str]		)						-> Tuple:
        _UpperCamelCase					:		Union[str, Any]  =				self.dummy_prior
        _UpperCamelCase					:		str  =				self.dummy_image_encoder
        _UpperCamelCase					:		Optional[int]  =				self.dummy_text_encoder
        _UpperCamelCase					:		int  =				self.dummy_tokenizer
        _UpperCamelCase					:		Dict  =				self.dummy_image_processor

        _UpperCamelCase					:		Any  =				UnCLIPScheduler(
            variance_type="fixed_small_log"    ,      prediction_type="sample"    ,      num_train_timesteps=1000    ,      clip_sample=__a    ,      clip_sample_range=10.0    ,      )

        _UpperCamelCase					:		Any  =				{
            'prior': prior,
            'image_encoder': image_encoder,
            'text_encoder': text_encoder,
            'tokenizer': tokenizer,
            'scheduler': scheduler,
            'image_processor': image_processor,
        }

        return components

  def       __SCREAMING_SNAKE_CASE (						self      :				Any    ,      __a      :				Dict    ,      __a      :				Any=0		)						-> List[str]:




        if str(__a		).startswith("mps"		):
              _UpperCamelCase					:		str  =				torch.manual_seed(__a		)
        else:
              _UpperCamelCase					:		Dict  =				torch.Generator(device=__a		).manual_seed(__a		)

        _UpperCamelCase					:		int  =				floats_tensor((1, 3, 64, 64)    ,      rng=random.Random(__a		)		).to(__a		)
        _UpperCamelCase					:		Optional[Any]  =				image.cpu().permute(0    ,      2    ,      3    ,      1		)[0]
        _UpperCamelCase					:		List[Any]  =				Image.fromarray(np.uinta(__a		)		).convert("RGB"		).resize((256, 256)		)

        _UpperCamelCase					:		Optional[int]  =				{
            'prompt': 'horse',
            'image': init_image,
            'strength': 0.5,
            'generator': generator,
            'guidance_scale': 4.0,
            'num_inference_steps': 2,
            'output_type': 'np',
        }
        return inputs

  def       __SCREAMING_SNAKE_CASE (						self      :				Optional[int]		)						-> Tuple:
        _UpperCamelCase					:		int  =				'cpu'

        _UpperCamelCase					:		Union[str, Any]  =				self.get_dummy_components()

        _UpperCamelCase					:		Union[str, Any]  =				self.pipeline_class(**__a		)
        _UpperCamelCase					:		Tuple  =				pipe.to(__a		)

        pipe.set_progress_bar_config(disable=__a		)

        _UpperCamelCase					:		List[str]  =				pipe(**self.get_dummy_inputs(__a		)		)
        _UpperCamelCase					:		List[str]  =				output.image_embeds

        _UpperCamelCase					:		Union[str, Any]  =				pipe(
            **self.get_dummy_inputs(__a		)    ,      return_dict=__a    ,      )[0]

        _UpperCamelCase					:		Optional[Any]  =				image[0, -10:]
        _UpperCamelCase					:		Any  =				image_from_tuple[0, -10:]

        assert image.shape == (1, 32)

        _UpperCamelCase					:		Tuple  =				np.array(
            [
                0.1_07_12_84,
                1.3_33_02_71,
                0.61_26_02_23,
                -0.6_69_10_65,
                -0.3_84_68_52,
                -1.0_30_36_61,
                0.22_71_61_11,
                0.03_34_89_01,
                0.30_04_06_75,
                -0.24_80_50_29,
            ]		)

        assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-2
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1e-2

  @skip_mps
  def       __SCREAMING_SNAKE_CASE (						self      :				Union[str, Any]		)						-> Optional[int]:
        _UpperCamelCase					:		int  =				torch_device == 'cpu'
        _UpperCamelCase					:		Any  =				True
        _UpperCamelCase					:		Union[str, Any]  =				False

        self._test_inference_batch_single_identical(
            test_max_difference=__a    ,      relax_max_difference=__a    ,      test_mean_pixel_difference=__a    ,      )


  @skip_mps
  def       __SCREAMING_SNAKE_CASE (						self      :				Union[str, Any]		)						-> List[Any]:
        _UpperCamelCase					:		List[Any]  =				torch_device == 'cpu'
        _UpperCamelCase					:		int  =				False

        self._test_attention_slicing_forward_pass(
            test_max_difference=__a    ,      test_mean_pixel_difference=__a    ,      )








"""simple docstring"""


import csv
import os

import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.csv import CsvDatasetReader, CsvDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases




def 			snake_case  (A_      :Dict	,			A_      :Tuple      ):







  '''simple docstring'''

  assert isinstance(A_	,			A_      )
  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 			snake_case  (A_      :str	,			A_      :List[Any]	,			A_      :Optional[Any]      ):







  '''simple docstring'''

  a     :			List[Any]				=   tmp_path / 'cache'
  a     :			str				=   {'col_1': 'int64', 'col_2': 'int64', 'col_3': 'float64'}
  with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
    a     :			Union[str, Any]				=   CsvDatasetReader(A_	,			cache_dir=A_	,			keep_in_memory=A_      ).read()
  _check_csv_dataset(A_	,			A_      )




@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 			snake_case  (A_      :List[str]	,			A_      :Union[str, Any]	,			A_      :Union[str, Any]      ):







  '''simple docstring'''

  a     :			str				=   tmp_path / 'cache'
  # CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
  a     :			Optional[int]				=   {'col_1': 'int64', 'col_2': 'int64', 'col_3': 'float64'}
  a     :			Dict				=   features.copy() if features else default_expected_features
  a     :			Dict				=   (
      Features({feature: Value(A_      ) for feature, dtype in features.items()}      ) if features is not None else None
  )
  a     :			Tuple				=   CsvDatasetReader(A_	,			features=A_	,			cache_dir=A_      ).read()
  _check_csv_dataset(A_	,			A_      )




@pytest.mark.parametrize('split'	,			[None, NamedSplit('train'      ), 'train', 'test']      )
def 			snake_case  (A_      :List[str]	,			A_      :str	,			A_      :int      ):







  '''simple docstring'''

  a     :			str				=   tmp_path / 'cache'
  a     :			Tuple				=   {'col_1': 'int64', 'col_2': 'int64', 'col_3': 'float64'}
  a     :			List[str]				=   CsvDatasetReader(A_	,			cache_dir=A_	,			split=A_      ).read()
  _check_csv_dataset(A_	,			A_      )
  assert dataset.split == split if split else "train"




@pytest.mark.parametrize('path_type'	,			[str, list]      )
def 			snake_case  (A_      :int	,			A_      :List[Any]	,			A_      :str      ):







  '''simple docstring'''

  if issubclass(A_	,			A_      ):
    a     :			Dict				=   csv_path
  elif issubclass(A_	,			A_      ):
    a     :			Optional[Any]				=   [csv_path]
  a     :			Tuple				=   tmp_path / 'cache'
  a     :			str				=   {'col_1': 'int64', 'col_2': 'int64', 'col_3': 'float64'}
  a     :			Any				=   CsvDatasetReader(A_	,			cache_dir=A_      ).read()
  _check_csv_dataset(A_	,			A_      )




def 			snake_case  (A_      :List[str]	,			A_      :Dict	,			A_      :Optional[Any]=("train",)      ):







  '''simple docstring'''

  assert isinstance(A_	,			A_      )
  for split in splits:
    a     :			List[Any]				=   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 			snake_case  (A_      :str	,			A_      :Any	,			A_      :List[str]      ):







  '''simple docstring'''

  a     :			Any				=   tmp_path / 'cache'
  a     :			int				=   {'col_1': 'int64', 'col_2': 'int64', 'col_3': 'float64'}
  with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
    a     :			str				=   CsvDatasetReader({'train': csv_path}	,			cache_dir=A_	,			keep_in_memory=A_      ).read()
  _check_csv_datasetdict(A_	,			A_      )




@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 			snake_case  (A_      :Optional[int]	,			A_      :List[Any]	,			A_      :Tuple      ):







  '''simple docstring'''

  a     :			Dict				=   tmp_path / 'cache'
  # CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
  a     :			Any				=   {'col_1': 'int64', 'col_2': 'int64', 'col_3': 'float64'}
  a     :			Tuple				=   features.copy() if features else default_expected_features
  a     :			Optional[int]				=   (
      Features({feature: Value(A_      ) for feature, dtype in features.items()}      ) if features is not None else None
  )
  a     :			Any				=   CsvDatasetReader({'train': csv_path}	,			features=A_	,			cache_dir=A_      ).read()
  _check_csv_datasetdict(A_	,			A_      )




@pytest.mark.parametrize('split'	,			[None, NamedSplit('train'      ), 'train', 'test']      )
def 			snake_case  (A_      :List[str]	,			A_      :Optional[Any]	,			A_      :str      ):







  '''simple docstring'''

  if split:
    a     :			List[Any]				=   {split: csv_path}
  else:
    a     :			str				=   {'train': csv_path, 'test': csv_path}
  a     :			Union[str, Any]				=   tmp_path / 'cache'
  a     :			Tuple				=   {'col_1': 'int64', 'col_2': 'int64', 'col_3': 'float64'}
  a     :			List[str]				=   CsvDatasetReader(A_	,			cache_dir=A_      ).read()
  _check_csv_datasetdict(A_	,			A_	,			splits=list(path.keys()      )      )
  assert all(dataset[split].split == split for split in path.keys()      )




def 			snake_case  (A_      :Dict      ):







  '''simple docstring'''

  with open(A_	,			encoding='utf-8'      ) as csvfile:
    yield from csv.reader(A_      )




def 			snake_case  (A_      :int	,			A_      :Optional[int]      ):







  '''simple docstring'''

  a     :			str				=   tmp_path / 'cache'
  a     :			List[str]				=   os.path.join(A_	,			'tmp.csv'      )
  a     :			Union[str, Any]				=   CsvDatasetReader({'train': csv_path}	,			cache_dir=A_      ).read()
  CsvDatasetWriter(dataset['train']	,			A_	,			num_proc=1      ).write()

  a     :			int				=   iter_csv_file(A_      )
  a     :			Optional[Any]				=   iter_csv_file(A_      )

  for rowa, rowa in zip(A_	,			A_      ):
    assert rowa == rowa




def 			snake_case  (A_      :Tuple	,			A_      :Tuple      ):







  '''simple docstring'''

  a     :			Union[str, Any]				=   tmp_path / 'cache'
  a     :			List[Any]				=   os.path.join(A_	,			'tmp.csv'      )
  a     :			Union[str, Any]				=   CsvDatasetReader({'train': csv_path}	,			cache_dir=A_      ).read()
  CsvDatasetWriter(dataset['train']	,			A_	,			num_proc=2      ).write()

  a     :			Dict				=   iter_csv_file(A_      )
  a     :			List[Any]				=   iter_csv_file(A_      )

  for rowa, rowa in zip(A_	,			A_      ):
    assert rowa == rowa






def 			snake_case  (A_      :Union[str, Any]	,			A_      :List[Any]      ):







  '''simple docstring'''

  a     :			List[Any]				=   tmp_path / 'cache'
  a     :			Dict				=   os.path.join(A_	,			'tmp.csv'      )
  a     :			Union[str, Any]				=   CsvDatasetReader({'train': csv_path}	,			cache_dir=A_      ).read()
  with pytest.raises(A_      ):
    CsvDatasetWriter(dataset['train']	,			A_	,			num_proc=0      )




