


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





  """simple docstring"""




  def __init__(    self			,  UpperCamelCase			,  UpperCamelCase=2			,  UpperCamelCase=True			,  UpperCamelCase=False			,  UpperCamelCase=10			,  UpperCamelCase=3			,  UpperCamelCase=32 * 8			,  UpperCamelCase=32 * 8			,  UpperCamelCase=4			,  UpperCamelCase=64			,  ):







   """simple docstring"""
   lowerCamelCase_      =  parent
   lowerCamelCase_      =  batch_size
   lowerCamelCase_      =  is_training
   lowerCamelCase_      =  use_auxiliary_loss
   lowerCamelCase_      =  num_queries
   lowerCamelCase_      =  num_channels
   lowerCamelCase_      =  min_size
   lowerCamelCase_      =  max_size
   lowerCamelCase_      =  num_labels
   lowerCamelCase_      =  hidden_dim
   lowerCamelCase_      =  hidden_dim






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]     ).to(
       UpperCamelCase     )

   lowerCamelCase_      =  torch.ones([self.batch_size, self.min_size, self.max_size]			,  device=UpperCamelCase     )

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

   lowerCamelCase_      =  self.get_config()
   return config, pixel_values, pixel_mask, mask_labels, class_labels






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  MaskaFormerConfig(
       hidden_size=self.hidden_dim			,  )
   lowerCamelCase_      =  self.num_queries
   lowerCamelCase_      =  self.num_labels

   lowerCamelCase_      =  [1, 1, 1, 1]
   lowerCamelCase_      =  self.num_channels

   lowerCamelCase_      =  64
   lowerCamelCase_      =  128
   lowerCamelCase_      =  self.hidden_dim
   lowerCamelCase_      =  self.hidden_dim
   lowerCamelCase_      =  self.hidden_dim
   return config






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_      =  self.prepare_config_and_inputs()
   lowerCamelCase_      =  {"pixel_values": pixel_values, "pixel_mask": pixel_mask}
   return config, inputs_dict






  def   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase     ):







   """simple docstring"""
   lowerCamelCase_      =  output.encoder_hidden_states
   lowerCamelCase_      =  output.pixel_decoder_hidden_states
   lowerCamelCase_      =  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   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase=False     ):







   """simple docstring"""
   with torch.no_grad():
    lowerCamelCase_      =  MaskaFormerModel(config=UpperCamelCase     )
    model.to(UpperCamelCase     )
    model.eval()

    lowerCamelCase_      =  model(pixel_values=UpperCamelCase			,  pixel_mask=UpperCamelCase     )
    lowerCamelCase_      =  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   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase     ):







   """simple docstring"""
   lowerCamelCase_      =  MaskaFormerForUniversalSegmentation(config=UpperCamelCase     )
   model.to(UpperCamelCase     )
   model.eval()

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

    comm_check_on_output(UpperCamelCase     )

    lowerCamelCase_      =  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      snake_case    (    lowercase			,						lowercase			,						unittest.TestCase			):





  """simple docstring"""




  _lowerCamelCase		    =  (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
  _lowerCamelCase		    =  {"feature-extraction": MaskaFormerModel} if is_torch_available() else {}

  _lowerCamelCase		    =  False
  _lowerCamelCase		    =  False
  _lowerCamelCase		    =  False
  _lowerCamelCase		    =  False






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  MaskaFormerModelTester(self     )
   lowerCamelCase_      =  ConfigTester(self			,  config_class=UpperCamelCase			,  has_text_modality=UpperCamelCase     )






  def   snake_case		(    self     ):







   """simple docstring"""
   self.config_tester.run_common_tests()






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_ ,lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs_for_common()
   self.model_tester.create_and_check_maskaformer_model(UpperCamelCase			,  **UpperCamelCase			,  output_hidden_states=UpperCamelCase     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  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   snake_case		(    self     ):







   """simple docstring"""
   pass






  @unittest.skip(reason="Mask2Former does not have a get_input_embeddings method"     )
  def   snake_case		(    self     ):







   """simple docstring"""
   pass






  @unittest.skip(reason="Mask2Former is not a generative model"     )
  def   snake_case		(    self     ):







   """simple docstring"""
   pass






  @unittest.skip(reason="Mask2Former does not use token embeddings"     )
  def   snake_case		(    self     ):







   """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   snake_case		(    self     ):







   """simple docstring"""
   pass






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







   """simple docstring"""
   pass






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_ ,lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs_for_common()

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

    lowerCamelCase_      =  ["pixel_values"]
    self.assertListEqual(arg_names[:1]			,  UpperCamelCase     )






  @slow
  def   snake_case		(    self     ):







   """simple docstring"""
   for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
    lowerCamelCase_      =  MaskaFormerModel.from_pretrained(UpperCamelCase     )
    self.assertIsNotNone(UpperCamelCase     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  (self.model_tester.min_size,) * 2
   lowerCamelCase_      =  {
       "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(),
   }
   lowerCamelCase_      =  self.model_tester.get_config()

   lowerCamelCase_      =  MaskaFormerForUniversalSegmentation(UpperCamelCase     ).to(UpperCamelCase     )
   lowerCamelCase_      =  model(**UpperCamelCase     )
   self.assertTrue(outputs.loss is not None     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_ ,lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs_for_common()
   self.model_tester.create_and_check_maskaformer_model(UpperCamelCase			,  **UpperCamelCase			,  output_hidden_states=UpperCamelCase     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_ ,lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    lowerCamelCase_      =  model_class(UpperCamelCase     ).to(UpperCamelCase     )
    lowerCamelCase_      =  model(**UpperCamelCase			,  output_attentions=UpperCamelCase     )
    self.assertTrue(outputs.attentions is not None     )






  def   snake_case		(    self     ):







   """simple docstring"""
   if not self.model_tester.is_training:
    return

   lowerCamelCase_      =  self.all_model_classes[1]
   lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs()

   lowerCamelCase_      =  model_class(UpperCamelCase     )
   model.to(UpperCamelCase     )
   model.train()

   lowerCamelCase_      =  model(UpperCamelCase			,  mask_labels=UpperCamelCase			,  class_labels=UpperCamelCase     ).loss
   loss.backward()







  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.all_model_classes[1]
   lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs()
   lowerCamelCase_      =  True
   lowerCamelCase_      =  True

   lowerCamelCase_      =  model_class(UpperCamelCase     ).to(UpperCamelCase     )
   model.train()

   lowerCamelCase_      =  model(UpperCamelCase			,  mask_labels=UpperCamelCase			,  class_labels=UpperCamelCase     )

   lowerCamelCase_      =  outputs.encoder_hidden_states[0]
   encoder_hidden_states.retain_grad()

   lowerCamelCase_      =  outputs.pixel_decoder_hidden_states[0]
   pixel_decoder_hidden_states.retain_grad()

   lowerCamelCase_      =  outputs.transformer_decoder_hidden_states[0]
   transformer_decoder_hidden_states.retain_grad()

   lowerCamelCase_      =  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     )



a_  :				Any						= 1e-4





def __snake_case					(      ):
 lowerCamelCase_      =  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"		)
 return image







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





  """simple docstring"""




  @cached_property
  def   snake_case		(    self     ):







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






  @cached_property
  def   snake_case		(    self     ):







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






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  MaskaFormerModel.from_pretrained(self.model_checkpoints     ).to(UpperCamelCase     )
   lowerCamelCase_      =  self.default_image_processor
   lowerCamelCase_      =  prepare_img()
   lowerCamelCase_      =  image_processor(UpperCamelCase			,  return_tensors="pt"     ).to(UpperCamelCase     )
   lowerCamelCase_      =  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():
    lowerCamelCase_      =  model(**UpperCamelCase     )

   lowerCamelCase_      =  torch.tensor(
       [[-0.2_790, -1.0_717, -1.1_668], [-0.5_128, -0.3_128, -0.4_987], [-0.5_832, 0.1_971, -0.0_197]]     ).to(UpperCamelCase     )
   self.assertTrue(
       torch.allclose(
           outputs.encoder_last_hidden_state[0, 0, :3, :3]			,  UpperCamelCase			,  atol=UpperCamelCase     )     )

   lowerCamelCase_      =  torch.tensor(
       [[0.8_973, 1.1_847, 1.1_776], [1.1_934, 1.5_040, 1.5_128], [1.1_153, 1.4_486, 1.4_951]]     ).to(UpperCamelCase     )
   self.assertTrue(
       torch.allclose(
           outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]			,  UpperCamelCase			,  atol=UpperCamelCase     )     )

   lowerCamelCase_      =  torch.tensor(
       [[2.1_152, 1.7_000, -0.8_603], [1.5_808, 1.8_004, -0.9_353], [1.6_043, 1.7_495, -0.5_999]]     ).to(UpperCamelCase     )
   self.assertTrue(
       torch.allclose(
           outputs.transformer_decoder_last_hidden_state[0, :3, :3]			,  UpperCamelCase			,  atol=UpperCamelCase     )     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints     ).to(UpperCamelCase     ).eval()
   lowerCamelCase_      =  self.default_image_processor
   lowerCamelCase_      =  prepare_img()
   lowerCamelCase_      =  image_processor(UpperCamelCase			,  return_tensors="pt"     ).to(UpperCamelCase     )
   lowerCamelCase_      =  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():
    lowerCamelCase_      =  model(**UpperCamelCase     )
   # masks_queries_logits
   lowerCamelCase_      =  outputs.masks_queries_logits
   self.assertEqual(
       masks_queries_logits.shape			,  (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)     )
   lowerCamelCase_      =  [
       [-8.7_839, -9.0_056, -8.8_121],
       [-7.4_104, -7.0_313, -6.5_401],
       [-6.6_105, -6.3_427, -6.4_675],
   ]
   lowerCamelCase_      =  torch.tensor(UpperCamelCase     ).to(UpperCamelCase     )
   self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]			,  UpperCamelCase			,  atol=UpperCamelCase     )     )
   # class_queries_logits
   lowerCamelCase_      =  outputs.class_queries_logits
   self.assertEqual(class_queries_logits.shape			,  (1, model.config.num_queries, model.config.num_labels + 1)     )
   lowerCamelCase_      =  torch.tensor(
       [
           [1.8_324, -8.0_835, -4.1_922],
           [0.8_450, -9.0_050, -3.6_053],
           [0.3_045, -7.7_293, -3.0_275],
       ]     ).to(UpperCamelCase     )
   self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]			,  UpperCamelCase			,  atol=UpperCamelCase     )     )







  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints     ).to(UpperCamelCase     ).eval()
   lowerCamelCase_      =  self.default_image_processor

   lowerCamelCase_      =  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"			,  )

   lowerCamelCase_      =  inputs["pixel_values"].to(UpperCamelCase     )
   lowerCamelCase_      =  [el.to(UpperCamelCase     ) for el in inputs["mask_labels"]]
   lowerCamelCase_      =  [el.to(UpperCamelCase     ) for el in inputs["class_labels"]]

   with torch.no_grad():
    lowerCamelCase_      =  model(**UpperCamelCase     )

   self.assertTrue(outputs.loss is not None     )







'''simple docstring'''

import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()





class      snake_case    (    lowercase			,						lowercase			,						lowercase			,						unittest.TestCase			):





  """simple docstring"""




  _lowerCamelCase		    =  StableUnCLIPPipeline
  _lowerCamelCase		    =  TEXT_TO_IMAGE_PARAMS
  _lowerCamelCase		    =  TEXT_TO_IMAGE_BATCH_PARAMS
  _lowerCamelCase		    =  TEXT_TO_IMAGE_IMAGE_PARAMS
  _lowerCamelCase		    =  TEXT_TO_IMAGE_IMAGE_PARAMS

  # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
  _lowerCamelCase		    =  False






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  32
   lowerCamelCase_      =  embedder_hidden_size

   # prior components

   torch.manual_seed(0     )
   lowerCamelCase_      =  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

   torch.manual_seed(0     )
   lowerCamelCase_      =  CLIPTextModelWithProjection(
       CLIPTextConfig(
           bos_token_id=0			,  eos_token_id=2			,  hidden_size=UpperCamelCase			,  projection_dim=UpperCamelCase			,  intermediate_size=37			,  layer_norm_eps=1e-05			,  num_attention_heads=4			,  num_hidden_layers=5			,  pad_token_id=1			,  vocab_size=1000			,  )     )

   torch.manual_seed(0     )
   lowerCamelCase_      =  PriorTransformer(
       num_attention_heads=2			,  attention_head_dim=12			,  embedding_dim=UpperCamelCase			,  num_layers=1			,  )

   torch.manual_seed(0     )
   lowerCamelCase_      =  DDPMScheduler(
       variance_type="fixed_small_log"			,  prediction_type="sample"			,  num_train_timesteps=1000			,  clip_sample=UpperCamelCase			,  clip_sample_range=5.0			,  beta_schedule="squaredcos_cap_v2"			,  )

   # regular denoising components

   torch.manual_seed(0     )
   lowerCamelCase_      =  StableUnCLIPImageNormalizer(embedding_dim=UpperCamelCase     )
   lowerCamelCase_      =  DDPMScheduler(beta_schedule="squaredcos_cap_v2"     )

   torch.manual_seed(0     )
   lowerCamelCase_      =  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

   torch.manual_seed(0     )
   lowerCamelCase_      =  CLIPTextModel(
       CLIPTextConfig(
           bos_token_id=0			,  eos_token_id=2			,  hidden_size=UpperCamelCase			,  projection_dim=32			,  intermediate_size=37			,  layer_norm_eps=1e-05			,  num_attention_heads=4			,  num_hidden_layers=5			,  pad_token_id=1			,  vocab_size=1000			,  )     )

   torch.manual_seed(0     )
   lowerCamelCase_      =  UNetaDConditionModel(
       sample_size=32			,  in_channels=4			,  out_channels=4			,  down_block_types=("CrossAttnDownBlock2D", "DownBlock2D")			,  up_block_types=("UpBlock2D", "CrossAttnUpBlock2D")			,  block_out_channels=(32, 64)			,  attention_head_dim=(2, 4)			,  class_embed_type="projection"			,  projection_class_embeddings_input_dim=embedder_projection_dim * 2			,  cross_attention_dim=UpperCamelCase			,  layers_per_block=1			,  upcast_attention=UpperCamelCase			,  use_linear_projection=UpperCamelCase			,  )

   torch.manual_seed(0     )
   lowerCamelCase_      =  DDIMScheduler(
       beta_schedule="scaled_linear"			,  beta_start=0.00_085			,  beta_end=0.012			,  prediction_type="v_prediction"			,  set_alpha_to_one=UpperCamelCase			,  steps_offset=1			,  )

   torch.manual_seed(0     )
   lowerCamelCase_      =  AutoencoderKL()

   lowerCamelCase_      =  {
       # prior components
       "prior_tokenizer": prior_tokenizer,
       "prior_text_encoder": prior_text_encoder,
       "prior": prior,
       "prior_scheduler": prior_scheduler,
       # image noising components
       "image_normalizer": image_normalizer,
       "image_noising_scheduler": image_noising_scheduler,
       # regular denoising components
       "tokenizer": tokenizer,
       "text_encoder": text_encoder,
       "unet": unet,
       "scheduler": scheduler,
       "vae": vae,
   }

   return components






  def   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase=0     ):







   """simple docstring"""
   if str(UpperCamelCase     ).startswith("mps"     ):
    lowerCamelCase_      =  torch.manual_seed(UpperCamelCase     )
   else:
    lowerCamelCase_      =  torch.Generator(device=UpperCamelCase     ).manual_seed(UpperCamelCase     )
   lowerCamelCase_      =  {
       "prompt": "A painting of a squirrel eating a burger",
       "generator": generator,
       "num_inference_steps": 2,
       "prior_num_inference_steps": 2,
       "output_type": "numpy",
   }
   return inputs






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  torch_device == "cpu"

   self._test_attention_slicing_forward_pass(test_max_difference=UpperCamelCase     )







  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  torch_device in ["cpu", "mps"]

   self._test_inference_batch_single_identical(test_max_difference=UpperCamelCase     )





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





  """simple docstring"""




  def   snake_case		(    self     ):







   """simple docstring"""
   # clean up the VRAM after each test
   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  load_numpy(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"     )

   lowerCamelCase_      =  StableUnCLIPPipeline.from_pretrained("fusing/stable-unclip-2-1-l"			,  torch_dtype=torch.floataa     )
   pipe.to(UpperCamelCase     )
   pipe.set_progress_bar_config(disable=UpperCamelCase     )
   # stable unclip will oom when integration tests are run on a V100,
   # so turn on memory savings
   pipe.enable_attention_slicing()
   pipe.enable_sequential_cpu_offload()

   lowerCamelCase_      =  torch.Generator(device="cpu"     ).manual_seed(0     )
   lowerCamelCase_      =  pipe("anime turle"			,  generator=UpperCamelCase			,  output_type="np"     )

   lowerCamelCase_      =  output.images[0]

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

   assert_mean_pixel_difference(UpperCamelCase			,  UpperCamelCase     )







  def   snake_case		(    self     ):







   """simple docstring"""
   torch.cuda.empty_cache()
   torch.cuda.reset_max_memory_allocated()
   torch.cuda.reset_peak_memory_stats()

   lowerCamelCase_      =  StableUnCLIPPipeline.from_pretrained("fusing/stable-unclip-2-1-l"			,  torch_dtype=torch.floataa     )
   lowerCamelCase_      =  pipe.to(UpperCamelCase     )
   pipe.set_progress_bar_config(disable=UpperCamelCase     )
   pipe.enable_attention_slicing()
   pipe.enable_sequential_cpu_offload()

   lowerCamelCase_      =  pipe(
       "anime turtle"			,  prior_num_inference_steps=2			,  num_inference_steps=2			,  output_type="np"			,  )

   lowerCamelCase_      =  torch.cuda.max_memory_allocated()
   # make sure that less than 7 GB is allocated
   assert mem_bytes < 7 * 10**9



