


"""simple docstring"""




import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionInstructPixaPixPipeline,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils import floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
class 					lowerCAmelCase	(					lowerCamelCase_	,    lowerCamelCase_	,    lowerCamelCase_	,    unittest.TestCase    ):




		'''simple docstring'''







		SCREAMING_SNAKE_CASE_						:     int								=     StableDiffusionInstructPixaPixPipeline
		SCREAMING_SNAKE_CASE_						:     Union[str, Any]								=     TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width""", """cross_attention_kwargs"""}
		SCREAMING_SNAKE_CASE_						:     str								=     TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
		SCREAMING_SNAKE_CASE_						:     Dict								=     IMAGE_TO_IMAGE_IMAGE_PARAMS
		SCREAMING_SNAKE_CASE_						:     Dict								=     IMAGE_TO_IMAGE_IMAGE_PARAMS







		def 							__A						(	self				)     ->							Optional[int]:
								torch.manual_seed(0				)
								SCREAMING_SNAKE_CASE													=					UNetaDConditionModel(
								    block_out_channels=(32, 64)						,       layers_per_block=2						,       sample_size=32						,       in_channels=8						,       out_channels=4						,       down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')						,       up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')						,       cross_attention_dim=32						,       )
								SCREAMING_SNAKE_CASE													=					PNDMScheduler(skip_prk_steps=lowerCAmelCase__				)
								torch.manual_seed(0				)
								SCREAMING_SNAKE_CASE													=					AutoencoderKL(
								    block_out_channels=[32, 64]						,       in_channels=3						,       out_channels=3						,       down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']						,       up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']						,       latent_channels=4						,       )
								torch.manual_seed(0				)
								SCREAMING_SNAKE_CASE													=					CLIPTextConfig(
								    bos_token_id=0						,       eos_token_id=2						,       hidden_size=32						,       intermediate_size=37						,       layer_norm_eps=1e-05						,       num_attention_heads=4						,       num_hidden_layers=5						,       pad_token_id=1						,       vocab_size=1_000						,       )
								SCREAMING_SNAKE_CASE													=					CLIPTextModel(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'				)

								SCREAMING_SNAKE_CASE													=					{
								    'unet': unet,
								    'scheduler': scheduler,
								    'vae': vae,
								    'text_encoder': text_encoder,
								    'tokenizer': tokenizer,
								    'safety_checker': None,
								    'feature_extractor': None,
								}
								return components







		def 							__A						(	self						,       lowerCAmelCase__						,       lowerCAmelCase__=0				)     ->							List[str]:
								SCREAMING_SNAKE_CASE													=					floats_tensor((1, 3, 32, 32)						,       rng=random.Random(lowerCAmelCase__				)				).to(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					image.cpu().permute(0						,       2						,       3						,       1				)[0]
								SCREAMING_SNAKE_CASE													=					Image.fromarray(np.uinta(lowerCAmelCase__				)				).convert('RGB'				)
								if str(lowerCAmelCase__				).startswith('mps'				):
														SCREAMING_SNAKE_CASE													=					torch.manual_seed(lowerCAmelCase__				)
								else:
														SCREAMING_SNAKE_CASE													=					torch.Generator(device=lowerCAmelCase__				).manual_seed(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					{
								    'prompt': 'A painting of a squirrel eating a burger',
								    'image': image,
								    'generator': generator,
								    'num_inference_steps': 2,
								    'guidance_scale': 6.0,
								    'image_guidance_scale': 1,
								    'output_type': 'numpy',
								}
								return inputs







		def 							__A						(	self				)     ->							Union[str, Any]:
								SCREAMING_SNAKE_CASE													=					'cpu'  # ensure determinism for the device-dependent torch.Generator
								SCREAMING_SNAKE_CASE													=					self.get_dummy_components()
								SCREAMING_SNAKE_CASE													=					StableDiffusionInstructPixaPixPipeline(**lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					sd_pipe.to(lowerCAmelCase__				)
								sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE													=					self.get_dummy_inputs(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					sd_pipe(**lowerCAmelCase__				).images
								SCREAMING_SNAKE_CASE													=					image[0, -3:, -3:, -1]
								assert image.shape == (1, 32, 32, 3)
								SCREAMING_SNAKE_CASE													=					np.array([0.75_26, 0.37_50, 0.45_47, 0.61_17, 0.58_66, 0.50_16, 0.43_27, 0.56_42, 0.48_15]				)

								assert np.abs(image_slice.flatten() - expected_slice				).max() < 1e-3







		def 							__A						(	self				)     ->							int:
								SCREAMING_SNAKE_CASE													=					'cpu'  # ensure determinism for the device-dependent torch.Generator
								SCREAMING_SNAKE_CASE													=					self.get_dummy_components()
								SCREAMING_SNAKE_CASE													=					StableDiffusionInstructPixaPixPipeline(**lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					sd_pipe.to(lowerCAmelCase__				)
								sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE													=					self.get_dummy_inputs(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					'french fries'
								SCREAMING_SNAKE_CASE													=					sd_pipe(**lowerCAmelCase__						,       negative_prompt=lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					output.images
								SCREAMING_SNAKE_CASE													=					image[0, -3:, -3:, -1]

								assert image.shape == (1, 32, 32, 3)
								SCREAMING_SNAKE_CASE													=					np.array([0.75_11, 0.36_42, 0.45_53, 0.62_36, 0.57_97, 0.50_13, 0.43_43, 0.56_11, 0.48_31]				)

								assert np.abs(image_slice.flatten() - expected_slice				).max() < 1e-3







		def 							__A						(	self				)     ->							Optional[int]:
								SCREAMING_SNAKE_CASE													=					'cpu'  # ensure determinism for the device-dependent torch.Generator
								SCREAMING_SNAKE_CASE													=					self.get_dummy_components()
								SCREAMING_SNAKE_CASE													=					StableDiffusionInstructPixaPixPipeline(**lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					sd_pipe.to(lowerCAmelCase__				)
								sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE													=					self.get_dummy_inputs(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					[inputs['prompt']] * 2

								SCREAMING_SNAKE_CASE													=					np.array(inputs['image']				).astype(np.floataa				) / 2_55.0
								SCREAMING_SNAKE_CASE													=					torch.from_numpy(lowerCAmelCase__				).unsqueeze(0				).to(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					image / 2 + 0.5
								SCREAMING_SNAKE_CASE													=					image.permute(0						,       3						,       1						,       2				)
								SCREAMING_SNAKE_CASE													=					image.repeat(2						,       1						,       1						,       1				)

								SCREAMING_SNAKE_CASE													=					sd_pipe(**lowerCAmelCase__				).images
								SCREAMING_SNAKE_CASE													=					image[-1, -3:, -3:, -1]

								assert image.shape == (2, 32, 32, 3)
								SCREAMING_SNAKE_CASE													=					np.array([0.58_12, 0.57_48, 0.52_22, 0.59_08, 0.56_95, 0.71_74, 0.68_04, 0.55_23, 0.55_79]				)

								assert np.abs(image_slice.flatten() - expected_slice				).max() < 1e-3







		def 							__A						(	self				)     ->							List[str]:
								SCREAMING_SNAKE_CASE													=					'cpu'  # ensure determinism for the device-dependent torch.Generator
								SCREAMING_SNAKE_CASE													=					self.get_dummy_components()
								SCREAMING_SNAKE_CASE													=					EulerAncestralDiscreteScheduler(
								    beta_start=0.0_00_85						,       beta_end=0.0_12						,       beta_schedule='scaled_linear'				)
								SCREAMING_SNAKE_CASE													=					StableDiffusionInstructPixaPixPipeline(**lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					sd_pipe.to(lowerCAmelCase__				)
								sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE													=					self.get_dummy_inputs(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					sd_pipe(**lowerCAmelCase__				).images
								SCREAMING_SNAKE_CASE													=					image[0, -3:, -3:, -1]

								SCREAMING_SNAKE_CASE													=					[round(lowerCAmelCase__						,       4				) for x in image_slice.flatten().tolist()]
								print(','.join([str(lowerCAmelCase__				) for x in slice]				)				)

								assert image.shape == (1, 32, 32, 3)
								SCREAMING_SNAKE_CASE													=					np.array([0.74_17, 0.38_42, 0.47_32, 0.57_76, 0.58_91, 0.51_39, 0.40_52, 0.56_73, 0.49_86]				)

								assert np.abs(image_slice.flatten() - expected_slice				).max() < 1e-3







		def 							__A						(	self				)     ->							Dict:
								super().test_inference_batch_single_identical(expected_max_diff=3e-3				)







		def 							__A						(	self				)     ->							List[str]:
								SCREAMING_SNAKE_CASE													=					self.get_dummy_components()
								SCREAMING_SNAKE_CASE													=					StableDiffusionInstructPixaPixPipeline(**lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					VaeImageProcessor(do_resize=lowerCAmelCase__						,       do_normalize=lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					pipe.to(lowerCAmelCase__				)
								pipe.set_progress_bar_config(disable=lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE													=					pipe(**self.get_dummy_inputs_by_type(lowerCAmelCase__						,       input_image_type='pt'				)				)[0]

								SCREAMING_SNAKE_CASE													=					components['vae']
								SCREAMING_SNAKE_CASE													=					self.get_dummy_inputs_by_type(lowerCAmelCase__						,       input_image_type='pt'				)

								for image_param in self.image_latents_params:
														if image_param in inputs.keys():
																				SCREAMING_SNAKE_CASE													=					vae.encode(inputs[image_param]				).latent_dist.mode()

								SCREAMING_SNAKE_CASE													=					pipe(**lowerCAmelCase__				)[0]

								SCREAMING_SNAKE_CASE													=					np.abs(out - out_latents_inputs				).max()
								self.assertLess(lowerCAmelCase__						,       1e-4						,       'passing latents as image input generate different result from passing image'				)





@slow
@require_torch_gpu
class 					lowerCAmelCase	(					unittest.TestCase    ):




		'''simple docstring'''







		def 							__A						(	self				)     ->							str:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







		def 							__A						(	self						,       lowerCAmelCase__=0				)     ->							Any:
								SCREAMING_SNAKE_CASE													=					torch.manual_seed(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					load_image(
								    'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg'				)
								SCREAMING_SNAKE_CASE													=					{
								    'prompt': 'turn him into a cyborg',
								    'image': image,
								    'generator': generator,
								    'num_inference_steps': 3,
								    'guidance_scale': 7.5,
								    'image_guidance_scale': 1.0,
								    'output_type': 'numpy',
								}
								return inputs







		def 							__A						(	self				)     ->							Union[str, Any]:
								SCREAMING_SNAKE_CASE													=					StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    'timbrooks/instruct-pix2pix'						,       safety_checker=lowerCAmelCase__				)
								pipe.to(lowerCAmelCase__				)
								pipe.set_progress_bar_config(disable=lowerCAmelCase__				)
								pipe.enable_attention_slicing()

								SCREAMING_SNAKE_CASE													=					self.get_inputs()
								SCREAMING_SNAKE_CASE													=					pipe(**lowerCAmelCase__				).images
								SCREAMING_SNAKE_CASE													=					image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 512, 512, 3)
								SCREAMING_SNAKE_CASE													=					np.array([0.59_02, 0.60_15, 0.60_27, 0.59_83, 0.60_92, 0.60_61, 0.57_65, 0.57_85, 0.55_55]				)

								assert np.abs(expected_slice - image_slice				).max() < 1e-3







		def 							__A						(	self				)     ->							Optional[Any]:
								SCREAMING_SNAKE_CASE													=					StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    'timbrooks/instruct-pix2pix'						,       safety_checker=lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					LMSDiscreteScheduler.from_config(pipe.scheduler.config				)
								pipe.to(lowerCAmelCase__				)
								pipe.set_progress_bar_config(disable=lowerCAmelCase__				)
								pipe.enable_attention_slicing()

								SCREAMING_SNAKE_CASE													=					self.get_inputs()
								SCREAMING_SNAKE_CASE													=					pipe(**lowerCAmelCase__				).images
								SCREAMING_SNAKE_CASE													=					image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 512, 512, 3)
								SCREAMING_SNAKE_CASE													=					np.array([0.65_78, 0.68_17, 0.69_72, 0.67_61, 0.68_56, 0.69_16, 0.64_28, 0.65_16, 0.63_01]				)

								assert np.abs(expected_slice - image_slice				).max() < 1e-3







		def 							__A						(	self				)     ->							Optional[Any]:
								SCREAMING_SNAKE_CASE													=					StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    'timbrooks/instruct-pix2pix'						,       safety_checker=lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					DDIMScheduler.from_config(pipe.scheduler.config				)
								pipe.to(lowerCAmelCase__				)
								pipe.set_progress_bar_config(disable=lowerCAmelCase__				)
								pipe.enable_attention_slicing()

								SCREAMING_SNAKE_CASE													=					self.get_inputs()
								SCREAMING_SNAKE_CASE													=					pipe(**lowerCAmelCase__				).images
								SCREAMING_SNAKE_CASE													=					image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 512, 512, 3)
								SCREAMING_SNAKE_CASE													=					np.array([0.38_28, 0.38_34, 0.38_18, 0.37_92, 0.38_65, 0.37_52, 0.37_92, 0.38_47, 0.37_53]				)

								assert np.abs(expected_slice - image_slice				).max() < 1e-3







		def 							__A						(	self				)     ->							Optional[int]:
								SCREAMING_SNAKE_CASE													=					0

								def callback_fn(lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__				) -> None:
														SCREAMING_SNAKE_CASE													=					True
														nonlocal number_of_steps
														number_of_steps += 1
														if step == 1:
																				SCREAMING_SNAKE_CASE													=					latents.detach().cpu().numpy()
																				assert latents.shape == (1, 4, 64, 64)
																				SCREAMING_SNAKE_CASE													=					latents[0, -3:, -3:, -1]
																				SCREAMING_SNAKE_CASE													=					np.array([-0.24_63, -0.46_44, -0.97_56, 1.51_76, 1.44_14, 0.78_66, 0.98_97, 0.85_21, 0.79_83]				)

																				assert np.abs(latents_slice.flatten() - expected_slice				).max() < 5e-2
														elif step == 2:
																				SCREAMING_SNAKE_CASE													=					latents.detach().cpu().numpy()
																				assert latents.shape == (1, 4, 64, 64)
																				SCREAMING_SNAKE_CASE													=					latents[0, -3:, -3:, -1]
																				SCREAMING_SNAKE_CASE													=					np.array([-0.26_44, -0.46_26, -0.96_53, 1.51_76, 1.45_51, 0.76_86, 0.98_05, 0.84_52, 0.81_15]				)

																				assert np.abs(latents_slice.flatten() - expected_slice				).max() < 5e-2

								SCREAMING_SNAKE_CASE													=					False

								SCREAMING_SNAKE_CASE													=					StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    'timbrooks/instruct-pix2pix'						,       safety_checker=lowerCAmelCase__						,       torch_dtype=torch.floataa				)
								SCREAMING_SNAKE_CASE													=					pipe.to(lowerCAmelCase__				)
								pipe.set_progress_bar_config(disable=lowerCAmelCase__				)
								pipe.enable_attention_slicing()

								SCREAMING_SNAKE_CASE													=					self.get_inputs()
								pipe(**lowerCAmelCase__						,       callback=lowerCAmelCase__						,       callback_steps=1				)
								assert callback_fn.has_been_called
								assert number_of_steps == 3







		def 							__A						(	self				)     ->							Optional[int]:
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

								SCREAMING_SNAKE_CASE													=					StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    'timbrooks/instruct-pix2pix'						,       safety_checker=lowerCAmelCase__						,       torch_dtype=torch.floataa				)
								SCREAMING_SNAKE_CASE													=					pipe.to(lowerCAmelCase__				)
								pipe.set_progress_bar_config(disable=lowerCAmelCase__				)
								pipe.enable_attention_slicing(1				)
								pipe.enable_sequential_cpu_offload()

								SCREAMING_SNAKE_CASE													=					self.get_inputs()
								SCREAMING_SNAKE_CASE													=					pipe(**lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE													=					torch.cuda.max_memory_allocated()
								# make sure that less than 2.2 GB is allocated
								assert mem_bytes < 2.2 * 10**9







		def 							__A						(	self				)     ->							Dict:
								SCREAMING_SNAKE_CASE													=					self.get_inputs()
								# resize to resolution that is divisible by 8 but not 16 or 32
								SCREAMING_SNAKE_CASE													=					inputs['image'].resize((504, 504)				)

								SCREAMING_SNAKE_CASE													=					'timbrooks/instruct-pix2pix'
								SCREAMING_SNAKE_CASE													=					StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    lowerCAmelCase__						,       safety_checker=lowerCAmelCase__						,       )
								pipe.to(lowerCAmelCase__				)
								pipe.set_progress_bar_config(disable=lowerCAmelCase__				)
								pipe.enable_attention_slicing()

								SCREAMING_SNAKE_CASE													=					pipe(**lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					output.images[0]

								SCREAMING_SNAKE_CASE													=					image[255:258, 383:386, -1]

								assert image.shape == (504, 504, 3)
								SCREAMING_SNAKE_CASE													=					np.array([0.27_26, 0.25_29, 0.26_64, 0.26_55, 0.26_41, 0.26_42, 0.25_91, 0.26_49, 0.25_90]				)

								assert np.abs(image_slice.flatten() - expected_slice				).max() < 5e-3






"""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


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

