

"""simple docstring"""




import collections
import inspect
import unittest
from typing import Dict, List, Tuple

from transformers import MaskFormerSwinConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, torch_device
from transformers.utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch
      from torch import nn

      from transformers import MaskFormerSwinBackbone
      from transformers.models.maskformer import MaskFormerSwinModel







class   __lowercase      :

  def __init__(	self		:      Dict       ,A		:      Union[str, Any]       ,A		:      Dict=13       ,A		:      int=32       ,A		:      str=2       ,A		:      Union[str, Any]=3       ,A		:      Union[str, Any]=16       ,A		:      Union[str, Any]=[1, 2, 1]       ,A		:      Optional[int]=[2, 2, 4]       ,A		:      str=2       ,A		:      int=2.0       ,A		:      Union[str, Any]=True       ,A		:      Any=0.0       ,A		:      Optional[Any]=0.0       ,A		:      str=0.1       ,A		:      List[str]="gelu"       ,A		:      Any=False       ,A		:      List[Any]=True       ,A		:      List[str]=0.0_2       ,A		:      Union[str, Any]=1e-5       ,A		:      List[Any]=True       ,A		:      Optional[Any]=None       ,A		:      Union[str, Any]=True       ,A		:      Any=10       ,A		:      str=8       ,A		:      str=["stage1", "stage2", "stage3"]       ,A		:      List[Any]=[1, 2, 3]       ,):






   '''simple docstring'''
   UpperCAmelCase__       :	List[str]        =       parent
   UpperCAmelCase__       :	Tuple        =       batch_size
   UpperCAmelCase__       :	Dict        =       image_size
   UpperCAmelCase__       :	List[Any]        =       patch_size
   UpperCAmelCase__       :	List[str]        =       num_channels
   UpperCAmelCase__       :	str        =       embed_dim
   UpperCAmelCase__       :	Union[str, Any]        =       depths
   UpperCAmelCase__       :	Optional[Any]        =       num_heads
   UpperCAmelCase__       :	int        =       window_size
   UpperCAmelCase__       :	Tuple        =       mlp_ratio
   UpperCAmelCase__       :	Optional[Any]        =       qkv_bias
   UpperCAmelCase__       :	List[Any]        =       hidden_dropout_prob
   UpperCAmelCase__       :	List[Any]        =       attention_probs_dropout_prob
   UpperCAmelCase__       :	Union[str, Any]        =       drop_path_rate
   UpperCAmelCase__       :	Union[str, Any]        =       hidden_act
   UpperCAmelCase__       :	List[Any]        =       use_absolute_embeddings
   UpperCAmelCase__       :	Optional[int]        =       patch_norm
   UpperCAmelCase__       :	Union[str, Any]        =       layer_norm_eps
   UpperCAmelCase__       :	Optional[int]        =       initializer_range
   UpperCAmelCase__       :	List[str]        =       is_training
   UpperCAmelCase__       :	str        =       scope
   UpperCAmelCase__       :	Union[str, Any]        =       use_labels
   UpperCAmelCase__       :	Union[str, Any]        =       type_sequence_label_size
   UpperCAmelCase__       :	Optional[Any]        =       encoder_stride
   UpperCAmelCase__       :	List[Any]        =       out_features
   UpperCAmelCase__       :	Union[str, Any]        =       out_indices

  def  __lowercase							(	self		:      Tuple     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Dict        =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

   UpperCAmelCase__       :	Tuple        =       None
   if self.use_labels:
    UpperCAmelCase__       :	int        =       ids_tensor([self.batch_size]       ,self.type_sequence_label_size     )

   UpperCAmelCase__       :	Any        =       self.get_config()

   return config, pixel_values, labels

  def  __lowercase							(	self		:      Tuple     ):






   '''simple docstring'''
   return MaskFormerSwinConfig(
       image_size=self.image_size       ,patch_size=self.patch_size       ,num_channels=self.num_channels       ,embed_dim=self.embed_dim       ,depths=self.depths       ,num_heads=self.num_heads       ,window_size=self.window_size       ,mlp_ratio=self.mlp_ratio       ,qkv_bias=self.qkv_bias       ,hidden_dropout_prob=self.hidden_dropout_prob       ,attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,drop_path_rate=self.drop_path_rate       ,hidden_act=self.hidden_act       ,use_absolute_embeddings=self.use_absolute_embeddings       ,path_norm=self.patch_norm       ,layer_norm_eps=self.layer_norm_eps       ,initializer_range=self.initializer_range       ,encoder_stride=self.encoder_stride       ,out_features=self.out_features       ,out_indices=self.out_indices       ,)

  def  __lowercase							(	self		:      Any       ,A		:      Tuple       ,A		:      Union[str, Any]       ,A		:      Dict     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[int]        =       MaskFormerSwinModel(config=a__     )
   model.to(a__     )
   model.eval()
   UpperCAmelCase__       :	Union[str, Any]        =       model(a__     )

   UpperCAmelCase__       :	Any        =       ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths     ) - 1))
   UpperCAmelCase__       :	str        =       int(config.embed_dim * 2 ** (len(config.depths     ) - 1)     )

   self.parent.assertEqual(result.last_hidden_state.shape       ,(self.batch_size, expected_seq_len, expected_dim)     )

  def  __lowercase							(	self		:      Tuple       ,A		:      Tuple       ,A		:      List[str]       ,A		:      List[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Any        =       MaskFormerSwinBackbone(config=a__     )
   model.to(a__     )
   model.eval()
   UpperCAmelCase__       :	List[str]        =       model(a__     )

   # verify feature maps
   self.parent.assertEqual(len(result.feature_maps     )       ,len(config.out_features     )     )
   self.parent.assertListEqual(list(result.feature_maps[0].shape     )       ,[13, 16, 16, 16]     )

   # verify channels
   self.parent.assertEqual(len(model.channels     )       ,len(config.out_features     )     )
   self.parent.assertListEqual(model.channels       ,[16, 32, 64]     )

   # verify ValueError
   with self.parent.assertRaises(a__     ):
    UpperCAmelCase__       :	Optional[Any]        =       ["stem"]
    UpperCAmelCase__       :	Dict        =       MaskFormerSwinBackbone(config=a__     )




  def  __lowercase							(	self		:      List[str]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Dict        =       self.prepare_config_and_inputs()
   UpperCAmelCase__       :	Dict        =       config_and_inputs
   UpperCAmelCase__       :	List[str]        =       {"pixel_values": pixel_values}
   return config, inputs_dict







@require_torch
class   __lowercase      (     UpperCamelCase_						,      UpperCamelCase_						,      unittest.TestCase       ):
  snake_case_												=							(
      (
          MaskFormerSwinModel,
          MaskFormerSwinBackbone,
      )
      if is_torch_available()
      else ()
  )
  snake_case_												=							{"""feature-extraction""": MaskFormerSwinModel} if is_torch_available() else {}
  snake_case_												=							False
  snake_case_												=							False
  snake_case_												=							False
  snake_case_												=							False
  snake_case_												=							False

  def  __lowercase							(	self		:      List[str]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[int]        =       MaskFormerSwinModelTester(self     )
   UpperCAmelCase__       :	List[Any]        =       ConfigTester(self       ,config_class=a__       ,embed_dim=37     )

  @require_torch_multi_gpu
  @unittest.skip(
      reason=(
          """`MaskFormerSwinModel` outputs `hidden_states_spatial_dimensions` which doesn't work well with"""
          """ `nn.DataParallel`"""
      )     )
  def  __lowercase							(	self		:      Dict     ):






   '''simple docstring'''
   pass

  def  __lowercase							(	self		:      Any     ):






   '''simple docstring'''
   self.create_and_test_config_common_properties()
   self.config_tester.create_and_test_config_to_json_string()
   self.config_tester.create_and_test_config_to_json_file()
   self.config_tester.create_and_test_config_from_and_save_pretrained()
   self.config_tester.create_and_test_config_with_num_labels()
   self.config_tester.check_config_can_be_init_without_params()
   self.config_tester.check_config_arguments_init()

  def  __lowercase							(	self		:      int     ):






   '''simple docstring'''
   return

  def  __lowercase							(	self		:      List[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Any        =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*a__     )

  def  __lowercase							(	self		:      Union[str, Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[int]        =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_backbone(*a__     )

  @unittest.skip("""Swin does not use inputs_embeds"""     )
  def  __lowercase							(	self		:      Union[str, Any]     ):






   '''simple docstring'''
   pass

  @unittest.skip("""Swin does not support feedforward chunking"""     )
  def  __lowercase							(	self		:      Optional[int]     ):






   '''simple docstring'''
   pass

  def  __lowercase							(	self		:      List[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	List[Any]        =       self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    UpperCAmelCase__       :	List[Any]        =       model_class(a__     )
    self.assertIsInstance(model.get_input_embeddings()       ,(nn.Module)     )
    UpperCAmelCase__       :	str        =       model.get_output_embeddings()
    self.assertTrue(x is None or isinstance(a__       ,nn.Linear     )     )

  def  __lowercase							(	self		:      Optional[int]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Any        =       self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    UpperCAmelCase__       :	Optional[int]        =       model_class(a__     )
    UpperCAmelCase__       :	str        =       inspect.signature(model.forward     )
    # signature.parameters is an OrderedDict => so arg_names order is deterministic
    UpperCAmelCase__       :	str        =       [*signature.parameters.keys()]

    UpperCAmelCase__       :	Optional[int]        =       ["pixel_values"]
    self.assertListEqual(arg_names[:1]       ,a__     )

  @unittest.skip(reason="""MaskFormerSwin is only used as backbone and doesn't support output_attentions"""     )
  def  __lowercase							(	self		:      Optional[int]     ):






   '''simple docstring'''
   pass

  @unittest.skip(reason="""MaskFormerSwin is only used as an internal backbone"""     )
  def  __lowercase							(	self		:      Tuple     ):






   '''simple docstring'''
   pass

  def  __lowercase							(	self		:      List[str]       ,A		:      Tuple       ,A		:      int       ,A		:      Optional[Any]       ,A		:      Optional[int]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Tuple        =       model_class(a__     )
   model.to(a__     )
   model.eval()

   with torch.no_grad():
    UpperCAmelCase__       :	Union[str, Any]        =       model(**self._prepare_for_class(a__       ,a__     )     )

   UpperCAmelCase__       :	Dict        =       outputs.hidden_states

   UpperCAmelCase__       :	Tuple        =       getattr(
       self.model_tester       ,"""expected_num_hidden_layers"""       ,len(self.model_tester.depths     ) + 1     )
   self.assertEqual(len(a__     )       ,a__     )

   # Swin has a different seq_length
   UpperCAmelCase__       :	Any        =       (
       config.patch_size
       if isinstance(config.patch_size       ,collections.abc.Iterable     )
       else (config.patch_size, config.patch_size)
   )

   UpperCAmelCase__       :	int        =       (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

   self.assertListEqual(
       list(hidden_states[0].shape[-2:]     )       ,[num_patches, self.model_tester.embed_dim]       ,)

  def  __lowercase							(	self		:      Dict     ):






   '''simple docstring'''
   UpperCAmelCase__       :	List[Any]        =       self.model_tester.prepare_config_and_inputs_for_common()

   UpperCAmelCase__       :	Union[str, Any]        =       (
       self.model_tester.image_size
       if isinstance(self.model_tester.image_size       ,collections.abc.Iterable     )
       else (self.model_tester.image_size, self.model_tester.image_size)
   )

   for model_class in self.all_model_classes:
    UpperCAmelCase__       :	int        =       True
    self.check_hidden_states_output(a__       ,a__       ,a__       ,a__     )

    # check that output_hidden_states also work using config
    del inputs_dict["output_hidden_states"]
    UpperCAmelCase__       :	List[Any]        =       True

    self.check_hidden_states_output(a__       ,a__       ,a__       ,a__     )

  def  __lowercase							(	self		:      Dict     ):






   '''simple docstring'''
   UpperCAmelCase__       :	List[Any]        =       self.model_tester.prepare_config_and_inputs_for_common()
   UpperCAmelCase__       :	Dict        =       3

   UpperCAmelCase__       :	str        =       (
       self.model_tester.image_size
       if isinstance(self.model_tester.image_size       ,collections.abc.Iterable     )
       else (self.model_tester.image_size, self.model_tester.image_size)
   )
   UpperCAmelCase__       :	str        =       (
       config.patch_size
       if isinstance(config.patch_size       ,collections.abc.Iterable     )
       else (config.patch_size, config.patch_size)
   )

   UpperCAmelCase__       :	Union[str, Any]        =       image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
   UpperCAmelCase__       :	str        =       image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

   for model_class in self.all_model_classes:
    UpperCAmelCase__       :	List[Any]        =       True
    self.check_hidden_states_output(a__       ,a__       ,a__       ,(padded_height, padded_width)     )

    # check that output_hidden_states also work using config
    del inputs_dict["output_hidden_states"]
    UpperCAmelCase__       :	Optional[Any]        =       True
    self.check_hidden_states_output(a__       ,a__       ,a__       ,(padded_height, padded_width)     )

  @unittest.skip(reason="""MaskFormerSwin doesn't have pretrained checkpoints"""     )
  def  __lowercase							(	self		:      str     ):






   '''simple docstring'''
   pass

  @unittest.skip(reason="""This will be fixed once MaskFormerSwin is replaced by native Swin"""     )
  def  __lowercase							(	self		:      Any     ):






   '''simple docstring'''
   pass

  @unittest.skip(reason="""This will be fixed once MaskFormerSwin is replaced by native Swin"""     )
  def  __lowercase							(	self		:      List[str]     ):






   '''simple docstring'''
   pass




  def  __lowercase							(	self		:      str     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[Any]        =       self.model_tester.prepare_config_and_inputs_for_common()

   def set_nan_tensor_to_zero(A		:      str     ):
    UpperCAmelCase__       :	Optional[Any]        =       0
    return t

   def check_equivalence(A		:      Union[str, Any]       ,A		:      Dict       ,A		:      Optional[Any]       ,A		:      int={}     ):
    with torch.no_grad():
     UpperCAmelCase__       :	List[str]        =       model(**a__       ,return_dict=a__       ,**a__     )
     UpperCAmelCase__       :	List[Any]        =       model(**a__       ,return_dict=a__       ,**a__     ).to_tuple()

     def recursive_check(A		:      int       ,A		:      List[Any]     ):
      if isinstance(a__       ,(List, Tuple)     ):
       for tuple_iterable_value, dict_iterable_value in zip(a__       ,a__     ):
        recursive_check(a__       ,a__     )
      elif isinstance(a__       ,a__     ):
       for tuple_iterable_value, dict_iterable_value in zip(
           tuple_object.values()       ,dict_object.values()     ):
        recursive_check(a__       ,a__     )
      elif tuple_object is None:
       return
      else:
       self.assertTrue(
           torch.allclose(
               set_nan_tensor_to_zero(a__     )       ,set_nan_tensor_to_zero(a__     )       ,atol=1e-5     )       ,msg=(
               """Tuple and dict output are not equal. Difference:"""
               f" {torch.max(torch.abs(tuple_object - dict_object     )     )}. Tuple has `nan`:"
               f" {torch.isnan(a__     ).any()} and `inf`: {torch.isinf(a__     )}. Dict has"
               f" `nan`: {torch.isnan(a__     ).any()} and `inf`: {torch.isinf(a__     )}."
           )       ,)

     recursive_check(a__       ,a__     )

   for model_class in self.all_model_classes:
    UpperCAmelCase__       :	Optional[int]        =       model_class(a__     )
    model.to(a__     )
    model.eval()

    UpperCAmelCase__       :	List[str]        =       self._prepare_for_class(a__       ,a__     )
    UpperCAmelCase__       :	List[str]        =       self._prepare_for_class(a__       ,a__     )
    check_equivalence(a__       ,a__       ,a__     )

    UpperCAmelCase__       :	Tuple        =       self._prepare_for_class(a__       ,a__       ,return_labels=a__     )
    UpperCAmelCase__       :	List[str]        =       self._prepare_for_class(a__       ,a__       ,return_labels=a__     )
    check_equivalence(a__       ,a__       ,a__     )

    UpperCAmelCase__       :	Optional[int]        =       self._prepare_for_class(a__       ,a__     )
    UpperCAmelCase__       :	List[str]        =       self._prepare_for_class(a__       ,a__     )
    check_equivalence(a__       ,a__       ,a__       ,{"""output_hidden_states""": True}     )

    UpperCAmelCase__       :	Union[str, Any]        =       self._prepare_for_class(a__       ,a__       ,return_labels=a__     )
    UpperCAmelCase__       :	int        =       self._prepare_for_class(a__       ,a__       ,return_labels=a__     )
    check_equivalence(a__       ,a__       ,a__       ,{"""output_hidden_states""": True}     )







@require_torch
class   __lowercase      (     unittest.TestCase						,      UpperCamelCase_       ):
  snake_case_												=							(MaskFormerSwinBackbone,) if is_torch_available() else ()
  snake_case_												=							MaskFormerSwinConfig

  def  __lowercase							(	self		:      List[str]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	List[str]        =       MaskFormerSwinModelTester(self     )




  def  __lowercase							(	self		:      List[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[Any]        =       self.model_tester.prepare_config_and_inputs_for_common()
   UpperCAmelCase__       :	List[str]        =       inputs_dict["pixel_values"].shape[0]

   for backbone_class in self.all_model_classes:
    UpperCAmelCase__       :	int        =       backbone_class(a__     )
    backbone.to(a__     )
    backbone.eval()

    UpperCAmelCase__       :	Dict        =       backbone(**a__     )

    # Test default outputs and verify feature maps
    self.assertIsInstance(outputs.feature_maps       ,a__     )
    self.assertTrue(len(outputs.feature_maps     ) == len(backbone.channels     )     )
    for feature_map, n_channels in zip(outputs.feature_maps       ,backbone.channels     ):
     self.assertTrue(feature_map.shape[:2]       ,(batch_size, n_channels)     )
    self.assertIsNone(outputs.hidden_states     )
    self.assertIsNone(outputs.attentions     )

    # Test output_hidden_states=True
    UpperCAmelCase__       :	Tuple        =       backbone(**a__       ,output_hidden_states=a__     )
    self.assertIsNotNone(outputs.hidden_states     )
    self.assertTrue(len(outputs.hidden_states     )       ,len(backbone.stage_names     )     )
    # We skip the stem layer
    for hidden_states, n_channels in zip(outputs.hidden_states[1:]       ,backbone.channels     ):
     for hidden_state in hidden_states:
      # Hidden states are in the format (batch_size, (height * width), n_channels)
      UpperCAmelCase__       :	Any        =       hidden_state.shape
      self.assertTrue((h_batch_size, h_n_channels)       ,(batch_size, n_channels)     )

             # Test output_attentions=True
    if self.has_attentions:
     UpperCAmelCase__       :	Any        =       backbone(**a__       ,output_attentions=a__     )
     self.assertIsNotNone(outputs.attentions     )







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


 def    _A	(    self  : Union[str, Any]    ):
  '''simple docstring'''

  torch.manual_seed(0    )
  lowerCAmelCase__     :   List[str]     		=	TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"    )

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Optional[int]     		=	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"    )

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Any     		=	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    )
  lowerCAmelCase__     :   Union[str, Any]     		=	DDPMScheduler(
      num_train_timesteps=1000			,							beta_schedule="squaredcos_cap_v2"			,							beta_start=0.0001			,							beta_end=0.02			,							thresholding=a__			,							dynamic_thresholding_ratio=0.95			,							sample_max_value=1.0			,							prediction_type="epsilon"			,							variance_type="learned_range"			,							)

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Tuple     		=	IFWatermarker()

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


 def    _A	(    self  : Dict    ):
  '''simple docstring'''

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Any     		=	TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"    )

  torch.manual_seed(0    )
  lowerCAmelCase__     :   List[Any]     		=	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"    )

  torch.manual_seed(0    )
  lowerCAmelCase__     :   int     		=	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    )
  lowerCAmelCase__     :   Tuple     		=	DDPMScheduler(
      num_train_timesteps=1000			,							beta_schedule="squaredcos_cap_v2"			,							beta_start=0.0001			,							beta_end=0.02			,							thresholding=a__			,							dynamic_thresholding_ratio=0.95			,							sample_max_value=1.0			,							prediction_type="epsilon"			,							variance_type="learned_range"			,							)

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Tuple     		=	DDPMScheduler(
      num_train_timesteps=1000			,							beta_schedule="squaredcos_cap_v2"			,							beta_start=0.0001			,							beta_end=0.02			,							)

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Dict     		=	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    _A	(    self  : Optional[Any]    ):
  '''simple docstring'''

  lowerCAmelCase__     :   Union[str, Any]     		=	self.get_dummy_components()
  lowerCAmelCase__     :   Dict     		=	self.pipeline_class(**a__    )
  pipe.to(a__    )
  pipe.set_progress_bar_config(disable=a__    )

  lowerCAmelCase__     :   Tuple     		=	self.get_dummy_inputs(a__    )

  lowerCAmelCase__     :   str     		=	inputs["prompt"]
  lowerCAmelCase__     :   Optional[Any]     		=	inputs["generator"]
  lowerCAmelCase__     :   str     		=	inputs["num_inference_steps"]
  lowerCAmelCase__     :   int     		=	inputs["output_type"]

  if "image" in inputs:
   lowerCAmelCase__     :   Optional[int]     		=	inputs["image"]
  else:
   lowerCAmelCase__     :   str     		=	None

  if "mask_image" in inputs:
   lowerCAmelCase__     :   List[str]     		=	inputs["mask_image"]
  else:
   lowerCAmelCase__     :   str     		=	None

  if "original_image" in inputs:
   lowerCAmelCase__     :   Optional[int]     		=	inputs["original_image"]
  else:
   lowerCAmelCase__     :   Union[str, Any]     		=	None

  lowerCAmelCase__		,     lowerCAmelCase__     :   Dict     		=	pipe.encode_prompt(a__    )

  # inputs with prompt converted to embeddings
  lowerCAmelCase__     :   Dict     		=	{
      "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:
   lowerCAmelCase__     :   Optional[Any]     		=	image

  if mask_image is not None:
   lowerCAmelCase__     :   Optional[int]     		=	mask_image

  if original_image is not None:
   lowerCAmelCase__     :   Optional[Any]     		=	original_image

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

  lowerCAmelCase__     :   Optional[int]     		=	pipe(**a__    )[0]

  with tempfile.TemporaryDirectory() as tmpdir:
   pipe.save_pretrained(a__    )
   lowerCAmelCase__     :   Tuple     		=	self.pipeline_class.from_pretrained(a__    )
   pipe_loaded.to(a__    )
   pipe_loaded.set_progress_bar_config(disable=a__    )

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

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

  lowerCAmelCase__     :   Union[str, Any]     		=	self.get_dummy_inputs(a__    )

  lowerCAmelCase__     :   Any     		=	inputs["generator"]
  lowerCAmelCase__     :   Union[str, Any]     		=	inputs["num_inference_steps"]
  lowerCAmelCase__     :   str     		=	inputs["output_type"]

  # inputs with prompt converted to embeddings
  lowerCAmelCase__     :   int     		=	{
      "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:
   lowerCAmelCase__     :   Tuple     		=	image

  if mask_image is not None:
   lowerCAmelCase__     :   int     		=	mask_image

  if original_image is not None:
   lowerCAmelCase__     :   List[Any]     		=	original_image

  lowerCAmelCase__     :   Dict     		=	pipe_loaded(**a__    )[0]

  lowerCAmelCase__     :   int     		=	np.abs(to_np(a__    ) - to_np(a__    )    ).max()
  self.assertLess(a__			,							1e-4    )


 def    _A	(    self  : Any    ):
  '''simple docstring'''

  lowerCAmelCase__     :   str     		=	self.get_dummy_components()
  lowerCAmelCase__     :   int     		=	self.pipeline_class(**a__    )
  pipe.to(a__    )
  pipe.set_progress_bar_config(disable=a__    )

  lowerCAmelCase__     :   Any     		=	self.get_dummy_inputs(a__    )
  lowerCAmelCase__     :   Union[str, Any]     		=	pipe(**a__    )[0]

  with tempfile.TemporaryDirectory() as tmpdir:
   pipe.save_pretrained(a__    )
   lowerCAmelCase__     :   str     		=	self.pipeline_class.from_pretrained(a__    )
   pipe_loaded.to(a__    )
   pipe_loaded.set_progress_bar_config(disable=a__    )

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

  lowerCAmelCase__     :   str     		=	self.get_dummy_inputs(a__    )
  lowerCAmelCase__     :   Union[str, Any]     		=	pipe_loaded(**a__    )[0]

  lowerCAmelCase__     :   Union[str, Any]     		=	np.abs(to_np(a__    ) - to_np(a__    )    ).max()
  self.assertLess(a__			,							1e-4    )





