
'''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      :     Any							,_a      :     Union[str, Any]							,_a      :     Any=2							,_a      :     Dict=True							,_a      :     Dict=False							,_a      :     str=10							,_a      :     Any=3							,_a      :     List[Any]=32 * 8							,_a      :     List[Any]=32 * 8							,_a      :     str=4							,_a      :     List[str]=64							,):


     '''simple docstring'''






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


  def   __lowercase		(		self      :     Dict       ):


     '''simple docstring'''






     _a				:	Optional[Any]			=   floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]       ).to(
         UpperCamelCase__       )

     _a				:	int			=   torch.ones([self.batch_size, self.min_size, self.max_size]							,device=UpperCamelCase__       )

     _a				:	Dict			=   (
         torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]							,device=UpperCamelCase__       ) > 0.5
     ).float()
     _a				:	Union[str, Any]			=   (torch.rand((self.batch_size, self.num_labels)							,device=UpperCamelCase__       ) > 0.5).long()

     _a				:	Optional[Any]			=   self.get_config()
     return config, pixel_values, pixel_mask, mask_labels, class_labels


  def   __lowercase		(		self      :     Optional[int]       ):


     '''simple docstring'''






     _a				:	Optional[int]			=   MaskaFormerConfig(
         hidden_size=self.hidden_dim							,)
     _a				:	Union[str, Any]			=   self.num_queries
     _a				:	Optional[int]			=   self.num_labels

     _a				:	Optional[int]			=   [1, 1, 1, 1]
     _a				:	str			=   self.num_channels

     _a				:	List[Any]			=   64
     _a				:	str			=   128
     _a				:	Any			=   self.hidden_dim
     _a				:	List[str]			=   self.hidden_dim
     _a				:	int			=   self.hidden_dim
     return config


  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     _a				:	str			=   self.prepare_config_and_inputs()
     _a				:	str			=   {"pixel_values": pixel_values, "pixel_mask": pixel_mask}
     return config, inputs_dict


  def   __lowercase		(		self      :     Optional[int]							,_a      :     Optional[int]							,_a      :     Tuple       ):


     '''simple docstring'''






     _a				:	Optional[Any]			=   output.encoder_hidden_states
     _a				:	Optional[Any]			=   output.pixel_decoder_hidden_states
     _a				:	str			=   output.transformer_decoder_hidden_states

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


  def   __lowercase		(		self      :     Optional[Any]							,_a      :     List[Any]							,_a      :     Any							,_a      :     int							,_a      :     List[str]=False       ):


     '''simple docstring'''






     with torch.no_grad():
        _a				:	List[str]			=   MaskaFormerModel(config=UpperCamelCase__       )
        model.to(UpperCamelCase__       )
        model.eval()

        _a				:	int			=   model(pixel_values=UpperCamelCase__							,pixel_mask=UpperCamelCase__       )
        _a				:	Union[str, Any]			=   model(UpperCamelCase__							,output_hidden_states=UpperCamelCase__       )

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


  def   __lowercase		(		self      :     Any							,_a      :     Dict							,_a      :     str							,_a      :     Tuple							,_a      :     Any							,_a      :     List[str]       ):


     '''simple docstring'''






     _a				:	str			=   MaskaFormerForUniversalSegmentation(config=UpperCamelCase__       )
     model.to(UpperCamelCase__       )
     model.eval()

     def comm_check_on_output(_a      :     Tuple       ):
        # 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():
        _a				:	Tuple			=   model(pixel_values=UpperCamelCase__							,pixel_mask=UpperCamelCase__       )
        _a				:	int			=   model(UpperCamelCase__       )

        comm_check_on_output(UpperCamelCase__       )

        _a				:	List[str]			=   model(
            pixel_values=UpperCamelCase__							,pixel_mask=UpperCamelCase__							,mask_labels=UpperCamelCase__							,class_labels=UpperCamelCase__       )

     comm_check_on_output(UpperCamelCase__       )

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







@require_torch
class UpperCAmelCase__	( __snake_case    ,				__snake_case    ,				unittest.TestCase	):





  """simple docstring"""

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

  __UpperCAmelCase				:				Any   =					False
  __UpperCAmelCase				:				Union[str, Any]   =					False
  __UpperCAmelCase				:				int   =					False
  __UpperCAmelCase				:				Union[str, Any]   =					False


  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     _a				:	List[Any]			=   MaskaFormerModelTester(self       )
     _a				:	Optional[int]			=   ConfigTester(self							,config_class=UpperCamelCase__							,has_text_modality=UpperCamelCase__       )


  def   __lowercase		(		self      :     Optional[int]       ):


     '''simple docstring'''






     self.config_tester.run_common_tests()


  def   __lowercase		(		self      :     Dict       ):


     '''simple docstring'''






     _a				:	List[Any]			=   self.model_tester.prepare_config_and_inputs_for_common()
     self.model_tester.create_and_check_maskaformer_model(UpperCamelCase__							,**UpperCamelCase__							,output_hidden_states=UpperCamelCase__       )


  def   __lowercase		(		self      :     List[Any]       ):


     '''simple docstring'''






     _a				:	Dict			=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*UpperCamelCase__       )


  @unittest.skip(reason='Mask2Former does not use inputs_embeds'       )
  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     pass


  @unittest.skip(reason='Mask2Former does not have a get_input_embeddings method'       )
  def   __lowercase		(		self      :     Union[str, Any]       ):


     '''simple docstring'''






     pass


  @unittest.skip(reason='Mask2Former is not a generative model'       )
  def   __lowercase		(		self      :     str       ):


     '''simple docstring'''






     pass


  @unittest.skip(reason='Mask2Former does not use token embeddings'       )
  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     pass


  @require_torch_multi_gpu
  @unittest.skip(
      reason='Mask2Former has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'       )
  def   __lowercase		(		self      :     str       ):


     '''simple docstring'''






     pass


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


     '''simple docstring'''






     pass


  def   __lowercase		(		self      :     List[Any]       ):


     '''simple docstring'''






     _a				:	Optional[Any]			=   self.model_tester.prepare_config_and_inputs_for_common()

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

        _a				:	str			=   ["pixel_values"]
        self.assertListEqual(arg_names[:1]							,UpperCamelCase__       )


  @slow
  def   __lowercase		(		self      :     List[Any]       ):


     '''simple docstring'''






     for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
        _a				:	Union[str, Any]			=   MaskaFormerModel.from_pretrained(UpperCamelCase__       )
        self.assertIsNotNone(UpperCamelCase__       )


  def   __lowercase		(		self      :     Optional[int]       ):


     '''simple docstring'''






     _a				:	Optional[Any]			=   (self.model_tester.min_size,) * 2
     _a				:	Dict			=   {
         "pixel_values": torch.randn((2, 3, *size)							,device=UpperCamelCase__       ),
         "mask_labels": torch.randn((2, 10, *size)							,device=UpperCamelCase__       ),
         "class_labels": torch.zeros(2							,10							,device=UpperCamelCase__       ).long(),
     }
     _a				:	Dict			=   self.model_tester.get_config()

     _a				:	int			=   MaskaFormerForUniversalSegmentation(UpperCamelCase__       ).to(UpperCamelCase__       )
     _a				:	str			=   model(**UpperCamelCase__       )
     self.assertTrue(outputs.loss is not None       )


  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     _a				:	Tuple			=   self.model_tester.prepare_config_and_inputs_for_common()
     self.model_tester.create_and_check_maskaformer_model(UpperCamelCase__							,**UpperCamelCase__							,output_hidden_states=UpperCamelCase__       )


  def   __lowercase		(		self      :     Optional[Any]       ):


     '''simple docstring'''






     _a				:	Tuple			=   self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        _a				:	Dict			=   model_class(UpperCamelCase__       ).to(UpperCamelCase__       )
        _a				:	List[str]			=   model(**UpperCamelCase__							,output_attentions=UpperCamelCase__       )
        self.assertTrue(outputs.attentions is not None       )


  def   __lowercase		(		self      :     Dict       ):


     '''simple docstring'''






     if not self.model_tester.is_training:
        return

     _a				:	Dict			=   self.all_model_classes[1]
     _a				:	str			=   self.model_tester.prepare_config_and_inputs()

     _a				:	Tuple			=   model_class(UpperCamelCase__       )
     model.to(UpperCamelCase__       )
     model.train()

     _a				:	Union[str, Any]			=   model(UpperCamelCase__							,mask_labels=UpperCamelCase__							,class_labels=UpperCamelCase__       ).loss
     loss.backward()


  def   __lowercase		(		self      :     List[Any]       ):


     '''simple docstring'''






     _a				:	List[Any]			=   self.all_model_classes[1]
     _a				:	str			=   self.model_tester.prepare_config_and_inputs()
     _a				:	Dict			=   True
     _a				:	List[Any]			=   True

     _a				:	int			=   model_class(UpperCamelCase__       ).to(UpperCamelCase__       )
     model.train()

     _a				:	Tuple			=   model(UpperCamelCase__							,mask_labels=UpperCamelCase__							,class_labels=UpperCamelCase__       )

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

     _a				:	Dict			=   outputs.pixel_decoder_hidden_states[0]
     pixel_decoder_hidden_states.retain_grad()

     _a				:	Any			=   outputs.transformer_decoder_hidden_states[0]
     transformer_decoder_hidden_states.retain_grad()

     _a				:	Optional[int]			=   outputs.attentions[0]
     attentions.retain_grad()

     outputs.loss.backward(retain_graph=UpperCamelCase__       )

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


__lowerCAmelCase         =    1e-4





def 					UpperCAmelCase_		():






   """simple docstring"""

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







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





  """simple docstring"""

  @cached_property
  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     return "facebook/mask2former-swin-small-coco-instance"


  @cached_property
  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints       ) if is_vision_available() else None


  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     _a				:	Union[str, Any]			=   MaskaFormerModel.from_pretrained(self.model_checkpoints       ).to(UpperCamelCase__       )
     _a				:	Optional[Any]			=   self.default_image_processor
     _a				:	List[str]			=   prepare_img()
     _a				:	List[str]			=   image_processor(UpperCamelCase__							,return_tensors='pt'       ).to(UpperCamelCase__       )
     _a				:	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(UpperCamelCase__							,(1, 3, 384, 384)       )

     with torch.no_grad():
        _a				:	List[str]			=   model(**UpperCamelCase__       )

     _a				:	str			=   torch.tensor(
         [[-0.2790, -1.0717, -1.1668], [-0.5128, -0.3128, -0.4987], [-0.5832, 0.1971, -0.0197]]       ).to(UpperCamelCase__       )
     self.assertTrue(
         torch.allclose(
             outputs.encoder_last_hidden_state[0, 0, :3, :3]							,UpperCamelCase__							,atol=UpperCamelCase__       )       )

     _a				:	int			=   torch.tensor(
         [[0.8973, 1.1847, 1.1776], [1.1934, 1.5040, 1.5128], [1.1153, 1.4486, 1.4951]]       ).to(UpperCamelCase__       )
     self.assertTrue(
         torch.allclose(
             outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]							,UpperCamelCase__							,atol=UpperCamelCase__       )       )

     _a				:	str			=   torch.tensor(
         [[2.1152, 1.7000, -0.8603], [1.5808, 1.8004, -0.9353], [1.6043, 1.7495, -0.5999]]       ).to(UpperCamelCase__       )
     self.assertTrue(
         torch.allclose(
             outputs.transformer_decoder_last_hidden_state[0, :3, :3]							,UpperCamelCase__							,atol=UpperCamelCase__       )       )


  def   __lowercase		(		self      :     int       ):


     '''simple docstring'''






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

     with torch.no_grad():
        _a				:	Optional[Any]			=   model(**UpperCamelCase__       )
     # masks_queries_logits
     _a				:	int			=   outputs.masks_queries_logits
     self.assertEqual(
         masks_queries_logits.shape							,(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)       )
     _a				:	str			=   [
         [-8.7839, -9.0056, -8.8121],
         [-7.4104, -7.0313, -6.5401],
         [-6.6105, -6.3427, -6.4675],
     ]
     _a				:	Optional[int]			=   torch.tensor(UpperCamelCase__       ).to(UpperCamelCase__       )
     self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]							,UpperCamelCase__							,atol=UpperCamelCase__       )       )
     # class_queries_logits
     _a				:	List[str]			=   outputs.class_queries_logits
     self.assertEqual(class_queries_logits.shape							,(1, model.config.num_queries, model.config.num_labels + 1)       )
     _a				:	str			=   torch.tensor(
         [
             [1.8324, -8.0835, -4.1922],
             [0.8450, -9.0050, -3.6053],
             [0.3045, -7.7293, -3.0275],
         ]       ).to(UpperCamelCase__       )
     self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]							,UpperCamelCase__							,atol=UpperCamelCase__       )       )


  def   __lowercase		(		self      :     str       ):


     '''simple docstring'''






     _a				:	Any			=   MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints       ).to(UpperCamelCase__       ).eval()
     _a				:	Any			=   self.default_image_processor

     _a				:	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'							,)

     _a				:	Tuple			=   inputs["pixel_values"].to(UpperCamelCase__       )
     _a				:	Dict			=   [el.to(UpperCamelCase__       ) for el in inputs["mask_labels"]]
     _a				:	Dict			=   [el.to(UpperCamelCase__       ) for el in inputs["class_labels"]]

     with torch.no_grad():
        _a				:	int			=   model(**UpperCamelCase__       )

     self.assertTrue(outputs.loss is not None       )






"""simple docstring"""


import gc
import random
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
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


enable_full_determinism()





class     snake_case	(				unittest.TestCase  ):

				def 							lowercase_     (	self							:				Optional[int])->   int:



						'''simple docstring'''


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

				@property
				def 							lowercase_     (	self							:				int)->   str:



						'''simple docstring'''


						__lowerCAmelCase:       str          =			1
						__lowerCAmelCase:       Union[str, Any]          =			3
						__lowerCAmelCase:       Union[str, Any]          =			(3_2, 3_2)

						__lowerCAmelCase:       Union[str, Any]          =			floats_tensor((batch_size, num_channels) + sizes     ,  rng=random.Random(0)).to(UpperCamelCase__)
						return image

				@property
				def 							lowercase_     (	self							:				Tuple)->   str:



						'''simple docstring'''


						torch.manual_seed(0)
						__lowerCAmelCase:       Optional[Any]          =			UNetaDConditionModel(
						    block_out_channels=(3_2, 3_2, 6_4)     ,  layers_per_block=2     ,  sample_size=3_2     ,  in_channels=7     ,  out_channels=4     ,  down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D")     ,  up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D")     ,  cross_attention_dim=3_2     ,  attention_head_dim=8     ,  use_linear_projection=UpperCamelCase__     ,  only_cross_attention=(True, True, False)     ,  num_class_embeds=1_0_0     ,  )
						return model

				@property
				def 							lowercase_     (	self							:				Any)->   Optional[Any]:



						'''simple docstring'''


						torch.manual_seed(0)
						__lowerCAmelCase:       Tuple          =			AutoencoderKL(
						    block_out_channels=[3_2, 3_2, 6_4]     ,  in_channels=3     ,  out_channels=3     ,  down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"]     ,  up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"]     ,  latent_channels=4     ,  )
						return model

				@property
				def 							lowercase_     (	self							:				Any)->   Optional[Any]:



						'''simple docstring'''


						torch.manual_seed(0)
						__lowerCAmelCase:       List[str]          =			CLIPTextConfig(
						    bos_token_id=0     ,  eos_token_id=2     ,  hidden_size=3_2     ,  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     ,  hidden_act="gelu"     ,  projection_dim=5_1_2     ,  )
						return CLIPTextModel(UpperCamelCase__)

				def 							lowercase_     (	self							:				List[str])->   Dict:



						'''simple docstring'''


						__lowerCAmelCase:       Tuple          =			"cpu"  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase:       int          =			self.dummy_cond_unet_upscale
						__lowerCAmelCase:       int          =			DDPMScheduler()
						__lowerCAmelCase:       List[str]          =			DDIMScheduler(prediction_type="v_prediction")
						__lowerCAmelCase:       Tuple          =			self.dummy_vae
						__lowerCAmelCase:       Optional[Any]          =			self.dummy_text_encoder
						__lowerCAmelCase:       Any          =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

						__lowerCAmelCase:       Tuple          =			self.dummy_image.cpu().permute(0     ,  2     ,  3     ,  1)[0]
						__lowerCAmelCase:       List[Any]          =			Image.fromarray(np.uinta(UpperCamelCase__)).convert("RGB").resize((6_4, 6_4))

						# make sure here that pndm scheduler skips prk
						__lowerCAmelCase:       Optional[int]          =			StableDiffusionUpscalePipeline(
						    unet=UpperCamelCase__     ,  low_res_scheduler=UpperCamelCase__     ,  scheduler=UpperCamelCase__     ,  vae=UpperCamelCase__     ,  text_encoder=UpperCamelCase__     ,  tokenizer=UpperCamelCase__     ,  max_noise_level=3_5_0     ,  )
						__lowerCAmelCase:       Tuple          =			sd_pipe.to(UpperCamelCase__)
						sd_pipe.set_progress_bar_config(disable=UpperCamelCase__)

						__lowerCAmelCase:       Any          =			"A painting of a squirrel eating a burger"
						__lowerCAmelCase:       str          =			torch.Generator(device=UpperCamelCase__).manual_seed(0)
						__lowerCAmelCase:       Optional[int]          =			sd_pipe(
						    [prompt]     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  guidance_scale=6.0     ,  noise_level=2_0     ,  num_inference_steps=2     ,  output_type="np"     ,  )

						__lowerCAmelCase:       List[str]          =			output.images

						__lowerCAmelCase:       Union[str, Any]          =			torch.Generator(device=UpperCamelCase__).manual_seed(0)
						__lowerCAmelCase:       List[str]          =			sd_pipe(
						    [prompt]     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  guidance_scale=6.0     ,  noise_level=2_0     ,  num_inference_steps=2     ,  output_type="np"     ,  return_dict=UpperCamelCase__     ,  )[0]

						__lowerCAmelCase:       int          =			image[0, -3:, -3:, -1]
						__lowerCAmelCase:       Dict          =			image_from_tuple[0, -3:, -3:, -1]

						__lowerCAmelCase:       Dict          =			low_res_image.size[0] * 4
						assert image.shape == (1, expected_height_width, expected_height_width, 3)
						__lowerCAmelCase:       List[Any]          =			np.array([0.3113, 0.3910, 0.4272, 0.4859, 0.5061, 0.4652, 0.5362, 0.5715, 0.5661])

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

				def 							lowercase_     (	self							:				List[str])->   Optional[Any]:



						'''simple docstring'''


						__lowerCAmelCase:       Union[str, Any]          =			"cpu"  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase:       Dict          =			self.dummy_cond_unet_upscale
						__lowerCAmelCase:       List[str]          =			DDPMScheduler()
						__lowerCAmelCase:       Union[str, Any]          =			DDIMScheduler(prediction_type="v_prediction")
						__lowerCAmelCase:       Optional[int]          =			self.dummy_vae
						__lowerCAmelCase:       List[Any]          =			self.dummy_text_encoder
						__lowerCAmelCase:       Optional[int]          =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

						__lowerCAmelCase:       List[Any]          =			self.dummy_image.cpu().permute(0     ,  2     ,  3     ,  1)[0]
						__lowerCAmelCase:       str          =			Image.fromarray(np.uinta(UpperCamelCase__)).convert("RGB").resize((6_4, 6_4))

						# make sure here that pndm scheduler skips prk
						__lowerCAmelCase:       Optional[int]          =			StableDiffusionUpscalePipeline(
						    unet=UpperCamelCase__     ,  low_res_scheduler=UpperCamelCase__     ,  scheduler=UpperCamelCase__     ,  vae=UpperCamelCase__     ,  text_encoder=UpperCamelCase__     ,  tokenizer=UpperCamelCase__     ,  max_noise_level=3_5_0     ,  )
						__lowerCAmelCase:       Optional[int]          =			sd_pipe.to(UpperCamelCase__)
						sd_pipe.set_progress_bar_config(disable=UpperCamelCase__)

						__lowerCAmelCase:       List[str]          =			"A painting of a squirrel eating a burger"
						__lowerCAmelCase:       List[Any]          =			sd_pipe(
						    2 * [prompt]     ,  image=2 * [low_res_image]     ,  guidance_scale=6.0     ,  noise_level=2_0     ,  num_inference_steps=2     ,  output_type="np"     ,  )
						__lowerCAmelCase:       List[Any]          =			output.images
						assert image.shape[0] == 2

						__lowerCAmelCase:       Dict          =			torch.Generator(device=UpperCamelCase__).manual_seed(0)
						__lowerCAmelCase:       Optional[Any]          =			sd_pipe(
						    [prompt]     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  num_images_per_prompt=2     ,  guidance_scale=6.0     ,  noise_level=2_0     ,  num_inference_steps=2     ,  output_type="np"     ,  )
						__lowerCAmelCase:       List[Any]          =			output.images
						assert image.shape[0] == 2




				@unittest.skipIf(torch_device != "cuda"     ,  "This test requires a GPU")
				def 							lowercase_     (	self							:				Tuple)->   Any:



						'''simple docstring'''


						__lowerCAmelCase:       Union[str, Any]          =			self.dummy_cond_unet_upscale
						__lowerCAmelCase:       int          =			DDPMScheduler()
						__lowerCAmelCase:       int          =			DDIMScheduler(prediction_type="v_prediction")
						__lowerCAmelCase:       Dict          =			self.dummy_vae
						__lowerCAmelCase:       int          =			self.dummy_text_encoder
						__lowerCAmelCase:       List[Any]          =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

						__lowerCAmelCase:       List[Any]          =			self.dummy_image.cpu().permute(0     ,  2     ,  3     ,  1)[0]
						__lowerCAmelCase:       Optional[int]          =			Image.fromarray(np.uinta(UpperCamelCase__)).convert("RGB").resize((6_4, 6_4))

						# put models in fp16, except vae as it overflows in fp16
						__lowerCAmelCase:       List[Any]          =			unet.half()
						__lowerCAmelCase:       List[str]          =			text_encoder.half()

						# make sure here that pndm scheduler skips prk
						__lowerCAmelCase:       List[Any]          =			StableDiffusionUpscalePipeline(
						    unet=UpperCamelCase__     ,  low_res_scheduler=UpperCamelCase__     ,  scheduler=UpperCamelCase__     ,  vae=UpperCamelCase__     ,  text_encoder=UpperCamelCase__     ,  tokenizer=UpperCamelCase__     ,  max_noise_level=3_5_0     ,  )
						__lowerCAmelCase:       str          =			sd_pipe.to(UpperCamelCase__)
						sd_pipe.set_progress_bar_config(disable=UpperCamelCase__)

						__lowerCAmelCase:       Optional[Any]          =			"A painting of a squirrel eating a burger"
						__lowerCAmelCase:       str          =			torch.manual_seed(0)
						__lowerCAmelCase:       Dict          =			sd_pipe(
						    [prompt]     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  num_inference_steps=2     ,  output_type="np"     ,  ).images

						__lowerCAmelCase:       Optional[Any]          =			low_res_image.size[0] * 4
						assert image.shape == (1, expected_height_width, expected_height_width, 3)





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

				def 							lowercase_     (	self							:				Tuple)->   Tuple:



						'''simple docstring'''


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

				def 							lowercase_     (	self							:				List[Any])->   Tuple:



						'''simple docstring'''


						__lowerCAmelCase:       Dict          =			load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/sd2-upscale/low_res_cat.png")
						__lowerCAmelCase:       Optional[Any]          =			load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
						    "/upsampled_cat.npy")

						__lowerCAmelCase:       str          =			"stabilityai/stable-diffusion-x4-upscaler"
						__lowerCAmelCase:       Optional[int]          =			StableDiffusionUpscalePipeline.from_pretrained(UpperCamelCase__)
						pipe.to(UpperCamelCase__)
						pipe.set_progress_bar_config(disable=UpperCamelCase__)
						pipe.enable_attention_slicing()

						__lowerCAmelCase:       Tuple          =			"a cat sitting on a park bench"

						__lowerCAmelCase:       int          =			torch.manual_seed(0)
						__lowerCAmelCase:       List[Any]          =			pipe(
						    prompt=UpperCamelCase__     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  output_type="np"     ,  )
						__lowerCAmelCase:       Dict          =			output.images[0]

						assert image.shape == (5_1_2, 5_1_2, 3)
						assert np.abs(expected_image - image).max() < 1e-3

				def 							lowercase_     (	self							:				Optional[int])->   Any:



						'''simple docstring'''


						__lowerCAmelCase:       Dict          =			load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/sd2-upscale/low_res_cat.png")
						__lowerCAmelCase:       Tuple          =			load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
						    "/upsampled_cat_fp16.npy")

						__lowerCAmelCase:       Optional[Any]          =			"stabilityai/stable-diffusion-x4-upscaler"
						__lowerCAmelCase:       Tuple          =			StableDiffusionUpscalePipeline.from_pretrained(
						    UpperCamelCase__     ,  torch_dtype=torch.floataa     ,  )
						pipe.to(UpperCamelCase__)
						pipe.set_progress_bar_config(disable=UpperCamelCase__)
						pipe.enable_attention_slicing()

						__lowerCAmelCase:       str          =			"a cat sitting on a park bench"

						__lowerCAmelCase:       List[str]          =			torch.manual_seed(0)
						__lowerCAmelCase:       Optional[Any]          =			pipe(
						    prompt=UpperCamelCase__     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  output_type="np"     ,  )
						__lowerCAmelCase:       Union[str, Any]          =			output.images[0]

						assert image.shape == (5_1_2, 5_1_2, 3)
						assert np.abs(expected_image - image).max() < 5e-1




				def 							lowercase_     (	self							:				Optional[int])->   Dict:



						'''simple docstring'''


						torch.cuda.empty_cache()
						torch.cuda.reset_max_memory_allocated()
						torch.cuda.reset_peak_memory_stats()

						__lowerCAmelCase:       Tuple          =			load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/sd2-upscale/low_res_cat.png")

						__lowerCAmelCase:       Union[str, Any]          =			"stabilityai/stable-diffusion-x4-upscaler"
						__lowerCAmelCase:       Any          =			StableDiffusionUpscalePipeline.from_pretrained(
						    UpperCamelCase__     ,  torch_dtype=torch.floataa     ,  )
						pipe.to(UpperCamelCase__)
						pipe.set_progress_bar_config(disable=UpperCamelCase__)
						pipe.enable_attention_slicing(1)
						pipe.enable_sequential_cpu_offload()

						__lowerCAmelCase:       int          =			"a cat sitting on a park bench"

						__lowerCAmelCase:       Dict          =			torch.manual_seed(0)
						__lowerCAmelCase:       Dict          =			pipe(
						    prompt=UpperCamelCase__     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  num_inference_steps=5     ,  output_type="np"     ,  )

						__lowerCAmelCase:       Optional[int]          =			torch.cuda.max_memory_allocated()
						# make sure that less than 2.9 GB is allocated
						assert mem_bytes < 2.9 * 1_0**9


