




'''simple docstring'''







import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()





class   lowerCamelCase_							( UpperCAmelCase_, unittest.TestCase							):




					'''simple docstring'''






					lowercase__					=				KandinskyPipeline
					lowercase__					=				[
					    """prompt""",
					    """image_embeds""",
					    """negative_image_embeds""",
					]
					lowercase__					=				["""prompt""", """negative_prompt""", """image_embeds""", """negative_image_embeds"""]
					lowercase__					=				[
					    """generator""",
					    """height""",
					    """width""",
					    """latents""",
					    """guidance_scale""",
					    """negative_prompt""",
					    """num_inference_steps""",
					    """return_dict""",
					    """guidance_scale""",
					    """num_images_per_prompt""",
					    """output_type""",
					    """return_dict""",
					]
					lowercase__					=				False






					@property
					def        __SCREAMING_SNAKE_CASE			(					self				):
							return 32






					@property
					def        __SCREAMING_SNAKE_CASE			(					self				):
							return 32






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






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






					@property
					def        __SCREAMING_SNAKE_CASE			(					self				):
							return 100






					@property
					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = XLMRobertaTokenizerFast.from_pretrained("YiYiXu/tiny-random-mclip-base"				)
							return tokenizer






					@property
					def        __SCREAMING_SNAKE_CASE			(					self				):
							torch.manual_seed(0				)
							_snake_case			 = MCLIPConfig(
							    numDims=self.cross_attention_dim,						transformerDimensions=self.text_embedder_hidden_size,						hidden_size=self.text_embedder_hidden_size,						intermediate_size=37,						num_attention_heads=4,						num_hidden_layers=5,						vocab_size=1_005,						)

							_snake_case			 = MultilingualCLIP(__UpperCamelCase				)
							_snake_case			 = text_encoder.eval()

							return text_encoder






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

							_snake_case			 = {
							    "in_channels": 4,
							    # Out channels is double in channels because predicts mean and variance
							    "out_channels": 8,
							    "addition_embed_type": "text_image",
							    "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
							    "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
							    "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
							    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
							    "layers_per_block": 1,
							    "encoder_hid_dim": self.text_embedder_hidden_size,
							    "encoder_hid_dim_type": "text_image_proj",
							    "cross_attention_dim": self.cross_attention_dim,
							    "attention_head_dim": 4,
							    "resnet_time_scale_shift": "scale_shift",
							    "class_embed_type": None,
							}

							_snake_case			 = UNetaDConditionModel(**__UpperCamelCase				)
							return model






					@property
					def        __SCREAMING_SNAKE_CASE			(					self				):
							return {
							    "block_out_channels": [32, 64],
							    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
							    "in_channels": 3,
							    "latent_channels": 4,
							    "layers_per_block": 1,
							    "norm_num_groups": 8,
							    "norm_type": "spatial",
							    "num_vq_embeddings": 12,
							    "out_channels": 3,
							    "up_block_types": [
							        "AttnUpDecoderBlock2D",
							        "UpDecoderBlock2D",
							    ],
							    "vq_embed_dim": 4,
							}






					@property
					def        __SCREAMING_SNAKE_CASE			(					self				):
							torch.manual_seed(0				)
							_snake_case			 = VQModel(**self.dummy_movq_kwargs				)
							return model






					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = self.dummy_text_encoder
							_snake_case			 = self.dummy_tokenizer
							_snake_case			 = self.dummy_unet
							_snake_case			 = self.dummy_movq

							_snake_case			 = DDIMScheduler(
							    num_train_timesteps=1_000,						beta_schedule="linear",						beta_start=0.0_00_85,						beta_end=0.0_12,						clip_sample=__UpperCamelCase,						set_alpha_to_one=__UpperCamelCase,						steps_offset=1,						prediction_type="epsilon",						thresholding=__UpperCamelCase,						)

							_snake_case			 = {
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "unet": unet,
							    "scheduler": scheduler,
							    "movq": movq,
							}
							return components






					def        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase,						__UpperCamelCase=0				):
							_snake_case			 = floats_tensor((1, self.cross_attention_dim),						rng=random.Random(__UpperCamelCase				)				).to(__UpperCamelCase				)
							_snake_case			 = floats_tensor((1, self.cross_attention_dim),						rng=random.Random(seed + 1				)				).to(__UpperCamelCase				)
							if str(__UpperCamelCase				).startswith("mps"				):
									_snake_case			 = torch.manual_seed(__UpperCamelCase				)
							else:
									_snake_case			 = torch.Generator(device=__UpperCamelCase				).manual_seed(__UpperCamelCase				)
							_snake_case			 = {
							    "prompt": "horse",
							    "image_embeds": image_embeds,
							    "negative_image_embeds": negative_image_embeds,
							    "generator": generator,
							    "height": 64,
							    "width": 64,
							    "guidance_scale": 4.0,
							    "num_inference_steps": 2,
							    "output_type": "np",
							}
							return inputs






					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = "cpu"

							_snake_case			 = self.get_dummy_components()

							_snake_case			 = self.pipeline_class(**__UpperCamelCase				)
							_snake_case			 = pipe.to(__UpperCamelCase				)

							pipe.set_progress_bar_config(disable=__UpperCamelCase				)

							_snake_case			 = pipe(**self.get_dummy_inputs(__UpperCamelCase				)				)
							_snake_case			 = output.images

							_snake_case			 = pipe(
							    **self.get_dummy_inputs(__UpperCamelCase				),						return_dict=__UpperCamelCase,						)[0]

							_snake_case			 = image[0, -3:, -3:, -1]
							_snake_case			 = image_from_tuple[0, -3:, -3:, -1]

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

							_snake_case			 = np.array(
							    [0.32_86_63, 1.0, 0.23_21_68_73, 1.0, 0.92_71_75_64, 0.4_63_90_46, 0.96_89_47_77, 0.31_71_33_78, 0.6_29_39_53]				)

							assert (
							    np.abs(image_slice.flatten() - expected_slice				).max() < 1E-2
							), f''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''
							assert (
							    np.abs(image_from_tuple_slice.flatten() - expected_slice				).max() < 1E-2
							), f''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''





