





'''simple docstring'''




from __future__ import annotations

import inspect
import unittest
from math import floor

import numpy as np

from transformers import CvtConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
 import tensorflow as tf

 from transformers import TFCvtForImageClassification, TFCvtModel
 from transformers.models.cvt.modeling_tf_cvt import TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor
class 			_lowerCAmelCase		(			A__	):






     """simple docstring"""






     def      lowerCAmelCase			(   self							:			Tuple						)->	Optional[Any]:
        snake_case 						=   self.config_class(**self.inputs_dict						)
        self.parent.assertTrue(hasattr(__snake_case   ,       """embed_dim"""						)						)
        self.parent.assertTrue(hasattr(__snake_case   ,       """num_heads"""						)						)




class 			_lowerCAmelCase		:






     """simple docstring"""






     def __init__(   self							:			Optional[Any]   ,       __snake_case							:			int   ,       __snake_case							:			Optional[Any]=13   ,       __snake_case							:			Union[str, Any]=64   ,       __snake_case							:			str=3   ,       __snake_case							:			Tuple=[16, 48, 96]   ,       __snake_case							:			int=[1, 3, 6]   ,       __snake_case							:			Optional[Any]=[1, 2, 10]   ,       __snake_case							:			str=[7, 3, 3]   ,       __snake_case							:			List[str]=[4, 2, 2]   ,       __snake_case							:			Optional[int]=[2, 1, 1]   ,       __snake_case							:			str=[2, 2, 2]   ,       __snake_case							:			Optional[int]=[False, False, True]   ,       __snake_case							:			Any=[0.0, 0.0, 0.0]   ,       __snake_case							:			Optional[int]=0.02   ,       __snake_case							:			Optional[Any]=1e-12   ,       __snake_case							:			Tuple=True   ,       __snake_case							:			List[str]=True   ,       __snake_case							:			Tuple=2   ,       )->	Optional[int]:
        snake_case 						=   parent
        snake_case 						=   batch_size
        snake_case 						=   image_size
        snake_case 						=   patch_sizes
        snake_case 						=   patch_stride
        snake_case 						=   patch_padding
        snake_case 						=   is_training
        snake_case 						=   use_labels
        snake_case 						=   num_labels
        snake_case 						=   num_channels
        snake_case 						=   embed_dim
        snake_case 						=   num_heads
        snake_case 						=   stride_kv
        snake_case 						=   depth
        snake_case 						=   cls_token
        snake_case 						=   attention_drop_rate
        snake_case 						=   initializer_range
        snake_case 						=   layer_norm_eps



     def      lowerCAmelCase			(   self							:			str						)->	Optional[int]:
        snake_case 						=   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

        snake_case 						=   None
        if self.use_labels:
           # create a random int32 tensor of given shape
           snake_case 						=   ids_tensor([self.batch_size]   ,       self.num_labels						)

        snake_case 						=   self.get_config()
        return config, pixel_values, labels



     def      lowerCAmelCase			(   self							:			Any						)->	int:
        return CvtConfig(
            image_size=self.image_size   ,       num_labels=self.num_labels   ,       num_channels=self.num_channels   ,       embed_dim=self.embed_dim   ,       num_heads=self.num_heads   ,       patch_sizes=self.patch_sizes   ,       patch_padding=self.patch_padding   ,       patch_stride=self.patch_stride   ,       stride_kv=self.stride_kv   ,       depth=self.depth   ,       cls_token=self.cls_token   ,       attention_drop_rate=self.attention_drop_rate   ,       initializer_range=self.initializer_range   ,       )



     def      lowerCAmelCase			(   self							:			List[str]   ,       __snake_case							:			Union[str, Any]   ,       __snake_case							:			int   ,       __snake_case							:			Union[str, Any]						)->	List[Any]:
        snake_case 						=   TFCvtModel(config=__snake_case						)
        snake_case 						=   model(__snake_case   ,       training=__snake_case						)
        snake_case 						=   (self.image_size, self.image_size)
        snake_case						,    snake_case 						=   image_size[0], image_size[1]
        for i in range(len(self.depth						)						):
           snake_case 						=   floor(((height + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1						)
           snake_case 						=   floor(((width + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1						)
        self.parent.assertEqual(result.last_hidden_state.shape   ,       (self.batch_size, self.embed_dim[-1], height, width)						)



     def      lowerCAmelCase			(   self							:			int   ,       __snake_case							:			Union[str, Any]   ,       __snake_case							:			Optional[int]   ,       __snake_case							:			Optional[int]						)->	Union[str, Any]:
        snake_case 						=   self.num_labels
        snake_case 						=   TFCvtForImageClassification(__snake_case						)
        snake_case 						=   model(__snake_case   ,       labels=__snake_case   ,       training=__snake_case						)
        self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, self.num_labels)						)



     def      lowerCAmelCase			(   self							:			int						)->	Optional[Any]:
        snake_case 						=   self.prepare_config_and_inputs()
        snake_case						,    snake_case						,    snake_case 						=   config_and_inputs
        snake_case 						=   {"""pixel_values""": pixel_values}
        return config, inputs_dict




@require_tf
class 			_lowerCAmelCase		(			A__							,     A__							,     unittest.TestCase	):






     """simple docstring"""






     snake_case_						      =   (TFCvtModel, TFCvtForImageClassification) if is_tf_available() else ()
     snake_case_						      =   (
         {"feature-extraction": TFCvtModel, "image-classification": TFCvtForImageClassification}
         if is_tf_available()
         else {}
     )
     snake_case_						      =   False
     snake_case_						      =   False
     snake_case_						      =   False
     snake_case_						      =   False
     snake_case_						      =   False



     def      lowerCAmelCase			(   self							:			Union[str, Any]						)->	Tuple:
        snake_case 						=   TFCvtModelTester(self						)
        snake_case 						=   TFCvtConfigTester(self   ,       config_class=__snake_case   ,       has_text_modality=__snake_case   ,       hidden_size=37						)



     def      lowerCAmelCase			(   self							:			int						)->	Tuple:
        self.config_tester.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()



     @unittest.skip(reason="""Cvt does not output attentions"""						)
     def      lowerCAmelCase			(   self							:			Union[str, Any]						)->	str:
        pass



     @unittest.skip(reason="""Cvt does not use inputs_embeds"""						)
     def      lowerCAmelCase			(   self							:			Dict						)->	int:
        pass



     @unittest.skip(reason="""Cvt does not support input and output embeddings"""						)
     def      lowerCAmelCase			(   self							:			int						)->	Tuple:
        pass



     @unittest.skipIf(
         not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""						)						) == 0   ,       reason="""TF does not support backprop for grouped convolutions on CPU."""   ,       )
     def      lowerCAmelCase			(   self							:			List[Any]						)->	List[str]:
        super().test_dataset_conversion()



     @unittest.skipIf(
         not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""						)						) == 0   ,       reason="""TF does not support backprop for grouped convolutions on CPU."""   ,       )
     @slow
     def      lowerCAmelCase			(   self							:			List[Any]						)->	Tuple:
        super().test_keras_fit()



     @unittest.skip(reason="""Get `Failed to determine best cudnn convolution algo.` error after using TF 2.12+cuda 11.8"""						)
     def      lowerCAmelCase			(   self							:			Tuple						)->	Dict:
        snake_case 						=   tf.keras.mixed_precision.Policy("""mixed_float16"""						)
        tf.keras.mixed_precision.set_global_policy(__snake_case						)
        super().test_keras_fit()
        tf.keras.mixed_precision.set_global_policy("""float32"""						)



     def      lowerCAmelCase			(   self							:			Dict						)->	Union[str, Any]:
        snake_case						,    snake_case 						=   self.model_tester.prepare_config_and_inputs_for_common()

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

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



     def      lowerCAmelCase			(   self							:			List[str]						)->	Tuple:
        def check_hidden_states_output(__snake_case							:			Optional[int]   ,       __snake_case							:			Optional[int]   ,       __snake_case							:			Union[str, Any]						):
           snake_case 						=   model_class(__snake_case						)

           snake_case 						=   model(**self._prepare_for_class(__snake_case   ,       __snake_case						)						)
           snake_case 						=   outputs.hidden_states

           snake_case 						=   len(self.model_tester.depth						)
           self.assertEqual(len(__snake_case						)   ,       __snake_case						)

           # verify the first hidden states (first block)
           self.assertListEqual(
               list(hidden_states[0].shape[-3:]						)   ,       [
                   self.model_tester.embed_dim[0],
                   self.model_tester.image_size // 4,
                   self.model_tester.image_size // 4,
               ]   ,       )

        snake_case						,    snake_case 						=   self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           snake_case 						=   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"]
           snake_case 						=   True

           check_hidden_states_output(__snake_case   ,       __snake_case   ,       __snake_case						)



     def      lowerCAmelCase			(   self							:			Any						)->	Tuple:
        snake_case 						=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__snake_case						)



     def      lowerCAmelCase			(   self							:			List[Any]						)->	List[Any]:
        snake_case 						=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*__snake_case						)



     @slow
     def      lowerCAmelCase			(   self							:			Any						)->	Dict:
        for model_name in TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           snake_case 						=   TFCvtModel.from_pretrained(__snake_case						)
           self.assertIsNotNone(__snake_case						)





