



'''simple docstring'''

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 						a__  :


	def __init__(self			:			List[str],       __UpperCAmelCase			:			List[Any],       __UpperCAmelCase			:			int=13,       __UpperCAmelCase			:			List[str]=32,       __UpperCAmelCase			:			Tuple=2,       __UpperCAmelCase			:			Optional[Any]=3,       __UpperCAmelCase			:			int=16,       __UpperCAmelCase			:			int=[32, 64, 128],       __UpperCAmelCase			:			str=[1, 2, 1],       __UpperCAmelCase			:			Union[str, Any]=[2, 2, 4],       __UpperCAmelCase			:			Union[str, Any]=2,       __UpperCAmelCase			:			int=2.0,       __UpperCAmelCase			:			Tuple=True,       __UpperCAmelCase			:			Optional[int]=0.0,       __UpperCAmelCase			:			List[str]=0.0,       __UpperCAmelCase			:			List[Any]=0.1,       __UpperCAmelCase			:			List[Any]="gelu",       __UpperCAmelCase			:			Optional[int]=False,       __UpperCAmelCase			:			Tuple=True,       __UpperCAmelCase			:			Tuple=0.02,       __UpperCAmelCase			:			List[Any]=1e-5,       __UpperCAmelCase			:			Dict=True,       __UpperCAmelCase			:			List[Any]=None,       __UpperCAmelCase			:			Tuple=True,       __UpperCAmelCase			:			Optional[int]=10,       __UpperCAmelCase			:			Dict=8,       __UpperCAmelCase			:			List[Any]=["stage1", "stage2"],       __UpperCAmelCase			:			Optional[Any]=[1, 2],       )  ->      Tuple:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    Dict          =			parent
					SCREAMING_SNAKE_CASE	:    str          =			batch_size
					SCREAMING_SNAKE_CASE	:    Tuple          =			image_size
					SCREAMING_SNAKE_CASE	:    List[Any]          =			patch_size
					SCREAMING_SNAKE_CASE	:    Any          =			num_channels
					SCREAMING_SNAKE_CASE	:    List[Any]          =			embed_dim
					SCREAMING_SNAKE_CASE	:    Optional[int]          =			hidden_sizes
					SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			depths
					SCREAMING_SNAKE_CASE	:    Tuple          =			num_heads
					SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			window_size
					SCREAMING_SNAKE_CASE	:    str          =			mlp_ratio
					SCREAMING_SNAKE_CASE	:    Any          =			qkv_bias
					SCREAMING_SNAKE_CASE	:    Tuple          =			hidden_dropout_prob
					SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			attention_probs_dropout_prob
					SCREAMING_SNAKE_CASE	:    int          =			drop_path_rate
					SCREAMING_SNAKE_CASE	:    Any          =			hidden_act
					SCREAMING_SNAKE_CASE	:    str          =			use_absolute_embeddings
					SCREAMING_SNAKE_CASE	:    Tuple          =			patch_norm
					SCREAMING_SNAKE_CASE	:    Tuple          =			layer_norm_eps
					SCREAMING_SNAKE_CASE	:    List[Any]          =			initializer_range
					SCREAMING_SNAKE_CASE	:    List[str]          =			is_training
					SCREAMING_SNAKE_CASE	:    Tuple          =			scope
					SCREAMING_SNAKE_CASE	:    str          =			use_labels
					SCREAMING_SNAKE_CASE	:    Optional[int]          =			type_sequence_label_size
					SCREAMING_SNAKE_CASE	:    Optional[int]          =			encoder_stride
					SCREAMING_SNAKE_CASE	:    Tuple          =			out_features
					SCREAMING_SNAKE_CASE	:    Any          =			out_indices


	def 					lowercase__	(self			:			Optional[Any]   )  ->      str:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    Optional[Any]          =			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

					SCREAMING_SNAKE_CASE	:    int          =			None
					if self.use_labels:
									SCREAMING_SNAKE_CASE	:    Dict          =			ids_tensor([self.batch_size],       self.type_sequence_label_size   )

					SCREAMING_SNAKE_CASE	:    Tuple          =			self.get_config()

					return config, pixel_values, labels


	def 					lowercase__	(self			:			Tuple   )  ->      int:

					"""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 					lowercase__	(self			:			Any,       __UpperCAmelCase			:			str,       __UpperCAmelCase			:			str,       __UpperCAmelCase			:			int   )  ->      Optional[Any]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    str          =			FocalNetModel(config=lowerCamelCase__   )
					model.to(lowerCamelCase__   )
					model.eval()
					SCREAMING_SNAKE_CASE	:    int          =			model(lowerCamelCase__   )

					SCREAMING_SNAKE_CASE	:    Dict          =			((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths   ) - 1))
					SCREAMING_SNAKE_CASE	:    List[str]          =			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 					lowercase__	(self			:			List[Any],       __UpperCAmelCase			:			Tuple,       __UpperCAmelCase			:			Optional[int],       __UpperCAmelCase			:			Dict   )  ->      List[str]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    Any          =			FocalNetBackbone(config=lowerCamelCase__   )
					model.to(lowerCamelCase__   )
					model.eval()
					SCREAMING_SNAKE_CASE	:    Optional[int]          =			model(lowerCamelCase__   )

					# 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
					SCREAMING_SNAKE_CASE	:    str          =			None
					SCREAMING_SNAKE_CASE	:    str          =			FocalNetBackbone(config=lowerCamelCase__   )
					model.to(lowerCamelCase__   )
					model.eval()
					SCREAMING_SNAKE_CASE	:    str          =			model(lowerCamelCase__   )

					# 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 					lowercase__	(self			:			Optional[int],       __UpperCAmelCase			:			int,       __UpperCAmelCase			:			Any,       __UpperCAmelCase			:			str   )  ->      str:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    Dict          =			FocalNetForMaskedImageModeling(config=lowerCamelCase__   )
					model.to(lowerCamelCase__   )
					model.eval()
					SCREAMING_SNAKE_CASE	:    List[str]          =			model(lowerCamelCase__   )
					self.parent.assertEqual(
					    result.reconstruction.shape,       (self.batch_size, self.num_channels, self.image_size, self.image_size)   )

					# test greyscale images
					SCREAMING_SNAKE_CASE	:    Dict          =			1
					SCREAMING_SNAKE_CASE	:    Optional[Any]          =			FocalNetForMaskedImageModeling(lowerCamelCase__   )
					model.to(lowerCamelCase__   )
					model.eval()

					SCREAMING_SNAKE_CASE	:    Optional[int]          =			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
					SCREAMING_SNAKE_CASE	:    int          =			model(lowerCamelCase__   )
					self.parent.assertEqual(result.reconstruction.shape,       (self.batch_size, 1, self.image_size, self.image_size)   )


	def 					lowercase__	(self			:			str,       __UpperCAmelCase			:			Union[str, Any],       __UpperCAmelCase			:			List[Any],       __UpperCAmelCase			:			Dict   )  ->      int:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    int          =			self.type_sequence_label_size
					SCREAMING_SNAKE_CASE	:    Dict          =			FocalNetForImageClassification(lowerCamelCase__   )
					model.to(lowerCamelCase__   )
					model.eval()
					SCREAMING_SNAKE_CASE	:    Optional[int]          =			model(lowerCamelCase__,       labels=lowerCamelCase__   )
					self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.type_sequence_label_size)   )

					# test greyscale images
					SCREAMING_SNAKE_CASE	:    List[str]          =			1
					SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			FocalNetForImageClassification(lowerCamelCase__   )
					model.to(lowerCamelCase__   )
					model.eval()

					SCREAMING_SNAKE_CASE	:    List[Any]          =			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
					SCREAMING_SNAKE_CASE	:    str          =			model(lowerCamelCase__   )
					self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.type_sequence_label_size)   )






	def 					lowercase__	(self			:			Union[str, Any]   )  ->      Any:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    List[str]          =			self.prepare_config_and_inputs()

					SCREAMING_SNAKE_CASE	:    Tuple          =			config_and_inputs
					SCREAMING_SNAKE_CASE	:    int          =			{"pixel_values": pixel_values}
					return config, inputs_dict



@require_torch
class 						a__  (					UpperCAmelCase__,      UpperCAmelCase__,      unittest.TestCase						):
	__magic_name__		:     List[Any]						     =  (
	    (
	        FocalNetModel,
	        FocalNetForImageClassification,
	        FocalNetForMaskedImageModeling,
	        FocalNetBackbone,
	    )
	    if is_torch_available()
	    else ()
	)
	__magic_name__		:     Optional[Any]						     =  (
	    {'feature-extraction': FocalNetModel, 'image-classification': FocalNetForImageClassification}
	    if is_torch_available()
	    else {}
	)
	__magic_name__		:     int						     =  False

	__magic_name__		:     Optional[int]						     =  False
	__magic_name__		:     Any						     =  False
	__magic_name__		:     List[Any]						     =  False
	__magic_name__		:     Dict						     =  False


	def 					lowercase__	(self			:			int   )  ->      List[Any]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    int          =			FocalNetModelTester(self   )
					SCREAMING_SNAKE_CASE	:    str          =			ConfigTester(self,       config_class=lowerCamelCase__,       embed_dim=37,       has_text_modality=lowerCamelCase__   )


	def 					lowercase__	(self			:			List[Any]   )  ->      Optional[int]:

					"""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 					lowercase__	(self			:			str   )  ->      List[str]:

					"""simple docstring"""

					return


	def 					lowercase__	(self			:			Optional[int]   )  ->      Tuple:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    List[str]          =			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*lowerCamelCase__   )


	def 					lowercase__	(self			:			str   )  ->      Dict:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    str          =			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_backbone(*lowerCamelCase__   )


	def 					lowercase__	(self			:			str   )  ->      Tuple:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    List[str]          =			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_image_modeling(*lowerCamelCase__   )


	def 					lowercase__	(self			:			List[Any]   )  ->      List[str]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    Optional[int]          =			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_image_classification(*lowerCamelCase__   )


	@unittest.skip(reason='''FocalNet does not use inputs_embeds'''   )
	def 					lowercase__	(self			:			Optional[int]   )  ->      Optional[Any]:

					"""simple docstring"""

					pass


	@unittest.skip(reason='''FocalNet does not use feedforward chunking'''   )
	def 					lowercase__	(self			:			Optional[Any]   )  ->      int:

					"""simple docstring"""

					pass


	def 					lowercase__	(self			:			Tuple   )  ->      str:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    str          =			self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes[:-1]:
									SCREAMING_SNAKE_CASE	:    int          =			model_class(lowerCamelCase__   )
									self.assertIsInstance(model.get_input_embeddings(),       (nn.Module)   )
									SCREAMING_SNAKE_CASE	:    List[Any]          =			model.get_output_embeddings()
									self.assertTrue(x is None or isinstance(lowerCamelCase__,       nn.Linear   )   )


	def 					lowercase__	(self			:			Any   )  ->      int:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    Optional[Any]          =			self.model_tester.prepare_config_and_inputs_for_common()

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

									SCREAMING_SNAKE_CASE	:    Optional[Any]          =			["pixel_values"]
									self.assertListEqual(arg_names[:1],       lowerCamelCase__   )


	def 					lowercase__	(self			:			Optional[Any],       __UpperCAmelCase			:			Optional[int],       __UpperCAmelCase			:			str,       __UpperCAmelCase			:			Optional[int],       __UpperCAmelCase			:			Optional[Any]   )  ->      Any:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    str          =			model_class(lowerCamelCase__   )
					model.to(lowerCamelCase__   )
					model.eval()

					with torch.no_grad():
									SCREAMING_SNAKE_CASE	:    List[str]          =			model(**self._prepare_for_class(lowerCamelCase__,       lowerCamelCase__   )   )

					SCREAMING_SNAKE_CASE	:    Tuple          =			outputs.hidden_states

					SCREAMING_SNAKE_CASE	:    Optional[int]          =			getattr(
					    self.model_tester,       '''expected_num_hidden_layers''',       len(self.model_tester.depths   ) + 1   )
					self.assertEqual(len(lowerCamelCase__   ),       lowerCamelCase__   )

					# FocalNet has a different seq_length
					SCREAMING_SNAKE_CASE	:    List[Any]          =			(
					    config.patch_size
					    if isinstance(config.patch_size,       collections.abc.Iterable   )
					    else (config.patch_size, config.patch_size)
					)

					SCREAMING_SNAKE_CASE	:    int          =			(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],       )

					SCREAMING_SNAKE_CASE	:    Optional[int]          =			outputs.reshaped_hidden_states
					self.assertEqual(len(lowerCamelCase__   ),       lowerCamelCase__   )

					SCREAMING_SNAKE_CASE	:    List[str]          =			reshaped_hidden_states[0].shape
					SCREAMING_SNAKE_CASE	:    Optional[Any]          =			(
					    reshaped_hidden_states[0].view(lowerCamelCase__,       lowerCamelCase__,       height * width   ).permute(0,       2,       1   )
					)
					self.assertListEqual(
					    list(reshaped_hidden_states.shape[-2:]   ),       [num_patches, self.model_tester.embed_dim],       )


	def 					lowercase__	(self			:			int   )  ->      Optional[Any]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    Any          =			self.model_tester.prepare_config_and_inputs_for_common()

					SCREAMING_SNAKE_CASE	:    Dict          =			(
					    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]:
									SCREAMING_SNAKE_CASE	:    str          =			True
									self.check_hidden_states_output(lowerCamelCase__,       lowerCamelCase__,       lowerCamelCase__,       lowerCamelCase__   )

									# check that output_hidden_states also work using config
									del inputs_dict["output_hidden_states"]
									SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			True

									self.check_hidden_states_output(lowerCamelCase__,       lowerCamelCase__,       lowerCamelCase__,       lowerCamelCase__   )


	def 					lowercase__	(self			:			Any   )  ->      Tuple:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    Tuple          =			self.model_tester.prepare_config_and_inputs_for_common()
					SCREAMING_SNAKE_CASE	:    Any          =			3

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

					SCREAMING_SNAKE_CASE	:    List[str]          =			image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
					SCREAMING_SNAKE_CASE	:    int          =			image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

					for model_class in self.all_model_classes[:-1]:
									SCREAMING_SNAKE_CASE	:    Dict          =			True
									self.check_hidden_states_output(lowerCamelCase__,       lowerCamelCase__,       lowerCamelCase__,       (padded_height, padded_width)   )

									# check that output_hidden_states also work using config
									del inputs_dict["output_hidden_states"]
									SCREAMING_SNAKE_CASE	:    int          =			True
									self.check_hidden_states_output(lowerCamelCase__,       lowerCamelCase__,       lowerCamelCase__,       (padded_height, padded_width)   )


	@slow
	def 					lowercase__	(self			:			List[Any]   )  ->      Union[str, Any]:

					"""simple docstring"""

					for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									SCREAMING_SNAKE_CASE	:    List[Any]          =			FocalNetModel.from_pretrained(lowerCamelCase__   )
									self.assertIsNotNone(lowerCamelCase__   )






	def 					lowercase__	(self			:			Union[str, Any]   )  ->      Tuple:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    List[Any]          =			self.model_tester.prepare_config_and_inputs_for_common()

					SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			_config_zero_init(lowerCamelCase__   )
					for model_class in self.all_model_classes:
									SCREAMING_SNAKE_CASE	:    List[str]          =			model_class(config=lowerCamelCase__   )
									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 						a__  (					unittest.TestCase						):


	@cached_property
	def 					lowercase__	(self			:			Any   )  ->      Tuple:

					"""simple docstring"""

					return AutoImageProcessor.from_pretrained('''microsoft/focalnet-tiny'''   ) if is_vision_available() else None






	@slow
	def 					lowercase__	(self			:			str   )  ->      str:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    Tuple          =			FocalNetForImageClassification.from_pretrained('''microsoft/focalnet-tiny'''   ).to(lowerCamelCase__   )
					SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			self.default_image_processor

					SCREAMING_SNAKE_CASE	:    Any          =			Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''   )
					SCREAMING_SNAKE_CASE	:    int          =			image_processor(images=lowerCamelCase__,       return_tensors='''pt'''   ).to(lowerCamelCase__   )

					# forward pass
					with torch.no_grad():
									SCREAMING_SNAKE_CASE	:    Optional[Any]          =			model(**lowerCamelCase__   )

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



