

import inspect
import unittest

from transformers import ConvNextVaConfig
from transformers.models.auto import get_values
from transformers.models.auto.modeling_auto import MODEL_FOR_BACKBONE_MAPPING_NAMES, MODEL_MAPPING_NAMES
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 transformers import ConvNextVaBackbone, ConvNextVaForImageClassification, ConvNextVaModel
      from transformers.models.convnextva.modeling_convnextva import CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import AutoImageProcessor

class   _a					:




      """simple docstring"""

      def __init__(   self							:				List[str]						,			UpperCAmelCase							:				Union[str, Any]						,			UpperCAmelCase							:				Optional[Any]=13						,			UpperCAmelCase							:				Optional[Any]=32						,			UpperCAmelCase							:				Tuple=3						,			UpperCAmelCase							:				str=4						,			UpperCAmelCase							:				List[str]=[10, 20, 30, 40]						,			UpperCAmelCase							:				Optional[Any]=[2, 2, 3, 2]						,			UpperCAmelCase							:				Any=True						,			UpperCAmelCase							:				List[Any]=True						,			UpperCAmelCase							:				List[Any]=37						,			UpperCAmelCase							:				Any="gelu"						,			UpperCAmelCase							:				Optional[Any]=10						,			UpperCAmelCase							:				List[Any]=0.02						,			UpperCAmelCase							:				Dict=["stage2", "stage3", "stage4"]						,			UpperCAmelCase							:				List[Any]=[2, 3, 4]						,			UpperCAmelCase							:				Dict=None						,			):
            A_						  =				parent
            A_						  =				batch_size
            A_						  =				image_size
            A_						  =				num_channels
            A_						  =				num_stages
            A_						  =				hidden_sizes
            A_						  =				depths
            A_						  =				is_training
            A_						  =				use_labels
            A_						  =				intermediate_size
            A_						  =				hidden_act
            A_						  =				num_labels
            A_						  =				initializer_range
            A_						  =				out_features
            A_						  =				out_indices
            A_						  =				scope





      def 		__A	(   self							:				Optional[int]	):
            A_						  =				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

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

            A_						  =				self.get_config()

            return config, pixel_values, labels





      def 		__A	(   self							:				List[str]	):
            return ConvNextVaConfig(
                num_channels=self.num_channels						,			hidden_sizes=self.hidden_sizes						,			depths=self.depths						,			num_stages=self.num_stages						,			hidden_act=self.hidden_act						,			is_decoder=a__						,			initializer_range=self.initializer_range						,			out_features=self.out_features						,			out_indices=self.out_indices						,			num_labels=self.num_labels						,			)





      def 		__A	(   self							:				int						,			UpperCAmelCase							:				Dict						,			UpperCAmelCase							:				List[Any]						,			UpperCAmelCase							:				Tuple	):
            A_						  =				ConvNextVaModel(config=a__	)
            model.to(a__	)
            model.eval()
            A_						  =				model(a__	)
            # expected last hidden states: B, C, H // 32, W // 32
            self.parent.assertEqual(
                result.last_hidden_state.shape						,			(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)						,			)





      def 		__A	(   self							:				Any						,			UpperCAmelCase							:				Any						,			UpperCAmelCase							:				Optional[Any]						,			UpperCAmelCase							:				str	):
            A_						  =				ConvNextVaForImageClassification(a__	)
            model.to(a__	)
            model.eval()
            A_						  =				model(a__						,			labels=a__	)
            self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.num_labels)	)





      def 		__A	(   self							:				Any						,			UpperCAmelCase							:				Optional[Any]						,			UpperCAmelCase							:				Tuple						,			UpperCAmelCase							:				Union[str, Any]	):
            A_						  =				ConvNextVaBackbone(config=a__	)
            model.to(a__	)
            model.eval()
            A_						  =				model(a__	)

            # verify hidden states
            self.parent.assertEqual(len(result.feature_maps	)						,			len(config.out_features	)	)
            self.parent.assertListEqual(list(result.feature_maps[0].shape	)						,			[self.batch_size, self.hidden_sizes[1], 4, 4]	)

            # 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
            A_						  =				None
            A_						  =				ConvNextVaBackbone(config=a__	)
            model.to(a__	)
            model.eval()
            A_						  =				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.hidden_sizes[-1], 1, 1]	)

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





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







      def 		__A	(   self							:				Optional[Any]	):
            A_						  =				self.prepare_config_and_inputs()
            A_						  =				config_and_inputs
            A_						  =				{"""pixel_values""": pixel_values, """labels""": labels}
            return config, inputs_dict




