




"""simple docstring"""


from typing import List, Optional, Tuple, Union

import torch
from torch import nn
from torch.nn import CrossEntropyLoss

from ... import AutoBackbone
from ...modeling_outputs import SemanticSegmenterOutput
from ...modeling_utils import PreTrainedModel
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
from ...utils.backbone_utils import BackboneMixin
from .configuration_upernet import UperNetConfig


A_      			=       [
    '''openmmlab/upernet-convnext-tiny''',
    # See all UperNet models at https://huggingface.co/models?filter=upernet
]

# General docstring
A_      			=       '''UperNetConfig'''






class       lowercase(  nn.Module				):







 '''simple docstring'''





 def __init__(       self:	str,   a_:	int,   a_:	int,   a_:	Union[int, Tuple[int, int]],   a_:	Union[int, Tuple[int, int], str] = 0,   a_:	bool = False,   a_:	Union[int, Tuple[int, int]] = 1,   ):




   '''simple docstring'''



   super().__init__()
   _snake_case							:	Union[str, Any]			=						nn.Convad(
       in_channels=a_,   out_channels=a_,   kernel_size=a_,   padding=a_,   bias=a_,   dilation=a_,   )
   _snake_case							:	Optional[int]			=						nn.BatchNormad(a_	)
   _snake_case							:	Tuple			=						nn.ReLU()


 def 				UpperCamelCase_     (       self:	Optional[Any],   a_:	torch.Tensor	):




   '''simple docstring'''



   _snake_case							:	Union[str, Any]			=						self.conv(a_	)
   _snake_case							:	Optional[int]			=						self.batch_norm(a_	)
   _snake_case							:	Dict			=						self.activation(a_	)

   return output







class       lowercase(  nn.Module				):







 '''simple docstring'''





 def __init__(       self:	List[str],   a_:	int,   a_:	int,   a_:	int	):




   '''simple docstring'''



   super().__init__()
   _snake_case							:	List[str]			=						[
       nn.AdaptiveAvgPoolad(a_	),
       UperNetConvModule(a_,   a_,   kernel_size=1	),
   ]
   for i, layer in enumerate(self.layers	):
     self.add_module(str(a_	),   a_	)


 def 				UpperCamelCase_     (       self:	List[Any],   a_:	torch.Tensor	):




   '''simple docstring'''



   _snake_case							:	List[Any]			=						input
   for layer in self.layers:
     _snake_case							:	str			=						layer(a_	)
   return hidden_state







class       lowercase(  nn.Module				):







 '''simple docstring'''





 def __init__(       self:	List[str],   a_:	Tuple[int, ...],   a_:	int,   a_:	int,   a_:	bool	):




   '''simple docstring'''



   super().__init__()
   _snake_case							:	List[str]			=						pool_scales
   _snake_case							:	Optional[Any]			=						align_corners
   _snake_case							:	List[Any]			=						in_channels
   _snake_case							:	Dict			=						channels
   _snake_case							:	Optional[int]			=						[]
   for i, pool_scale in enumerate(a_	):
     _snake_case							:	Optional[int]			=						UperNetPyramidPoolingBlock(pool_scale=a_,   in_channels=a_,   channels=a_	)
     self.blocks.append(a_	)
     self.add_module(str(a_	),   a_	)


 def 				UpperCamelCase_     (       self:	Union[str, Any],   a_:	torch.Tensor	):




   '''simple docstring'''



   _snake_case							:	str			=						[]
   for ppm in self.blocks:
     _snake_case							:	Any			=						ppm(a_	)
     _snake_case							:	Dict			=						nn.functional.interpolate(
         a_,   size=x.size()[2:],   mode="""bilinear""",   align_corners=self.align_corners	)
     ppm_outs.append(a_	)
   return ppm_outs







