





import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
  import jax.numpy as jnp

  from transformers.models.albert.modeling_flax_albert import (
      FlaxAlbertForMaskedLM,
      FlaxAlbertForMultipleChoice,
      FlaxAlbertForPreTraining,
      FlaxAlbertForQuestionAnswering,
      FlaxAlbertForSequenceClassification,
      FlaxAlbertForTokenClassification,
      FlaxAlbertModel,
  )






class 							__lowercase			(     unittest.TestCase       ):






      def __init__(		self			,							lowercase_			,							lowercase_=1_3			,							lowercase_=7			,							lowercase_=True			,							lowercase_=True			,							lowercase_=True			,							lowercase_=True			,							lowercase_=9_9			,							lowercase_=3_2			,							lowercase_=5			,							lowercase_=4			,							lowercase_=3_7			,							lowercase_="gelu"			,							lowercase_=0.1			,							lowercase_=0.1			,							lowercase_=5_1_2			,							lowercase_=1_6			,							lowercase_=2			,							lowercase_=0.02			,							lowercase_=4			,							)     ->       List[str]:
       __snake_case         =			parent
       __snake_case         =			batch_size
       __snake_case         =			seq_length
       __snake_case         =			is_training
       __snake_case         =			use_attention_mask
       __snake_case         =			use_token_type_ids
       __snake_case         =			use_labels
       __snake_case         =			vocab_size
       __snake_case         =			hidden_size
       __snake_case         =			num_hidden_layers
       __snake_case         =			num_attention_heads
       __snake_case         =			intermediate_size
       __snake_case         =			hidden_act
       __snake_case         =			hidden_dropout_prob
       __snake_case         =			attention_probs_dropout_prob
       __snake_case         =			max_position_embeddings
       __snake_case         =			type_vocab_size
       __snake_case         =			type_sequence_label_size
       __snake_case         =			initializer_range
       __snake_case         =			num_choices






      def 			_a						(		self)     ->       Optional[int]:
       __snake_case         =			ids_tensor([self.batch_size, self.seq_length]			,							self.vocab_size)

       __snake_case         =			None
       if self.use_attention_mask:
        __snake_case         =			random_attention_mask([self.batch_size, self.seq_length])

       __snake_case         =			None
       if self.use_token_type_ids:
        __snake_case         =			ids_tensor([self.batch_size, self.seq_length]			,							self.type_vocab_size)

       __snake_case         =			AlbertConfig(
           vocab_size=self.vocab_size			,							hidden_size=self.hidden_size			,							num_hidden_layers=self.num_hidden_layers			,							num_attention_heads=self.num_attention_heads			,							intermediate_size=self.intermediate_size			,							hidden_act=self.hidden_act			,							hidden_dropout_prob=self.hidden_dropout_prob			,							attention_probs_dropout_prob=self.attention_probs_dropout_prob			,							max_position_embeddings=self.max_position_embeddings			,							type_vocab_size=self.type_vocab_size			,							is_decoder=_UpperCAmelCase			,							initializer_range=self.initializer_range			,							)

       return config, input_ids, token_type_ids, attention_mask







      def 			_a						(		self)     ->       Tuple:
       __snake_case         =			self.prepare_config_and_inputs()
       __snake_case   ,    __snake_case   ,    __snake_case   ,    __snake_case         =			config_and_inputs
       __snake_case         =			{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
       return config, inputs_dict






@require_flax
class 							__lowercase			(     snake_case__   ,							unittest.TestCase       ):
      __UpperCAmelCase		      =						(
          (
              FlaxAlbertModel,
              FlaxAlbertForPreTraining,
              FlaxAlbertForMaskedLM,
              FlaxAlbertForMultipleChoice,
              FlaxAlbertForQuestionAnswering,
              FlaxAlbertForSequenceClassification,
              FlaxAlbertForTokenClassification,
              FlaxAlbertForQuestionAnswering,
          )
          if is_flax_available()
          else ()
      )






      def 			_a						(		self)     ->       Dict:
       __snake_case         =			FlaxAlbertModelTester(self)







      @slow
      def 			_a						(		self)     ->       Any:
       for model_class_name in self.all_model_classes:
        __snake_case         =			model_class_name.from_pretrained('albert-base-v2')
        __snake_case         =			model(np.ones((1, 1)))
        self.assertIsNotNone(_UpperCAmelCase)






@require_flax
class 							__lowercase			(     unittest.TestCase       ):






      @slow
      def 			_a						(		self)     ->       Union[str, Any]:
       __snake_case         =			FlaxAlbertModel.from_pretrained('albert-base-v2')
       __snake_case         =			np.array([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]])
       __snake_case         =			np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
       __snake_case         =			model(_UpperCAmelCase			,							attention_mask=_UpperCAmelCase)[0]
       __snake_case         =			(1, 1_1, 7_6_8)
       self.assertEqual(output.shape			,							_UpperCAmelCase)
       __snake_case         =			np.array(
           [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]])

       self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]			,							_UpperCAmelCase			,							atol=1e-4))






