






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




    def __init__(			self  :       List[Any]  ,     a  :       Dict  ,     a  :       List[str]=2  ,     a  :       List[Any]=True  ,     a  :       Union[str, Any]=False  ,     a  :       int=1_0  ,     a  :       Tuple=3  ,     a  :       Any=3_2 * 4  ,     a  :       str=3_2 * 6  ,     a  :       str=4  ,     a  :       Union[str, Any]=3_2  ,     ):






       '''simple docstring'''






       lowercase_	:    List[Any]           =     parent
       lowercase_	:    Dict           =     batch_size
       lowercase_	:    int           =     is_training
       lowercase_	:    Dict           =     use_auxiliary_loss
       lowercase_	:    Dict           =     num_queries
       lowercase_	:    Tuple           =     num_channels
       lowercase_	:    List[Any]           =     min_size
       lowercase_	:    Optional[int]           =     max_size
       lowercase_	:    List[str]           =     num_labels
       lowercase_	:    Union[str, Any]           =     mask_feature_size




    def 				lowerCAmelCase__				(			self  :       int			):






       '''simple docstring'''






       lowercase_	:    List[Any]           =     floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]			).to(
           _lowerCAmelCase			)

       lowercase_	:    Union[str, Any]           =     torch.ones([self.batch_size, self.min_size, self.max_size]  ,     device=_lowerCAmelCase			)

       lowercase_	:    str           =     (
           torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]  ,     device=_lowerCAmelCase			) > 0.5
       ).float()
       lowercase_	:    Optional[Any]           =     (torch.rand((self.batch_size, self.num_labels)  ,     device=_lowerCAmelCase			) > 0.5).long()

       lowercase_	:    int           =     self.get_config()
       return config, pixel_values, pixel_mask, mask_labels, class_labels




    def 				lowerCAmelCase__				(			self  :       Tuple			):






       '''simple docstring'''






       return MaskFormerConfig.from_backbone_and_decoder_configs(
           backbone_config=SwinConfig(
               depths=[1, 1, 1, 1]  ,     )  ,     decoder_config=DetrConfig(
               decoder_ffn_dim=1_2_8  ,     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 				lowerCAmelCase__				(			self  :       Tuple			):






       '''simple docstring'''






       lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_	:    Tuple           =     self.prepare_config_and_inputs()
       lowercase_	:    Any           =     {"pixel_values": pixel_values, "pixel_mask": pixel_mask}
       return config, inputs_dict




    def 				lowerCAmelCase__				(			self  :       List[Any]  ,     a  :       int  ,     a  :       Dict			):






       '''simple docstring'''






       lowercase_	:    Union[str, Any]           =     output.encoder_hidden_states
       lowercase_	:    List[str]           =     output.pixel_decoder_hidden_states
       lowercase_	:    Optional[Any]           =     output.transformer_decoder_hidden_states

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




    def 				lowerCAmelCase__				(			self  :       Optional[Any]  ,     a  :       Optional[Any]  ,     a  :       Dict  ,     a  :       Any  ,     a  :       List[Any]=False			):






       '''simple docstring'''






       with torch.no_grad():
          lowercase_	:    str           =     MaskFormerModel(config=_lowerCAmelCase			)
          model.to(_lowerCAmelCase			)
          model.eval()

          lowercase_	:    int           =     model(pixel_values=_lowerCAmelCase  ,     pixel_mask=_lowerCAmelCase			)
          lowercase_	:    Optional[int]           =     model(_lowerCAmelCase  ,     output_hidden_states=_lowerCAmelCase			)
       # 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(_lowerCAmelCase  ,     _lowerCAmelCase			)





    def 				lowerCAmelCase__				(			self  :       Optional[Any]  ,     a  :       int  ,     a  :       Any  ,     a  :       int  ,     a  :       Any  ,     a  :       int			):






       '''simple docstring'''






       lowercase_	:    Dict           =     MaskFormerForInstanceSegmentation(config=_lowerCAmelCase			)
       model.to(_lowerCAmelCase			)
       model.eval()

       def comm_check_on_output(a  :       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():
          lowercase_	:    Any           =     model(pixel_values=_lowerCAmelCase  ,     pixel_mask=_lowerCAmelCase			)
          lowercase_	:    List[Any]           =     model(_lowerCAmelCase			)

          comm_check_on_output(_lowerCAmelCase			)

          lowercase_	:    Optional[Any]           =     model(
              pixel_values=_lowerCAmelCase  ,     pixel_mask=_lowerCAmelCase  ,     mask_labels=_lowerCAmelCase  ,     class_labels=_lowerCAmelCase			)

       comm_check_on_output(_lowerCAmelCase			)

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


@require_torch
class        _UpperCAmelCase						(		_lowerCAmelCase				,					_lowerCAmelCase				,					unittest.TestCase			):
    __lowerCamelCase:   str		       =    (MaskFormerModel, MaskFormerForInstanceSegmentation) if is_torch_available() else ()
    __lowerCamelCase:   Tuple		       =    (
        {'feature-extraction': MaskFormerModel, 'image-segmentation': MaskFormerForInstanceSegmentation}
        if is_torch_available()
        else {}
    )

    __lowerCamelCase:   Tuple		       =    False
    __lowerCamelCase:   List[str]		       =    False
    __lowerCamelCase:   List[str]		       =    False
    __lowerCamelCase:   List[str]		       =    False




    def 				lowerCAmelCase__				(			self  :       Dict			):






       '''simple docstring'''






       lowercase_	:    str           =     MaskFormerModelTester(self			)
       lowercase_	:    Optional[Any]           =     ConfigTester(self  ,     config_class=_lowerCAmelCase  ,     has_text_modality=_lowerCAmelCase			)




    def 				lowerCAmelCase__				(			self  :       Tuple			):






       '''simple docstring'''






       self.config_tester.run_common_tests()




    def 				lowerCAmelCase__				(			self  :       str			):






       '''simple docstring'''






       lowercase_    ,		lowercase_	:    Optional[Any]           =     self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.create_and_check_maskformer_model(_lowerCAmelCase  ,     **_lowerCAmelCase  ,     output_hidden_states=_lowerCAmelCase			)




    def 				lowerCAmelCase__				(			self  :       List[Any]			):






       '''simple docstring'''






       lowercase_	:    str           =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_maskformer_instance_segmentation_head_model(*_lowerCAmelCase			)




    @unittest.skip(reason="MaskFormer does not use inputs_embeds"			)
    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       pass




    @unittest.skip(reason="MaskFormer does not have a get_input_embeddings method"			)
    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       pass




    @unittest.skip(reason="MaskFormer is not a generative model"			)
    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       pass




    @unittest.skip(reason="MaskFormer does not use token embeddings"			)
    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       pass




    @require_torch_multi_gpu
    @unittest.skip(
        reason="MaskFormer has some layers using `add_module` which doesn't work well with `nn.DataParallel`"			)
    def 				lowerCAmelCase__				(			self  :       Any			):






       '''simple docstring'''






       pass




    @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."			)
    def 				lowerCAmelCase__				(			self  :       Optional[Any]			):






       '''simple docstring'''






       pass




    def 				lowerCAmelCase__				(			self  :       int			):






       '''simple docstring'''






       lowercase_    ,		lowercase_	:    int           =     self.model_tester.prepare_config_and_inputs_for_common()

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

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




    @slow
    def 				lowerCAmelCase__				(			self  :       Any			):






       '''simple docstring'''






       for model_name in ["facebook/maskformer-swin-small-coco"]:
          lowercase_	:    Tuple           =     MaskFormerModel.from_pretrained(_lowerCAmelCase			)
          self.assertIsNotNone(_lowerCAmelCase			)




    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       lowercase_	:    Optional[int]           =     (self.model_tester.min_size,) * 2
       lowercase_	:    List[str]           =     {
           "pixel_values": torch.randn((2, 3, *size)  ,     device=_lowerCAmelCase			),
           "mask_labels": torch.randn((2, 1_0, *size)  ,     device=_lowerCAmelCase			),
           "class_labels": torch.zeros(2  ,     1_0  ,     device=_lowerCAmelCase			).long(),
       }

       lowercase_	:    Dict           =     MaskFormerForInstanceSegmentation(MaskFormerConfig()			).to(_lowerCAmelCase			)
       lowercase_	:    int           =     model(**_lowerCAmelCase			)
       self.assertTrue(outputs.loss is not None			)




    def 				lowerCAmelCase__				(			self  :       Optional[Any]			):






       '''simple docstring'''






       lowercase_    ,		lowercase_	:    str           =     self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.create_and_check_maskformer_model(_lowerCAmelCase  ,     **_lowerCAmelCase  ,     output_hidden_states=_lowerCAmelCase			)




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






       '''simple docstring'''






       lowercase_    ,		lowercase_	:    str           =     self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
          lowercase_	:    Optional[Any]           =     model_class(_lowerCAmelCase			).to(_lowerCAmelCase			)
          lowercase_	:    Dict           =     model(**_lowerCAmelCase  ,     output_attentions=_lowerCAmelCase			)
          self.assertTrue(outputs.attentions is not None			)




    def 				lowerCAmelCase__				(			self  :       Dict			):






       '''simple docstring'''






       if not self.model_tester.is_training:
          return
       # only MaskFormerForInstanceSegmentation has the loss
       lowercase_	:    Union[str, Any]           =     self.all_model_classes[1]
       lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_	:    List[str]           =     self.model_tester.prepare_config_and_inputs()

       lowercase_	:    Dict           =     model_class(_lowerCAmelCase			)
       model.to(_lowerCAmelCase			)
       model.train()

       lowercase_	:    List[str]           =     model(_lowerCAmelCase  ,     mask_labels=_lowerCAmelCase  ,     class_labels=_lowerCAmelCase			).loss
       loss.backward()





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






       '''simple docstring'''






       lowercase_	:    List[str]           =     self.all_model_classes[1]
       lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_	:    Any           =     self.model_tester.prepare_config_and_inputs()
       lowercase_	:    List[Any]           =     True
       lowercase_	:    Optional[int]           =     True

       lowercase_	:    Union[str, Any]           =     model_class(_lowerCAmelCase			)
       model.to(_lowerCAmelCase			)
       model.train()

       lowercase_	:    List[str]           =     model(_lowerCAmelCase  ,     mask_labels=_lowerCAmelCase  ,     class_labels=_lowerCAmelCase			)

       lowercase_	:    Optional[Any]           =     outputs.encoder_hidden_states[0]
       encoder_hidden_states.retain_grad()

       lowercase_	:    Optional[Any]           =     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
       lowercase_	:    Tuple           =     outputs.transformer_decoder_hidden_states[0]
       transformer_decoder_hidden_states.retain_grad()

       lowercase_	:    str           =     outputs.attentions[0]
       attentions.retain_grad()

       outputs.loss.backward(retain_graph=_lowerCAmelCase			)

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


UpperCamelCase__ =			1e-4




def 		__SCREAMING_SNAKE_CASE					(    ):


   """simple docstring"""




   lowercase_	:    List[str]           =     Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"							)
   return image





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




    @cached_property
    def 				lowerCAmelCase__				(			self  :       Dict			):






       '''simple docstring'''






       return (
           MaskFormerImageProcessor.from_pretrained("facebook/maskformer-swin-small-coco"			)
           if is_vision_available()
           else None
       )




    def 				lowerCAmelCase__				(			self  :       str			):






       '''simple docstring'''






       lowercase_	:    Union[str, Any]           =     MaskFormerModel.from_pretrained("facebook/maskformer-swin-small-coco"			).to(_lowerCAmelCase			)
       lowercase_	:    str           =     self.default_image_processor
       lowercase_	:    int           =     prepare_img()
       lowercase_	:    Tuple           =     image_processor(_lowerCAmelCase  ,     return_tensors="pt"			).to(_lowerCAmelCase			)
       lowercase_	:    Union[str, Any]           =     inputs["pixel_values"].shape
       # check size is divisible by 32
       self.assertTrue((inputs_shape[-1] % 3_2) == 0 and (inputs_shape[-2] % 3_2) == 0			)
       # check size
       self.assertEqual(_lowerCAmelCase  ,     (1, 3, 8_0_0, 1_0_8_8)			)

       with torch.no_grad():
          lowercase_	:    int           =     model(**_lowerCAmelCase			)

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

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

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




    def 				lowerCAmelCase__				(			self  :       int			):






       '''simple docstring'''






       lowercase_	:    List[Any]           =     (
           MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-small-coco"			)
           .to(_lowerCAmelCase			)
           .eval()
       )
       lowercase_	:    Union[str, Any]           =     self.default_image_processor
       lowercase_	:    str           =     prepare_img()
       lowercase_	:    List[Any]           =     image_processor(_lowerCAmelCase  ,     return_tensors="pt"			).to(_lowerCAmelCase			)
       lowercase_	:    List[str]           =     inputs["pixel_values"].shape
       # check size is divisible by 32
       self.assertTrue((inputs_shape[-1] % 3_2) == 0 and (inputs_shape[-2] % 3_2) == 0			)
       # check size
       self.assertEqual(_lowerCAmelCase  ,     (1, 3, 8_0_0, 1_0_8_8)			)

       with torch.no_grad():
          lowercase_	:    str           =     model(**_lowerCAmelCase			)
       # masks_queries_logits
       lowercase_	:    int           =     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)  ,     )
       lowercase_	:    List[str]           =     [
           [-1.373_7124, -1.772_4937, -1.936_4233],
           [-1.597_7281, -1.986_7939, -2.152_3695],
           [-1.579_5398, -1.926_9832, -2.09_3942],
       ]
       lowercase_	:    List[str]           =     torch.tensor(_lowerCAmelCase			).to(_lowerCAmelCase			)
       self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]  ,     _lowerCAmelCase  ,     atol=_lowerCAmelCase			)			)
       # class_queries_logits
       lowercase_	:    Dict           =     outputs.class_queries_logits
       self.assertEqual(
           class_queries_logits.shape  ,     (1, model.config.decoder_config.num_queries, model.config.num_labels + 1)			)
       lowercase_	:    List[Any]           =     torch.tensor(
           [
               [1.6512e00, -5.2572e00, -3.3519e00],
               [3.6169e-02, -5.9025e00, -2.9313e00],
               [1.0766e-04, -7.7630e00, -5.1263e00],
           ]			).to(_lowerCAmelCase			)
       self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]  ,     _lowerCAmelCase  ,     atol=_lowerCAmelCase			)			)




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






       '''simple docstring'''






       lowercase_	:    Optional[Any]           =     (
           MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-resnet101-coco-stuff"			)
           .to(_lowerCAmelCase			)
           .eval()
       )
       lowercase_	:    Optional[int]           =     self.default_image_processor
       lowercase_	:    List[str]           =     prepare_img()
       lowercase_	:    Tuple           =     image_processor(_lowerCAmelCase  ,     return_tensors="pt"			).to(_lowerCAmelCase			)
       lowercase_	:    List[str]           =     inputs["pixel_values"].shape
       # check size is divisible by 32
       self.assertTrue((inputs_shape[-1] % 3_2) == 0 and (inputs_shape[-2] % 3_2) == 0			)
       # check size
       self.assertEqual(_lowerCAmelCase  ,     (1, 3, 8_0_0, 1_0_8_8)			)

       with torch.no_grad():
          lowercase_	:    Any           =     model(**_lowerCAmelCase			)
       # masks_queries_logits
       lowercase_	:    Optional[Any]           =     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)  ,     )
       lowercase_	:    str           =     [[-0.9046, -2.6366, -4.6062], [-3.4179, -5.7890, -8.8057], [-4.9179, -7.6560, -10.7711]]
       lowercase_	:    Union[str, Any]           =     torch.tensor(_lowerCAmelCase			).to(_lowerCAmelCase			)
       self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]  ,     _lowerCAmelCase  ,     atol=_lowerCAmelCase			)			)
       # class_queries_logits
       lowercase_	:    Optional[Any]           =     outputs.class_queries_logits
       self.assertEqual(
           class_queries_logits.shape  ,     (1, model.config.decoder_config.num_queries, model.config.num_labels + 1)			)
       lowercase_	:    Optional[int]           =     torch.tensor(
           [[4.7188, -3.2585, -2.8857], [6.6871, -2.9181, -1.2487], [7.2449, -2.2764, -2.1874]]			).to(_lowerCAmelCase			)
       self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]  ,     _lowerCAmelCase  ,     atol=_lowerCAmelCase			)			)





    def 				lowerCAmelCase__				(			self  :       List[Any]			):






       '''simple docstring'''






       lowercase_	:    Optional[Any]           =     (
           MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-small-coco"			)
           .to(_lowerCAmelCase			)
           .eval()
       )
       lowercase_	:    int           =     self.default_image_processor

       lowercase_	:    int           =     image_processor(
           [np.zeros((3, 8_0_0, 1_3_3_3)			), np.zeros((3, 8_0_0, 1_3_3_3)			)]  ,     segmentation_maps=[np.zeros((3_8_4, 3_8_4)			).astype(np.floataa			), np.zeros((3_8_4, 3_8_4)			).astype(np.floataa			)]  ,     return_tensors="pt"  ,     )

       lowercase_	:    Optional[int]           =     inputs["pixel_values"].to(_lowerCAmelCase			)
       lowercase_	:    Optional[int]           =     [el.to(_lowerCAmelCase			) for el in inputs["mask_labels"]]
       lowercase_	:    Tuple           =     [el.to(_lowerCAmelCase			) for el in inputs["class_labels"]]

       with torch.no_grad():
          lowercase_	:    str           =     model(**_lowerCAmelCase			)

       self.assertTrue(outputs.loss is not None			)