class       lowercase(  nn.Module				):







 '''simple docstring'''





 def __init__(       self:	Any,   a_:	Any,   a_:	Dict	):




   '''simple docstring'''



   super().__init__()

   _snake_case							:	Tuple			=						config
   _snake_case							:	Tuple			=						config.pool_scales  # e.g. (1, 2, 3, 6)
   _snake_case							:	str			=						in_channels
   _snake_case							:	Dict			=						config.hidden_size
   _snake_case							:	Any			=						False
   _snake_case							:	List[Any]			=						nn.Convad(self.channels,   config.num_labels,   kernel_size=1	)

   # PSP Module
   _snake_case							:	List[Any]			=						UperNetPyramidPoolingModule(
       self.pool_scales,   self.in_channels[-1],   self.channels,   align_corners=self.align_corners,   )
   _snake_case							:	int			=						UperNetConvModule(
       self.in_channels[-1] + len(self.pool_scales	) * self.channels,   self.channels,   kernel_size=3,   padding=1,   )
   # FPN Module
   _snake_case							:	str			=						nn.ModuleList()
   _snake_case							:	Tuple			=						nn.ModuleList()
   for in_channels in self.in_channels[:-1]:  # skip the top layer
     _snake_case							:	List[Any]			=						UperNetConvModule(a_,   self.channels,   kernel_size=1	)
     _snake_case							:	Optional[Any]			=						UperNetConvModule(self.channels,   self.channels,   kernel_size=3,   padding=1	)
     self.lateral_convs.append(a_	)
     self.fpn_convs.append(a_	)

   _snake_case							:	str			=						UperNetConvModule(
       len(self.in_channels	) * self.channels,   self.channels,   kernel_size=3,   padding=1,   )


 def 				UpperCamelCase_     (       self:	List[Any]	):




   '''simple docstring'''



   self.apply(self._init_weights	)


 def 				UpperCamelCase_     (       self:	Optional[Any],   a_:	List[Any]	):




   '''simple docstring'''



   if isinstance(a_,   nn.Convad	):
     module.weight.data.normal_(mean=0.0,   std=self.config.initializer_range	)
     if module.bias is not None:
       module.bias.data.zero_()


 def 				UpperCamelCase_     (       self:	List[Any],   a_:	Any	):




   '''simple docstring'''



   _snake_case							:	Optional[int]			=						inputs[-1]
   _snake_case							:	Tuple			=						[x]
   psp_outs.extend(self.psp_modules(a_	)	)
   _snake_case							:	Tuple			=						torch.cat(a_,   dim=1	)
   _snake_case							:	Optional[Any]			=						self.bottleneck(a_	)

   return output


 def 				UpperCamelCase_     (       self:	Optional[Any],   a_:	torch.Tensor	):




   '''simple docstring'''



   _snake_case							:	int			=						[lateral_conv(encoder_hidden_states[i]	) for i, lateral_conv in enumerate(self.lateral_convs	)]

   laterals.append(self.psp_forward(a_	)	)

   # build top-down path
   _snake_case							:	List[str]			=						len(a_	)
   for i in range(used_backbone_levels - 1,   0,   -1	):
     _snake_case							:	Dict			=						laterals[i - 1].shape[2:]
     _snake_case							:	Tuple			=						laterals[i - 1] + nn.functional.interpolate(
         laterals[i],   size=a_,   mode="""bilinear""",   align_corners=self.align_corners	)

   # build outputs
   _snake_case							:	Optional[int]			=						[self.fpn_convs[i](laterals[i]	) for i in range(used_backbone_levels - 1	)]
   # append psp feature
   fpn_outs.append(laterals[-1]	)

   for i in range(used_backbone_levels - 1,   0,   -1	):
     _snake_case							:	int			=						nn.functional.interpolate(
         fpn_outs[i],   size=fpn_outs[0].shape[2:],   mode="""bilinear""",   align_corners=self.align_corners	)
   _snake_case							:	Optional[Any]			=						torch.cat(a_,   dim=1	)
   _snake_case							:	Any			=						self.fpn_bottleneck(a_	)
   _snake_case							:	Dict			=						self.classifier(a_	)

   return output







