





'''simple docstring'''






import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import (
          MODEL_FOR_PRETRAINING_MAPPING,
          MegatronBertForCausalLM,
          MegatronBertForMaskedLM,
          MegatronBertForMultipleChoice,
          MegatronBertForNextSentencePrediction,
          MegatronBertForPreTraining,
          MegatronBertForQuestionAnswering,
          MegatronBertForSequenceClassification,
          MegatronBertForTokenClassification,
          MegatronBertModel,
      )







class 		UpperCAmelCase   :


  """simple docstring"""







  def __init__(				self    :					List[Any]							,	UpperCamelCase__    :					Any							,	UpperCamelCase__    :					List[str]=13							,	UpperCamelCase__    :					Tuple=7							,	UpperCamelCase__    :					Any=True							,	UpperCamelCase__    :					Union[str, Any]=True							,	UpperCamelCase__    :					str=True							,	UpperCamelCase__    :					List[str]=True							,	UpperCamelCase__    :					Any=99							,	UpperCamelCase__    :					str=64							,	UpperCamelCase__    :					str=32							,	UpperCamelCase__    :					Tuple=5							,	UpperCamelCase__    :					Dict=4							,	UpperCamelCase__    :					Optional[int]=37							,	UpperCamelCase__    :					List[str]="gelu"							,	UpperCamelCase__    :					Dict=0.1							,	UpperCamelCase__    :					Optional[int]=0.1							,	UpperCamelCase__    :					Optional[Any]=512							,	UpperCamelCase__    :					List[Any]=16							,	UpperCamelCase__    :					Any=2							,	UpperCamelCase__    :					List[str]=0.02							,	UpperCamelCase__    :					Union[str, Any]=3							,	UpperCamelCase__    :					Dict=4							,	UpperCamelCase__    :					Optional[Any]=None							,	)     ->						str:
    _UpperCamelCase       			=parent
    _UpperCamelCase       			=batch_size
    _UpperCamelCase       			=seq_length
    _UpperCamelCase       			=is_training
    _UpperCamelCase       			=use_input_mask
    _UpperCamelCase       			=use_token_type_ids
    _UpperCamelCase       			=use_labels
    _UpperCamelCase       			=vocab_size
    _UpperCamelCase       			=hidden_size
    _UpperCamelCase       			=embedding_size
    _UpperCamelCase       			=num_hidden_layers
    _UpperCamelCase       			=num_attention_heads
    _UpperCamelCase       			=intermediate_size
    _UpperCamelCase       			=hidden_act
    _UpperCamelCase       			=hidden_dropout_prob
    _UpperCamelCase       			=attention_probs_dropout_prob
    _UpperCamelCase       			=max_position_embeddings
    _UpperCamelCase       			=type_vocab_size
    _UpperCamelCase       			=type_sequence_label_size
    _UpperCamelCase       			=initializer_range
    _UpperCamelCase       			=num_labels
    _UpperCamelCase       			=num_choices
    _UpperCamelCase       			=scope


  def   UpperCamelCase__      (				self    :					List[str]  )     ->						Optional[int]:
    _UpperCamelCase       			=ids_tensor([self.batch_size, self.seq_length]							,	self.vocab_size  )

    _UpperCamelCase       			=None
    if self.use_input_mask:
      _UpperCamelCase       			=random_attention_mask([self.batch_size, self.seq_length]  )

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

    _UpperCamelCase       			=None
    _UpperCamelCase       			=None
    _UpperCamelCase       			=None
    if self.use_labels:
      _UpperCamelCase       			=ids_tensor([self.batch_size]							,	self.type_sequence_label_size  )
      _UpperCamelCase       			=ids_tensor([self.batch_size, self.seq_length]							,	self.num_labels  )
      _UpperCamelCase       			=ids_tensor([self.batch_size]							,	self.num_choices  )

    _UpperCamelCase       			=self.get_config()

    return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


  def   UpperCamelCase__      (				self    :					List[str]  )     ->						Optional[Any]:
    return MegatronBertConfig(
        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							,	embedding_size=self.embedding_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=__A							,	initializer_range=self.initializer_range							,	)


  def   UpperCamelCase__      (				self    :					Optional[int]							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Tuple							,	UpperCamelCase__    :					List[Any]							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					Optional[Any]							,	UpperCamelCase__    :					List[Any]  )     ->						Union[str, Any]:
    _UpperCamelCase       			=MegatronBertModel(config=__A  )
    model.to(__A  )
    model.eval()
    _UpperCamelCase       			=model(__A							,	attention_mask=__A							,	token_type_ids=__A  )
    _UpperCamelCase       			=model(__A							,	token_type_ids=__A  )
    _UpperCamelCase       			=model(__A  )

    self.parent.assertEqual(result.last_hidden_state.shape							,	(self.batch_size, self.seq_length, self.hidden_size)  )
    self.parent.assertEqual(result.pooler_output.shape							,	(self.batch_size, self.hidden_size)  )


  def   UpperCamelCase__      (				self    :					int							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Union[str, Any]							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					List[str]  )     ->						Optional[int]:
    _UpperCamelCase       			=MegatronBertForMaskedLM(config=__A  )
    model.to(__A  )
    model.eval()
    _UpperCamelCase       			=model(__A							,	attention_mask=__A							,	token_type_ids=__A							,	labels=__A  )
    self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.seq_length, self.vocab_size)  )


  def   UpperCamelCase__      (				self    :					Any							,	UpperCamelCase__    :					Optional[Any]							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					Union[str, Any]							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					Optional[Any]  )     ->						int:
    _UpperCamelCase       			=MegatronBertForCausalLM(config=__A  )
    model.to(__A  )
    model.eval()
    _UpperCamelCase       			=model(__A							,	attention_mask=__A							,	token_type_ids=__A							,	labels=__A  )
    self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.seq_length, self.vocab_size)  )


  def   UpperCamelCase__      (				self    :					str							,	UpperCamelCase__    :					Tuple							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					Tuple							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					Union[str, Any]							,	UpperCamelCase__    :					Tuple							,	UpperCamelCase__    :					List[Any]  )     ->						str:
    _UpperCamelCase       			=MegatronBertForNextSentencePrediction(config=__A  )
    model.to(__A  )
    model.eval()
    _UpperCamelCase       			=model(
        __A							,	attention_mask=__A							,	token_type_ids=__A							,	labels=__A							,	)
    self.parent.assertEqual(result.logits.shape							,	(self.batch_size, 2)  )


  def   UpperCamelCase__      (				self    :					List[Any]							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					Union[str, Any]							,	UpperCamelCase__    :					Any							,	UpperCamelCase__    :					Any							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					Any  )     ->						List[str]:
    _UpperCamelCase       			=MegatronBertForPreTraining(config=__A  )
    model.to(__A  )
    model.eval()
    _UpperCamelCase       			=model(
        __A							,	attention_mask=__A							,	token_type_ids=__A							,	labels=__A							,	next_sentence_label=__A							,	)
    self.parent.assertEqual(result.prediction_logits.shape							,	(self.batch_size, self.seq_length, self.vocab_size)  )
    self.parent.assertEqual(result.seq_relationship_logits.shape							,	(self.batch_size, 2)  )


  def   UpperCamelCase__      (				self    :					Optional[Any]							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					Tuple							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Optional[Any]							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Tuple  )     ->						Union[str, Any]:
    _UpperCamelCase       			=MegatronBertForQuestionAnswering(config=__A  )
    model.to(__A  )
    model.eval()
    _UpperCamelCase       			=model(
        __A							,	attention_mask=__A							,	token_type_ids=__A							,	start_positions=__A							,	end_positions=__A							,	)
    self.parent.assertEqual(result.start_logits.shape							,	(self.batch_size, self.seq_length)  )
    self.parent.assertEqual(result.end_logits.shape							,	(self.batch_size, self.seq_length)  )


  def   UpperCamelCase__      (				self    :					Union[str, Any]							,	UpperCamelCase__    :					Union[str, Any]							,	UpperCamelCase__    :					Union[str, Any]							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Optional[Any]  )     ->						Union[str, Any]:
    _UpperCamelCase       			=self.num_labels
    _UpperCamelCase       			=MegatronBertForSequenceClassification(__A  )
    model.to(__A  )
    model.eval()
    _UpperCamelCase       			=model(__A							,	attention_mask=__A							,	token_type_ids=__A							,	labels=__A  )
    self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.num_labels)  )


  def   UpperCamelCase__      (				self    :					Union[str, Any]							,	UpperCamelCase__    :					List[Any]							,	UpperCamelCase__    :					Tuple							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					Any							,	UpperCamelCase__    :					str							,	UpperCamelCase__    :					int  )     ->						Union[str, Any]:
    _UpperCamelCase       			=self.num_labels
    _UpperCamelCase       			=MegatronBertForTokenClassification(config=__A  )
    model.to(__A  )
    model.eval()
    _UpperCamelCase       			=model(__A							,	attention_mask=__A							,	token_type_ids=__A							,	labels=__A  )
    self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.seq_length, self.num_labels)  )


  def   UpperCamelCase__      (				self    :					str							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					str							,	UpperCamelCase__    :					str  )     ->						Dict:
    _UpperCamelCase       			=self.num_choices
    _UpperCamelCase       			=MegatronBertForMultipleChoice(config=__A  )
    model.to(__A  )
    model.eval()
    _UpperCamelCase       			=input_ids.unsqueeze(1  ).expand(-1							,	self.num_choices							,	-1  ).contiguous()
    _UpperCamelCase       			=token_type_ids.unsqueeze(1  ).expand(-1							,	self.num_choices							,	-1  ).contiguous()
    _UpperCamelCase       			=input_mask.unsqueeze(1  ).expand(-1							,	self.num_choices							,	-1  ).contiguous()
    _UpperCamelCase       			=model(
        __A							,	attention_mask=__A							,	token_type_ids=__A							,	labels=__A							,	)
    self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.num_choices)  )




  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						Any:
    _UpperCamelCase       			=self.prepare_config_and_inputs()
    (
        (
        _UpperCamelCase
    )			,     (
        _UpperCamelCase
    )			,     (
        _UpperCamelCase
    )			,     (
        _UpperCamelCase
    )			,     (
        _UpperCamelCase
    )			,     (
        _UpperCamelCase
    )			,     (
        _UpperCamelCase
    )			,     
    )       			=config_and_inputs
    _UpperCamelCase       			={'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
    return config, inputs_dict







@require_torch
class 		UpperCAmelCase   (				lowercase_   ,     lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				(
      (
          MegatronBertModel,
          MegatronBertForMaskedLM,
          MegatronBertForCausalLM,
          MegatronBertForMultipleChoice,
          MegatronBertForNextSentencePrediction,
          MegatronBertForPreTraining,
          MegatronBertForQuestionAnswering,
          MegatronBertForSequenceClassification,
          MegatronBertForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  lowerCAmelCase_					=				(
      {
          """feature-extraction""": MegatronBertModel,
          """fill-mask""": MegatronBertForMaskedLM,
          """question-answering""": MegatronBertForQuestionAnswering,
          """text-classification""": MegatronBertForSequenceClassification,
          """text-generation""": MegatronBertForCausalLM,
          """token-classification""": MegatronBertForTokenClassification,
          """zero-shot""": MegatronBertForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  lowerCAmelCase_					=				True
  # test_resize_embeddings = False
  lowerCAmelCase_					=				False


  def   UpperCamelCase__      (				self    :					List[Any]							,	UpperCamelCase__    :					Optional[Any]							,	UpperCamelCase__    :					Union[str, Any]							,	UpperCamelCase__    :					Any=False  )     ->						Dict:
    _UpperCamelCase       			=super()._prepare_for_class(__A							,	__A							,	return_labels=__A  )

    if return_labels:
      if model_class in get_values(__A  ):
        _UpperCamelCase       			=torch.zeros(
            (self.model_tester.batch_size, self.model_tester.seq_length)							,	dtype=torch.long							,	device=__A  )
        _UpperCamelCase       			=torch.zeros(
            self.model_tester.batch_size							,	dtype=torch.long							,	device=__A  )
    return inputs_dict


  def   UpperCamelCase__      (				self    :					List[str]  )     ->						str:
    _UpperCamelCase       			=MegatronBertModelTester(self  )
    _UpperCamelCase       			=ConfigTester(self							,	config_class=__A							,	hidden_size=37  )


  def   UpperCamelCase__      (				self    :					int  )     ->						Any:
    self.config_tester.run_common_tests()


  def   UpperCamelCase__      (				self    :					List[str]  )     ->						List[Any]:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_model(*__A  )


  def   UpperCamelCase__      (				self    :					int  )     ->						Union[str, Any]:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__A  )


  def   UpperCamelCase__      (				self    :					Tuple  )     ->						Optional[Any]:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__A  )


  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						Optional[Any]:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__A  )


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Optional[int]:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_pretraining(*__A  )


  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						List[str]:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_question_answering(*__A  )


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Optional[int]:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__A  )




  def   UpperCamelCase__      (				self    :					int  )     ->						Dict:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_token_classification(*__A  )




def 					_a   (__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  return torch.tensor(
      _lowercase  ,      dtype=torch.long  ,      device=_lowercase  ,      )


__lowerCamelCase						: Union[str, Any]		=   1E-4







@require_torch
@require_sentencepiece
@require_tokenizers
class 		UpperCAmelCase   (				unittest.TestCase):


  """simple docstring"""







  @slow
  @unittest.skip('''Model is not available.'''  )
  def   UpperCamelCase__      (				self    :					List[str]  )     ->						Dict:
    _UpperCamelCase       			='''nvidia/megatron-bert-uncased-345m'''
    if "MYDIR" in os.environ:
      _UpperCamelCase       			=os.path.join(os.environ['''MYDIR''']							,	__A  )
    _UpperCamelCase       			=MegatronBertModel.from_pretrained(__A  )
    model.to(__A  )
    model.half()
    _UpperCamelCase       			=_long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]  )
    with torch.no_grad():
      _UpperCamelCase       			=model(__A  )[0]
    _UpperCamelCase       			=torch.Size((1, 9, 1024)  )
    self.assertEqual(output.shape							,	__A  )

    _UpperCamelCase       			=[-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
    for ii in range(3  ):
      for jj in range(3  ):
        _UpperCamelCase       			=output[0, ii, jj]
        _UpperCamelCase       			=expected[3 * ii + jj]
        _UpperCamelCase       			='''ii={} jj={} a={} b={}'''.format(__A							,	__A							,	__A							,	__A  )
        self.assertTrue(math.isclose(__A							,	__A							,	rel_tol=__A							,	abs_tol=__A  )							,	msg=__A  )













'''simple docstring'''






import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np







class 		UpperCAmelCase   :


  """simple docstring"""







  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Union[str, Any]:
    torch.manual_seed(0  )
    _UpperCamelCase       			=TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )

    torch.manual_seed(0  )
    _UpperCamelCase       			=AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )

    torch.manual_seed(0  )
    _UpperCamelCase       			=UNetaDConditionModel(
        sample_size=32							,	layers_per_block=1							,	block_out_channels=[32, 64]							,	down_block_types=[
            '''ResnetDownsampleBlock2D''',
            '''SimpleCrossAttnDownBlock2D''',
        ]							,	mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''							,	up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']							,	in_channels=3							,	out_channels=6							,	cross_attention_dim=32							,	encoder_hid_dim=32							,	attention_head_dim=8							,	addition_embed_type='''text'''							,	addition_embed_type_num_heads=2							,	cross_attention_norm='''group_norm'''							,	resnet_time_scale_shift='''scale_shift'''							,	act_fn='''gelu'''							,	)
    unet.set_attn_processor(AttnAddedKVProcessor()  )  # For reproducibility tests

    torch.manual_seed(0  )
    _UpperCamelCase       			=DDPMScheduler(
        num_train_timesteps=1000							,	beta_schedule='''squaredcos_cap_v2'''							,	beta_start=0.0001							,	beta_end=0.02							,	thresholding=UpperCamelCase__							,	dynamic_thresholding_ratio=0.95							,	sample_max_value=1.0							,	prediction_type='''epsilon'''							,	variance_type='''learned_range'''							,	)

    torch.manual_seed(0  )
    _UpperCamelCase       			=IFWatermarker()

    return {
        "text_encoder": text_encoder,
        "tokenizer": tokenizer,
        "unet": unet,
        "scheduler": scheduler,
        "watermarker": watermarker,
        "safety_checker": None,
        "feature_extractor": None,
    }


  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						Optional[int]:
    torch.manual_seed(0  )
    _UpperCamelCase       			=TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )

    torch.manual_seed(0  )
    _UpperCamelCase       			=AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )

    torch.manual_seed(0  )
    _UpperCamelCase       			=UNetaDConditionModel(
        sample_size=32							,	layers_per_block=[1, 2]							,	block_out_channels=[32, 64]							,	down_block_types=[
            '''ResnetDownsampleBlock2D''',
            '''SimpleCrossAttnDownBlock2D''',
        ]							,	mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''							,	up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']							,	in_channels=6							,	out_channels=6							,	cross_attention_dim=32							,	encoder_hid_dim=32							,	attention_head_dim=8							,	addition_embed_type='''text'''							,	addition_embed_type_num_heads=2							,	cross_attention_norm='''group_norm'''							,	resnet_time_scale_shift='''scale_shift'''							,	act_fn='''gelu'''							,	class_embed_type='''timestep'''							,	mid_block_scale_factor=1.414							,	time_embedding_act_fn='''gelu'''							,	time_embedding_dim=32							,	)
    unet.set_attn_processor(AttnAddedKVProcessor()  )  # For reproducibility tests

    torch.manual_seed(0  )
    _UpperCamelCase       			=DDPMScheduler(
        num_train_timesteps=1000							,	beta_schedule='''squaredcos_cap_v2'''							,	beta_start=0.0001							,	beta_end=0.02							,	thresholding=UpperCamelCase__							,	dynamic_thresholding_ratio=0.95							,	sample_max_value=1.0							,	prediction_type='''epsilon'''							,	variance_type='''learned_range'''							,	)

    torch.manual_seed(0  )
    _UpperCamelCase       			=DDPMScheduler(
        num_train_timesteps=1000							,	beta_schedule='''squaredcos_cap_v2'''							,	beta_start=0.0001							,	beta_end=0.02							,	)

    torch.manual_seed(0  )
    _UpperCamelCase       			=IFWatermarker()

    return {
        "text_encoder": text_encoder,
        "tokenizer": tokenizer,
        "unet": unet,
        "scheduler": scheduler,
        "image_noising_scheduler": image_noising_scheduler,
        "watermarker": watermarker,
        "safety_checker": None,
        "feature_extractor": None,
    }


  def   UpperCamelCase__      (				self    :					Any  )     ->						Any:
    _UpperCamelCase       			=self.get_dummy_components()
    _UpperCamelCase       			=self.pipeline_class(**UpperCamelCase__  )
    pipe.to(UpperCamelCase__  )
    pipe.set_progress_bar_config(disable=UpperCamelCase__  )

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )

    _UpperCamelCase       			=inputs['''prompt''']
    _UpperCamelCase       			=inputs['''generator''']
    _UpperCamelCase       			=inputs['''num_inference_steps''']
    _UpperCamelCase       			=inputs['''output_type''']

    if "image" in inputs:
      _UpperCamelCase       			=inputs['''image''']
    else:
      _UpperCamelCase       			=None

    if "mask_image" in inputs:
      _UpperCamelCase       			=inputs['''mask_image''']
    else:
      _UpperCamelCase       			=None

    if "original_image" in inputs:
      _UpperCamelCase       			=inputs['''original_image''']
    else:
      _UpperCamelCase       			=None

    _UpperCamelCase			,     _UpperCamelCase       			=pipe.encode_prompt(UpperCamelCase__  )

    # inputs with prompt converted to embeddings
    _UpperCamelCase       			={
        '''prompt_embeds''': prompt_embeds,
        '''negative_prompt_embeds''': negative_prompt_embeds,
        '''generator''': generator,
        '''num_inference_steps''': num_inference_steps,
        '''output_type''': output_type,
    }

    if image is not None:
      _UpperCamelCase       			=image

    if mask_image is not None:
      _UpperCamelCase       			=mask_image

    if original_image is not None:
      _UpperCamelCase       			=original_image

    # set all optional components to None
    for optional_component in pipe._optional_components:
      setattr(UpperCamelCase__							,	UpperCamelCase__							,	UpperCamelCase__  )

    _UpperCamelCase       			=pipe(**UpperCamelCase__  )[0]

    with tempfile.TemporaryDirectory() as tmpdir:
      pipe.save_pretrained(UpperCamelCase__  )
      _UpperCamelCase       			=self.pipeline_class.from_pretrained(UpperCamelCase__  )
      pipe_loaded.to(UpperCamelCase__  )
      pipe_loaded.set_progress_bar_config(disable=UpperCamelCase__  )

    pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()  )  # For reproducibility tests

    for optional_component in pipe._optional_components:
      self.assertTrue(
          getattr(UpperCamelCase__							,	UpperCamelCase__  ) is None							,	F'''`{optional_component}` did not stay set to None after loading.'''							,	)

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )

    _UpperCamelCase       			=inputs['''generator''']
    _UpperCamelCase       			=inputs['''num_inference_steps''']
    _UpperCamelCase       			=inputs['''output_type''']

    # inputs with prompt converted to embeddings
    _UpperCamelCase       			={
        '''prompt_embeds''': prompt_embeds,
        '''negative_prompt_embeds''': negative_prompt_embeds,
        '''generator''': generator,
        '''num_inference_steps''': num_inference_steps,
        '''output_type''': output_type,
    }

    if image is not None:
      _UpperCamelCase       			=image

    if mask_image is not None:
      _UpperCamelCase       			=mask_image

    if original_image is not None:
      _UpperCamelCase       			=original_image

    _UpperCamelCase       			=pipe_loaded(**UpperCamelCase__  )[0]

    _UpperCamelCase       			=np.abs(to_np(UpperCamelCase__  ) - to_np(UpperCamelCase__  )  ).max()
    self.assertLess(UpperCamelCase__							,	1E-4  )




  def   UpperCamelCase__      (				self    :					int  )     ->						str:
    _UpperCamelCase       			=self.get_dummy_components()
    _UpperCamelCase       			=self.pipeline_class(**UpperCamelCase__  )
    pipe.to(UpperCamelCase__  )
    pipe.set_progress_bar_config(disable=UpperCamelCase__  )

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )
    _UpperCamelCase       			=pipe(**UpperCamelCase__  )[0]

    with tempfile.TemporaryDirectory() as tmpdir:
      pipe.save_pretrained(UpperCamelCase__  )
      _UpperCamelCase       			=self.pipeline_class.from_pretrained(UpperCamelCase__  )
      pipe_loaded.to(UpperCamelCase__  )
      pipe_loaded.set_progress_bar_config(disable=UpperCamelCase__  )

    pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()  )  # For reproducibility tests

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )
    _UpperCamelCase       			=pipe_loaded(**UpperCamelCase__  )[0]

    _UpperCamelCase       			=np.abs(to_np(UpperCamelCase__  ) - to_np(UpperCamelCase__  )  ).max()
    self.assertLess(UpperCamelCase__							,	1E-4  )






