


'''simple docstring'''







import copy
import inspect
import unittest

from transformers import PretrainedConfig, SwiftFormerConfig
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_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch
       from torch import nn

       from transformers import SwiftFormerForImageClassification, SwiftFormerModel
       from transformers.models.swiftformer.modeling_swiftformer import SWIFTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
       from PIL import Image

       from transformers import ViTImageProcessor
class 						UpperCamelCase__       :



     """simple docstring"""






     def __init__(			self   ,						snake_case__   ,						snake_case__=13   ,						snake_case__=3   ,						snake_case__=True   ,						snake_case__=True   ,						snake_case__=0.1   ,						snake_case__=0.1   ,						snake_case__=224   ,						snake_case__=1000   ,						snake_case__=[3, 3, 6, 4]   ,						snake_case__=[48, 56, 112, 220]   ,						):

            '''simple docstring'''

            _lowerCAmelCase     :				Dict       =							parent
            _lowerCAmelCase     :				Union[str, Any]       =							batch_size
            _lowerCAmelCase     :				Any       =							num_channels
            _lowerCAmelCase     :				List[str]       =							is_training
            _lowerCAmelCase     :				Optional[int]       =							use_labels
            _lowerCAmelCase     :				Dict       =							hidden_dropout_prob
            _lowerCAmelCase     :				List[Any]       =							attention_probs_dropout_prob
            _lowerCAmelCase     :				List[str]       =							num_labels
            _lowerCAmelCase     :				Any       =							image_size
            _lowerCAmelCase     :				List[str]       =							layer_depths
            _lowerCAmelCase     :				List[Any]       =							embed_dims







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase     :				Optional[int]       =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

            _lowerCAmelCase     :				List[str]       =							None
            if self.use_labels:
                   _lowerCAmelCase     :				Union[str, Any]       =							ids_tensor([self.batch_size]   ,						self.num_labels )

            _lowerCAmelCase     :				Optional[int]       =							self.get_config()

            return config, pixel_values, labels







     def a						(			self ):

            '''simple docstring'''

            return SwiftFormerConfig(
                depths=self.layer_depths   ,						embed_dims=self.embed_dims   ,						mlp_ratio=4   ,						downsamples=[True, True, True, True]   ,						hidden_act='gelu'   ,						num_labels=self.num_labels   ,						down_patch_size=3   ,						down_stride=2   ,						down_pad=1   ,						drop_rate=0.0   ,						drop_path_rate=0.0   ,						use_layer_scale=snake_case__   ,						layer_scale_init_value=1E-5   ,						)







     def a						(			self   ,						snake_case__   ,						snake_case__   ,						snake_case__ ):

            '''simple docstring'''

            _lowerCAmelCase     :				List[str]       =							SwiftFormerModel(config=snake_case__ )
            model.to(snake_case__ )
            model.eval()
            _lowerCAmelCase     :				Any       =							model(snake_case__ )
            self.parent.assertEqual(result.last_hidden_state.shape   ,						(self.batch_size, self.embed_dims[-1], 7, 7) )







     def a						(			self   ,						snake_case__   ,						snake_case__   ,						snake_case__ ):

            '''simple docstring'''

            _lowerCAmelCase     :				Union[str, Any]       =							self.num_labels
            _lowerCAmelCase     :				List[Any]       =							SwiftFormerForImageClassification(snake_case__ )
            model.to(snake_case__ )
            model.eval()
            _lowerCAmelCase     :				Union[str, Any]       =							model(snake_case__   ,						labels=snake_case__ )
            self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.num_labels) )

            _lowerCAmelCase     :				Optional[Any]       =							SwiftFormerForImageClassification(snake_case__ )
            model.to(snake_case__ )
            model.eval()

            _lowerCAmelCase     :				Tuple       =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )
            _lowerCAmelCase     :				Optional[int]       =							model(snake_case__ )
            self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.num_labels) )







     def a						(			self ):

            '''simple docstring'''

            ((_lowerCAmelCase)      ,  (_lowerCAmelCase)      ,  (_lowerCAmelCase))     :				Optional[Any]       =							self.prepare_config_and_inputs()
            _lowerCAmelCase     :				Union[str, Any]       =							{'pixel_values': pixel_values}
            return config, inputs_dict
