



"""simple docstring"""
import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


_lowerCAmelCase    :int     =					logging.get_logger(__name__)

enable_full_determinism()
class 							_UpperCAmelCase							( SCREAMING_SNAKE_CASE_		,SCREAMING_SNAKE_CASE_		,unittest.TestCase      ):





     '''simple docstring'''



     a__									=UNetaDModel
     a__									='''sample'''





     @property
     def  __lowerCAmelCase	(					self		)							->					str:
         _UpperCAmelCase							:  List[str]											=    4
         _UpperCAmelCase							:  str											=    3
         _UpperCAmelCase							:  List[str]											=    (3_2, 3_2)

         _UpperCAmelCase							:  Optional[Any]											=    floats_tensor((batch_size, num_channels) + sizes		).to(__UpperCAmelCase		)
         _UpperCAmelCase							:  Tuple											=    torch.tensor([1_0]		).to(__UpperCAmelCase		)

         return {"sample": noise, "timestep": time_step}





     @property
     def  __lowerCAmelCase	(					self		)							->					List[Any]:
         return (3, 3_2, 3_2)





     @property
     def  __lowerCAmelCase	(					self		)							->					int:
         return (3, 3_2, 3_2)





     def  __lowerCAmelCase	(					self		)							->					Dict:
         _UpperCAmelCase							:  List[str]											=    {
             """block_out_channels""": (3_2, 6_4),
             """down_block_types""": ("""DownBlock2D""", """AttnDownBlock2D"""),
             """up_block_types""": ("""AttnUpBlock2D""", """UpBlock2D"""),
             """attention_head_dim""": 3,
             """out_channels""": 3,
             """in_channels""": 3,
             """layers_per_block""": 2,
             """sample_size""": 3_2,
         }
         _UpperCAmelCase							:  str											=    self.dummy_input
         return init_dict, inputs_dict



