

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 				UpperCAmelCase_		(			__lowercase				,       unittest.TestCase   ):
							"""simple docstring"""







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


							@property
							def 				__lowercase				(			self       )    ->    Optional[int]:
												return 3_2


							@property
							def 				__lowercase				(			self       )    ->    List[Any]:
												return 3_2


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


							@property
							def 				__lowercase				(			self       )    ->    int:
												return self.time_input_dim * 4


							@property
							def 				__lowercase				(			self       )    ->    Tuple:
												return 1_0_0


							@property
							def 				__lowercase				(			self       )    ->    str:
												_a	:    int         =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       )
												return tokenizer


							@property
							def 				__lowercase				(			self       )    ->    str:
												torch.manual_seed(0       )
												_a	:    Any         =		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-0_5       ,							num_attention_heads=4       ,							num_hidden_layers=5       ,							pad_token_id=1       ,							vocab_size=1_0_0_0       ,							)
												return CLIPTextModelWithProjection(_a       )


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

												_a	:    Dict         =		{
												    '''num_attention_heads''': 2,
												    '''attention_head_dim''': 1_2,
												    '''embedding_dim''': self.text_embedder_hidden_size,
												    '''num_layers''': 1,
												}

												_a	:    Any         =		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	:    Union[str, Any]         =		nn.Parameter(torch.ones(model.clip_std.shape       )       )
												return model


							@property
							def 				__lowercase				(			self       )    ->    List[Any]:
												torch.manual_seed(0       )
												_a	:    Tuple         =		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       ,							)

												_a	:    List[Any]         =		CLIPVisionModelWithProjection(_a       )
												return model


							@property
							def 				__lowercase				(			self       )    ->    Union[str, Any]:
												_a	:    Optional[Any]         =		CLIPImageProcessor(
												    crop_size=2_2_4       ,							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=2_2_4       ,							)

												return image_processor


							def 				__lowercase				(			self       )    ->    Optional[Any]:
												_a	:    Optional[Any]         =		self.dummy_prior
												_a	:    int         =		self.dummy_image_encoder
												_a	:    Optional[Any]         =		self.dummy_text_encoder
												_a	:    str         =		self.dummy_tokenizer
												_a	:    List[str]         =		self.dummy_image_processor

												_a	:    List[str]         =		UnCLIPScheduler(
												    variance_type='''fixed_small_log'''       ,							prediction_type='''sample'''       ,							num_train_timesteps=1_0_0_0       ,							clip_sample=_a       ,							clip_sample_range=10.0       ,							)

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

												return components


							def 				__lowercase				(			self       ,							_a       ,							_a=0       )    ->    Tuple:
												if str(_a       ).startswith('''mps'''       ):
																	_a	:    List[Any]         =		torch.manual_seed(_a       )
												else:
																	_a	:    Union[str, Any]         =		torch.Generator(device=_a       ).manual_seed(_a       )
												_a	:    List[Any]         =		{
												    '''prompt''': '''horse''',
												    '''generator''': generator,
												    '''guidance_scale''': 4.0,
												    '''num_inference_steps''': 2,
												    '''output_type''': '''np''',
												}
												return inputs


							def 				__lowercase				(			self       )    ->    Union[str, Any]:
												_a	:    Dict         =		'''cpu'''

												_a	:    Dict         =		self.get_dummy_components()

												_a	:    Tuple         =		self.pipeline_class(**_a       )
												_a	:    Any         =		pipe.to(_a       )

												pipe.set_progress_bar_config(disable=_a       )

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

												_a	:    Optional[Any]         =		pipe(
												    **self.get_dummy_inputs(_a       )       ,							return_dict=_a       ,							)[0]

												_a	:    str         =		image[0, -1_0:]
												_a	:    List[Any]         =		image_from_tuple[0, -1_0:]

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

												_a	:    int         =		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 				__lowercase				(			self       )    ->    Tuple:
												_a	:    List[Any]         =		torch_device == '''cpu'''
												_a	:    Optional[int]         =		True
												_a	:    Optional[int]         =		False

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




							@skip_mps
							def 				__lowercase				(			self       )    ->    List[str]:
												_a	:    List[str]         =		torch_device == '''cpu'''
												_a	:    int         =		False

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






from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


a__				=		logging.get_logger(__name__)

a__				=		{
    '''google/bit-50''': '''https://huggingface.co/google/bit-50/resolve/main/config.json''',
}

class 				UpperCAmelCase_		(			__lowercase				,       __lowercase   ):
							"""simple docstring"""







							UpperCAmelCase__       :						Any	    =    "bit"
							UpperCAmelCase__       :						Optional[int]	    =    ["preactivation", "bottleneck"]
							UpperCAmelCase__       :						Optional[Any]	    =    ["SAME", "VALID"]


							def __init__(			self       ,							_a=3       ,							_a=6_4       ,							_a=[2_5_6, 5_1_2, 1_0_2_4, 2_0_4_8]       ,							_a=[3, 4, 6, 3]       ,							_a="preactivation"       ,							_a="relu"       ,							_a=None       ,							_a=3_2       ,							_a=0.0       ,							_a=False       ,							_a=3_2       ,							_a=1       ,							_a=None       ,							_a=None       ,							**_a       ,							)    ->    Union[str, Any]:
												super().__init__(**_a       )
												if layer_type not in self.layer_types:
																	raise ValueError(F"""layer_type={layer_type} is not one of {','.join(self.layer_types       )}"""       )
												if global_padding is not None:
																	if global_padding.upper() in self.supported_padding:
																						_a	:    Any         =		global_padding.upper()
																	else:
																						raise ValueError(F"""Padding strategy {global_padding} not supported"""       )
												_a	:    Optional[int]         =		num_channels
												_a	:    List[Any]         =		embedding_size
												_a	:    Any         =		hidden_sizes
												_a	:    int         =		depths
												_a	:    Dict         =		layer_type
												_a	:    int         =		hidden_act
												_a	:    Optional[Any]         =		global_padding
												_a	:    Optional[Any]         =		num_groups
												_a	:    Union[str, Any]         =		drop_path_rate
												_a	:    Tuple         =		embedding_dynamic_padding
												_a	:    Union[str, Any]         =		output_stride
												_a	:    Any         =		width_factor

												_a	:    Any         =		['''stem'''] + [F"""stage{idx}""" for idx in range(1       ,							len(_a       ) + 1       )]
												_a      ,					_a	:    List[str]         =		get_aligned_output_features_output_indices(
												    out_features=_a       ,							out_indices=_a       ,							stage_names=self.stage_names       )



