'''simple docstring'''





import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, 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 MaskaFormerForUniversalSegmentation, MaskaFormerModel

 if is_vision_available():
  from transformers import MaskaFormerImageProcessor

if is_vision_available():
 from PIL import Image






class 				UpperCAmelCase :




    '''simple docstring'''







    def __init__(       self    ,							lowercase__    ,							lowercase__=2    ,							lowercase__=True    ,							lowercase__=False    ,							lowercase__=10    ,							lowercase__=3    ,							lowercase__=32 * 8    ,							lowercase__=32 * 8    ,							lowercase__=4    ,							lowercase__=64    ,							)							-> Optional[int]:
         SCREAMING_SNAKE_CASE					:			Any    	=       parent
         SCREAMING_SNAKE_CASE					:			List[str]    	=       batch_size
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       is_training
         SCREAMING_SNAKE_CASE					:			int    	=       use_auxiliary_loss
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       num_queries
         SCREAMING_SNAKE_CASE					:			List[str]    	=       num_channels
         SCREAMING_SNAKE_CASE					:			str    	=       min_size
         SCREAMING_SNAKE_CASE					:			List[str]    	=       max_size
         SCREAMING_SNAKE_CASE					:			str    	=       num_labels
         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       hidden_dim
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       hidden_dim






    def       _UpperCamelCase    (       self )							-> int:
         SCREAMING_SNAKE_CASE					:			int    	=       floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size] ).to(
             lowercase__ )

         SCREAMING_SNAKE_CASE					:			Tuple    	=       torch.ones([self.batch_size, self.min_size, self.max_size]    ,							device=lowercase__ )

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

         SCREAMING_SNAKE_CASE					:			Dict    	=       self.get_config()
         return config, pixel_values, pixel_mask, mask_labels, class_labels






    def       _UpperCamelCase    (       self )							-> Dict:
         SCREAMING_SNAKE_CASE					:			Tuple    	=       MaskaFormerConfig(
             hidden_size=self.hidden_dim    ,							)
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       self.num_queries
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       self.num_labels

         SCREAMING_SNAKE_CASE					:			int    	=       [1, 1, 1, 1]
         SCREAMING_SNAKE_CASE					:			str    	=       self.num_channels

         SCREAMING_SNAKE_CASE					:			int    	=       64
         SCREAMING_SNAKE_CASE					:			Tuple    	=       128
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       self.hidden_dim
         SCREAMING_SNAKE_CASE					:			List[str]    	=       self.hidden_dim
         SCREAMING_SNAKE_CASE					:			Any    	=       self.hidden_dim
         return config






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






    def       _UpperCamelCase    (       self    ,							lowercase__    ,							lowercase__ )							-> Any:
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       output.encoder_hidden_states
         SCREAMING_SNAKE_CASE					:			str    	=       output.pixel_decoder_hidden_states
         SCREAMING_SNAKE_CASE					:			List[str]    	=       output.transformer_decoder_hidden_states

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






    def       _UpperCamelCase    (       self    ,							lowercase__    ,							lowercase__    ,							lowercase__    ,							lowercase__=False )							-> Optional[Any]:
         with torch.no_grad():
              SCREAMING_SNAKE_CASE					:			Any    	=       MaskaFormerModel(config=lowercase__ )
              model.to(lowercase__ )
              model.eval()

              SCREAMING_SNAKE_CASE					:			List[str]    	=       model(pixel_values=lowercase__    ,							pixel_mask=lowercase__ )
              SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       model(lowercase__    ,							output_hidden_states=lowercase__ )

         self.parent.assertEqual(
             output.transformer_decoder_last_hidden_state.shape    ,							(self.batch_size, self.num_queries, self.hidden_dim)    ,							)
         # 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(lowercase__    ,							lowercase__ )






    def       _UpperCamelCase    (       self    ,							lowercase__    ,							lowercase__    ,							lowercase__    ,							lowercase__    ,							lowercase__ )							-> Any:
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       MaskaFormerForUniversalSegmentation(config=lowercase__ )
         model.to(lowercase__ )
         model.eval()

         def comm_check_on_output(lowercase__ ):
              # 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():
              SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       model(pixel_values=lowercase__    ,							pixel_mask=lowercase__ )
              SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       model(lowercase__ )

              comm_check_on_output(lowercase__ )

              SCREAMING_SNAKE_CASE					:			Any    	=       model(
                  pixel_values=lowercase__    ,							pixel_mask=lowercase__    ,							mask_labels=lowercase__    ,							class_labels=lowercase__ )

         comm_check_on_output(lowercase__ )

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






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




    '''simple docstring'''







    snake_case__							:   Optional[Any]							=						(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
    snake_case__							:   Optional[int]							=						{"feature-extraction": MaskaFormerModel} if is_torch_available() else {}

    snake_case__							:   Optional[int]							=						False
    snake_case__							:   int							=						False
    snake_case__							:   Optional[int]							=						False
    snake_case__							:   Tuple							=						False






    def       _UpperCamelCase    (       self )							-> Optional[int]:
         SCREAMING_SNAKE_CASE					:			Tuple    	=       MaskaFormerModelTester(self )
         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       ConfigTester(self    ,							config_class=lowercase__    ,							has_text_modality=lowercase__ )






    def       _UpperCamelCase    (       self )							-> Tuple:
         self.config_tester.run_common_tests()






    def       _UpperCamelCase    (       self )							-> List[Any]:
         SCREAMING_SNAKE_CASE	,   SCREAMING_SNAKE_CASE					:			Dict    	=       self.model_tester.prepare_config_and_inputs_for_common()
         self.model_tester.create_and_check_maskaformer_model(lowercase__    ,							**lowercase__    ,							output_hidden_states=lowercase__ )






    def       _UpperCamelCase    (       self )							-> int:
         SCREAMING_SNAKE_CASE					:			List[str]    	=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*lowercase__ )






    @unittest.skip(reason='Mask2Former does not use inputs_embeds' )
    def       _UpperCamelCase    (       self )							-> List[Any]:
         pass






    @unittest.skip(reason='Mask2Former does not have a get_input_embeddings method' )
    def       _UpperCamelCase    (       self )							-> str:
         pass






    @unittest.skip(reason='Mask2Former is not a generative model' )
    def       _UpperCamelCase    (       self )							-> Any:
         pass






    @unittest.skip(reason='Mask2Former does not use token embeddings' )
    def       _UpperCamelCase    (       self )							-> str:
         pass






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






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






    def       _UpperCamelCase    (       self )							-> Tuple:
         SCREAMING_SNAKE_CASE	,   SCREAMING_SNAKE_CASE					:			str    	=       self.model_tester.prepare_config_and_inputs_for_common()

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

              SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       ['pixel_values']
              self.assertListEqual(arg_names[:1]    ,							lowercase__ )






    @slow
    def       _UpperCamelCase    (       self )							-> Any:
         for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
              SCREAMING_SNAKE_CASE					:			List[str]    	=       MaskaFormerModel.from_pretrained(lowercase__ )
              self.assertIsNotNone(lowercase__ )






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

         SCREAMING_SNAKE_CASE					:			Tuple    	=       MaskaFormerForUniversalSegmentation(lowercase__ ).to(lowercase__ )
         SCREAMING_SNAKE_CASE					:			Dict    	=       model(**lowercase__ )
         self.assertTrue(outputs.loss is not None )






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






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

         for model_class in self.all_model_classes:
              SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       model_class(lowercase__ ).to(lowercase__ )
              SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       model(**lowercase__    ,							output_attentions=lowercase__ )
              self.assertTrue(outputs.attentions is not None )






    def       _UpperCamelCase    (       self )							-> Any:
         if not self.model_tester.is_training:
              return

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

         SCREAMING_SNAKE_CASE					:			str    	=       model_class(lowercase__ )
         model.to(lowercase__ )
         model.train()

         SCREAMING_SNAKE_CASE					:			str    	=       model(lowercase__    ,							mask_labels=lowercase__    ,							class_labels=lowercase__ ).loss
         loss.backward()






    def       _UpperCamelCase    (       self )							-> Optional[Any]:
         SCREAMING_SNAKE_CASE					:			Dict    	=       self.all_model_classes[1]
         SCREAMING_SNAKE_CASE	,   SCREAMING_SNAKE_CASE	,   SCREAMING_SNAKE_CASE	,   SCREAMING_SNAKE_CASE	,   SCREAMING_SNAKE_CASE					:			Tuple    	=       self.model_tester.prepare_config_and_inputs()
         SCREAMING_SNAKE_CASE					:			Dict    	=       True
         SCREAMING_SNAKE_CASE					:			List[str]    	=       True

         SCREAMING_SNAKE_CASE					:			List[str]    	=       model_class(lowercase__ ).to(lowercase__ )
         model.train()

         SCREAMING_SNAKE_CASE					:			int    	=       model(lowercase__    ,							mask_labels=lowercase__    ,							class_labels=lowercase__ )

         SCREAMING_SNAKE_CASE					:			List[str]    	=       outputs.encoder_hidden_states[0]
         encoder_hidden_states.retain_grad()

         SCREAMING_SNAKE_CASE					:			List[str]    	=       outputs.pixel_decoder_hidden_states[0]
         pixel_decoder_hidden_states.retain_grad()

         SCREAMING_SNAKE_CASE					:			str    	=       outputs.transformer_decoder_hidden_states[0]
         transformer_decoder_hidden_states.retain_grad()

         SCREAMING_SNAKE_CASE					:			Tuple    	=       outputs.attentions[0]
         attentions.retain_grad()

         outputs.loss.backward(retain_graph=lowercase__ )

         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  :List[Any]  				=    1E-4