class 							_UpperCAmelCase							( SCREAMING_SNAKE_CASE_		,SCREAMING_SNAKE_CASE_		,unittest.TestCase      ):





     '''simple docstring'''



     a__									=UNetaDModel
     a__									='''sample'''





     @property
     def  __lowerCAmelCase	(					self		)							->					Any:
         _UpperCAmelCase							:  Any											=    4
         _UpperCAmelCase							:  Dict											=    4
         _UpperCAmelCase							:  str											=    (3_2, 3_2)

         _UpperCAmelCase							:  Any											=    floats_tensor((batch_size, num_channels) + sizes		).to(__UpperCAmelCase		)
         _UpperCAmelCase							:  Union[str, Any]											=    torch.tensor([1_0]		).to(__UpperCAmelCase		)

         return {"sample": noise, "timestep": time_step}





     @property
     def  __lowerCAmelCase	(					self		)							->					Any:
         return (4, 3_2, 3_2)





     @property
     def  __lowerCAmelCase	(					self		)							->					Union[str, Any]:
         return (4, 3_2, 3_2)





     def  __lowerCAmelCase	(					self		)							->					Optional[Any]:
         _UpperCAmelCase							:  str											=    {
             """sample_size""": 3_2,
             """in_channels""": 4,
             """out_channels""": 4,
             """layers_per_block""": 2,
             """block_out_channels""": (3_2, 6_4),
             """attention_head_dim""": 3_2,
             """down_block_types""": ("""DownBlock2D""", """DownBlock2D"""),
             """up_block_types""": ("""UpBlock2D""", """UpBlock2D"""),
         }
         _UpperCAmelCase							:  List[str]											=    self.dummy_input
         return init_dict, inputs_dict





     def  __lowerCAmelCase	(					self		)							->					List[str]:
         _UpperCAmelCase							:  Tuple											=    UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''       ,	output_loading_info=__UpperCAmelCase		)

         self.assertIsNotNone(__UpperCAmelCase		)
         self.assertEqual(len(loading_info['''missing_keys''']		)       ,	0		)

         model.to(__UpperCAmelCase		)
         _UpperCAmelCase							:  Any											=    model(**self.dummy_input		).sample

         assert image is not None, "Make sure output is not None"





     @unittest.skipIf(torch_device != '''cuda'''       ,	'''This test is supposed to run on GPU'''		)
     def  __lowerCAmelCase	(					self		)							->					List[str]:
         _UpperCAmelCase							:  Any											=    UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''       ,	output_loading_info=__UpperCAmelCase		)
         model.to(__UpperCAmelCase		)
         _UpperCAmelCase							:  Optional[int]											=    model(**self.dummy_input		).sample

         assert image is not None, "Make sure output is not None"





     @unittest.skipIf(torch_device != '''cuda'''       ,	'''This test is supposed to run on GPU'''		)
     def  __lowerCAmelCase	(					self		)							->					Optional[Any]:
         # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
         _UpperCAmelCase							:  List[Any]											=    UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''       ,	output_loading_info=__UpperCAmelCase		)
         model_accelerate.to(__UpperCAmelCase		)
         model_accelerate.eval()

         _UpperCAmelCase							:  Optional[Any]											=    torch.randn(
             1       ,	model_accelerate.config.in_channels       ,	model_accelerate.config.sample_size       ,	model_accelerate.config.sample_size       ,	generator=torch.manual_seed(0		)       ,	)
         _UpperCAmelCase							:  List[str]											=    noise.to(__UpperCAmelCase		)
         _UpperCAmelCase							:  int											=    torch.tensor([1_0] * noise.shape[0]		).to(__UpperCAmelCase		)

         _UpperCAmelCase							:  List[Any]											=    model_accelerate(__UpperCAmelCase       ,	__UpperCAmelCase		)["""sample"""]

         # two models don't need to stay in the device at the same time
         del model_accelerate
         torch.cuda.empty_cache()
         gc.collect()

         _UpperCAmelCase							:  Optional[Any]											=    UNetaDModel.from_pretrained(
             '''fusing/unet-ldm-dummy-update'''       ,	output_loading_info=__UpperCAmelCase       ,	low_cpu_mem_usage=__UpperCAmelCase		)
         model_normal_load.to(__UpperCAmelCase		)
         model_normal_load.eval()
         _UpperCAmelCase							:  Dict											=    model_normal_load(__UpperCAmelCase       ,	__UpperCAmelCase		)["""sample"""]

         assert torch_all_close(__UpperCAmelCase       ,	__UpperCAmelCase       ,	rtol=1E-3		)





     def  __lowerCAmelCase	(					self		)							->					str:
         _UpperCAmelCase							:  Optional[int]											=    UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''		)
         model.eval()
         model.to(__UpperCAmelCase		)

         _UpperCAmelCase							:  int											=    torch.randn(
             1       ,	model.config.in_channels       ,	model.config.sample_size       ,	model.config.sample_size       ,	generator=torch.manual_seed(0		)       ,	)
         _UpperCAmelCase							:  Optional[int]											=    noise.to(__UpperCAmelCase		)
         _UpperCAmelCase							:  str											=    torch.tensor([1_0] * noise.shape[0]		).to(__UpperCAmelCase		)

         with torch.no_grad():
             _UpperCAmelCase							:  List[str]											=    model(__UpperCAmelCase       ,	__UpperCAmelCase		).sample

         _UpperCAmelCase							:  Optional[Any]											=    output[0, -1, -3:, -3:].flatten().cpu()
         # fmt: off
         _UpperCAmelCase							:  int											=    torch.tensor([-1_3.3_2_5_8, -2_0.1_1_0_0, -1_5.9_8_7_3, -1_7.6_6_1_7, -2_3.0_5_9_6, -1_7.9_4_1_9, -1_3.3_6_7_5, -1_6.1_8_8_9, -1_2.3_8_0_0]		)
         # fmt: on

         self.assertTrue(torch_all_close(__UpperCAmelCase       ,	__UpperCAmelCase       ,	rtol=1E-3		)		)



