




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





				'''simple docstring'''



				__a							: List[str]										=KandinskyVaaPipeline
				__a							: Optional[Any]										=[
				    """image_embeds""",
				    """negative_image_embeds""",
				]
				__a							: Optional[Any]										=["""image_embeds""", """negative_image_embeds"""]
				__a							: Any										=[
				    """generator""",
				    """height""",
				    """width""",
				    """latents""",
				    """guidance_scale""",
				    """num_inference_steps""",
				    """return_dict""",
				    """guidance_scale""",
				    """num_images_per_prompt""",
				    """output_type""",
				    """return_dict""",
				]
				__a							: Dict										=False



				@property
				def       __snake_case		( self						):
											return 32



				@property
				def       __snake_case		( self						):
											return 32



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



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



				@property
				def       __snake_case		( self						):
											return 1_00



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

											lowerCAmelCase   =		{
											    '''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,
											}

											lowerCAmelCase   =		UNetaDConditionModel(**UpperCAmelCase_						)
											return model



				@property
				def       __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       __snake_case		( self						):
											torch.manual_seed(0						)
											lowerCAmelCase   =		VQModel(**self.dummy_movq_kwargs						)
											return model



				def       __snake_case		( self						):
											lowerCAmelCase   =		self.dummy_unet
											lowerCAmelCase   =		self.dummy_movq

											lowerCAmelCase   =		DDIMScheduler(
											    num_train_timesteps=10_00       , beta_schedule='''linear'''       , beta_start=0.00085       , beta_end=0.012       , clip_sample=UpperCAmelCase_       , set_alpha_to_one=UpperCAmelCase_       , steps_offset=1       , prediction_type='''epsilon'''       , thresholding=UpperCAmelCase_       , )

											lowerCAmelCase   =		{
											    '''unet''': unet,
											    '''scheduler''': scheduler,
											    '''movq''': movq,
											}
											return components



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_=0						):
											lowerCAmelCase   =		floats_tensor((1, self.text_embedder_hidden_size)       , rng=random.Random(UpperCAmelCase_						)						).to(UpperCAmelCase_						)
											lowerCAmelCase   =		floats_tensor((1, self.text_embedder_hidden_size)       , rng=random.Random(seed + 1						)						).to(
											    UpperCAmelCase_						)
											if str(UpperCAmelCase_						).startswith('''mps'''						):
																		lowerCAmelCase   =		torch.manual_seed(UpperCAmelCase_						)
											else:
																		lowerCAmelCase   =		torch.Generator(device=UpperCAmelCase_						).manual_seed(UpperCAmelCase_						)
											lowerCAmelCase   =		{
											    '''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       __snake_case		( self						):
											lowerCAmelCase   =		'''cpu'''

											lowerCAmelCase   =		self.get_dummy_components()

											lowerCAmelCase   =		self.pipeline_class(**UpperCAmelCase_						)
											lowerCAmelCase   =		pipe.to(UpperCAmelCase_						)

											pipe.set_progress_bar_config(disable=UpperCAmelCase_						)

											lowerCAmelCase   =		pipe(**self.get_dummy_inputs(UpperCAmelCase_						)						)
											lowerCAmelCase   =		output.images

											lowerCAmelCase   =		pipe(
											    **self.get_dummy_inputs(UpperCAmelCase_						)       , return_dict=UpperCAmelCase_       , )[0]

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

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

											lowerCAmelCase   =		np.array(
											    [0.6237976, 1.0, 0.36441332, 1.0, 0.70639634, 0.29877186, 0.85652125, 0.5216843, 0.54454046]						)

											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  ):





				'''simple docstring'''



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






				def       __snake_case		( self						):
											lowerCAmelCase   =		load_numpy(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
											    '''/kandinskyv22/kandinskyv22_text2img_cat_fp16.npy'''						)

											lowerCAmelCase   =		KandinskyVaaPriorPipeline.from_pretrained(
											    '''kandinsky-community/kandinsky-2-2-prior'''       , torch_dtype=torch.floataa						)
											pipe_prior.to(UpperCAmelCase_						)

											lowerCAmelCase   =		KandinskyVaaPipeline.from_pretrained(
											    '''kandinsky-community/kandinsky-2-2-decoder'''       , torch_dtype=torch.floataa						)
											lowerCAmelCase   =		pipeline.to(UpperCAmelCase_						)
											pipeline.set_progress_bar_config(disable=UpperCAmelCase_						)

											lowerCAmelCase   =		'''red cat, 4k photo'''

											lowerCAmelCase   =		torch.Generator(device='''cuda'''						).manual_seed(0						)
											lowerCAmelCase   ,     lowerCAmelCase   =		pipe_prior(
											    UpperCAmelCase_       , generator=UpperCAmelCase_       , num_inference_steps=5       , negative_prompt=''''''       , ).to_tuple()

											lowerCAmelCase   =		torch.Generator(device='''cuda'''						).manual_seed(0						)
											lowerCAmelCase   =		pipeline(
											    image_embeds=UpperCAmelCase_       , negative_image_embeds=UpperCAmelCase_       , generator=UpperCAmelCase_       , num_inference_steps=1_00       , output_type='''np'''       , )

											lowerCAmelCase   =		output.images[0]

											assert image.shape == (5_12, 5_12, 3)

											assert_mean_pixel_difference(UpperCAmelCase_       , UpperCAmelCase_						)









