





from typing import Tuple, Union

from ...modeling_outputs import BackboneOutput
from ...modeling_utils import PreTrainedModel
from ...utils import is_timm_available, is_torch_available, requires_backends
from ...utils.backbone_utils import BackboneMixin
from .configuration_timm_backbone import TimmBackboneConfig


if is_timm_available():
     import timm


if is_torch_available():
     from torch import Tensor



class 							_UpperCAmelCase       (       _UpperCamelCase						,       _UpperCamelCase       ):

  """simple docstring"""


  a_              =					"""pixel_values"""
  a_              =					False
  a_              =					TimmBackboneConfig






  def __init__(							self					:						Tuple   ,    lowerCAmelCase_					:						Any   ,    **lowerCAmelCase_					:						Optional[int]	)				->		Optional[Any]:
     requires_backends(self   ,    'timm'	)
     super().__init__(lowerCAmelCase_	)
     __lowerCAmelCase				      =      config

     if config.backbone is None:
        raise ValueError('backbone is not set in the config. Please set it to a timm model name.'	)

     if config.backbone not in timm.list_models():
        raise ValueError(f"""backbone {config.backbone} is not supported by timm."""	)

     if hasattr(lowerCAmelCase_   ,    'out_features'	) and config.out_features is not None:
        raise ValueError('out_features is not supported by TimmBackbone. Please use out_indices instead.'	)

     __lowerCAmelCase				      =      getattr(lowerCAmelCase_   ,    'use_pretrained_backbone'   ,    lowerCAmelCase_	)
     if pretrained is None:
        raise ValueError('use_pretrained_backbone is not set in the config. Please set it to True or False.'	)

     # We just take the final layer by default. This matches the default for the transformers models.
     __lowerCAmelCase				      =      config.out_indices if getattr(lowerCAmelCase_   ,    'out_indices'   ,    lowerCAmelCase_	) is not None else (-1,)

     __lowerCAmelCase				      =      timm.create_model(
         config.backbone   ,    pretrained=lowerCAmelCase_   ,    features_only=config.features_only   ,    in_chans=config.num_channels   ,    out_indices=lowerCAmelCase_   ,    **lowerCAmelCase_   ,    )
     # These are used to control the output of the model when called. If output_hidden_states is True, then
     # return_layers is modified to include all layers.
     __lowerCAmelCase				      =      self._backbone.return_layers
     __lowerCAmelCase				      =      {layer['module']: str(lowerCAmelCase_	) for i, layer in enumerate(self._backbone.feature_info.info	)}
     super()._init_backbone(lowerCAmelCase_	)






  @classmethod
  def 							lowercase						(							cls					:						int   ,    lowerCAmelCase_					:						Dict   ,    *lowerCAmelCase_					:						Dict   ,    **lowerCAmelCase_					:						Union[str, Any]	)				->		Optional[int]:
     requires_backends(cls   ,    ['vision', 'timm']	)
     from ...models.timm_backbone import TimmBackboneConfig

     __lowerCAmelCase				      =      kwargs.pop('config'   ,    TimmBackboneConfig()	)

     __lowerCAmelCase				      =      kwargs.pop('use_timm_backbone'   ,    lowerCAmelCase_	)
     if not use_timm:
        raise ValueError('use_timm_backbone must be True for timm backbones'	)

     __lowerCAmelCase				      =      kwargs.pop('num_channels'   ,    config.num_channels	)
     __lowerCAmelCase				      =      kwargs.pop('features_only'   ,    config.features_only	)
     __lowerCAmelCase				      =      kwargs.pop('use_pretrained_backbone'   ,    config.use_pretrained_backbone	)
     __lowerCAmelCase				      =      kwargs.pop('out_indices'   ,    config.out_indices	)
     __lowerCAmelCase				      =      TimmBackboneConfig(
         backbone=lowerCAmelCase_   ,    num_channels=lowerCAmelCase_   ,    features_only=lowerCAmelCase_   ,    use_pretrained_backbone=lowerCAmelCase_   ,    out_indices=lowerCAmelCase_   ,    )
     return super()._from_config(lowerCAmelCase_   ,    **lowerCAmelCase_	)






  def 							lowercase						(							self					:						Tuple   ,    lowerCAmelCase_					:						int	)				->		Dict:
     pass






  def 							lowercase						(							self					:						Optional[int]   ,    lowerCAmelCase_					:						Union[str, Any]   ,    lowerCAmelCase_					:						List[Any]=None   ,    lowerCAmelCase_					:						List[str]=None   ,    lowerCAmelCase_					:						Any=None   ,    **lowerCAmelCase_					:						Dict	)				->		Union[BackboneOutput, Tuple[Tensor, ...]]:
     __lowerCAmelCase				      =      return_dict if return_dict is not None else self.config.use_return_dict
     __lowerCAmelCase				      =      (
         output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
     )
     __lowerCAmelCase				      =      output_attentions if output_attentions is not None else self.config.output_attentions

     if output_attentions:
        raise ValueError('Cannot output attentions for timm backbones at the moment'	)

     if output_hidden_states:
        # We modify the return layers to include all the stages of the backbone
        __lowerCAmelCase				      =      self._all_layers
        __lowerCAmelCase				      =      self._backbone(lowerCAmelCase_   ,    **lowerCAmelCase_	)
        __lowerCAmelCase				      =      self._return_layers
        __lowerCAmelCase				      =      tuple(hidden_states[i] for i in self.out_indices	)
     else:
        __lowerCAmelCase				      =      self._backbone(lowerCAmelCase_   ,    **lowerCAmelCase_	)
        __lowerCAmelCase				      =      None

     __lowerCAmelCase				      =      tuple(lowerCAmelCase_	)
     __lowerCAmelCase				      =      tuple(lowerCAmelCase_	) if hidden_states is not None else None

     if not return_dict:
        __lowerCAmelCase				      =      (feature_maps,)
        if output_hidden_states:
           __lowerCAmelCase				      =      output + (hidden_states,)
        return output

     return BackboneOutput(feature_maps=lowerCAmelCase_   ,    hidden_states=lowerCAmelCase_   ,    attentions=lowerCAmelCase_	)









