





from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import (
        TFDeiTForImageClassification,
        TFDeiTForImageClassificationWithTeacher,
        TFDeiTForMaskedImageModeling,
        TFDeiTModel,
    )
    from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import DeiTImageProcessor

class    lowerCAmelCase							:
 def __init__(      self  :      List[str]			,     UpperCAmelCase  :      Dict			,     UpperCAmelCase  :      str=13			,     UpperCAmelCase  :      Tuple=30			,     UpperCAmelCase  :      List[Any]=2			,     UpperCAmelCase  :      Union[str, Any]=3			,     UpperCAmelCase  :      int=True			,     UpperCAmelCase  :      str=True			,     UpperCAmelCase  :      Tuple=32			,     UpperCAmelCase  :      Optional[int]=2			,     UpperCAmelCase  :      Optional[Any]=4			,     UpperCAmelCase  :      Tuple=37			,     UpperCAmelCase  :      Optional[Any]="gelu"			,     UpperCAmelCase  :      Dict=0.1			,     UpperCAmelCase  :      Tuple=0.1			,     UpperCAmelCase  :      Dict=10			,     UpperCAmelCase  :      List[Any]=0.0_2			,     UpperCAmelCase  :      int=3			,     UpperCAmelCase  :      List[str]=None			,     UpperCAmelCase  :      int=2			,     )			->	Dict:
     lowerCamelCase__    :      int					 =		parent
     lowerCamelCase__    :      Optional[Any]					 =		batch_size
     lowerCamelCase__    :      Dict					 =		image_size
     lowerCamelCase__    :      Tuple					 =		patch_size
     lowerCamelCase__    :      Tuple					 =		num_channels
     lowerCamelCase__    :      Any					 =		is_training
     lowerCamelCase__    :      Any					 =		use_labels
     lowerCamelCase__    :      Optional[Any]					 =		hidden_size
     lowerCamelCase__    :      Tuple					 =		num_hidden_layers
     lowerCamelCase__    :      Optional[int]					 =		num_attention_heads
     lowerCamelCase__    :      Optional[int]					 =		intermediate_size
     lowerCamelCase__    :      Optional[Any]					 =		hidden_act
     lowerCamelCase__    :      List[str]					 =		hidden_dropout_prob
     lowerCamelCase__    :      Optional[Any]					 =		attention_probs_dropout_prob
     lowerCamelCase__    :      str					 =		type_sequence_label_size
     lowerCamelCase__    :      Union[str, Any]					 =		initializer_range
     lowerCamelCase__    :      List[Any]					 =		scope
     lowerCamelCase__    :      Optional[int]					 =		encoder_stride

     # in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
     lowerCamelCase__    :      Tuple					 =		(image_size // patch_size) ** 2
     lowerCamelCase__    :      Tuple					 =		num_patches + 2
 def      A_	(      self  :      List[str]	)			->	Union[str, Any]:
     lowerCamelCase__    :      Tuple					 =		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

     lowerCamelCase__    :      Dict					 =		None
     if self.use_labels:
         lowerCamelCase__    :      Optional[Any]					 =		ids_tensor([self.batch_size]			,     self.type_sequence_label_size	)

     lowerCamelCase__    :      str					 =		self.get_config()

     return config, pixel_values, labels
 def      A_	(      self  :      Optional[Any]	)			->	List[Any]:
     return DeiTConfig(
         image_size=self.image_size			,     patch_size=self.patch_size			,     num_channels=self.num_channels			,     hidden_size=self.hidden_size			,     num_hidden_layers=self.num_hidden_layers			,     num_attention_heads=self.num_attention_heads			,     intermediate_size=self.intermediate_size			,     hidden_act=self.hidden_act			,     hidden_dropout_prob=self.hidden_dropout_prob			,     attention_probs_dropout_prob=self.attention_probs_dropout_prob			,     is_decoder=UpperCAmelCase			,     initializer_range=self.initializer_range			,     encoder_stride=self.encoder_stride			,     )
 def      A_	(      self  :      str			,     UpperCAmelCase  :      str			,     UpperCAmelCase  :      Optional[int]			,     UpperCAmelCase  :      Dict	)			->	Union[str, Any]:
     lowerCamelCase__    :      List[Any]					 =		TFDeiTModel(config=UpperCAmelCase	)
     lowerCamelCase__    :      Union[str, Any]					 =		model(UpperCAmelCase	)
     self.parent.assertEqual(result.last_hidden_state.shape			,     (self.batch_size, self.seq_length, self.hidden_size)	)
 def      A_	(      self  :      List[Any]			,     UpperCAmelCase  :      Optional[Any]			,     UpperCAmelCase  :      Dict			,     UpperCAmelCase  :      List[Any]	)			->	Tuple:
     lowerCamelCase__    :      int					 =		TFDeiTForMaskedImageModeling(config=UpperCAmelCase	)
     lowerCamelCase__    :      Union[str, Any]					 =		model(UpperCAmelCase	)
     self.parent.assertEqual(
         result.reconstruction.shape			,     (self.batch_size, self.num_channels, self.image_size, self.image_size)	)

     # test greyscale images
     lowerCamelCase__    :      Union[str, Any]					 =		1
     lowerCamelCase__    :      int					 =		TFDeiTForMaskedImageModeling(UpperCAmelCase	)

     lowerCamelCase__    :      Optional[int]					 =		floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
     lowerCamelCase__    :      Optional[Any]					 =		model(UpperCAmelCase	)
     self.parent.assertEqual(result.reconstruction.shape			,     (self.batch_size, 1, self.image_size, self.image_size)	)
 def      A_	(      self  :      Tuple			,     UpperCAmelCase  :      Any			,     UpperCAmelCase  :      str			,     UpperCAmelCase  :      List[str]	)			->	Optional[int]:
     lowerCamelCase__    :      Union[str, Any]					 =		self.type_sequence_label_size
     lowerCamelCase__    :      str					 =		TFDeiTForImageClassification(UpperCAmelCase	)
     lowerCamelCase__    :      Optional[Any]					 =		model(UpperCAmelCase			,     labels=UpperCAmelCase	)
     self.parent.assertEqual(result.logits.shape			,     (self.batch_size, self.type_sequence_label_size)	)

     # test greyscale images
     lowerCamelCase__    :      Optional[int]					 =		1
     lowerCamelCase__    :      Union[str, Any]					 =		TFDeiTForImageClassification(UpperCAmelCase	)

     lowerCamelCase__    :      Union[str, Any]					 =		floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
     lowerCamelCase__    :      Dict					 =		model(UpperCAmelCase			,     labels=UpperCAmelCase	)
     self.parent.assertEqual(result.logits.shape			,     (self.batch_size, self.type_sequence_label_size)	)







 def      A_	(      self  :      List[Any]	)			->	Union[str, Any]:
     lowerCamelCase__    :      List[str]					 =		self.prepare_config_and_inputs()
     lowerCamelCase__ ,							lowerCamelCase__ ,							lowerCamelCase__    :      int					 =		config_and_inputs
     lowerCamelCase__    :      Any					 =		{'pixel_values': pixel_values}
     return config, inputs_dict


@require_tf
class    lowerCAmelCase							(					__UpperCamelCase,  __UpperCamelCase,  unittest.TestCase			):
 UpperCAmelCase__													=		(
     (
         TFDeiTModel,
         TFDeiTForImageClassification,
         TFDeiTForImageClassificationWithTeacher,
         TFDeiTForMaskedImageModeling,
     )
     if is_tf_available()
     else ()
 )
 UpperCAmelCase__													=		(
     {
         """feature-extraction""": TFDeiTModel,
         """image-classification""": (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
     }
     if is_tf_available()
     else {}
 )

 UpperCAmelCase__													=		False
 UpperCAmelCase__													=		False
 UpperCAmelCase__													=		False
 UpperCAmelCase__													=		False
 def      A_	(      self  :      List[Any]	)			->	List[str]:
     lowerCamelCase__    :      List[Any]					 =		TFDeiTModelTester(self	)
     lowerCamelCase__    :      Union[str, Any]					 =		ConfigTester(self			,     config_class=UpperCAmelCase			,     has_text_modality=UpperCAmelCase			,     hidden_size=37	)
 def      A_	(      self  :      Tuple	)			->	Union[str, Any]:
     self.config_tester.run_common_tests()
 @unittest.skip(reason='DeiT does not use inputs_embeds'	)
 def      A_	(      self  :      int	)			->	List[Any]:
     pass
 def      A_	(      self  :      Optional[Any]	)			->	Optional[Any]:
     lowerCamelCase__ ,							lowerCamelCase__    :      List[Any]					 =		self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         lowerCamelCase__    :      Tuple					 =		model_class(UpperCAmelCase	)
         self.assertIsInstance(model.get_input_embeddings()			,     (tf.keras.layers.Layer)	)
         lowerCamelCase__    :      Tuple					 =		model.get_output_embeddings()
         self.assertTrue(x is None or isinstance(UpperCAmelCase			,     tf.keras.layers.Dense	)	)
 def      A_	(      self  :      List[str]	)			->	str:
     lowerCamelCase__ ,							lowerCamelCase__    :      int					 =		self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         lowerCamelCase__    :      List[Any]					 =		model_class(UpperCAmelCase	)
         lowerCamelCase__    :      Optional[int]					 =		inspect.signature(model.call	)
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         lowerCamelCase__    :      Any					 =		[*signature.parameters.keys()]

         lowerCamelCase__    :      List[Any]					 =		['pixel_values']
         self.assertListEqual(arg_names[:1]			,     UpperCAmelCase	)
 def      A_	(      self  :      List[Any]	)			->	Dict:
     lowerCamelCase__    :      List[Any]					 =		self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*UpperCAmelCase	)
 def      A_	(      self  :      List[str]	)			->	int:
     lowerCamelCase__    :      Union[str, Any]					 =		self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_image_modeling(*UpperCAmelCase	)
 def      A_	(      self  :      int	)			->	str:
     lowerCamelCase__    :      int					 =		self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase	)
 def      A_	(      self  :      List[str]			,     UpperCAmelCase  :      Dict			,     UpperCAmelCase  :      Optional[int]			,     UpperCAmelCase  :      Any=False	)			->	Union[str, Any]:
     lowerCamelCase__    :      List[str]					 =		super()._prepare_for_class(UpperCAmelCase			,     UpperCAmelCase			,     return_labels=UpperCAmelCase	)

     if return_labels:
         if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call	).parameters:
             del inputs_dict["labels"]

     return inputs_dict







 @slow
 def      A_	(      self  :      Optional[int]	)			->	Tuple:
     for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         lowerCamelCase__    :      List[str]					 =		TFDeiTModel.from_pretrained(UpperCAmelCase	)
         self.assertIsNotNone(UpperCAmelCase	)


