



import inspect
import unittest

from transformers import ConvNextConfig
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, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import ConvNextBackbone, ConvNextForImageClassification, ConvNextModel
   from transformers.models.convnext.modeling_convnext import CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
   from PIL import Image

   from transformers import AutoImageProcessor


class   snake_case_ :




       """simple docstring"""






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

             """simple docstring"""





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



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	List[str]				      =		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

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

             UpperCAmelCase_  :	int				      =		self.get_config()
             return config, pixel_values, labels



       def  A_   (     self):

             """simple docstring"""





             return ConvNextConfig(
                 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=lowercase     ,initializer_range=self.initializer_range     ,out_features=self.out_features     ,out_indices=self.out_indices     ,num_labels=self.num_labels     ,)



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	Tuple				      =		ConvNextModel(config=lowercase)
             model.to(lowercase)
             model.eval()
             UpperCAmelCase_  :	Dict				      =		model(lowercase)
             # 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     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		ConvNextForImageClassification(lowercase)
             model.to(lowercase)
             model.eval()
             UpperCAmelCase_  :	Any				      =		model(lowercase     ,labels=lowercase)
             self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.num_labels))



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	List[str]				      =		ConvNextBackbone(config=lowercase)
             model.to(lowercase)
             model.eval()
             UpperCAmelCase_  :	Optional[Any]				      =		model(lowercase)

             # 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
             UpperCAmelCase_  :	Tuple				      =		None
             UpperCAmelCase_  :	Tuple				      =		ConvNextBackbone(config=lowercase)
             model.to(lowercase)
             model.eval()
             UpperCAmelCase_  :	Optional[int]				      =		model(lowercase)

             # 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):

             """simple docstring"""





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


