



import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    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 PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class        __snake_case		(  UpperCamelCase_		,UpperCamelCase_		,unittest.TestCase		):
							_a       =				StableDiffusionParadigmsPipeline
							_a       =				TEXT_TO_IMAGE_PARAMS
							_a       =				TEXT_TO_IMAGE_BATCH_PARAMS
							_a       =				TEXT_TO_IMAGE_IMAGE_PARAMS
							_a       =				TEXT_TO_IMAGE_IMAGE_PARAMS




							def      UpperCAmelCase__    ( self  : Optional[Any]):
									torch.manual_seed(0)
									lowerCAmelCase_    :  Any             =				UNetaDConditionModel(
									    block_out_channels=(3_2, 6_4)       ,       layers_per_block=2       ,       sample_size=3_2       ,       in_channels=4       ,       out_channels=4       ,       down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')       ,       up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')       ,       cross_attention_dim=3_2       ,       attention_head_dim=(2, 4)       ,       use_linear_projection=A_       ,       )
									lowerCAmelCase_    :  Union[str, Any]             =				DDIMParallelScheduler(
									    beta_start=0.0_0085       ,       beta_end=0.012       ,       beta_schedule='''scaled_linear'''       ,       clip_sample=A_       ,       set_alpha_to_one=A_       ,       )
									torch.manual_seed(0)
									lowerCAmelCase_    :  List[str]             =				AutoencoderKL(
									    block_out_channels=[3_2, 6_4]       ,       in_channels=3       ,       out_channels=3       ,       down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']       ,       up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']       ,       latent_channels=4       ,       sample_size=1_2_8       ,       )
									torch.manual_seed(0)
									lowerCAmelCase_    :  Optional[int]             =				CLIPTextConfig(
									    bos_token_id=0       ,       eos_token_id=2       ,       hidden_size=3_2       ,       intermediate_size=3_7       ,       layer_norm_eps=1e-05       ,       num_attention_heads=4       ,       num_hidden_layers=5       ,       pad_token_id=1       ,       vocab_size=1_0_0_0       ,       hidden_act='''gelu'''       ,       projection_dim=5_1_2       ,       )
									lowerCAmelCase_    :  Dict             =				CLIPTextModel(A_)
									lowerCAmelCase_    :  str             =				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')

									lowerCAmelCase_    :  Union[str, Any]             =				{
									    '''unet''': unet,
									    '''scheduler''': scheduler,
									    '''vae''': vae,
									    '''text_encoder''': text_encoder,
									    '''tokenizer''': tokenizer,
									    '''safety_checker''': None,
									    '''feature_extractor''': None,
									}
									return components




							def      UpperCAmelCase__    ( self  : List[str]       ,       A_  : int       ,       A_  : Union[str, Any]=0):
									if str(A_).startswith('''mps'''):
											lowerCAmelCase_    :  List[Any]             =				torch.manual_seed(A_)
									else:
											lowerCAmelCase_    :  List[str]             =				torch.Generator(device=A_).manual_seed(A_)
									lowerCAmelCase_    :  Any             =				{
									    '''prompt''': '''a photograph of an astronaut riding a horse''',
									    '''generator''': generator,
									    '''num_inference_steps''': 1_0,
									    '''guidance_scale''': 6.0,
									    '''output_type''': '''numpy''',
									    '''parallel''': 3,
									    '''debug''': True,
									}
									return inputs




							def      UpperCAmelCase__    ( self  : Tuple):
									lowerCAmelCase_    :  int             =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									lowerCAmelCase_    :  List[Any]             =				self.get_dummy_components()
									lowerCAmelCase_    :  str             =				StableDiffusionParadigmsPipeline(**A_)
									lowerCAmelCase_    :  Any             =				sd_pipe.to(A_)
									sd_pipe.set_progress_bar_config(disable=A_)

									lowerCAmelCase_    :  Any             =				self.get_dummy_inputs(A_)
									lowerCAmelCase_    :  Any             =				sd_pipe(**A_).images
									lowerCAmelCase_    :  List[str]             =				image[0, -3:, -3:, -1]
									assert image.shape == (1, 6_4, 6_4, 3)

									lowerCAmelCase_    :  int             =				np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023])

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




							def      UpperCAmelCase__    ( self  : List[Any]):
									lowerCAmelCase_    :  Union[str, Any]             =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									lowerCAmelCase_    :  Optional[Any]             =				self.get_dummy_components()
									torch.manual_seed(0)
									lowerCAmelCase_    :  str             =				DDPMParallelScheduler()
									torch.manual_seed(0)
									lowerCAmelCase_    :  Any             =				StableDiffusionParadigmsPipeline(**A_)
									lowerCAmelCase_    :  List[str]             =				sd_pipe.to(A_)
									sd_pipe.set_progress_bar_config(disable=A_)

									lowerCAmelCase_    :  Tuple             =				self.get_dummy_inputs(A_)
									lowerCAmelCase_    :  int             =				sd_pipe(**A_).images
									lowerCAmelCase_    :  List[str]             =				image[0, -3:, -3:, -1]
									assert image.shape == (1, 6_4, 6_4, 3)

									lowerCAmelCase_    :  Optional[Any]             =				np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468])

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




							def      UpperCAmelCase__    ( self  : List[Any]):
									super().test_inference_batch_consistent(batch_sizes=[1, 2])




							def      UpperCAmelCase__    ( self  : Optional[Any]):
									super().test_inference_batch_single_identical(batch_size=2       ,       expected_max_diff=3e-3)







							def      UpperCAmelCase__    ( self  : Optional[int]):
									lowerCAmelCase_    :  Union[str, Any]             =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									lowerCAmelCase_    :  Any             =				self.get_dummy_components()
									lowerCAmelCase_    :  List[Any]             =				StableDiffusionParadigmsPipeline(**A_)
									lowerCAmelCase_    :  Dict             =				sd_pipe.to(A_)
									sd_pipe.set_progress_bar_config(disable=A_)

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

									assert image.shape == (1, 6_4, 6_4, 3)

									lowerCAmelCase_    :  str             =				np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015])

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





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




							def      UpperCAmelCase__    ( self  : List[str]):
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()




							def      UpperCAmelCase__    ( self  : Optional[int]       ,       A_  : Dict=0):
									lowerCAmelCase_    :  Union[str, Any]             =				torch.Generator(device=A_).manual_seed(A_)
									lowerCAmelCase_    :  int             =				{
									    '''prompt''': '''a photograph of an astronaut riding a horse''',
									    '''generator''': generator,
									    '''num_inference_steps''': 1_0,
									    '''guidance_scale''': 7.5,
									    '''output_type''': '''numpy''',
									    '''parallel''': 3,
									    '''debug''': True,
									}
									return inputs







							def      UpperCAmelCase__    ( self  : Tuple):
									lowerCAmelCase_    :  int             =				'''stabilityai/stable-diffusion-2-base'''
									lowerCAmelCase_    :  Optional[Any]             =				DDIMParallelScheduler.from_pretrained(A_       ,       subfolder='''scheduler''')
									lowerCAmelCase_    :  Any             =				StableDiffusionParadigmsPipeline.from_pretrained(A_       ,       scheduler=A_       ,       safety_checker=A_)
									pipe.to(A_)
									pipe.set_progress_bar_config(disable=A_)
									pipe.enable_attention_slicing()

									lowerCAmelCase_    :  List[str]             =				self.get_inputs()
									lowerCAmelCase_    :  Any             =				pipe(**A_).images
									lowerCAmelCase_    :  Any             =				image[0, -3:, -3:, -1].flatten()

									assert image.shape == (1, 5_1_2, 5_1_2, 3)

									lowerCAmelCase_    :  Dict             =				np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741])

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










