





'''simple docstring'''






import unittest

import numpy as np
import timeout_decorator  # noqa

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

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
      import os

      # The slow tests are often failing with OOM error on GPU
      # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
      # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
      __lowerCamelCase						: List[str]		=   'platform'

      import jax
      import jax.numpy as jnp

      from transformers.models.blenderbot_small.modeling_flax_blenderbot_small import (
          FlaxBlenderbotSmallForConditionalGeneration,
          FlaxBlenderbotSmallModel,
          shift_tokens_right,
      )


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE=None  ,      __SCREAMING_SNAKE_CASE=None  ,      __SCREAMING_SNAKE_CASE=None  ,      __SCREAMING_SNAKE_CASE=None  ,      __SCREAMING_SNAKE_CASE=None  ,      __SCREAMING_SNAKE_CASE=None  ,      ):


  """simple docstring"""

  if attention_mask is None:
    _UpperCamelCase       			=np.where(input_ids != config.pad_token_id  ,      1  ,      0   )
  if decoder_attention_mask is None:
    _UpperCamelCase       			=np.where(decoder_input_ids != config.pad_token_id  ,      1  ,      0   )
  if head_mask is None:
    _UpperCamelCase       			=np.ones((config.encoder_layers, config.encoder_attention_heads)   )
  if decoder_head_mask is None:
    _UpperCamelCase       			=np.ones((config.decoder_layers, config.decoder_attention_heads)   )
  if cross_attn_head_mask is None:
    _UpperCamelCase       			=np.ones((config.decoder_layers, config.decoder_attention_heads)   )
  return {
      "input_ids": input_ids,
      "decoder_input_ids": decoder_input_ids,
      "attention_mask": attention_mask,
      "decoder_attention_mask": attention_mask,
  }







