



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 				lowerCAmelCase_		:
       def __init__(			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_=2,			SCREAMING_SNAKE_CASE_=True,			SCREAMING_SNAKE_CASE_=False,			SCREAMING_SNAKE_CASE_=10,			SCREAMING_SNAKE_CASE_=3,			SCREAMING_SNAKE_CASE_=32 * 8,			SCREAMING_SNAKE_CASE_=32 * 8,			SCREAMING_SNAKE_CASE_=4,			SCREAMING_SNAKE_CASE_=64,			)     ->		Optional[Any]:
          UpperCamelCase    :  str							=	parent
          UpperCamelCase    :  List[Any]							=	batch_size
          UpperCamelCase    :  Any							=	is_training
          UpperCamelCase    :  Optional[int]							=	use_auxiliary_loss
          UpperCamelCase    :  str							=	num_queries
          UpperCamelCase    :  Optional[Any]							=	num_channels
          UpperCamelCase    :  str							=	min_size
          UpperCamelCase    :  int							=	max_size
          UpperCamelCase    :  List[str]							=	num_labels
          UpperCamelCase    :  Dict							=	hidden_dim
          UpperCamelCase    :  List[str]							=	hidden_dim
       def 					snake_case_       (			self       )     ->		Tuple:
          UpperCamelCase    :  Any							=	floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]       ).to(
              SCREAMING_SNAKE_CASE_       )

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

          UpperCamelCase    :  Dict							=	(
              torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size],			device=SCREAMING_SNAKE_CASE_       ) > 0.5
          ).float()
          UpperCamelCase    :  Tuple							=	(torch.rand((self.batch_size, self.num_labels),			device=SCREAMING_SNAKE_CASE_       ) > 0.5).long()

          UpperCamelCase    :  List[str]							=	self.get_config()
          return config, pixel_values, pixel_mask, mask_labels, class_labels
       def 					snake_case_       (			self       )     ->		Tuple:
          UpperCamelCase    :  Dict							=	MaskaFormerConfig(
              hidden_size=self.hidden_dim,			)
          UpperCamelCase    :  Optional[Any]							=	self.num_queries
          UpperCamelCase    :  Optional[Any]							=	self.num_labels

          UpperCamelCase    :  List[Any]							=	[1, 1, 1, 1]
          UpperCamelCase    :  Optional[int]							=	self.num_channels

          UpperCamelCase    :  str							=	64
          UpperCamelCase    :  int							=	128
          UpperCamelCase    :  Any							=	self.hidden_dim
          UpperCamelCase    :  Tuple							=	self.hidden_dim
          UpperCamelCase    :  Optional[int]							=	self.hidden_dim
          return config
       def 					snake_case_       (			self       )     ->		Optional[Any]:
          UpperCamelCase	,     UpperCamelCase	,     UpperCamelCase	,     UpperCamelCase	,     UpperCamelCase    :  Optional[Any]							=	self.prepare_config_and_inputs()
          UpperCamelCase    :  Any							=	{'pixel_values': pixel_values, 'pixel_mask': pixel_mask}
          return config, inputs_dict
       def 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       )     ->		str:
          UpperCamelCase    :  Union[str, Any]							=	output.encoder_hidden_states
          UpperCamelCase    :  Any							=	output.pixel_decoder_hidden_states
          UpperCamelCase    :  List[Any]							=	output.transformer_decoder_hidden_states

          self.parent.assertTrue(len(SCREAMING_SNAKE_CASE_       ),			len(config.backbone_config.depths       )       )
          self.parent.assertTrue(len(SCREAMING_SNAKE_CASE_       ),			len(config.backbone_config.depths       )       )
          self.parent.assertTrue(len(SCREAMING_SNAKE_CASE_       ),			config.decoder_layers       )
       def 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_=False       )     ->		List[Any]:
          with torch.no_grad():
             UpperCamelCase    :  Tuple							=	MaskaFormerModel(config=SCREAMING_SNAKE_CASE_       )
             model.to(SCREAMING_SNAKE_CASE_       )
             model.eval()

             UpperCamelCase    :  Optional[Any]							=	model(pixel_values=SCREAMING_SNAKE_CASE_,			pixel_mask=SCREAMING_SNAKE_CASE_       )
             UpperCamelCase    :  Union[str, Any]							=	model(SCREAMING_SNAKE_CASE_,			output_hidden_states=SCREAMING_SNAKE_CASE_       )

          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(SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       )






       def 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       )     ->		List[str]:
          UpperCamelCase    :  Tuple							=	MaskaFormerForUniversalSegmentation(config=SCREAMING_SNAKE_CASE_       )
          model.to(SCREAMING_SNAKE_CASE_       )
          model.eval()

          def comm_check_on_output(SCREAMING_SNAKE_CASE_       ):
             # 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    :  int							=	model(pixel_values=SCREAMING_SNAKE_CASE_,			pixel_mask=SCREAMING_SNAKE_CASE_       )
             UpperCamelCase    :  Any							=	model(SCREAMING_SNAKE_CASE_       )

             comm_check_on_output(SCREAMING_SNAKE_CASE_       )

             UpperCamelCase    :  int							=	model(
                 pixel_values=SCREAMING_SNAKE_CASE_,			pixel_mask=SCREAMING_SNAKE_CASE_,			mask_labels=SCREAMING_SNAKE_CASE_,			class_labels=SCREAMING_SNAKE_CASE_       )

          comm_check_on_output(SCREAMING_SNAKE_CASE_       )

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