'''simple docstring'''
import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class        _UpperCAmelCase						(		snake_case				,					unittest.TestCase			):
    __lowerCamelCase:   Dict		       =    KandinskyVaaPriorPipeline
    __lowerCamelCase:   Optional[int]		       =    ['prompt']
    __lowerCamelCase:   Any		       =    ['prompt', 'negative_prompt']
    __lowerCamelCase:   List[Any]		       =    [
        'num_images_per_prompt',
        'generator',
        'num_inference_steps',
        'latents',
        'negative_prompt',
        'guidance_scale',
        'output_type',
        'return_dict',
    ]
    __lowerCamelCase:   List[Any]		       =    False




    @property
    def 				lowerCAmelCase__				(			self  :       Optional[Any]			):






       '''simple docstring'''






       return 3_2




    @property
    def 				lowerCAmelCase__				(			self  :       Any			):






       '''simple docstring'''






       return 3_2




    @property
    def 				lowerCAmelCase__				(			self  :       Any			):






       '''simple docstring'''






       return self.time_input_dim




    @property
    def 				lowerCAmelCase__				(			self  :       str			):






       '''simple docstring'''






       return self.time_input_dim * 4




    @property
    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       return 1_0_0




    @property
    def 				lowerCAmelCase__				(			self  :       List[Any]			):






       '''simple docstring'''






       lowercase_	:    Any           =     CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			)
       return tokenizer




    @property
    def 				lowerCAmelCase__				(			self  :       Optional[Any]			):






       '''simple docstring'''






       torch.manual_seed(0			)
       lowercase_	:    Tuple           =     CLIPTextConfig(
           bos_token_id=0  ,     eos_token_id=2  ,     hidden_size=self.text_embedder_hidden_size  ,     projection_dim=self.text_embedder_hidden_size  ,     intermediate_size=3_7  ,     layer_norm_eps=1e-05  ,     num_attention_heads=4  ,     num_hidden_layers=5  ,     pad_token_id=1  ,     vocab_size=1_0_0_0  ,     )
       return CLIPTextModelWithProjection(a			)




    @property
    def 				lowerCAmelCase__				(			self  :       Tuple			):






       '''simple docstring'''






       torch.manual_seed(0			)

       lowercase_	:    List[str]           =     {
           "num_attention_heads": 2,
           "attention_head_dim": 1_2,
           "embedding_dim": self.text_embedder_hidden_size,
           "num_layers": 1,
       }

       lowercase_	:    Union[str, Any]           =     PriorTransformer(**a			)
       # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
       lowercase_	:    List[Any]           =     nn.Parameter(torch.ones(model.clip_std.shape			)			)
       return model




    @property
    def 				lowerCAmelCase__				(			self  :       Tuple			):






       '''simple docstring'''






       torch.manual_seed(0			)
       lowercase_	:    Dict           =     CLIPVisionConfig(
           hidden_size=self.text_embedder_hidden_size  ,     image_size=2_2_4  ,     projection_dim=self.text_embedder_hidden_size  ,     intermediate_size=3_7  ,     num_attention_heads=4  ,     num_channels=3  ,     num_hidden_layers=5  ,     patch_size=1_4  ,     )

       lowercase_	:    Optional[Any]           =     CLIPVisionModelWithProjection(a			)
       return model




    @property
    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       lowercase_	:    List[str]           =     CLIPImageProcessor(
           crop_size=2_2_4  ,     do_center_crop=a  ,     do_normalize=a  ,     do_resize=a  ,     image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]  ,     image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]  ,     resample=3  ,     size=2_2_4  ,     )

       return image_processor




    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       lowercase_	:    Any           =     self.dummy_prior
       lowercase_	:    Optional[Any]           =     self.dummy_image_encoder
       lowercase_	:    List[Any]           =     self.dummy_text_encoder
       lowercase_	:    Any           =     self.dummy_tokenizer
       lowercase_	:    Optional[Any]           =     self.dummy_image_processor

       lowercase_	:    List[str]           =     UnCLIPScheduler(
           variance_type="fixed_small_log"  ,     prediction_type="sample"  ,     num_train_timesteps=1_0_0_0  ,     clip_sample=a  ,     clip_sample_range=10.0  ,     )

       lowercase_	:    List[Any]           =     {
           "prior": prior,
           "image_encoder": image_encoder,
           "text_encoder": text_encoder,
           "tokenizer": tokenizer,
           "scheduler": scheduler,
           "image_processor": image_processor,
       }

       return components




    def 				lowerCAmelCase__				(			self  :       Any  ,     a  :       Dict  ,     a  :       Dict=0			):






       '''simple docstring'''






       if str(a			).startswith("mps"			):
          lowercase_	:    int           =     torch.manual_seed(a			)
       else:
          lowercase_	:    Optional[Any]           =     torch.Generator(device=a			).manual_seed(a			)
       lowercase_	:    Any           =     {
           "prompt": "horse",
           "generator": generator,
           "guidance_scale": 4.0,
           "num_inference_steps": 2,
           "output_type": "np",
       }
       return inputs




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






       '''simple docstring'''






       lowercase_	:    str           =     "cpu"

       lowercase_	:    Any           =     self.get_dummy_components()

       lowercase_	:    int           =     self.pipeline_class(**a			)
       lowercase_	:    Any           =     pipe.to(a			)

       pipe.set_progress_bar_config(disable=a			)

       lowercase_	:    Any           =     pipe(**self.get_dummy_inputs(a			)			)
       lowercase_	:    List[Any]           =     output.image_embeds

       lowercase_	:    str           =     pipe(
           **self.get_dummy_inputs(a			)  ,     return_dict=a  ,     )[0]

       lowercase_	:    Any           =     image[0, -1_0:]
       lowercase_	:    Dict           =     image_from_tuple[0, -1_0:]

       assert image.shape == (1, 3_2)

       lowercase_	:    int           =     np.array(
           [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156]			)

       assert np.abs(image_slice.flatten() - expected_slice			).max() < 1e-2
       assert np.abs(image_from_tuple_slice.flatten() - expected_slice			).max() < 1e-2




    @skip_mps
    def 				lowerCAmelCase__				(			self  :       Optional[Any]			):






       '''simple docstring'''






       lowercase_	:    int           =     torch_device == "cpu"
       lowercase_	:    Tuple           =     True
       lowercase_	:    str           =     False

       self._test_inference_batch_single_identical(
           test_max_difference=a  ,     relax_max_difference=a  ,     test_mean_pixel_difference=a  ,     )





    @skip_mps
    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       lowercase_	:    Any           =     torch_device == "cpu"
       lowercase_	:    int           =     False

       self._test_attention_slicing_forward_pass(
           test_max_difference=a  ,     test_mean_pixel_difference=a  ,     )




