





'''simple docstring'''




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionAttendAndExcitePipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_numpy, skip_mps, slow
from diffusers.utils.testing_utils import require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


lowercase__   :				Union[str, Any]      =   False


@skip_mps
class 			SCREAMING_SNAKE_CASE   (a__ ,		a__ ,		a__ ,		unittest.TestCase							):
     lowerCAmelCase		      =	StableDiffusionAttendAndExcitePipeline
     lowerCAmelCase		      =	False
     lowerCAmelCase		      =	TEXT_TO_IMAGE_PARAMS
     lowerCAmelCase		      =	TEXT_TO_IMAGE_BATCH_PARAMS.union({'''token_indices'''}							)
     lowerCAmelCase		      =	TEXT_TO_IMAGE_IMAGE_PARAMS
     lowerCAmelCase		      =	TEXT_TO_IMAGE_IMAGE_PARAMS


     @classmethod
     def 							SCREAMING_SNAKE_CASE  (						cls):

          '''simple docstring'''






          super().setUpClass()
          torch.use_deterministic_algorithms(_UpperCAmelCase)


     @classmethod
     def 							SCREAMING_SNAKE_CASE  (						cls):

          '''simple docstring'''






          super().tearDownClass()
          torch.use_deterministic_algorithms(_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          torch.manual_seed(0)
          __A						:			int								=	UNetaDConditionModel(
              block_out_channels=(32, 64)						,     layers_per_block=1						,     sample_size=32						,     in_channels=4						,     out_channels=4						,     down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')						,     up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')						,     cross_attention_dim=32						,     attention_head_dim=(2, 4)						,     use_linear_projection=_UpperCAmelCase						,     )
          __A						:			List[str]								=	DDIMScheduler(
              beta_start=0.00085						,     beta_end=0.012						,     beta_schedule='scaled_linear'						,     clip_sample=_UpperCAmelCase						,     set_alpha_to_one=_UpperCAmelCase						,     )
          torch.manual_seed(0)
          __A						:			Any								=	AutoencoderKL(
              block_out_channels=[32, 64]						,     in_channels=3						,     out_channels=3						,     down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']						,     up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']						,     latent_channels=4						,     sample_size=128						,     )
          torch.manual_seed(0)
          __A						:			List[Any]								=	CLIPTextConfig(
              bos_token_id=0						,     eos_token_id=2						,     hidden_size=32						,     intermediate_size=37						,     layer_norm_eps=1e-0_5						,     num_attention_heads=4						,     num_hidden_layers=5						,     pad_token_id=1						,     vocab_size=1000						,     hidden_act='gelu'						,     projection_dim=512						,     )
          __A						:			List[Any]								=	CLIPTextModel(_UpperCAmelCase)
          __A						:			List[Any]								=	CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

          __A						:			int								=	{
              'unet': unet,
              'scheduler': scheduler,
              'vae': vae,
              'text_encoder': text_encoder,
              'tokenizer': tokenizer,
              'safety_checker': None,
              'feature_extractor': None,
          }

          return components


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase						,     _UpperCAmelCase=0):

          '''simple docstring'''






          if str(_UpperCAmelCase).startswith('mps'):
               __A						:			Any								=	torch.manual_seed(_UpperCAmelCase)
          else:
               __A						:			List[str]								=	torch.Generator(device=_UpperCAmelCase).manual_seed(_UpperCAmelCase)
          __A						:			Any								=	{
              'prompt': 'a cat and a frog',
              'token_indices': [2, 5],
              'generator': generator,
              'num_inference_steps': 1,
              'guidance_scale': 6.0,
              'output_type': 'numpy',
              'max_iter_to_alter': 2,
              'thresholds': {0: 0.7},
          }
          return inputs


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Tuple								=	'cpu'

          __A						:			Tuple								=	self.get_dummy_components()
          __A						:			Any								=	self.pipeline_class(**_UpperCAmelCase)
          pipe.to(_UpperCAmelCase)
          pipe.set_progress_bar_config(disable=_UpperCAmelCase)

          __A						:			Union[str, Any]								=	self.get_dummy_inputs(_UpperCAmelCase)
          __A						:			str								=	pipe(**_UpperCAmelCase).images
          __A						:			List[str]								=	image[0, -3:, -3:, -1]

          self.assertEqual(image.shape						,     (1, 64, 64, 3))
          __A						:			List[Any]								=	np.array(
              [0.63905364, 0.62897307, 0.48599017, 0.5133624, 0.5550048, 0.45769516, 0.50326973, 0.5023139, 0.45384496])
          __A						:			Optional[Any]								=	np.abs(image_slice.flatten() - expected_slice).max()
          self.assertLessEqual(_UpperCAmelCase						,     1e-3)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          super().test_cpu_offload_forward_pass(expected_max_diff=5e-4)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          self._test_inference_batch_consistent(batch_sizes=[1, 2])


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          self._test_inference_batch_single_identical(batch_size=2						,     expected_max_diff=7e-4)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5e-4)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          super().test_save_load_local(expected_max_difference=5e-4)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          super().test_save_load_optional_components(expected_max_difference=4e-4)




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


     @classmethod
     def 							SCREAMING_SNAKE_CASE  (						cls):

          '''simple docstring'''






          super().setUpClass()
          torch.use_deterministic_algorithms(_UpperCAmelCase)


     @classmethod
     def 							SCREAMING_SNAKE_CASE  (						cls):

          '''simple docstring'''






          super().tearDownClass()
          torch.use_deterministic_algorithms(_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Any								=	torch.manual_seed(51)

          __A						:			Dict								=	StableDiffusionAttendAndExcitePipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,     safety_checker=_UpperCAmelCase						,     torch_dtype=torch.floataa)
          pipe.to('cuda')

          __A						:			List[str]								=	'a painting of an elephant with glasses'
          __A						:			Union[str, Any]								=	[5, 7]

          __A						:			int								=	pipe(
              prompt=_UpperCAmelCase						,     token_indices=_UpperCAmelCase						,     guidance_scale=7.5						,     generator=_UpperCAmelCase						,     num_inference_steps=5						,     max_iter_to_alter=5						,     output_type='numpy'						,     ).images[0]

          __A						:			List[Any]								=	load_numpy(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/attend-and-excite/elephant_glasses.npy')
          assert np.abs((expected_image - image).max()) < 5e-1






'''simple docstring'''




from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


lowercase__   :				Optional[Any]      =   logging.get_logger(__name__)

lowercase__   :				str      =   {
    '''andreasmadsen/efficient_mlm_m0.40''': (
        '''https://huggingface.co/andreasmadsen/efficient_mlm_m0.40/resolve/main/config.json'''
    ),
}


class 			SCREAMING_SNAKE_CASE   (a__							):
     lowerCAmelCase		      =	'''roberta-prelayernorm'''


     def __init__(						self						,     _UpperCAmelCase=5_0265						,     _UpperCAmelCase=768						,     _UpperCAmelCase=12						,     _UpperCAmelCase=12						,     _UpperCAmelCase=3072						,     _UpperCAmelCase="gelu"						,     _UpperCAmelCase=0.1						,     _UpperCAmelCase=0.1						,     _UpperCAmelCase=512						,     _UpperCAmelCase=2						,     _UpperCAmelCase=0.02						,     _UpperCAmelCase=1e-1_2						,     _UpperCAmelCase=1						,     _UpperCAmelCase=0						,     _UpperCAmelCase=2						,     _UpperCAmelCase="absolute"						,     _UpperCAmelCase=True						,     _UpperCAmelCase=None						,     **_UpperCAmelCase						,     ):

          '''simple docstring'''






          super().__init__(pad_token_id=_UpperCAmelCase						,     bos_token_id=_UpperCAmelCase						,     eos_token_id=_UpperCAmelCase						,     **_UpperCAmelCase)

          __A						:			Union[str, Any]								=	vocab_size
          __A						:			List[str]								=	hidden_size
          __A						:			Tuple								=	num_hidden_layers
          __A						:			List[str]								=	num_attention_heads
          __A						:			Tuple								=	hidden_act
          __A						:			Optional[Any]								=	intermediate_size
          __A						:			List[Any]								=	hidden_dropout_prob
          __A						:			Union[str, Any]								=	attention_probs_dropout_prob
          __A						:			Dict								=	max_position_embeddings
          __A						:			Any								=	type_vocab_size
          __A						:			Optional[int]								=	initializer_range
          __A						:			List[str]								=	layer_norm_eps
          __A						:			Optional[int]								=	position_embedding_type
          __A						:			Dict								=	use_cache
          __A						:			List[str]								=	classifier_dropout




class 			SCREAMING_SNAKE_CASE   (a__							):


     @property
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          if self.task == "multiple-choice":
               __A						:			List[str]								=	{0: 'batch', 1: 'choice', 2: 'sequence'}
          else:
               __A						:			Tuple								=	{0: 'batch', 1: 'sequence'}
          return OrderedDict(
              [
                  ('input_ids', dynamic_axis),
                  ('attention_mask', dynamic_axis),
              ])