

"""simple docstring"""


import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    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	(				A_,							A_,							unittest.TestCase  ):
				SCREAMING_SNAKE_CASE_   :      str         =  StableDiffusionSAGPipeline
				SCREAMING_SNAKE_CASE_   :      Union[str, Any]         =  TEXT_TO_IMAGE_PARAMS
				SCREAMING_SNAKE_CASE_   :      Tuple         =  TEXT_TO_IMAGE_BATCH_PARAMS
				SCREAMING_SNAKE_CASE_   :      List[Any]         =  TEXT_TO_IMAGE_IMAGE_PARAMS
				SCREAMING_SNAKE_CASE_   :      str         =  TEXT_TO_IMAGE_IMAGE_PARAMS
				SCREAMING_SNAKE_CASE_   :      Optional[int]         =  False

				def 							lowercase_     (	self							:				Tuple)->   List[str]:



						'''simple docstring'''


						torch.manual_seed(0)
						__lowerCAmelCase:       Dict          =			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     ,  )
						__lowerCAmelCase:       Dict          =			DDIMScheduler(
						    beta_start=0.00085     ,  beta_end=0.012     ,  beta_schedule="scaled_linear"     ,  clip_sample=_snake_case     ,  set_alpha_to_one=_snake_case     ,  )
						torch.manual_seed(0)
						__lowerCAmelCase:       Dict          =			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     ,  )
						torch.manual_seed(0)
						__lowerCAmelCase:       List[Any]          =			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     ,  )
						__lowerCAmelCase:       str          =			CLIPTextModel(_snake_case)
						__lowerCAmelCase:       Optional[int]          =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

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

				def 							lowercase_     (	self							:				Dict     ,  UpperCamelCase__							:				List[str]     ,  UpperCamelCase__							:				str=0)->   List[Any]:



						'''simple docstring'''


						if str(_snake_case).startswith("mps"):
								__lowerCAmelCase:       Optional[Any]          =			torch.manual_seed(_snake_case)
						else:
								__lowerCAmelCase:       int          =			torch.Generator(device=_snake_case).manual_seed(_snake_case)
						__lowerCAmelCase:       Any          =			{
						    '''prompt''': '''.''',
						    '''generator''': generator,
						    '''num_inference_steps''': 2,
						    '''guidance_scale''': 1.0,
						    '''sag_scale''': 1.0,
						    '''output_type''': '''numpy''',
						}
						return inputs




				def 							lowercase_     (	self							:				str)->   Optional[Any]:



						'''simple docstring'''


						super().test_inference_batch_single_identical(expected_max_diff=3e-3)





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

				def 							lowercase_     (	self							:				Optional[Any])->   List[Any]:



						'''simple docstring'''


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

				def 							lowercase_     (	self							:				Optional[Any])->   Dict:



						'''simple docstring'''


						__lowerCAmelCase:       str          =			StableDiffusionSAGPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
						__lowerCAmelCase:       Any          =			sag_pipe.to(_snake_case)
						sag_pipe.set_progress_bar_config(disable=_snake_case)

						__lowerCAmelCase:       str          =			'''.'''
						__lowerCAmelCase:       Optional[Any]          =			torch.manual_seed(0)
						__lowerCAmelCase:       Any          =			sag_pipe(
						    [prompt]     ,  generator=_snake_case     ,  guidance_scale=7.5     ,  sag_scale=1.0     ,  num_inference_steps=2_0     ,  output_type="np")

						__lowerCAmelCase:       Any          =			output.images

						__lowerCAmelCase:       Any          =			image[0, -3:, -3:, -1]

						assert image.shape == (1, 5_1_2, 5_1_2, 3)
						__lowerCAmelCase:       int          =			np.array([0.1568, 0.1738, 0.1695, 0.1693, 0.1507, 0.1705, 0.1547, 0.1751, 0.1949])

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

				def 							lowercase_     (	self							:				int)->   Any:



						'''simple docstring'''


						__lowerCAmelCase:       Dict          =			StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base")
						__lowerCAmelCase:       Optional[Any]          =			sag_pipe.to(_snake_case)
						sag_pipe.set_progress_bar_config(disable=_snake_case)

						__lowerCAmelCase:       Any          =			'''.'''
						__lowerCAmelCase:       Optional[Any]          =			torch.manual_seed(0)
						__lowerCAmelCase:       int          =			sag_pipe(
						    [prompt]     ,  generator=_snake_case     ,  guidance_scale=7.5     ,  sag_scale=1.0     ,  num_inference_steps=2_0     ,  output_type="np")

						__lowerCAmelCase:       Dict          =			output.images

						__lowerCAmelCase:       Dict          =			image[0, -3:, -3:, -1]

						assert image.shape == (1, 5_1_2, 5_1_2, 3)
						__lowerCAmelCase:       int          =			np.array([0.3459, 0.2876, 0.2537, 0.3002, 0.2671, 0.2160, 0.3026, 0.2262, 0.2371])

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




				def 							lowercase_     (	self							:				Any)->   Optional[Any]:



						'''simple docstring'''


						__lowerCAmelCase:       int          =			StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base")
						__lowerCAmelCase:       Tuple          =			sag_pipe.to(_snake_case)
						sag_pipe.set_progress_bar_config(disable=_snake_case)

						__lowerCAmelCase:       Tuple          =			'''.'''
						__lowerCAmelCase:       List[Any]          =			torch.manual_seed(0)
						__lowerCAmelCase:       Dict          =			sag_pipe(
						    [prompt]     ,  width=7_6_8     ,  height=5_1_2     ,  generator=_snake_case     ,  guidance_scale=7.5     ,  sag_scale=1.0     ,  num_inference_steps=2_0     ,  output_type="np"     ,  )

						__lowerCAmelCase:       List[str]          =			output.images

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







