






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


     UpperCamelCase_							:						str								=		KandinskyVaaPriorEmbaEmbPipeline
     UpperCamelCase_							:						List[Any]								=		['''prompt''', '''image''']
     UpperCamelCase_							:						List[Any]								=		['''prompt''', '''image''']
     UpperCamelCase_							:						List[Any]								=		[
         '''num_images_per_prompt''',
         '''strength''',
         '''generator''',
         '''num_inference_steps''',
         '''latents''',
         '''negative_prompt''',
         '''guidance_scale''',
         '''output_type''',
         '''return_dict''',
     ]
     UpperCamelCase_							:						Tuple								=		False


     @property
     def  _A				( self				:		Optional[int]					):
       return 32


     @property
     def  _A				( self				:		List[Any]					):
       return 32


     @property
     def  _A				( self				:		Optional[Any]					):
       return self.time_input_dim


     @property
     def  _A				( self				:		List[Any]					):
       return self.time_input_dim * 4


     @property
     def  _A				( self				:		Optional[Any]					):
       return 100


     @property
     def  _A				( self				:		Optional[Any]					):
       SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					)
       return tokenizer


     @property
     def  _A				( self				:		Optional[Any]					):
       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE				:					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-05    ,       num_attention_heads=4    ,       num_hidden_layers=5    ,       pad_token_id=1    ,       vocab_size=1000    ,       )
       return CLIPTextModelWithProjection(UpperCAmelCase_					)


     @property
     def  _A				( self				:		Union[str, Any]					):
       torch.manual_seed(0					)

       SCREAMING_SNAKE_CASE				:					Any		  =			{
           "num_attention_heads": 2,
           "attention_head_dim": 12,
           "embedding_dim": self.text_embedder_hidden_size,
           "num_layers": 1,
       }

       SCREAMING_SNAKE_CASE				:					List[Any]		  =			PriorTransformer(**UpperCAmelCase_					)
       # 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
       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			nn.Parameter(torch.ones(model.clip_std.shape					)					)
       return model


     @property
     def  _A				( self				:		str					):
       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE				:					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    ,       )

       SCREAMING_SNAKE_CASE				:					List[Any]		  =			CLIPVisionModelWithProjection(UpperCAmelCase_					)
       return model


     @property
     def  _A				( self				:		List[Any]					):
       SCREAMING_SNAKE_CASE				:					Any		  =			CLIPImageProcessor(
           crop_size=224    ,       do_center_crop=UpperCAmelCase_    ,       do_normalize=UpperCAmelCase_    ,       do_resize=UpperCAmelCase_    ,       image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073]    ,       image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711]    ,       resample=3    ,       size=224    ,       )

       return image_processor


     def  _A				( self				:		Optional[Any]					):
       SCREAMING_SNAKE_CASE				:					List[str]		  =			self.dummy_prior
       SCREAMING_SNAKE_CASE				:					Tuple		  =			self.dummy_image_encoder
       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			self.dummy_text_encoder
       SCREAMING_SNAKE_CASE				:					List[Any]		  =			self.dummy_tokenizer
       SCREAMING_SNAKE_CASE				:					List[str]		  =			self.dummy_image_processor

       SCREAMING_SNAKE_CASE				:					List[Any]		  =			UnCLIPScheduler(
           variance_type="fixed_small_log"    ,       prediction_type="sample"    ,       num_train_timesteps=1000    ,       clip_sample=UpperCAmelCase_    ,       clip_sample_range=10.0    ,       )

       SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			{
           "prior": prior,
           "image_encoder": image_encoder,
           "text_encoder": text_encoder,
           "tokenizer": tokenizer,
           "scheduler": scheduler,
           "image_processor": image_processor,
       }

       return components


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

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

       SCREAMING_SNAKE_CASE				:					Any		  =			{
           "prompt": "horse",
           "image": init_image,
           "strength": 0.5,
           "generator": generator,
           "guidance_scale": 4.0,
           "num_inference_steps": 2,
           "output_type": "np",
       }
       return inputs


     def  _A				( self				:		Union[str, Any]					):
       SCREAMING_SNAKE_CASE				:					int		  =			"cpu"

       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			self.get_dummy_components()

       SCREAMING_SNAKE_CASE				:					List[Any]		  =			self.pipeline_class(**UpperCAmelCase_					)
       SCREAMING_SNAKE_CASE				:					List[str]		  =			pipe.to(UpperCAmelCase_					)

       pipe.set_progress_bar_config(disable=UpperCAmelCase_					)

       SCREAMING_SNAKE_CASE				:					str		  =			pipe(**self.get_dummy_inputs(UpperCAmelCase_					)					)
       SCREAMING_SNAKE_CASE				:					str		  =			output.image_embeds

       SCREAMING_SNAKE_CASE				:					Any		  =			pipe(
           **self.get_dummy_inputs(UpperCAmelCase_					)    ,       return_dict=UpperCAmelCase_    ,       )[0]

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

       assert image.shape == (1, 32)

       SCREAMING_SNAKE_CASE				:					List[Any]		  =			np.array(
           [
               0.1_071_284,
               1.3_330_271,
               0.61_260_223,
               -0.6_691_065,
               -0.3_846_852,
               -1.0_303_661,
               0.22_716_111,
               0.03_348_901,
               0.30_040_675,
               -0.24_805_029,
           ]					)

       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  _A				( self				:		List[Any]					):
       SCREAMING_SNAKE_CASE				:					str		  =			torch_device == "cpu"
       SCREAMING_SNAKE_CASE				:					Dict		  =			True
       SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			False

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






     @skip_mps
     def  _A				( self				:		str					):
       SCREAMING_SNAKE_CASE				:					Any		  =			torch_device == "cpu"
       SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			False

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














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      lowerCamelCase__	(   lowercase  ,   lowercase ):

  """simple docstring"""






  assert isinstance(lowercase  ,   lowercase )
  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      lowerCamelCase__	(   lowercase  ,   lowercase  ,   lowercase ):

  """simple docstring"""






  SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			tmp_path / "cache"
  SCREAMING_SNAKE_CASE				:					Dict		  =			{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
  with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
    SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			CsvDatasetReader(lowercase  ,   cache_dir=lowercase  ,   keep_in_memory=lowercase ).read()
  _check_csv_dataset(lowercase  ,   lowercase )


@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      lowerCamelCase__	(   lowercase  ,   lowercase  ,   lowercase ):

  """simple docstring"""






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


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

  """simple docstring"""






  SCREAMING_SNAKE_CASE				:					Dict		  =			tmp_path / "cache"
  SCREAMING_SNAKE_CASE				:					Dict		  =			{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
  SCREAMING_SNAKE_CASE				:					List[Any]		  =			CsvDatasetReader(lowercase  ,   cache_dir=lowercase  ,   split=lowercase ).read()
  _check_csv_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				:					Tuple		  =			csv_path
  elif issubclass(lowercase  ,   lowercase ):
    SCREAMING_SNAKE_CASE				:					int		  =			[csv_path]
  SCREAMING_SNAKE_CASE				:					Any		  =			tmp_path / "cache"
  SCREAMING_SNAKE_CASE				:					Optional[int]		  =			{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
  SCREAMING_SNAKE_CASE				:					Tuple		  =			CsvDatasetReader(lowercase  ,   cache_dir=lowercase ).read()
  _check_csv_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 == 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      lowerCamelCase__	(   lowercase  ,   lowercase  ,   lowercase ):

  """simple docstring"""






  SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			tmp_path / "cache"
  SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
  with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
    SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			CsvDatasetReader({"train": csv_path}  ,   cache_dir=lowercase  ,   keep_in_memory=lowercase ).read()
  _check_csv_datasetdict(lowercase  ,   lowercase )


@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      lowerCamelCase__	(   lowercase  ,   lowercase  ,   lowercase ):

  """simple docstring"""






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


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

  """simple docstring"""






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


def      lowerCamelCase__	(   lowercase ):

  """simple docstring"""






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


def      lowerCamelCase__	(   lowercase  ,   lowercase ):

  """simple docstring"""






  SCREAMING_SNAKE_CASE				:					List[Any]		  =			tmp_path / "cache"
  SCREAMING_SNAKE_CASE				:					Optional[int]		  =			os.path.join(lowercase  ,   "tmp.csv" )
  SCREAMING_SNAKE_CASE				:					Optional[int]		  =			CsvDatasetReader({"train": csv_path}  ,   cache_dir=lowercase ).read()
  CsvDatasetWriter(dataset["train"]  ,   lowercase  ,   num_proc=1 ).write()

  SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			iter_csv_file(lowercase )
  SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			iter_csv_file(lowercase )

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


def      lowerCamelCase__	(   lowercase  ,   lowercase ):

  """simple docstring"""






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

  SCREAMING_SNAKE_CASE				:					Dict		  =			iter_csv_file(lowercase )
  SCREAMING_SNAKE_CASE				:					Tuple		  =			iter_csv_file(lowercase )

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


def      lowerCamelCase__	(   lowercase  ,   lowercase ):

  """simple docstring"""






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






