



"""simple docstring"""






import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

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


enable_full_determinism()

@skip_mps
class 			__SCREAMING_SNAKE_CASE      ( lowerCAmelCase_						,  lowerCAmelCase_						,  lowerCAmelCase_						,  unittest.TestCase  ):


					'''simple docstring'''





					_a       =				StableDiffusionModelEditingPipeline
					_a       =				TEXT_TO_IMAGE_PARAMS
					_a       =				TEXT_TO_IMAGE_BATCH_PARAMS
					_a       =				TEXT_TO_IMAGE_IMAGE_PARAMS
					_a       =				TEXT_TO_IMAGE_IMAGE_PARAMS



					def 						snake_case			(			self				:    Optional[int]	)->   str:
								torch.manual_seed(0	)
								lowerCamelCase__						:						Optional[int]						      =UNetaDConditionModel(
								    block_out_channels=(32, 64),    layers_per_block=2,    sample_size=32,    in_channels=4,    out_channels=4,    down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D'''),    up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D'''),    cross_attention_dim=32,    )
								lowerCamelCase__						:						Any						      =DDIMScheduler()
								torch.manual_seed(0	)
								lowerCamelCase__						:						Union[str, 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,    )
								torch.manual_seed(0	)
								lowerCamelCase__						:						List[Any]						      =CLIPTextConfig(
								    bos_token_id=0,    eos_token_id=2,    hidden_size=32,    intermediate_size=37,    layer_norm_eps=1E-05,    num_attention_heads=4,    num_hidden_layers=5,    pad_token_id=1,    vocab_size=1000,    )
								lowerCamelCase__						:						List[Any]						      =CLIPTextModel(lowerCamelCase	)
								lowerCamelCase__						:						Tuple						      =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''	)

								lowerCamelCase__						:						Any						      ={
								    '''unet''': unet,
								    '''scheduler''': scheduler,
								    '''vae''': vae,
								    '''text_encoder''': text_encoder,
								    '''tokenizer''': tokenizer,
								    '''safety_checker''': None,
								    '''feature_extractor''': None,
								}
								return components



					def 						snake_case			(			self				:    Tuple,    lowerCamelCase				:    Optional[int],    lowerCamelCase				:    List[Any]=0	)->   Dict:
								lowerCamelCase__						:						Union[str, Any]						      =torch.manual_seed(lowerCamelCase	)
								lowerCamelCase__						:						Tuple						      ={
								    '''prompt''': '''A field of roses''',
								    '''generator''': generator,
								    # Setting height and width to None to prevent OOMs on CPU.
								    '''height''': None,
								    '''width''': None,
								    '''num_inference_steps''': 2,
								    '''guidance_scale''': 6.0,
								    '''output_type''': '''numpy''',
								}
								return inputs



					def 						snake_case			(			self				:    List[Any]	)->   int:
								lowerCamelCase__						:						str						      ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								lowerCamelCase__						:						int						      =self.get_dummy_components()
								lowerCamelCase__						:						Any						      =StableDiffusionModelEditingPipeline(**lowerCamelCase	)
								lowerCamelCase__						:						List[Any]						      =sd_pipe.to(lowerCamelCase	)
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase	)

								lowerCamelCase__						:						Optional[int]						      =self.get_dummy_inputs(lowerCamelCase	)
								lowerCamelCase__						:						Dict						      =sd_pipe(**lowerCamelCase	).images
								lowerCamelCase__						:						List[Any]						      =image[0, -3:, -3:, -1]
								assert image.shape == (1, 64, 64, 3)

								lowerCamelCase__						:						Optional[Any]						      =np.array([0.4_755, 0.5_132, 0.4_976, 0.3_904, 0.3_554, 0.4_765, 0.5_139, 0.5_158, 0.4_889]	)

								assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-2



					def 						snake_case			(			self				:    int	)->   List[Any]:
								lowerCamelCase__						:						Dict						      ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								lowerCamelCase__						:						List[Any]						      =self.get_dummy_components()
								lowerCamelCase__						:						int						      =StableDiffusionModelEditingPipeline(**lowerCamelCase	)
								lowerCamelCase__						:						Optional[int]						      =sd_pipe.to(lowerCamelCase	)
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase	)

								lowerCamelCase__						:						Optional[int]						      =self.get_dummy_inputs(lowerCamelCase	)
								lowerCamelCase__						:						Union[str, Any]						      ='''french fries'''
								lowerCamelCase__						:						Optional[int]						      =sd_pipe(**lowerCamelCase,    negative_prompt=lowerCamelCase	)
								lowerCamelCase__						:						Optional[Any]						      =output.images
								lowerCamelCase__						:						str						      =image[0, -3:, -3:, -1]

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

								lowerCamelCase__						:						Union[str, Any]						      =np.array([0.4_992, 0.5_101, 0.5_004, 0.3_949, 0.3_604, 0.4_735, 0.5_216, 0.5_204, 0.4_913]	)

								assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-2



					def 						snake_case			(			self				:    List[Any]	)->   Tuple:
								lowerCamelCase__						:						int						      ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								lowerCamelCase__						:						Dict						      =self.get_dummy_components()
								lowerCamelCase__						:						Any						      =EulerAncestralDiscreteScheduler(
								    beta_start=0.00_085,    beta_end=0.012,    beta_schedule='''scaled_linear'''	)
								lowerCamelCase__						:						List[Any]						      =StableDiffusionModelEditingPipeline(**lowerCamelCase	)
								lowerCamelCase__						:						Union[str, Any]						      =sd_pipe.to(lowerCamelCase	)
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase	)

								lowerCamelCase__						:						Optional[Any]						      =self.get_dummy_inputs(lowerCamelCase	)
								lowerCamelCase__						:						int						      =sd_pipe(**lowerCamelCase	).images
								lowerCamelCase__						:						Tuple						      =image[0, -3:, -3:, -1]

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

								lowerCamelCase__						:						List[str]						      =np.array([0.4_747, 0.5_372, 0.4_779, 0.4_982, 0.5_543, 0.4_816, 0.5_238, 0.4_904, 0.5_027]	)

								assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-2



					def 						snake_case			(			self				:    List[str]	)->   Dict:
								lowerCamelCase__						:						Any						      ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								lowerCamelCase__						:						Union[str, Any]						      =self.get_dummy_components()
								lowerCamelCase__						:						str						      =PNDMScheduler()
								lowerCamelCase__						:						Dict						      =StableDiffusionModelEditingPipeline(**lowerCamelCase	)
								lowerCamelCase__						:						Optional[Any]						      =sd_pipe.to(lowerCamelCase	)
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase	)

								lowerCamelCase__						:						Tuple						      =self.get_dummy_inputs(lowerCamelCase	)
								# the pipeline does not expect pndm so test if it raises error.
								with self.assertRaises(lowerCamelCase	):
											lowerCamelCase__						:						Union[str, Any]						      =sd_pipe(**lowerCamelCase	).images



					def 						snake_case			(			self				:    Optional[int]	)->   Optional[Any]:
								super().test_inference_batch_single_identical(expected_max_diff=5E-3	)




					def 						snake_case			(			self				:    Any	)->   str:
								super().test_attention_slicing_forward_pass(expected_max_diff=5E-3	)





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


					'''simple docstring'''





					def 						snake_case			(			self				:    Any	)->   int:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()



					def 						snake_case			(			self				:    Union[str, Any],    lowerCamelCase				:    Dict=0	)->   List[Any]:
								lowerCamelCase__						:						Any						      =torch.manual_seed(lowerCamelCase	)
								lowerCamelCase__						:						List[Any]						      ={
								    '''prompt''': '''A field of roses''',
								    '''generator''': generator,
								    '''num_inference_steps''': 3,
								    '''guidance_scale''': 7.5,
								    '''output_type''': '''numpy''',
								}
								return inputs



					def 						snake_case			(			self				:    int	)->   Any:
								lowerCamelCase__						:						Tuple						      ='''CompVis/stable-diffusion-v1-4'''
								lowerCamelCase__						:						Tuple						      =StableDiffusionModelEditingPipeline.from_pretrained(lowerCamelCase,    safety_checker=lowerCamelCase	)
								pipe.to(lowerCamelCase	)
								pipe.set_progress_bar_config(disable=lowerCamelCase	)
								pipe.enable_attention_slicing()

								lowerCamelCase__						:						Tuple						      =self.get_inputs()
								lowerCamelCase__						:						List[str]						      =pipe(**lowerCamelCase	).images
								lowerCamelCase__						:						Optional[int]						      =image[0, -3:, -3:, -1].flatten()

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

								lowerCamelCase__						:						Tuple						      =np.array(
								    [0.6_749_496, 0.6_386_453, 0.51_443_267, 0.66_094_905, 0.61_921_215, 0.5_491_332, 0.5_744_417, 0.58_075_106, 0.5_174_658]	)

								assert np.abs(expected_slice - image_slice	).max() < 1E-2

								# make sure image changes after editing
								pipe.edit_model('''A pack of roses''',    '''A pack of blue roses'''	)

								lowerCamelCase__						:						Any						      =pipe(**lowerCamelCase	).images
								lowerCamelCase__						:						List[str]						      =image[0, -3:, -3:, -1].flatten()

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

								assert np.abs(expected_slice - image_slice	).max() > 1E-1




					def 						snake_case			(			self				:    Union[str, Any]	)->   Tuple:
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

								lowerCamelCase__						:						Any						      ='''CompVis/stable-diffusion-v1-4'''
								lowerCamelCase__						:						Tuple						      =DDIMScheduler.from_pretrained(lowerCamelCase,    subfolder='''scheduler'''	)
								lowerCamelCase__						:						Tuple						      =StableDiffusionModelEditingPipeline.from_pretrained(
								    lowerCamelCase,    scheduler=lowerCamelCase,    safety_checker=lowerCamelCase	)
								lowerCamelCase__						:						int						      =pipe.to(lowerCamelCase	)
								pipe.set_progress_bar_config(disable=lowerCamelCase	)
								pipe.enable_attention_slicing(1	)
								pipe.enable_sequential_cpu_offload()

								lowerCamelCase__						:						int						      =self.get_inputs()
								lowerCamelCase__						:						Optional[int]						      =pipe(**lowerCamelCase	)

								lowerCamelCase__						:						int						      =torch.cuda.max_memory_allocated()
								# make sure that less than 4.4 GB is allocated
								assert mem_bytes < 4.4 * 10**9