# General docstring
__UpperCamelCase						=      '''UperNetConfig'''
class 					lowerCAmelCase	(					nn.Module    ):




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__ = 0						,       lowerCAmelCase__ = False						,       lowerCAmelCase__ = 1						,       )     ->							None:
								super().__init__()
								SCREAMING_SNAKE_CASE													=					nn.Convad(
								    in_channels=lowerCAmelCase__						,       out_channels=lowerCAmelCase__						,       kernel_size=lowerCAmelCase__						,       padding=lowerCAmelCase__						,       bias=lowerCAmelCase__						,       dilation=lowerCAmelCase__						,       )
								SCREAMING_SNAKE_CASE													=					nn.BatchNormad(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					nn.ReLU()







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							torch.Tensor:
								SCREAMING_SNAKE_CASE													=					self.conv(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					self.batch_norm(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					self.activation(lowerCAmelCase__				)

								return output





class 					lowerCAmelCase	(					nn.Module    ):




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__				)     ->							None:
								super().__init__()
								SCREAMING_SNAKE_CASE													=					[
								    nn.AdaptiveAvgPoolad(lowerCAmelCase__				),
								    UperNetConvModule(lowerCAmelCase__						,       lowerCAmelCase__						,       kernel_size=1				),
								]
								for i, layer in enumerate(self.layers				):
														self.add_module(str(lowerCAmelCase__				)						,       lowerCAmelCase__				)







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							torch.Tensor:
								SCREAMING_SNAKE_CASE													=					input
								for layer in self.layers:
														SCREAMING_SNAKE_CASE													=					layer(lowerCAmelCase__				)
								return hidden_state





class 					lowerCAmelCase	(					nn.Module    ):




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__				)     ->							None:
								super().__init__()
								SCREAMING_SNAKE_CASE													=					pool_scales
								SCREAMING_SNAKE_CASE													=					align_corners
								SCREAMING_SNAKE_CASE													=					in_channels
								SCREAMING_SNAKE_CASE													=					channels
								SCREAMING_SNAKE_CASE													=					[]
								for i, pool_scale in enumerate(lowerCAmelCase__				):
														SCREAMING_SNAKE_CASE													=					UperNetPyramidPoolingBlock(pool_scale=lowerCAmelCase__						,       in_channels=lowerCAmelCase__						,       channels=lowerCAmelCase__				)
														self.blocks.append(lowerCAmelCase__				)
														self.add_module(str(lowerCAmelCase__				)						,       lowerCAmelCase__				)







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							List[torch.Tensor]:
								SCREAMING_SNAKE_CASE													=					[]
								for ppm in self.blocks:
														SCREAMING_SNAKE_CASE													=					ppm(lowerCAmelCase__				)
														SCREAMING_SNAKE_CASE													=					nn.functional.interpolate(
														    lowerCAmelCase__						,       size=x.size()[2:]						,       mode='bilinear'						,       align_corners=self.align_corners				)
														ppm_outs.append(lowerCAmelCase__				)
								return ppm_outs





class 					lowerCAmelCase	(					nn.Module    ):




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__						,       lowerCAmelCase__				)     ->							List[str]:
								super().__init__()

								SCREAMING_SNAKE_CASE													=					config
								SCREAMING_SNAKE_CASE													=					config.pool_scales  # e.g. (1, 2, 3, 6)
								SCREAMING_SNAKE_CASE													=					in_channels
								SCREAMING_SNAKE_CASE													=					config.hidden_size
								SCREAMING_SNAKE_CASE													=					False
								SCREAMING_SNAKE_CASE													=					nn.Convad(self.channels						,       config.num_labels						,       kernel_size=1				)

								# PSP Module
								SCREAMING_SNAKE_CASE													=					UperNetPyramidPoolingModule(
								    self.pool_scales						,       self.in_channels[-1]						,       self.channels						,       align_corners=self.align_corners						,       )
								SCREAMING_SNAKE_CASE													=					UperNetConvModule(
								    self.in_channels[-1] + len(self.pool_scales				) * self.channels						,       self.channels						,       kernel_size=3						,       padding=1						,       )
								# FPN Module
								SCREAMING_SNAKE_CASE													=					nn.ModuleList()
								SCREAMING_SNAKE_CASE													=					nn.ModuleList()
								for in_channels in self.in_channels[:-1]:  # skip the top layer
														SCREAMING_SNAKE_CASE													=					UperNetConvModule(lowerCAmelCase__						,       self.channels						,       kernel_size=1				)
														SCREAMING_SNAKE_CASE													=					UperNetConvModule(self.channels						,       self.channels						,       kernel_size=3						,       padding=1				)
														self.lateral_convs.append(lowerCAmelCase__				)
														self.fpn_convs.append(lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE													=					UperNetConvModule(
								    len(self.in_channels				) * self.channels						,       self.channels						,       kernel_size=3						,       padding=1						,       )







		def 							__A						(	self				)     ->							int:
								self.apply(self._init_weights				)







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							Tuple:
								if isinstance(lowerCAmelCase__						,       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 							__A						(	self						,       lowerCAmelCase__				)     ->							Optional[int]:
								SCREAMING_SNAKE_CASE													=					inputs[-1]
								SCREAMING_SNAKE_CASE													=					[x]
								psp_outs.extend(self.psp_modules(lowerCAmelCase__				)				)
								SCREAMING_SNAKE_CASE													=					torch.cat(lowerCAmelCase__						,       dim=1				)
								SCREAMING_SNAKE_CASE													=					self.bottleneck(lowerCAmelCase__				)

								return output







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							torch.Tensor:
								# build laterals
								SCREAMING_SNAKE_CASE													=					[lateral_conv(encoder_hidden_states[i]				) for i, lateral_conv in enumerate(self.lateral_convs				)]

								laterals.append(self.psp_forward(lowerCAmelCase__				)				)

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

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

								return output





class 					lowerCAmelCase	(					nn.Module    ):




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__						,       lowerCAmelCase__ = 2						,       lowerCAmelCase__ = 3						,       lowerCAmelCase__ = 1				)     ->							None:
								super().__init__()

								SCREAMING_SNAKE_CASE													=					config
								SCREAMING_SNAKE_CASE													=					config.auxiliary_in_channels
								SCREAMING_SNAKE_CASE													=					config.auxiliary_channels
								SCREAMING_SNAKE_CASE													=					config.auxiliary_num_convs
								SCREAMING_SNAKE_CASE													=					config.auxiliary_concat_input
								SCREAMING_SNAKE_CASE													=					in_index

								SCREAMING_SNAKE_CASE													=					(kernel_size // 2) * dilation
								SCREAMING_SNAKE_CASE													=					[]
								convs.append(
								    UperNetConvModule(
								        self.in_channels						,       self.channels						,       kernel_size=lowerCAmelCase__						,       padding=lowerCAmelCase__						,       dilation=lowerCAmelCase__				)				)
								for i in range(self.num_convs - 1				):
														convs.append(
														    UperNetConvModule(
														        self.channels						,       self.channels						,       kernel_size=lowerCAmelCase__						,       padding=lowerCAmelCase__						,       dilation=lowerCAmelCase__				)				)
								if self.num_convs == 0:
														SCREAMING_SNAKE_CASE													=					nn.Identity()
								else:
														SCREAMING_SNAKE_CASE													=					nn.Sequential(*lowerCAmelCase__				)
								if self.concat_input:
														SCREAMING_SNAKE_CASE													=					UperNetConvModule(
														    self.in_channels + self.channels						,       self.channels						,       kernel_size=lowerCAmelCase__						,       padding=kernel_size // 2				)

								SCREAMING_SNAKE_CASE													=					nn.Convad(self.channels						,       config.num_labels						,       kernel_size=1				)







		def 							__A						(	self				)     ->							Dict:
								self.apply(self._init_weights				)







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							Dict:
								if isinstance(lowerCAmelCase__						,       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 							__A						(	self						,       lowerCAmelCase__				)     ->							torch.Tensor:
								# just take the relevant feature maps
								SCREAMING_SNAKE_CASE													=					encoder_hidden_states[self.in_index]
								SCREAMING_SNAKE_CASE													=					self.convs(lowerCAmelCase__				)
								if self.concat_input:
														SCREAMING_SNAKE_CASE													=					self.conv_cat(torch.cat([hidden_states, output]						,       dim=1				)				)
								SCREAMING_SNAKE_CASE													=					self.classifier(lowerCAmelCase__				)
								return output





class 					lowerCAmelCase	(					lowerCamelCase_    ):




		'''simple docstring'''







		SCREAMING_SNAKE_CASE_						:     int								=     UperNetConfig
		SCREAMING_SNAKE_CASE_						:     Optional[Any]								=     """pixel_values"""
		SCREAMING_SNAKE_CASE_						:     Optional[int]								=     True







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							List[str]:
								if isinstance(lowerCAmelCase__						,       lowerCAmelCase__				):
														module.backbone.init_weights()
														module.decode_head.init_weights()
														module.auxiliary_head.init_weights()







		def 							__A						(	self				)     ->							int:
								self.backbone.init_weights()
								self.decode_head.init_weights()
								self.auxiliary_head.init_weights()







		def 							__A						(	self						,       lowerCAmelCase__						,       lowerCAmelCase__=False				)     ->							str:
								if isinstance(lowerCAmelCase__						,       lowerCAmelCase__				):
														SCREAMING_SNAKE_CASE													=					value





__UpperCamelCase						=      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.
'''

__UpperCamelCase						=      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."""	,    lowerCamelCase_	,    )
class 					lowerCAmelCase	(					lowerCamelCase_    ):




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__				)     ->							Optional[int]:
								super().__init__(lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE													=					AutoBackbone.from_config(config.backbone_config				)

								# Semantic segmentation head(s)
								SCREAMING_SNAKE_CASE													=					UperNetHead(lowerCAmelCase__						,       in_channels=self.backbone.channels				)
								SCREAMING_SNAKE_CASE													=					UperNetFCNHead(lowerCAmelCase__				) 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=lowerCAmelCase__						,       config_class=_CONFIG_FOR_DOC				)
		def 							__A						(	self						,       lowerCAmelCase__ = None						,       lowerCAmelCase__ = None						,       lowerCAmelCase__ = None						,       lowerCAmelCase__ = None						,       lowerCAmelCase__ = None						,       )     ->							Union[tuple, SemanticSegmenterOutput]:
								SCREAMING_SNAKE_CASE													=					return_dict if return_dict is not None else self.config.use_return_dict
								SCREAMING_SNAKE_CASE													=					(
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)
								SCREAMING_SNAKE_CASE													=					output_attentions if output_attentions is not None else self.config.output_attentions

								SCREAMING_SNAKE_CASE													=					self.backbone.forward_with_filtered_kwargs(
								    lowerCAmelCase__						,       output_hidden_states=lowerCAmelCase__						,       output_attentions=lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					outputs.feature_maps

								SCREAMING_SNAKE_CASE													=					self.decode_head(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					nn.functional.interpolate(lowerCAmelCase__						,       size=pixel_values.shape[2:]						,       mode='bilinear'						,       align_corners=lowerCAmelCase__				)

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

								SCREAMING_SNAKE_CASE													=					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
																				SCREAMING_SNAKE_CASE													=					CrossEntropyLoss(ignore_index=self.config.loss_ignore_index				)
																				SCREAMING_SNAKE_CASE													=					loss_fct(lowerCAmelCase__						,       lowerCAmelCase__				)
																				SCREAMING_SNAKE_CASE													=					loss_fct(lowerCAmelCase__						,       lowerCAmelCase__				)
																				SCREAMING_SNAKE_CASE													=					main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

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

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