'''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


SCREAMING_SNAKE_CASE__     :						List[Any]	=  logging.get_logger(__name__)

enable_full_determinism()

class 							a__( snake_case__   ,      snake_case__   ,      unittest.TestCase     ):
			a_	: Dict									=      UNetaDModel
			a_	: List[Any]									=      '''sample'''






			@property
			def 						_lowercase  (     self						)							->							Tuple:
					snake_case__	     =4
					snake_case__	     =3
					snake_case__	     =(32, 32)

					snake_case__	     =floats_tensor((batch_size, num_channels) + sizes						).to(_UpperCAmelCase						)
					snake_case__	     =torch.tensor([10]						).to(_UpperCAmelCase						)

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






			@property
			def 						_lowercase  (     self						)							->							Optional[int]:
					return (3, 32, 32)






			@property
			def 						_lowercase  (     self						)							->							Optional[int]:
					return (3, 32, 32)






			def 						_lowercase  (     self						)							->							Union[str, Any]:
					snake_case__	     ={
					    'block_out_channels': (32, 64),
					    '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': 32,
					}
					snake_case__	     =self.dummy_input
					return init_dict, inputs_dict





class 							a__( snake_case__   ,      snake_case__   ,      unittest.TestCase     ):
			a_	: Union[str, Any]									=      UNetaDModel
			a_	: Optional[Any]									=      '''sample'''






			@property
			def 						_lowercase  (     self						)							->							Union[str, Any]:
					snake_case__	     =4
					snake_case__	     =4
					snake_case__	     =(32, 32)

					snake_case__	     =floats_tensor((batch_size, num_channels) + sizes						).to(_UpperCAmelCase						)
					snake_case__	     =torch.tensor([10]						).to(_UpperCAmelCase						)

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






			@property
			def 						_lowercase  (     self						)							->							Optional[int]:
					return (4, 32, 32)






			@property
			def 						_lowercase  (     self						)							->							Dict:
					return (4, 32, 32)






			def 						_lowercase  (     self						)							->							str:
					snake_case__	     ={
					    'sample_size': 32,
					    'in_channels': 4,
					    'out_channels': 4,
					    'layers_per_block': 2,
					    'block_out_channels': (32, 64),
					    'attention_head_dim': 32,
					    'down_block_types': ('DownBlock2D', 'DownBlock2D'),
					    'up_block_types': ('UpBlock2D', 'UpBlock2D'),
					}
					snake_case__	     =self.dummy_input
					return init_dict, inputs_dict






			def 						_lowercase  (     self						)							->							Dict:
					snake_case__     ,	snake_case__	     =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						)
					snake_case__	     =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 						_lowercase  (     self						)							->							Optional[Any]:
					snake_case__     ,	snake_case__	     =UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'					,   output_loading_info=_UpperCAmelCase						)
					model.to(_UpperCAmelCase						)
					snake_case__	     =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 						_lowercase  (     self						)							->							Optional[Any]:
					# by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
					snake_case__     ,	snake_case__	     =UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'					,   output_loading_info=_UpperCAmelCase						)
					model_accelerate.to(_UpperCAmelCase						)
					model_accelerate.eval()

					snake_case__	     =torch.randn(
					    1					,   model_accelerate.config.in_channels					,   model_accelerate.config.sample_size					,   model_accelerate.config.sample_size					,   generator=torch.manual_seed(0						)					,   )
					snake_case__	     =noise.to(_UpperCAmelCase						)
					snake_case__	     =torch.tensor([10] * noise.shape[0]						).to(_UpperCAmelCase						)

					snake_case__	     =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()

					snake_case__     ,	snake_case__	     =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()
					snake_case__	     =model_normal_load(_UpperCAmelCase					,   _UpperCAmelCase						)['sample']

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






			def 						_lowercase  (     self						)							->							Optional[Any]:
					snake_case__	     =UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'						)
					model.eval()
					model.to(_UpperCAmelCase						)

					snake_case__	     =torch.randn(
					    1					,   model.config.in_channels					,   model.config.sample_size					,   model.config.sample_size					,   generator=torch.manual_seed(0						)					,   )
					snake_case__	     =noise.to(_UpperCAmelCase						)
					snake_case__	     =torch.tensor([10] * noise.shape[0]						).to(_UpperCAmelCase						)

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

					snake_case__	     =output[0, -1, -3:, -3:].flatten().cpu()
					# fmt: off
					snake_case__	     =torch.tensor([-13.3_258, -20.1_100, -15.9_873, -17.6_617, -23.0_596, -17.9_419, -13.3_675, -16.1_889, -12.3_800]						)
					# fmt: on

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





class 							a__( snake_case__   ,      snake_case__   ,      unittest.TestCase     ):
			a_	: List[str]									=      UNetaDModel
			a_	: Optional[int]									=      '''sample'''






			@property
			def 						_lowercase  (     self					,   _UpperCAmelCase=(32, 32)						)							->							Tuple:
					snake_case__	     =4
					snake_case__	     =3

					snake_case__	     =floats_tensor((batch_size, num_channels) + sizes						).to(_UpperCAmelCase						)
					snake_case__	     =torch.tensor(batch_size * [10]						).to(dtype=torch.intaa					,   device=_UpperCAmelCase						)

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






			@property
			def 						_lowercase  (     self						)							->							Union[str, Any]:
					return (3, 32, 32)






			@property
			def 						_lowercase  (     self						)							->							Optional[Any]:
					return (3, 32, 32)






			def 						_lowercase  (     self						)							->							str:
					snake_case__	     ={
					    'block_out_channels': [32, 64, 64, 64],
					    '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',
					    ],
					}
					snake_case__	     =self.dummy_input
					return init_dict, inputs_dict






			@slow
			def 						_lowercase  (     self						)							->							List[Any]:
					snake_case__     ,	snake_case__	     =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						)
					snake_case__	     =self.dummy_input
					snake_case__	     =floats_tensor((4, 3) + (256, 256)						).to(_UpperCAmelCase						)
					snake_case__	     =noise
					snake_case__	     =model(**_UpperCAmelCase						)

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






			@slow
			def 						_lowercase  (     self						)							->							Union[str, Any]:
					snake_case__	     =UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'						)
					model.to(_UpperCAmelCase						)

					snake_case__	     =4
					snake_case__	     =3
					snake_case__	     =(256, 256)

					snake_case__	     =torch.ones((batch_size, num_channels) + sizes						).to(_UpperCAmelCase						)
					snake_case__	     =torch.tensor(batch_size * [1E-4]						).to(_UpperCAmelCase						)

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

					snake_case__	     =output[0, -3:, -3:, -1].flatten().cpu()
					# fmt: off
					snake_case__	     =torch.tensor([-4_842.8_691, -6_499.6_631, -3_800.1_953, -7_978.2_686, -10_980.7_129, -20_028.8_535, 8_148.2_822, 2_342.2_905, 567.7_608]						)
					# fmt: on

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






			def 						_lowercase  (     self						)							->							List[Any]:
					snake_case__	     =UNetaDModel.from_pretrained('fusing/ncsnpp-ffhq-ve-dummy-update'						)
					model.to(_UpperCAmelCase						)

					snake_case__	     =4
					snake_case__	     =3
					snake_case__	     =(32, 32)

					snake_case__	     =torch.ones((batch_size, num_channels) + sizes						).to(_UpperCAmelCase						)
					snake_case__	     =torch.tensor(batch_size * [1E-4]						).to(_UpperCAmelCase						)

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

					snake_case__	     =output[0, -3:, -3:, -1].flatten().cpu()
					# fmt: off
					snake_case__	     =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 						_lowercase  (     self						)							->							Optional[Any]:
					# not required for this model
					pass




