






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

       '''simple docstring'''



       def __init__(				self      ,    a__      ,    a__=13      ,    a__=32      ,    a__=3      ,    a__=4      ,    a__=[10, 20, 30, 40]      ,    a__=[2, 2, 3, 2]      ,    a__=True      ,    a__=True      ,    a__=37      ,    a__="gelu"      ,    a__=10      ,    a__=0.02      ,    a__=["stage2", "stage3", "stage4"]      ,    a__=[2, 3, 4]      ,    a__=None      ,    )			->   Optional[Any]:






            '''simple docstring'''





            __snake_case      :int 							=    parent
            __snake_case      :int 							=    batch_size
            __snake_case      :List[str] 							=    image_size
            __snake_case      :Union[str, Any] 							=    num_channels
            __snake_case      :Union[str, Any] 							=    num_stages
            __snake_case      :Any 							=    hidden_sizes
            __snake_case      :List[Any] 							=    depths
            __snake_case      :Any 							=    is_training
            __snake_case      :int 							=    use_labels
            __snake_case      :Dict 							=    intermediate_size
            __snake_case      :Optional[int] 							=    hidden_act
            __snake_case      :int 							=    num_labels
            __snake_case      :Tuple 							=    initializer_range
            __snake_case      :Optional[int] 							=    out_features
            __snake_case      :Optional[int] 							=    out_indices
            __snake_case      :Dict 							=    scope



       def        __lowercase    (				self     )			->   List[str]:






            '''simple docstring'''





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

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

            __snake_case      :Optional[Any] 							=    self.get_config()

            return config, pixel_values, labels



       def        __lowercase    (				self     )			->   List[str]:






            '''simple docstring'''





            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        __lowercase    (				self      ,    a__      ,    a__      ,    a__     )			->   Tuple:






            '''simple docstring'''





            __snake_case      :Optional[Any] 							=    ConvNextVaModel(config=a__     )
            model.to(a__     )
            model.eval()
            __snake_case      :List[Any] 							=    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        __lowercase    (				self      ,    a__      ,    a__      ,    a__     )			->   List[str]:






            '''simple docstring'''





            __snake_case      :Any 							=    ConvNextVaForImageClassification(a__     )
            model.to(a__     )
            model.eval()
            __snake_case      :Optional[int] 							=    model(a__      ,    labels=a__     )
            self.parent.assertEqual(result.logits.shape      ,    (self.batch_size, self.num_labels)     )



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






            '''simple docstring'''





            __snake_case      :Optional[Any] 							=    ConvNextVaBackbone(config=a__     )
            model.to(a__     )
            model.eval()
            __snake_case      :List[Any] 							=    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
            __snake_case      :List[Any] 							=    None
            __snake_case      :List[Any] 							=    ConvNextVaBackbone(config=a__     )
            model.to(a__     )
            model.eval()
            __snake_case      :Any 							=    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        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





            __snake_case      :Optional[int] 							=    self.prepare_config_and_inputs()
            __snake_case							,			__snake_case							,			__snake_case      :int 							=    config_and_inputs
            __snake_case      :Tuple 							=    {"""pixel_values""": pixel_values}
            return config, inputs_dict



       def        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    self.prepare_config_and_inputs()
            __snake_case							,			__snake_case							,			__snake_case      :Union[str, Any] 							=    config_and_inputs
            __snake_case      :Optional[Any] 							=    {"""pixel_values""": pixel_values, """labels""": labels}
            return config, inputs_dict




