
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_							=     logging.get_logger(__name__)

enable_full_determinism()




class 	__A(				_lowerCAmelCase				,      _lowerCAmelCase				,      unittest.TestCase ):

  """simple docstring"""
  SCREAMING_SNAKE_CASE__			    =			UNetaDModel
  SCREAMING_SNAKE_CASE__			    =			"""sample"""







  @property
  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				4
       UpperCamelCase__				=				3
       UpperCamelCase__				=				(32, 32)

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

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







  @property
  def        UpperCAmelCase_					(self						):
       return (3, 32, 32)







  @property
  def        UpperCAmelCase_					(self						):
       return (3, 32, 32)







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				{
           """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,
       }
       UpperCamelCase__				=				self.dummy_input
       return init_dict, inputs_dict






class 	__A(				_lowerCAmelCase				,      _lowerCAmelCase				,      unittest.TestCase ):

  """simple docstring"""
  SCREAMING_SNAKE_CASE__			    =			UNetaDModel
  SCREAMING_SNAKE_CASE__			    =			"""sample"""







  @property
  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				4
       UpperCamelCase__				=				4
       UpperCamelCase__				=				(32, 32)

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

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







  @property
  def        UpperCAmelCase_					(self						):
       return (4, 32, 32)







  @property
  def        UpperCAmelCase_					(self						):
       return (4, 32, 32)







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				{
           """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"""),
       }
       UpperCamelCase__				=				self.dummy_input
       return init_dict, inputs_dict







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""       ,			output_loading_info=SCREAMING_SNAKE_CASE_						)

       self.assertIsNotNone(SCREAMING_SNAKE_CASE_						)
       self.assertEqual(len(loading_info["""missing_keys"""]						)       ,			0						)

       model.to(SCREAMING_SNAKE_CASE_						)
       UpperCamelCase__				=				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        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""       ,			output_loading_info=SCREAMING_SNAKE_CASE_						)
       model.to(SCREAMING_SNAKE_CASE_						)
       UpperCamelCase__				=				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        UpperCAmelCase_					(self						):
       # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
       UpperCamelCase__				=				UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""       ,			output_loading_info=SCREAMING_SNAKE_CASE_						)
       model_accelerate.to(SCREAMING_SNAKE_CASE_						)
       model_accelerate.eval()

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

       UpperCamelCase__				=				model_accelerate(SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_						)["""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__				=				UNetaDModel.from_pretrained(
           """fusing/unet-ldm-dummy-update"""       ,			output_loading_info=SCREAMING_SNAKE_CASE_       ,			low_cpu_mem_usage=SCREAMING_SNAKE_CASE_						)
       model_normal_load.to(SCREAMING_SNAKE_CASE_						)
       model_normal_load.eval()
       UpperCamelCase__				=				model_normal_load(SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_						)["""sample"""]

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







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""						)
       model.eval()
       model.to(SCREAMING_SNAKE_CASE_						)

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

       with torch.no_grad():
            UpperCamelCase__				=				model(SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_						).sample

       UpperCamelCase__				=				output[0, -1, -3:, -3:].flatten().cpu()
       # fmt: off
       UpperCamelCase__				=				torch.tensor([-13.32_58, -20.11_00, -15.98_73, -17.66_17, -23.05_96, -17.94_19, -13.36_75, -16.18_89, -12.38_00]						)
       # fmt: on

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






