





'''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 		__UpperCamelCase		(       _UpperCAmelCase,    _UpperCAmelCase   ):
		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 		__UpperCamelCase		(       _UpperCAmelCase,    _UpperCAmelCase,    _UpperCAmelCase   ):
		__UpperCAmelCase     :				int					   =			tmp_path / "cache"
		__UpperCAmelCase     :				Optional[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():
				__UpperCAmelCase     :				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 		__UpperCamelCase		(       _UpperCAmelCase,    _UpperCAmelCase,    _UpperCAmelCase   ):
		__UpperCAmelCase     :				str					   =			tmp_path / "cache"
		# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
		__UpperCAmelCase     :				str					   =			{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
		__UpperCAmelCase     :				List[str]					   =			features.copy() if features else default_expected_features
		__UpperCAmelCase     :				Union[str, Any]					   =			(
		    Features({feature: Value(A__   ) for feature, dtype in features.items()}   ) if features is not None else None
		)
		__UpperCAmelCase     :				Tuple					   =			CsvDatasetReader(A__,    features=A__,    cache_dir=A__   ).read()
		_check_csv_dataset(A__,    A__   )


@pytest.mark.parametrize("split",    [None, NamedSplit("train"   ), "train", "test"]   )
def 		__UpperCamelCase		(       _UpperCAmelCase,    _UpperCAmelCase,    _UpperCAmelCase   ):
		__UpperCAmelCase     :				List[str]					   =			tmp_path / "cache"
		__UpperCAmelCase     :				Dict					   =			{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
		__UpperCAmelCase     :				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 		__UpperCamelCase		(       _UpperCAmelCase,    _UpperCAmelCase,    _UpperCAmelCase   ):

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


def 		__UpperCamelCase		(       _UpperCAmelCase,    _UpperCAmelCase,    _UpperCAmelCase=("train",)   ):
		assert isinstance(A__,    A__   )
		for split in splits:
				__UpperCAmelCase     :				Tuple					   =			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 		__UpperCamelCase		(       _UpperCAmelCase,    _UpperCAmelCase,    _UpperCAmelCase   ):
		__UpperCAmelCase     :				Dict					   =			tmp_path / "cache"
		__UpperCAmelCase     :				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():
				__UpperCAmelCase     :				Optional[Any]					   =			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 		__UpperCamelCase		(       _UpperCAmelCase,    _UpperCAmelCase,    _UpperCAmelCase   ):
		__UpperCAmelCase     :				Optional[int]					   =			tmp_path / "cache"
		# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
		__UpperCAmelCase     :				Tuple					   =			{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
		__UpperCAmelCase     :				Tuple					   =			features.copy() if features else default_expected_features
		__UpperCAmelCase     :				Optional[int]					   =			(
		    Features({feature: Value(A__   ) for feature, dtype in features.items()}   ) if features is not None else None
		)
		__UpperCAmelCase     :				List[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 		__UpperCamelCase		(       _UpperCAmelCase,    _UpperCAmelCase,    _UpperCAmelCase   ):

		if split:
				__UpperCAmelCase     :				int					   =			{split: csv_path}
		else:
				__UpperCAmelCase     :				List[str]					   =			{"train": csv_path, "test": csv_path}
		__UpperCAmelCase     :				Optional[int]					   =			tmp_path / "cache"
		__UpperCAmelCase     :				Optional[Any]					   =			{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
		__UpperCAmelCase     :				Any					   =			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 		__UpperCamelCase		(       _UpperCAmelCase   ):

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


def 		__UpperCamelCase		(       _UpperCAmelCase,    _UpperCAmelCase   ):
		__UpperCAmelCase     :				Any					   =			tmp_path / "cache"
		__UpperCAmelCase     :				int					   =			os.path.join(A__,    "tmp.csv"   )
		__UpperCAmelCase     :				Union[str, Any]					   =			CsvDatasetReader({"train": csv_path},    cache_dir=A__   ).read()
		CsvDatasetWriter(dataset["train"],    A__,    num_proc=1   ).write()

		__UpperCAmelCase     :				Any					   =			iter_csv_file(A__   )
		__UpperCAmelCase     :				List[str]					   =			iter_csv_file(A__   )

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


def 		__UpperCamelCase		(       _UpperCAmelCase,    _UpperCAmelCase   ):
		__UpperCAmelCase     :				List[str]					   =			tmp_path / "cache"
		__UpperCAmelCase     :				int					   =			os.path.join(A__,    "tmp.csv"   )
		__UpperCAmelCase     :				str					   =			CsvDatasetReader({"train": csv_path},    cache_dir=A__   ).read()
		CsvDatasetWriter(dataset["train"],    A__,    num_proc=2   ).write()

		__UpperCAmelCase     :				Optional[Any]					   =			iter_csv_file(A__   )
		__UpperCAmelCase     :				List[str]					   =			iter_csv_file(A__   )

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






def 		__UpperCamelCase		(       _UpperCAmelCase,    _UpperCAmelCase   ):
		__UpperCAmelCase     :				str					   =			tmp_path / "cache"
		__UpperCAmelCase     :				Tuple					   =			os.path.join(A__,    "tmp.csv"   )
		__UpperCAmelCase     :				Tuple					   =			CsvDatasetReader({"train": csv_path},    cache_dir=A__   ).read()
		with pytest.raises(A__   ):
				CsvDatasetWriter(dataset["train"],    A__,    num_proc=0   )






import unittest

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

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

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class      lowerCamelCase		(  A_							, unittest.TestCase  ):
     UpperCAmelCase__				:							Tuple           =		KandinskyPriorPipeline
     UpperCAmelCase__				:							List[str]           =		["prompt"]
     UpperCAmelCase__				:							List[str]           =		["prompt", "negative_prompt"]
     UpperCAmelCase__				:							List[str]           =		[
         "num_images_per_prompt",
         "generator",
         "num_inference_steps",
         "latents",
         "negative_prompt",
         "guidance_scale",
         "output_type",
         "return_dict",
     ]
     UpperCAmelCase__				:							str           =		False





     @property
     def 		UpperCAmelCase(self      :    List[str]   )						->						Union[str, Any]:
         return 3_2





     @property
     def 		UpperCAmelCase(self      :    Any   )						->						Any:
         return 3_2





     @property
     def 		UpperCAmelCase(self      :    int   )						->						Dict:
         return self.time_input_dim





     @property
     def 		UpperCAmelCase(self      :    str   )						->						Tuple:
         return self.time_input_dim * 4





     @property
     def 		UpperCAmelCase(self      :    Dict   )						->						Dict:
         return 1_0_0





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





     @property
     def 		UpperCAmelCase(self      :    Optional[Any]   )						->						Optional[Any]:
         torch.manual_seed(0   )
         snake_case      =					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=3_7				,    layer_norm_eps=1E-05				,    num_attention_heads=4				,    num_hidden_layers=5				,    pad_token_id=1				,    vocab_size=1_0_0_0				,    )
         return CLIPTextModelWithProjection(_A   )





     @property
     def 		UpperCAmelCase(self      :    List[Any]   )						->						Tuple:
         torch.manual_seed(0   )

         snake_case      =					{
             "num_attention_heads": 2,
             "attention_head_dim": 1_2,
             "embedding_dim": self.text_embedder_hidden_size,
             "num_layers": 1,
         }

         snake_case      =					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
         snake_case      =					nn.Parameter(torch.ones(model.clip_std.shape   )   )
         return model





     @property
     def 		UpperCAmelCase(self      :    str   )						->						Union[str, Any]:
         torch.manual_seed(0   )
         snake_case      =					CLIPVisionConfig(
             hidden_size=self.text_embedder_hidden_size				,    image_size=2_2_4				,    projection_dim=self.text_embedder_hidden_size				,    intermediate_size=3_7				,    num_attention_heads=4				,    num_channels=3				,    num_hidden_layers=5				,    patch_size=1_4				,    )

         snake_case      =					CLIPVisionModelWithProjection(_A   )
         return model





     @property
     def 		UpperCAmelCase(self      :    List[Any]   )						->						int:
         snake_case      =					CLIPImageProcessor(
             crop_size=2_2_4				,    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=2_2_4				,    )

         return image_processor





     def 		UpperCAmelCase(self      :    int   )						->						Union[str, Any]:
         snake_case      =					self.dummy_prior
         snake_case      =					self.dummy_image_encoder
         snake_case      =					self.dummy_text_encoder
         snake_case      =					self.dummy_tokenizer
         snake_case      =					self.dummy_image_processor

         snake_case      =					UnCLIPScheduler(
             variance_type="fixed_small_log"				,    prediction_type="sample"				,    num_train_timesteps=1_0_0_0				,    clip_sample=_A				,    clip_sample_range=10.0				,    )

         snake_case      =					{
             "prior": prior,
             "image_encoder": image_encoder,
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "scheduler": scheduler,
             "image_processor": image_processor,
         }

         return components





     def 		UpperCAmelCase(self      :    Optional[int]				,    _A      :    Tuple				,    _A      :    str=0   )						->						Union[str, Any]:
         if str(_A   ).startswith("mps"   ):
             snake_case      =					torch.manual_seed(_A   )
         else:
             snake_case      =					torch.Generator(device=_A   ).manual_seed(_A   )
         snake_case      =					{
             "prompt": "horse",
             "generator": generator,
             "guidance_scale": 4.0,
             "num_inference_steps": 2,
             "output_type": "np",
         }
         return inputs





     def 		UpperCAmelCase(self      :    List[Any]   )						->						int:
         snake_case      =					"cpu"

         snake_case      =					self.get_dummy_components()

         snake_case      =					self.pipeline_class(**_A   )
         snake_case      =					pipe.to(_A   )

         pipe.set_progress_bar_config(disable=_A   )

         snake_case      =					pipe(**self.get_dummy_inputs(_A   )   )
         snake_case      =					output.image_embeds

         snake_case      =					pipe(
             **self.get_dummy_inputs(_A   )				,    return_dict=_A				,    )[0]

         snake_case      =					image[0, -1_0:]
         snake_case      =					image_from_tuple[0, -1_0:]

         assert image.shape == (1, 3_2)

         snake_case      =					np.array(
             [-0.05_32, 1.71_20, 0.36_56, -1.08_52, -0.89_46, -1.17_56, 0.43_48, 0.24_82, 0.51_46, -0.11_56]   )

         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 		UpperCAmelCase(self      :    List[Any]   )						->						str:
         snake_case      =					torch_device == "cpu"
         snake_case      =					True
         snake_case      =					False

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





     @skip_mps
     def 		UpperCAmelCase(self      :    Optional[int]   )						->						List[str]:
         snake_case      =					torch_device == "cpu"
         snake_case      =					False

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