"""simple docstring"""







from ...configuration_utils import PretrainedConfig


lowerCAmelCase_   						=   {
    'google/tapas-base-finetuned-sqa': (
        'https://huggingface.co/google/tapas-base-finetuned-sqa/resolve/main/config.json'
    ),
    'google/tapas-base-finetuned-wtq': (
        'https://huggingface.co/google/tapas-base-finetuned-wtq/resolve/main/config.json'
    ),
    'google/tapas-base-finetuned-wikisql-supervised': (
        'https://huggingface.co/google/tapas-base-finetuned-wikisql-supervised/resolve/main/config.json'
    ),
    'google/tapas-base-finetuned-tabfact': (
        'https://huggingface.co/google/tapas-base-finetuned-tabfact/resolve/main/config.json'
    ),
}
class        __A       (	A_    ):







		'''simple docstring'''





		lowerCAmelCase			:					str   	=				"tapas"







		def __init__(	self						:  List[Any]			,_snake_case						:  Dict=30_522			,_snake_case						:  Union[str, Any]=768			,_snake_case						:  int=12			,_snake_case						:  Union[str, Any]=12			,_snake_case						:  Union[str, Any]=3_072			,_snake_case						:  List[Any]="gelu"			,_snake_case						:  Optional[int]=0.1			,_snake_case						:  Tuple=0.1			,_snake_case						:  List[Any]=1_024			,_snake_case						:  Any=[3, 256, 256, 2, 256, 256, 10]			,_snake_case						:  List[Any]=0.02			,_snake_case						:  Union[str, Any]=1e-12			,_snake_case						:  str=0			,_snake_case						:  Any=10.0			,_snake_case						:  int=0			,_snake_case						:  Optional[Any]=1.0			,_snake_case						:  List[str]=None			,_snake_case						:  Tuple=1.0			,_snake_case						:  Tuple=False			,_snake_case						:  List[Any]=None			,_snake_case						:  int=1.0			,_snake_case						:  List[Any]=1.0			,_snake_case						:  Optional[int]=False			,_snake_case						:  Optional[int]=False			,_snake_case						:  Optional[int]="ratio"			,_snake_case						:  Any=None			,_snake_case						:  Union[str, Any]=None			,_snake_case						:  List[str]=64			,_snake_case						:  Optional[Any]=32			,_snake_case						:  Optional[Any]=False			,_snake_case						:  Optional[int]=True			,_snake_case						:  Dict=False			,_snake_case						:  Tuple=False			,_snake_case						:  int=True			,_snake_case						:  List[str]=False			,_snake_case						:  Dict=None			,_snake_case						:  Optional[int]=None			,**_snake_case						:  int			,)					->		List[Any]:






				"""simple docstring"""



				super().__init__(pad_token_id=_snake_case			,**_snake_case						)

				# BERT hyperparameters (with updated max_position_embeddings and type_vocab_sizes)
				lowercase__   :    Optional[int]				=			vocab_size
				lowercase__   :    List[str]				=			hidden_size
				lowercase__   :    Any				=			num_hidden_layers
				lowercase__   :    Optional[Any]				=			num_attention_heads
				lowercase__   :    Optional[int]				=			hidden_act
				lowercase__   :    List[Any]				=			intermediate_size
				lowercase__   :    List[Any]				=			hidden_dropout_prob
				lowercase__   :    Dict				=			attention_probs_dropout_prob
				lowercase__   :    str				=			max_position_embeddings
				lowercase__   :    Dict				=			type_vocab_sizes
				lowercase__   :    Optional[Any]				=			initializer_range
				lowercase__   :    Dict				=			layer_norm_eps

				# Fine-tuning task hyperparameters
				lowercase__   :    Any				=			positive_label_weight
				lowercase__   :    int				=			num_aggregation_labels
				lowercase__   :    List[str]				=			aggregation_loss_weight
				lowercase__   :    Optional[int]				=			use_answer_as_supervision
				lowercase__   :    Optional[Any]				=			answer_loss_importance
				lowercase__   :    Union[str, Any]				=			use_normalized_answer_loss
				lowercase__   :    str				=			huber_loss_delta
				lowercase__   :    str				=			temperature
				lowercase__   :    int				=			aggregation_temperature
				lowercase__   :    List[Any]				=			use_gumbel_for_cells
				lowercase__   :    Tuple				=			use_gumbel_for_aggregation
				lowercase__   :    Union[str, Any]				=			average_approximation_function
				lowercase__   :    Union[str, Any]				=			cell_selection_preference
				lowercase__   :    Any				=			answer_loss_cutoff
				lowercase__   :    List[Any]				=			max_num_rows
				lowercase__   :    str				=			max_num_columns
				lowercase__   :    int				=			average_logits_per_cell
				lowercase__   :    str				=			select_one_column
				lowercase__   :    str				=			allow_empty_column_selection
				lowercase__   :    Any				=			init_cell_selection_weights_to_zero
				lowercase__   :    Optional[int]				=			reset_position_index_per_cell
				lowercase__   :    Union[str, Any]				=			disable_per_token_loss

				# Aggregation hyperparameters
				lowercase__   :    Optional[Any]				=			aggregation_labels
				lowercase__   :    List[Any]				=			no_aggregation_label_index

				if isinstance(self.aggregation_labels			,_snake_case						):
						lowercase__   :    Union[str, Any]				=			{int(_snake_case						): v for k, v in aggregation_labels.items()}