@require_torch
class 						a__  (					UpperCAmelCase__,      unittest.TestCase						):
	__magic_name__		:     Any						     =  (FocalNetBackbone,) if is_torch_available() else ()
	__magic_name__		:     Union[str, Any]						     =  FocalNetConfig

	__magic_name__		:     Any						     =  False


	def 					lowercase__	(self			:			Union[str, Any]   )  ->      Optional[int]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE	:    Any          =			FocalNetModelTester(self   )





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 					_lowercase    (     UpperCAmelCase__			):






    '''simple docstring'''
    SCREAMING_SNAKE_CASE:				torch.FloatTensor





class 					_lowercase    (     UpperCAmelCase__ ,  UpperCAmelCase__			):






    '''simple docstring'''




    @register_to_config
    def __init__(    self       ,						lowerCamelCase__ = 65_536       ,						lowerCamelCase__ = None       ,						lowerCamelCase__ = 2       ,						lowerCamelCase__ = 2       ,						lowerCamelCase__ = 0       ,						lowerCamelCase__ = "fourier"       ,						lowerCamelCase__ = True       ,						lowerCamelCase__ = False       ,						lowerCamelCase__ = 0.0       ,						lowerCamelCase__ = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")       ,						lowerCamelCase__ = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")       ,						lowerCamelCase__ = "UNetMidBlock1D"       ,						lowerCamelCase__ = None       ,						lowerCamelCase__ = (32, 32, 64)       ,						lowerCamelCase__ = None       ,						lowerCamelCase__ = 8       ,						lowerCamelCase__ = 1       ,						lowerCamelCase__ = False       ,						):
       super().__init__()
       lowerCAmelCase_:							Optional[Any]       =	sample_size

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

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

       lowerCAmelCase_:							str       =	nn.ModuleList([] )
       lowerCAmelCase_:							Dict       =	None
       lowerCAmelCase_:							Optional[Any]       =	nn.ModuleList([] )
       lowerCAmelCase_:							int       =	None

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

          if i == 0:
             input_channel += extra_in_channels

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

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

       # mid
       lowerCAmelCase_:							Optional[int]       =	get_mid_block(
           lowerCamelCase__       ,						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=lowerCamelCase__       ,						add_downsample=lowerCamelCase__       ,						)

       # up
       lowerCAmelCase_:							Dict       =	list(reversed(lowerCamelCase__ ) )
       lowerCAmelCase_:							Any       =	reversed_block_out_channels[0]
       if out_block_type is None:
          lowerCAmelCase_:							str       =	out_channels
       else:
          lowerCAmelCase_:							Any       =	block_out_channels[0]

       for i, up_block_type in enumerate(lowerCamelCase__ ):
          lowerCAmelCase_:							Dict       =	output_channel
          lowerCAmelCase_:							int       =	(
              reversed_block_out_channels[i + 1] if i < len(lowerCamelCase__ ) - 1 else final_upsample_channels
          )

          lowerCAmelCase_:							Optional[int]       =	i == len(lowerCamelCase__ ) - 1

          lowerCAmelCase_:							Union[str, Any]       =	get_up_block(
              lowerCamelCase__       ,						num_layers=lowerCamelCase__       ,						in_channels=lowerCamelCase__       ,						out_channels=lowerCamelCase__       ,						temb_channels=block_out_channels[0]       ,						add_upsample=not is_final_block       ,						)
          self.up_blocks.append(lowerCamelCase__ )
          lowerCAmelCase_:							str       =	output_channel

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




    def      _a    (    self       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__ = True       ,						):
       lowerCAmelCase_:							Any       =	timestep
       if not torch.is_tensor(lowerCamelCase__ ):
          lowerCAmelCase_:							Optional[Any]       =	torch.tensor([timesteps]       ,						dtype=torch.long       ,						device=sample.device )
       elif torch.is_tensor(lowerCamelCase__ ) and len(timesteps.shape ) == 0:
          lowerCAmelCase_:							List[Any]       =	timesteps[None].to(sample.device )

       lowerCAmelCase_:							Union[str, Any]       =	self.time_proj(lowerCamelCase__ )
       if self.config.use_timestep_embedding:
          lowerCAmelCase_:							Any       =	self.time_mlp(lowerCamelCase__ )
       else:
          lowerCAmelCase_:							Any       =	timestep_embed[..., None]
          lowerCAmelCase_:							str       =	timestep_embed.repeat([1, 1, sample.shape[2]] ).to(sample.dtype )
          lowerCAmelCase_:							Dict       =	timestep_embed.broadcast_to((sample.shape[:1] + timestep_embed.shape[1:]) )

       # 2. down
       lowerCAmelCase_:							Dict       =	()
       for downsample_block in self.down_blocks:
          lowerCAmelCase_  ,	lowerCAmelCase_:							Optional[int]       =	downsample_block(hidden_states=lowerCamelCase__       ,						temb=lowerCamelCase__ )
          down_block_res_samples += res_samples

       # 3. mid
       if self.mid_block:
          lowerCAmelCase_:							int       =	self.mid_block(lowerCamelCase__       ,						lowerCamelCase__ )

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

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

       if not return_dict:
          return (sample,)

       return UNetaDOutput(sample=lowerCamelCase__ )