class 							_UpperCAmelCase							( SCREAMING_SNAKE_CASE_		,SCREAMING_SNAKE_CASE_		,unittest.TestCase      ):





     '''simple docstring'''



     a__									=UNetaDModel
     a__									='''sample'''





     @property
     def  __lowerCAmelCase	(					self       ,	A=(3_2, 3_2)		)							->					Dict:
         _UpperCAmelCase							:  Any											=    4
         _UpperCAmelCase							:  Optional[Any]											=    3

         _UpperCAmelCase							:  Union[str, Any]											=    floats_tensor((batch_size, num_channels) + sizes		).to(__UpperCAmelCase		)
         _UpperCAmelCase							:  str											=    torch.tensor(batch_size * [1_0]		).to(dtype=torch.intaa       ,	device=__UpperCAmelCase		)

         return {"sample": noise, "timestep": time_step}





     @property
     def  __lowerCAmelCase	(					self		)							->					Dict:
         return (3, 3_2, 3_2)





     @property
     def  __lowerCAmelCase	(					self		)							->					Any:
         return (3, 3_2, 3_2)





     def  __lowerCAmelCase	(					self		)							->					Optional[int]:
         _UpperCAmelCase							:  Union[str, Any]											=    {
             """block_out_channels""": [3_2, 6_4, 6_4, 6_4],
             """in_channels""": 3,
             """layers_per_block""": 1,
             """out_channels""": 3,
             """time_embedding_type""": """fourier""",
             """norm_eps""": 1E-6,
             """mid_block_scale_factor""": math.sqrt(2.0		),
             """norm_num_groups""": None,
             """down_block_types""": [
                 """SkipDownBlock2D""",
                 """AttnSkipDownBlock2D""",
                 """SkipDownBlock2D""",
                 """SkipDownBlock2D""",
             ],
             """up_block_types""": [
                 """SkipUpBlock2D""",
                 """SkipUpBlock2D""",
                 """AttnSkipUpBlock2D""",
                 """SkipUpBlock2D""",
             ],
         }
         _UpperCAmelCase							:  Tuple											=    self.dummy_input
         return init_dict, inputs_dict





     @slow
     def  __lowerCAmelCase	(					self		)							->					str:
         _UpperCAmelCase							:  str											=    UNetaDModel.from_pretrained('''google/ncsnpp-celebahq-256'''       ,	output_loading_info=__UpperCAmelCase		)
         self.assertIsNotNone(__UpperCAmelCase		)
         self.assertEqual(len(loading_info['''missing_keys''']		)       ,	0		)

         model.to(__UpperCAmelCase		)
         _UpperCAmelCase							:  Union[str, Any]											=    self.dummy_input
         _UpperCAmelCase							:  List[Any]											=    floats_tensor((4, 3) + (2_5_6, 2_5_6)		).to(__UpperCAmelCase		)
         _UpperCAmelCase							:  Tuple											=    noise
         _UpperCAmelCase							:  Union[str, Any]											=    model(**__UpperCAmelCase		)

         assert image is not None, "Make sure output is not None"





     @slow
     def  __lowerCAmelCase	(					self		)							->					str:
         _UpperCAmelCase							:  Optional[Any]											=    UNetaDModel.from_pretrained('''google/ncsnpp-celebahq-256'''		)
         model.to(__UpperCAmelCase		)

         _UpperCAmelCase							:  Optional[int]											=    4
         _UpperCAmelCase							:  Any											=    3
         _UpperCAmelCase							:  Optional[int]											=    (2_5_6, 2_5_6)

         _UpperCAmelCase							:  int											=    torch.ones((batch_size, num_channels) + sizes		).to(__UpperCAmelCase		)
         _UpperCAmelCase							:  List[Any]											=    torch.tensor(batch_size * [1E-4]		).to(__UpperCAmelCase		)

         with torch.no_grad():
             _UpperCAmelCase							:  Dict											=    model(__UpperCAmelCase       ,	__UpperCAmelCase		).sample

         _UpperCAmelCase							:  int											=    output[0, -3:, -3:, -1].flatten().cpu()
         # fmt: off
         _UpperCAmelCase							:  Any											=    torch.tensor([-4_8_4_2.8_6_9_1, -6_4_9_9.6_6_3_1, -3_8_0_0.1_9_5_3, -7_9_7_8.2_6_8_6, -1_0_9_8_0.7_1_2_9, -2_0_0_2_8.8_5_3_5, 8_1_4_8.2_8_2_2, 2_3_4_2.2_9_0_5, 5_6_7.7_6_0_8]		)
         # fmt: on

         self.assertTrue(torch_all_close(__UpperCAmelCase       ,	__UpperCAmelCase       ,	rtol=1E-2		)		)





     def  __lowerCAmelCase	(					self		)							->					Tuple:
         _UpperCAmelCase							:  Dict											=    UNetaDModel.from_pretrained('''fusing/ncsnpp-ffhq-ve-dummy-update'''		)
         model.to(__UpperCAmelCase		)

         _UpperCAmelCase							:  int											=    4
         _UpperCAmelCase							:  List[str]											=    3
         _UpperCAmelCase							:  Union[str, Any]											=    (3_2, 3_2)

         _UpperCAmelCase							:  int											=    torch.ones((batch_size, num_channels) + sizes		).to(__UpperCAmelCase		)
         _UpperCAmelCase							:  Optional[Any]											=    torch.tensor(batch_size * [1E-4]		).to(__UpperCAmelCase		)

         with torch.no_grad():
             _UpperCAmelCase							:  Any											=    model(__UpperCAmelCase       ,	__UpperCAmelCase		).sample

         _UpperCAmelCase							:  str											=    output[0, -3:, -3:, -1].flatten().cpu()
         # fmt: off
         _UpperCAmelCase							:  str											=    torch.tensor([-0.0_325, -0.0_900, -0.0_869, -0.0_332, -0.0_725, -0.0_270, -0.0_101, 0.0_227, 0.0_256]		)
         # fmt: on

         self.assertTrue(torch_all_close(__UpperCAmelCase       ,	__UpperCAmelCase       ,	rtol=1E-2		)		)





     def  __lowerCAmelCase	(					self		)							->					str:
         # not required for this model
         pass