def        SCREAMING_SNAKE_CASE      (       )			->     Optional[Any]:
    lowerCamelCase__    :      Union[str, Any]					 =		Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png' )
    return image


@require_tf
@require_vision
class    lowerCAmelCase							(					unittest.TestCase			):
 @cached_property
 def      A_	(      self  :      Optional[int]	)			->	List[Any]:
     return (
         DeiTImageProcessor.from_pretrained('facebook/deit-base-distilled-patch16-224'	)
         if is_vision_available()
         else None
     )







 @slow
 def      A_	(      self  :      List[Any]	)			->	List[str]:
     lowerCamelCase__    :      Optional[Any]					 =		TFDeiTForImageClassificationWithTeacher.from_pretrained('facebook/deit-base-distilled-patch16-224'	)

     lowerCamelCase__    :      int					 =		self.default_image_processor
     lowerCamelCase__    :      Dict					 =		prepare_img()
     lowerCamelCase__    :      Optional[Any]					 =		image_processor(images=UpperCAmelCase			,     return_tensors='tf'	)

     # forward pass
     lowerCamelCase__    :      int					 =		model(**UpperCAmelCase	)

     # verify the logits
     lowerCamelCase__    :      int					 =		tf.TensorShape((1, 1000)	)
     self.assertEqual(outputs.logits.shape			,     UpperCAmelCase	)

     lowerCamelCase__    :      Union[str, Any]					 =		tf.constant([-1.0_2_6_6, 0.1_9_1_2, -1.2_8_6_1]	)

     self.assertTrue(np.allclose(outputs.logits[0, :3]			,     UpperCAmelCase			,     atol=1e-4	)	)