from collections import OrderedDict
from typing import Mapping

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


UpperCAmelCase_       : Union[str, Any]							=    logging.get_logger(__name__)

UpperCAmelCase_       : int							=    {
    """google/electra-small-generator""": """https://huggingface.co/google/electra-small-generator/resolve/main/config.json""",
    """google/electra-base-generator""": """https://huggingface.co/google/electra-base-generator/resolve/main/config.json""",
    """google/electra-large-generator""": """https://huggingface.co/google/electra-large-generator/resolve/main/config.json""",
    """google/electra-small-discriminator""": (
        """https://huggingface.co/google/electra-small-discriminator/resolve/main/config.json"""
    ),
    """google/electra-base-discriminator""": (
        """https://huggingface.co/google/electra-base-discriminator/resolve/main/config.json"""
    ),
    """google/electra-large-discriminator""": (
        """https://huggingface.co/google/electra-large-discriminator/resolve/main/config.json"""
    ),
}





class   lowercase__     (				__A	):
  __UpperCamelCase						  =	"""electra"""







  def __init__(						self     ,    _lowercase=30_522     ,    _lowercase=128     ,    _lowercase=256     ,    _lowercase=12     ,    _lowercase=4     ,    _lowercase=1_024     ,    _lowercase="gelu"     ,    _lowercase=0.1     ,    _lowercase=0.1     ,    _lowercase=512     ,    _lowercase=2     ,    _lowercase=0.02     ,    _lowercase=1e-12     ,    _lowercase="first"     ,    _lowercase=True     ,    _lowercase="gelu"     ,    _lowercase=0.1     ,    _lowercase=0     ,    _lowercase="absolute"     ,    _lowercase=True     ,    _lowercase=None     ,    **_lowercase     ,    ):
       super().__init__(pad_token_id=_lowercase     ,    **_lowercase    )

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

       lowerCAmelCase_       :							str			 =		summary_type
       lowerCAmelCase_       :							Optional[int]			 =		summary_use_proj
       lowerCAmelCase_       :							Optional[int]			 =		summary_activation
       lowerCAmelCase_       :							int			 =		summary_last_dropout
       lowerCAmelCase_       :							str			 =		position_embedding_type
       lowerCAmelCase_       :							Optional[int]			 =		use_cache
       lowerCAmelCase_       :							Tuple			 =		classifier_dropout





class   lowercase__     (				__A	):







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