



"""simple docstring"""

import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionTextToImagePipeline
from diffusers.utils.testing_utils import nightly, require_torch_gpu, torch_device


A_						:    Union[str, Any]				=      False
class     lowerCAmelCase__						(	unittest.TestCase ):


	'''simple docstring'''







	pass




@nightly
@require_torch_gpu
class     lowerCAmelCase__						(	unittest.TestCase ):


	'''simple docstring'''







	def 			_lowerCAmelCase     (      self				:       Any    )						->	List[str]:





					"""simple docstring"""


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





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





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Tuple							=     VersatileDiffusionTextToImagePipeline.from_pretrained('shi-labs/versatile-diffusion'    )
					# remove text_unet
					pipe.remove_unused_weights()
					pipe.to(_SCREAMING_SNAKE_CASE    )
					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     'A painting of a squirrel eating a burger '
					SCREAMING_SNAKE_CASE						:							List[str]							=     torch.manual_seed(0    )
					SCREAMING_SNAKE_CASE						:							int							=     pipe(
					    prompt=_SCREAMING_SNAKE_CASE     ,       generator=_SCREAMING_SNAKE_CASE     ,       guidance_scale=7.5     ,       num_inference_steps=2     ,       output_type='numpy'    ).images

					with tempfile.TemporaryDirectory() as tmpdirname:
									pipe.save_pretrained(_SCREAMING_SNAKE_CASE    )
									SCREAMING_SNAKE_CASE						:							Dict							=     VersatileDiffusionTextToImagePipeline.from_pretrained(_SCREAMING_SNAKE_CASE    )
					pipe.to(_SCREAMING_SNAKE_CASE    )
					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     generator.manual_seed(0    )
					SCREAMING_SNAKE_CASE						:							Tuple							=     pipe(
					    prompt=_SCREAMING_SNAKE_CASE     ,       generator=_SCREAMING_SNAKE_CASE     ,       guidance_scale=7.5     ,       num_inference_steps=2     ,       output_type='numpy'    ).images

					assert np.abs(image - new_image    ).sum() < 1E-5, "Models don't have the same forward pass"







	def 			_lowerCAmelCase     (      self				:       Dict    )						->	Dict:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							List[Any]							=     VersatileDiffusionTextToImagePipeline.from_pretrained(
					    'shi-labs/versatile-diffusion'     ,       torch_dtype=torch.floataa    )
					pipe.to(_SCREAMING_SNAKE_CASE    )
					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

					SCREAMING_SNAKE_CASE						:							Optional[int]							=     'A painting of a squirrel eating a burger '
					SCREAMING_SNAKE_CASE						:							Dict							=     torch.manual_seed(0    )
					SCREAMING_SNAKE_CASE						:							Any							=     pipe(
					    prompt=_SCREAMING_SNAKE_CASE     ,       generator=_SCREAMING_SNAKE_CASE     ,       guidance_scale=7.5     ,       num_inference_steps=50     ,       output_type='numpy'    ).images

					SCREAMING_SNAKE_CASE						:							Tuple							=     image[0, 253:256, 253:256, -1]

					assert image.shape == (1, 512, 512, 3)
					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     np.array([0.3_3_6_7, 0.3_1_6_9, 0.2_6_5_6, 0.3_8_7_0, 0.4_7_9_0, 0.3_7_9_6, 0.4_0_0_9, 0.4_8_7_8, 0.4_7_7_8]    )

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










"""simple docstring"""

from collections import OrderedDict
from typing import Mapping

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


A_						:    Any				=      logging.get_logger(__name__)