'''simple docstring'''





from collections import OrderedDict
from typing import Mapping

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


_lowerCAmelCase      						=	logging.get_logger(__name__)

_lowerCAmelCase      						=	{
    '''facebook/xlm-roberta-xl''': '''https://huggingface.co/facebook/xlm-roberta-xl/resolve/main/config.json''',
    '''facebook/xlm-roberta-xxl''': '''https://huggingface.co/facebook/xlm-roberta-xxl/resolve/main/config.json''',
    # See all XLM-RoBERTa-XL models at https://huggingface.co/models?filter=xlm-roberta-xl
}






class        lowerCAmelCase_(				SCREAMING_SNAKE_CASE_     ):
  '''simple docstring'''




  __lowercase							:      int				=      '''xlm-roberta-xl'''




  def __init__(    self  ,__UpperCAmelCase=25_0880  ,__UpperCAmelCase=2560  ,__UpperCAmelCase=36  ,__UpperCAmelCase=32  ,__UpperCAmelCase=1_0240  ,__UpperCAmelCase="gelu"  ,__UpperCAmelCase=0.1  ,__UpperCAmelCase=0.1  ,__UpperCAmelCase=514  ,__UpperCAmelCase=1  ,__UpperCAmelCase=0.0_2  ,__UpperCAmelCase=1E-05  ,__UpperCAmelCase=1  ,__UpperCAmelCase=0  ,__UpperCAmelCase=2  ,__UpperCAmelCase="absolute"  ,__UpperCAmelCase=True  ,__UpperCAmelCase=None  ,**__UpperCAmelCase  ,)    ->		str:
      super().__init__(pad_token_id=__UpperCAmelCase  ,bos_token_id=__UpperCAmelCase  ,eos_token_id=__UpperCAmelCase  ,**__UpperCAmelCase      )
      lowerCAmelCase__ :  List[Any]							      =					vocab_size
      lowerCAmelCase__ :  int							      =					hidden_size
      lowerCAmelCase__ :  int							      =					num_hidden_layers
      lowerCAmelCase__ :  str							      =					num_attention_heads
      lowerCAmelCase__ :  int							      =					hidden_act
      lowerCAmelCase__ :  Dict							      =					intermediate_size
      lowerCAmelCase__ :  List[Any]							      =					hidden_dropout_prob
      lowerCAmelCase__ :  str							      =					attention_probs_dropout_prob
      lowerCAmelCase__ :  Optional[int]							      =					max_position_embeddings
      lowerCAmelCase__ :  List[str]							      =					type_vocab_size
      lowerCAmelCase__ :  List[Any]							      =					initializer_range
      lowerCAmelCase__ :  Tuple							      =					layer_norm_eps
      lowerCAmelCase__ :  int							      =					position_embedding_type
      lowerCAmelCase__ :  Optional[Any]							      =					use_cache
      lowerCAmelCase__ :  Optional[Any]							      =					classifier_dropout







class        lowerCAmelCase_(				SCREAMING_SNAKE_CASE_     ):
  '''simple docstring'''




  @property
  def 							UpperCAmelCase_	(    self      )    ->		Mapping[str, Mapping[int, str]]:
      if self.task == "multiple-choice":
          lowerCAmelCase__ :  Dict							      =					{0: """batch""", 1: """choice""", 2: """sequence"""}
      else:
          lowerCAmelCase__ :  Any							      =					{0: """batch""", 1: """sequence"""}
      return OrderedDict(
          [
              ("""input_ids""", dynamic_axis),
              ("""attention_mask""", dynamic_axis),
          ]      )