from collections import OrderedDict
from typing import Mapping

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


A__						:						Dict          =      logging.get_logger(__name__)

A__						:						List[str]          =      {
    '''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        __snake_case		(  UpperCamelCase_		):
							_a       =				'''electra'''




							def __init__( self  : Dict       ,       A_  : Dict=3_0_5_2_2       ,       A_  : List[Any]=1_2_8       ,       A_  : Dict=2_5_6       ,       A_  : Any=1_2       ,       A_  : Optional[Any]=4       ,       A_  : Optional[int]=1_0_2_4       ,       A_  : Optional[int]="gelu"       ,       A_  : str=0.1       ,       A_  : Optional[int]=0.1       ,       A_  : str=5_1_2       ,       A_  : Optional[int]=2       ,       A_  : Dict=0.02       ,       A_  : Union[str, Any]=1e-12       ,       A_  : List[Any]="first"       ,       A_  : List[str]=True       ,       A_  : List[str]="gelu"       ,       A_  : Any=0.1       ,       A_  : int=0       ,       A_  : Optional[Any]="absolute"       ,       A_  : Optional[int]=True       ,       A_  : str=None       ,       **A_  : List[Any]       ,       ):
									super().__init__(pad_token_id=A_       ,       **A_)

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

									lowerCAmelCase_    :  Union[str, Any]             =				summary_type
									lowerCAmelCase_    :  str             =				summary_use_proj
									lowerCAmelCase_    :  str             =				summary_activation
									lowerCAmelCase_    :  List[Any]             =				summary_last_dropout
									lowerCAmelCase_    :  Optional[Any]             =				position_embedding_type
									lowerCAmelCase_    :  List[Any]             =				use_cache
									lowerCAmelCase_    :  Any             =				classifier_dropout





class        __snake_case		(  UpperCamelCase_		):




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





