


"""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 					__lowerCAmelCase					(   __SCREAMING_SNAKE_CASE     ,			unittest.TestCase       ):


		'''simple docstring'''



		__UpperCAmelCase		:			List[str]										=				KandinskyVaaPriorPipeline
		__UpperCAmelCase		:			Optional[int]										=				['prompt']
		__UpperCAmelCase		:			int										=				['prompt', 'negative_prompt']
		__UpperCAmelCase		:			Optional[Any]										=				[
		    'num_images_per_prompt',
		    'generator',
		    'num_inference_steps',
		    'latents',
		    'negative_prompt',
		    'guidance_scale',
		    'output_type',
		    'return_dict',
		]
		__UpperCAmelCase		:			Any										=				False

		@property
		def 			__UpperCAmelCase       (  self      ):
				return 32

		@property
		def 			__UpperCAmelCase       (  self      ):
				return 32

		@property
		def 			__UpperCAmelCase       (  self      ):
				return self.time_input_dim

		@property
		def 			__UpperCAmelCase       (  self      ):
				return self.time_input_dim * 4

		@property
		def 			__UpperCAmelCase       (  self      ):
				return 100

		@property
		def 			__UpperCAmelCase       (  self      ):
				__a           =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      )
				return tokenizer

		@property
		def 			__UpperCAmelCase       (  self      ):
				torch.manual_seed(0      )
				__a           =			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=1_000							, )
				return CLIPTextModelWithProjection(_a      )

		@property
		def 			__UpperCAmelCase       (  self      ):
				torch.manual_seed(0      )

				__a           =			{
				    '''num_attention_heads''': 2,
				    '''attention_head_dim''': 12,
				    '''embedding_dim''': self.text_embedder_hidden_size,
				    '''num_layers''': 1,
				}

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

		@property
		def 			__UpperCAmelCase       (  self      ):
				torch.manual_seed(0      )
				__a           =			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							, )

				__a           =			CLIPVisionModelWithProjection(_a      )
				return model

		@property
		def 			__UpperCAmelCase       (  self      ):
				__a           =			CLIPImageProcessor(
				    crop_size=224							, do_center_crop=_a							, do_normalize=_a							, do_resize=_a							, 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 			__UpperCAmelCase       (  self      ):
				__a           =			self.dummy_prior
				__a           =			self.dummy_image_encoder
				__a           =			self.dummy_text_encoder
				__a           =			self.dummy_tokenizer
				__a           =			self.dummy_image_processor

				__a           =			UnCLIPScheduler(
				    variance_type='''fixed_small_log'''							, prediction_type='''sample'''							, num_train_timesteps=1_000							, clip_sample=_a							, clip_sample_range=10.0							, )

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

				return components

		def 			__UpperCAmelCase       (  self							, _a							, _a=0      ):
				if str(_a      ).startswith('''mps'''      ):
						__a           =			torch.manual_seed(_a      )
				else:
						__a           =			torch.Generator(device=_a      ).manual_seed(_a      )
				__a           =			{
				    '''prompt''': '''horse''',
				    '''generator''': generator,
				    '''guidance_scale''': 4.0,
				    '''num_inference_steps''': 2,
				    '''output_type''': '''np''',
				}
				return inputs

		def 			__UpperCAmelCase       (  self      ):
				__a           =			'''cpu'''

				__a           =			self.get_dummy_components()

				__a           =			self.pipeline_class(**_a      )
				__a           =			pipe.to(_a      )

				pipe.set_progress_bar_config(disable=_a      )

				__a           =			pipe(**self.get_dummy_inputs(_a      )      )
				__a           =			output.image_embeds

				__a           =			pipe(
				    **self.get_dummy_inputs(_a      )							, return_dict=_a							, )[0]

				__a           =			image[0, -10:]
				__a           =			image_from_tuple[0, -10:]

				assert image.shape == (1, 32)

				__a           =			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 			__UpperCAmelCase       (  self      ):
				__a           =			torch_device == '''cpu'''
				__a           =			True
				__a           =			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      ):
				__a           =			torch_device == '''cpu'''
				__a           =			False

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









from .imports import is_tqdm_available


if is_tqdm_available():
  from tqdm.auto import tqdm as _tqdm

from ..state import PartialState
def      a      (			A__	:   bool = True      ,   *A__	:   int      ,   **A__	:   Union[str, Any]     )    ->      List[str]:





      """simple docstring"""



      if not is_tqdm_available():
            raise ImportError('Accelerate\'s `tqdm` module requires `tqdm` to be installed. Please run `pip install tqdm`.'     )
      _lowercase	       =False
      if main_process_only:
            _lowercase	       =PartialState().local_process_index == 0
      return _tqdm(*A__      ,   **A__      ,   disable=A__     )