@require_torch
class 				lowerCAmelCase_		(  a__       ,    a__       ,    unittest.TestCase							):
       UpperCAmelCase__				:					Union[str, Any]        =  (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
       UpperCAmelCase__				:					Any        =  {"feature-extraction": MaskaFormerModel} if is_torch_available() else {}

       UpperCAmelCase__				:					Tuple        =  False
       UpperCAmelCase__				:					Optional[int]        =  False
       UpperCAmelCase__				:					List[Any]        =  False
       UpperCAmelCase__				:					int        =  False
       def 					snake_case_       (			self       )     ->		int:
          UpperCamelCase    :  int							=	MaskaFormerModelTester(self       )
          UpperCamelCase    :  Dict							=	ConfigTester(self,			config_class=SCREAMING_SNAKE_CASE_,			has_text_modality=SCREAMING_SNAKE_CASE_       )
       def 					snake_case_       (			self       )     ->		Optional[int]:
          self.config_tester.run_common_tests()
       def 					snake_case_       (			self       )     ->		Optional[Any]:
          UpperCamelCase	,     UpperCamelCase    :  int							=	self.model_tester.prepare_config_and_inputs_for_common()
          self.model_tester.create_and_check_maskaformer_model(SCREAMING_SNAKE_CASE_,			**SCREAMING_SNAKE_CASE_,			output_hidden_states=SCREAMING_SNAKE_CASE_       )
       def 					snake_case_       (			self       )     ->		List[str]:
          UpperCamelCase    :  Tuple							=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*SCREAMING_SNAKE_CASE_       )
       @unittest.skip(reason='Mask2Former does not use inputs_embeds'       )
       def 					snake_case_       (			self       )     ->		Union[str, Any]:
          pass
       @unittest.skip(reason='Mask2Former does not have a get_input_embeddings method'       )
       def 					snake_case_       (			self       )     ->		Union[str, Any]:
          pass
       @unittest.skip(reason='Mask2Former is not a generative model'       )
       def 					snake_case_       (			self       )     ->		Union[str, Any]:
          pass
       @unittest.skip(reason='Mask2Former does not use token embeddings'       )
       def 					snake_case_       (			self       )     ->		Any:
          pass
       @require_torch_multi_gpu
       @unittest.skip(
           reason='Mask2Former has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'       )
       def 					snake_case_       (			self       )     ->		Optional[Any]:
          pass
       @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'       )
       def 					snake_case_       (			self       )     ->		str:
          pass
       def 					snake_case_       (			self       )     ->		Optional[Any]:
          UpperCamelCase	,     UpperCamelCase    :  Union[str, Any]							=	self.model_tester.prepare_config_and_inputs_for_common()

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

             UpperCamelCase    :  List[Any]							=	['pixel_values']
             self.assertListEqual(arg_names[:1],			SCREAMING_SNAKE_CASE_       )
       @slow
       def 					snake_case_       (			self       )     ->		List[Any]:
          for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
             UpperCamelCase    :  str							=	MaskaFormerModel.from_pretrained(SCREAMING_SNAKE_CASE_       )
             self.assertIsNotNone(SCREAMING_SNAKE_CASE_       )
       def 					snake_case_       (			self       )     ->		str:
          UpperCamelCase    :  str							=	(self.model_tester.min_size,) * 2
          UpperCamelCase    :  Dict							=	{
              'pixel_values': torch.randn((2, 3, *size),			device=SCREAMING_SNAKE_CASE_       ),
              'mask_labels': torch.randn((2, 10, *size),			device=SCREAMING_SNAKE_CASE_       ),
              'class_labels': torch.zeros(2,			10,			device=SCREAMING_SNAKE_CASE_       ).long(),
          }
          UpperCamelCase    :  Any							=	self.model_tester.get_config()

          UpperCamelCase    :  str							=	MaskaFormerForUniversalSegmentation(SCREAMING_SNAKE_CASE_       ).to(SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  Optional[Any]							=	model(**SCREAMING_SNAKE_CASE_       )
          self.assertTrue(outputs.loss is not None       )
       def 					snake_case_       (			self       )     ->		int:
          UpperCamelCase	,     UpperCamelCase    :  Optional[int]							=	self.model_tester.prepare_config_and_inputs_for_common()
          self.model_tester.create_and_check_maskaformer_model(SCREAMING_SNAKE_CASE_,			**SCREAMING_SNAKE_CASE_,			output_hidden_states=SCREAMING_SNAKE_CASE_       )
       def 					snake_case_       (			self       )     ->		List[str]:
          UpperCamelCase	,     UpperCamelCase    :  Optional[int]							=	self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             UpperCamelCase    :  Union[str, Any]							=	model_class(SCREAMING_SNAKE_CASE_       ).to(SCREAMING_SNAKE_CASE_       )
             UpperCamelCase    :  Optional[Any]							=	model(**SCREAMING_SNAKE_CASE_,			output_attentions=SCREAMING_SNAKE_CASE_       )
             self.assertTrue(outputs.attentions is not None       )
       def 					snake_case_       (			self       )     ->		Any:
          if not self.model_tester.is_training:
             return

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

          UpperCamelCase    :  Optional[int]							=	model_class(SCREAMING_SNAKE_CASE_       )
          model.to(SCREAMING_SNAKE_CASE_       )
          model.train()

          UpperCamelCase    :  Dict							=	model(SCREAMING_SNAKE_CASE_,			mask_labels=SCREAMING_SNAKE_CASE_,			class_labels=SCREAMING_SNAKE_CASE_       ).loss
          loss.backward()






       def 					snake_case_       (			self       )     ->		int:
          UpperCamelCase    :  Union[str, Any]							=	self.all_model_classes[1]
          UpperCamelCase	,     UpperCamelCase	,     UpperCamelCase	,     UpperCamelCase	,     UpperCamelCase    :  int							=	self.model_tester.prepare_config_and_inputs()
          UpperCamelCase    :  List[str]							=	True
          UpperCamelCase    :  Optional[Any]							=	True

          UpperCamelCase    :  Any							=	model_class(SCREAMING_SNAKE_CASE_       ).to(SCREAMING_SNAKE_CASE_       )
          model.train()

          UpperCamelCase    :  str							=	model(SCREAMING_SNAKE_CASE_,			mask_labels=SCREAMING_SNAKE_CASE_,			class_labels=SCREAMING_SNAKE_CASE_       )

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

          UpperCamelCase    :  Union[str, Any]							=	outputs.pixel_decoder_hidden_states[0]
          pixel_decoder_hidden_states.retain_grad()

          UpperCamelCase    :  Tuple							=	outputs.transformer_decoder_hidden_states[0]
          transformer_decoder_hidden_states.retain_grad()

          UpperCamelCase    :  List[Any]							=	outputs.attentions[0]
          attentions.retain_grad()

          outputs.loss.backward(retain_graph=SCREAMING_SNAKE_CASE_       )

          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 		UpperCamelCase	(   )	->     List[str]:
   UpperCamelCase    :  str							=	Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'	)
   return image



@require_vision
@slow
class 				lowerCAmelCase_		(  unittest.TestCase							):
       @cached_property
       def 					snake_case_       (			self       )     ->		int:
          return "facebook/mask2former-swin-small-coco-instance"
       @cached_property
       def 					snake_case_       (			self       )     ->		Tuple:
          return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints       ) if is_vision_available() else None
       def 					snake_case_       (			self       )     ->		str:
          UpperCamelCase    :  Optional[Any]							=	MaskaFormerModel.from_pretrained(self.model_checkpoints       ).to(SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  Tuple							=	self.default_image_processor
          UpperCamelCase    :  List[Any]							=	prepare_img()
          UpperCamelCase    :  List[Any]							=	image_processor(SCREAMING_SNAKE_CASE_,			return_tensors='pt'       ).to(SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  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(SCREAMING_SNAKE_CASE_,			(1, 3, 384, 384)       )

          with torch.no_grad():
             UpperCamelCase    :  Union[str, Any]							=	model(**SCREAMING_SNAKE_CASE_       )

          UpperCamelCase    :  List[str]							=	torch.tensor(
              [[-0.27_90, -1.07_17, -1.16_68], [-0.51_28, -0.31_28, -0.49_87], [-0.58_32, 0.19_71, -0.01_97]]       ).to(SCREAMING_SNAKE_CASE_       )
          self.assertTrue(
              torch.allclose(
                  outputs.encoder_last_hidden_state[0, 0, :3, :3],			SCREAMING_SNAKE_CASE_,			atol=SCREAMING_SNAKE_CASE_       )       )

          UpperCamelCase    :  str							=	torch.tensor(
              [[0.89_73, 1.18_47, 1.17_76], [1.19_34, 1.50_40, 1.51_28], [1.11_53, 1.44_86, 1.49_51]]       ).to(SCREAMING_SNAKE_CASE_       )
          self.assertTrue(
              torch.allclose(
                  outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3],			SCREAMING_SNAKE_CASE_,			atol=SCREAMING_SNAKE_CASE_       )       )

          UpperCamelCase    :  Optional[Any]							=	torch.tensor(
              [[2.11_52, 1.70_00, -0.86_03], [1.58_08, 1.80_04, -0.93_53], [1.60_43, 1.74_95, -0.59_99]]       ).to(SCREAMING_SNAKE_CASE_       )
          self.assertTrue(
              torch.allclose(
                  outputs.transformer_decoder_last_hidden_state[0, :3, :3],			SCREAMING_SNAKE_CASE_,			atol=SCREAMING_SNAKE_CASE_       )       )
       def 					snake_case_       (			self       )     ->		List[Any]:
          UpperCamelCase    :  int							=	MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints       ).to(SCREAMING_SNAKE_CASE_       ).eval()
          UpperCamelCase    :  List[str]							=	self.default_image_processor
          UpperCamelCase    :  List[Any]							=	prepare_img()
          UpperCamelCase    :  List[str]							=	image_processor(SCREAMING_SNAKE_CASE_,			return_tensors='pt'       ).to(SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  str							=	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(SCREAMING_SNAKE_CASE_,			(1, 3, 384, 384)       )

          with torch.no_grad():
             UpperCamelCase    :  Optional[Any]							=	model(**SCREAMING_SNAKE_CASE_       )
          # masks_queries_logits
          UpperCamelCase    :  Any							=	outputs.masks_queries_logits
          self.assertEqual(
              masks_queries_logits.shape,			(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)       )
          UpperCamelCase    :  Tuple							=	[
              [-8.78_39, -9.00_56, -8.81_21],
              [-7.41_04, -7.03_13, -6.54_01],
              [-6.61_05, -6.34_27, -6.46_75],
          ]
          UpperCamelCase    :  Union[str, Any]							=	torch.tensor(SCREAMING_SNAKE_CASE_       ).to(SCREAMING_SNAKE_CASE_       )
          self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3],			SCREAMING_SNAKE_CASE_,			atol=SCREAMING_SNAKE_CASE_       )       )
          # class_queries_logits
          UpperCamelCase    :  Any							=	outputs.class_queries_logits
          self.assertEqual(class_queries_logits.shape,			(1, model.config.num_queries, model.config.num_labels + 1)       )
          UpperCamelCase    :  List[str]							=	torch.tensor(
              [
                  [1.83_24, -8.08_35, -4.19_22],
                  [0.84_50, -9.00_50, -3.60_53],
                  [0.30_45, -7.72_93, -3.02_75],
              ]       ).to(SCREAMING_SNAKE_CASE_       )
          self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3],			SCREAMING_SNAKE_CASE_,			atol=SCREAMING_SNAKE_CASE_       )       )






       def 					snake_case_       (			self       )     ->		Optional[int]:
          UpperCamelCase    :  List[Any]							=	MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints       ).to(SCREAMING_SNAKE_CASE_       ).eval()
          UpperCamelCase    :  Union[str, Any]							=	self.default_image_processor

          UpperCamelCase    :  List[str]							=	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    :  str							=	inputs['pixel_values'].to(SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  Any							=	[el.to(SCREAMING_SNAKE_CASE_       ) for el in inputs['mask_labels']]
          UpperCamelCase    :  Optional[int]							=	[el.to(SCREAMING_SNAKE_CASE_       ) for el in inputs['class_labels']]

          with torch.no_grad():
             UpperCamelCase    :  Any							=	model(**SCREAMING_SNAKE_CASE_       )

          self.assertTrue(outputs.loss is not None       )








import gc
import random
import unittest

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

import diffusers
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
def 		UpperCamelCase	(   snake_case__       :    Tuple	)	->     List[str]:
   UpperCamelCase    :  Optional[Any]							=	[tensor.shape for tensor in tensor_list]
   return all(shape == shapes[0] for shape in shapes[1:]	)



class 				lowerCAmelCase_		(  a__       ,    a__       ,    a__       ,    unittest.TestCase							):
       UpperCAmelCase__				:					str        =  StableDiffusionLatentUpscalePipeline
       UpperCAmelCase__				:					Optional[Any]        =  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
           "height",
           "width",
           "cross_attention_kwargs",
           "negative_prompt_embeds",
           "prompt_embeds",
       }
       UpperCAmelCase__				:					str        =  PipelineTesterMixin.required_optional_params - {"num_images_per_prompt"}
       UpperCAmelCase__				:					Any        =  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
       UpperCAmelCase__				:					List[str]        =  frozenset(
           []							)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
       UpperCAmelCase__				:					Union[str, Any]        =  frozenset([]							)

       UpperCAmelCase__				:					Optional[int]        =  True
       @property
       def 					snake_case_       (			self       )     ->		Union[str, Any]:
          UpperCamelCase    :  List[Any]							=	1
          UpperCamelCase    :  List[str]							=	4
          UpperCamelCase    :  List[str]							=	(16, 16)

          UpperCamelCase    :  Optional[int]							=	floats_tensor((batch_size, num_channels) + sizes,			rng=random.Random(0       )       ).to(SCREAMING_SNAKE_CASE_       )
          return image
       def 					snake_case_       (			self       )     ->		int:
          torch.manual_seed(0       )
          UpperCamelCase    :  int							=	UNetaDConditionModel(
              act_fn='gelu',			attention_head_dim=8,			norm_num_groups=SCREAMING_SNAKE_CASE_,			block_out_channels=[32, 32, 64, 64],			time_cond_proj_dim=160,			conv_in_kernel=1,			conv_out_kernel=1,			cross_attention_dim=32,			down_block_types=(
                  'KDownBlock2D',
                  'KCrossAttnDownBlock2D',
                  'KCrossAttnDownBlock2D',
                  'KCrossAttnDownBlock2D',
              ),			in_channels=8,			mid_block_type=SCREAMING_SNAKE_CASE_,			only_cross_attention=SCREAMING_SNAKE_CASE_,			out_channels=5,			resnet_time_scale_shift='scale_shift',			time_embedding_type='fourier',			timestep_post_act='gelu',			up_block_types=('KCrossAttnUpBlock2D', 'KCrossAttnUpBlock2D', 'KCrossAttnUpBlock2D', 'KUpBlock2D'),			)
          UpperCamelCase    :  Tuple							=	AutoencoderKL(
              block_out_channels=[32, 32, 64, 64],			in_channels=3,			out_channels=3,			down_block_types=[
                  'DownEncoderBlock2D',
                  'DownEncoderBlock2D',
                  'DownEncoderBlock2D',
                  'DownEncoderBlock2D',
              ],			up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D'],			latent_channels=4,			)
          UpperCamelCase    :  Tuple							=	EulerDiscreteScheduler(prediction_type='sample'       )
          UpperCamelCase    :  Optional[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=1000,			hidden_act='quick_gelu',			projection_dim=512,			)
          UpperCamelCase    :  Any							=	CLIPTextModel(SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  Union[str, Any]							=	CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

          UpperCamelCase    :  Optional[int]							=	{
              'unet': model.eval(),
              'vae': vae.eval(),
              'scheduler': scheduler,
              'text_encoder': text_encoder,
              'tokenizer': tokenizer,
          }

          return components
       def 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_=0       )     ->		List[str]:
          if str(SCREAMING_SNAKE_CASE_       ).startswith('mps'       ):
             UpperCamelCase    :  List[str]							=	torch.manual_seed(SCREAMING_SNAKE_CASE_       )
          else:
             UpperCamelCase    :  Optional[Any]							=	torch.Generator(device=SCREAMING_SNAKE_CASE_       ).manual_seed(SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  List[str]							=	{
              'prompt': 'A painting of a squirrel eating a burger',
              'image': self.dummy_image.cpu(),
              'generator': generator,
              'num_inference_steps': 2,
              'output_type': 'numpy',
          }
          return inputs
       def 					snake_case_       (			self       )     ->		List[Any]:
          UpperCamelCase    :  str							=	'cpu'

          UpperCamelCase    :  Tuple							=	self.get_dummy_components()
          UpperCamelCase    :  Dict							=	self.pipeline_class(**SCREAMING_SNAKE_CASE_       )
          pipe.to(SCREAMING_SNAKE_CASE_       )
          pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_       )

          UpperCamelCase    :  Union[str, Any]							=	self.get_dummy_inputs(SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  List[Any]							=	pipe(**SCREAMING_SNAKE_CASE_       ).images
          UpperCamelCase    :  str							=	image[0, -3:, -3:, -1]

          self.assertEqual(image.shape,			(1, 256, 256, 3)       )
          UpperCamelCase    :  List[Any]							=	np.array(
              [0.47_22_24_12, 0.41_92_16_33, 0.44_71_74_34, 0.46_87_41_92, 0.42_58_82_58, 0.46_15_07_26, 0.4_67_75_34, 0.45_58_38_32, 0.48_57_90_55]       )
          UpperCamelCase    :  Dict							=	np.abs(image_slice.flatten() - expected_slice       ).max()
          self.assertLessEqual(SCREAMING_SNAKE_CASE_,			1e-3       )
       def 					snake_case_       (			self       )     ->		List[Any]:
          super().test_attention_slicing_forward_pass(expected_max_diff=7e-3       )
       def 					snake_case_       (			self       )     ->		Tuple:
          super().test_cpu_offload_forward_pass(expected_max_diff=3e-3       )
       def 					snake_case_       (			self       )     ->		Dict:
          super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3       )
       def 					snake_case_       (			self       )     ->		Any:
          super().test_inference_batch_single_identical(expected_max_diff=7e-3       )
       def 					snake_case_       (			self       )     ->		List[str]:
          super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3e-3       )
       def 					snake_case_       (			self       )     ->		Optional[int]:
          super().test_save_load_local(expected_max_difference=3e-3       )
       def 					snake_case_       (			self       )     ->		int:
          super().test_save_load_optional_components(expected_max_difference=3e-3       )






       def 					snake_case_       (			self       )     ->		Optional[Any]:
          UpperCamelCase    :  Optional[Any]							=	[
              'DDIMScheduler',
              'DDPMScheduler',
              'PNDMScheduler',
              'HeunDiscreteScheduler',
              'EulerAncestralDiscreteScheduler',
              'KDPM2DiscreteScheduler',
              'KDPM2AncestralDiscreteScheduler',
              'DPMSolverSDEScheduler',
          ]
          UpperCamelCase    :  Tuple							=	self.get_dummy_components()
          UpperCamelCase    :  Tuple							=	self.pipeline_class(**SCREAMING_SNAKE_CASE_       )

          # make sure that PNDM does not need warm-up
          pipe.scheduler.register_to_config(skip_prk_steps=SCREAMING_SNAKE_CASE_       )

          pipe.to(SCREAMING_SNAKE_CASE_       )
          pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  int							=	self.get_dummy_inputs(SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  Optional[Any]							=	2

          UpperCamelCase    :  Optional[Any]							=	[]
          for scheduler_enum in KarrasDiffusionSchedulers:
             if scheduler_enum.name in skip_schedulers:
                # no sigma schedulers are not supported
                # no schedulers
                continue

             UpperCamelCase    :  List[Any]							=	getattr(SCREAMING_SNAKE_CASE_,			scheduler_enum.name       )
             UpperCamelCase    :  List[str]							=	scheduler_cls.from_config(pipe.scheduler.config       )
             UpperCamelCase    :  List[str]							=	pipe(**SCREAMING_SNAKE_CASE_       )[0]
             outputs.append(SCREAMING_SNAKE_CASE_       )

          assert check_same_shape(SCREAMING_SNAKE_CASE_       )






@require_torch_gpu
@slow
class 				lowerCAmelCase_		(  unittest.TestCase							):
       def 					snake_case_       (			self       )     ->		Union[str, Any]:
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()
       def 					snake_case_       (			self       )     ->		List[Any]:
          UpperCamelCase    :  Dict							=	torch.manual_seed(33       )

          UpperCamelCase    :  Union[str, Any]							=	StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4',			torch_dtype=torch.floataa       )
          pipe.to('cuda'       )

          UpperCamelCase    :  Tuple							=	StableDiffusionLatentUpscalePipeline.from_pretrained(
              'stabilityai/sd-x2-latent-upscaler',			torch_dtype=torch.floataa       )
          upscaler.to('cuda'       )

          UpperCamelCase    :  Union[str, Any]							=	'a photo of an astronaut high resolution, unreal engine, ultra realistic'

          UpperCamelCase    :  int							=	pipe(SCREAMING_SNAKE_CASE_,			generator=SCREAMING_SNAKE_CASE_,			output_type='latent'       ).images

          UpperCamelCase    :  List[str]							=	upscaler(
              prompt=SCREAMING_SNAKE_CASE_,			image=SCREAMING_SNAKE_CASE_,			num_inference_steps=20,			guidance_scale=0,			generator=SCREAMING_SNAKE_CASE_,			output_type='np',			).images[0]

          UpperCamelCase    :  Optional[int]							=	load_numpy(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/astronaut_1024.npy'       )
          assert np.abs((expected_image - image).mean()       ) < 5e-2






       def 					snake_case_       (			self       )     ->		int:
          UpperCamelCase    :  List[Any]							=	torch.manual_seed(33       )

          UpperCamelCase    :  Union[str, Any]							=	StableDiffusionLatentUpscalePipeline.from_pretrained(
              'stabilityai/sd-x2-latent-upscaler',			torch_dtype=torch.floataa       )
          upscaler.to('cuda'       )

          UpperCamelCase    :  Dict							=	'the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas'

          UpperCamelCase    :  List[str]							=	load_image(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_512.png'       )

          UpperCamelCase    :  str							=	upscaler(
              prompt=SCREAMING_SNAKE_CASE_,			image=SCREAMING_SNAKE_CASE_,			num_inference_steps=20,			guidance_scale=0,			generator=SCREAMING_SNAKE_CASE_,			output_type='np',			).images[0]

          UpperCamelCase    :  Optional[Any]							=	load_numpy(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_1024.npy'       )
          assert np.abs((expected_image - image).max()       ) < 5e-2