from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase_						=logging.get_logger(__name__)

UpperCAmelCase_						={
    """google/switch-base-8""": """https://huggingface.co/google/switch-base-8/blob/main/config.json""",
}


class 							__UpperCamelCase			( __UpperCAmelCase  ):





				'''simple docstring'''



				__a							: Any										="""switch_transformers"""
				__a							: Union[str, Any]										=["""past_key_values"""]
				__a							: Dict										={"""hidden_size""": """d_model""", """num_attention_heads""": """num_heads""", """num_hidden_layers""": """num_layers"""}



				def __init__( self       , UpperCAmelCase_=3_21_28       , UpperCAmelCase_=7_68       , UpperCAmelCase_=64       , UpperCAmelCase_=20_48       , UpperCAmelCase_=64       , UpperCAmelCase_=12       , UpperCAmelCase_=3       , UpperCAmelCase_=12       , UpperCAmelCase_=3       , UpperCAmelCase_=12       , UpperCAmelCase_=8       , UpperCAmelCase_=False       , UpperCAmelCase_=0.01       , UpperCAmelCase_="float32"       , UpperCAmelCase_=False       , UpperCAmelCase_=32       , UpperCAmelCase_=1_28       , UpperCAmelCase_=0.1       , UpperCAmelCase_=1E-6       , UpperCAmelCase_=0.001       , UpperCAmelCase_=0.001       , UpperCAmelCase_=1.0       , UpperCAmelCase_="relu"       , UpperCAmelCase_=True       , UpperCAmelCase_=False       , UpperCAmelCase_=True       , UpperCAmelCase_=0       , UpperCAmelCase_=1       , **UpperCAmelCase_       , ):
											lowerCAmelCase   =		vocab_size
											lowerCAmelCase   =		d_model
											lowerCAmelCase   =		d_kv
											lowerCAmelCase   =		d_ff

											lowerCAmelCase   =		num_sparse_encoder_layers

											lowerCAmelCase   =		num_layers
											lowerCAmelCase   =		(
											    num_decoder_layers if num_decoder_layers is not None else self.num_layers
											)  # default = symmetry
											lowerCAmelCase   =		num_sparse_decoder_layers

											# This tells us, each how many encoder layer we'll have to set a sparse layer.
											if self.num_sparse_encoder_layers > 0:
																		lowerCAmelCase   =		self.num_layers // self.num_sparse_encoder_layers
											else:
																		lowerCAmelCase   =		self.num_layers  # HACK: this will create 0 sparse layers

											# This tells us, each how many encoder layer we'll have to set a sparse layer.
											if self.num_sparse_decoder_layers > 0:
																		lowerCAmelCase   =		self.num_decoder_layers // self.num_sparse_decoder_layers
											else:
																		lowerCAmelCase   =		self.num_decoder_layers  # HACK: this will create 0 sparse layers

											lowerCAmelCase   =		num_heads
											lowerCAmelCase   =		num_experts
											lowerCAmelCase   =		expert_capacity
											lowerCAmelCase   =		router_bias
											lowerCAmelCase   =		router_jitter_noise
											if router_dtype not in ["float32", "float16", "bfloat16"]:
																		raise ValueError(F"""`router_dtype` must be one of 'float32', 'float16' or 'bfloat16', got {router_dtype}"""						)
											lowerCAmelCase   =		router_dtype

											lowerCAmelCase   =		router_ignore_padding_tokens
											lowerCAmelCase   =		relative_attention_num_buckets
											lowerCAmelCase   =		relative_attention_max_distance

											lowerCAmelCase   =		dropout_rate
											lowerCAmelCase   =		layer_norm_epsilon
											lowerCAmelCase   =		initializer_factor
											lowerCAmelCase   =		feed_forward_proj
											lowerCAmelCase   =		use_cache
											lowerCAmelCase   =		add_router_probs

											lowerCAmelCase   =		router_z_loss_coef
											lowerCAmelCase   =		router_aux_loss_coef

											lowerCAmelCase   =		self.feed_forward_proj.split('''-'''						)
											lowerCAmelCase   =		act_info[-1]
											lowerCAmelCase   =		act_info[0] == '''gated'''

											if len(UpperCAmelCase_						) > 1 and act_info[0] != "gated" or len(UpperCAmelCase_						) > 2:
																		raise ValueError(
																		    F"""`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."""
																		    '''Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. '''
																		    '''\'gated-gelu\' or \'relu\''''						)

											# for backwards compatibility
											if feed_forward_proj == "gated-gelu":
																		lowerCAmelCase   =		'''gelu_new'''

											super().__init__(
											    pad_token_id=UpperCAmelCase_       , eos_token_id=UpperCAmelCase_       , is_encoder_decoder=UpperCAmelCase_       , **UpperCAmelCase_       , )



