



import collections
import inspect
import unittest

from transformers import FocalNetConfig
from transformers.testing_utils import 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 (
	    FocalNetBackbone,
	    FocalNetForImageClassification,
	    FocalNetForMaskedImageModeling,
	    FocalNetModel,
	)
	from transformers.models.focalnet.modeling_focalnet import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
	from PIL import Image

	from transformers import AutoImageProcessor


class 	__snake_case   :
		def __init__(							self    ,a_    ,a_=13    ,a_=32    ,a_=2    ,a_=3    ,a_=16    ,a_=[32, 64, 128]    ,a_=[1, 2, 1]    ,a_=[2, 2, 4]    ,a_=2    ,a_=2.0    ,a_=True    ,a_=0.0    ,a_=0.0    ,a_=0.1    ,a_="gelu"    ,a_=False    ,a_=True    ,a_=0.02    ,a_=1e-5    ,a_=True    ,a_=None    ,a_=True    ,a_=10    ,a_=8    ,a_=["stage1", "stage2"]    ,a_=[1, 2]    ,):







						"""simple docstring"""
						lowerCAmelCase__		  =  parent
						lowerCAmelCase__		  =  batch_size
						lowerCAmelCase__		  =  image_size
						lowerCAmelCase__		  =  patch_size
						lowerCAmelCase__		  =  num_channels
						lowerCAmelCase__		  =  embed_dim
						lowerCAmelCase__		  =  hidden_sizes
						lowerCAmelCase__		  =  depths
						lowerCAmelCase__		  =  num_heads
						lowerCAmelCase__		  =  window_size
						lowerCAmelCase__		  =  mlp_ratio
						lowerCAmelCase__		  =  qkv_bias
						lowerCAmelCase__		  =  hidden_dropout_prob
						lowerCAmelCase__		  =  attention_probs_dropout_prob
						lowerCAmelCase__		  =  drop_path_rate
						lowerCAmelCase__		  =  hidden_act
						lowerCAmelCase__		  =  use_absolute_embeddings
						lowerCAmelCase__		  =  patch_norm
						lowerCAmelCase__		  =  layer_norm_eps
						lowerCAmelCase__		  =  initializer_range
						lowerCAmelCase__		  =  is_training
						lowerCAmelCase__		  =  scope
						lowerCAmelCase__		  =  use_labels
						lowerCAmelCase__		  =  type_sequence_label_size
						lowerCAmelCase__		  =  encoder_stride
						lowerCAmelCase__		  =  out_features
						lowerCAmelCase__		  =  out_indices
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

						lowerCAmelCase__		  =  None
						if self.use_labels:
										lowerCAmelCase__		  =  ids_tensor([self.batch_size]    ,self.type_sequence_label_size						)

						lowerCAmelCase__		  =  self.get_config()

						return config, pixel_values, labels
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						return FocalNetConfig(
						    image_size=self.image_size    ,patch_size=self.patch_size    ,num_channels=self.num_channels    ,embed_dim=self.embed_dim    ,hidden_sizes=self.hidden_sizes    ,depths=self.depths    ,num_heads=self.num_heads    ,window_size=self.window_size    ,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    ,use_absolute_embeddings=self.use_absolute_embeddings    ,path_norm=self.patch_norm    ,layer_norm_eps=self.layer_norm_eps    ,initializer_range=self.initializer_range    ,encoder_stride=self.encoder_stride    ,out_features=self.out_features    ,out_indices=self.out_indices    ,)
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  FocalNetModel(config=a_						)
						model.to(a_						)
						model.eval()
						lowerCAmelCase__		  =  model(a_						)

						lowerCAmelCase__		  =  ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths						) - 1))
						lowerCAmelCase__		  =  int(config.embed_dim * 2 ** (len(config.depths						) - 1)						)

						self.parent.assertEqual(result.last_hidden_state.shape    ,(self.batch_size, expected_seq_len, expected_dim)						)
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  FocalNetBackbone(config=a_						)
						model.to(a_						)
						model.eval()
						lowerCAmelCase__		  =  model(a_						)

						# verify feature maps
						self.parent.assertEqual(len(result.feature_maps						)    ,len(config.out_features						)						)
						self.parent.assertListEqual(list(result.feature_maps[0].shape						)    ,[self.batch_size, self.image_size, 8, 8]						)

						# verify channels
						self.parent.assertEqual(len(model.channels						)    ,len(config.out_features						)						)
						self.parent.assertListEqual(model.channels    ,config.hidden_sizes[:-1]						)

						# verify backbone works with out_features=None
						lowerCAmelCase__		  =  None
						lowerCAmelCase__		  =  FocalNetBackbone(config=a_						)
						model.to(a_						)
						model.eval()
						lowerCAmelCase__		  =  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, self.image_size * 2, 4, 4]						)

						# verify channels
						self.parent.assertEqual(len(model.channels						)    ,1						)
						self.parent.assertListEqual(model.channels    ,[config.hidden_sizes[-1]]						)
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  FocalNetForMaskedImageModeling(config=a_						)
						model.to(a_						)
						model.eval()
						lowerCAmelCase__		  =  model(a_						)
						self.parent.assertEqual(
						    result.reconstruction.shape    ,(self.batch_size, self.num_channels, self.image_size, self.image_size)						)

						# test greyscale images
						lowerCAmelCase__		  =  1
						lowerCAmelCase__		  =  FocalNetForMaskedImageModeling(a_						)
						model.to(a_						)
						model.eval()

						lowerCAmelCase__		  =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
						lowerCAmelCase__		  =  model(a_						)
						self.parent.assertEqual(result.reconstruction.shape    ,(self.batch_size, 1, self.image_size, self.image_size)						)
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.type_sequence_label_size
						lowerCAmelCase__		  =  FocalNetForImageClassification(a_						)
						model.to(a_						)
						model.eval()
						lowerCAmelCase__		  =  model(a_    ,labels=a_						)
						self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.type_sequence_label_size)						)

						# test greyscale images
						lowerCAmelCase__		  =  1
						lowerCAmelCase__		  =  FocalNetForImageClassification(a_						)
						model.to(a_						)
						model.eval()

						lowerCAmelCase__		  =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
						lowerCAmelCase__		  =  model(a_						)
						self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.type_sequence_label_size)						)



		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.prepare_config_and_inputs()

						lowerCAmelCase__  ,   lowerCAmelCase__  ,   lowerCAmelCase__		  =  config_and_inputs
						lowerCAmelCase__		  =  {'pixel_values': pixel_values}
						return config, inputs_dict