@slow
@require_torch_gpu
class   lowerCamelCase_							( unittest.TestCase							):




					'''simple docstring'''






					def        __SCREAMING_SNAKE_CASE			(					self				):
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()






					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = load_numpy(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
							    "/kandinsky/kandinsky_text2img_cat_fp16.npy"				)

							_snake_case			 = KandinskyPriorPipeline.from_pretrained(
							    "kandinsky-community/kandinsky-2-1-prior",						torch_dtype=torch.floataa				)
							pipe_prior.to(__UpperCamelCase				)

							_snake_case			 = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1",						torch_dtype=torch.floataa				)
							_snake_case			 = pipeline.to(__UpperCamelCase				)
							pipeline.set_progress_bar_config(disable=__UpperCamelCase				)

							_snake_case			 = "red cat, 4k photo"

							_snake_case			 = torch.Generator(device="cuda"				).manual_seed(0				)
							_snake_case ,					_snake_case			 = pipe_prior(
							    __UpperCamelCase,						generator=__UpperCamelCase,						num_inference_steps=5,						negative_prompt="",						).to_tuple()

							_snake_case			 = torch.Generator(device="cuda"				).manual_seed(0				)
							_snake_case			 = pipeline(
							    __UpperCamelCase,						image_embeds=__UpperCamelCase,						negative_image_embeds=__UpperCamelCase,						generator=__UpperCamelCase,						num_inference_steps=100,						output_type="np",						)

							_snake_case			 = output.images[0]

							assert image.shape == (512, 512, 3)

							assert_mean_pixel_difference(__UpperCamelCase,						__UpperCamelCase				)











'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_lowerCamelCase							:		Any										=						{'configuration_pegasus': ['PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP', 'PegasusConfig']}

try:
						if not is_sentencepiece_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCamelCase							:		Optional[Any]										=						['PegasusTokenizer']

try:
						if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCamelCase							:		Optional[Any]										=						['PegasusTokenizerFast']

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCamelCase							:		List[str]										=						[
						    'PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST',
						    'PegasusForCausalLM',
						    'PegasusForConditionalGeneration',
						    'PegasusModel',
						    'PegasusPreTrainedModel',
						]

try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCamelCase							:		Union[str, Any]										=						[
						    'TFPegasusForConditionalGeneration',
						    'TFPegasusModel',
						    'TFPegasusPreTrainedModel',
						]

try:
						if not is_flax_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCamelCase							:		List[Any]										=						[
						    'FlaxPegasusForConditionalGeneration',
						    'FlaxPegasusModel',
						    'FlaxPegasusPreTrainedModel',
						]


if TYPE_CHECKING:
						from .configuration_pegasus import PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP, PegasusConfig

						try:
												if not is_sentencepiece_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_pegasus import PegasusTokenizer

						try:
												if not is_tokenizers_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_pegasus_fast import PegasusTokenizerFast

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_pegasus import (
												    PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST,
												    PegasusForCausalLM,
												    PegasusForConditionalGeneration,
												    PegasusModel,
												    PegasusPreTrainedModel,
												)

						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tf_pegasus import TFPegasusForConditionalGeneration, TFPegasusModel, TFPegasusPreTrainedModel

						try:
												if not is_flax_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_flax_pegasus import (
												    FlaxPegasusForConditionalGeneration,
												    FlaxPegasusModel,
												    FlaxPegasusPreTrainedModel,
												)

else:
						import sys

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





