
"""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      :
     def __init__(     self  ,     snake_case__  ,     snake_case__=2  ,     snake_case__=True  ,     snake_case__=False  ,     snake_case__=10  ,     snake_case__=3  ,     snake_case__=32 * 8  ,     snake_case__=32 * 8  ,     snake_case__=4  ,     snake_case__=64  ,     ):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				str							    = parent
       _SCREAMING_SNAKE_CASE   :				List[str]							    = batch_size
       _SCREAMING_SNAKE_CASE   :				Tuple							    = is_training
       _SCREAMING_SNAKE_CASE   :				str							    = use_auxiliary_loss
       _SCREAMING_SNAKE_CASE   :				Any							    = num_queries
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = num_channels
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = min_size
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = max_size
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = num_labels
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = hidden_dim
       _SCREAMING_SNAKE_CASE   :				str							    = hidden_dim
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				int							    = floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]				).to(
           A_				)

       _SCREAMING_SNAKE_CASE   :				Dict							    = torch.ones([self.batch_size, self.min_size, self.max_size]  ,     device=A_				)

       _SCREAMING_SNAKE_CASE   :				Tuple							    = (
           torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]  ,     device=A_				) > 0.5
       ).float()
       _SCREAMING_SNAKE_CASE   :				str							    = (torch.rand((self.batch_size, self.num_labels)  ,     device=A_				) > 0.5).long()

       _SCREAMING_SNAKE_CASE   :				List[Any]							    = self.get_config()
       return config, pixel_values, pixel_mask, mask_labels, class_labels
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Tuple							    = MaskaFormerConfig(
           hidden_size=self.hidden_dim  ,     )
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = self.num_queries
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = self.num_labels

       _SCREAMING_SNAKE_CASE   :				Any							    = [1, 1, 1, 1]
       _SCREAMING_SNAKE_CASE   :				Dict							    = self.num_channels

       _SCREAMING_SNAKE_CASE   :				str							    = 64
       _SCREAMING_SNAKE_CASE   :				Dict							    = 128
       _SCREAMING_SNAKE_CASE   :				Tuple							    = self.hidden_dim
       _SCREAMING_SNAKE_CASE   :				str							    = self.hidden_dim
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = self.hidden_dim
       return config
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.prepare_config_and_inputs()
       _SCREAMING_SNAKE_CASE   :				int							    = {"pixel_values": pixel_values, "pixel_mask": pixel_mask}
       return config, inputs_dict
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Any							    = output.encoder_hidden_states
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = output.pixel_decoder_hidden_states
       _SCREAMING_SNAKE_CASE   :				str							    = output.transformer_decoder_hidden_states

       self.parent.assertTrue(len(A_				)  ,     len(config.backbone_config.depths				)				)
       self.parent.assertTrue(len(A_				)  ,     len(config.backbone_config.depths				)				)
       self.parent.assertTrue(len(A_				)  ,     config.decoder_layers				)
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__=False				):




       """simple docstring"""


       with torch.no_grad():
         _SCREAMING_SNAKE_CASE   :				Optional[int]							    = MaskaFormerModel(config=A_				)
         model.to(A_				)
         model.eval()

         _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = model(pixel_values=A_  ,     pixel_mask=A_				)
         _SCREAMING_SNAKE_CASE   :				Optional[int]							    = model(A_  ,     output_hidden_states=A_				)

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




     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				str							    = MaskaFormerForUniversalSegmentation(config=A_				)
       model.to(A_				)
       model.eval()

       def comm_check_on_output(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():
         _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = model(pixel_values=A_  ,     pixel_mask=A_				)
         _SCREAMING_SNAKE_CASE   :				int							    = model(A_				)

         comm_check_on_output(A_				)

         _SCREAMING_SNAKE_CASE   :				List[str]							    = model(
             pixel_values=A_  ,     pixel_mask=A_  ,     mask_labels=A_  ,     class_labels=A_				)

       comm_check_on_output(A_				)

       self.parent.assertTrue(result.loss is not None				)
       self.parent.assertEqual(result.loss.shape  ,     torch.Size([1]				)				)
@require_torch
class 			UpperCamelCase      (							_lowercase      ,       _lowercase      ,       unittest.TestCase  ):
     A__      =			(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
     A__      =			{'''feature-extraction''': MaskaFormerModel} if is_torch_available() else {}

     A__      =			False
     A__      =			False
     A__      =			False
     A__      =			False
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = MaskaFormerModelTester(self				)
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = ConfigTester(self  ,     config_class=A_  ,     has_text_modality=A_				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       self.config_tester.run_common_tests()
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				Dict							    = self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.create_and_check_maskaformer_model(A_  ,     **A_  ,     output_hidden_states=A_				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Dict							    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*A_				)
     @unittest.skip(reason="Mask2Former does not use inputs_embeds"				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


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




       """simple docstring"""


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




       """simple docstring"""


       pass
     @unittest.skip(reason="Mask2Former does not use token embeddings"				)
     def 				__SCREAMING_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 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


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




       """simple docstring"""


       pass
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				List[str]							    = self.model_tester.prepare_config_and_inputs_for_common()

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

         _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = ["pixel_values"]
         self.assertListEqual(arg_names[:1]  ,     A_				)
     @slow
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
         _SCREAMING_SNAKE_CASE   :				str							    = MaskaFormerModel.from_pretrained(A_				)
         self.assertIsNotNone(A_				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Any							    = (self.model_tester.min_size,) * 2
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = {
           "pixel_values": torch.randn((2, 3, *size)  ,     device=A_				),
           "mask_labels": torch.randn((2, 10, *size)  ,     device=A_				),
           "class_labels": torch.zeros(2  ,     10  ,     device=A_				).long(),
       }
       _SCREAMING_SNAKE_CASE   :				Tuple							    = self.model_tester.get_config()

       _SCREAMING_SNAKE_CASE   :				int							    = MaskaFormerForUniversalSegmentation(A_				).to(A_				)
       _SCREAMING_SNAKE_CASE   :				Any							    = model(**A_				)
       self.assertTrue(outputs.loss is not None				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				int							    = self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.create_and_check_maskaformer_model(A_  ,     **A_  ,     output_hidden_states=A_				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				Any							    = self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
         _SCREAMING_SNAKE_CASE   :				List[str]							    = model_class(A_				).to(A_				)
         _SCREAMING_SNAKE_CASE   :				Tuple							    = model(**A_  ,     output_attentions=A_				)
         self.assertTrue(outputs.attentions is not None				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       if not self.model_tester.is_training:
         return

       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = self.all_model_classes[1]
       _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				Tuple							    = self.model_tester.prepare_config_and_inputs()

       _SCREAMING_SNAKE_CASE   :				Any							    = model_class(A_				)
       model.to(A_				)
       model.train()

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




     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


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

       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = model_class(A_				).to(A_				)
       model.train()

       _SCREAMING_SNAKE_CASE   :				Any							    = model(A_  ,     mask_labels=A_  ,     class_labels=A_				)

       _SCREAMING_SNAKE_CASE   :				List[Any]							    = outputs.encoder_hidden_states[0]
       encoder_hidden_states.retain_grad()

       _SCREAMING_SNAKE_CASE   :				List[Any]							    = outputs.pixel_decoder_hidden_states[0]
       pixel_decoder_hidden_states.retain_grad()

       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = outputs.transformer_decoder_hidden_states[0]
       transformer_decoder_hidden_states.retain_grad()

       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = outputs.attentions[0]
       attentions.retain_grad()

       outputs.loss.backward(retain_graph=A_				)

       self.assertIsNotNone(encoder_hidden_states.grad				)
       self.assertIsNotNone(pixel_decoder_hidden_states.grad				)
       self.assertIsNotNone(transformer_decoder_hidden_states.grad				)
       self.assertIsNotNone(attentions.grad				)
lowercase_       :						int     					=	1e-4





def    _lowerCAmelCase    (     )  ->    int:
  _SCREAMING_SNAKE_CASE   :				int							    = Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"				)
  return image





@require_vision
@slow
class 			UpperCamelCase      (							unittest.TestCase  ):
     @cached_property
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       return "facebook/mask2former-swin-small-coco-instance"
     @cached_property
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


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




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[str]							    = MaskaFormerModel.from_pretrained(self.model_checkpoints				).to(A_				)
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = self.default_image_processor
       _SCREAMING_SNAKE_CASE   :				Dict							    = prepare_img()
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = image_processor(A_  ,     return_tensors="pt"				).to(A_				)
       _SCREAMING_SNAKE_CASE   :				int							    = inputs["pixel_values"].shape
       # check size is divisible by 32
       self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0				)
       # check size
       self.assertEqual(A_  ,     (1, 3, 384, 384)				)

       with torch.no_grad():
         _SCREAMING_SNAKE_CASE   :				Any							    = model(**A_				)

       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = 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(A_				)
       self.assertTrue(
           torch.allclose(
               outputs.encoder_last_hidden_state[0, 0, :3, :3]  ,     A_  ,     atol=A_				)				)

       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = 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(A_				)
       self.assertTrue(
           torch.allclose(
               outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]  ,     A_  ,     atol=A_				)				)

       _SCREAMING_SNAKE_CASE   :				List[Any]							    = 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(A_				)
       self.assertTrue(
           torch.allclose(
               outputs.transformer_decoder_last_hidden_state[0, :3, :3]  ,     A_  ,     atol=A_				)				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints				).to(A_				).eval()
       _SCREAMING_SNAKE_CASE   :				Any							    = self.default_image_processor
       _SCREAMING_SNAKE_CASE   :				int							    = prepare_img()
       _SCREAMING_SNAKE_CASE   :				Any							    = image_processor(A_  ,     return_tensors="pt"				).to(A_				)
       _SCREAMING_SNAKE_CASE   :				Optional[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(A_  ,     (1, 3, 384, 384)				)

       with torch.no_grad():
         _SCREAMING_SNAKE_CASE   :				List[Any]							    = model(**A_				)
       # masks_queries_logits
       _SCREAMING_SNAKE_CASE   :				Any							    = outputs.masks_queries_logits
       self.assertEqual(
           masks_queries_logits.shape  ,     (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)				)
       _SCREAMING_SNAKE_CASE   :				str							    = [
           [-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],
       ]
       _SCREAMING_SNAKE_CASE   :				int							    = torch.tensor(A_				).to(A_				)
       self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]  ,     A_  ,     atol=A_				)				)
       # class_queries_logits
       _SCREAMING_SNAKE_CASE   :				str							    = outputs.class_queries_logits
       self.assertEqual(class_queries_logits.shape  ,     (1, model.config.num_queries, model.config.num_labels + 1)				)
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = 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(A_				)
       self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]  ,     A_  ,     atol=A_				)				)




     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				int							    = MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints				).to(A_				).eval()
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = self.default_image_processor

       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = 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"  ,     )

       _SCREAMING_SNAKE_CASE   :				str							    = inputs["pixel_values"].to(A_				)
       _SCREAMING_SNAKE_CASE   :				int							    = [el.to(A_				) for el in inputs["mask_labels"]]
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = [el.to(A_				) for el in inputs["class_labels"]]

       with torch.no_grad():
         _SCREAMING_SNAKE_CASE   :				int							    = model(**A_				)

       self.assertTrue(outputs.loss is not None				)







"""simple docstring"""
import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE      ,       unittest.TestCase  ):
     A__      =			DiTPipeline
     A__      =			CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
     A__      =			PipelineTesterMixin.required_optional_params - {
         """latents""",
         """num_images_per_prompt""",
         """callback""",
         """callback_steps""",
     }
     A__      =			CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
     A__      =			False
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       torch.manual_seed(0				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = TransformeraDModel(
           sample_size=16  ,     num_layers=2  ,     patch_size=4  ,     attention_head_dim=8  ,     num_attention_heads=2  ,     in_channels=4  ,     out_channels=8  ,     attention_bias=snake_case__  ,     activation_fn="gelu-approximate"  ,     num_embeds_ada_norm=1000  ,     norm_type="ada_norm_zero"  ,     norm_elementwise_affine=snake_case__  ,     )
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = AutoencoderKL()
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = DDIMScheduler()
       _SCREAMING_SNAKE_CASE   :				Dict							    = {"transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler}
       return components
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__=0				):




       """simple docstring"""


       if str(snake_case__				).startswith("mps"				):
         _SCREAMING_SNAKE_CASE   :				Tuple							    = torch.manual_seed(snake_case__				)
       else:
         _SCREAMING_SNAKE_CASE   :				List[Any]							    = torch.Generator(device=snake_case__				).manual_seed(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = {
           "class_labels": [1],
           "generator": generator,
           "num_inference_steps": 2,
           "output_type": "numpy",
       }
       return inputs
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Dict							    = "cpu"

       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.get_dummy_components()
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = self.pipeline_class(**snake_case__				)
       pipe.to(snake_case__				)
       pipe.set_progress_bar_config(disable=snake_case__				)

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

       self.assertEqual(image.shape  ,     (1, 16, 16, 3)				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = np.array([0.2_946, 0.6_601, 0.4_329, 0.3_296, 0.4_144, 0.5_319, 0.7_273, 0.5_013, 0.4_457]				)
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = np.abs(image_slice.flatten() - expected_slice				).max()
       self.assertLessEqual(snake_case__  ,     1E-3				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       self._test_inference_batch_single_identical(relax_max_difference=snake_case__  ,     expected_max_diff=1E-3				)




     @unittest.skipIf(
         torch_device != "cuda" or not is_xformers_available()  ,     reason="XFormers attention is only available with CUDA and `xformers` installed"  ,     )
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3				)
@require_torch_gpu
@slow
class 			UpperCamelCase      (							unittest.TestCase  ):
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				str							    = torch.manual_seed(0				)

       _SCREAMING_SNAKE_CASE   :				str							    = DiTPipeline.from_pretrained("facebook/DiT-XL-2-256"				)
       pipe.to("cuda"				)

       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = ["vase", "umbrella", "white shark", "white wolf"]
       _SCREAMING_SNAKE_CASE   :				Any							    = pipe.get_label_ids(snake_case__				)

       _SCREAMING_SNAKE_CASE   :				List[str]							    = pipe(snake_case__  ,     generator=snake_case__  ,     num_inference_steps=40  ,     output_type="np"				).images

       for word, image in zip(snake_case__  ,     snake_case__				):
         _SCREAMING_SNAKE_CASE   :				int							    = load_numpy(
             F'''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy'''				)
         assert np.abs((expected_image - image).max()				) < 1E-2




     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Any							    = DiTPipeline.from_pretrained("facebook/DiT-XL-2-512"				)
       _SCREAMING_SNAKE_CASE   :				Dict							    = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config				)
       pipe.to("cuda"				)

       _SCREAMING_SNAKE_CASE   :				List[Any]							    = ["vase", "umbrella"]
       _SCREAMING_SNAKE_CASE   :				int							    = pipe.get_label_ids(snake_case__				)

       _SCREAMING_SNAKE_CASE   :				str							    = torch.manual_seed(0				)
       _SCREAMING_SNAKE_CASE   :				Tuple							    = pipe(snake_case__  ,     generator=snake_case__  ,     num_inference_steps=25  ,     output_type="np"				).images

       for word, image in zip(snake_case__  ,     snake_case__				):
         _SCREAMING_SNAKE_CASE   :				Tuple							    = load_numpy(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
             F'''/dit/{word}_512.npy'''				)

         assert np.abs((expected_image - image).max()				) < 1E-1