@require_torch
class       snake_case__ (				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        __lowercase    (				self     )			->   Tuple:






            '''simple docstring'''





            __snake_case      :Optional[Any] 							=    ConvNextVaModelTester(self     )
            __snake_case      :List[Any] 							=    ConfigTester(self      ,    config_class=a__      ,    has_text_modality=a__      ,    hidden_size=37     )



       def        __lowercase    (				self     )			->   Any:






            '''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     )			->   Tuple:






            '''simple docstring'''





            return



       @unittest.skip(reason="""ConvNextV2 does not use inputs_embeds"""     )
       def        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





            pass



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






            '''simple docstring'''





            pass



       @unittest.skip(reason="""ConvNextV2 does not use feedforward chunking"""     )
       def        __lowercase    (				self     )			->   List[Any]:






            '''simple docstring'''





            pass



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






            '''simple docstring'''





            if not self.model_tester.is_training:
                 return

            for model_class in self.all_model_classes:
                 __snake_case							,			__snake_case      :List[str] 							=    self.model_tester.prepare_config_and_inputs_with_labels()
                 __snake_case      :Union[str, Any] 							=    True

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

                 __snake_case      :Optional[Any] 							=    model_class(a__     )
                 model.to(a__     )
                 model.train()
                 __snake_case      :Dict 							=    self._prepare_for_class(a__      ,    a__      ,    return_labels=a__     )
                 __snake_case      :Dict 							=    model(**a__     ).loss
                 loss.backward()



       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            if not self.model_tester.is_training:
                 return

            for model_class in self.all_model_classes:
                 __snake_case							,			__snake_case      :int 							=    self.model_tester.prepare_config_and_inputs_with_labels()
                 __snake_case      :Optional[int] 							=    False
                 __snake_case      :str 							=    True

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

                 __snake_case      :List[str] 							=    model_class(a__     )
                 model.to(a__     )
                 model.gradient_checkpointing_enable()
                 model.train()
                 __snake_case      :Optional[Any] 							=    self._prepare_for_class(a__      ,    a__      ,    return_labels=a__     )
                 __snake_case      :Dict 							=    model(**a__     ).loss
                 loss.backward()



       def        __lowercase    (				self     )			->   str:






            '''simple docstring'''





            __snake_case							,			__snake_case      :str 							=    self.model_tester.prepare_config_and_inputs_for_common()

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

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



       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*a__     )



       def        __lowercase    (				self     )			->   Tuple:






            '''simple docstring'''





            def check_hidden_states_output(a__      ,    a__      ,    a__     ):
                 __snake_case      :Dict 							=    model_class(a__     )
                 model.to(a__     )
                 model.eval()

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

                 __snake_case      :Dict 							=    outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                 __snake_case      :Union[str, Any] 							=    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]      ,    )

            __snake_case							,			__snake_case      :int 							=    self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                 __snake_case      :List[Any] 							=    True
                 check_hidden_states_output(a__      ,    a__      ,    a__     )

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

                 check_hidden_states_output(a__      ,    a__      ,    a__     )



       def        __lowercase    (				self     )			->   str:






            '''simple docstring'''





            __snake_case      :Dict 							=    self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_image_classification(*a__     )



       @slow
       def        __lowercase    (				self     )			->   List[str]:






            '''simple docstring'''





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






def 	UpperCamelCase	(							):




     '''simple docstring'''






     __snake_case      :Optional[Any] 							=    Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""	)
     return image




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

       '''simple docstring'''



       @cached_property
       def        __lowercase    (				self     )			->   Optional[int]:






            '''simple docstring'''





            return AutoImageProcessor.from_pretrained("""facebook/convnextv2-tiny-1k-224"""     ) if is_vision_available() else None



       @slow
       def        __lowercase    (				self     )			->   int:






            '''simple docstring'''





            __snake_case      :Optional[Any] 							=    ConvNextVaForImageClassification.from_pretrained("""facebook/convnextv2-tiny-1k-224"""     ).to(a__     )

            __snake_case      :Optional[int] 							=    self.default_image_processor
            __snake_case      :Union[str, Any] 							=    prepare_img()
            __snake_case      :Any 							=    preprocessor(images=a__      ,    return_tensors="""pt"""     ).to(a__     )

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

            # verify the logits
            __snake_case      :Optional[Any] 							=    torch.Size((1, 10_00)     )
            self.assertEqual(outputs.logits.shape      ,    a__     )

            __snake_case      :Tuple 							=    torch.tensor([0.99_96, 0.19_66, -0.43_86]     ).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