from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
from .unet_ad_blocks import get_down_block, get_mid_block, get_out_block, get_up_block

@dataclass
class    lowerCAmelCase							(					__UpperCamelCase			):
 UpperCAmelCase__													=		42


class    lowerCAmelCase							(					__UpperCamelCase,  __UpperCamelCase			):
 @register_to_config
 def __init__(      self  :      List[str]			,     UpperCAmelCase  :      int = 65536			,     UpperCAmelCase  :      Optional[int] = None			,     UpperCAmelCase  :      int = 2			,     UpperCAmelCase  :      int = 2			,     UpperCAmelCase  :      int = 0			,     UpperCAmelCase  :      str = "fourier"			,     UpperCAmelCase  :      bool = True			,     UpperCAmelCase  :      bool = False			,     UpperCAmelCase  :      float = 0.0			,     UpperCAmelCase  :      Tuple[str] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")			,     UpperCAmelCase  :      Tuple[str] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")			,     UpperCAmelCase  :      Tuple[str] = "UNetMidBlock1D"			,     UpperCAmelCase  :      str = None			,     UpperCAmelCase  :      Tuple[int] = (32, 32, 64)			,     UpperCAmelCase  :      str = None			,     UpperCAmelCase  :      int = 8			,     UpperCAmelCase  :      int = 1			,     UpperCAmelCase  :      bool = False			,     )			->	List[Any]:
     super().__init__()
     lowerCamelCase__    :      Optional[int]					 =		sample_size

     # time
     if time_embedding_type == "fourier":
         lowerCamelCase__    :      Optional[Any]					 =		GaussianFourierProjection(
             embedding_size=8			,     set_W_to_weight=UpperCAmelCase			,     log=UpperCAmelCase			,     flip_sin_to_cos=UpperCAmelCase	)
         lowerCamelCase__    :      Any					 =		2 * block_out_channels[0]
     elif time_embedding_type == "positional":
         lowerCamelCase__    :      List[Any]					 =		Timesteps(
             block_out_channels[0]			,     flip_sin_to_cos=UpperCAmelCase			,     downscale_freq_shift=UpperCAmelCase	)
         lowerCamelCase__    :      Dict					 =		block_out_channels[0]

     if use_timestep_embedding:
         lowerCamelCase__    :      str					 =		block_out_channels[0] * 4
         lowerCamelCase__    :      List[Any]					 =		TimestepEmbedding(
             in_channels=UpperCAmelCase			,     time_embed_dim=UpperCAmelCase			,     act_fn=UpperCAmelCase			,     out_dim=block_out_channels[0]			,     )

     lowerCamelCase__    :      Any					 =		nn.ModuleList([]	)
     lowerCamelCase__    :      Tuple					 =		None
     lowerCamelCase__    :      List[str]					 =		nn.ModuleList([]	)
     lowerCamelCase__    :      Optional[int]					 =		None

     # down
     lowerCamelCase__    :      Optional[int]					 =		in_channels
     for i, down_block_type in enumerate(UpperCAmelCase	):
         lowerCamelCase__    :      Union[str, Any]					 =		output_channel
         lowerCamelCase__    :      Tuple					 =		block_out_channels[i]

         if i == 0:
             input_channel += extra_in_channels

         lowerCamelCase__    :      Union[str, Any]					 =		i == len(UpperCAmelCase	) - 1

         lowerCamelCase__    :      Optional[int]					 =		get_down_block(
             UpperCAmelCase			,     num_layers=UpperCAmelCase			,     in_channels=UpperCAmelCase			,     out_channels=UpperCAmelCase			,     temb_channels=block_out_channels[0]			,     add_downsample=not is_final_block or downsample_each_block			,     )
         self.down_blocks.append(UpperCAmelCase	)

     # mid
     lowerCamelCase__    :      Optional[int]					 =		get_mid_block(
         UpperCAmelCase			,     in_channels=block_out_channels[-1]			,     mid_channels=block_out_channels[-1]			,     out_channels=block_out_channels[-1]			,     embed_dim=block_out_channels[0]			,     num_layers=UpperCAmelCase			,     add_downsample=UpperCAmelCase			,     )

     # up
     lowerCamelCase__    :      Optional[int]					 =		list(reversed(UpperCAmelCase	)	)
     lowerCamelCase__    :      Optional[int]					 =		reversed_block_out_channels[0]
     if out_block_type is None:
         lowerCamelCase__    :      List[str]					 =		out_channels
     else:
         lowerCamelCase__    :      Any					 =		block_out_channels[0]

     for i, up_block_type in enumerate(UpperCAmelCase	):
         lowerCamelCase__    :      Tuple					 =		output_channel
         lowerCamelCase__    :      Union[str, Any]					 =		(
             reversed_block_out_channels[i + 1] if i < len(UpperCAmelCase	) - 1 else final_upsample_channels
         )

         lowerCamelCase__    :      List[str]					 =		i == len(UpperCAmelCase	) - 1

         lowerCamelCase__    :      Dict					 =		get_up_block(
             UpperCAmelCase			,     num_layers=UpperCAmelCase			,     in_channels=UpperCAmelCase			,     out_channels=UpperCAmelCase			,     temb_channels=block_out_channels[0]			,     add_upsample=not is_final_block			,     )
         self.up_blocks.append(UpperCAmelCase	)
         lowerCamelCase__    :      int					 =		output_channel

     # out
     lowerCamelCase__    :      int					 =		norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4			,     32	)
     lowerCamelCase__    :      List[Any]					 =		get_out_block(
         out_block_type=UpperCAmelCase			,     num_groups_out=UpperCAmelCase			,     embed_dim=block_out_channels[0]			,     out_channels=UpperCAmelCase			,     act_fn=UpperCAmelCase			,     fc_dim=block_out_channels[-1] // 4			,     )







 def      A_	(      self  :      List[Any]			,     UpperCAmelCase  :      torch.FloatTensor			,     UpperCAmelCase  :      Union[torch.Tensor, float, int]			,     UpperCAmelCase  :      bool = True			,     )			->	Union[UNetaDOutput, Tuple]:
     lowerCamelCase__    :      Optional[Any]					 =		timestep
     if not torch.is_tensor(UpperCAmelCase	):
         lowerCamelCase__    :      Optional[int]					 =		torch.tensor([timesteps]			,     dtype=torch.long			,     device=sample.device	)
     elif torch.is_tensor(UpperCAmelCase	) and len(timesteps.shape	) == 0:
         lowerCamelCase__    :      List[str]					 =		timesteps[None].to(sample.device	)

     lowerCamelCase__    :      Optional[int]					 =		self.time_proj(UpperCAmelCase	)
     if self.config.use_timestep_embedding:
         lowerCamelCase__    :      str					 =		self.time_mlp(UpperCAmelCase	)
     else:
         lowerCamelCase__    :      List[str]					 =		timestep_embed[..., None]
         lowerCamelCase__    :      str					 =		timestep_embed.repeat([1, 1, sample.shape[2]]	).to(sample.dtype	)
         lowerCamelCase__    :      str					 =		timestep_embed.broadcast_to((sample.shape[:1] + timestep_embed.shape[1:])	)

     # 2. down
     lowerCamelCase__    :      str					 =		()
     for downsample_block in self.down_blocks:
         lowerCamelCase__ ,							lowerCamelCase__    :      Union[str, Any]					 =		downsample_block(hidden_states=UpperCAmelCase			,     temb=UpperCAmelCase	)
         down_block_res_samples += res_samples

     # 3. mid
     if self.mid_block:
         lowerCamelCase__    :      Optional[Any]					 =		self.mid_block(UpperCAmelCase			,     UpperCAmelCase	)

     # 4. up
     for i, upsample_block in enumerate(self.up_blocks	):
         lowerCamelCase__    :      Dict					 =		down_block_res_samples[-1:]
         lowerCamelCase__    :      Optional[Any]					 =		down_block_res_samples[:-1]
         lowerCamelCase__    :      Any					 =		upsample_block(UpperCAmelCase			,     res_hidden_states_tuple=UpperCAmelCase			,     temb=UpperCAmelCase	)

     # 5. post-process
     if self.out_block:
         lowerCamelCase__    :      Any					 =		self.out_block(UpperCAmelCase			,     UpperCAmelCase	)

     if not return_dict:
         return (sample,)

     return UNetaDOutput(sample=UpperCAmelCase	)