'''simple docstring'''







import functools
import gc
import inspect

import torch

from .imports import is_npu_available, is_xpu_available


def       __UpperCamelCase					(      *lowercase_       :  Union[str, Any]   ):
       """simple docstring"""




       if not isinstance(lowercase_   ,       lowercase_   ):
              a_						     =			list(lowercase_   )
       for i in range(len(lowercase_   )   ):
              a_						     =			None
       gc.collect()
       if is_xpu_available():
              torch.xpu.empty_cache()
       elif is_npu_available():
              torch.npu.empty_cache()
       else:
              torch.cuda.empty_cache()
       return objects


def       __UpperCamelCase					(      lowercase_       :  Exception   ):
       """simple docstring"""




       a_						     =			[
           'CUDA out of memory.',  # CUDA OOM
           'cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.',  # CUDNN SNAFU
           'DefaultCPUAllocator: can\'t allocate memory',  # CPU OOM
       ]
       if isinstance(lowercase_   ,       lowercase_   ) and len(exception.args   ) == 1:
              return any(err in exception.args[0] for err in _statements   )
       return False


def       __UpperCamelCase					(      lowercase_       :  callable = None   ,       lowercase_       :  int = 128   ):
       """simple docstring"""




       if function is None:
              return functools.partial(lowercase_   ,       starting_batch_size=lowercase_   )

       a_						     =			starting_batch_size

       def decorator(*lowercase_       :  str   ,       **lowercase_       :  Dict   ):
              nonlocal batch_size
              gc.collect()
              if is_xpu_available():
                     torch.xpu.empty_cache()
              elif is_npu_available():
                     torch.npu.empty_cache()
              else:
                     torch.cuda.empty_cache()
              a_						     =			list(inspect.signature(lowercase_   ).parameters.keys()   )
              # Guard against user error
              if len(lowercase_   ) < (len(lowercase_   ) + 1):
                     a_						     =			', '.join([F'{arg}={value}' for arg, value in zip(params[1:]   ,       args[1:]   )]   )
                     raise TypeError(
                         F'Batch size was passed into `{function.__name__}` as the first argument when called.'
                         F'Remove this as the decorator already does so: `{function.__name__}({arg_str})`'   )
              while True:
                     if batch_size == 0:
                            raise RuntimeError('No executable batch size found, reached zero.'   )
                     try:
                            return function(lowercase_   ,       *lowercase_   ,       **lowercase_   )
                     except Exception as e:
                            if should_reduce_batch_size(lowercase_   ):
                                   gc.collect()
                                   if is_xpu_available():
                                          torch.xpu.empty_cache()
                                   elif is_npu_available():
                                          torch.npu.empty_cache()
                                   else:
                                          torch.cuda.empty_cache()
                                   batch_size //= 2
                            else:
                                   raise

       return decorator




