


from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


__UpperCAmelCase : str 			=				logging.get_logger(__name__)  # pylint: disable=invalid-name

__UpperCAmelCase : Optional[int] 			=				"\n    Examples:\n        ```py\n        >>> import torch\n        >>> import numpy as np\n\n        >>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22ControlnetPipeline\n        >>> from transformers import pipeline\n        >>> from diffusers.utils import load_image\n\n\n        >>> def make_hint(image, depth_estimator):\n        ...     image = depth_estimator(image)[\"depth\"]\n        ...     image = np.array(image)\n        ...     image = image[:, :, None]\n        ...     image = np.concatenate([image, image, image], axis=2)\n        ...     detected_map = torch.from_numpy(image).float() / 255.0\n        ...     hint = detected_map.permute(2, 0, 1)\n        ...     return hint\n\n\n        >>> depth_estimator = pipeline(\"depth-estimation\")\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\n        ...     \"kandinsky-community/kandinsky-2-2-prior\", torch_dtype=torch.float16\n        ... )\n        >>> pipe_prior = pipe_prior.to(\"cuda\")\n\n        >>> pipe = KandinskyV22ControlnetPipeline.from_pretrained(\n        ...     \"kandinsky-community/kandinsky-2-2-controlnet-depth\", torch_dtype=torch.float16\n        ... )\n        >>> pipe = pipe.to(\"cuda\")\n\n\n        >>> img = load_image(\n        ...     \"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main\"\n        ...     \"/kandinsky/cat.png\"\n        ... ).resize((768, 768))\n\n        >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to(\"cuda\")\n\n        >>> prompt = \"A robot, 4k photo\"\n        >>> negative_prior_prompt = \"lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature\"\n\n        >>> generator = torch.Generator(device=\"cuda\").manual_seed(43)\n\n        >>> image_emb, zero_image_emb = pipe_prior(\n        ...     prompt=prompt, negative_prompt=negative_prior_prompt, generator=generator\n        ... ).to_tuple()\n\n        >>> images = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     hint=hint,\n        ...     num_inference_steps=50,\n        ...     generator=generator,\n        ...     height=768,\n        ...     width=768,\n        ... ).images\n\n        >>> images[0].save(\"robot_cat.png\")\n        ```\n"


def      a	(							SCREAMING_SNAKE_CASE_      :		Tuple							,			SCREAMING_SNAKE_CASE_      :		Optional[int]							,			SCREAMING_SNAKE_CASE_      :		str=8       ):






       """simple docstring"""





       UpperCamelCase    :	List[str]   	=		height // scale_factor**2
       if height % scale_factor**2 != 0:
              new_height += 1
       UpperCamelCase    :	List[str]   	=		width // scale_factor**2
       if width % scale_factor**2 != 0:
              new_width += 1
       return new_height * scale_factor, new_width * scale_factor