class 		UpperCAmelCase   :


  """simple docstring"""







  def __init__(				self    :					int							,	UpperCamelCase__    :					Any							,	UpperCamelCase__    :					Any=13							,	UpperCamelCase__    :					Optional[Any]=7							,	UpperCamelCase__    :					Tuple=True							,	UpperCamelCase__    :					Dict=False							,	UpperCamelCase__    :					Union[str, Any]=99							,	UpperCamelCase__    :					List[Any]=16							,	UpperCamelCase__    :					Optional[Any]=2							,	UpperCamelCase__    :					List[Any]=4							,	UpperCamelCase__    :					Optional[int]=4							,	UpperCamelCase__    :					Any="gelu"							,	UpperCamelCase__    :					List[Any]=0.1							,	UpperCamelCase__    :					Union[str, Any]=0.1							,	UpperCamelCase__    :					Dict=32							,	UpperCamelCase__    :					Optional[int]=2							,	UpperCamelCase__    :					Dict=1							,	UpperCamelCase__    :					int=0							,	UpperCamelCase__    :					Optional[Any]=0.02							,	)     ->						Dict:
    _UpperCamelCase       			=parent
    _UpperCamelCase       			=batch_size
    _UpperCamelCase       			=seq_length
    _UpperCamelCase       			=is_training
    _UpperCamelCase       			=use_labels
    _UpperCamelCase       			=vocab_size
    _UpperCamelCase       			=hidden_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       			=eos_token_id
    _UpperCamelCase       			=pad_token_id
    _UpperCamelCase       			=bos_token_id
    _UpperCamelCase       			=initializer_range


  def   UpperCamelCase__      (				self    :					Dict  )     ->						int:
    _UpperCamelCase       			=np.clip(ids_tensor([self.batch_size, self.seq_length - 1]							,	self.vocab_size  )							,	3							,	self.vocab_size  )
    _UpperCamelCase       			=np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)							,	dtype=np.intaa  ))							,	-1  )

    _UpperCamelCase       			=shift_tokens_right(UpperCamelCase__							,	1							,	2  )

    _UpperCamelCase       			=BlenderbotSmallConfig(
        vocab_size=self.vocab_size							,	d_model=self.hidden_size							,	encoder_layers=self.num_hidden_layers							,	decoder_layers=self.num_hidden_layers							,	encoder_attention_heads=self.num_attention_heads							,	decoder_attention_heads=self.num_attention_heads							,	encoder_ffn_dim=self.intermediate_size							,	decoder_ffn_dim=self.intermediate_size							,	dropout=self.hidden_dropout_prob							,	attention_dropout=self.attention_probs_dropout_prob							,	max_position_embeddings=self.max_position_embeddings							,	eos_token_id=self.eos_token_id							,	bos_token_id=self.bos_token_id							,	pad_token_id=self.pad_token_id							,	initializer_range=self.initializer_range							,	use_cache=UpperCamelCase__							,	)
    _UpperCamelCase       			=prepare_blenderbot_inputs_dict(UpperCamelCase__							,	UpperCamelCase__							,	UpperCamelCase__  )
    return config, inputs_dict


  def   UpperCamelCase__      (				self    :					str  )     ->						Tuple:
    _UpperCamelCase			,     _UpperCamelCase       			=self.prepare_config_and_inputs()
    return config, inputs_dict


  def   UpperCamelCase__      (				self    :					str							,	UpperCamelCase__    :					Optional[Any]							,	UpperCamelCase__    :					Tuple							,	UpperCamelCase__    :					Optional[Any]  )     ->						Union[str, Any]:
    _UpperCamelCase       			=20
    _UpperCamelCase       			=model_class_name(UpperCamelCase__  )

    _UpperCamelCase       			=model.encode(inputs_dict['''input_ids''']  )

    _UpperCamelCase			,     _UpperCamelCase       			=(
        inputs_dict['''decoder_input_ids'''],
        inputs_dict['''decoder_attention_mask'''],
    )

    _UpperCamelCase       			=model.init_cache(decoder_input_ids.shape[0]							,	UpperCamelCase__							,	UpperCamelCase__  )
    _UpperCamelCase       			=jnp.ones((decoder_input_ids.shape[0], max_decoder_length)							,	dtype='''i4'''  )

    _UpperCamelCase       			=jnp.broadcast_to(
        jnp.arange(decoder_input_ids.shape[-1] - 1  )[None, :]							,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)							,	)
    _UpperCamelCase       			=model.decode(
        decoder_input_ids[:, :-1]							,	UpperCamelCase__							,	decoder_attention_mask=UpperCamelCase__							,	past_key_values=UpperCamelCase__							,	decoder_position_ids=UpperCamelCase__							,	)

    _UpperCamelCase       			=jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]							,	dtype='''i4'''  )
    _UpperCamelCase       			=model.decode(
        decoder_input_ids[:, -1:]							,	UpperCamelCase__							,	decoder_attention_mask=UpperCamelCase__							,	past_key_values=outputs_cache.past_key_values							,	decoder_position_ids=UpperCamelCase__							,	)

    _UpperCamelCase       			=model.decode(UpperCamelCase__							,	UpperCamelCase__  )

    _UpperCamelCase       			=np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])  )  )
    self.parent.assertTrue(diff < 1E-3							,	msg=F'''Max diff is {diff}'''  )




  def   UpperCamelCase__      (				self    :					List[str]							,	UpperCamelCase__    :					Optional[Any]							,	UpperCamelCase__    :					List[Any]							,	UpperCamelCase__    :					List[Any]  )     ->						List[Any]:
    _UpperCamelCase       			=20
    _UpperCamelCase       			=model_class_name(UpperCamelCase__  )

    _UpperCamelCase       			=model.encode(inputs_dict['''input_ids''']  )

    _UpperCamelCase			,     _UpperCamelCase       			=(
        inputs_dict['''decoder_input_ids'''],
        inputs_dict['''decoder_attention_mask'''],
    )

    _UpperCamelCase       			=jnp.concatenate(
        [
            decoder_attention_mask,
            jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])  ),
        ]							,	axis=-1							,	)

    _UpperCamelCase       			=model.init_cache(decoder_input_ids.shape[0]							,	UpperCamelCase__							,	UpperCamelCase__  )
    _UpperCamelCase       			=jnp.broadcast_to(
        jnp.arange(decoder_input_ids.shape[-1] - 1  )[None, :]							,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)							,	)

    _UpperCamelCase       			=model.decode(
        decoder_input_ids[:, :-1]							,	UpperCamelCase__							,	decoder_attention_mask=UpperCamelCase__							,	past_key_values=UpperCamelCase__							,	decoder_position_ids=UpperCamelCase__							,	)
    _UpperCamelCase       			=jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]							,	dtype='''i4'''  )
    _UpperCamelCase       			=model.decode(
        decoder_input_ids[:, -1:]							,	UpperCamelCase__							,	past_key_values=outputs_cache.past_key_values							,	decoder_attention_mask=UpperCamelCase__							,	decoder_position_ids=UpperCamelCase__							,	)

    _UpperCamelCase       			=model.decode(UpperCamelCase__							,	UpperCamelCase__							,	decoder_attention_mask=UpperCamelCase__  )

    _UpperCamelCase       			=np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])  )  )
    self.parent.assertTrue(diff < 1E-3							,	msg=F'''Max diff is {diff}'''  )







