"""simple docstring"""






from typing import TYPE_CHECKING

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


lowerCAmelCase__  :  List[Any]											=     {
    'configuration_utils': ['GenerationConfig'],
    'streamers': ['TextIteratorStreamer', 'TextStreamer'],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCAmelCase__  :  Dict											=     [
						    'Constraint',
						    'ConstraintListState',
						    'DisjunctiveConstraint',
						    'PhrasalConstraint',
						]
						lowerCAmelCase__  :  Tuple											=     [
						    'BeamHypotheses',
						    'BeamScorer',
						    'BeamSearchScorer',
						    'ConstrainedBeamSearchScorer',
						]
						lowerCAmelCase__  :  int											=     [
						    'EpsilonLogitsWarper',
						    'EtaLogitsWarper',
						    'ForcedBOSTokenLogitsProcessor',
						    'ForcedEOSTokenLogitsProcessor',
						    'HammingDiversityLogitsProcessor',
						    'InfNanRemoveLogitsProcessor',
						    'LogitsProcessor',
						    'LogitsProcessorList',
						    'LogitsWarper',
						    'MinLengthLogitsProcessor',
						    'MinNewTokensLengthLogitsProcessor',
						    'NoBadWordsLogitsProcessor',
						    'NoRepeatNGramLogitsProcessor',
						    'PrefixConstrainedLogitsProcessor',
						    'RepetitionPenaltyLogitsProcessor',
						    'SequenceBiasLogitsProcessor',
						    'EncoderRepetitionPenaltyLogitsProcessor',
						    'TemperatureLogitsWarper',
						    'TopKLogitsWarper',
						    'TopPLogitsWarper',
						    'TypicalLogitsWarper',
						    'EncoderNoRepeatNGramLogitsProcessor',
						    'ExponentialDecayLengthPenalty',
						    'LogitNormalization',
						]
						lowerCAmelCase__  :  Dict											=     [
						    'MaxNewTokensCriteria',
						    'MaxLengthCriteria',
						    'MaxTimeCriteria',
						    'StoppingCriteria',
						    'StoppingCriteriaList',
						    'validate_stopping_criteria',
						]
						lowerCAmelCase__  :  Tuple											=     [
						    'GenerationMixin',
						    'top_k_top_p_filtering',
						    'GreedySearchEncoderDecoderOutput',
						    'GreedySearchDecoderOnlyOutput',
						    'SampleEncoderDecoderOutput',
						    'SampleDecoderOnlyOutput',
						    'BeamSearchEncoderDecoderOutput',
						    'BeamSearchDecoderOnlyOutput',
						    'BeamSampleEncoderDecoderOutput',
						    'BeamSampleDecoderOnlyOutput',
						    'ContrastiveSearchEncoderDecoderOutput',
						    'ContrastiveSearchDecoderOnlyOutput',
						]

try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCAmelCase__  :  Union[str, Any]											=     [
						    'TFForcedBOSTokenLogitsProcessor',
						    'TFForcedEOSTokenLogitsProcessor',
						    'TFLogitsProcessor',
						    'TFLogitsProcessorList',
						    'TFLogitsWarper',
						    'TFMinLengthLogitsProcessor',
						    'TFNoBadWordsLogitsProcessor',
						    'TFNoRepeatNGramLogitsProcessor',
						    'TFRepetitionPenaltyLogitsProcessor',
						    'TFTemperatureLogitsWarper',
						    'TFTopKLogitsWarper',
						    'TFTopPLogitsWarper',
						    'TFForceTokensLogitsProcessor',
						    'TFSuppressTokensAtBeginLogitsProcessor',
						    'TFSuppressTokensLogitsProcessor',
						]
						lowerCAmelCase__  :  Dict											=     [
						    'TFGenerationMixin',
						    'tf_top_k_top_p_filtering',
						    'TFGreedySearchDecoderOnlyOutput',
						    'TFGreedySearchEncoderDecoderOutput',
						    'TFSampleEncoderDecoderOutput',
						    'TFSampleDecoderOnlyOutput',
						    'TFBeamSearchEncoderDecoderOutput',
						    'TFBeamSearchDecoderOnlyOutput',
						    'TFBeamSampleEncoderDecoderOutput',
						    'TFBeamSampleDecoderOnlyOutput',
						    'TFContrastiveSearchEncoderDecoderOutput',
						    'TFContrastiveSearchDecoderOnlyOutput',
						]

try:
						if not is_flax_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCAmelCase__  :  Tuple											=     [
						    'FlaxForcedBOSTokenLogitsProcessor',
						    'FlaxForcedEOSTokenLogitsProcessor',
						    'FlaxLogitsProcessor',
						    'FlaxLogitsProcessorList',
						    'FlaxLogitsWarper',
						    'FlaxMinLengthLogitsProcessor',
						    'FlaxTemperatureLogitsWarper',
						    'FlaxTopKLogitsWarper',
						    'FlaxTopPLogitsWarper',
						]
						lowerCAmelCase__  :  Any											=     [
						    'FlaxGenerationMixin',
						    'FlaxGreedySearchOutput',
						    'FlaxSampleOutput',
						    'FlaxBeamSearchOutput',
						]

if TYPE_CHECKING:
						from .configuration_utils import GenerationConfig
						from .streamers import TextIteratorStreamer, TextStreamer

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .beam_constraints import Constraint, ConstraintListState, DisjunctiveConstraint, PhrasalConstraint
												from .beam_search import BeamHypotheses, BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
												from .logits_process import (
												    EncoderNoRepeatNGramLogitsProcessor,
												    EncoderRepetitionPenaltyLogitsProcessor,
												    EpsilonLogitsWarper,
												    EtaLogitsWarper,
												    ExponentialDecayLengthPenalty,
												    ForcedBOSTokenLogitsProcessor,
												    ForcedEOSTokenLogitsProcessor,
												    HammingDiversityLogitsProcessor,
												    InfNanRemoveLogitsProcessor,
												    LogitNormalization,
												    LogitsProcessor,
												    LogitsProcessorList,
												    LogitsWarper,
												    MinLengthLogitsProcessor,
												    MinNewTokensLengthLogitsProcessor,
												    NoBadWordsLogitsProcessor,
												    NoRepeatNGramLogitsProcessor,
												    PrefixConstrainedLogitsProcessor,
												    RepetitionPenaltyLogitsProcessor,
												    SequenceBiasLogitsProcessor,
												    TemperatureLogitsWarper,
												    TopKLogitsWarper,
												    TopPLogitsWarper,
												    TypicalLogitsWarper,
												)
												from .stopping_criteria import (
												    MaxLengthCriteria,
												    MaxNewTokensCriteria,
												    MaxTimeCriteria,
												    StoppingCriteria,
												    StoppingCriteriaList,
												    validate_stopping_criteria,
												)
												from .utils import (
												    BeamSampleDecoderOnlyOutput,
												    BeamSampleEncoderDecoderOutput,
												    BeamSearchDecoderOnlyOutput,
												    BeamSearchEncoderDecoderOutput,
												    ContrastiveSearchDecoderOnlyOutput,
												    ContrastiveSearchEncoderDecoderOutput,
												    GenerationMixin,
												    GreedySearchDecoderOnlyOutput,
												    GreedySearchEncoderDecoderOutput,
												    SampleDecoderOnlyOutput,
												    SampleEncoderDecoderOutput,
												    top_k_top_p_filtering,
												)

						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tf_logits_process import (
												    TFForcedBOSTokenLogitsProcessor,
												    TFForcedEOSTokenLogitsProcessor,
												    TFForceTokensLogitsProcessor,
												    TFLogitsProcessor,
												    TFLogitsProcessorList,
												    TFLogitsWarper,
												    TFMinLengthLogitsProcessor,
												    TFNoBadWordsLogitsProcessor,
												    TFNoRepeatNGramLogitsProcessor,
												    TFRepetitionPenaltyLogitsProcessor,
												    TFSuppressTokensAtBeginLogitsProcessor,
												    TFSuppressTokensLogitsProcessor,
												    TFTemperatureLogitsWarper,
												    TFTopKLogitsWarper,
												    TFTopPLogitsWarper,
												)
												from .tf_utils import (
												    TFBeamSampleDecoderOnlyOutput,
												    TFBeamSampleEncoderDecoderOutput,
												    TFBeamSearchDecoderOnlyOutput,
												    TFBeamSearchEncoderDecoderOutput,
												    TFContrastiveSearchDecoderOnlyOutput,
												    TFContrastiveSearchEncoderDecoderOutput,
												    TFGenerationMixin,
												    TFGreedySearchDecoderOnlyOutput,
												    TFGreedySearchEncoderDecoderOutput,
												    TFSampleDecoderOnlyOutput,
												    TFSampleEncoderDecoderOutput,
												    tf_top_k_top_p_filtering,
												)

						try:
												if not is_flax_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .flax_logits_process import (
												    FlaxForcedBOSTokenLogitsProcessor,
												    FlaxForcedEOSTokenLogitsProcessor,
												    FlaxLogitsProcessor,
												    FlaxLogitsProcessorList,
												    FlaxLogitsWarper,
												    FlaxMinLengthLogitsProcessor,
												    FlaxTemperatureLogitsWarper,
												    FlaxTopKLogitsWarper,
												    FlaxTopPLogitsWarper,
												)
												from .flax_utils import FlaxBeamSearchOutput, FlaxGenerationMixin, FlaxGreedySearchOutput, FlaxSampleOutput
else:
						import sys

						lowerCAmelCase__  :  List[Any]											=     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




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




					"""simple docstring"""

					snake_case__       =			KandinskyPriorPipeline
					snake_case__       =			["prompt"]
					snake_case__       =			["prompt", "negative_prompt"]
					snake_case__       =			[
					    "num_images_per_prompt",
					    "generator",
					    "num_inference_steps",
					    "latents",
					    "negative_prompt",
					    "guidance_scale",
					    "output_type",
					    "return_dict",
					]
					snake_case__       =			False







					@property
					def 		__lowerCAmelCase						(    self       :     Optional[int]      ):
									return 32







					@property
					def 		__lowerCAmelCase						(    self       :     int      ):
									return 32







					@property
					def 		__lowerCAmelCase						(    self       :     Tuple      ):
									return self.time_input_dim







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







					@property
					def 		__lowerCAmelCase						(    self       :     Optional[int]      ):
									return 100







					@property
					def 		__lowerCAmelCase						(    self       :     List[str]      ):
									UpperCAmelCase__					=							CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      )
									return tokenizer







					@property
					def 		__lowerCAmelCase						(    self       :     int      ):
									torch.manual_seed(0      )
									UpperCAmelCase__					=							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(lowerCamelCase__      )







					@property
					def 		__lowerCAmelCase						(    self       :     Tuple      ):
									torch.manual_seed(0      )

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

									UpperCAmelCase__					=							PriorTransformer(**lowerCamelCase__      )
									# 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
									UpperCAmelCase__					=							nn.Parameter(torch.ones(model.clip_std.shape      )      )
									return model







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

									UpperCAmelCase__					=							CLIPVisionModelWithProjection(lowerCamelCase__      )
									return model







					@property
					def 		__lowerCAmelCase						(    self       :     Dict      ):
									UpperCAmelCase__					=							CLIPImageProcessor(
									    crop_size=224   ,do_center_crop=lowerCamelCase__   ,do_normalize=lowerCamelCase__   ,do_resize=lowerCamelCase__   ,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=224   ,)

									return image_processor







					def 		__lowerCAmelCase						(    self       :     List[Any]      ):
									UpperCAmelCase__					=							self.dummy_prior
									UpperCAmelCase__					=							self.dummy_image_encoder
									UpperCAmelCase__					=							self.dummy_text_encoder
									UpperCAmelCase__					=							self.dummy_tokenizer
									UpperCAmelCase__					=							self.dummy_image_processor

									UpperCAmelCase__					=							UnCLIPScheduler(
									    variance_type='fixed_small_log'   ,prediction_type='sample'   ,num_train_timesteps=1_000   ,clip_sample=lowerCamelCase__   ,clip_sample_range=1_0.0   ,)

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

									return components







					def 		__lowerCAmelCase						(    self       :     Dict   ,lowerCamelCase__       :     Union[str, Any]   ,lowerCamelCase__       :     Optional[int]=0      ):
									if str(lowerCamelCase__      ).startswith('mps'      ):
													UpperCAmelCase__					=							torch.manual_seed(lowerCamelCase__      )
									else:
													UpperCAmelCase__					=							torch.Generator(device=lowerCamelCase__      ).manual_seed(lowerCamelCase__      )
									UpperCAmelCase__					=							{
									    'prompt': 'horse',
									    'generator': generator,
									    'guidance_scale': 4.0,
									    'num_inference_steps': 2,
									    'output_type': 'np',
									}
									return inputs







					def 		__lowerCAmelCase						(    self       :     List[str]      ):
									UpperCAmelCase__					=							'cpu'

									UpperCAmelCase__					=							self.get_dummy_components()

									UpperCAmelCase__					=							self.pipeline_class(**lowerCamelCase__      )
									UpperCAmelCase__					=							pipe.to(lowerCamelCase__      )

									pipe.set_progress_bar_config(disable=lowerCamelCase__      )

									UpperCAmelCase__					=							pipe(**self.get_dummy_inputs(lowerCamelCase__      )      )
									UpperCAmelCase__					=							output.image_embeds

									UpperCAmelCase__					=							pipe(
									    **self.get_dummy_inputs(lowerCamelCase__      )   ,return_dict=lowerCamelCase__   ,)[0]

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

									assert image.shape == (1, 32)

									UpperCAmelCase__					=							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 		__lowerCAmelCase						(    self       :     Optional[int]      ):
									UpperCAmelCase__					=							torch_device == 'cpu'
									UpperCAmelCase__					=							True
									UpperCAmelCase__					=							False

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







					@skip_mps
					def 		__lowerCAmelCase						(    self       :     Optional[int]      ):
									UpperCAmelCase__					=							torch_device == 'cpu'
									UpperCAmelCase__					=							False

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



