





'''simple docstring'''






import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import DetrConfig, MaskFormerConfig, SwinConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import MaskFormerForInstanceSegmentation, MaskFormerModel

      if is_vision_available():
            from transformers import MaskFormerImageProcessor

if is_vision_available():
      from PIL import Image







class 		UpperCAmelCase   :


  """simple docstring"""







  def __init__(				self    :					List[str]							,	UpperCamelCase__    :					Tuple							,	UpperCamelCase__    :					List[str]=2							,	UpperCamelCase__    :					Dict=True							,	UpperCamelCase__    :					Union[str, Any]=False							,	UpperCamelCase__    :					Optional[int]=10							,	UpperCamelCase__    :					str=3							,	UpperCamelCase__    :					Tuple=32 * 4							,	UpperCamelCase__    :					Any=32 * 6							,	UpperCamelCase__    :					Union[str, Any]=4							,	UpperCamelCase__    :					List[Any]=32							,	)     ->						Any:
    _UpperCamelCase       			=parent
    _UpperCamelCase       			=batch_size
    _UpperCamelCase       			=is_training
    _UpperCamelCase       			=use_auxiliary_loss
    _UpperCamelCase       			=num_queries
    _UpperCamelCase       			=num_channels
    _UpperCamelCase       			=min_size
    _UpperCamelCase       			=max_size
    _UpperCamelCase       			=num_labels
    _UpperCamelCase       			=mask_feature_size


  def   UpperCamelCase__      (				self    :					str  )     ->						Any:
    _UpperCamelCase       			=floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]  ).to(
        UpperCamelCase__  )

    _UpperCamelCase       			=torch.ones([self.batch_size, self.min_size, self.max_size]							,	device=UpperCamelCase__  )

    _UpperCamelCase       			=(
        torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]							,	device=UpperCamelCase__  ) > 0.5
    ).float()
    _UpperCamelCase       			=(torch.rand((self.batch_size, self.num_labels)							,	device=UpperCamelCase__  ) > 0.5).long()

    _UpperCamelCase       			=self.get_config()
    return config, pixel_values, pixel_mask, mask_labels, class_labels


  def   UpperCamelCase__      (				self    :					int  )     ->						int:
    return MaskFormerConfig.from_backbone_and_decoder_configs(
        backbone_config=SwinConfig(
            depths=[1, 1, 1, 1]							,	)							,	decoder_config=DetrConfig(
            decoder_ffn_dim=128							,	num_queries=self.num_queries							,	decoder_attention_heads=2							,	d_model=self.mask_feature_size							,	)							,	mask_feature_size=self.mask_feature_size							,	fpn_feature_size=self.mask_feature_size							,	num_channels=self.num_channels							,	num_labels=self.num_labels							,	)


  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Any:
    _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase       			=self.prepare_config_and_inputs()
    _UpperCamelCase       			={'''pixel_values''': pixel_values, '''pixel_mask''': pixel_mask}
    return config, inputs_dict


  def   UpperCamelCase__      (				self    :					Optional[int]							,	UpperCamelCase__    :					List[Any]							,	UpperCamelCase__    :					int  )     ->						int:
    _UpperCamelCase       			=output.encoder_hidden_states
    _UpperCamelCase       			=output.pixel_decoder_hidden_states
    _UpperCamelCase       			=output.transformer_decoder_hidden_states

    self.parent.assertTrue(len(UpperCamelCase__  )							,	len(config.backbone_config.depths  )  )
    self.parent.assertTrue(len(UpperCamelCase__  )							,	len(config.backbone_config.depths  )  )
    self.parent.assertTrue(len(UpperCamelCase__  )							,	config.decoder_config.decoder_layers  )


  def   UpperCamelCase__      (				self    :					Any							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					Union[str, Any]=False  )     ->						List[Any]:
    with torch.no_grad():
      _UpperCamelCase       			=MaskFormerModel(config=UpperCamelCase__  )
      model.to(UpperCamelCase__  )
      model.eval()

      _UpperCamelCase       			=model(pixel_values=UpperCamelCase__							,	pixel_mask=UpperCamelCase__  )
      _UpperCamelCase       			=model(UpperCamelCase__							,	output_hidden_states=UpperCamelCase__  )
    # the correct shape of output.transformer_decoder_hidden_states ensure the correcteness of the
    # encoder and pixel decoder
    self.parent.assertEqual(
        output.transformer_decoder_last_hidden_state.shape							,	(self.batch_size, self.num_queries, self.mask_feature_size)							,	)
    # let's ensure the other two hidden state exists
    self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None  )
    self.parent.assertTrue(output.encoder_last_hidden_state is not None  )

    if output_hidden_states:
      self.check_output_hidden_state(UpperCamelCase__							,	UpperCamelCase__  )




  def   UpperCamelCase__      (				self    :					List[str]							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					Any							,	UpperCamelCase__    :					Optional[Any]							,	UpperCamelCase__    :					List[str]  )     ->						Optional[Any]:
    _UpperCamelCase       			=MaskFormerForInstanceSegmentation(config=UpperCamelCase__  )
    model.to(UpperCamelCase__  )
    model.eval()

    def comm_check_on_output(UpperCamelCase__    :					List[Any]  ):
      # let's still check that all the required stuff is there
      self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None  )
      self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None  )
      self.parent.assertTrue(result.encoder_last_hidden_state is not None  )
      # okay, now we need to check the logits shape
      # due to the encoder compression, masks have a //4 spatial size
      self.parent.assertEqual(
          result.masks_queries_logits.shape							,	(self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)							,	)
      # + 1 for null class
      self.parent.assertEqual(
          result.class_queries_logits.shape							,	(self.batch_size, self.num_queries, self.num_labels + 1)  )

    with torch.no_grad():
      _UpperCamelCase       			=model(pixel_values=UpperCamelCase__							,	pixel_mask=UpperCamelCase__  )
      _UpperCamelCase       			=model(UpperCamelCase__  )

      comm_check_on_output(UpperCamelCase__  )

      _UpperCamelCase       			=model(
          pixel_values=UpperCamelCase__							,	pixel_mask=UpperCamelCase__							,	mask_labels=UpperCamelCase__							,	class_labels=UpperCamelCase__  )

    comm_check_on_output(UpperCamelCase__  )

    self.parent.assertTrue(result.loss is not None  )
    self.parent.assertEqual(result.loss.shape							,	torch.Size([1]  )  )







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


  """simple docstring"""







  lowerCAmelCase_					=				(MaskFormerModel, MaskFormerForInstanceSegmentation) if is_torch_available() else ()
  lowerCAmelCase_					=				(
      {"""feature-extraction""": MaskFormerModel, """image-segmentation""": MaskFormerForInstanceSegmentation}
      if is_torch_available()
      else {}
  )

  lowerCAmelCase_					=				False
  lowerCAmelCase_					=				False
  lowerCAmelCase_					=				False
  lowerCAmelCase_					=				False


  def   UpperCamelCase__      (				self    :					Tuple  )     ->						Optional[int]:
    _UpperCamelCase       			=MaskFormerModelTester(self  )
    _UpperCamelCase       			=ConfigTester(self							,	config_class=UpperCamelCase__							,	has_text_modality=UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						str:
    self.config_tester.run_common_tests()


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						List[Any]:
    _UpperCamelCase			,     _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs_for_common()
    self.model_tester.create_and_check_maskformer_model(UpperCamelCase__							,	**UpperCamelCase__							,	output_hidden_states=UpperCamelCase__  )


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


  @unittest.skip(reason='''MaskFormer does not use inputs_embeds'''  )
  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						int:
    pass


  @unittest.skip(reason='''MaskFormer does not have a get_input_embeddings method'''  )
  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						List[str]:
    pass


  @unittest.skip(reason='''MaskFormer is not a generative model'''  )
  def   UpperCamelCase__      (				self    :					int  )     ->						List[Any]:
    pass


  @unittest.skip(reason='''MaskFormer does not use token embeddings'''  )
  def   UpperCamelCase__      (				self    :					Dict  )     ->						Optional[int]:
    pass


  @require_torch_multi_gpu
  @unittest.skip(
      reason='''MaskFormer has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''  )
  def   UpperCamelCase__      (				self    :					int  )     ->						Dict:
    pass


  @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''  )
  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						Union[str, Any]:
    pass


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

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

      _UpperCamelCase       			=['''pixel_values''']
      self.assertListEqual(arg_names[:1]							,	UpperCamelCase__  )


  @slow
  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						Union[str, Any]:
    for model_name in ["facebook/maskformer-swin-small-coco"]:
      _UpperCamelCase       			=MaskFormerModel.from_pretrained(UpperCamelCase__  )
      self.assertIsNotNone(UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						List[str]:
    _UpperCamelCase       			=(self.model_tester.min_size,) * 2
    _UpperCamelCase       			={
        '''pixel_values''': torch.randn((2, 3, *size)							,	device=UpperCamelCase__  ),
        '''mask_labels''': torch.randn((2, 10, *size)							,	device=UpperCamelCase__  ),
        '''class_labels''': torch.zeros(2							,	10							,	device=UpperCamelCase__  ).long(),
    }

    _UpperCamelCase       			=MaskFormerForInstanceSegmentation(MaskFormerConfig()  ).to(UpperCamelCase__  )
    _UpperCamelCase       			=model(**UpperCamelCase__  )
    self.assertTrue(outputs.loss is not None  )


  def   UpperCamelCase__      (				self    :					List[str]  )     ->						Dict:
    _UpperCamelCase			,     _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs_for_common()
    self.model_tester.create_and_check_maskformer_model(UpperCamelCase__							,	**UpperCamelCase__							,	output_hidden_states=UpperCamelCase__  )


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

    for model_class in self.all_model_classes:
      _UpperCamelCase       			=model_class(UpperCamelCase__  ).to(UpperCamelCase__  )
      _UpperCamelCase       			=model(**UpperCamelCase__							,	output_attentions=UpperCamelCase__  )
      self.assertTrue(outputs.attentions is not None  )


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						int:
    if not self.model_tester.is_training:
      return
    # only MaskFormerForInstanceSegmentation has the loss
    _UpperCamelCase       			=self.all_model_classes[1]
    _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()

    _UpperCamelCase       			=model_class(UpperCamelCase__  )
    model.to(UpperCamelCase__  )
    model.train()

    _UpperCamelCase       			=model(UpperCamelCase__							,	mask_labels=UpperCamelCase__							,	class_labels=UpperCamelCase__  ).loss
    loss.backward()




  def   UpperCamelCase__      (				self    :					Dict  )     ->						Optional[int]:
    # only MaskFormerForInstanceSegmentation has the loss
    _UpperCamelCase       			=self.all_model_classes[1]
    _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    _UpperCamelCase       			=True
    _UpperCamelCase       			=True

    _UpperCamelCase       			=model_class(UpperCamelCase__  )
    model.to(UpperCamelCase__  )
    model.train()

    _UpperCamelCase       			=model(UpperCamelCase__							,	mask_labels=UpperCamelCase__							,	class_labels=UpperCamelCase__  )

    _UpperCamelCase       			=outputs.encoder_hidden_states[0]
    encoder_hidden_states.retain_grad()

    _UpperCamelCase       			=outputs.pixel_decoder_hidden_states[0]
    pixel_decoder_hidden_states.retain_grad()
    # we requires_grad=True in inputs_embeds (line 2152), the original implementation don't
    _UpperCamelCase       			=outputs.transformer_decoder_hidden_states[0]
    transformer_decoder_hidden_states.retain_grad()

    _UpperCamelCase       			=outputs.attentions[0]
    attentions.retain_grad()

    outputs.loss.backward(retain_graph=UpperCamelCase__  )

    self.assertIsNotNone(encoder_hidden_states.grad  )
    self.assertIsNotNone(pixel_decoder_hidden_states.grad  )
    self.assertIsNotNone(transformer_decoder_hidden_states.grad  )
    self.assertIsNotNone(attentions.grad  )




__lowerCamelCase						: Any		=   1E-4


def 					_a   ():


  """simple docstring"""

  _UpperCamelCase       			=Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''   )
  return image







@require_vision
@slow
class 		UpperCAmelCase   (				unittest.TestCase):


  """simple docstring"""







  @cached_property
  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Dict:
    return (
        MaskFormerImageProcessor.from_pretrained('''facebook/maskformer-swin-small-coco'''  )
        if is_vision_available()
        else None
    )


  def   UpperCamelCase__      (				self    :					Any  )     ->						Dict:
    _UpperCamelCase       			=MaskFormerModel.from_pretrained('''facebook/maskformer-swin-small-coco'''  ).to(UpperCamelCase__  )
    _UpperCamelCase       			=self.default_image_processor
    _UpperCamelCase       			=prepare_img()
    _UpperCamelCase       			=image_processor(UpperCamelCase__							,	return_tensors='''pt'''  ).to(UpperCamelCase__  )
    _UpperCamelCase       			=inputs['''pixel_values'''].shape
    # check size is divisible by 32
    self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0  )
    # check size
    self.assertEqual(UpperCamelCase__							,	(1, 3, 800, 1088)  )

    with torch.no_grad():
      _UpperCamelCase       			=model(**UpperCamelCase__  )

    _UpperCamelCase       			=torch.tensor(
        [[-0.0482, 0.9228, 0.4951], [-0.2547, 0.8017, 0.8527], [-0.0069, 0.3385, -0.0089]]  ).to(UpperCamelCase__  )
    self.assertTrue(
        torch.allclose(
            outputs.encoder_last_hidden_state[0, 0, :3, :3]							,	UpperCamelCase__							,	atol=UpperCamelCase__  )  )

    _UpperCamelCase       			=torch.tensor(
        [[-0.8422, -0.8434, -0.9718], [-1.0144, -0.5565, -0.4195], [-1.0038, -0.4484, -0.1961]]  ).to(UpperCamelCase__  )
    self.assertTrue(
        torch.allclose(
            outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]							,	UpperCamelCase__							,	atol=UpperCamelCase__  )  )

    _UpperCamelCase       			=torch.tensor(
        [[0.2852, -0.0159, 0.9735], [0.6254, 0.1858, 0.8529], [-0.0680, -0.4116, 1.8413]]  ).to(UpperCamelCase__  )
    self.assertTrue(
        torch.allclose(
            outputs.transformer_decoder_last_hidden_state[0, :3, :3]							,	UpperCamelCase__							,	atol=UpperCamelCase__  )  )


  def   UpperCamelCase__      (				self    :					int  )     ->						Optional[int]:
    _UpperCamelCase       			=(
        MaskFormerForInstanceSegmentation.from_pretrained('''facebook/maskformer-swin-small-coco'''  )
        .to(UpperCamelCase__  )
        .eval()
    )
    _UpperCamelCase       			=self.default_image_processor
    _UpperCamelCase       			=prepare_img()
    _UpperCamelCase       			=image_processor(UpperCamelCase__							,	return_tensors='''pt'''  ).to(UpperCamelCase__  )
    _UpperCamelCase       			=inputs['''pixel_values'''].shape
    # check size is divisible by 32
    self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0  )
    # check size
    self.assertEqual(UpperCamelCase__							,	(1, 3, 800, 1088)  )

    with torch.no_grad():
      _UpperCamelCase       			=model(**UpperCamelCase__  )
    # masks_queries_logits
    _UpperCamelCase       			=outputs.masks_queries_logits
    self.assertEqual(
        masks_queries_logits.shape							,	(1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)							,	)
    _UpperCamelCase       			=[
        [-1.3737124, -1.7724937, -1.9364233],
        [-1.5977281, -1.9867939, -2.1523695],
        [-1.5795398, -1.9269832, -2.093942],
    ]
    _UpperCamelCase       			=torch.tensor(UpperCamelCase__  ).to(UpperCamelCase__  )
    self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]							,	UpperCamelCase__							,	atol=UpperCamelCase__  )  )
    # class_queries_logits
    _UpperCamelCase       			=outputs.class_queries_logits
    self.assertEqual(
        class_queries_logits.shape							,	(1, model.config.decoder_config.num_queries, model.config.num_labels + 1)  )
    _UpperCamelCase       			=torch.tensor(
        [
            [1.6512E00, -5.2572E00, -3.3519E00],
            [3.6169E-02, -5.9025E00, -2.9313E00],
            [1.0766E-04, -7.7630E00, -5.1263E00],
        ]  ).to(UpperCamelCase__  )
    self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]							,	UpperCamelCase__							,	atol=UpperCamelCase__  )  )


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Tuple:
    _UpperCamelCase       			=(
        MaskFormerForInstanceSegmentation.from_pretrained('''facebook/maskformer-resnet101-coco-stuff'''  )
        .to(UpperCamelCase__  )
        .eval()
    )
    _UpperCamelCase       			=self.default_image_processor
    _UpperCamelCase       			=prepare_img()
    _UpperCamelCase       			=image_processor(UpperCamelCase__							,	return_tensors='''pt'''  ).to(UpperCamelCase__  )
    _UpperCamelCase       			=inputs['''pixel_values'''].shape
    # check size is divisible by 32
    self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0  )
    # check size
    self.assertEqual(UpperCamelCase__							,	(1, 3, 800, 1088)  )

    with torch.no_grad():
      _UpperCamelCase       			=model(**UpperCamelCase__  )
    # masks_queries_logits
    _UpperCamelCase       			=outputs.masks_queries_logits
    self.assertEqual(
        masks_queries_logits.shape							,	(1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)							,	)
    _UpperCamelCase       			=[[-0.9046, -2.6366, -4.6062], [-3.4179, -5.7890, -8.8057], [-4.9179, -7.6560, -10.7711]]
    _UpperCamelCase       			=torch.tensor(UpperCamelCase__  ).to(UpperCamelCase__  )
    self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]							,	UpperCamelCase__							,	atol=UpperCamelCase__  )  )
    # class_queries_logits
    _UpperCamelCase       			=outputs.class_queries_logits
    self.assertEqual(
        class_queries_logits.shape							,	(1, model.config.decoder_config.num_queries, model.config.num_labels + 1)  )
    _UpperCamelCase       			=torch.tensor(
        [[4.7188, -3.2585, -2.8857], [6.6871, -2.9181, -1.2487], [7.2449, -2.2764, -2.1874]]  ).to(UpperCamelCase__  )
    self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]							,	UpperCamelCase__							,	atol=UpperCamelCase__  )  )




  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						int:
    _UpperCamelCase       			=(
        MaskFormerForInstanceSegmentation.from_pretrained('''facebook/maskformer-swin-small-coco'''  )
        .to(UpperCamelCase__  )
        .eval()
    )
    _UpperCamelCase       			=self.default_image_processor

    _UpperCamelCase       			=image_processor(
        [np.zeros((3, 800, 1333)  ), np.zeros((3, 800, 1333)  )]							,	segmentation_maps=[np.zeros((384, 384)  ).astype(np.floataa  ), np.zeros((384, 384)  ).astype(np.floataa  )]							,	return_tensors='''pt'''							,	)

    _UpperCamelCase       			=inputs['''pixel_values'''].to(UpperCamelCase__  )
    _UpperCamelCase       			=[el.to(UpperCamelCase__  ) for el in inputs['''mask_labels''']]
    _UpperCamelCase       			=[el.to(UpperCamelCase__  ) for el in inputs['''class_labels''']]

    with torch.no_grad():
      _UpperCamelCase       			=model(**UpperCamelCase__  )

    self.assertTrue(outputs.loss is not None  )













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






