
"""simple docstring"""




import collections
import inspect
import unittest

from transformers import DinatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch
   from torch import nn

   from transformers import DinatBackbone, DinatForImageClassification, DinatModel
   from transformers.models.dinat.modeling_dinat import DINAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
   from PIL import Image

   from transformers import AutoImageProcessor

class 		a__						:
   def __init__(      self     ,   _a     ,   _a=13     ,   _a=64     ,   _a=4     ,   _a=3     ,   _a=16     ,   _a=[1, 2, 1]     ,   _a=[2, 4, 8]     ,   _a=3     ,   _a=[[3], [1, 2], [1]]     ,   _a=2.0     ,   _a=True     ,   _a=0.0     ,   _a=0.0     ,   _a=0.1     ,   _a="gelu"     ,   _a=True     ,   _a=0.0_2     ,   _a=1E-5     ,   _a=True     ,   _a=None     ,   _a=True     ,   _a=10     ,   _a=["stage1", "stage2"]     ,   _a=[1, 2]     ,   ):
          lowercase   :    List[Any]			=							parent
          lowercase   :    List[Any]			=							batch_size
          lowercase   :    Any			=							image_size
          lowercase   :    List[Any]			=							patch_size
          lowercase   :    str			=							num_channels
          lowercase   :    List[str]			=							embed_dim
          lowercase   :    Optional[Any]			=							depths
          lowercase   :    List[str]			=							num_heads
          lowercase   :    List[str]			=							kernel_size
          lowercase   :    Dict			=							dilations
          lowercase   :    Optional[int]			=							mlp_ratio
          lowercase   :    int			=							qkv_bias
          lowercase   :    int			=							hidden_dropout_prob
          lowercase   :    Union[str, Any]			=							attention_probs_dropout_prob
          lowercase   :    Optional[int]			=							drop_path_rate
          lowercase   :    List[str]			=							hidden_act
          lowercase   :    Tuple			=							patch_norm
          lowercase   :    Optional[Any]			=							layer_norm_eps
          lowercase   :    Optional[int]			=							initializer_range
          lowercase   :    str			=							is_training
          lowercase   :    List[Any]			=							scope
          lowercase   :    List[Any]			=							use_labels
          lowercase   :    Tuple			=							num_labels
          lowercase   :    Tuple			=							out_features
          lowercase   :    Optional[Any]			=							out_indices
   def       __magic_name__      (      self    ):
          lowercase   :    Optional[int]			=							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

          lowercase   :    Dict			=							None
          if self.use_labels:
                 lowercase   :    Dict			=							ids_tensor([self.batch_size]     ,   self.num_labels    )

          lowercase   :    str			=							self.get_config()

          return config, pixel_values, labels
   def       __magic_name__      (      self    ):
          return DinatConfig(
              num_labels=self.num_labels     ,   image_size=self.image_size     ,   patch_size=self.patch_size     ,   num_channels=self.num_channels     ,   embed_dim=self.embed_dim     ,   depths=self.depths     ,   num_heads=self.num_heads     ,   kernel_size=self.kernel_size     ,   dilations=self.dilations     ,   mlp_ratio=self.mlp_ratio     ,   qkv_bias=self.qkv_bias     ,   hidden_dropout_prob=self.hidden_dropout_prob     ,   attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,   drop_path_rate=self.drop_path_rate     ,   hidden_act=self.hidden_act     ,   patch_norm=self.patch_norm     ,   layer_norm_eps=self.layer_norm_eps     ,   initializer_range=self.initializer_range     ,   out_features=self.out_features     ,   out_indices=self.out_indices     ,   )
   def       __magic_name__      (      self     ,   _a     ,   _a     ,   _a    ):
          lowercase   :    Optional[Any]			=							DinatModel(config=_a    )
          model.to(_a    )
          model.eval()
          lowercase   :    Optional[int]			=							model(_a    )

          lowercase   :    Union[str, Any]			=							(config.image_size // config.patch_size) // (2 ** (len(config.depths    ) - 1))
          lowercase   :    List[Any]			=							int(config.embed_dim * 2 ** (len(config.depths    ) - 1)    )

          self.parent.assertEqual(
              result.last_hidden_state.shape     ,   (self.batch_size, expected_height, expected_width, expected_dim)    )
   def       __magic_name__      (      self     ,   _a     ,   _a     ,   _a    ):
          lowercase   :    Optional[int]			=							DinatForImageClassification(_a    )
          model.to(_a    )
          model.eval()
          lowercase   :    List[Any]			=							model(_a     ,   labels=_a    )
          self.parent.assertEqual(result.logits.shape     ,   (self.batch_size, self.num_labels)    )

          # test greyscale images
          lowercase   :    str			=							1
          lowercase   :    Optional[Any]			=							DinatForImageClassification(_a    )
          model.to(_a    )
          model.eval()

          lowercase   :    Tuple			=							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
          lowercase   :    List[str]			=							model(_a    )
          self.parent.assertEqual(result.logits.shape     ,   (self.batch_size, self.num_labels)    )
   def       __magic_name__      (      self     ,   _a     ,   _a     ,   _a    ):
          lowercase   :    Optional[Any]			=							DinatBackbone(config=_a    )
          model.to(_a    )
          model.eval()
          lowercase   :    Any			=							model(_a    )

          # verify hidden states
          self.parent.assertEqual(len(result.feature_maps    )     ,   len(config.out_features    )    )
          self.parent.assertListEqual(list(result.feature_maps[0].shape    )     ,   [self.batch_size, model.channels[0], 16, 16]    )

          # verify channels
          self.parent.assertEqual(len(model.channels    )     ,   len(config.out_features    )    )

          # verify backbone works with out_features=None
          lowercase   :    Any			=							None
          lowercase   :    Tuple			=							DinatBackbone(config=_a    )
          model.to(_a    )
          model.eval()
          lowercase   :    str			=							model(_a    )

          # verify feature maps
          self.parent.assertEqual(len(result.feature_maps    )     ,   1    )
          self.parent.assertListEqual(list(result.feature_maps[0].shape    )     ,   [self.batch_size, model.channels[-1], 4, 4]    )

          # verify channels
          self.parent.assertEqual(len(model.channels    )     ,   1    )

   def       __magic_name__      (      self    ):
          lowercase   :    Dict			=							self.prepare_config_and_inputs()
          lowercase		,   lowercase		,   lowercase   :    Any			=							config_and_inputs
          lowercase   :    Union[str, Any]			=							{"pixel_values": pixel_values}
          return config, inputs_dict


@require_natten
@require_torch
class 		a__						(				a_,      a_,      unittest.TestCase						):
   __lowerCAmelCase						=    (
       (
           DinatModel,
           DinatForImageClassification,
           DinatBackbone,
       )
       if is_torch_available()
       else ()
   )
   __lowerCAmelCase						=    (
       {"""feature-extraction""": DinatModel, """image-classification""": DinatForImageClassification}
       if is_torch_available()
       else {}
   )
   __lowerCAmelCase						=    False

   __lowerCAmelCase						=    False
   __lowerCAmelCase						=    False
   __lowerCAmelCase						=    False
   __lowerCAmelCase						=    False
   def       __magic_name__      (      self    ):
          lowercase   :    Any			=							DinatModelTester(self    )
          lowercase   :    Optional[int]			=							ConfigTester(self     ,   config_class=_a     ,   embed_dim=37    )
   def       __magic_name__      (      self    ):
          self.create_and_test_config_common_properties()
          self.config_tester.create_and_test_config_to_json_string()
          self.config_tester.create_and_test_config_to_json_file()
          self.config_tester.create_and_test_config_from_and_save_pretrained()
          self.config_tester.create_and_test_config_with_num_labels()
          self.config_tester.check_config_can_be_init_without_params()
          self.config_tester.check_config_arguments_init()
   def       __magic_name__      (      self    ):
          return
   def       __magic_name__      (      self    ):
          lowercase   :    Any			=							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*_a    )
   def       __magic_name__      (      self    ):
          lowercase   :    List[str]			=							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*_a    )
   def       __magic_name__      (      self    ):
          lowercase   :    Union[str, Any]			=							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_backbone(*_a    )
   @unittest.skip(reason="Dinat does not use inputs_embeds"    )
   def       __magic_name__      (      self    ):
          pass
   @unittest.skip(reason="Dinat does not use feedforward chunking"    )
   def       __magic_name__      (      self    ):
          pass
   def       __magic_name__      (      self    ):
          lowercase		,   lowercase   :    str			=							self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
                 lowercase   :    Any			=							model_class(_a    )
                 self.assertIsInstance(model.get_input_embeddings()     ,   (nn.Module)    )
                 lowercase   :    Optional[int]			=							model.get_output_embeddings()
                 self.assertTrue(x is None or isinstance(_a     ,   nn.Linear    )    )
   def       __magic_name__      (      self    ):
          lowercase		,   lowercase   :    str			=							self.model_tester.prepare_config_and_inputs_for_common()

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

                 lowercase   :    Optional[Any]			=							["pixel_values"]
                 self.assertListEqual(arg_names[:1]     ,   _a    )
   def       __magic_name__      (      self    ):
          self.skipTest("Dinat's attention operation is handled entirely by NATTEN."    )
   def       __magic_name__      (      self     ,   _a     ,   _a     ,   _a     ,   _a    ):
          lowercase   :    Dict			=							model_class(_a    )
          model.to(_a    )
          model.eval()

          with torch.no_grad():
                 lowercase   :    Tuple			=							model(**self._prepare_for_class(_a     ,   _a    )    )

          lowercase   :    Any			=							outputs.hidden_states

          lowercase   :    Any			=							getattr(
              self.model_tester     ,   "expected_num_hidden_layers"     ,   len(self.model_tester.depths    ) + 1    )
          self.assertEqual(len(_a    )     ,   _a    )

          # Dinat has a different seq_length
          lowercase   :    Optional[Any]			=							(
              config.patch_size
              if isinstance(config.patch_size     ,   collections.abc.Iterable    )
              else (config.patch_size, config.patch_size)
          )

          lowercase   :    Any			=							image_size[0] // patch_size[0]
          lowercase   :    Any			=							image_size[1] // patch_size[1]

          self.assertListEqual(
              list(hidden_states[0].shape[-3:]    )     ,   [height, width, self.model_tester.embed_dim]     ,   )

          if model_class.__name__ != "DinatBackbone":
                 lowercase   :    Dict			=							outputs.reshaped_hidden_states
                 self.assertEqual(len(_a    )     ,   _a    )

                 lowercase		,   lowercase		,   lowercase		,   lowercase   :    Dict			=							reshaped_hidden_states[0].shape
                 lowercase   :    Union[str, Any]			=							(
                     reshaped_hidden_states[0].view(_a     ,   _a     ,   _a     ,   _a    ).permute(0     ,   2     ,   3     ,   1    )
                 )
                 self.assertListEqual(
                     list(reshaped_hidden_states.shape[-3:]    )     ,   [height, width, self.model_tester.embed_dim]     ,   )
   def       __magic_name__      (      self    ):
          lowercase		,   lowercase   :    Any			=							self.model_tester.prepare_config_and_inputs_for_common()

          lowercase   :    List[Any]			=							(
              self.model_tester.image_size
              if isinstance(self.model_tester.image_size     ,   collections.abc.Iterable    )
              else (self.model_tester.image_size, self.model_tester.image_size)
          )

          for model_class in self.all_model_classes:
                 lowercase   :    int			=							True
                 self.check_hidden_states_output(_a     ,   _a     ,   _a     ,   _a    )

                 # check that output_hidden_states also work using config
                 del inputs_dict["output_hidden_states"]
                 lowercase   :    List[Any]			=							True

                 self.check_hidden_states_output(_a     ,   _a     ,   _a     ,   _a    )
   @slow
   def       __magic_name__      (      self    ):
          for model_name in DINAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowercase   :    Any			=							DinatModel.from_pretrained(_a    )
                 self.assertIsNotNone(_a    )

   def       __magic_name__      (      self    ):
          lowercase		,   lowercase   :    List[Any]			=							self.model_tester.prepare_config_and_inputs_for_common()

          lowercase   :    Any			=							_config_zero_init(_a    )
          for model_class in self.all_model_classes:
                 lowercase   :    Tuple			=							model_class(config=_a    )
                 for name, param in model.named_parameters():
                        if "embeddings" not in name and param.requires_grad:
                               self.assertIn(
                                   ((param.data.mean() * 1E9).round() / 1E9).item()     ,   [0.0, 1.0]     ,   msg=f"""Parameter {name} of model {model_class} seems not properly initialized"""     ,   )