class       lowercase(  nn.Module				):







 '''simple docstring'''





 def __init__(       self:	Dict,   a_:	int,   a_:	int = 2,   a_:	int = 3,   a_:	Union[int, Tuple[int, int]] = 1	):




   '''simple docstring'''



   super().__init__()

   _snake_case							:	List[Any]			=						config
   _snake_case							:	Dict			=						config.auxiliary_in_channels
   _snake_case							:	Dict			=						config.auxiliary_channels
   _snake_case							:	int			=						config.auxiliary_num_convs
   _snake_case							:	Any			=						config.auxiliary_concat_input
   _snake_case							:	Union[str, Any]			=						in_index

   _snake_case							:	int			=						(kernel_size // 2) * dilation
   _snake_case							:	Tuple			=						[]
   convs.append(
       UperNetConvModule(
           self.in_channels,   self.channels,   kernel_size=a_,   padding=a_,   dilation=a_	)	)
   for i in range(self.num_convs - 1	):
     convs.append(
         UperNetConvModule(
             self.channels,   self.channels,   kernel_size=a_,   padding=a_,   dilation=a_	)	)
   if self.num_convs == 0:
     _snake_case							:	str			=						nn.Identity()
   else:
     _snake_case							:	List[Any]			=						nn.Sequential(*a_	)
   if self.concat_input:
     _snake_case							:	Dict			=						UperNetConvModule(
         self.in_channels + self.channels,   self.channels,   kernel_size=a_,   padding=kernel_size // 2	)

   _snake_case							:	Optional[Any]			=						nn.Convad(self.channels,   config.num_labels,   kernel_size=1	)


 def 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



   self.apply(self._init_weights	)


 def 				UpperCamelCase_     (       self:	Dict,   a_:	List[str]	):




   '''simple docstring'''



   if isinstance(a_,   nn.Convad	):
     module.weight.data.normal_(mean=0.0,   std=self.config.initializer_range	)
     if module.bias is not None:
       module.bias.data.zero_()


 def 				UpperCamelCase_     (       self:	List[Any],   a_:	torch.Tensor	):




   '''simple docstring'''



   _snake_case							:	Any			=						encoder_hidden_states[self.in_index]
   _snake_case							:	Optional[Any]			=						self.convs(a_	)
   if self.concat_input:
     _snake_case							:	Any			=						self.conv_cat(torch.cat([hidden_states, output],   dim=1	)	)
   _snake_case							:	str			=						self.classifier(a_	)
   return output







class       lowercase(  __a				):







 '''simple docstring'''





 lowercase__				      =      UperNetConfig
 lowercase__				      =      "pixel_values"
 lowercase__				      =      True


 def 				UpperCamelCase_     (       self:	Dict,   a_:	Optional[int]	):




   '''simple docstring'''



   if isinstance(a_,   a_	):
     module.backbone.init_weights()
     module.decode_head.init_weights()
     module.auxiliary_head.init_weights()


 def 				UpperCamelCase_     (       self:	str	):




   '''simple docstring'''



   self.backbone.init_weights()
   self.decode_head.init_weights()
   self.auxiliary_head.init_weights()


 def 				UpperCamelCase_     (       self:	int,   a_:	Union[str, Any],   a_:	str=False	):




   '''simple docstring'''



   if isinstance(a_,   a_	):
     _snake_case							:	Union[str, Any]			=						value







A_      			=       r'''
    Parameters:
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.
        config ([`UperNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
'''

A_      			=       r'''
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
            [`AutoImageProcessor`]. See [`SegformerImageProcessor.__call__`] for details.
        output_attentions (`bool`, *optional*):
            Whether or not to return the attentions tensors of all attention layers in case the backbone has them. See
            `attentions` under returned tensors for more detail.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers of the backbone. See `hidden_states` under
            returned tensors for more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
'''






@add_start_docstrings(
    "UperNet framework leveraging any vision backbone e.g. for ADE20k, CityScapes." ,    __a ,    )
class       lowercase(  __a				):







 '''simple docstring'''





 def __init__(       self:	Union[str, Any],   a_:	Any	):




   '''simple docstring'''



   super().__init__(a_	)

   _snake_case							:	Tuple			=						AutoBackbone.from_config(config.backbone_config	)

   # Semantic segmentation head(s)
   _snake_case							:	int			=						UperNetHead(a_,   in_channels=self.backbone.channels	)
   _snake_case							:	Tuple			=						UperNetFCNHead(a_	) if config.use_auxiliary_head else None

   # Initialize weights and apply final processing
   self.post_init()


 @add_start_docstrings_to_model_forward(UPERNET_INPUTS_DOCSTRING.format("""batch_size, sequence_length"""	)	)
 @replace_return_docstrings(output_type=a_,   config_class=_CONFIG_FOR_DOC	)
 def 				UpperCamelCase_     (       self:	Union[str, Any],   a_:	Optional[torch.Tensor] = None,   a_:	Optional[bool] = None,   a_:	Optional[bool] = None,   a_:	Optional[torch.Tensor] = None,   a_:	Optional[bool] = None,   ):




   '''simple docstring'''



   _snake_case							:	Dict			=						return_dict if return_dict is not None else self.config.use_return_dict
   _snake_case							:	Optional[int]			=						(
       output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
   )
   _snake_case							:	Tuple			=						output_attentions if output_attentions is not None else self.config.output_attentions

   _snake_case							:	str			=						self.backbone.forward_with_filtered_kwargs(
       a_,   output_hidden_states=a_,   output_attentions=a_	)
   _snake_case							:	Dict			=						outputs.feature_maps

   _snake_case							:	Optional[Any]			=						self.decode_head(a_	)
   _snake_case							:	int			=						nn.functional.interpolate(a_,   size=pixel_values.shape[2:],   mode="""bilinear""",   align_corners=a_	)

   _snake_case							:	Dict			=						None
   if self.auxiliary_head is not None:
     _snake_case							:	Optional[Any]			=						self.auxiliary_head(a_	)
     _snake_case							:	List[str]			=						nn.functional.interpolate(
         a_,   size=pixel_values.shape[2:],   mode="""bilinear""",   align_corners=a_	)

   _snake_case							:	Any			=						None
   if labels is not None:
     if self.config.num_labels == 1:
       raise ValueError("""The number of labels should be greater than one"""	)
     else:
       # compute weighted loss
       _snake_case							:	str			=						CrossEntropyLoss(ignore_index=self.config.loss_ignore_index	)
       _snake_case							:	Any			=						loss_fct(a_,   a_	)
       _snake_case							:	List[str]			=						loss_fct(a_,   a_	)
       _snake_case							:	Optional[Any]			=						main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

   if not return_dict:
     if output_hidden_states:
       _snake_case							:	List[str]			=						(logits,) + outputs[1:]
     else:
       _snake_case							:	str			=						(logits,) + outputs[2:]
     return ((loss,) + output) if loss is not None else output

   return SemanticSegmenterOutput(
       loss=a_,   logits=a_,   hidden_states=outputs.hidden_states,   attentions=outputs.attentions,   )








"""simple docstring"""


import os
import sys
import unittest


A_      			=       os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, '''utils'''))

