

"""simple docstring"""



import inspect
import unittest

from transformers import EfficientNetConfig
from transformers.testing_utils import is_pipeline_test, 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 EfficientNetForImageClassification, EfficientNetModel
 from transformers.models.efficientnet.modeling_efficientnet import EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor



class        a    :



    """simple docstring"""



    def __init__(  self:			str     ,    UpperCamelCase:			Dict     ,    UpperCamelCase:			Optional[Any]=13     ,    UpperCamelCase:			Optional[int]=32     ,    UpperCamelCase:			str=3     ,    UpperCamelCase:			Optional[Any]=[3, 3, 5]     ,    UpperCamelCase:			int=[32, 16, 24]     ,    UpperCamelCase:			Optional[int]=[16, 24, 40]     ,    UpperCamelCase:			List[Any]=[1, 1, 2]     ,    UpperCamelCase:			Any=[1, 1, 2]     ,    UpperCamelCase:			str=[1, 6, 6]     ,    UpperCamelCase:			str=True     ,    UpperCamelCase:			int=True     ,    UpperCamelCase:			List[Any]=37     ,    UpperCamelCase:			str="gelu"     ,    UpperCamelCase:			Optional[Any]=10     ,    ):






         """simple docstring"""


         A__				      =	parent
         A__				      =	batch_size
         A__				      =	image_size
         A__				      =	num_channels
         A__				      =	kernel_sizes
         A__				      =	in_channels
         A__				      =	out_channels
         A__				      =	strides
         A__				      =	num_block_repeats
         A__				      =	expand_ratios
         A__				      =	is_training
         A__				      =	hidden_act
         A__				      =	num_labels
         A__				      =	use_labels

    def 					UpperCamelCase						(  self:			Tuple							):






         """simple docstring"""


         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 					UpperCamelCase						(  self:			Optional[Any]							):






         """simple docstring"""


         return EfficientNetConfig(
             num_channels=self.num_channels     ,    kernel_sizes=self.kernel_sizes     ,    in_channels=self.in_channels     ,    out_channels=self.out_channels     ,    strides=self.strides     ,    num_block_repeats=self.num_block_repeats     ,    expand_ratios=self.expand_ratios     ,    hidden_act=self.hidden_act     ,    num_labels=self.num_labels     ,    )

    def 					UpperCamelCase						(  self:			List[Any]     ,    UpperCamelCase:			List[str]     ,    UpperCamelCase:			Union[str, Any]     ,    UpperCamelCase:			List[str]							):






         """simple docstring"""


         A__				      =	EfficientNetModel(config=UpperCamelCase							)
         model.to(UpperCamelCase							)
         model.eval()
         A__				      =	model(UpperCamelCase							)
         # expected last hidden states: B, C, H // 4, W // 4
         self.parent.assertEqual(
             result.last_hidden_state.shape     ,    (self.batch_size, config.hidden_dim, self.image_size // 4, self.image_size // 4)     ,    )

    def 					UpperCamelCase						(  self:			str     ,    UpperCamelCase:			int     ,    UpperCamelCase:			Any     ,    UpperCamelCase:			List[str]							):






         """simple docstring"""


         A__				      =	EfficientNetForImageClassification(UpperCamelCase							)
         model.to(UpperCamelCase							)
         model.eval()
         A__				      =	model(UpperCamelCase     ,    labels=UpperCamelCase							)
         self.parent.assertEqual(result.logits.shape     ,    (self.batch_size, self.num_labels)							)





    def 					UpperCamelCase						(  self:			int							):






         """simple docstring"""


         A__				      =	self.prepare_config_and_inputs()
         A__      ,					A__      ,					A__				      =	config_and_inputs
         A__				      =	{"""pixel_values""": pixel_values}
         return config, inputs_dict



@require_torch
class        a    (							__snake_case,						__snake_case,						unittest.TestCase					):



    """simple docstring"""



    UpperCAmelCase    			=       (EfficientNetModel, EfficientNetForImageClassification) if is_torch_available() else ()
    UpperCAmelCase    			=       (
        {'feature-extraction': EfficientNetModel, 'image-classification': EfficientNetForImageClassification}
        if is_torch_available()
        else {}
    )

    UpperCAmelCase    			=       False
    UpperCAmelCase    			=       False
    UpperCAmelCase    			=       False
    UpperCAmelCase    			=       False
    UpperCAmelCase    			=       False

    def 					UpperCamelCase						(  self:			int							):






         """simple docstring"""


         A__				      =	EfficientNetModelTester(self							)
         A__				      =	ConfigTester(
             self     ,    config_class=UpperCamelCase     ,    has_text_modality=UpperCamelCase     ,    hidden_size=37							)

    def 					UpperCamelCase						(  self:			List[str]							):






         """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 					UpperCamelCase						(  self:			List[str]							):






         """simple docstring"""


         return

    @unittest.skip(reason="""EfficientNet does not use inputs_embeds"""							)
    def 					UpperCamelCase						(  self:			Tuple							):






         """simple docstring"""


         pass

    @unittest.skip(reason="""EfficientNet does not support input and output embeddings"""							)
    def 					UpperCamelCase						(  self:			Dict							):






         """simple docstring"""


         pass

    @unittest.skip(reason="""EfficientNet does not use feedforward chunking"""							)
    def 					UpperCamelCase						(  self:			int							):






         """simple docstring"""


         pass

    def 					UpperCamelCase						(  self:			Optional[Any]							):






         """simple docstring"""


         A__      ,					A__				      =	self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
              A__				      =	model_class(UpperCamelCase							)
              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]     ,    UpperCamelCase							)

    def 					UpperCamelCase						(  self:			int							):






         """simple docstring"""


         A__				      =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*UpperCamelCase							)

    def 					UpperCamelCase						(  self:			int							):






         """simple docstring"""


         def check_hidden_states_output(UpperCamelCase:			Optional[int]     ,    UpperCamelCase:			List[Any]     ,    UpperCamelCase:			Union[str, Any]							):
              A__				      =	model_class(UpperCamelCase							)
              model.to(UpperCamelCase							)
              model.eval()

              with torch.no_grad():
                   A__				      =	model(**self._prepare_for_class(UpperCamelCase     ,    UpperCamelCase							)							)

              A__				      =	outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states
              A__				      =	sum(config.num_block_repeats							) * 4
              self.assertEqual(len(UpperCamelCase							)     ,    UpperCamelCase							)

              # EfficientNet'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 // 2, self.model_tester.image_size // 2]     ,    )

         A__      ,					A__				      =	self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
              A__				      =	True
              check_hidden_states_output(UpperCamelCase     ,    UpperCamelCase     ,    UpperCamelCase							)

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

              check_hidden_states_output(UpperCamelCase     ,    UpperCamelCase     ,    UpperCamelCase							)

    def 					UpperCamelCase						(  self:			Optional[int]							):






         """simple docstring"""


         A__				      =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*UpperCamelCase							)

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






         """simple docstring"""


         pass

    @slow
    def 					UpperCamelCase						(  self:			Dict							):






         """simple docstring"""


         for model_name in EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              A__				      =	EfficientNetModel.from_pretrained(UpperCamelCase							)
              self.assertIsNotNone(UpperCamelCase							)





    @is_pipeline_test
    @require_vision
    @slow
    def 					UpperCamelCase						(  self:			str							):






         """simple docstring"""


         super().test_pipeline_image_classification()





def 			_snake_case  (     ):
     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 					UpperCamelCase						(  self:			Optional[Any]							):






         """simple docstring"""


         return AutoImageProcessor.from_pretrained("""google/efficientnet-b7"""							) if is_vision_available() else None





    @slow
    def 					UpperCamelCase						(  self:			Union[str, Any]							):






         """simple docstring"""


         A__				      =	EfficientNetForImageClassification.from_pretrained("""google/efficientnet-b7"""							).to(UpperCamelCase							)

         A__				      =	self.default_image_processor
         A__				      =	prepare_img()
         A__				      =	image_processor(images=UpperCamelCase     ,    return_tensors="""pt"""							).to(UpperCamelCase							)

         # forward pass
         with torch.no_grad():
              A__				      =	model(**UpperCamelCase							)

         # verify the logits
         A__				      =	torch.Size((1, 10_00)							)
         self.assertEqual(outputs.logits.shape     ,    UpperCamelCase							)

         A__				      =	torch.tensor([0.0_001, 0.0_002, 0.0_002]							).to(UpperCamelCase							)
         self.assertTrue(torch.allclose(outputs.logits[0, :3]     ,    UpperCamelCase     ,    atol=1e-4							)							)









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

import flax
import jax.numpy as jnp
from scipy import integrate

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils_flax import (
    CommonSchedulerState,
    FlaxKarrasDiffusionSchedulers,
    FlaxSchedulerMixin,
    FlaxSchedulerOutput,
    broadcast_to_shape_from_left,
)







@flax.struct.dataclass
class 					__lowerCamelCase							:
      lowerCamelCase_      : CommonSchedulerState

      # setable values
      lowerCamelCase_      : jnp.ndarray
      lowerCamelCase_      : jnp.ndarray
      lowerCamelCase_      : jnp.ndarray
      lowerCamelCase_      : Optional[int]  			=   None

      # running values
      lowerCamelCase_      : Optional[jnp.ndarray]  			=   None







      @classmethod
      def       lowerCAmelCase_       (	cls				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )		->			Union[str, Any]:
        return cls(common=lowerCamelCase				,    init_noise_sigma=lowerCamelCase				,    timesteps=lowerCamelCase				,    sigmas=lowerCamelCase  )







@dataclass
class 					__lowerCamelCase							(      __snake_case      ):
      lowerCamelCase_      : LMSDiscreteSchedulerState







class 					__lowerCamelCase							(      __snake_case    ,				__snake_case      ):
      lowerCamelCase_      : str  			=   [e.name for e in FlaxKarrasDiffusionSchedulers]

      lowerCamelCase_      : jnp.dtype







      @property
      def       lowerCAmelCase_       (	self  )		->			List[str]:
        return True







      @register_to_config
      def __init__(	self				,    lowerCamelCase = 1000				,    lowerCamelCase = 0.0001				,    lowerCamelCase = 0.02				,    lowerCamelCase = "linear"				,    lowerCamelCase = None				,    lowerCamelCase = "epsilon"				,    lowerCamelCase = jnp.floataa				,    )		->			Optional[Any]:
        snake_case_					= dtype







      def       lowerCAmelCase_       (	self				,    lowerCamelCase = None  )		->			LMSDiscreteSchedulerState:
        if common is None:
          snake_case_					= CommonSchedulerState.create(self  )

        snake_case_					= jnp.arange(0				,    self.config.num_train_timesteps  ).round()[::-1]
        snake_case_					= ((1 - common.alphas_cumprod) / common.alphas_cumprod) ** 0.5

        # standard deviation of the initial noise distribution
        snake_case_					= sigmas.max()

        return LMSDiscreteSchedulerState.create(
            common=lowerCamelCase				,    init_noise_sigma=lowerCamelCase				,    timesteps=lowerCamelCase				,    sigmas=lowerCamelCase				,    )







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )		->			jnp.ndarray:
        ((snake_case_)		,							)					= jnp.where(state.timesteps == timestep				,    size=1  )
        snake_case_					= step_index[0]

        snake_case_					= state.sigmas[step_index]
        snake_case_					= sample / ((sigma**2 + 1) ** 0.5)
        return sample







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )		->			Any:

        def lms_derivative(lowerCamelCase  ):
          snake_case_					= 1.0
          for k in range(lowerCamelCase  ):
            if current_order == k:
              continue
            prod *= (tau - state.sigmas[t - k]) / (state.sigmas[t - current_order] - state.sigmas[t - k])
          return prod

        snake_case_					= integrate.quad(lowerCamelCase				,    state.sigmas[t]				,    state.sigmas[t + 1]				,    epsrel=1e-4  )[0]

        return integrated_coeff







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase = ()  )		->			LMSDiscreteSchedulerState:
        snake_case_					= jnp.linspace(self.config.num_train_timesteps - 1				,    0				,    lowerCamelCase				,    dtype=self.dtype  )

        snake_case_					= jnp.floor(lowerCamelCase  ).astype(jnp.intaa  )
        snake_case_					= jnp.ceil(lowerCamelCase  ).astype(jnp.intaa  )

        snake_case_					= jnp.mod(lowerCamelCase				,    1.0  )

        snake_case_					= ((1 - state.common.alphas_cumprod) / state.common.alphas_cumprod) ** 0.5
        snake_case_					= (1 - frac) * sigmas[low_idx] + frac * sigmas[high_idx]
        snake_case_					= jnp.concatenate([sigmas, jnp.array([0.0]				,    dtype=self.dtype  )]  )

        snake_case_					= timesteps.astype(jnp.intaa  )

        # initial running values
        snake_case_					= jnp.zeros((0,) + shape				,    dtype=self.dtype  )

        return state.replace(
            timesteps=lowerCamelCase				,    sigmas=lowerCamelCase				,    num_inference_steps=lowerCamelCase				,    derivatives=lowerCamelCase				,    )







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase = 4				,    lowerCamelCase = True				,    )		->			Union[FlaxLMSSchedulerOutput, Tuple]:
        if state.num_inference_steps is None:
          raise ValueError(
              """Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"""  )

        snake_case_					= state.sigmas[timestep]

        # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
        if self.config.prediction_type == "epsilon":
          snake_case_					= sample - sigma * model_output
        elif self.config.prediction_type == "v_prediction":
          # * c_out + input * c_skip
          snake_case_					= model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))
        else:
          raise ValueError(
              f'''prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`'''  )

        # 2. Convert to an ODE derivative
        snake_case_					= (sample - pred_original_sample) / sigma
        snake_case_					= state.replace(derivatives=jnp.append(state.derivatives				,    lowerCamelCase  )  )
        if len(state.derivatives  ) > order:
          snake_case_					= state.replace(derivatives=jnp.delete(state.derivatives				,    0  )  )

        # 3. Compute linear multistep coefficients
        snake_case_					= min(timestep + 1				,    lowerCamelCase  )
        snake_case_					= [self.get_lms_coefficient(lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  ) for curr_order in range(lowerCamelCase  )]

        # 4. Compute previous sample based on the derivatives path
        snake_case_					= sample + sum(
            coeff * derivative for coeff, derivative in zip(lowerCamelCase				,    reversed(state.derivatives  )  )  )

        if not return_dict:
          return (prev_sample, state)

        return FlaxLMSSchedulerOutput(prev_sample=lowerCamelCase				,    state=lowerCamelCase  )







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    )		->			jnp.ndarray:
        snake_case_					= state.sigmas[timesteps].flatten()
        snake_case_					= broadcast_to_shape_from_left(lowerCamelCase				,    noise.shape  )

        snake_case_					= original_samples + noise * sigma

        return noisy_samples







      def __len__(	self  )		->			Optional[Any]:
        return self.config.num_train_timesteps