



import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np



class     _UpperCAmelCase			:




     def 		snake_case_		(				self):
         torch.manual_seed(0)
         A__		 =					TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5''')

         torch.manual_seed(0)
         A__		 =					AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5''')

         torch.manual_seed(0)
         A__		 =					UNetaDConditionModel(
             sample_size=3_2 ,						layers_per_block=1 ,						block_out_channels=[3_2, 6_4] ,						down_block_types=[
                 '''ResnetDownsampleBlock2D''',
                 '''SimpleCrossAttnDownBlock2D''',
             ] ,						mid_block_type='''UNetMidBlock2DSimpleCrossAttn''' ,						up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''] ,						in_channels=3 ,						out_channels=6 ,						cross_attention_dim=3_2 ,						encoder_hid_dim=3_2 ,						attention_head_dim=8 ,						addition_embed_type='''text''' ,						addition_embed_type_num_heads=2 ,						cross_attention_norm='''group_norm''' ,						resnet_time_scale_shift='''scale_shift''' ,						act_fn='''gelu''' ,						)
         unet.set_attn_processor(AttnAddedKVProcessor())  # For reproducibility tests

         torch.manual_seed(0)
         A__		 =					DDPMScheduler(
             num_train_timesteps=1_0_0_0 ,						beta_schedule='''squaredcos_cap_v2''' ,						beta_start=0.0_0_0_1 ,						beta_end=0.0_2 ,						thresholding=snake_case__ ,						dynamic_thresholding_ratio=0.9_5 ,						sample_max_value=1.0 ,						prediction_type='''epsilon''' ,						variance_type='''learned_range''' ,						)

         torch.manual_seed(0)
         A__		 =					IFWatermarker()

         return {
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "unet": unet,
             "scheduler": scheduler,
             "watermarker": watermarker,
             "safety_checker": None,
             "feature_extractor": None,
         }




     def 		snake_case_		(				self):
         torch.manual_seed(0)
         A__		 =					TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5''')

         torch.manual_seed(0)
         A__		 =					AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5''')

         torch.manual_seed(0)
         A__		 =					UNetaDConditionModel(
             sample_size=3_2 ,						layers_per_block=[1, 2] ,						block_out_channels=[3_2, 6_4] ,						down_block_types=[
                 '''ResnetDownsampleBlock2D''',
                 '''SimpleCrossAttnDownBlock2D''',
             ] ,						mid_block_type='''UNetMidBlock2DSimpleCrossAttn''' ,						up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''] ,						in_channels=6 ,						out_channels=6 ,						cross_attention_dim=3_2 ,						encoder_hid_dim=3_2 ,						attention_head_dim=8 ,						addition_embed_type='''text''' ,						addition_embed_type_num_heads=2 ,						cross_attention_norm='''group_norm''' ,						resnet_time_scale_shift='''scale_shift''' ,						act_fn='''gelu''' ,						class_embed_type='''timestep''' ,						mid_block_scale_factor=1.4_1_4 ,						time_embedding_act_fn='''gelu''' ,						time_embedding_dim=3_2 ,						)
         unet.set_attn_processor(AttnAddedKVProcessor())  # For reproducibility tests

         torch.manual_seed(0)
         A__		 =					DDPMScheduler(
             num_train_timesteps=1_0_0_0 ,						beta_schedule='''squaredcos_cap_v2''' ,						beta_start=0.0_0_0_1 ,						beta_end=0.0_2 ,						thresholding=snake_case__ ,						dynamic_thresholding_ratio=0.9_5 ,						sample_max_value=1.0 ,						prediction_type='''epsilon''' ,						variance_type='''learned_range''' ,						)

         torch.manual_seed(0)
         A__		 =					DDPMScheduler(
             num_train_timesteps=1_0_0_0 ,						beta_schedule='''squaredcos_cap_v2''' ,						beta_start=0.0_0_0_1 ,						beta_end=0.0_2 ,						)

         torch.manual_seed(0)
         A__		 =					IFWatermarker()

         return {
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "unet": unet,
             "scheduler": scheduler,
             "image_noising_scheduler": image_noising_scheduler,
             "watermarker": watermarker,
             "safety_checker": None,
             "feature_extractor": None,
         }




     def 		snake_case_		(				self):
         A__		 =					self.get_dummy_components()
         A__		 =					self.pipeline_class(**snake_case__)
         pipe.to(snake_case__)
         pipe.set_progress_bar_config(disable=snake_case__)

         A__		 =					self.get_dummy_inputs(snake_case__)

         A__		 =					inputs["prompt"]
         A__		 =					inputs["generator"]
         A__		 =					inputs["num_inference_steps"]
         A__		 =					inputs["output_type"]

         if "image" in inputs:
             A__		 =					inputs["image"]
         else:
             A__		 =					None

         if "mask_image" in inputs:
             A__		 =					inputs["mask_image"]
         else:
             A__		 =					None

         if "original_image" in inputs:
             A__		 =					inputs["original_image"]
         else:
             A__		 =					None

         A__		 =					pipe.encode_prompt(snake_case__)

         # inputs with prompt converted to embeddings
         A__		 =					{
             "prompt_embeds": prompt_embeds,
             "negative_prompt_embeds": negative_prompt_embeds,
             "generator": generator,
             "num_inference_steps": num_inference_steps,
             "output_type": output_type,
         }

         if image is not None:
             A__		 =					image

         if mask_image is not None:
             A__		 =					mask_image

         if original_image is not None:
             A__		 =					original_image

         # set all optional components to None
         for optional_component in pipe._optional_components:
             setattr(snake_case__ ,						snake_case__ ,						snake_case__)

         A__		 =					pipe(**snake_case__)[0]

         with tempfile.TemporaryDirectory() as tmpdir:
             pipe.save_pretrained(snake_case__)
             A__		 =					self.pipeline_class.from_pretrained(snake_case__)
             pipe_loaded.to(snake_case__)
             pipe_loaded.set_progress_bar_config(disable=snake_case__)

         pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor())  # For reproducibility tests

         for optional_component in pipe._optional_components:
             self.assertTrue(
                 getattr(snake_case__ ,						snake_case__) is None ,						F"`{optional_component}` did not stay set to None after loading." ,						)

         A__		 =					self.get_dummy_inputs(snake_case__)

         A__		 =					inputs["generator"]
         A__		 =					inputs["num_inference_steps"]
         A__		 =					inputs["output_type"]

         # inputs with prompt converted to embeddings
         A__		 =					{
             "prompt_embeds": prompt_embeds,
             "negative_prompt_embeds": negative_prompt_embeds,
             "generator": generator,
             "num_inference_steps": num_inference_steps,
             "output_type": output_type,
         }

         if image is not None:
             A__		 =					image

         if mask_image is not None:
             A__		 =					mask_image

         if original_image is not None:
             A__		 =					original_image

         A__		 =					pipe_loaded(**snake_case__)[0]

         A__		 =					np.abs(to_np(snake_case__) - to_np(snake_case__)).max()
         self.assertLess(snake_case__ ,						1e-4)




     def 		snake_case_		(				self):
         A__		 =					self.get_dummy_components()
         A__		 =					self.pipeline_class(**snake_case__)
         pipe.to(snake_case__)
         pipe.set_progress_bar_config(disable=snake_case__)

         A__		 =					self.get_dummy_inputs(snake_case__)
         A__		 =					pipe(**snake_case__)[0]

         with tempfile.TemporaryDirectory() as tmpdir:
             pipe.save_pretrained(snake_case__)
             A__		 =					self.pipeline_class.from_pretrained(snake_case__)
             pipe_loaded.to(snake_case__)
             pipe_loaded.set_progress_bar_config(disable=snake_case__)

         pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor())  # For reproducibility tests

         A__		 =					self.get_dummy_inputs(snake_case__)
         A__		 =					pipe_loaded(**snake_case__)[0]

         A__		 =					np.abs(to_np(snake_case__) - to_np(snake_case__)).max()
         self.assertLess(snake_case__ ,						1e-4)