def     __lowerCAmelCase						(       )      ->	Dict:





     '''simple docstring'''





     SCREAMING_SNAKE_CASE					:			Tuple    	=       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 )							-> Tuple:
         return "facebook/mask2former-swin-small-coco-instance"






    @cached_property
    def       _UpperCamelCase    (       self )							-> int:
         return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints ) if is_vision_available() else None






    def       _UpperCamelCase    (       self )							-> Optional[int]:
         SCREAMING_SNAKE_CASE					:			int    	=       MaskaFormerModel.from_pretrained(self.model_checkpoints ).to(lowercase__ )
         SCREAMING_SNAKE_CASE					:			List[str]    	=       self.default_image_processor
         SCREAMING_SNAKE_CASE					:			Dict    	=       prepare_img()
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       image_processor(lowercase__    ,							return_tensors='pt' ).to(lowercase__ )
         SCREAMING_SNAKE_CASE					:			int    	=       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(lowercase__    ,							(1, 3, 384, 384) )

         with torch.no_grad():
              SCREAMING_SNAKE_CASE					:			Any    	=       model(**lowercase__ )

         SCREAMING_SNAKE_CASE					:			Any    	=       torch.tensor(
             [[-0.2_7_9_0, -1.0_7_1_7, -1.1_6_6_8], [-0.5_1_2_8, -0.3_1_2_8, -0.4_9_8_7], [-0.5_8_3_2, 0.1_9_7_1, -0.0_1_9_7]] ).to(lowercase__ )
         self.assertTrue(
             torch.allclose(
                 outputs.encoder_last_hidden_state[0, 0, :3, :3]    ,							lowercase__    ,							atol=lowercase__ ) )

         SCREAMING_SNAKE_CASE					:			Tuple    	=       torch.tensor(
             [[0.8_9_7_3, 1.1_8_4_7, 1.1_7_7_6], [1.1_9_3_4, 1.5_0_4_0, 1.5_1_2_8], [1.1_1_5_3, 1.4_4_8_6, 1.4_9_5_1]] ).to(lowercase__ )
         self.assertTrue(
             torch.allclose(
                 outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]    ,							lowercase__    ,							atol=lowercase__ ) )

         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       torch.tensor(
             [[2.1_1_5_2, 1.7_0_0_0, -0.8_6_0_3], [1.5_8_0_8, 1.8_0_0_4, -0.9_3_5_3], [1.6_0_4_3, 1.7_4_9_5, -0.5_9_9_9]] ).to(lowercase__ )
         self.assertTrue(
             torch.allclose(
                 outputs.transformer_decoder_last_hidden_state[0, :3, :3]    ,							lowercase__    ,							atol=lowercase__ ) )






    def       _UpperCamelCase    (       self )							-> Any:
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints ).to(lowercase__ ).eval()
         SCREAMING_SNAKE_CASE					:			Any    	=       self.default_image_processor
         SCREAMING_SNAKE_CASE					:			List[str]    	=       prepare_img()
         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       image_processor(lowercase__    ,							return_tensors='pt' ).to(lowercase__ )
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       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(lowercase__    ,							(1, 3, 384, 384) )

         with torch.no_grad():
              SCREAMING_SNAKE_CASE					:			Tuple    	=       model(**lowercase__ )
         # masks_queries_logits
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       outputs.masks_queries_logits
         self.assertEqual(
             masks_queries_logits.shape    ,							(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4) )
         SCREAMING_SNAKE_CASE					:			Tuple    	=       [
             [-8.7_8_3_9, -9.0_0_5_6, -8.8_1_2_1],
             [-7.4_1_0_4, -7.0_3_1_3, -6.5_4_0_1],
             [-6.6_1_0_5, -6.3_4_2_7, -6.4_6_7_5],
         ]
         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       torch.tensor(lowercase__ ).to(lowercase__ )
         self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]    ,							lowercase__    ,							atol=lowercase__ ) )
         # class_queries_logits
         SCREAMING_SNAKE_CASE					:			Tuple    	=       outputs.class_queries_logits
         self.assertEqual(class_queries_logits.shape    ,							(1, model.config.num_queries, model.config.num_labels + 1) )
         SCREAMING_SNAKE_CASE					:			int    	=       torch.tensor(
             [
                 [1.8_3_2_4, -8.0_8_3_5, -4.1_9_2_2],
                 [0.8_4_5_0, -9.0_0_5_0, -3.6_0_5_3],
                 [0.3_0_4_5, -7.7_2_9_3, -3.0_2_7_5],
             ] ).to(lowercase__ )
         self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]    ,							lowercase__    ,							atol=lowercase__ ) )






    def       _UpperCamelCase    (       self )							-> Optional[int]:
         SCREAMING_SNAKE_CASE					:			Tuple    	=       MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints ).to(lowercase__ ).eval()
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       self.default_image_processor

         SCREAMING_SNAKE_CASE					:			Any    	=       image_processor(
             [np.zeros((3, 800, 1_333) ), np.zeros((3, 800, 1_333) )]    ,							segmentation_maps=[np.zeros((384, 384) ).astype(np.floataa ), np.zeros((384, 384) ).astype(np.floataa )]    ,							return_tensors='pt'    ,							)

         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       inputs['pixel_values'].to(lowercase__ )
         SCREAMING_SNAKE_CASE					:			Tuple    	=       [el.to(lowercase__ ) for el in inputs['mask_labels']]
         SCREAMING_SNAKE_CASE					:			Dict    	=       [el.to(lowercase__ ) for el in inputs['class_labels']]

         with torch.no_grad():
              SCREAMING_SNAKE_CASE					:			Any    	=       model(**lowercase__ )

         self.assertTrue(outputs.loss is not None )