@require_flax
class 		UpperCAmelCase   (				unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				99


  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Tuple:
    _UpperCamelCase       			=np.array(
        [
            [71, 82, 18, 33, 46, 91, 2],
            [68, 34, 26, 58, 30, 82, 2],
            [5, 97, 17, 39, 94, 40, 2],
            [76, 83, 94, 25, 70, 78, 2],
            [87, 59, 41, 35, 48, 66, 2],
            [55, 13, 16, 58, 5, 2, 1],  # note padding
            [64, 27, 31, 51, 12, 75, 2],
            [52, 64, 86, 17, 83, 39, 2],
            [48, 61, 9, 24, 71, 82, 2],
            [26, 1, 60, 48, 22, 13, 2],
            [21, 5, 62, 28, 14, 76, 2],
            [45, 98, 37, 86, 59, 48, 2],
            [70, 70, 50, 9, 28, 0, 2],
        ]							,	dtype=np.intaa							,	)

    _UpperCamelCase       			=input_ids.shape[0]
    _UpperCamelCase       			=BlenderbotSmallConfig(
        vocab_size=self.vocab_size							,	d_model=24							,	encoder_layers=2							,	decoder_layers=2							,	encoder_attention_heads=2							,	decoder_attention_heads=2							,	encoder_ffn_dim=32							,	decoder_ffn_dim=32							,	max_position_embeddings=48							,	eos_token_id=2							,	pad_token_id=1							,	bos_token_id=0							,	)
    return config, input_ids, batch_size


  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Dict:
    _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase       			=self._get_config_and_data()
    _UpperCamelCase       			=FlaxBlenderbotSmallForConditionalGeneration(UpperCamelCase__  )
    _UpperCamelCase       			=lm_model(input_ids=UpperCamelCase__  )
    _UpperCamelCase       			=(batch_size, input_ids.shape[1], config.vocab_size)
    self.assertEqual(outputs['''logits'''].shape							,	UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						str:
    _UpperCamelCase       			=BlenderbotSmallConfig(
        vocab_size=self.vocab_size							,	d_model=14							,	encoder_layers=2							,	decoder_layers=2							,	encoder_attention_heads=2							,	decoder_attention_heads=2							,	encoder_ffn_dim=8							,	decoder_ffn_dim=8							,	max_position_embeddings=48							,	)
    _UpperCamelCase       			=FlaxBlenderbotSmallForConditionalGeneration(UpperCamelCase__  )
    _UpperCamelCase       			=np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]							,	dtype=np.intaa  )
    _UpperCamelCase       			=np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]							,	dtype=np.intaa  )
    _UpperCamelCase       			=lm_model(input_ids=UpperCamelCase__							,	decoder_input_ids=UpperCamelCase__  )
    _UpperCamelCase       			=(*summary.shape, config.vocab_size)
    self.assertEqual(outputs['''logits'''].shape							,	UpperCamelCase__  )




  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						Optional[Any]:
    _UpperCamelCase       			=np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]							,	dtype=np.intaa  )
    _UpperCamelCase       			=shift_tokens_right(UpperCamelCase__							,	1							,	2  )
    _UpperCamelCase       			=np.equal(UpperCamelCase__							,	1  ).astype(np.floataa  ).sum()
    _UpperCamelCase       			=np.equal(UpperCamelCase__							,	1  ).astype(np.floataa  ).sum()
    self.assertEqual(shifted.shape							,	input_ids.shape  )
    self.assertEqual(UpperCamelCase__							,	n_pad_before - 1  )
    self.assertTrue(np.equal(shifted[:, 0]							,	2  ).all()  )







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


  """simple docstring"""







  lowerCAmelCase_					=				True
  lowerCAmelCase_					=				(
      (
          FlaxBlenderbotSmallModel,
          FlaxBlenderbotSmallForConditionalGeneration,
      )
      if is_flax_available()
      else ()
  )
  lowerCAmelCase_					=				(FlaxBlenderbotSmallForConditionalGeneration,) if is_flax_available() else ()


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Tuple:
    _UpperCamelCase       			=FlaxBlenderbotSmallModelTester(self  )


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						List[str]:
    _UpperCamelCase			,     _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    for model_class in self.all_model_classes:
      self.model_tester.check_use_cache_forward(UpperCamelCase__							,	UpperCamelCase__							,	UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Tuple:
    _UpperCamelCase			,     _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    for model_class in self.all_model_classes:
      self.model_tester.check_use_cache_forward_with_attn_mask(UpperCamelCase__							,	UpperCamelCase__							,	UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					int  )     ->						str:
    _UpperCamelCase			,     _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
      with self.subTest(model_class.__name__  ):
        _UpperCamelCase       			=self._prepare_for_class(UpperCamelCase__							,	UpperCamelCase__  )
        _UpperCamelCase       			=model_class(UpperCamelCase__  )

        @jax.jit
        def encode_jitted(UpperCamelCase__    :					Optional[Any]							,	UpperCamelCase__    :					Any=None							,	**UpperCamelCase__    :					Union[str, Any]  ):
          return model.encode(input_ids=UpperCamelCase__							,	attention_mask=UpperCamelCase__  )

        with self.subTest('''JIT Enabled'''  ):
          _UpperCamelCase       			=encode_jitted(**UpperCamelCase__  ).to_tuple()

        with self.subTest('''JIT Disabled'''  ):
          with jax.disable_jit():
            _UpperCamelCase       			=encode_jitted(**UpperCamelCase__  ).to_tuple()

        self.assertEqual(len(UpperCamelCase__  )							,	len(UpperCamelCase__  )  )
        for jitted_output, output in zip(UpperCamelCase__							,	UpperCamelCase__  ):
          self.assertEqual(jitted_output.shape							,	output.shape  )


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Tuple:
    _UpperCamelCase			,     _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
      with self.subTest(model_class.__name__  ):
        _UpperCamelCase       			=model_class(UpperCamelCase__  )
        _UpperCamelCase       			=model.encode(inputs_dict['''input_ids''']							,	inputs_dict['''attention_mask''']  )

        _UpperCamelCase       			={
            '''decoder_input_ids''': inputs_dict['''decoder_input_ids'''],
            '''decoder_attention_mask''': inputs_dict['''decoder_attention_mask'''],
            '''encoder_outputs''': encoder_outputs,
        }

        @jax.jit
        def decode_jitted(UpperCamelCase__    :					Optional[Any]							,	UpperCamelCase__    :					List[Any]							,	UpperCamelCase__    :					Union[str, Any]  ):
          return model.decode(
              decoder_input_ids=UpperCamelCase__							,	decoder_attention_mask=UpperCamelCase__							,	encoder_outputs=UpperCamelCase__							,	)

        with self.subTest('''JIT Enabled'''  ):
          _UpperCamelCase       			=decode_jitted(**UpperCamelCase__  ).to_tuple()

        with self.subTest('''JIT Disabled'''  ):
          with jax.disable_jit():
            _UpperCamelCase       			=decode_jitted(**UpperCamelCase__  ).to_tuple()

        self.assertEqual(len(UpperCamelCase__  )							,	len(UpperCamelCase__  )  )
        for jitted_output, output in zip(UpperCamelCase__							,	UpperCamelCase__  ):
          self.assertEqual(jitted_output.shape							,	output.shape  )




  @slow
  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						Optional[Any]:
    for model_class_name in self.all_model_classes:
      _UpperCamelCase       			=model_class_name.from_pretrained('''facebook/blenderbot_small-90M'''  )
      # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
      _UpperCamelCase       			=np.ones((1, 1)  ) * model.config.eos_token_id
      _UpperCamelCase       			=model(UpperCamelCase__  )
      self.assertIsNotNone(UpperCamelCase__  )













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