from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig


_lowercase						  =		{
    "albert-base-v1": "https://huggingface.co/albert-base-v1/resolve/main/config.json",
    "albert-large-v1": "https://huggingface.co/albert-large-v1/resolve/main/config.json",
    "albert-xlarge-v1": "https://huggingface.co/albert-xlarge-v1/resolve/main/config.json",
    "albert-xxlarge-v1": "https://huggingface.co/albert-xxlarge-v1/resolve/main/config.json",
    "albert-base-v2": "https://huggingface.co/albert-base-v2/resolve/main/config.json",
    "albert-large-v2": "https://huggingface.co/albert-large-v2/resolve/main/config.json",
    "albert-xlarge-v2": "https://huggingface.co/albert-xlarge-v2/resolve/main/config.json",
    "albert-xxlarge-v2": "https://huggingface.co/albert-xxlarge-v2/resolve/main/config.json",
}



class     _UpperCAmelCase			(       A__     ):
     UpperCamelCase__		    =				'''albert'''




     def __init__(				self ,						a__=3_0_0_0_0 ,						a__=1_2_8 ,						a__=4_0_9_6 ,						a__=1_2 ,						a__=1 ,						a__=6_4 ,						a__=1_6_3_8_4 ,						a__=1 ,						a__="gelu_new" ,						a__=0 ,						a__=0 ,						a__=5_1_2 ,						a__=2 ,						a__=0.0_2 ,						a__=1e-12 ,						a__=0.1 ,						a__="absolute" ,						a__=0 ,						a__=2 ,						a__=3 ,						**a__ ,						):
         super().__init__(pad_token_id=a__ ,						bos_token_id=a__ ,						eos_token_id=a__ ,						**a__)

         A__		 =					vocab_size
         A__		 =					embedding_size
         A__		 =					hidden_size
         A__		 =					num_hidden_layers
         A__		 =					num_hidden_groups
         A__		 =					num_attention_heads
         A__		 =					inner_group_num
         A__		 =					hidden_act
         A__		 =					intermediate_size
         A__		 =					hidden_dropout_prob
         A__		 =					attention_probs_dropout_prob
         A__		 =					max_position_embeddings
         A__		 =					type_vocab_size
         A__		 =					initializer_range
         A__		 =					layer_norm_eps
         A__		 =					classifier_dropout_prob
         A__		 =					position_embedding_type







class     _UpperCAmelCase			(       A__     ):




     @property
     def 		snake_case_		(				self):
         if self.task == "multiple-choice":
             A__		 =					{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
         else:
             A__		 =					{0: '''batch''', 1: '''sequence'''}
         return OrderedDict(
             [
                 ('''input_ids''', dynamic_axis),
                 ('''attention_mask''', dynamic_axis),
                 ('''token_type_ids''', dynamic_axis),
             ])