@require_torch
class   snake_case_ (lowercase__ , lowercase__ , unittest.TestCase  ):




       """simple docstring"""






       _lowerCamelCase			    =				(
           (
               ConvNextModel,
               ConvNextForImageClassification,
               ConvNextBackbone,
           )
           if is_torch_available()
           else ()
       )
       _lowerCamelCase			    =				(
           {"""feature-extraction""": ConvNextModel, """image-classification""": ConvNextForImageClassification}
           if is_torch_available()
           else {}
       )

       _lowerCamelCase			    =				True
       _lowerCamelCase			    =				False
       _lowerCamelCase			    =				False
       _lowerCamelCase			    =				False
       _lowerCamelCase			    =				False



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Dict				      =		ConvNextModelTester(self)
             UpperCAmelCase_  :	Union[str, Any]				      =		ConfigTester(self     ,config_class=lowercase     ,has_text_modality=lowercase     ,hidden_size=37)



       def  A_   (     self):

             """simple docstring"""





             self.create_and_test_config_common_properties()
             self.config_tester.create_and_test_config_to_json_string()
             self.config_tester.create_and_test_config_to_json_file()
             self.config_tester.create_and_test_config_from_and_save_pretrained()
             self.config_tester.create_and_test_config_with_num_labels()
             self.config_tester.check_config_can_be_init_without_params()
             self.config_tester.check_config_arguments_init()



       def  A_   (     self):

             """simple docstring"""





             return



       @unittest.skip(reason="ConvNext does not use inputs_embeds")
       def  A_   (     self):

             """simple docstring"""





             pass



       @unittest.skip(reason="ConvNext does not support input and output embeddings")
       def  A_   (     self):

             """simple docstring"""





             pass



       @unittest.skip(reason="ConvNext does not use feedforward chunking")
       def  A_   (     self):

             """simple docstring"""





             pass



       def  A_   (     self):

             """simple docstring"""





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

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

                   UpperCAmelCase_  :	str				      =		["pixel_values"]
                   self.assertListEqual(arg_names[:1]     ,lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Dict				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	str				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_backbone(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             def check_hidden_states_output(lowercase     ,lowercase     ,lowercase):
                   UpperCAmelCase_  :	Optional[int]				      =		model_class(lowercase)
                   model.to(lowercase)
                   model.eval()

                   with torch.no_grad():
                         UpperCAmelCase_  :	Any				      =		model(**self._prepare_for_class(lowercase     ,lowercase))

                   UpperCAmelCase_  :	Union[str, Any]				      =		outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                   UpperCAmelCase_  :	Tuple				      =		self.model_tester.num_stages
                   self.assertEqual(len(lowercase)     ,expected_num_stages + 1)

                   # ConvNext'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]     ,)

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

             for model_class in self.all_model_classes:
                   UpperCAmelCase_  :	int				      =		True
                   check_hidden_states_output(lowercase     ,lowercase     ,lowercase)

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

                   check_hidden_states_output(lowercase     ,lowercase     ,lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Any				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_image_classification(*lowercase)



       @slow
       def  A_   (     self):

             """simple docstring"""





             for model_name in CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   UpperCAmelCase_  :	Tuple				      =		ConvNextModel.from_pretrained(lowercase)
                   self.assertIsNotNone(lowercase)







def    _snake_case (       )     ->     str:






      '''simple docstring'''







      UpperCAmelCase_  :	List[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  A_   (     self):

             """simple docstring"""





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



       @slow
       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Union[str, Any]				      =		ConvNextForImageClassification.from_pretrained("facebook/convnext-tiny-224").to(lowercase)

             UpperCAmelCase_  :	Tuple				      =		self.default_image_processor
             UpperCAmelCase_  :	Any				      =		prepare_img()
             UpperCAmelCase_  :	Tuple				      =		image_processor(images=lowercase     ,return_tensors="pt").to(lowercase)

             # forward pass
             with torch.no_grad():
                   UpperCAmelCase_  :	List[str]				      =		model(**lowercase)

             # verify the logits
             UpperCAmelCase_  :	Any				      =		torch.Size((1, 1000))
             self.assertEqual(outputs.logits.shape     ,lowercase)

             UpperCAmelCase_  :	str				      =		torch.tensor([-0.0260, -0.4739, 0.1911]).to(lowercase)

             self.assertTrue(torch.allclose(outputs.logits[0, :3]     ,lowercase     ,atol=1E-4))


@require_torch
class   snake_case_ (unittest.TestCase , lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				(ConvNextBackbone,) if is_torch_available() else ()
       _lowerCamelCase			    =				ConvNextConfig

       _lowerCamelCase			    =				False



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	List[Any]				      =		ConvNextModelTester(self)





# 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			    =				42
       _lowerCamelCase			    =				42


class   snake_case_ (lowercase__ , lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				1



       @register_to_config
       def __init__(     self     ,lowercase = 2000     ,lowercase = 0.15     ,lowercase = 0.01     ,lowercase = 1348.0     ,lowercase = 1E-5     ,lowercase = 1     ,):

             """simple docstring"""





             UpperCAmelCase_  :	Optional[int]				      =		sigma_max

             # setable values
             UpperCAmelCase_  :	Optional[int]				      =		None

             self.set_sigmas(lowercase     ,lowercase     ,lowercase     ,lowercase)



       def  A_   (     self     ,lowercase     ,lowercase = None):

             """simple docstring"""





             return sample



       def  A_   (     self     ,lowercase     ,lowercase = None     ,lowercase = None):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		sampling_eps if sampling_eps is not None else self.config.sampling_eps

             UpperCAmelCase_  :	List[Any]				      =		torch.linspace(1     ,lowercase     ,lowercase     ,device=lowercase)



       def  A_   (     self     ,lowercase     ,lowercase = None     ,lowercase = None     ,lowercase = None):

             """simple docstring"""





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

             UpperCAmelCase_  :	Union[str, Any]				      =		sigma_min * (sigma_max / sigma_min) ** (self.timesteps / sampling_eps)
             UpperCAmelCase_  :	Optional[int]				      =		torch.exp(torch.linspace(math.log(lowercase)     ,math.log(lowercase)     ,lowercase))
             UpperCAmelCase_  :	Any				      =		torch.tensor([sigma_min * (sigma_max / sigma_min) ** t for t in self.timesteps])



       def  A_   (     self     ,lowercase     ,lowercase):

             """simple docstring"""





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



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase = None     ,lowercase = True     ,):

             """simple docstring"""





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

             UpperCAmelCase_  :	Optional[int]				      =		timestep * torch.ones(
                 sample.shape[0]     ,device=sample.device)  # torch.repeat_interleave(timestep, sample.shape[0])
             UpperCAmelCase_  :	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
             UpperCAmelCase_  :	Optional[int]				      =		timesteps.to(self.discrete_sigmas.device)

             UpperCAmelCase_  :	Optional[Any]				      =		self.discrete_sigmas[timesteps].to(sample.device)
             UpperCAmelCase_  :	Optional[Any]				      =		self.get_adjacent_sigma(lowercase     ,lowercase).to(sample.device)
             UpperCAmelCase_  :	Any				      =		torch.zeros_like(lowercase)
             UpperCAmelCase_  :	Dict				      =		(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
             UpperCAmelCase_  :	Dict				      =		diffusion.flatten()
             while len(diffusion.shape) < len(sample.shape):
                   UpperCAmelCase_  :	List[str]				      =		diffusion.unsqueeze(-1)
             UpperCAmelCase_  :	List[Any]				      =		drift - diffusion**2 * model_output

             #  equation 6: sample noise for the diffusion term of
             UpperCAmelCase_  :	Union[str, Any]				      =		randn_tensor(
                 sample.shape     ,layout=sample.layout     ,generator=lowercase     ,device=sample.device     ,dtype=sample.dtype)
             UpperCAmelCase_  :	Any				      =		sample - drift  # subtract because `dt` is a small negative timestep
             # TODO is the variable diffusion the correct scaling term for the noise?
             UpperCAmelCase_  :	Tuple				      =		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=lowercase     ,prev_sample_mean=lowercase)



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase = None     ,lowercase = True     ,):

             """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
             UpperCAmelCase_  :	int				      =		randn_tensor(sample.shape     ,layout=sample.layout     ,generator=lowercase).to(sample.device)

             # compute step size from the model_output, the noise, and the snr
             UpperCAmelCase_  :	Union[str, Any]				      =		torch.norm(model_output.reshape(model_output.shape[0]     ,-1)     ,dim=-1).mean()
             UpperCAmelCase_  :	Optional[Any]				      =		torch.norm(noise.reshape(noise.shape[0]     ,-1)     ,dim=-1).mean()
             UpperCAmelCase_  :	List[Any]				      =		(self.config.snr * noise_norm / grad_norm) ** 2 * 2
             UpperCAmelCase_  :	Optional[Any]				      =		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
             UpperCAmelCase_  :	Any				      =		step_size.flatten()
             while len(step_size.shape) < len(sample.shape):
                   UpperCAmelCase_  :	Tuple				      =		step_size.unsqueeze(-1)
             UpperCAmelCase_  :	Dict				      =		sample + step_size * model_output
             UpperCAmelCase_  :	int				      =		prev_sample_mean + ((step_size * 2) ** 0.5) * noise

             if not return_dict:
                   return (prev_sample,)

             return SchedulerOutput(prev_sample=lowercase)



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,):

             """simple docstring"""





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



       def __len__(     self):

             """simple docstring"""





             return self.config.num_train_timesteps