@require_natten
@require_vision
@require_torch
class 		a__						(				unittest.TestCase						):
   @cached_property
   def       __magic_name__      (      self    ):
          return AutoImageProcessor.from_pretrained("shi-labs/dinat-mini-in1k-224"    ) if is_vision_available() else None

   @slow
   def       __magic_name__      (      self    ):
          lowercase   :    str			=							DinatForImageClassification.from_pretrained("shi-labs/dinat-mini-in1k-224"    ).to(_a    )
          lowercase   :    int			=							self.default_image_processor

          lowercase   :    Tuple			=							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"    )
          lowercase   :    Tuple			=							image_processor(images=_a     ,   return_tensors="pt"    ).to(_a    )

          # forward pass
          with torch.no_grad():
                 lowercase   :    Any			=							model(**_a    )

          # verify the logits
          lowercase   :    str			=							torch.Size((1, 1_000)    )
          self.assertEqual(outputs.logits.shape     ,   _a    )
          lowercase   :    int			=							torch.tensor([-0.1_5_4_5, -0.7_6_6_7, 0.4_6_4_2]    ).to(_a    )
          self.assertTrue(torch.allclose(outputs.logits[0, :3]     ,   _a     ,   atol=1E-4    )    )


@require_torch
@require_natten
class 		a__						(				unittest.TestCase,      a_						):
   __lowerCAmelCase						=    (DinatBackbone,) if is_torch_available() else ()
   __lowerCAmelCase						=    DinatConfig
   def       __magic_name__      (      self    ):
          lowercase   :    int			=							DinatModelTester(self    )