@require_torch
class 	__snake_case   ( SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE							, unittest.TestCase ):
		SCREAMING_SNAKE_CASE__								=	(
		    (
		        FocalNetModel,
		        FocalNetForImageClassification,
		        FocalNetForMaskedImageModeling,
		        FocalNetBackbone,
		    )
		    if is_torch_available()
		    else ()
		)
		SCREAMING_SNAKE_CASE__								=	(
		    {'feature-extraction': FocalNetModel, 'image-classification': FocalNetForImageClassification}
		    if is_torch_available()
		    else {}
		)
		SCREAMING_SNAKE_CASE__								=	False

		SCREAMING_SNAKE_CASE__								=	False
		SCREAMING_SNAKE_CASE__								=	False
		SCREAMING_SNAKE_CASE__								=	False
		SCREAMING_SNAKE_CASE__								=	False
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  FocalNetModelTester(self						)
						lowerCAmelCase__		  =  ConfigTester(self    ,config_class=a_    ,embed_dim=37    ,has_text_modality=a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						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 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						return
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_backbone(*a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_image_modeling(*a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*a_						)
		@unittest.skip(reason='FocalNet does not use inputs_embeds'						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						pass
		@unittest.skip(reason='FocalNet does not use feedforward chunking'						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						pass
		def 						SCREAMING_SNAKE_CASE_					(							self						):







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

						for model_class in self.all_model_classes[:-1]:
										lowerCAmelCase__		  =  model_class(a_						)
										self.assertIsInstance(model.get_input_embeddings()    ,(nn.Module)						)
										lowerCAmelCase__		  =  model.get_output_embeddings()
										self.assertTrue(x is None or isinstance(a_    ,nn.Linear						)						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







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

						for model_class in self.all_model_classes[:-1]:
										lowerCAmelCase__		  =  model_class(a_						)
										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]    ,a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  model_class(a_						)
						model.to(a_						)
						model.eval()

						with torch.no_grad():
										lowerCAmelCase__		  =  model(**self._prepare_for_class(a_    ,a_						)						)

						lowerCAmelCase__		  =  outputs.hidden_states

						lowerCAmelCase__		  =  getattr(
						    self.model_tester    ,'expected_num_hidden_layers'    ,len(self.model_tester.depths						) + 1						)
						self.assertEqual(len(a_						)    ,a_						)

						# FocalNet has a different seq_length
						lowerCAmelCase__		  =  (
						    config.patch_size
						    if isinstance(config.patch_size    ,collections.abc.Iterable						)
						    else (config.patch_size, config.patch_size)
						)

						lowerCAmelCase__		  =  (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

						self.assertListEqual(
						    list(hidden_states[0].shape[-2:]						)    ,[num_patches, self.model_tester.embed_dim]    ,)

						lowerCAmelCase__		  =  outputs.reshaped_hidden_states
						self.assertEqual(len(a_						)    ,a_						)

						lowerCAmelCase__  ,   lowerCAmelCase__  ,   lowerCAmelCase__  ,   lowerCAmelCase__		  =  reshaped_hidden_states[0].shape
						lowerCAmelCase__		  =  (
						    reshaped_hidden_states[0].view(a_    ,a_    ,height * width						).permute(0    ,2    ,1						)
						)
						self.assertListEqual(
						    list(reshaped_hidden_states.shape[-2:]						)    ,[num_patches, self.model_tester.embed_dim]    ,)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







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

						lowerCAmelCase__		  =  (
						    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[:-1]:
										lowerCAmelCase__		  =  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"]
										lowerCAmelCase__		  =  True

										self.check_hidden_states_output(a_    ,a_    ,a_    ,a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__  ,   lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs_for_common()
						lowerCAmelCase__		  =  3

						lowerCAmelCase__		  =  (
						    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)
						)
						lowerCAmelCase__		  =  (
						    config.patch_size
						    if isinstance(config.patch_size    ,collections.abc.Iterable						)
						    else (config.patch_size, config.patch_size)
						)

						lowerCAmelCase__		  =  image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
						lowerCAmelCase__		  =  image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

						for model_class in self.all_model_classes[:-1]:
										lowerCAmelCase__		  =  True
										self.check_hidden_states_output(a_    ,a_    ,a_    ,(padded_height, padded_width)						)

										# check that output_hidden_states also work using config
										del inputs_dict["output_hidden_states"]
										lowerCAmelCase__		  =  True
										self.check_hidden_states_output(a_    ,a_    ,a_    ,(padded_height, padded_width)						)
		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCAmelCase__		  =  FocalNetModel.from_pretrained(a_						)
										self.assertIsNotNone(a_						)



		def 						SCREAMING_SNAKE_CASE_					(							self						):







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

						lowerCAmelCase__		  =  _config_zero_init(a_						)
						for model_class in self.all_model_classes:
										lowerCAmelCase__		  =  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_vision
@require_torch
class 	__snake_case   ( unittest.TestCase ):
		@cached_property
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						# TODO update organization
						return AutoImageProcessor.from_pretrained('microsoft/focalnet-tiny'						) if is_vision_available() else None



		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  FocalNetForImageClassification.from_pretrained('microsoft/focalnet-tiny'						).to(a_						)
						lowerCAmelCase__		  =  self.default_image_processor

						lowerCAmelCase__		  =  Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'						)
						lowerCAmelCase__		  =  image_processor(images=a_    ,return_tensors='pt'						).to(a_						)

						# forward pass
						with torch.no_grad():
										lowerCAmelCase__		  =  model(**a_						)

						# verify the logits
						lowerCAmelCase__		  =  torch.Size((1, 1000)						)
						self.assertEqual(outputs.logits.shape    ,a_						)
						lowerCAmelCase__		  =  torch.tensor([0.2166, -0.4368, 0.2191]						).to(a_						)
						self.assertTrue(torch.allclose(outputs.logits[0, :3]    ,a_    ,atol=1e-4						)						)
						self.assertTrue(outputs.logits.argmax(dim=-1						).item()    ,281						)



@require_torch
class 	__snake_case   ( SCREAMING_SNAKE_CASE							, unittest.TestCase ):
		SCREAMING_SNAKE_CASE__								=	(FocalNetBackbone,) if is_torch_available() else ()
		SCREAMING_SNAKE_CASE__								=	FocalNetConfig

		SCREAMING_SNAKE_CASE__								=	False
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  FocalNetModelTester(self						)





from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin


@dataclass
class 	__snake_case   ( SCREAMING_SNAKE_CASE ):
		SCREAMING_SNAKE_CASE__								=	42



class 	__snake_case   ( SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE ):
		@register_to_config
		def __init__(							self    ,a_ = 16    ,a_ = 88    ,a_ = None    ,a_ = None    ,a_ = 1    ,a_ = 0.0    ,a_ = 32    ,a_ = None    ,a_ = False    ,a_ = None    ,a_ = "geglu"    ,a_ = True    ,a_ = True    ,):







						"""simple docstring"""
						super().__init__()
						lowerCAmelCase__		  =  num_attention_heads
						lowerCAmelCase__		  =  attention_head_dim
						lowerCAmelCase__		  =  num_attention_heads * attention_head_dim

						lowerCAmelCase__		  =  in_channels

						lowerCAmelCase__		  =  torch.nn.GroupNorm(num_groups=a_    ,num_channels=a_    ,eps=1e-6    ,affine=a_						)
						lowerCAmelCase__		  =  nn.Linear(a_    ,a_						)

						# 3. Define transformers blocks
						lowerCAmelCase__		  =  nn.ModuleList(
						    [
						        BasicTransformerBlock(
						            a_    ,a_    ,a_    ,dropout=a_    ,cross_attention_dim=a_    ,activation_fn=a_    ,attention_bias=a_    ,double_self_attention=a_    ,norm_elementwise_affine=a_    ,)
						        for d in range(a_						)
						    ]						)

						lowerCAmelCase__		  =  nn.Linear(a_    ,a_						)



		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_=None    ,a_=None    ,a_=None    ,a_=1    ,a_=None    ,a_ = True    ,):







						"""simple docstring"""
						lowerCAmelCase__  ,   lowerCAmelCase__  ,   lowerCAmelCase__  ,   lowerCAmelCase__		  =  hidden_states.shape
						lowerCAmelCase__		  =  batch_frames // num_frames

						lowerCAmelCase__		  =  hidden_states

						lowerCAmelCase__		  =  hidden_states[None, :].reshape(a_    ,a_    ,a_    ,a_    ,a_						)
						lowerCAmelCase__		  =  hidden_states.permute(0    ,2    ,1    ,3    ,4						)

						lowerCAmelCase__		  =  self.norm(a_						)
						lowerCAmelCase__		  =  hidden_states.permute(0    ,3    ,4    ,2    ,1						).reshape(batch_size * height * width    ,a_    ,a_						)

						lowerCAmelCase__		  =  self.proj_in(a_						)

						# 2. Blocks
						for block in self.transformer_blocks:
										lowerCAmelCase__		  =  block(
										    a_    ,encoder_hidden_states=a_    ,timestep=a_    ,cross_attention_kwargs=a_    ,class_labels=a_    ,)

						# 3. Output
						lowerCAmelCase__		  =  self.proj_out(a_						)
						lowerCAmelCase__		  =  (
						    hidden_states[None, None, :]
						    .reshape(a_    ,a_    ,a_    ,a_    ,a_						)
						    .permute(0    ,3    ,4    ,1    ,2						)
						    .contiguous()
						)
						lowerCAmelCase__		  =  hidden_states.reshape(a_    ,a_    ,a_    ,a_						)

						lowerCAmelCase__		  =  hidden_states + residual

						if not return_dict:
										return (output,)

						return TransformerTemporalModelOutput(sample=a_						)