class 	__A(				_lowerCAmelCase				,      _lowerCAmelCase				,      unittest.TestCase ):

  """simple docstring"""
  SCREAMING_SNAKE_CASE__			    =			UNetaDModel
  SCREAMING_SNAKE_CASE__			    =			"""sample"""







  @property
  def        UpperCAmelCase_					(self       ,			SCREAMING_SNAKE_CASE_=(32, 32)						):
       UpperCamelCase__				=				4
       UpperCamelCase__				=				3

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

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







  @property
  def        UpperCAmelCase_					(self						):
       return (3, 32, 32)







  @property
  def        UpperCAmelCase_					(self						):
       return (3, 32, 32)







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				{
           """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""",
           ],
       }
       UpperCamelCase__				=				self.dummy_input
       return init_dict, inputs_dict







  @slow
  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""       ,			output_loading_info=SCREAMING_SNAKE_CASE_						)
       self.assertIsNotNone(SCREAMING_SNAKE_CASE_						)
       self.assertEqual(len(loading_info["""missing_keys"""]						)       ,			0						)

       model.to(SCREAMING_SNAKE_CASE_						)
       UpperCamelCase__				=				self.dummy_input
       UpperCamelCase__				=				floats_tensor((4, 3) + (2_56, 2_56)						).to(SCREAMING_SNAKE_CASE_						)
       UpperCamelCase__				=				noise
       UpperCamelCase__				=				model(**SCREAMING_SNAKE_CASE_						)

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







  @slow
  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""						)
       model.to(SCREAMING_SNAKE_CASE_						)

       UpperCamelCase__				=				4
       UpperCamelCase__				=				3
       UpperCamelCase__				=				(2_56, 2_56)

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

       with torch.no_grad():
            UpperCamelCase__				=				model(SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_						).sample

       UpperCamelCase__				=				output[0, -3:, -3:, -1].flatten().cpu()
       # fmt: off
       UpperCamelCase__				=				torch.tensor([-48_42.86_91, -64_99.66_31, -38_00.19_53, -79_78.26_86, -1_09_80.71_29, -2_00_28.85_35, 81_48.28_22, 23_42.29_05, 5_67.76_08]						)
       # fmt: on

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







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				UNetaDModel.from_pretrained("""fusing/ncsnpp-ffhq-ve-dummy-update"""						)
       model.to(SCREAMING_SNAKE_CASE_						)

       UpperCamelCase__				=				4
       UpperCamelCase__				=				3
       UpperCamelCase__				=				(32, 32)

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

       with torch.no_grad():
            UpperCamelCase__				=				model(SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_						).sample

       UpperCamelCase__				=				output[0, -3:, -3:, -1].flatten().cpu()
       # fmt: off
       UpperCamelCase__				=				torch.tensor([-0.0325, -0.0900, -0.0869, -0.0332, -0.0725, -0.0270, -0.0101, 0.0227, 0.0256]						)
       # fmt: on

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







  def        UpperCAmelCase_					(self						):
       # not required for this model
       pass










"""simple docstring"""
import unittest

import numpy as np

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

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
						from transformers.models.bert.modeling_flax_bert import (
						    FlaxBertForMaskedLM,
						    FlaxBertForMultipleChoice,
						    FlaxBertForNextSentencePrediction,
						    FlaxBertForPreTraining,
						    FlaxBertForQuestionAnswering,
						    FlaxBertForSequenceClassification,
						    FlaxBertForTokenClassification,
						    FlaxBertModel,
						)