@require_torch
class 						UpperCamelCase__       (			SCREAMING_SNAKE_CASE_     ,       SCREAMING_SNAKE_CASE_     ,       unittest.TestCase						):



     """simple docstring"""






     __magic_name__           =     (SwiftFormerModel, SwiftFormerForImageClassification) if is_torch_available() else ()
     __magic_name__           =     (
         {"feature-extraction": SwiftFormerModel, "image-classification": SwiftFormerForImageClassification}
         if is_torch_available()
         else {}
     )

     __magic_name__           =     False
     __magic_name__           =     False
     __magic_name__           =     False
     __magic_name__           =     False
     __magic_name__           =     False







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase     :				List[str]       =							SwiftFormerModelTester(self )
            _lowerCAmelCase     :				List[str]       =							ConfigTester(
                self   ,						config_class=snake_case__   ,						has_text_modality=snake_case__   ,						hidden_size=37   ,						num_attention_heads=12   ,						num_hidden_layers=12   ,						)







     def a						(			self ):

            '''simple docstring'''

            self.config_tester.run_common_tests()







     @unittest.skip(reason='SwiftFormer does not use inputs_embeds' )
     def a						(			self ):

            '''simple docstring'''

            pass







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase      ,  _lowerCAmelCase     :				Union[str, Any]       =							self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                   _lowerCAmelCase     :				Any       =							model_class(snake_case__ )
                   _lowerCAmelCase     :				List[str]       =							model.get_output_embeddings()
                   self.assertTrue(x is None or isinstance(snake_case__   ,						nn.Linear ) )







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase      ,  _lowerCAmelCase     :				Tuple       =							self.model_tester.prepare_config_and_inputs_for_common()

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

                   _lowerCAmelCase     :				Union[str, Any]       =							['pixel_values']
                   self.assertListEqual(arg_names[:1]   ,						snake_case__ )







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase     :				Any       =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*snake_case__ )







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase     :				Optional[int]       =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_image_classification(*snake_case__ )







     @slow
     def a						(			self ):

            '''simple docstring'''

            for model_name in SWIFTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   _lowerCAmelCase     :				Dict       =							SwiftFormerModel.from_pretrained(snake_case__ )
                   self.assertIsNotNone(snake_case__ )







     @unittest.skip(reason='SwiftFormer does not output attentions' )
     def a						(			self ):

            '''simple docstring'''

            pass







     def a						(			self ):

            '''simple docstring'''

            def check_hidden_states_output(snake_case__   ,						snake_case__   ,						snake_case__ ):
                   _lowerCAmelCase     :				Any       =							model_class(snake_case__ )
                   model.to(snake_case__ )
                   model.eval()

                   with torch.no_grad():
                          _lowerCAmelCase     :				Optional[int]       =							model(**self._prepare_for_class(snake_case__   ,						snake_case__ ) )

                   _lowerCAmelCase     :				Optional[int]       =							outputs.hidden_states

                   _lowerCAmelCase     :				Tuple       =							8
                   self.assertEqual(len(snake_case__ )   ,						snake_case__ )  # TODO

                   # SwiftFormer's feature maps are of shape (batch_size, embed_dims, height, width)
                   # with the width and height being successively divided by 2, after every 2 blocks
                   for i in range(len(snake_case__ ) ):
                          self.assertEqual(
                              hidden_states[i].shape   ,						torch.Size(
                                  [
                                      self.model_tester.batch_size,
                                      self.model_tester.embed_dims[i // 2],
                                      (self.model_tester.image_size // 4) // 2 ** (i // 2),
                                      (self.model_tester.image_size // 4) // 2 ** (i // 2),
                                  ] )   ,						)

            _lowerCAmelCase      ,  _lowerCAmelCase     :				Dict       =							self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                   _lowerCAmelCase     :				Optional[int]       =							True
                   check_hidden_states_output(snake_case__   ,						snake_case__   ,						snake_case__ )

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

                   check_hidden_states_output(snake_case__   ,						snake_case__   ,						snake_case__ )







     def a						(			self ):

            '''simple docstring'''

            def _config_zero_init(snake_case__ ):
                   _lowerCAmelCase     :				List[str]       =							copy.deepcopy(snake_case__ )
                   for key in configs_no_init.__dict__.keys():
                          if "_range" in key or "_std" in key or "initializer_factor" in key or "layer_scale" in key:
                                 setattr(snake_case__   ,						snake_case__   ,						1E-10 )
                          if isinstance(getattr(snake_case__   ,						snake_case__   ,						snake_case__ )   ,						snake_case__ ):
                                 _lowerCAmelCase     :				int       =							_config_zero_init(getattr(snake_case__   ,						snake_case__ ) )
                                 setattr(snake_case__   ,						snake_case__   ,						snake_case__ )
                   return configs_no_init

            _lowerCAmelCase      ,  _lowerCAmelCase     :				Optional[int]       =							self.model_tester.prepare_config_and_inputs_for_common()

            _lowerCAmelCase     :				List[Any]       =							_config_zero_init(snake_case__ )
            for model_class in self.all_model_classes:
                   _lowerCAmelCase     :				List[str]       =							model_class(config=snake_case__ )
                   for name, param in model.named_parameters():
                          if param.requires_grad:
                                 self.assertIn(
                                     ((param.data.mean() * 1E9) / 1E9).round().item()   ,						[0.0, 1.0]   ,						msg=F'Parameter {name} of model {model_class} seems not properly initialized'   ,						)







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

            '''simple docstring'''

            pass






def       lowercase					():







       """simple docstring"""







       _lowerCAmelCase     :				List[str]       =							Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'       )
       return image
@require_torch
@require_vision
class 						UpperCamelCase__       (			unittest.TestCase						):



     """simple docstring"""






     @cached_property
     def a						(			self ):

            '''simple docstring'''

            return ViTImageProcessor.from_pretrained('MBZUAI/swiftformer-xs' ) if is_vision_available() else None







     @slow
     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase     :				int       =							SwiftFormerForImageClassification.from_pretrained('MBZUAI/swiftformer-xs' ).to(snake_case__ )

            _lowerCAmelCase     :				Dict       =							self.default_image_processor
            _lowerCAmelCase     :				str       =							prepare_img()
            _lowerCAmelCase     :				Optional[int]       =							image_processor(images=snake_case__   ,						return_tensors='pt' ).to(snake_case__ )

            # forward pass
            with torch.no_grad():
                   _lowerCAmelCase     :				Dict       =							model(**snake_case__ )

            # verify the logits
            _lowerCAmelCase     :				Optional[Any]       =							torch.Size((1, 1000) )
            self.assertEqual(outputs.logits.shape   ,						snake_case__ )

            _lowerCAmelCase     :				Optional[int]       =							torch.tensor([[-2.1703E00, 2.1107E00, -2.0811E00]] ).to(snake_case__ )
            self.assertTrue(torch.allclose(outputs.logits[0, :3]   ,						snake_case__   ,						atol=1E-4 ) )





'''simple docstring'''







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 						UpperCamelCase__       (			SCREAMING_SNAKE_CASE_						):



     """simple docstring"""






     __magic_name__           =     42
class 						UpperCamelCase__       (			SCREAMING_SNAKE_CASE_     ,       SCREAMING_SNAKE_CASE_						):



     """simple docstring"""






     @register_to_config
     def __init__(			self   ,						snake_case__ = 6_5536   ,						snake_case__ = None   ,						snake_case__ = 2   ,						snake_case__ = 2   ,						snake_case__ = 0   ,						snake_case__ = "fourier"   ,						snake_case__ = True   ,						snake_case__ = False   ,						snake_case__ = 0.0   ,						snake_case__ = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")   ,						snake_case__ = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")   ,						snake_case__ = "UNetMidBlock1D"   ,						snake_case__ = None   ,						snake_case__ = (32, 32, 64)   ,						snake_case__ = None   ,						snake_case__ = 8   ,						snake_case__ = 1   ,						snake_case__ = False   ,						):

            '''simple docstring'''

            super().__init__()
            _lowerCAmelCase     :				Optional[int]       =							sample_size

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

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

            _lowerCAmelCase     :				Tuple       =							nn.ModuleList([] )
            _lowerCAmelCase     :				Tuple       =							None
            _lowerCAmelCase     :				Optional[Any]       =							nn.ModuleList([] )
            _lowerCAmelCase     :				str       =							None

            # down
            _lowerCAmelCase     :				List[str]       =							in_channels
            for i, down_block_type in enumerate(snake_case__ ):
                   _lowerCAmelCase     :				str       =							output_channel
                   _lowerCAmelCase     :				Any       =							block_out_channels[i]

                   if i == 0:
                          input_channel += extra_in_channels

                   _lowerCAmelCase     :				Tuple       =							i == len(snake_case__ ) - 1

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

            # mid
            _lowerCAmelCase     :				List[str]       =							get_mid_block(
                snake_case__   ,						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=snake_case__   ,						add_downsample=snake_case__   ,						)

            # up
            _lowerCAmelCase     :				Union[str, Any]       =							list(reversed(snake_case__ ) )
            _lowerCAmelCase     :				List[Any]       =							reversed_block_out_channels[0]
            if out_block_type is None:
                   _lowerCAmelCase     :				Tuple       =							out_channels
            else:
                   _lowerCAmelCase     :				Tuple       =							block_out_channels[0]

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

                   _lowerCAmelCase     :				Union[str, Any]       =							i == len(snake_case__ ) - 1

                   _lowerCAmelCase     :				Tuple       =							get_up_block(
                       snake_case__   ,						num_layers=snake_case__   ,						in_channels=snake_case__   ,						out_channels=snake_case__   ,						temb_channels=block_out_channels[0]   ,						add_upsample=not is_final_block   ,						)
                   self.up_blocks.append(snake_case__ )
                   _lowerCAmelCase     :				List[str]       =							output_channel

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







     def a						(			self   ,						snake_case__   ,						snake_case__   ,						snake_case__ = True   ,						):

            '''simple docstring'''

            _lowerCAmelCase     :				Optional[int]       =							timestep
            if not torch.is_tensor(snake_case__ ):
                   _lowerCAmelCase     :				Optional[int]       =							torch.tensor([timesteps]   ,						dtype=torch.long   ,						device=sample.device )
            elif torch.is_tensor(snake_case__ ) and len(timesteps.shape ) == 0:
                   _lowerCAmelCase     :				List[Any]       =							timesteps[None].to(sample.device )

            _lowerCAmelCase     :				int       =							self.time_proj(snake_case__ )
            if self.config.use_timestep_embedding:
                   _lowerCAmelCase     :				Optional[int]       =							self.time_mlp(snake_case__ )
            else:
                   _lowerCAmelCase     :				List[str]       =							timestep_embed[..., None]
                   _lowerCAmelCase     :				Dict       =							timestep_embed.repeat([1, 1, sample.shape[2]] ).to(sample.dtype )
                   _lowerCAmelCase     :				Tuple       =							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=snake_case__   ,						temb=snake_case__ )
                   down_block_res_samples += res_samples

            # 3. mid
            if self.mid_block:
                   _lowerCAmelCase     :				Any       =							self.mid_block(snake_case__   ,						snake_case__ )

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

            # 5. post-process
            if self.out_block:
                   _lowerCAmelCase     :				Optional[Any]       =							self.out_block(snake_case__   ,						snake_case__ )

            if not return_dict:
                   return (sample,)

            return UNetaDOutput(sample=snake_case__ )

