import json
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
 import torch

 if is_vision_available():
  from transformers import OneFormerImageProcessor
  from transformers.models.oneformer.image_processing_oneformer import binary_mask_to_rle
  from transformers.models.oneformer.modeling_oneformer import OneFormerForUniversalSegmentationOutput

if is_vision_available():
 from PIL import Image






def 							UpperCAmelCase__(					__UpperCAmelCase			:     List[Any]     ,				__UpperCAmelCase			:     str="shi-labs/oneformer_demo"       ):
   with open(hf_hub_download(__UpperCAmelCase     ,				__UpperCAmelCase     ,				repo_type='dataset'       )     ,				'r'       ) as f:
      __snake_case   :			Optional[Any]										=						json.load(__UpperCAmelCase       )
   __snake_case   :			Any										=						{}
   __snake_case   :			Optional[Any]										=						[]
   __snake_case   :			int										=						[]
   for key, info in class_info.items():
      __snake_case   :			Tuple										=						info['''name''']
      class_names.append(info['name']       )
      if info["isthing"]:
         thing_ids.append(int(__UpperCAmelCase       )       )
   __snake_case   :			int										=						thing_ids
   __snake_case   :			List[Any]										=						class_names
   return metadata
class 					__SCREAMING_SNAKE_CASE       (       unittest.TestCase):






       """simple docstring"""

       def __init__( self			,		_UpperCAmelCase			,		_UpperCAmelCase=7			,		_UpperCAmelCase=3			,		_UpperCAmelCase=30			,		_UpperCAmelCase=400			,		_UpperCAmelCase=None			,		_UpperCAmelCase=True			,		_UpperCAmelCase=True			,		_UpperCAmelCase=[0.5, 0.5, 0.5]			,		_UpperCAmelCase=[0.5, 0.5, 0.5]			,		_UpperCAmelCase=10			,		_UpperCAmelCase=False			,		_UpperCAmelCase=255			,		_UpperCAmelCase="shi-labs/oneformer_demo"			,		_UpperCAmelCase="ade20k_panoptic.json"			,		_UpperCAmelCase=10			,		):
          __snake_case   :			Dict										=						parent
          __snake_case   :			Any										=						batch_size
          __snake_case   :			List[str]										=						num_channels
          __snake_case   :			Union[str, Any]										=						min_resolution
          __snake_case   :			str										=						max_resolution
          __snake_case   :			Tuple										=						do_resize
          __snake_case   :			Union[str, Any]										=						{'''shortest_edge''': 32, '''longest_edge''': 1_333} if size is None else size
          __snake_case   :			str										=						do_normalize
          __snake_case   :			List[str]										=						image_mean
          __snake_case   :			Dict										=						image_std
          __snake_case   :			Optional[int]										=						class_info_file
          __snake_case   :			Optional[int]										=						prepare_metadata(_lowercase			,		_lowercase				)
          __snake_case   :			str										=						num_text
          __snake_case   :			Optional[Any]										=						repo_path

          # for the post_process_functions
          __snake_case   :			Optional[Any]										=						2
          __snake_case   :			Optional[Any]										=						10
          __snake_case   :			Tuple										=						10
          __snake_case   :			str										=						3
          __snake_case   :			str										=						4
          __snake_case   :			Optional[int]										=						num_labels
          __snake_case   :			Union[str, Any]										=						do_reduce_labels
          __snake_case   :			Optional[int]										=						ignore_index







       def      lowercase_   ( self				):
          return {
              "do_resize": self.do_resize,
              "size": self.size,
              "do_normalize": self.do_normalize,
              "image_mean": self.image_mean,
              "image_std": self.image_std,
              "num_labels": self.num_labels,
              "do_reduce_labels": self.do_reduce_labels,
              "ignore_index": self.ignore_index,
              "class_info_file": self.class_info_file,
              "metadata": self.metadata,
              "num_text": self.num_text,
          }







       def      lowercase_   ( self			,		_UpperCAmelCase			,		_UpperCAmelCase=False				):



          if not batched:
             __snake_case   :			Optional[int]										=						image_inputs[0]
             if isinstance(_lowercase			,		Image.Image				):
                __snake_case   :			Optional[int]										=						image.size
             else:
                __snake_case   :			Union[str, Any]										=						image.shape[1], image.shape[2]
             if w < h:
                __snake_case   :			Tuple										=						int(self.size['shortest_edge'] * h / w				)
                __snake_case   :			Optional[Any]										=						self.size['''shortest_edge''']
             elif w > h:
                __snake_case   :			List[Any]										=						self.size['''shortest_edge''']
                __snake_case   :			Any										=						int(self.size['shortest_edge'] * w / h				)
             else:
                __snake_case   :			Optional[int]										=						self.size['''shortest_edge''']
                __snake_case   :			Tuple										=						self.size['''shortest_edge''']

          else:
             __snake_case   :			Dict										=						[]
             for image in image_inputs:
                __snake_case   :			List[Any]										=						self.get_expected_values([image]				)
                expected_values.append((expected_height, expected_width)				)
             __snake_case   :			Any										=						max(_lowercase			,		key=lambda _UpperCAmelCase				: item[0]				)[0]
             __snake_case   :			Union[str, Any]										=						max(_lowercase			,		key=lambda _UpperCAmelCase				: item[1]				)[1]

          return expected_height, expected_width







       def      lowercase_   ( self				):
          return OneFormerForUniversalSegmentationOutput(
              # +1 for null class
              class_queries_logits=torch.randn((self.batch_size, self.num_queries, self.num_classes + 1)				)			,		masks_queries_logits=torch.randn((self.batch_size, self.num_queries, self.height, self.width)				)			,		)