class 				lowerCamelCase  (			unittest.TestCase				):






					'''simple docstring'''





					def __init__(       self:							List[Any]						,     snake_case:							List[str]						,     snake_case:							Optional[Any]=13						,     snake_case:							List[str]=7						,     snake_case:							Dict=True						,     snake_case:							List[str]=True						,     snake_case:							Optional[int]=True						,     snake_case:							Any=True						,     snake_case:							Optional[Any]=99						,     snake_case:							Tuple=32						,     snake_case:							Tuple=5						,     snake_case:							Dict=4						,     snake_case:							Optional[Any]=37						,     snake_case:							Union[str, Any]="gelu"						,     snake_case:							Tuple=0.1						,     snake_case:							List[Any]=0.1						,     snake_case:							List[str]=512						,     snake_case:							Optional[int]=16						,     snake_case:							int=2						,     snake_case:							List[Any]=0.0_2						,     snake_case:							Union[str, Any]=4						,     )   ->	List[str]:
											snake_case_ :Dict						=			parent
											snake_case_ :Any						=			batch_size
											snake_case_ :Any						=			seq_length
											snake_case_ :List[str]						=			is_training
											snake_case_ :Optional[Any]						=			use_attention_mask
											snake_case_ :Dict						=			use_token_type_ids
											snake_case_ :Union[str, Any]						=			use_labels
											snake_case_ :str						=			vocab_size
											snake_case_ :int						=			hidden_size
											snake_case_ :List[str]						=			num_hidden_layers
											snake_case_ :Dict						=			num_attention_heads
											snake_case_ :Any						=			intermediate_size
											snake_case_ :Tuple						=			hidden_act
											snake_case_ :int						=			hidden_dropout_prob
											snake_case_ :Optional[Any]						=			attention_probs_dropout_prob
											snake_case_ :Any						=			max_position_embeddings
											snake_case_ :Union[str, Any]						=			type_vocab_size
											snake_case_ :Optional[int]						=			type_sequence_label_size
											snake_case_ :Union[str, Any]						=			initializer_range
											snake_case_ :Tuple						=			num_choices






					def    lowerCAmelCase_				(       self:							Tuple						)   ->	str:
											snake_case_ :Optional[Any]						=			ids_tensor([self.batch_size, self.seq_length]						,     self.vocab_size						)

											snake_case_ :Union[str, Any]						=			None
											if self.use_attention_mask:
																	snake_case_ :str						=			random_attention_mask([self.batch_size, self.seq_length]						)

											snake_case_ :Any						=			None
											if self.use_token_type_ids:
																	snake_case_ :List[str]						=			ids_tensor([self.batch_size, self.seq_length]						,     self.type_vocab_size						)

											snake_case_ :int						=			BertConfig(
											    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=snake_case						,     initializer_range=self.initializer_range						,     )

											return config, input_ids, token_type_ids, attention_mask






					def    lowerCAmelCase_				(       self:							Optional[int]						)   ->	int:
											snake_case_ :str						=			self.prepare_config_and_inputs()
											snake_case_,       snake_case_,       snake_case_,       snake_case_ :Optional[int]						=			config_and_inputs
											snake_case_ :Union[str, Any]						=			{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
											return config, inputs_dict






					def    lowerCAmelCase_				(       self:							Optional[Any]						)   ->	Any:
											snake_case_ :int						=			self.prepare_config_and_inputs()
											snake_case_,       snake_case_,       snake_case_,       snake_case_ :Dict						=			config_and_inputs

											snake_case_ :Union[str, Any]						=			True
											snake_case_ :Optional[int]						=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
											snake_case_ :Tuple						=			ids_tensor([self.batch_size, self.seq_length]						,     vocab_size=2						)

											return (
											    config,
											    input_ids,
											    attention_mask,
											    encoder_hidden_states,
											    encoder_attention_mask,
											)

@require_flax
class 				lowerCamelCase  (			_lowerCAmelCase    ,						unittest.TestCase				):






					'''simple docstring'''





					_A		:			List[str]						=     True

					_A		:			Dict						=     (
					    (
					        FlaxBertModel,
					        FlaxBertForPreTraining,
					        FlaxBertForMaskedLM,
					        FlaxBertForMultipleChoice,
					        FlaxBertForQuestionAnswering,
					        FlaxBertForNextSentencePrediction,
					        FlaxBertForSequenceClassification,
					        FlaxBertForTokenClassification,
					        FlaxBertForQuestionAnswering,
					    )
					    if is_flax_available()
					    else ()
					)






					def    lowerCAmelCase_				(       self:							int						)   ->	List[str]:
											snake_case_ :Any						=			FlaxBertModelTester(self						)






					@slow
					def    lowerCAmelCase_				(       self:							List[str]						)   ->	Dict:
											# Only check this for base model, not necessary for all model classes.
											# This will also help speed-up tests.
											snake_case_ :Dict						=			FlaxBertModel.from_pretrained("""bert-base-cased"""						)
											snake_case_ :Dict						=			model(np.ones((1, 1)						)						)
											self.assertIsNotNone(snake_case						)