"""simple docstring"""




import unittest

import numpy as np
import torch

from diffusers import DDPMPipeline, DDPMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow, torch_device


enable_full_determinism()

class 		a__						(				unittest.TestCase						):
   @property
   def       __magic_name__      (      self    ):
          torch.manual_seed(0    )
          lowercase   :    int			=							UNetaDModel(
              block_out_channels=(32, 64)     ,   layers_per_block=2     ,   sample_size=32     ,   in_channels=3     ,   out_channels=3     ,   down_block_types=("DownBlock2D", "AttnDownBlock2D")     ,   up_block_types=("AttnUpBlock2D", "UpBlock2D")     ,   )
          return model
   def       __magic_name__      (      self    ):
          lowercase   :    int			=							"cpu"
          lowercase   :    Dict			=							self.dummy_uncond_unet
          lowercase   :    List[str]			=							DDPMScheduler()

          lowercase   :    Tuple			=							DDPMPipeline(unet=_a     ,   scheduler=_a    )
          ddpm.to(_a    )
          ddpm.set_progress_bar_config(disable=_a    )

          lowercase   :    Optional[int]			=							torch.Generator(device=_a    ).manual_seed(0    )
          lowercase   :    Optional[int]			=							ddpm(generator=_a     ,   num_inference_steps=2     ,   output_type="numpy"    ).images

          lowercase   :    List[str]			=							torch.Generator(device=_a    ).manual_seed(0    )
          lowercase   :    List[str]			=							ddpm(generator=_a     ,   num_inference_steps=2     ,   output_type="numpy"     ,   return_dict=_a    )[0]

          lowercase   :    List[str]			=							image[0, -3:, -3:, -1]
          lowercase   :    Any			=							image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 32, 32, 3)
          lowercase   :    List[str]			=							np.array(
              [9.956E-01, 5.785E-01, 4.675E-01, 9.930E-01, 0.0, 1.0_0_0, 1.199E-03, 2.648E-04, 5.101E-04]    )

          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       __magic_name__      (      self    ):
          lowercase   :    Any			=							self.dummy_uncond_unet
          lowercase   :    Dict			=							DDPMScheduler(prediction_type="sample"    )

          lowercase   :    List[Any]			=							DDPMPipeline(unet=_a     ,   scheduler=_a    )
          ddpm.to(_a    )
          ddpm.set_progress_bar_config(disable=_a    )

          lowercase   :    Any			=							torch.manual_seed(0    )
          lowercase   :    int			=							ddpm(generator=_a     ,   num_inference_steps=2     ,   output_type="numpy"    ).images

          lowercase   :    List[str]			=							torch.manual_seed(0    )
          lowercase   :    Optional[Any]			=							ddpm(generator=_a     ,   num_inference_steps=2     ,   output_type="numpy"    )[0]

          lowercase   :    Any			=							image[0, -3:, -3:, -1]
          lowercase   :    Dict			=							image_eps[0, -3:, -3:, -1]

          assert image.shape == (1, 32, 32, 3)
          lowercase   :    Tuple			=							1E-2 if torch_device != "mps" else 3E-2
          assert np.abs(image_slice.flatten() - image_eps_slice.flatten()    ).max() < tolerance


@slow
@require_torch_gpu
class 		a__						(				unittest.TestCase						):
   def       __magic_name__      (      self    ):
          lowercase   :    Optional[Any]			=							"google/ddpm-cifar10-32"

          lowercase   :    Dict			=							UNetaDModel.from_pretrained(_a    )
          lowercase   :    str			=							DDPMScheduler.from_pretrained(_a    )

          lowercase   :    int			=							DDPMPipeline(unet=_a     ,   scheduler=_a    )
          ddpm.to(_a    )
          ddpm.set_progress_bar_config(disable=_a    )

          lowercase   :    Dict			=							torch.manual_seed(0    )
          lowercase   :    List[str]			=							ddpm(generator=_a     ,   output_type="numpy"    ).images

          lowercase   :    int			=							image[0, -3:, -3:, -1]

          assert image.shape == (1, 32, 32, 3)
          lowercase   :    Optional[Any]			=							np.array([0.4_2_0_0, 0.3_5_8_8, 0.1_9_3_9, 0.3_8_4_7, 0.3_3_8_2, 0.2_6_4_7, 0.4_1_5_5, 0.3_5_8_2, 0.3_3_8_5]    )
          assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2