@require_torch
class   _a					(			lowercase_     ,				lowercase_     ,				unittest.TestCase				):




      """simple docstring"""

      _lowerCamelCase       :			Optional[int]      		=       (
          (
              ConvNextVaModel,
              ConvNextVaForImageClassification,
              ConvNextVaBackbone,
          )
          if is_torch_available()
          else ()
      )
      _lowerCamelCase       :			int      		=       (
          {"feature-extraction": ConvNextVaModel, "image-classification": ConvNextVaForImageClassification}
          if is_torch_available()
          else {}
      )

      _lowerCamelCase       :			Tuple      		=       False
      _lowerCamelCase       :			Tuple      		=       False
      _lowerCamelCase       :			int      		=       False
      _lowerCamelCase       :			List[str]      		=       False
      _lowerCamelCase       :			Optional[int]      		=       False





      def 		__A	(   self							:				Union[str, Any]	):
            A_						  =				ConvNextVaModelTester(self	)
            A_						  =				ConfigTester(self						,			config_class=a__						,			has_text_modality=a__						,			hidden_size=37	)





      def 		__A	(   self							:				Dict	):
            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 		__A	(   self							:				Union[str, Any]	):
            return





      @unittest.skip(reason="ConvNextV2 does not use inputs_embeds"	)
      def 		__A	(   self							:				Optional[int]	):
            pass





      @unittest.skip(reason="ConvNextV2 does not support input and output embeddings"	)
      def 		__A	(   self							:				List[str]	):
            pass





      @unittest.skip(reason="ConvNextV2 does not use feedforward chunking"	)
      def 		__A	(   self							:				Tuple	):
            pass





      def 		__A	(   self							:				Optional[int]	):





            if not self.model_tester.is_training:
                  return

            for model_class in self.all_model_classes:
                  A_						  =				self.model_tester.prepare_config_and_inputs_with_labels()
                  A_						  =				True

                  if model_class.__name__ in [
                      *get_values(a__	),
                      *get_values(a__	),
                  ]:
                        continue

                  A_						  =				model_class(a__	)
                  model.to(a__	)
                  model.train()
                  A_						  =				self._prepare_for_class(a__						,			a__						,			return_labels=a__	)
                  A_						  =				model(**a__	).loss
                  loss.backward()





      def 		__A	(   self							:				List[str]	):





            if not self.model_tester.is_training:
                  return

            for model_class in self.all_model_classes:
                  A_						  =				self.model_tester.prepare_config_and_inputs_with_labels()
                  A_						  =				False
                  A_						  =				True

                  if (
                      model_class.__name__
                      in [*get_values(a__	), *get_values(a__	)]
                      or not model_class.supports_gradient_checkpointing
                  ):
                        continue

                  A_						  =				model_class(a__	)
                  model.to(a__	)
                  model.gradient_checkpointing_enable()
                  model.train()
                  A_						  =				self._prepare_for_class(a__						,			a__						,			return_labels=a__	)
                  A_						  =				model(**a__	).loss
                  loss.backward()





      def 		__A	(   self							:				str	):
            A_						  =				self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                  A_						  =				model_class(a__	)
                  A_						  =				inspect.signature(model.forward	)
                  # signature.parameters is an OrderedDict => so arg_names order is deterministic
                  A_						  =				[*signature.parameters.keys()]

                  A_						  =				["""pixel_values"""]
                  self.assertListEqual(arg_names[:1]						,			a__	)





      def 		__A	(   self							:				Any	):
            A_						  =				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*a__	)





      def 		__A	(   self							:				Tuple	):





            def check_hidden_states_output(UpperCAmelCase							:				List[str]						,			UpperCAmelCase							:				str						,			UpperCAmelCase							:				List[str]	):
                  A_						  =				model_class(a__	)
                  model.to(a__	)
                  model.eval()

                  with torch.no_grad():
                        A_						  =				model(**self._prepare_for_class(a__						,			a__	)	)

                  A_						  =				outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                  A_						  =				self.model_tester.num_stages
                  self.assertEqual(len(a__	)						,			expected_num_stages + 1	)

                  # ConvNextV2's feature maps are of shape (batch_size, num_channels, height, width)
                  self.assertListEqual(
                      list(hidden_states[0].shape[-2:]	)						,			[self.model_tester.image_size // 4, self.model_tester.image_size // 4]						,			)

            A_						  =				self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                  A_						  =				True
                  check_hidden_states_output(a__						,			a__						,			a__	)

                  # check that output_hidden_states also work using config
                  del inputs_dict["output_hidden_states"]
                  A_						  =				True

                  check_hidden_states_output(a__						,			a__						,			a__	)





      def 		__A	(   self							:				int	):
            A_						  =				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_image_classification(*a__	)







      @slow
      def 		__A	(   self							:				Any	):





            for model_name in CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A_						  =				ConvNextVaModel.from_pretrained(a__	)
                  self.assertIsNotNone(a__	)




def       __snake_case						( ):



      """simple docstring"""

      A_						  =				Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"    )
      return image



@require_torch
@require_vision
class   _a					(			unittest.TestCase				):




      """simple docstring"""

      @cached_property
      def 		__A	(   self							:				List[Any]	):
            return AutoImageProcessor.from_pretrained("facebook/convnextv2-tiny-1k-224"	) if is_vision_available() else None







      @slow
      def 		__A	(   self							:				List[Any]	):
            A_						  =				ConvNextVaForImageClassification.from_pretrained("facebook/convnextv2-tiny-1k-224"	).to(a__	)

            A_						  =				self.default_image_processor
            A_						  =				prepare_img()
            A_						  =				preprocessor(images=a__						,			return_tensors="pt"	).to(a__	)

            # forward pass
            with torch.no_grad():
                  A_						  =				model(**a__	)

            # verify the logits
            A_						  =				torch.Size((1, 1000)	)
            self.assertEqual(outputs.logits.shape						,			a__	)

            A_						  =				torch.tensor([0.9_996, 0.1_966, -0.4_386]	).to(a__	)
            self.assertTrue(torch.allclose(outputs.logits[0, :3]						,			a__						,			atol=1E-4	)	)







# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch

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

import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, randn_tensor
from .scheduling_utils import SchedulerMixin, SchedulerOutput




@dataclass
class       snake_case__ (				lowercase_):

       '''simple docstring'''
       lowerCamelCase				:					torch.FloatTensor
       lowerCamelCase				:					torch.FloatTensor




class       snake_case__ (				lowercase_						,     lowercase_):

       '''simple docstring'''
       lowerCamelCase				:					int								=			1



       @register_to_config
       def __init__(				self      ,    a__ = 20_00      ,    a__ = 0.15      ,    a__ = 0.01      ,    a__ = 13_48.0      ,    a__ = 1e-5      ,    a__ = 1      ,    )			->   Any:






            '''simple docstring'''





            __snake_case      :int 							=    sigma_max

            # setable values
            __snake_case      :Any 							=    None

            self.set_sigmas(a__      ,    a__      ,    a__      ,    a__     )



       def        __lowercase    (				self      ,    a__      ,    a__ = None     )			->   torch.FloatTensor:






            '''simple docstring'''





            return sample



       def        __lowercase    (				self      ,    a__      ,    a__ = None      ,    a__ = None     )			->   Union[str, Any]:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    sampling_eps if sampling_eps is not None else self.config.sampling_eps

            __snake_case      :int 							=    torch.linspace(1      ,    a__      ,    a__      ,    device=a__     )



       def        __lowercase    (				self      ,    a__      ,    a__ = None      ,    a__ = None      ,    a__ = None     )			->   Union[str, Any]:






            '''simple docstring'''





            __snake_case      :Union[str, Any] 							=    sigma_min if sigma_min is not None else self.config.sigma_min
            __snake_case      :Any 							=    sigma_max if sigma_max is not None else self.config.sigma_max
            __snake_case      :Any 							=    sampling_eps if sampling_eps is not None else self.config.sampling_eps
            if self.timesteps is None:
                 self.set_timesteps(a__      ,    a__     )

            __snake_case      :int 							=    sigma_min * (sigma_max / sigma_min) ** (self.timesteps / sampling_eps)
            __snake_case      :List[str] 							=    torch.exp(torch.linspace(math.log(a__     )      ,    math.log(a__     )      ,    a__     )     )
            __snake_case      :Any 							=    torch.tensor([sigma_min * (sigma_max / sigma_min) ** t for t in self.timesteps]     )



       def        __lowercase    (				self      ,    a__      ,    a__     )			->   Tuple:






            '''simple docstring'''





            return torch.where(
                timesteps == 0      ,    torch.zeros_like(t.to(timesteps.device     )     )      ,    self.discrete_sigmas[timesteps - 1].to(timesteps.device     )      ,    )



       def        __lowercase    (				self      ,    a__      ,    a__      ,    a__      ,    a__ = None      ,    a__ = True      ,    )			->   Union[SdeVeOutput, Tuple]:






            '''simple docstring'''





            if self.timesteps is None:
                 raise ValueError(
                     """`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"""     )

            __snake_case      :Optional[int] 							=    timestep * torch.ones(
                sample.shape[0]      ,    device=sample.device     )  # torch.repeat_interleave(timestep, sample.shape[0])
            __snake_case      :Tuple 							=    (timestep * (len(self.timesteps     ) - 1)).long()

            # mps requires indices to be in the same device, so we use cpu as is the default with cuda
            __snake_case      :str 							=    timesteps.to(self.discrete_sigmas.device     )

            __snake_case      :str 							=    self.discrete_sigmas[timesteps].to(sample.device     )
            __snake_case      :Union[str, Any] 							=    self.get_adjacent_sigma(a__      ,    a__     ).to(sample.device     )
            __snake_case      :Dict 							=    torch.zeros_like(a__     )
            __snake_case      :str 							=    (sigma**2 - adjacent_sigma**2) ** 0.5

            # equation 6 in the paper: the model_output modeled by the network is grad_x log pt(x)
            # also equation 47 shows the analog from SDE models to ancestral sampling methods
            __snake_case      :Optional[Any] 							=    diffusion.flatten()
            while len(diffusion.shape     ) < len(sample.shape     ):
                 __snake_case      :int 							=    diffusion.unsqueeze(-1     )
            __snake_case      :Optional[Any] 							=    drift - diffusion**2 * model_output

            #  equation 6: sample noise for the diffusion term of
            __snake_case      :str 							=    randn_tensor(
                sample.shape      ,    layout=sample.layout      ,    generator=a__      ,    device=sample.device      ,    dtype=sample.dtype     )
            __snake_case      :str 							=    sample - drift  # subtract because `dt` is a small negative timestep
            # TODO is the variable diffusion the correct scaling term for the noise?
            __snake_case      :str 							=    prev_sample_mean + diffusion * noise  # add impact of diffusion field g

            if not return_dict:
                 return (prev_sample, prev_sample_mean)

            return SdeVeOutput(prev_sample=a__      ,    prev_sample_mean=a__     )



       def        __lowercase    (				self      ,    a__      ,    a__      ,    a__ = None      ,    a__ = True      ,    )			->   Union[SchedulerOutput, Tuple]:






            '''simple docstring'''





            if self.timesteps is None:
                 raise ValueError(
                     """`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"""     )

            # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z"
            # sample noise for correction
            __snake_case      :Any 							=    randn_tensor(sample.shape      ,    layout=sample.layout      ,    generator=a__     ).to(sample.device     )

            # compute step size from the model_output, the noise, and the snr
            __snake_case      :List[Any] 							=    torch.norm(model_output.reshape(model_output.shape[0]      ,    -1     )      ,    dim=-1     ).mean()
            __snake_case      :Tuple 							=    torch.norm(noise.reshape(noise.shape[0]      ,    -1     )      ,    dim=-1     ).mean()
            __snake_case      :Dict 							=    (self.config.snr * noise_norm / grad_norm) ** 2 * 2
            __snake_case      :int 							=    step_size * torch.ones(sample.shape[0]     ).to(sample.device     )
            # self.repeat_scalar(step_size, sample.shape[0])

            # compute corrected sample: model_output term and noise term
            __snake_case      :Tuple 							=    step_size.flatten()
            while len(step_size.shape     ) < len(sample.shape     ):
                 __snake_case      :Optional[int] 							=    step_size.unsqueeze(-1     )
            __snake_case      :List[str] 							=    sample + step_size * model_output
            __snake_case      :Tuple 							=    prev_sample_mean + ((step_size * 2) ** 0.5) * noise

            if not return_dict:
                 return (prev_sample,)

            return SchedulerOutput(prev_sample=a__     )



       def        __lowercase    (				self      ,    a__      ,    a__      ,    a__      ,    )			->   torch.FloatTensor:






            '''simple docstring'''





            __snake_case      :Tuple 							=    timesteps.to(original_samples.device     )
            __snake_case      :Optional[Any] 							=    self.discrete_sigmas.to(original_samples.device     )[timesteps]
            __snake_case      :int 							=    (
                noise * sigmas[:, None, None, None]
                if noise is not None
                else torch.randn_like(a__     ) * sigmas[:, None, None, None]
            )
            __snake_case      :str 							=    noise + original_samples
            return noisy_samples



       def __len__(				self     )			->   Optional[int]:






            '''simple docstring'''





            return self.config.num_train_timesteps