@require_torch
@require_vision
class 					__SCREAMING_SNAKE_CASE       (       __a						,	unittest.TestCase):






       """simple docstring"""

       __UpperCAmelCase       =					OneFormerImageProcessor if (is_vision_available() and is_torch_available()) else None
       # only for test_image_processing_common.test_image_proc_to_json_string
       __UpperCAmelCase       =					image_processing_class







       def      lowercase_   ( self				):
          __snake_case   :			int										=						OneFormerImageProcessorTester(self				)







       @property
       def      lowercase_   ( self				):
          return self.image_processing_tester.prepare_image_processor_dict()







       def      lowercase_   ( self				):
          __snake_case   :			List[str]										=						self.image_processing_class(**self.image_processor_dict				)
          self.assertTrue(hasattr(_lowercase			,		'image_mean'				)				)
          self.assertTrue(hasattr(_lowercase			,		'image_std'				)				)
          self.assertTrue(hasattr(_lowercase			,		'do_normalize'				)				)
          self.assertTrue(hasattr(_lowercase			,		'do_resize'				)				)
          self.assertTrue(hasattr(_lowercase			,		'size'				)				)
          self.assertTrue(hasattr(_lowercase			,		'ignore_index'				)				)
          self.assertTrue(hasattr(_lowercase			,		'class_info_file'				)				)
          self.assertTrue(hasattr(_lowercase			,		'num_text'				)				)
          self.assertTrue(hasattr(_lowercase			,		'repo_path'				)				)
          self.assertTrue(hasattr(_lowercase			,		'metadata'				)				)
          self.assertTrue(hasattr(_lowercase			,		'do_reduce_labels'				)				)







       def      lowercase_   ( self				):
          pass







       def      lowercase_   ( self				):
          __snake_case   :			Optional[int]										=						self.image_processing_class(**self.image_processor_dict				)
          # create random PIL images
          __snake_case   :			int										=						prepare_image_inputs(self.image_processing_tester			,		equal_resolution=_lowercase				)
          for image in image_inputs:
             self.assertIsInstance(_lowercase			,		Image.Image				)

          # Test not batched input
          __snake_case   :			Any										=						image_processor(image_inputs[0]			,		['semantic']			,		return_tensors='pt'				).pixel_values

          __snake_case   :			Dict										=						self.image_processing_tester.get_expected_values(_lowercase				)

          self.assertEqual(
              encoded_images.shape			,		(1, self.image_processing_tester.num_channels, expected_height, expected_width)			,		)

          # Test batched
          __snake_case   :			Union[str, Any]										=						self.image_processing_tester.get_expected_values(_lowercase			,		batched=_lowercase				)

          __snake_case   :			int										=						image_processor(
              _lowercase			,		['semantic'] * len(_lowercase				)			,		return_tensors='pt'				).pixel_values
          self.assertEqual(
              encoded_images.shape			,		(
                  self.image_processing_tester.batch_size,
                  self.image_processing_tester.num_channels,
                  expected_height,
                  expected_width,
              )			,		)







       def      lowercase_   ( self				):
          __snake_case   :			List[Any]										=						self.image_processing_class(**self.image_processor_dict				)
          # create random numpy tensors
          __snake_case   :			Tuple										=						prepare_image_inputs(self.image_processing_tester			,		equal_resolution=_lowercase			,		numpify=_lowercase				)
          for image in image_inputs:
             self.assertIsInstance(_lowercase			,		np.ndarray				)

          # Test not batched input
          __snake_case   :			Union[str, Any]										=						image_processor(image_inputs[0]			,		['semantic']			,		return_tensors='pt'				).pixel_values

          __snake_case   :			List[Any]										=						self.image_processing_tester.get_expected_values(_lowercase				)

          self.assertEqual(
              encoded_images.shape			,		(1, self.image_processing_tester.num_channels, expected_height, expected_width)			,		)

          # Test batched
          __snake_case   :			Tuple										=						self.image_processing_tester.get_expected_values(_lowercase			,		batched=_lowercase				)

          __snake_case   :			Union[str, Any]										=						image_processor(
              _lowercase			,		['semantic'] * len(_lowercase				)			,		return_tensors='pt'				).pixel_values
          self.assertEqual(
              encoded_images.shape			,		(
                  self.image_processing_tester.batch_size,
                  self.image_processing_tester.num_channels,
                  expected_height,
                  expected_width,
              )			,		)







       def      lowercase_   ( self				):
          __snake_case   :			Tuple										=						self.image_processing_class(**self.image_processor_dict				)
          # create random PyTorch tensors
          __snake_case   :			Any										=						prepare_image_inputs(self.image_processing_tester			,		equal_resolution=_lowercase			,		torchify=_lowercase				)
          for image in image_inputs:
             self.assertIsInstance(_lowercase			,		torch.Tensor				)

          # Test not batched input
          __snake_case   :			List[Any]										=						image_processor(image_inputs[0]			,		['semantic']			,		return_tensors='pt'				).pixel_values

          __snake_case   :			List[str]										=						self.image_processing_tester.get_expected_values(_lowercase				)

          self.assertEqual(
              encoded_images.shape			,		(1, self.image_processing_tester.num_channels, expected_height, expected_width)			,		)

          # Test batched
          __snake_case   :			Optional[int]										=						self.image_processing_tester.get_expected_values(_lowercase			,		batched=_lowercase				)

          __snake_case   :			Optional[Any]										=						image_processor(
              _lowercase			,		['semantic'] * len(_lowercase				)			,		return_tensors='pt'				).pixel_values
          self.assertEqual(
              encoded_images.shape			,		(
                  self.image_processing_tester.batch_size,
                  self.image_processing_tester.num_channels,
                  expected_height,
                  expected_width,
              )			,		)







       def      lowercase_   ( self			,		_UpperCAmelCase=False			,		_UpperCAmelCase=False			,		_UpperCAmelCase="np"				):
          __snake_case   :			Optional[Any]										=						self.image_processing_class(**self.image_processor_dict				)
          # prepare image and target
          __snake_case   :			Tuple										=						self.image_processing_tester.num_labels
          __snake_case   :			List[str]										=						None
          __snake_case   :			Any										=						None
          __snake_case   :			Optional[int]										=						prepare_image_inputs(self.image_processing_tester			,		equal_resolution=_lowercase				)
          if with_segmentation_maps:
             __snake_case   :			Optional[Any]										=						num_labels
             if is_instance_map:
                __snake_case   :			Any										=						list(range(_lowercase				)				) * 2
                __snake_case   :			List[str]										=						dict(enumerate(_lowercase				)				)
             __snake_case   :			Optional[int]										=						[
                 np.random.randint(0			,		high * 2			,		(img.size[1], img.size[0])				).astype(np.uinta				) for img in image_inputs
             ]
             if segmentation_type == "pil":
                __snake_case   :			Union[str, Any]										=						[Image.fromarray(_lowercase				) for annotation in annotations]

          __snake_case   :			int										=						image_processor(
              _lowercase			,		['semantic'] * len(_lowercase				)			,		_lowercase			,		return_tensors='pt'			,		instance_id_to_semantic_id=_lowercase			,		pad_and_return_pixel_mask=_lowercase			,		)

          return inputs







       def      lowercase_   ( self				):
          pass







       def      lowercase_   ( self				):



          def common(_UpperCAmelCase=False			,		_UpperCAmelCase=None				):
             __snake_case   :			int										=						self.comm_get_image_processor_inputs(
                 with_segmentation_maps=_lowercase			,		is_instance_map=_lowercase			,		segmentation_type=_lowercase				)

             __snake_case   :			Optional[int]										=						inputs['''mask_labels''']
             __snake_case   :			List[Any]										=						inputs['''class_labels''']
             __snake_case   :			Dict										=						inputs['''pixel_values''']
             __snake_case   :			Optional[int]										=						inputs['''text_inputs''']

             # check the batch_size
             for mask_label, class_label, text_input in zip(_lowercase			,		_lowercase			,		_lowercase				):
                self.assertEqual(mask_label.shape[0]			,		class_label.shape[0]				)
                # this ensure padding has happened
                self.assertEqual(mask_label.shape[1:]			,		pixel_values.shape[2:]				)
                self.assertEqual(len(_lowercase				)			,		self.image_processing_tester.num_text				)

          common()
          common(is_instance_map=_lowercase				)
          common(is_instance_map=_lowercase			,		segmentation_type='pil'				)
          common(is_instance_map=_lowercase			,		segmentation_type='pil'				)







       def      lowercase_   ( self				):
          __snake_case   :			List[str]										=						np.zeros((20, 50)				)
          __snake_case   :			List[Any]										=						1
          __snake_case   :			List[Any]										=						1
          __snake_case   :			Optional[Any]										=						1

          __snake_case   :			List[Any]										=						binary_mask_to_rle(_lowercase				)
          self.assertEqual(len(_lowercase				)			,		4				)
          self.assertEqual(rle[0]			,		21				)
          self.assertEqual(rle[1]			,		45				)







       def      lowercase_   ( self				):
          __snake_case   :			Dict										=						self.image_processing_class(
              num_labels=self.image_processing_tester.num_classes			,		max_seq_length=77			,		task_seq_length=77			,		class_info_file='ade20k_panoptic.json'			,		num_text=self.image_processing_tester.num_text			,		repo_path='shi-labs/oneformer_demo'			,		)
          __snake_case   :			List[Any]										=						self.image_processing_tester.get_fake_oneformer_outputs()

          __snake_case   :			Optional[Any]										=						fature_extractor.post_process_semantic_segmentation(_lowercase				)

          self.assertEqual(len(_lowercase				)			,		self.image_processing_tester.batch_size				)
          self.assertEqual(
              segmentation[0].shape			,		(
                  self.image_processing_tester.height,
                  self.image_processing_tester.width,
              )			,		)

          __snake_case   :			Optional[Any]										=						[(1, 4) for i in range(self.image_processing_tester.batch_size				)]
          __snake_case   :			List[str]										=						fature_extractor.post_process_semantic_segmentation(_lowercase			,		target_sizes=_lowercase				)

          self.assertEqual(segmentation[0].shape			,		target_sizes[0]				)







       def      lowercase_   ( self				):
          __snake_case   :			Any										=						self.image_processing_class(
              num_labels=self.image_processing_tester.num_classes			,		max_seq_length=77			,		task_seq_length=77			,		class_info_file='ade20k_panoptic.json'			,		num_text=self.image_processing_tester.num_text			,		repo_path='shi-labs/oneformer_demo'			,		)
          __snake_case   :			int										=						self.image_processing_tester.get_fake_oneformer_outputs()
          __snake_case   :			Tuple										=						image_processor.post_process_instance_segmentation(_lowercase			,		threshold=0				)

          self.assertTrue(len(_lowercase				) == self.image_processing_tester.batch_size				)
          for el in segmentation:
             self.assertTrue('segmentation' in el				)
             self.assertTrue('segments_info' in el				)
             self.assertEqual(type(el['segments_info']				)			,		_lowercase				)
             self.assertEqual(
                 el['segmentation'].shape			,		(self.image_processing_tester.height, self.image_processing_tester.width)				)







       def      lowercase_   ( self				):
          __snake_case   :			int										=						self.image_processing_class(
              num_labels=self.image_processing_tester.num_classes			,		max_seq_length=77			,		task_seq_length=77			,		class_info_file='ade20k_panoptic.json'			,		num_text=self.image_processing_tester.num_text			,		repo_path='shi-labs/oneformer_demo'			,		)
          __snake_case   :			Union[str, Any]										=						self.image_processing_tester.get_fake_oneformer_outputs()
          __snake_case   :			Optional[int]										=						image_processor.post_process_panoptic_segmentation(_lowercase			,		threshold=0				)

          self.assertTrue(len(_lowercase				) == self.image_processing_tester.batch_size				)
          for el in segmentation:
             self.assertTrue('segmentation' in el				)
             self.assertTrue('segments_info' in el				)
             self.assertEqual(type(el['segments_info']				)			,		_lowercase				)
             self.assertEqual(
                 el['segmentation'].shape			,		(self.image_processing_tester.height, self.image_processing_tester.width)				)




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class 			__magic_name__	(		__a   ,						__a   ,						unittest.TestCase     ):





   """simple docstring"""




   lowerCAmelCase   :     List[str]	      = StableDiffusionSAGPipeline
   lowerCAmelCase   :     List[Any]	      = TEXT_TO_IMAGE_PARAMS
   lowerCAmelCase   :     List[str]	      = TEXT_TO_IMAGE_BATCH_PARAMS
   lowerCAmelCase   :     str	      = TEXT_TO_IMAGE_IMAGE_PARAMS
   lowerCAmelCase   :     int	      = TEXT_TO_IMAGE_IMAGE_PARAMS
   lowerCAmelCase   :     int	      = False



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    torch.manual_seed(0				)
    _UpperCamelCase:  Dict    						=			UNetaDConditionModel(
        block_out_channels=(32, 64)		,	layers_per_block=2		,	sample_size=32		,	in_channels=4		,	out_channels=4		,	down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')		,	up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')		,	cross_attention_dim=32		,	)
    _UpperCamelCase:  List[Any]    						=			DDIMScheduler(
        beta_start=0.00085		,	beta_end=0.012		,	beta_schedule='''scaled_linear'''		,	clip_sample=_lowercase		,	set_alpha_to_one=_lowercase		,	)
    torch.manual_seed(0				)
    _UpperCamelCase:  Optional[int]    						=			AutoencoderKL(
        block_out_channels=[32, 64]		,	in_channels=3		,	out_channels=3		,	down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']		,	up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']		,	latent_channels=4		,	)
    torch.manual_seed(0				)
    _UpperCamelCase:  int    						=			CLIPTextConfig(
        bos_token_id=0		,	eos_token_id=2		,	hidden_size=32		,	intermediate_size=37		,	layer_norm_eps=1E-05		,	num_attention_heads=4		,	num_hidden_layers=5		,	pad_token_id=1		,	vocab_size=1_000		,	)
    _UpperCamelCase:  Any    						=			CLIPTextModel(_lowercase				)
    _UpperCamelCase:  List[str]    						=			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)

    _UpperCamelCase:  int    						=			{
        '''unet''': unet,
        '''scheduler''': scheduler,
        '''vae''': vae,
        '''text_encoder''': text_encoder,
        '''tokenizer''': tokenizer,
        '''safety_checker''': None,
        '''feature_extractor''': None,
    }
    return components



   def 	lowerCAmelCase    (						self       :	Dict		,	_lowercase       :	List[Any]		,	_lowercase       :	int=0				):



    """simple docstring"""



    if str(_lowercase				).startswith('''mps'''				):
     _UpperCamelCase:  Tuple    						=			torch.manual_seed(_lowercase				)
    else:
     _UpperCamelCase:  List[str]    						=			torch.Generator(device=_lowercase				).manual_seed(_lowercase				)
    _UpperCamelCase:  Dict    						=			{
        '''prompt''': '''.''',
        '''generator''': generator,
        '''num_inference_steps''': 2,
        '''guidance_scale''': 1.0,
        '''sag_scale''': 1.0,
        '''output_type''': '''numpy''',
    }
    return inputs



   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    super().test_inference_batch_single_identical(expected_max_diff=3E-3				)