A_						:    Optional[int]				=      {
    'facebook/xmod-base': 'https://huggingface.co/facebook/xmod-base/resolve/main/config.json',
    'facebook/xmod-large-prenorm': 'https://huggingface.co/facebook/xmod-large-prenorm/resolve/main/config.json',
    'facebook/xmod-base-13-125k': 'https://huggingface.co/facebook/xmod-base-13-125k/resolve/main/config.json',
    'facebook/xmod-base-30-125k': 'https://huggingface.co/facebook/xmod-base-30-125k/resolve/main/config.json',
    'facebook/xmod-base-30-195k': 'https://huggingface.co/facebook/xmod-base-30-195k/resolve/main/config.json',
    'facebook/xmod-base-60-125k': 'https://huggingface.co/facebook/xmod-base-60-125k/resolve/main/config.json',
    'facebook/xmod-base-60-265k': 'https://huggingface.co/facebook/xmod-base-60-265k/resolve/main/config.json',
    'facebook/xmod-base-75-125k': 'https://huggingface.co/facebook/xmod-base-75-125k/resolve/main/config.json',
    'facebook/xmod-base-75-269k': 'https://huggingface.co/facebook/xmod-base-75-269k/resolve/main/config.json',
}
class     lowerCAmelCase__						(	_lowerCamelCase ):


	'''simple docstring'''







	_SCREAMING_SNAKE_CASE	:		int  						= '''xmod'''





	def __init__(      self				:       Union[str, Any]     ,       _SCREAMING_SNAKE_CASE				:       Tuple=30_522     ,       _SCREAMING_SNAKE_CASE				:       Tuple=768     ,       _SCREAMING_SNAKE_CASE				:       Optional[int]=12     ,       _SCREAMING_SNAKE_CASE				:       List[str]=12     ,       _SCREAMING_SNAKE_CASE				:       Optional[Any]=3_072     ,       _SCREAMING_SNAKE_CASE				:       Any="gelu"     ,       _SCREAMING_SNAKE_CASE				:       List[Any]=0.1     ,       _SCREAMING_SNAKE_CASE				:       Dict=0.1     ,       _SCREAMING_SNAKE_CASE				:       str=512     ,       _SCREAMING_SNAKE_CASE				:       Tuple=2     ,       _SCREAMING_SNAKE_CASE				:       Union[str, Any]=0.0_2     ,       _SCREAMING_SNAKE_CASE				:       Union[str, Any]=1E-1_2     ,       _SCREAMING_SNAKE_CASE				:       List[str]=1     ,       _SCREAMING_SNAKE_CASE				:       Optional[Any]=0     ,       _SCREAMING_SNAKE_CASE				:       Optional[Any]=2     ,       _SCREAMING_SNAKE_CASE				:       Union[str, Any]="absolute"     ,       _SCREAMING_SNAKE_CASE				:       Union[str, Any]=True     ,       _SCREAMING_SNAKE_CASE				:       Optional[Any]=None     ,       _SCREAMING_SNAKE_CASE				:       str=False     ,       _SCREAMING_SNAKE_CASE				:       Dict=2     ,       _SCREAMING_SNAKE_CASE				:       str=False     ,       _SCREAMING_SNAKE_CASE				:       int=True     ,       _SCREAMING_SNAKE_CASE				:       Union[str, Any]=True     ,       _SCREAMING_SNAKE_CASE				:       Tuple=("en_XX",)     ,       _SCREAMING_SNAKE_CASE				:       Optional[int]=None     ,       **_SCREAMING_SNAKE_CASE				:       List[Any]     ,       )						->	List[str]:





					"""simple docstring"""


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

					SCREAMING_SNAKE_CASE						:							Any							=     vocab_size
					SCREAMING_SNAKE_CASE						:							Any							=     hidden_size
					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     num_hidden_layers
					SCREAMING_SNAKE_CASE						:							str							=     num_attention_heads
					SCREAMING_SNAKE_CASE						:							List[Any]							=     hidden_act
					SCREAMING_SNAKE_CASE						:							Any							=     intermediate_size
					SCREAMING_SNAKE_CASE						:							Any							=     hidden_dropout_prob
					SCREAMING_SNAKE_CASE						:							str							=     attention_probs_dropout_prob
					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     max_position_embeddings
					SCREAMING_SNAKE_CASE						:							str							=     type_vocab_size
					SCREAMING_SNAKE_CASE						:							List[str]							=     initializer_range
					SCREAMING_SNAKE_CASE						:							Optional[int]							=     layer_norm_eps
					SCREAMING_SNAKE_CASE						:							Tuple							=     position_embedding_type
					SCREAMING_SNAKE_CASE						:							int							=     use_cache
					SCREAMING_SNAKE_CASE						:							Optional[int]							=     classifier_dropout
					SCREAMING_SNAKE_CASE						:							List[Any]							=     pre_norm
					SCREAMING_SNAKE_CASE						:							int							=     adapter_reduction_factor
					SCREAMING_SNAKE_CASE						:							List[Any]							=     adapter_layer_norm
					SCREAMING_SNAKE_CASE						:							Any							=     adapter_reuse_layer_norm
					SCREAMING_SNAKE_CASE						:							Any							=     ln_before_adapter
					SCREAMING_SNAKE_CASE						:							Tuple							=     list(_SCREAMING_SNAKE_CASE    )
					SCREAMING_SNAKE_CASE						:							List[Any]							=     default_language




class     lowerCAmelCase__						(	_lowerCamelCase ):


	'''simple docstring'''







	@property
	def 			_lowerCAmelCase     (      self				:       Union[str, Any]    )						->	Mapping[str, Mapping[int, str]]:





					"""simple docstring"""


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