def 	__lowerCamelCase			(      )		->  Union[str, Any]:
   snake_case 						=   Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""		)
   return image


@require_tf
@require_vision
class 			_lowerCAmelCase		(			unittest.TestCase	):






     """simple docstring"""






     @cached_property
     def      lowerCAmelCase			(   self							:			int						)->	Tuple:
        return AutoImageProcessor.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]						)



     @slow
     def      lowerCAmelCase			(   self							:			Any						)->	Any:
        snake_case 						=   TFCvtForImageClassification.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]						)

        snake_case 						=   self.default_image_processor
        snake_case 						=   prepare_img()
        snake_case 						=   image_processor(images=__snake_case   ,       return_tensors="""tf"""						)

        # forward pass
        snake_case 						=   model(**__snake_case						)

        # verify the logits
        snake_case 						=   tf.TensorShape((1, 10_00)						)
        self.assertEqual(outputs.logits.shape   ,       __snake_case						)

        snake_case 						=   tf.constant([0.92_85, 0.90_15, -0.31_50]						)
        self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()   ,       __snake_case   ,       atol=1e-4						)						)








'''simple docstring'''




# 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 			_lowerCAmelCase		(			A__	):






     """simple docstring"""






     snake_case_						      =   42
     snake_case_						      =   42




class 			_lowerCAmelCase		(			A__							,     A__	):






     """simple docstring"""






     snake_case_						      =   1



     @register_to_config
     def __init__(   self							:			str   ,       __snake_case							:			int = 20_00   ,       __snake_case							:			float = 0.15   ,       __snake_case							:			float = 0.01   ,       __snake_case							:			float = 13_48.0   ,       __snake_case							:			float = 1e-5   ,       __snake_case							:			int = 1   ,       )->	str:
        # standard deviation of the initial noise distribution
        snake_case 						=   sigma_max

        # setable values
        snake_case 						=   None

        self.set_sigmas(__snake_case   ,       __snake_case   ,       __snake_case   ,       __snake_case						)



     def      lowerCAmelCase			(   self							:			List[Any]   ,       __snake_case							:			torch.FloatTensor   ,       __snake_case							:			Optional[int] = None						)->	torch.FloatTensor:
        return sample



     def      lowerCAmelCase			(   self							:			List[str]   ,       __snake_case							:			int   ,       __snake_case							:			float = None   ,       __snake_case							:			Union[str, torch.device] = None						)->	Optional[Any]:
        snake_case 						=   sampling_eps if sampling_eps is not None else self.config.sampling_eps

        snake_case 						=   torch.linspace(1   ,       __snake_case   ,       __snake_case   ,       device=__snake_case						)



     def      lowerCAmelCase			(   self							:			Union[str, Any]   ,       __snake_case							:			int   ,       __snake_case							:			float = None   ,       __snake_case							:			float = None   ,       __snake_case							:			float = None						)->	str:
        snake_case 						=   sigma_min if sigma_min is not None else self.config.sigma_min
        snake_case 						=   sigma_max if sigma_max is not None else self.config.sigma_max
        snake_case 						=   sampling_eps if sampling_eps is not None else self.config.sampling_eps
        if self.timesteps is None:
           self.set_timesteps(__snake_case   ,       __snake_case						)

        snake_case 						=   sigma_min * (sigma_max / sigma_min) ** (self.timesteps / sampling_eps)
        snake_case 						=   torch.exp(torch.linspace(math.log(__snake_case						)   ,       math.log(__snake_case						)   ,       __snake_case						)						)
        snake_case 						=   torch.tensor([sigma_min * (sigma_max / sigma_min) ** t for t in self.timesteps]						)



     def      lowerCAmelCase			(   self							:			str   ,       __snake_case							:			List[str]   ,       __snake_case							:			str						)->	Optional[int]:
        return torch.where(
            timesteps == 0   ,       torch.zeros_like(t.to(timesteps.device						)						)   ,       self.discrete_sigmas[timesteps - 1].to(timesteps.device						)   ,       )



     def      lowerCAmelCase			(   self							:			int   ,       __snake_case							:			torch.FloatTensor   ,       __snake_case							:			int   ,       __snake_case							:			torch.FloatTensor   ,       __snake_case							:			Optional[torch.Generator] = None   ,       __snake_case							:			bool = True   ,       )->	Union[SdeVeOutput, Tuple]:
        if self.timesteps is None:
           raise ValueError(
               """`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"""						)

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

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

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



     def      lowerCAmelCase			(   self							:			Optional[Any]   ,       __snake_case							:			torch.FloatTensor   ,       __snake_case							:			torch.FloatTensor   ,       __snake_case							:			Optional[torch.Generator] = None   ,       __snake_case							:			bool = True   ,       )->	Union[SchedulerOutput, Tuple]:
        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 						=   randn_tensor(sample.shape   ,       layout=sample.layout   ,       generator=__snake_case						).to(sample.device						)

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

        if not return_dict:
           return (prev_sample,)

        return SchedulerOutput(prev_sample=__snake_case						)



     def      lowerCAmelCase			(   self							:			str   ,       __snake_case							:			torch.FloatTensor   ,       __snake_case							:			torch.FloatTensor   ,       __snake_case							:			torch.FloatTensor   ,       )->	torch.FloatTensor:
        # Make sure sigmas and timesteps have the same device and dtype as original_samples
        snake_case 						=   timesteps.to(original_samples.device						)
        snake_case 						=   self.discrete_sigmas.to(original_samples.device						)[timesteps]
        snake_case 						=   (
            noise * sigmas[:, None, None, None]
            if noise is not None
            else torch.randn_like(__snake_case						) * sigmas[:, None, None, None]
        )
        snake_case 						=   noise + original_samples
        return noisy_samples



     def __len__(   self							:			str						)->	Any:
        return self.config.num_train_timesteps