@slow
@require_torch_gpu
class 			__magic_name__	(		unittest.TestCase     ):





   """simple docstring"""




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



    """simple docstring"""



    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()



   def 	lowerCAmelCase    (						self       :	int				):



    """simple docstring"""



    _UpperCamelCase:  Any    						=			StableDiffusionSAGPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''				)
    _UpperCamelCase:  Any    						=			sag_pipe.to(_lowercase				)
    sag_pipe.set_progress_bar_config(disable=_lowercase				)

    _UpperCamelCase:  Tuple    						=			'''.'''
    _UpperCamelCase:  Tuple    						=			torch.manual_seed(0				)
    _UpperCamelCase:  Union[str, Any]    						=			sag_pipe(
        [prompt]		,	generator=_lowercase		,	guidance_scale=7.5		,	sag_scale=1.0		,	num_inference_steps=20		,	output_type='''np'''				)

    _UpperCamelCase:  Optional[Any]    						=			output.images

    _UpperCamelCase:  Tuple    						=			image[0, -3:, -3:, -1]

    assert image.shape == (1, 512, 512, 3)
    _UpperCamelCase:  Any    						=			np.array([0.1568, 0.1738, 0.1695, 0.1693, 0.1507, 0.1705, 0.1547, 0.1751, 0.1949]				)

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



   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    _UpperCamelCase:  Optional[Any]    						=			StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''				)
    _UpperCamelCase:  int    						=			sag_pipe.to(_lowercase				)
    sag_pipe.set_progress_bar_config(disable=_lowercase				)

    _UpperCamelCase:  List[Any]    						=			'''.'''
    _UpperCamelCase:  str    						=			torch.manual_seed(0				)
    _UpperCamelCase:  Dict    						=			sag_pipe(
        [prompt]		,	generator=_lowercase		,	guidance_scale=7.5		,	sag_scale=1.0		,	num_inference_steps=20		,	output_type='''np'''				)

    _UpperCamelCase:  Union[str, Any]    						=			output.images

    _UpperCamelCase:  List[str]    						=			image[0, -3:, -3:, -1]

    assert image.shape == (1, 512, 512, 3)
    _UpperCamelCase:  int    						=			np.array([0.3459, 0.2876, 0.2537, 0.3002, 0.2671, 0.2160, 0.3026, 0.2262, 0.2371]				)

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



   def 	lowerCAmelCase    (						self       :	Any				):



    """simple docstring"""



    _UpperCamelCase:  Tuple    						=			StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''				)
    _UpperCamelCase:  int    						=			sag_pipe.to(_lowercase				)
    sag_pipe.set_progress_bar_config(disable=_lowercase				)

    _UpperCamelCase:  Any    						=			'''.'''
    _UpperCamelCase:  List[Any]    						=			torch.manual_seed(0				)
    _UpperCamelCase:  str    						=			sag_pipe(
        [prompt]		,	width=768		,	height=512		,	generator=_lowercase		,	guidance_scale=7.5		,	sag_scale=1.0		,	num_inference_steps=20		,	output_type='''np'''		,	)

    _UpperCamelCase:  str    						=			output.images

    assert image.shape == (1, 512, 768, 3)