import check_dummies  # noqa: E402
from check_dummies import create_dummy_files, create_dummy_object, find_backend, read_init  # noqa: E402


# Align TRANSFORMERS_PATH in check_dummies with the current path
A_      			=       os.path.join(git_repo_path, '''src''', '''diffusers''')






class       lowercase(  unittest.TestCase				):







 '''simple docstring'''





 def 				UpperCamelCase_     (       self:	Tuple	):




   '''simple docstring'''



   _snake_case							:	Optional[int]			=						find_backend("""    if not is_torch_available():"""	)
   self.assertEqual(a_,   """torch"""	)

   # backend_with_underscore = find_backend("    if not is_tensorflow_text_available():")
   # self.assertEqual(backend_with_underscore, "tensorflow_text")

   _snake_case							:	Any			=						find_backend("""    if not (is_torch_available() and is_transformers_available()):"""	)
   self.assertEqual(a_,   """torch_and_transformers"""	)

   # double_backend_with_underscore = find_backend(
   #    "    if not (is_sentencepiece_available() and is_tensorflow_text_available()):"
   # )
   # self.assertEqual(double_backend_with_underscore, "sentencepiece_and_tensorflow_text")

   _snake_case							:	Union[str, Any]			=						find_backend(
       """    if not (is_torch_available() and is_transformers_available() and is_onnx_available()):"""	)
   self.assertEqual(a_,   """torch_and_transformers_and_onnx"""	)


 def 				UpperCamelCase_     (       self:	List[Any]	):




   '''simple docstring'''



   _snake_case							:	Dict			=						read_init()
   # We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
   self.assertIn("""torch""",   a_	)
   self.assertIn("""torch_and_transformers""",   a_	)
   self.assertIn("""flax_and_transformers""",   a_	)
   self.assertIn("""torch_and_transformers_and_onnx""",   a_	)

   # Likewise, we can't assert on the exact content of a key
   self.assertIn("""UNet2DModel""",   objects["""torch"""]	)
   self.assertIn("""FlaxUNet2DConditionModel""",   objects["""flax"""]	)
   self.assertIn("""StableDiffusionPipeline""",   objects["""torch_and_transformers"""]	)
   self.assertIn("""FlaxStableDiffusionPipeline""",   objects["""flax_and_transformers"""]	)
   self.assertIn("""LMSDiscreteScheduler""",   objects["""torch_and_scipy"""]	)
   self.assertIn("""OnnxStableDiffusionPipeline""",   objects["""torch_and_transformers_and_onnx"""]	)


 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   _snake_case							:	Optional[Any]			=						create_dummy_object("""CONSTANT""",   """'torch'"""	)
   self.assertEqual(a_,   """\nCONSTANT = None\n"""	)

   _snake_case							:	Optional[int]			=						create_dummy_object("""function""",   """'torch'"""	)
   self.assertEqual(
       a_,   """\ndef function(*args, **kwargs):\n    requires_backends(function, 'torch')\n"""	)

   _snake_case							:	List[Any]			=						"""
class FakeClass(metaclass=DummyObject):
    _backends = 'torch'

    def __init__(self, *args, **kwargs):
        requires_backends(self, 'torch')

    @classmethod
    def from_config(cls, *args, **kwargs):
        requires_backends(cls, 'torch')

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        requires_backends(cls, 'torch')
"""
   _snake_case							:	Union[str, Any]			=						create_dummy_object("""FakeClass""",   """'torch'"""	)
   self.assertEqual(a_,   a_	)


 def 				UpperCamelCase_     (       self:	Union[str, Any]	):




   '''simple docstring'''



   _snake_case							:	Union[str, Any]			=						"""# This file is autogenerated by the command `make fix-copies`, do not edit.
from ..utils import DummyObject, requires_backends


CONSTANT = None


def function(*args, **kwargs):
    requires_backends(function, [\"torch\"])


class FakeClass(metaclass=DummyObject):
    _backends = [\"torch\"]

    def __init__(self, *args, **kwargs):
        requires_backends(self, [\"torch\"])

    @classmethod
    def from_config(cls, *args, **kwargs):
        requires_backends(cls, [\"torch\"])

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        requires_backends(cls, [\"torch\"])
"""
   _snake_case							:	List[Any]			=						create_dummy_files({"""torch""": ["""CONSTANT""", """function""", """FakeClass"""]}	)
   self.assertEqual(dummy_files["""torch"""],   a_	)