'''simple docstring'''





import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionTextToImagePipeline
from diffusers.utils.testing_utils import nightly, require_torch_gpu, torch_device


_lowerCAmelCase  :List[Any]  				=    False






class 				UpperCAmelCase (       unittest.TestCase ):




    '''simple docstring'''







    pass






@nightly
@require_torch_gpu
class 				UpperCAmelCase (       unittest.TestCase ):




    '''simple docstring'''







    def       _UpperCamelCase    (       self )							-> Any:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






    def       _UpperCamelCase    (       self )							-> Tuple:
         SCREAMING_SNAKE_CASE					:			Dict    	=       VersatileDiffusionTextToImagePipeline.from_pretrained('shi-labs/versatile-diffusion' )
         # remove text_unet
         pipe.remove_unused_weights()
         pipe.to(lowercase__ )
         pipe.set_progress_bar_config(disable=lowercase__ )

         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       'A painting of a squirrel eating a burger '
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       torch.manual_seed(0 )
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       pipe(
             prompt=lowercase__    ,							generator=lowercase__    ,							guidance_scale=7.5    ,							num_inference_steps=2    ,							output_type='numpy' ).images

         with tempfile.TemporaryDirectory() as tmpdirname:
              pipe.save_pretrained(lowercase__ )
              SCREAMING_SNAKE_CASE					:			int    	=       VersatileDiffusionTextToImagePipeline.from_pretrained(lowercase__ )
         pipe.to(lowercase__ )
         pipe.set_progress_bar_config(disable=lowercase__ )

         SCREAMING_SNAKE_CASE					:			List[str]    	=       generator.manual_seed(0 )
         SCREAMING_SNAKE_CASE					:			List[str]    	=       pipe(
             prompt=lowercase__    ,							generator=lowercase__    ,							guidance_scale=7.5    ,							num_inference_steps=2    ,							output_type='numpy' ).images

         assert np.abs(image - new_image ).sum() < 1E-5, "Models don't have the same forward pass"






    def       _UpperCamelCase    (       self )							-> Any:
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       VersatileDiffusionTextToImagePipeline.from_pretrained(
             'shi-labs/versatile-diffusion'    ,							torch_dtype=torch.floataa )
         pipe.to(lowercase__ )
         pipe.set_progress_bar_config(disable=lowercase__ )

         SCREAMING_SNAKE_CASE					:			int    	=       'A painting of a squirrel eating a burger '
         SCREAMING_SNAKE_CASE					:			Tuple    	=       torch.manual_seed(0 )
         SCREAMING_SNAKE_CASE					:			Tuple    	=       pipe(
             prompt=lowercase__    ,							generator=lowercase__    ,							guidance_scale=7.5    ,							num_inference_steps=50    ,							output_type='numpy' ).images

         SCREAMING_SNAKE_CASE					:			Dict    	=       image[0, 253:256, 253:256, -1]

         assert image.shape == (1, 512, 512, 3)
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       np.array([0.3_3_6_7, 0.3_1_6_9, 0.2_6_5_6, 0.3_8_7_0, 0.4_7_9_0, 0.3_7_9_6, 0.4_0_0_9, 0.4_8_7_8, 0.4_7_7_8] )

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



