



"""simple docstring"""






import unittest

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

from diffusers import KandinskyVaaPriorPipeline, 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      A__  (    __SCREAMING_SNAKE_CASE    ,						unittest.TestCase				):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE        =      KandinskyVaaPriorPipeline
	SCREAMING_SNAKE_CASE        =      ['prompt']
	SCREAMING_SNAKE_CASE        =      ['prompt', 'negative_prompt']
	SCREAMING_SNAKE_CASE        =      [
	    'num_images_per_prompt',
	    'generator',
	    'num_inference_steps',
	    'latents',
	    'negative_prompt',
	    'guidance_scale',
	    'output_type',
	    'return_dict',
	]
	SCREAMING_SNAKE_CASE        =      False

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						int)							->			Optional[int]:

				"""simple docstring"""

				return 32

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						List[str])							->			Optional[int]:

				"""simple docstring"""

				return 32

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						Optional[Any])							->			Union[str, Any]:

				"""simple docstring"""

				return self.time_input_dim

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						str)							->			str:

				"""simple docstring"""

				return self.time_input_dim * 4

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						List[Any])							->			List[Any]:

				"""simple docstring"""

				return 100

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						Union[str, Any])							->			Any:

				"""simple docstring"""

				__lowerCAmelCase :    str				=			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
				return tokenizer

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						str)							->			List[str]:

				"""simple docstring"""

				torch.manual_seed(0)
				__lowerCAmelCase :    Optional[int]				=			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(_SCREAMING_SNAKE_CASE)

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						Tuple)							->			Tuple:

				"""simple docstring"""

				torch.manual_seed(0)

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

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

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						List[str])							->			Optional[Any]:

				"""simple docstring"""

				torch.manual_seed(0)
				__lowerCAmelCase :    Tuple				=			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			,     )

				__lowerCAmelCase :    str				=			CLIPVisionModelWithProjection(_SCREAMING_SNAKE_CASE)
				return model

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						int)							->			Union[str, Any]:

				"""simple docstring"""

				__lowerCAmelCase :    int				=			CLIPImageProcessor(
				    crop_size=224			,     do_center_crop=_SCREAMING_SNAKE_CASE			,     do_normalize=_SCREAMING_SNAKE_CASE			,     do_resize=_SCREAMING_SNAKE_CASE			,     image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]			,     image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]			,     resample=3			,     size=224			,     )

				return image_processor

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

				"""simple docstring"""

				__lowerCAmelCase :    Tuple				=			self.dummy_prior
				__lowerCAmelCase :    Optional[Any]				=			self.dummy_image_encoder
				__lowerCAmelCase :    List[Any]				=			self.dummy_text_encoder
				__lowerCAmelCase :    Union[str, Any]				=			self.dummy_tokenizer
				__lowerCAmelCase :    Optional[Any]				=			self.dummy_image_processor

				__lowerCAmelCase :    int				=			UnCLIPScheduler(
				    variance_type="fixed_small_log"			,     prediction_type="sample"			,     num_train_timesteps=1000			,     clip_sample=_SCREAMING_SNAKE_CASE			,     clip_sample_range=10.0			,     )

				__lowerCAmelCase :    Optional[int]				=			{
				    "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:						Optional[int]			,     _SCREAMING_SNAKE_CASE:						List[Any]			,     _SCREAMING_SNAKE_CASE:						Tuple=0)							->			Optional[int]:

				"""simple docstring"""

				if str(_SCREAMING_SNAKE_CASE).startswith("mps"):
							__lowerCAmelCase :    Optional[int]				=			torch.manual_seed(_SCREAMING_SNAKE_CASE)
				else:
							__lowerCAmelCase :    int				=			torch.Generator(device=_SCREAMING_SNAKE_CASE).manual_seed(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Tuple				=			{
				    "prompt": "horse",
				    "generator": generator,
				    "guidance_scale": 4.0,
				    "num_inference_steps": 2,
				    "output_type": "np",
				}
				return inputs

	def 				_SCREAMING_SNAKE_CASE  (	self:						List[str])							->			Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase :    Tuple				=			"cpu"

				__lowerCAmelCase :    List[str]				=			self.get_dummy_components()

				__lowerCAmelCase :    Dict				=			self.pipeline_class(**_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    int				=			pipe.to(_SCREAMING_SNAKE_CASE)

				pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Any				=			pipe(**self.get_dummy_inputs(_SCREAMING_SNAKE_CASE))
				__lowerCAmelCase :    List[str]				=			output.image_embeds

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

				__lowerCAmelCase :    List[Any]				=			image[0, -10:]
				__lowerCAmelCase :    str				=			image_from_tuple[0, -10:]

				assert image.shape == (1, 32)

				__lowerCAmelCase :    Any				=			np.array(
				    [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156])

				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:						Dict)							->			Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase :    Dict				=			torch_device == "cpu"
				__lowerCAmelCase :    Any				=			True
				__lowerCAmelCase :    Optional[int]				=			False

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







	@skip_mps
	def 				_SCREAMING_SNAKE_CASE  (	self:						int)							->			Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase :    str				=			torch_device == "cpu"
				__lowerCAmelCase :    Tuple				=			False

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




"""simple docstring"""






import collections
import gzip
import os
import urllib

import numpy
from tensorflow.python.framework import dtypes, random_seed
from tensorflow.python.platform import gfile
from tensorflow.python.util.deprecation import deprecated

__snake_case				:  Union[str, Any] =       collections.namedtuple('_Datasets', ['train', 'validation', 'test'])

# CVDF mirror of http://yann.lecun.com/exdb/mnist/
__snake_case				:  Any =       'https://storage.googleapis.com/cvdf-datasets/mnist/'


def      _lowercase							(    __snake_case      )		->  List[Any]:
			__lowerCAmelCase :    int				=			numpy.dtype(numpy.uintaa      ).newbyteorder(">"      )
			return numpy.frombuffer(bytestream.read(4      )   ,dtype=__snake_case      )[0]


@deprecated(__snake_case   ,"Please use tf.data to implement this functionality."      )
def      _lowercase							(    __snake_case      )		->  Union[str, Any]:
			print("Extracting"   ,f.name      )
			with gzip.GzipFile(fileobj=__snake_case      ) as bytestream:
						__lowerCAmelCase :    Union[str, Any]				=			_readaa(__snake_case      )
						if magic != 2_051:
									raise ValueError(
									    "Invalid magic number %d in MNIST image file: %s" % (magic, f.name)      )
						__lowerCAmelCase :    List[str]				=			_readaa(__snake_case      )
						__lowerCAmelCase :    Union[str, Any]				=			_readaa(__snake_case      )
						__lowerCAmelCase :    int				=			_readaa(__snake_case      )
						__lowerCAmelCase :    int				=			bytestream.read(rows * cols * num_images      )
						__lowerCAmelCase :    Optional[Any]				=			numpy.frombuffer(__snake_case   ,dtype=numpy.uinta      )
						__lowerCAmelCase :    str				=			data.reshape(__snake_case   ,__snake_case   ,__snake_case   ,1      )
						return data


@deprecated(__snake_case   ,"Please use tf.one_hot on tensors."      )
def      _lowercase							(    __snake_case   ,__snake_case      )		->  Any:
			__lowerCAmelCase :    Union[str, Any]				=			labels_dense.shape[0]
			__lowerCAmelCase :    Optional[int]				=			numpy.arange(__snake_case      ) * num_classes
			__lowerCAmelCase :    int				=			numpy.zeros((num_labels, num_classes)      )
			__lowerCAmelCase :    str				=			1
			return labels_one_hot


@deprecated(__snake_case   ,"Please use tf.data to implement this functionality."      )
def      _lowercase							(    __snake_case   ,__snake_case=False   ,__snake_case=10      )		->  str:
			print("Extracting"   ,f.name      )
			with gzip.GzipFile(fileobj=__snake_case      ) as bytestream:
						__lowerCAmelCase :    List[str]				=			_readaa(__snake_case      )
						if magic != 2_049:
									raise ValueError(
									    "Invalid magic number %d in MNIST label file: %s" % (magic, f.name)      )
						__lowerCAmelCase :    Union[str, Any]				=			_readaa(__snake_case      )
						__lowerCAmelCase :    Union[str, Any]				=			bytestream.read(__snake_case      )
						__lowerCAmelCase :    Dict				=			numpy.frombuffer(__snake_case   ,dtype=numpy.uinta      )
						if one_hot:
									return _dense_to_one_hot(__snake_case   ,__snake_case      )
						return labels



class      A__  :




	'''simple docstring'''





	@deprecated(
	    _SCREAMING_SNAKE_CASE			,     "Please use alternatives such as official/mnist/_DataSet.py"
	    " from tensorflow/models."			,     )
	def __init__(	self:						Dict			,     _SCREAMING_SNAKE_CASE:						Any			,     _SCREAMING_SNAKE_CASE:						str			,     _SCREAMING_SNAKE_CASE:						Any=False			,     _SCREAMING_SNAKE_CASE:						int=False			,     _SCREAMING_SNAKE_CASE:						str=dtypes.floataa			,     _SCREAMING_SNAKE_CASE:						Any=True			,     _SCREAMING_SNAKE_CASE:						Union[str, Any]=None			,     )							->			Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase		,      __lowerCAmelCase :    Optional[Any]				=			random_seed.get_seed(_SCREAMING_SNAKE_CASE)
				# If op level seed is not set, use whatever graph level seed is returned
				numpy.random.seed(seeda if seed is None else seeda)
				__lowerCAmelCase :    Optional[Any]				=			dtypes.as_dtype(_SCREAMING_SNAKE_CASE).base_dtype
				if dtype not in (dtypes.uinta, dtypes.floataa):
							raise TypeError("Invalid image dtype %r, expected uint8 or float32" % dtype)
				if fake_data:
							__lowerCAmelCase :    Tuple				=			1_0000
							__lowerCAmelCase :    Optional[Any]				=			one_hot
				else:
							assert (
							    images.shape[0] == labels.shape[0]
							), F"""images.shape: {images.shape} labels.shape: {labels.shape}"""
							__lowerCAmelCase :    Union[str, Any]				=			images.shape[0]

							# Convert shape from [num examples, rows, columns, depth]
							# to [num examples, rows*columns] (assuming depth == 1)
							if reshape:
										assert images.shape[3] == 1
										__lowerCAmelCase :    List[str]				=			images.reshape(
										    images.shape[0]			,     images.shape[1] * images.shape[2])
							if dtype == dtypes.floataa:
										# Convert from [0, 255] -> [0.0, 1.0].
										__lowerCAmelCase :    Dict				=			images.astype(numpy.floataa)
										__lowerCAmelCase :    int				=			numpy.multiply(_SCREAMING_SNAKE_CASE			,     1.0 / 255.0)
				__lowerCAmelCase :    Optional[Any]				=			images
				__lowerCAmelCase :    int				=			labels
				__lowerCAmelCase :    Optional[Any]				=			0
				__lowerCAmelCase :    Optional[int]				=			0

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						Dict)							->			Tuple:

				"""simple docstring"""

				return self._images

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						Dict)							->			int:

				"""simple docstring"""

				return self._labels

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						int)							->			str:

				"""simple docstring"""

				return self._num_examples

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						Optional[Any])							->			List[str]:

				"""simple docstring"""

				return self._epochs_completed







	def 				_SCREAMING_SNAKE_CASE  (	self:						str			,     _SCREAMING_SNAKE_CASE:						List[str]			,     _SCREAMING_SNAKE_CASE:						int=False			,     _SCREAMING_SNAKE_CASE:						List[str]=True)							->			int:

				"""simple docstring"""

				if fake_data:
							__lowerCAmelCase :    Dict				=			[1] * 784
							__lowerCAmelCase :    str				=			[1] + [0] * 9 if self.one_hot else 0
							return (
							    [fake_image for _ in range(_SCREAMING_SNAKE_CASE)],
							    [fake_label for _ in range(_SCREAMING_SNAKE_CASE)],
							)
				__lowerCAmelCase :    Tuple				=			self._index_in_epoch
				# Shuffle for the first epoch
				if self._epochs_completed == 0 and start == 0 and shuffle:
							__lowerCAmelCase :    Any				=			numpy.arange(self._num_examples)
							numpy.random.shuffle(_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase :    Union[str, Any]				=			self.images[perma]
							__lowerCAmelCase :    Dict				=			self.labels[perma]
				# Go to the next epoch
				if start + batch_size > self._num_examples:
							# Finished epoch
							self._epochs_completed += 1
							# Get the rest examples in this epoch
							__lowerCAmelCase :    Tuple				=			self._num_examples - start
							__lowerCAmelCase :    List[str]				=			self._images[start : self._num_examples]
							__lowerCAmelCase :    Tuple				=			self._labels[start : self._num_examples]
							# Shuffle the data
							if shuffle:
										__lowerCAmelCase :    Tuple				=			numpy.arange(self._num_examples)
										numpy.random.shuffle(_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase :    Optional[Any]				=			self.images[perm]
										__lowerCAmelCase :    Dict				=			self.labels[perm]
							# Start next epoch
							__lowerCAmelCase :    str				=			0
							__lowerCAmelCase :    Dict				=			batch_size - rest_num_examples
							__lowerCAmelCase :    str				=			self._index_in_epoch
							__lowerCAmelCase :    Optional[Any]				=			self._images[start:end]
							__lowerCAmelCase :    Optional[Any]				=			self._labels[start:end]
							return (
							    numpy.concatenate((images_rest_part, images_new_part)			,     axis=0),
							    numpy.concatenate((labels_rest_part, labels_new_part)			,     axis=0),
							)
				else:
							self._index_in_epoch += batch_size
							__lowerCAmelCase :    int				=			self._index_in_epoch
							return self._images[start:end], self._labels[start:end]






@deprecated(__snake_case   ,"Please write your own downloading logic."      )
def      _lowercase							(    __snake_case   ,__snake_case   ,__snake_case      )		->  Any:
			if not gfile.Exists(__snake_case      ):
						gfile.MakeDirs(__snake_case      )
			__lowerCAmelCase :    Tuple				=			os.path.join(__snake_case   ,__snake_case      )
			if not gfile.Exists(__snake_case      ):
						urllib.request.urlretrieve(__snake_case   ,__snake_case      )  # noqa: S310
						with gfile.GFile(__snake_case      ) as f:
									__lowerCAmelCase :    Union[str, Any]				=			f.size()
						print("Successfully downloaded"   ,__snake_case   ,__snake_case   ,"bytes."      )
			return filepath





@deprecated(
    __snake_case   ,"Please use alternatives such as:" " tensorflow_datasets.load('mnist')"      )
def      _lowercase							(    __snake_case   ,__snake_case=False   ,__snake_case=False   ,__snake_case=dtypes.floataa   ,__snake_case=True   ,__snake_case=5_000   ,__snake_case=None   ,__snake_case=DEFAULT_SOURCE_URL   ,)		->  Tuple:
			if fake_data:

						def fake():
									return _DataSet(
									    []   ,[]   ,fake_data=__snake_case   ,one_hot=__snake_case   ,dtype=__snake_case   ,seed=__snake_case      )

						__lowerCAmelCase :    Union[str, Any]				=			fake()
						__lowerCAmelCase :    Optional[Any]				=			fake()
						__lowerCAmelCase :    List[Any]				=			fake()
						return _Datasets(train=__snake_case   ,validation=__snake_case   ,test=__snake_case      )

			if not source_url:  # empty string check
						__lowerCAmelCase :    Optional[Any]				=			DEFAULT_SOURCE_URL

			__lowerCAmelCase :    Dict				=			"train-images-idx3-ubyte.gz"
			__lowerCAmelCase :    int				=			"train-labels-idx1-ubyte.gz"
			__lowerCAmelCase :    List[str]				=			"t10k-images-idx3-ubyte.gz"
			__lowerCAmelCase :    Any				=			"t10k-labels-idx1-ubyte.gz"

			__lowerCAmelCase :    Any				=			_maybe_download(
			    __snake_case   ,__snake_case   ,source_url + train_images_file      )
			with gfile.Open(__snake_case   ,"rb"      ) as f:
						__lowerCAmelCase :    Union[str, Any]				=			_extract_images(__snake_case      )

			__lowerCAmelCase :    Optional[int]				=			_maybe_download(
			    __snake_case   ,__snake_case   ,source_url + train_labels_file      )
			with gfile.Open(__snake_case   ,"rb"      ) as f:
						__lowerCAmelCase :    Optional[int]				=			_extract_labels(__snake_case   ,one_hot=__snake_case      )

			__lowerCAmelCase :    Optional[int]				=			_maybe_download(
			    __snake_case   ,__snake_case   ,source_url + test_images_file      )
			with gfile.Open(__snake_case   ,"rb"      ) as f:
						__lowerCAmelCase :    List[Any]				=			_extract_images(__snake_case      )

			__lowerCAmelCase :    str				=			_maybe_download(
			    __snake_case   ,__snake_case   ,source_url + test_labels_file      )
			with gfile.Open(__snake_case   ,"rb"      ) as f:
						__lowerCAmelCase :    List[Any]				=			_extract_labels(__snake_case   ,one_hot=__snake_case      )

			if not 0 <= validation_size <= len(__snake_case      ):
						__lowerCAmelCase :    Tuple				=			(
						    "Validation size should be between 0 and "
						    F"""{len(__snake_case      )}. Received: {validation_size}."""
						)
						raise ValueError(__snake_case      )

			__lowerCAmelCase :    Any				=			train_images[:validation_size]
			__lowerCAmelCase :    Any				=			train_labels[:validation_size]
			__lowerCAmelCase :    List[str]				=			train_images[validation_size:]
			__lowerCAmelCase :    Optional[Any]				=			train_labels[validation_size:]

			__lowerCAmelCase :    Dict				=			{"dtype": dtype, "reshape": reshape, "seed": seed}

			__lowerCAmelCase :    str				=			_DataSet(__snake_case   ,__snake_case   ,**__snake_case      )
			__lowerCAmelCase :    Dict				=			_DataSet(__snake_case   ,__snake_case   ,**__snake_case      )
			__lowerCAmelCase :    Union[str, Any]				=			_DataSet(__snake_case   ,__snake_case   ,**__snake_case      )

			return _Datasets(train=__snake_case   ,validation=__snake_case   ,test=__snake_case      )