class    UpperCAmelCase_	(							_a):
       '''simple docstring'''






       def __init__(     self ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					):




              """simple docstring"""




              super().__init__()

              self.register_modules(
                  unet=__SCREAMING_SNAKE_CASE ,					scheduler=__SCREAMING_SNAKE_CASE ,					movq=__SCREAMING_SNAKE_CASE ,					)
              UpperCamelCase    :	Any   	=		2 ** (len(self.movq.config.block_out_channels						) - 1)






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              if latents is None:
                     UpperCamelCase    :	List[str]   	=		randn_tensor(__SCREAMING_SNAKE_CASE ,					generator=__SCREAMING_SNAKE_CASE ,					device=__SCREAMING_SNAKE_CASE ,					dtype=__SCREAMING_SNAKE_CASE						)
              else:
                     if latents.shape != shape:
                            raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {shape}"""						)
                     UpperCamelCase    :	Optional[Any]   	=		latents.to(__SCREAMING_SNAKE_CASE						)

              UpperCamelCase    :	Optional[Any]   	=		latents * scheduler.init_noise_sigma
              return latents






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE=0						):




              """simple docstring"""




              if is_accelerate_available():
                     from accelerate import cpu_offload
              else:
                     raise ImportError('''Please install accelerate via `pip install accelerate`'''						)

              UpperCamelCase    :	List[str]   	=		torch.device(f"""cuda:{gpu_id}"""						)

              UpperCamelCase    :	str   	=		[
                  self.unet,
                  self.movq,
              ]
              for cpu_offloaded_model in models:
                     if cpu_offloaded_model is not None:
                            cpu_offload(__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE						)






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE=0						):




              """simple docstring"""




              if is_accelerate_available() and is_accelerate_version('''>=''' ,					'''0.17.0.dev0'''						):
                     from accelerate import cpu_offload_with_hook
              else:
                     raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'''						)

              UpperCamelCase    :	Dict   	=		torch.device(f"""cuda:{gpu_id}"""						)

              if self.device.type != "cpu":
                     self.to('''cpu''' ,					silence_dtype_warnings=__SCREAMING_SNAKE_CASE						)
                     torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

              UpperCamelCase    :	Optional[int]   	=		None
              for cpu_offloaded_model in [self.unet, self.movq]:
                     UpperCamelCase    :	Dict   	=		cpu_offload_with_hook(__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					prev_module_hook=__SCREAMING_SNAKE_CASE						)

              # We'll offload the last model manually.
              UpperCamelCase    :	Tuple   	=		hook






       @property
       # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
       def  _lowercase   (     self						):




              """simple docstring"""




              if not hasattr(self.unet ,					'''_hf_hook'''						):
                     return self.device
              for module in self.unet.modules():
                     if (
                         hasattr(__SCREAMING_SNAKE_CASE ,					'''_hf_hook'''						)
                         and hasattr(module._hf_hook ,					'''execution_device'''						)
                         and module._hf_hook.execution_device is not None
                     ):
                            return torch.device(module._hf_hook.execution_device						)
              return self.device






       @torch.no_grad()
       @replace_example_docstring(__SCREAMING_SNAKE_CASE						)
       def __call__(     self ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE = 512 ,					__SCREAMING_SNAKE_CASE = 512 ,					__SCREAMING_SNAKE_CASE = 100 ,					__SCREAMING_SNAKE_CASE = 4.0 ,					__SCREAMING_SNAKE_CASE = 1 ,					__SCREAMING_SNAKE_CASE = None ,					__SCREAMING_SNAKE_CASE = None ,					__SCREAMING_SNAKE_CASE = "pil" ,					__SCREAMING_SNAKE_CASE = True ,					):




              """simple docstring"""




              UpperCamelCase    :	Dict   	=		self._execution_device

              UpperCamelCase    :	str   	=		guidance_scale > 1.0

              if isinstance(__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE						):
                     UpperCamelCase    :	List[Any]   	=		torch.cat(__SCREAMING_SNAKE_CASE ,					dim=0						)
              if isinstance(__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE						):
                     UpperCamelCase    :	Optional[int]   	=		torch.cat(__SCREAMING_SNAKE_CASE ,					dim=0						)
              if isinstance(__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE						):
                     UpperCamelCase    :	int   	=		torch.cat(__SCREAMING_SNAKE_CASE ,					dim=0						)

              UpperCamelCase    :	List[str]   	=		image_embeds.shape[0] * num_images_per_prompt

              if do_classifier_free_guidance:
                     UpperCamelCase    :	Union[str, Any]   	=		image_embeds.repeat_interleave(__SCREAMING_SNAKE_CASE ,					dim=0						)
                     UpperCamelCase    :	Tuple   	=		negative_image_embeds.repeat_interleave(__SCREAMING_SNAKE_CASE ,					dim=0						)
                     UpperCamelCase    :	Any   	=		hint.repeat_interleave(__SCREAMING_SNAKE_CASE ,					dim=0						)

              UpperCamelCase    :	List[Any]   	=		torch.cat([negative_image_embeds, image_embeds] ,					dim=0						).to(dtype=self.unet.dtype ,					device=__SCREAMING_SNAKE_CASE						)
              UpperCamelCase    :	int   	=		torch.cat([hint, hint] ,					dim=0						).to(dtype=self.unet.dtype ,					device=__SCREAMING_SNAKE_CASE						)
              self.scheduler.set_timesteps(__SCREAMING_SNAKE_CASE ,					device=__SCREAMING_SNAKE_CASE						)
              UpperCamelCase    :	int   	=		self.scheduler.timesteps

              UpperCamelCase    :	List[Any]   	=		self.movq.config.latent_channels

              UpperCamelCase    :	Any   	=		downscale_height_and_width(__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					self.movq_scale_factor						)

              # create initial latent
              UpperCamelCase    :	str   	=		self.prepare_latents(
                  (batch_size, num_channels_latents, height, width) ,					image_embeds.dtype ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					self.scheduler ,					)

              for i, t in enumerate(self.progress_bar(__SCREAMING_SNAKE_CASE						)						):
                     # expand the latents if we are doing classifier free guidance
                     UpperCamelCase    :	int   	=		torch.cat([latents] * 2						) if do_classifier_free_guidance else latents

                     UpperCamelCase    :	Union[str, Any]   	=		{'''image_embeds''': image_embeds, '''hint''': hint}
                     UpperCamelCase    :	Optional[int]   	=		self.unet(
                         sample=__SCREAMING_SNAKE_CASE ,					timestep=__SCREAMING_SNAKE_CASE ,					encoder_hidden_states=__SCREAMING_SNAKE_CASE ,					added_cond_kwargs=__SCREAMING_SNAKE_CASE ,					return_dict=__SCREAMING_SNAKE_CASE ,					)[0]

                     if do_classifier_free_guidance:
                            UpperCamelCase    :	int   	=		noise_pred.split(latents.shape[1] ,					dim=1						)
                            UpperCamelCase    :	List[Any]   	=		noise_pred.chunk(2						)
                            UpperCamelCase    :	Union[str, Any]   	=		variance_pred.chunk(2						)
                            UpperCamelCase    :	int   	=		noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                            UpperCamelCase    :	Dict   	=		torch.cat([noise_pred, variance_pred_text] ,					dim=1						)

                     if not (
                         hasattr(self.scheduler.config ,					'''variance_type'''						)
                         and self.scheduler.config.variance_type in ["learned", "learned_range"]
                     ):
                            UpperCamelCase    :	Dict   	=		noise_pred.split(latents.shape[1] ,					dim=1						)

                     # compute the previous noisy sample x_t -> x_t-1
                     UpperCamelCase    :	int   	=		self.scheduler.step(
                         __SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					generator=__SCREAMING_SNAKE_CASE ,					)[0]
              # post-processing
              UpperCamelCase    :	List[Any]   	=		self.movq.decode(__SCREAMING_SNAKE_CASE ,					force_not_quantize=__SCREAMING_SNAKE_CASE						)['''sample''']

              if output_type not in ["pt", "np", "pil"]:
                     raise ValueError(f"""Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"""						)

              if output_type in ["np", "pil"]:
                     UpperCamelCase    :	List[str]   	=		image * 0.5 + 0.5
                     UpperCamelCase    :	List[str]   	=		image.clamp(0 ,					1						)
                     UpperCamelCase    :	Any   	=		image.cpu().permute(0 ,					2 ,					3 ,					1						).float().numpy()

              if output_type == "pil":
                     UpperCamelCase    :	List[Any]   	=		self.numpy_to_pil(__SCREAMING_SNAKE_CASE						)

              if not return_dict:
                     return (image,)

              return ImagePipelineOutput(images=__SCREAMING_SNAKE_CASE						)





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


__UpperCAmelCase : List[Any] 			=				[
    "openmmlab/upernet-convnext-tiny",
    # See all UperNet models at https://huggingface.co/models?filter=upernet
]

# General docstring
__UpperCAmelCase : List[str] 			=				"UperNetConfig"



class    UpperCAmelCase_	(							nn.Module):
       '''simple docstring'''






       def __init__(     self ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE = 0 ,					__SCREAMING_SNAKE_CASE = False ,					__SCREAMING_SNAKE_CASE = 1 ,					):




              """simple docstring"""




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






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              UpperCamelCase    :	Any   	=		self.conv(__SCREAMING_SNAKE_CASE						)
              UpperCamelCase    :	Optional[Any]   	=		self.batch_norm(__SCREAMING_SNAKE_CASE						)
              UpperCamelCase    :	Optional[int]   	=		self.activation(__SCREAMING_SNAKE_CASE						)

              return output




class    UpperCAmelCase_	(							nn.Module):
       '''simple docstring'''






       def __init__(     self ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




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






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              UpperCamelCase    :	str   	=		input
              for layer in self.layers:
                     UpperCamelCase    :	int   	=		layer(__SCREAMING_SNAKE_CASE						)
              return hidden_state




class    UpperCAmelCase_	(							nn.Module):
       '''simple docstring'''






       def __init__(     self ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




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






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              UpperCamelCase    :	str   	=		[]
              for ppm in self.blocks:
                     UpperCamelCase    :	List[str]   	=		ppm(__SCREAMING_SNAKE_CASE						)
                     UpperCamelCase    :	Tuple   	=		nn.functional.interpolate(
                         __SCREAMING_SNAKE_CASE ,					size=x.size()[2:] ,					mode='''bilinear''' ,					align_corners=self.align_corners						)
                     ppm_outs.append(__SCREAMING_SNAKE_CASE						)
              return ppm_outs




class    UpperCAmelCase_	(							nn.Module):
       '''simple docstring'''






       def __init__(     self ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              super().__init__()

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

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

              UpperCamelCase    :	Optional[Any]   	=		UperNetConvModule(
                  len(self.in_channels						) * self.channels ,					self.channels ,					kernel_size=3 ,					padding=1 ,					)






       def  _lowercase   (     self						):




              """simple docstring"""




              self.apply(self._init_weights						)






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              if isinstance(__SCREAMING_SNAKE_CASE ,					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  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              UpperCamelCase    :	Any   	=		inputs[-1]
              UpperCamelCase    :	int   	=		[x]
              psp_outs.extend(self.psp_modules(__SCREAMING_SNAKE_CASE						)						)
              UpperCamelCase    :	Any   	=		torch.cat(__SCREAMING_SNAKE_CASE ,					dim=1						)
              UpperCamelCase    :	Union[str, Any]   	=		self.bottleneck(__SCREAMING_SNAKE_CASE						)

              return output






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




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

              laterals.append(self.psp_forward(__SCREAMING_SNAKE_CASE						)						)

              # build top-down path
              UpperCamelCase    :	int   	=		len(__SCREAMING_SNAKE_CASE						)
              for i in range(used_backbone_levels - 1 ,					0 ,					-1						):
                     UpperCamelCase    :	Optional[int]   	=		laterals[i - 1].shape[2:]
                     UpperCamelCase    :	Optional[Any]   	=		laterals[i - 1] + nn.functional.interpolate(
                         laterals[i] ,					size=__SCREAMING_SNAKE_CASE ,					mode='''bilinear''' ,					align_corners=self.align_corners						)

              # build outputs
              UpperCamelCase    :	str   	=		[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						):
                     UpperCamelCase    :	int   	=		nn.functional.interpolate(
                         fpn_outs[i] ,					size=fpn_outs[0].shape[2:] ,					mode='''bilinear''' ,					align_corners=self.align_corners						)
              UpperCamelCase    :	str   	=		torch.cat(__SCREAMING_SNAKE_CASE ,					dim=1						)
              UpperCamelCase    :	Tuple   	=		self.fpn_bottleneck(__SCREAMING_SNAKE_CASE						)
              UpperCamelCase    :	Union[str, Any]   	=		self.classifier(__SCREAMING_SNAKE_CASE						)

              return output




class    UpperCAmelCase_	(							nn.Module):
       '''simple docstring'''






       def __init__(     self ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE = 2 ,					__SCREAMING_SNAKE_CASE = 3 ,					__SCREAMING_SNAKE_CASE = 1						):




              """simple docstring"""




              super().__init__()

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

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

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






       def  _lowercase   (     self						):




              """simple docstring"""




              self.apply(self._init_weights						)






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              if isinstance(__SCREAMING_SNAKE_CASE ,					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  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              UpperCamelCase    :	Any   	=		encoder_hidden_states[self.in_index]
              UpperCamelCase    :	str   	=		self.convs(__SCREAMING_SNAKE_CASE						)
              if self.concat_input:
                     UpperCamelCase    :	int   	=		self.conv_cat(torch.cat([hidden_states, output] ,					dim=1						)						)
              UpperCamelCase    :	Union[str, Any]   	=		self.classifier(__SCREAMING_SNAKE_CASE						)
              return output




class    UpperCAmelCase_	(							_a):
       '''simple docstring'''
       __UpperCamelCase      :     Optional[Any]							=					UperNetConfig
       __UpperCamelCase      :     Optional[int]							=					"pixel_values"
       __UpperCamelCase      :     Dict							=					True






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




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






       def  _lowercase   (     self						):




              """simple docstring"""




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






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE=False						):




              """simple docstring"""




              if isinstance(__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE						):
                     UpperCamelCase    :	str   	=		value




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

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



@add_start_docstrings(
    "UperNet framework leveraging any vision backbone e.g. for ADE20k, CityScapes.",     _a,     )
class    UpperCAmelCase_	(							_a):
       '''simple docstring'''






       def __init__(     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              super().__init__(__SCREAMING_SNAKE_CASE						)

              UpperCamelCase    :	Dict   	=		AutoBackbone.from_config(config.backbone_config						)

              # Semantic segmentation head(s)
              UpperCamelCase    :	int   	=		UperNetHead(__SCREAMING_SNAKE_CASE ,					in_channels=self.backbone.channels						)
              UpperCamelCase    :	int   	=		UperNetFCNHead(__SCREAMING_SNAKE_CASE						) 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=__SCREAMING_SNAKE_CASE ,					config_class=_CONFIG_FOR_DOC						)
       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE = None ,					__SCREAMING_SNAKE_CASE = None ,					__SCREAMING_SNAKE_CASE = None ,					__SCREAMING_SNAKE_CASE = None ,					__SCREAMING_SNAKE_CASE = None ,					):




              """simple docstring"""




              UpperCamelCase    :	List[str]   	=		return_dict if return_dict is not None else self.config.use_return_dict
              UpperCamelCase    :	str   	=		(
                  output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
              )
              UpperCamelCase    :	Tuple   	=		output_attentions if output_attentions is not None else self.config.output_attentions

              UpperCamelCase    :	Tuple   	=		self.backbone.forward_with_filtered_kwargs(
                  __SCREAMING_SNAKE_CASE ,					output_hidden_states=__SCREAMING_SNAKE_CASE ,					output_attentions=__SCREAMING_SNAKE_CASE						)
              UpperCamelCase    :	List[str]   	=		outputs.feature_maps

              UpperCamelCase    :	Union[str, Any]   	=		self.decode_head(__SCREAMING_SNAKE_CASE						)
              UpperCamelCase    :	Optional[int]   	=		nn.functional.interpolate(__SCREAMING_SNAKE_CASE ,					size=pixel_values.shape[2:] ,					mode='''bilinear''' ,					align_corners=__SCREAMING_SNAKE_CASE						)

              UpperCamelCase    :	int   	=		None
              if self.auxiliary_head is not None:
                     UpperCamelCase    :	int   	=		self.auxiliary_head(__SCREAMING_SNAKE_CASE						)
                     UpperCamelCase    :	Tuple   	=		nn.functional.interpolate(
                         __SCREAMING_SNAKE_CASE ,					size=pixel_values.shape[2:] ,					mode='''bilinear''' ,					align_corners=__SCREAMING_SNAKE_CASE						)

              UpperCamelCase    :	Union[str, 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
                            UpperCamelCase    :	Optional[int]   	=		CrossEntropyLoss(ignore_index=self.config.loss_ignore_index						)
                            UpperCamelCase    :	Tuple   	=		loss_fct(__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE						)
                            UpperCamelCase    :	Optional[Any]   	=		loss_fct(__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE						)
                            UpperCamelCase    :	Union[str, Any]   	=		main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

              if not return_dict:
                     if output_hidden_states:
                            UpperCamelCase    :	Optional[Any]   	=		(logits,) + outputs[1:]
                     else:
                            UpperCamelCase    :	int   	=		(logits,) + outputs[2:]
                     return ((loss,) + output) if loss is not None else output

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

