



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        snake_case__    (A__ ,     unittest.TestCase		):







						"""simple docstring"""




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

						@property
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Union[str, Any]:
									"""simple docstring"""




									return 3_2

						@property
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Any:
									"""simple docstring"""




									return 3_2

						@property
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				str:
									"""simple docstring"""




									return self.time_input_dim

						@property
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				List[Any]:
									"""simple docstring"""




									return self.time_input_dim * 4

						@property
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				str:
									"""simple docstring"""




									return 1_0_0

						@property
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Optional[Any]:
									"""simple docstring"""




									a__       :						Union[str, Any]							=  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""  )
									return tokenizer

						@property
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Dict:
									"""simple docstring"""




									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-05 ,       num_attention_heads=4 ,       num_hidden_layers=5 ,       pad_token_id=1 ,       vocab_size=1_0_0_0 ,       )
									return CLIPTextModelWithProjection(__lowercase  )

						@property
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Dict:
									"""simple docstring"""




									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__       :						List[str]							=  PriorTransformer(**__lowercase  )
									# 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__       :						List[str]							=  nn.Parameter(torch.ones(model.clip_std.shape  )  )
									return model

						@property
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Optional[int]:
									"""simple docstring"""




									torch.manual_seed(0  )
									a__       :						Optional[int]							=  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__       :						Optional[int]							=  CLIPVisionModelWithProjection(__lowercase  )
									return model

						@property
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				List[str]:
									"""simple docstring"""




									a__       :						List[str]							=  CLIPImageProcessor(
									    crop_size=2_2_4 ,       do_center_crop=__lowercase ,       do_normalize=__lowercase ,       do_resize=__lowercase ,       image_mean=[0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3] ,       image_std=[0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1] ,       resample=3 ,       size=2_2_4 ,       )

									return image_processor

						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				List[str]:
									"""simple docstring"""




									a__       :						Dict							=  self.dummy_prior
									a__       :						Any							=  self.dummy_image_encoder
									a__       :						Any							=  self.dummy_text_encoder
									a__       :						Dict							=  self.dummy_tokenizer
									a__       :						int							=  self.dummy_image_processor

									a__       :						Tuple							=  UnCLIPScheduler(
									    variance_type="""fixed_small_log""" ,       prediction_type="""sample""" ,       num_train_timesteps=1_0_0_0 ,       clip_sample=__lowercase ,       clip_sample_range=1_0.0 ,       )

									a__       :						Any							=  {
									    """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 ,       __lowercase ,       __lowercase=0  ) ->				Any:
									"""simple docstring"""




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

						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Union[str, Any]:
									"""simple docstring"""




									a__       :						str							=  """cpu"""

									a__       :						Union[str, Any]							=  self.get_dummy_components()

									a__       :						Optional[Any]							=  self.pipeline_class(**__lowercase  )
									a__       :						Dict							=  pipe.to(__lowercase  )

									pipe.set_progress_bar_config(disable=__lowercase  )

									a__       :						Optional[int]							=  pipe(**self.get_dummy_inputs(__lowercase  )  )
									a__       :						Tuple							=  output.image_embeds

									a__       :						Dict							=  pipe(
									    **self.get_dummy_inputs(__lowercase  ) ,       return_dict=__lowercase ,       )[0]

									a__       :						Optional[int]							=  image[0, -1_0:]
									a__       :						Union[str, Any]							=  image_from_tuple[0, -1_0:]

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

									a__       :						Union[str, Any]							=  np.array(
									    [-0.0_5_3_2, 1.7_1_2_0, 0.3_6_5_6, -1.0_8_5_2, -0.8_9_4_6, -1.1_7_5_6, 0.4_3_4_8, 0.2_4_8_2, 0.5_1_4_6, -0.1_1_5_6]  )

									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:
									"""simple docstring"""




									a__       :						Tuple							=  torch_device == """cpu"""
									a__       :						Optional[Any]							=  True
									a__       :						str							=  False

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

						@skip_mps
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Dict:
									"""simple docstring"""




									a__       :						List[Any]							=  torch_device == """cpu"""
									a__       :						Union[str, Any]							=  False

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









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_torch_available


_lowercase		:				int      		={
    "configuration_longt5": ["LONGT5_PRETRAINED_CONFIG_ARCHIVE_MAP", "LongT5Config", "LongT5OnnxConfig"],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase		:				Optional[Any]      		=[
		    "LONGT5_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "LongT5EncoderModel",
		    "LongT5ForConditionalGeneration",
		    "LongT5Model",
		    "LongT5PreTrainedModel",
		]

try:
		if not is_flax_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase		:				List[Any]      		=[
		    "FlaxLongT5ForConditionalGeneration",
		    "FlaxLongT5Model",
		    "FlaxLongT5PreTrainedModel",
		]


if TYPE_CHECKING:
		from .configuration_longta import LONGT5_PRETRAINED_CONFIG_ARCHIVE_MAP, LongTaConfig, LongTaOnnxConfig

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_longta import (
				    LONGT5_PRETRAINED_MODEL_ARCHIVE_LIST,
				    LongTaEncoderModel,
				    LongTaForConditionalGeneration,
				    LongTaModel,
				    LongTaPreTrainedModel,
				)

		try:
				if not is_flax_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_flax_longta import (
				    FlaxLongTaForConditionalGeneration,
				    FlaxLongTaModel,
				    FlaxLongTaPreTrainedModel,
				)


else:
		import sys

		_lowercase		:				Any      		=_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




