






"""simple docstring"""
import gc
import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, KandinskyVaaPipeline, KandinskyVaaPriorPipeline, UNetaDConditionModel, VQModel
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     UpperCamelCase_ (_snake_case       ,		unittest.TestCase  ):
	__magic_name__			   =       KandinskyVaaPipeline
	__magic_name__			   =       [
	    """image_embeds""",
	    """negative_image_embeds""",
	]
	__magic_name__			   =       ["""image_embeds""", """negative_image_embeds"""]
	__magic_name__			   =       [
	    """generator""",
	    """height""",
	    """width""",
	    """latents""",
	    """guidance_scale""",
	    """num_inference_steps""",
	    """return_dict""",
	    """guidance_scale""",
	    """num_images_per_prompt""",
	    """output_type""",
	    """return_dict""",
	]
	__magic_name__			   =       False





	@property
	def   _SCREAMING_SNAKE_CASE   (   self      :				int		)		->    Optional[Any]:
				return 32





	@property
	def   _SCREAMING_SNAKE_CASE   (   self      :				Tuple		)		->    Any:
				return 32





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





	@property
	def   _SCREAMING_SNAKE_CASE   (   self      :				List[str]		)		->    Dict:
				return self.time_input_dim * 4





	@property
	def   _SCREAMING_SNAKE_CASE   (   self      :				Union[str, Any]		)		->    Optional[Any]:
				return 100





	@property
	def   _SCREAMING_SNAKE_CASE   (   self      :				List[str]		)		->    List[Any]:
				torch.manual_seed(0		)

				UpperCAmelCase_  :				int      =      {
				    "in_channels": 4,
				    # Out channels is double in channels because predicts mean and variance
				    "out_channels": 8,
				    "addition_embed_type": "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": "image_proj",
				    "cross_attention_dim": self.cross_attention_dim,
				    "attention_head_dim": 4,
				    "resnet_time_scale_shift": "scale_shift",
				    "class_embed_type": None,
				}

				UpperCAmelCase_  :				Optional[Any]      =      UNetaDConditionModel(**snake_case_		)
				return model





	@property
	def   _SCREAMING_SNAKE_CASE   (   self      :				Union[str, Any]		)		->    Any:
				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      :				Any		)		->    Dict:
				torch.manual_seed(0		)
				UpperCAmelCase_  :				Optional[Any]      =      VQModel(**self.dummy_movq_kwargs		)
				return model





	def   _SCREAMING_SNAKE_CASE   (   self      :				int		)		->    Tuple:
				UpperCAmelCase_  :				str      =      self.dummy_unet
				UpperCAmelCase_  :				Optional[Any]      =      self.dummy_movq

				UpperCAmelCase_  :				Optional[int]      =      DDIMScheduler(
				    num_train_timesteps=1_000			,			beta_schedule="linear"			,			beta_start=0.0_0_0_8_5			,			beta_end=0.0_1_2			,			clip_sample=snake_case_			,			set_alpha_to_one=snake_case_			,			steps_offset=1			,			prediction_type="epsilon"			,			thresholding=snake_case_			,			)

				UpperCAmelCase_  :				Optional[int]      =      {
				    "unet": unet,
				    "scheduler": scheduler,
				    "movq": movq,
				}
				return components





	def   _SCREAMING_SNAKE_CASE   (   self      :				Union[str, Any]			,			lowerCAmelCase_      :				Optional[int]			,			lowerCAmelCase_      :				Optional[int]=0		)		->    List[Any]:
				UpperCAmelCase_  :				Dict      =      floats_tensor((1, self.text_embedder_hidden_size)			,			rng=random.Random(snake_case_		)		).to(snake_case_		)
				UpperCAmelCase_  :				List[Any]      =      floats_tensor((1, self.text_embedder_hidden_size)			,			rng=random.Random(seed + 1		)		).to(
				    snake_case_		)
				if str(snake_case_		).startswith("mps"		):
							UpperCAmelCase_  :				Dict      =      torch.manual_seed(snake_case_		)
				else:
							UpperCAmelCase_  :				Optional[int]      =      torch.Generator(device=snake_case_		).manual_seed(snake_case_		)
				UpperCAmelCase_  :				Optional[int]      =      {
				    "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      :				Tuple		)		->    List[str]:
				UpperCAmelCase_  :				List[str]      =      "cpu"

				UpperCAmelCase_  :				str      =      self.get_dummy_components()

				UpperCAmelCase_  :				str      =      self.pipeline_class(**snake_case_		)
				UpperCAmelCase_  :				str      =      pipe.to(snake_case_		)

				pipe.set_progress_bar_config(disable=snake_case_		)

				UpperCAmelCase_  :				Any      =      pipe(**self.get_dummy_inputs(snake_case_		)		)
				UpperCAmelCase_  :				Any      =      output.images

				UpperCAmelCase_  :				Any      =      pipe(
				    **self.get_dummy_inputs(snake_case_		)			,			return_dict=snake_case_			,			)[0]

				UpperCAmelCase_  :				List[Any]      =      image[0, -3:, -3:, -1]
				UpperCAmelCase_  :				int      =      image_from_tuple[0, -3:, -3:, -1]

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

				UpperCAmelCase_  :				int      =      np.array(
				    [0.6_2_3_7_9_7_6, 1.0, 0.3_6_4_4_1_3_3_2, 1.0, 0.7_0_6_3_9_6_3_4, 0.2_9_8_7_7_1_8_6, 0.8_5_6_5_2_1_2_5, 0.5_2_1_6_8_4_3, 0.5_4_4_5_4_0_4_6]		)

				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     UpperCamelCase_ (unittest.TestCase  ):





	def   _SCREAMING_SNAKE_CASE   (   self      :				List[str]		)		->    Optional[int]:
				# clean up the VRAM after each test
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()





	def   _SCREAMING_SNAKE_CASE   (   self      :				Dict		)		->    List[Any]:
				UpperCAmelCase_  :				List[str]      =      load_numpy(
				    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
				    "/kandinskyv22/kandinskyv22_text2img_cat_fp16.npy"		)

				UpperCAmelCase_  :				Dict      =      KandinskyVaaPriorPipeline.from_pretrained(
				    "kandinsky-community/kandinsky-2-2-prior"			,			torch_dtype=torch.floataa		)
				pipe_prior.to(snake_case_		)

				UpperCAmelCase_  :				Dict      =      KandinskyVaaPipeline.from_pretrained(
				    "kandinsky-community/kandinsky-2-2-decoder"			,			torch_dtype=torch.floataa		)
				UpperCAmelCase_  :				Dict      =      pipeline.to(snake_case_		)
				pipeline.set_progress_bar_config(disable=snake_case_		)

				UpperCAmelCase_  :				Tuple      =      "red cat, 4k photo"

				UpperCAmelCase_  :				Optional[Any]      =      torch.Generator(device="cuda"		).manual_seed(0		)
				UpperCAmelCase_  :				Optional[int]      =      pipe_prior(
				    snake_case_			,			generator=snake_case_			,			num_inference_steps=5			,			negative_prompt=""			,			).to_tuple()

				UpperCAmelCase_  :				Any      =      torch.Generator(device="cuda"		).manual_seed(0		)
				UpperCAmelCase_  :				Any      =      pipeline(
				    image_embeds=snake_case_			,			negative_image_embeds=snake_case_			,			generator=snake_case_			,			num_inference_steps=100			,			output_type="np"			,			)

				UpperCAmelCase_  :				Optional[int]      =      output.images[0]

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

				assert_mean_pixel_difference(snake_case_			,			snake_case_		)











"""simple docstring"""
from typing import TYPE_CHECKING

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


lowerCamelCase_       			=   {
    '''configuration_funnel''': ['''FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''FunnelConfig'''],
    '''convert_funnel_original_tf_checkpoint_to_pytorch''': [],
    '''tokenization_funnel''': ['''FunnelTokenizer'''],
}

try:
			if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   ['''FunnelTokenizerFast''']

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   [
			    '''FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST''',
			    '''FunnelBaseModel''',
			    '''FunnelForMaskedLM''',
			    '''FunnelForMultipleChoice''',
			    '''FunnelForPreTraining''',
			    '''FunnelForQuestionAnswering''',
			    '''FunnelForSequenceClassification''',
			    '''FunnelForTokenClassification''',
			    '''FunnelModel''',
			    '''FunnelPreTrainedModel''',
			    '''load_tf_weights_in_funnel''',
			]

try:
			if not is_tf_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   [
			    '''TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST''',
			    '''TFFunnelBaseModel''',
			    '''TFFunnelForMaskedLM''',
			    '''TFFunnelForMultipleChoice''',
			    '''TFFunnelForPreTraining''',
			    '''TFFunnelForQuestionAnswering''',
			    '''TFFunnelForSequenceClassification''',
			    '''TFFunnelForTokenClassification''',
			    '''TFFunnelModel''',
			    '''TFFunnelPreTrainedModel''',
			]


if TYPE_CHECKING:
			from .configuration_funnel import FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP, FunnelConfig
			from .tokenization_funnel import FunnelTokenizer

			try:
						if not is_tokenizers_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_funnel_fast import FunnelTokenizerFast

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_funnel import (
						    FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
						    FunnelBaseModel,
						    FunnelForMaskedLM,
						    FunnelForMultipleChoice,
						    FunnelForPreTraining,
						    FunnelForQuestionAnswering,
						    FunnelForSequenceClassification,
						    FunnelForTokenClassification,
						    FunnelModel,
						    FunnelPreTrainedModel,
						    load_tf_weights_in_funnel,
						)

			try:
						if not is_tf_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_tf_funnel import (
						    TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TFFunnelBaseModel,
						    TFFunnelForMaskedLM,
						    TFFunnelForMultipleChoice,
						    TFFunnelForPreTraining,
						    TFFunnelForQuestionAnswering,
						    TFFunnelForSequenceClassification,
						    TFFunnelForTokenClassification,
						    TFFunnelModel,
						    TFFunnelPreTrainedModel,
						)

else:
			import sys

			lowerCamelCase_       			=   _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



