






"""simple docstring"""

import inspect
from typing import Callable, List, Optional, Union

import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextModel,
    CLIPTokenizer,
    WhisperForConditionalGeneration,
    WhisperProcessor,
)

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.utils import logging


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


class 		__lowerCamelCase				( snake_case_			):





		'''simple docstring'''






		def __init__(   self							:	Any						,			a_							:	List[Any]						,			a_							:	Union[str, Any]						,			a_							:	int						,			a_							:	Union[str, Any]						,			a_							:	List[Any]						,			a_							:	Tuple						,			a_							:	int						,			a_							:	int						,			a_							:	Tuple						,			):
				super().__init__()

				if safety_checker is None:
						logger.warning(
						    f'''You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure'''
						    " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
						    " results in services or applications open to the public. Both the diffusers team and Hugging Face"
						    " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
						    " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
						    " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."	)

				self.register_modules(
				    speech_model=a_						,			speech_processor=a_						,			vae=a_						,			text_encoder=a_						,			tokenizer=a_						,			unet=a_						,			scheduler=a_						,			feature_extractor=a_						,			)





		def    lowerCamelCase       (   self							:	Union[str, Any]						,			a_							:	Union[str, Any] = "auto"	):
				if slice_size == "auto":
						lowerCAmelCase_   :							int									=			self.unet.config.attention_head_dim // 2
				self.unet.set_attention_slice(a_	)





		def    lowerCamelCase       (   self							:	int	):
				self.enable_attention_slicing(a_	)





		@torch.no_grad()
		def __call__(   self							:	str						,			a_							:	Union[str, Any]						,			a_							:	Optional[int]=1_60_00						,			a_							:	Tuple = 5_12						,			a_							:	int = 5_12						,			a_							:	Optional[Any] = 50						,			a_							:	int = 7.5						,			a_							:	Tuple = None						,			a_							:	Any = 1						,			a_							:	Optional[int] = 0.0						,			a_							:	Any = None						,			a_							:	int = None						,			a_							:	str = "pil"						,			a_							:	Union[str, Any] = True						,			a_							:	Tuple = None						,			a_							:	List[Any] = 1						,			**a_							:	Optional[int]						,			):
				lowerCAmelCase_   :							List[Any]									=			self.speech_processor.feature_extractor(
				    a_						,			return_tensors="pt"						,			sampling_rate=a_	).input_features.to(self.device	)
				lowerCAmelCase_   :							Optional[Any]									=			self.speech_model.generate(a_						,			max_length=48_00_00	)

				lowerCAmelCase_   :							Any									=			self.speech_processor.tokenizer.batch_decode(a_						,			skip_special_tokens=a_						,			normalize=a_	)[
				    0
				]

				if isinstance(a_						,			a_	):
						lowerCAmelCase_   :							List[str]									=			1
				elif isinstance(a_						,			a_	):
						lowerCAmelCase_   :							List[str]									=			len(a_	)
				else:
						raise ValueError(f'''`prompt` has to be of type `str` or `list` but is {type(a_	)}'''	)

				if height % 8 != 0 or width % 8 != 0:
						raise ValueError(f'''`height` and `width` have to be divisible by 8 but are {height} and {width}.'''	)

				if (callback_steps is None) or (
				    callback_steps is not None and (not isinstance(a_						,			a_	) or callback_steps <= 0)
				):
						raise ValueError(
						    f'''`callback_steps` has to be a positive integer but is {callback_steps} of type'''
						    f''' {type(a_	)}.'''	)

				# get prompt text embeddings
				lowerCAmelCase_   :							List[str]									=			self.tokenizer(
				    a_						,			padding="max_length"						,			max_length=self.tokenizer.model_max_length						,			return_tensors="pt"						,			)
				lowerCAmelCase_   :							Tuple									=			text_inputs.input_ids

				if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
						lowerCAmelCase_   :							Optional[int]									=			self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]	)
						logger.warning(
						    "The following part of your input was truncated because CLIP can only handle sequences up to"
						    f''' {self.tokenizer.model_max_length} tokens: {removed_text}'''	)
						lowerCAmelCase_   :							List[Any]									=			text_input_ids[:, : self.tokenizer.model_max_length]
				lowerCAmelCase_   :							Optional[Any]									=			self.text_encoder(text_input_ids.to(self.device	)	)[0]

				# duplicate text embeddings for each generation per prompt, using mps friendly method
				lowerCAmelCase_				,	lowerCAmelCase_				,	lowerCAmelCase_   :							Optional[Any]									=			text_embeddings.shape
				lowerCAmelCase_   :							Tuple									=			text_embeddings.repeat(1						,			a_						,			1	)
				lowerCAmelCase_   :							Optional[Any]									=			text_embeddings.view(bs_embed * num_images_per_prompt						,			a_						,			-1	)

				# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
				# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
				# corresponds to doing no classifier free guidance.
				lowerCAmelCase_   :							Optional[int]									=			guidance_scale > 1.0
				# get unconditional embeddings for classifier free guidance
				if do_classifier_free_guidance:
						lowerCAmelCase_   :							int									=			42
						if negative_prompt is None:
								lowerCAmelCase_   :							Optional[int]									=			[""] * batch_size
						elif type(a_	) is not type(a_	):
								raise TypeError(
								    f'''`negative_prompt` should be the same type to `prompt`, but got {type(a_	)} !='''
								    f''' {type(a_	)}.'''	)
						elif isinstance(a_						,			a_	):
								lowerCAmelCase_   :							Union[str, Any]									=			[negative_prompt]
						elif batch_size != len(a_	):
								raise ValueError(
								    f'''`negative_prompt`: {negative_prompt} has batch size {len(a_	)}, but `prompt`:'''
								    f''' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'''
								    " the batch size of `prompt`."	)
						else:
								lowerCAmelCase_   :							str									=			negative_prompt

						lowerCAmelCase_   :							Optional[Any]									=			text_input_ids.shape[-1]
						lowerCAmelCase_   :							Dict									=			self.tokenizer(
						    a_						,			padding="max_length"						,			max_length=a_						,			truncation=a_						,			return_tensors="pt"						,			)
						lowerCAmelCase_   :							Dict									=			self.text_encoder(uncond_input.input_ids.to(self.device	)	)[0]

						# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
						lowerCAmelCase_   :							List[str]									=			uncond_embeddings.shape[1]
						lowerCAmelCase_   :							Tuple									=			uncond_embeddings.repeat(1						,			a_						,			1	)
						lowerCAmelCase_   :							Optional[Any]									=			uncond_embeddings.view(batch_size * num_images_per_prompt						,			a_						,			-1	)

						# For classifier free guidance, we need to do two forward passes.
						# Here we concatenate the unconditional and text embeddings into a single batch
						# to avoid doing two forward passes
						lowerCAmelCase_   :							str									=			torch.cat([uncond_embeddings, text_embeddings]	)

				# get the initial random noise unless the user supplied it

				# Unlike in other pipelines, latents need to be generated in the target device
				# for 1-to-1 results reproducibility with the CompVis implementation.
				# However this currently doesn't work in `mps`.
				lowerCAmelCase_   :							List[str]									=			(batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
				lowerCAmelCase_   :							List[str]									=			text_embeddings.dtype
				if latents is None:
						if self.device.type == "mps":
								# randn does not exist on mps
								lowerCAmelCase_   :							Dict									=			torch.randn(a_						,			generator=a_						,			device="cpu"						,			dtype=a_	).to(
								    self.device	)
						else:
								lowerCAmelCase_   :							Optional[int]									=			torch.randn(a_						,			generator=a_						,			device=self.device						,			dtype=a_	)
				else:
						if latents.shape != latents_shape:
								raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}'''	)
						lowerCAmelCase_   :							Tuple									=			latents.to(self.device	)

				# set timesteps
				self.scheduler.set_timesteps(a_	)

				# Some schedulers like PNDM have timesteps as arrays
				# It's more optimized to move all timesteps to correct device beforehand
				lowerCAmelCase_   :							Optional[Any]									=			self.scheduler.timesteps.to(self.device	)

				# scale the initial noise by the standard deviation required by the scheduler
				lowerCAmelCase_   :							Dict									=			latents * self.scheduler.init_noise_sigma

				# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
				# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
				# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
				# and should be between [0, 1]
				lowerCAmelCase_   :							Any									=			"eta" in set(inspect.signature(self.scheduler.step	).parameters.keys()	)
				lowerCAmelCase_   :							Optional[Any]									=			{}
				if accepts_eta:
						lowerCAmelCase_   :							Optional[int]									=			eta

				for i, t in enumerate(self.progress_bar(a_	)	):
						# expand the latents if we are doing classifier free guidance
						lowerCAmelCase_   :							str									=			torch.cat([latents] * 2	) if do_classifier_free_guidance else latents
						lowerCAmelCase_   :							Union[str, Any]									=			self.scheduler.scale_model_input(a_						,			a_	)

						# predict the noise residual
						lowerCAmelCase_   :							List[Any]									=			self.unet(a_						,			a_						,			encoder_hidden_states=a_	).sample

						# perform guidance
						if do_classifier_free_guidance:
								lowerCAmelCase_				,	lowerCAmelCase_   :							Any									=			noise_pred.chunk(2	)
								lowerCAmelCase_   :							Optional[Any]									=			noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

						# compute the previous noisy sample x_t -> x_t-1
						lowerCAmelCase_   :							List[str]									=			self.scheduler.step(a_						,			a_						,			a_						,			**a_	).prev_sample

						# call the callback, if provided
						if callback is not None and i % callback_steps == 0:
								callback(a_						,			a_						,			a_	)

				lowerCAmelCase_   :							int									=			1 / 0.18215 * latents
				lowerCAmelCase_   :							Tuple									=			self.vae.decode(a_	).sample

				lowerCAmelCase_   :							List[Any]									=			(image / 2 + 0.5).clamp(0						,			1	)

				# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
				lowerCAmelCase_   :							List[Any]									=			image.cpu().permute(0						,			2						,			3						,			1	).float().numpy()

				if output_type == "pil":
						lowerCAmelCase_   :							Optional[int]									=			self.numpy_to_pil(a_	)

				if not return_dict:
						return image

				return StableDiffusionPipelineOutput(images=a_						,			nsfw_content_detected=a_	)








import argparse
import json
from collections import OrderedDict
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    SegformerConfig,
    SegformerForImageClassification,
    SegformerForSemanticSegmentation,
    SegformerImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
lowercase  							=					logging.get_logger(__name__)





def 						__UpperCAmelCase			(     a_							,     a_=False):
							snake_case_				=							OrderedDict()
							for key, value in state_dict.items():
														if encoder_only and not key.startswith('head'):
																					snake_case_				=							'segformer.encoder.' + key
														if key.startswith('backbone'):
																					snake_case_				=							key.replace('backbone'							,     'segformer.encoder')
														if "patch_embed" in key:
																					# replace for example patch_embed1 by patch_embeddings.0
																					snake_case_				=							key[key.find('patch_embed') + len('patch_embed')]
																					snake_case_				=							key.replace(f'''patch_embed{idx}'''							,     f'''patch_embeddings.{int(a_)-1}''')
														if "norm" in key:
																					snake_case_				=							key.replace('norm'							,     'layer_norm')
														if "segformer.encoder.layer_norm" in key:
																					# replace for example layer_norm1 by layer_norm.0
																					snake_case_				=							key[key.find('segformer.encoder.layer_norm') + len('segformer.encoder.layer_norm')]
																					snake_case_				=							key.replace(f'''layer_norm{idx}'''							,     f'''layer_norm.{int(a_)-1}''')
														if "layer_norm1" in key:
																					snake_case_				=							key.replace('layer_norm1'							,     'layer_norm_1')
														if "layer_norm2" in key:
																					snake_case_				=							key.replace('layer_norm2'							,     'layer_norm_2')
														if "block" in key:
																					# replace for example block1 by block.0
																					snake_case_				=							key[key.find('block') + len('block')]
																					snake_case_				=							key.replace(f'''block{idx}'''							,     f'''block.{int(a_)-1}''')
														if "attn.q" in key:
																					snake_case_				=							key.replace('attn.q'							,     'attention.self.query')
														if "attn.proj" in key:
																					snake_case_				=							key.replace('attn.proj'							,     'attention.output.dense')
														if "attn" in key:
																					snake_case_				=							key.replace('attn'							,     'attention.self')
														if "fc1" in key:
																					snake_case_				=							key.replace('fc1'							,     'dense1')
														if "fc2" in key:
																					snake_case_				=							key.replace('fc2'							,     'dense2')
														if "linear_pred" in key:
																					snake_case_				=							key.replace('linear_pred'							,     'classifier')
														if "linear_fuse" in key:
																					snake_case_				=							key.replace('linear_fuse.conv'							,     'linear_fuse')
																					snake_case_				=							key.replace('linear_fuse.bn'							,     'batch_norm')
														if "linear_c" in key:
																					# replace for example linear_c4 by linear_c.3
																					snake_case_				=							key[key.find('linear_c') + len('linear_c')]
																					snake_case_				=							key.replace(f'''linear_c{idx}'''							,     f'''linear_c.{int(a_)-1}''')
														if key.startswith('head'):
																					snake_case_				=							key.replace('head'							,     'classifier')
														snake_case_				=							value

							return new_state_dict





def 						__UpperCAmelCase			(     a_							,     a_):
							# for each of the encoder blocks:
							for i in range(config.num_encoder_blocks):
														for j in range(config.depths[i]):
																					# read in weights + bias of keys and values (which is a single matrix in the original implementation)
																					snake_case_				=							state_dict.pop(f'''segformer.encoder.block.{i}.{j}.attention.self.kv.weight''')
																					snake_case_				=							state_dict.pop(f'''segformer.encoder.block.{i}.{j}.attention.self.kv.bias''')
																					# next, add keys and values (in that order) to the state dict
																					snake_case_				=							kv_weight[
																					    : config.hidden_sizes[i], :
																					]
																					snake_case_				=							kv_bias[: config.hidden_sizes[i]]
																					snake_case_				=							kv_weight[
																					    config.hidden_sizes[i] :, :
																					]
																					snake_case_				=							kv_bias[
																					    config.hidden_sizes[i] :
																					]





def 						__UpperCAmelCase			(     ):
							snake_case_				=							'http://images.cocodataset.org/val2017/000000039769.jpg'
							snake_case_				=							Image.open(requests.get(a_							,     stream=a_).raw)

							return image





@torch.no_grad()
def 						__UpperCAmelCase			(     a_							,     a_							,     a_):
							snake_case_				=							SegformerConfig()
							snake_case_				=							False

							# set attributes based on model_name
							snake_case_				=							'huggingface/label-files'
							if "segformer" in model_name:
														snake_case_				=							model_name[len('segformer.') : len('segformer.') + 2]
														if "ade" in model_name:
																					snake_case_				=							1_50
																					snake_case_				=							'ade20k-id2label.json'
																					snake_case_				=							(1, 1_50, 1_28, 1_28)
														elif "city" in model_name:
																					snake_case_				=							19
																					snake_case_				=							'cityscapes-id2label.json'
																					snake_case_				=							(1, 19, 1_28, 1_28)
														else:
																					raise ValueError(f'''Model {model_name} not supported''')
							elif "mit" in model_name:
														snake_case_				=							True
														snake_case_				=							model_name[4:6]
														snake_case_				=							10_00
														snake_case_				=							'imagenet-1k-id2label.json'
														snake_case_				=							(1, 10_00)
							else:
														raise ValueError(f'''Model {model_name} not supported''')

							# set config attributes
							snake_case_				=							json.load(open(hf_hub_download(a_							,     a_							,     repo_type='dataset')							,     'r'))
							snake_case_				=							{int(a_): v for k, v in idalabel.items()}
							snake_case_				=							idalabel
							snake_case_				=							{v: k for k, v in idalabel.items()}
							if size == "b0":
														pass
							elif size == "b1":
														snake_case_				=							[64, 1_28, 3_20, 5_12]
														snake_case_				=							2_56
							elif size == "b2":
														snake_case_				=							[64, 1_28, 3_20, 5_12]
														snake_case_				=							7_68
														snake_case_				=							[3, 4, 6, 3]
							elif size == "b3":
														snake_case_				=							[64, 1_28, 3_20, 5_12]
														snake_case_				=							7_68
														snake_case_				=							[3, 4, 18, 3]
							elif size == "b4":
														snake_case_				=							[64, 1_28, 3_20, 5_12]
														snake_case_				=							7_68
														snake_case_				=							[3, 8, 27, 3]
							elif size == "b5":
														snake_case_				=							[64, 1_28, 3_20, 5_12]
														snake_case_				=							7_68
														snake_case_				=							[3, 6, 40, 3]
							else:
														raise ValueError(f'''Size {size} not supported''')

							# load image processor (only resize + normalize)
							snake_case_				=							SegformerImageProcessor(
							    image_scale=(5_12, 5_12)							,     keep_ratio=a_							,     align=a_							,     do_random_crop=a_)

							# prepare image
							snake_case_				=							prepare_img()
							snake_case_				=							image_processor(images=a_							,     return_tensors='pt').pixel_values

							logger.info(f'''Converting model {model_name}...''')

							# load original state dict
							if encoder_only:
														snake_case_				=							torch.load(a_							,     map_location=torch.device('cpu'))
							else:
														snake_case_				=							torch.load(a_							,     map_location=torch.device('cpu'))['state_dict']

							# rename keys
							snake_case_				=							rename_keys(a_							,     encoder_only=a_)
							if not encoder_only:
														del state_dict["decode_head.conv_seg.weight"]
														del state_dict["decode_head.conv_seg.bias"]

							# key and value matrices need special treatment
							read_in_k_v(a_							,     a_)

							# create HuggingFace model and load state dict
							if encoder_only:
														snake_case_				=							False
														snake_case_				=							SegformerForImageClassification(a_)
							else:
														snake_case_				=							SegformerForSemanticSegmentation(a_)
							model.load_state_dict(a_)
							model.eval()

							# forward pass
							snake_case_				=							model(a_)
							snake_case_				=							outputs.logits

							# set expected_slice based on model name
							# ADE20k checkpoints
							if model_name == "segformer.b0.512x512.ade.160k":
														snake_case_				=							torch.tensor(
														    [
														        [[-4.63_10, -5.52_32, -6.23_56], [-5.19_21, -6.14_44, -6.59_96], [-5.44_24, -6.27_90, -6.75_74]],
														        [[-12.13_91, -13.31_22, -13.95_54], [-12.87_32, -13.93_52, -14.35_63], [-12.94_38, -13.82_26, -14.25_13]],
														        [[-12.51_34, -13.46_86, -14.49_15], [-12.86_69, -14.43_43, -14.77_58], [-13.25_23, -14.58_19, -15.06_94]],
														    ])
							elif model_name == "segformer.b1.512x512.ade.160k":
														snake_case_				=							torch.tensor(
														    [
														        [[-7.58_20, -8.72_31, -8.32_15], [-8.06_00, -10.35_29, -10.03_04], [-7.52_08, -9.41_03, -9.62_39]],
														        [[-12.69_18, -13.89_94, -13.71_37], [-13.31_96, -15.75_23, -15.47_89], [-12.93_43, -14.87_57, -14.96_89]],
														        [[-11.19_11, -11.94_21, -11.32_43], [-11.33_42, -13.68_39, -13.35_81], [-10.39_09, -12.18_32, -12.48_58]],
														    ])
							elif model_name == "segformer.b2.512x512.ade.160k":
														snake_case_				=							torch.tensor(
														    [
														        [[-11.81_73, -14.38_50, -16.31_28], [-14.56_48, -16.58_04, -18.65_68], [-14.72_23, -15.73_87, -18.42_18]],
														        [[-15.72_90, -17.91_71, -19.44_23], [-18.31_05, -19.94_48, -21.46_61], [-17.92_96, -18.64_97, -20.79_10]],
														        [[-15.07_83, -17.03_36, -18.27_89], [-16.87_71, -18.68_70, -20.16_12], [-16.24_54, -17.14_26, -19.50_55]],
														    ])
							elif model_name == "segformer.b3.512x512.ade.160k":
														snake_case_				=							torch.tensor(
														    [
														        [[-9.08_78, -10.20_81, -10.18_91], [-9.31_44, -10.79_41, -10.98_43], [-9.22_94, -10.38_55, -10.57_04]],
														        [[-12.23_16, -13.90_68, -13.61_02], [-12.91_61, -14.37_02, -14.32_35], [-12.52_33, -13.71_74, -13.79_32]],
														        [[-14.62_75, -15.24_90, -14.97_27], [-14.34_00, -15.96_87, -16.28_27], [-14.14_84, -15.40_33, -15.89_37]],
														    ])
							elif model_name == "segformer.b4.512x512.ade.160k":
														snake_case_				=							torch.tensor(
														    [
														        [[-12.31_44, -13.24_47, -14.08_02], [-13.36_14, -14.58_16, -15.61_17], [-13.33_40, -14.44_33, -16.22_19]],
														        [[-19.27_81, -20.41_28, -20.75_06], [-20.61_53, -21.65_66, -22.09_98], [-19.98_00, -21.04_30, -22.14_94]],
														        [[-18.87_39, -19.78_04, -21.18_34], [-20.12_33, -21.67_65, -23.29_44], [-20.03_15, -21.26_41, -23.69_44]],
														    ])
							elif model_name == "segformer.b5.640x640.ade.160k":
														snake_case_				=							torch.tensor(
														    [
														        [[-9.55_24, -12.08_35, -11.73_48], [-10.52_29, -13.64_46, -14.56_62], [-9.58_42, -12.88_51, -13.94_14]],
														        [[-15.34_32, -17.53_23, -17.08_18], [-16.33_30, -18.92_55, -19.21_01], [-15.13_40, -17.78_48, -18.39_71]],
														        [[-12.60_72, -14.94_86, -14.66_31], [-13.76_29, -17.09_07, -17.77_45], [-12.78_99, -16.16_95, -17.16_71]],
														    ])
							# Cityscapes checkpoints
							elif model_name == "segformer.b0.1024x1024.city.160k":
														snake_case_				=							torch.tensor(
														    [
														        [[-11.92_95, -13.40_57, -14.81_06], [-13.34_31, -14.81_79, -15.37_81], [-14.28_36, -15.59_42, -16.15_88]],
														        [[-11.49_06, -12.80_67, -13.65_64], [-13.11_89, -14.05_00, -14.15_43], [-13.87_48, -14.51_36, -14.87_89]],
														        [[0.53_74, 0.10_67, -0.47_42], [0.11_41, -0.22_55, -0.70_99], [-0.30_00, -0.59_24, -1.31_05]],
														    ])
							elif model_name == "segformer.b0.512x1024.city.160k":
														snake_case_				=							torch.tensor(
														    [
														        [[-7.82_17, -9.87_67, -10.17_17], [-9.44_38, -10.90_58, -11.40_47], [-9.79_39, -12.34_95, -12.10_79]],
														        [[-7.15_14, -9.53_36, -10.08_60], [-9.77_76, -11.68_22, -11.84_39], [-10.14_11, -12.76_55, -12.89_72]],
														        [[0.30_21, 0.08_05, -0.23_10], [-0.03_28, -0.16_05, -0.27_14], [-0.14_08, -0.54_77, -0.69_76]],
														    ])
							elif model_name == "segformer.b0.640x1280.city.160k":
														snake_case_				=							torch.tensor(
														    [
														        [
														            [-1.1_372E01, -1.2_787E01, -1.3_477E01],
														            [-1.2_536E01, -1.4_194E01, -1.4_409E01],
														            [-1.3_217E01, -1.4_888E01, -1.5_327E01],
														        ],
														        [
														            [-1.4_791E01, -1.7_122E01, -1.8_277E01],
														            [-1.7_163E01, -1.9_192E01, -1.9_533E01],
														            [-1.7_897E01, -1.9_991E01, -2.0_315E01],
														        ],
														        [
														            [7.6_723E-01, 4.1_921E-01, -7.7_878E-02],
														            [4.7_772E-01, 9.5_557E-03, -2.8_082E-01],
														            [3.6_032E-01, -2.4_826E-01, -5.1_168E-01],
														        ],
														    ])
							elif model_name == "segformer.b0.768x768.city.160k":
														snake_case_				=							torch.tensor(
														    [
														        [[-9.49_59, -11.30_87, -11.74_79], [-11.00_25, -12.65_40, -12.33_19], [-11.40_64, -13.04_87, -12.99_05]],
														        [[-9.89_05, -11.30_84, -12.08_54], [-11.17_26, -12.76_98, -12.95_83], [-11.59_85, -13.32_78, -14.17_74]],
														        [[0.22_13, 0.01_92, -0.24_66], [-0.17_31, -0.42_13, -0.48_74], [-0.31_26, -0.65_41, -1.13_89]],
														    ])
							elif model_name == "segformer.b1.1024x1024.city.160k":
														snake_case_				=							torch.tensor(
														    [
														        [[-13.57_48, -13.91_11, -12.65_00], [-14.35_00, -15.36_83, -14.23_28], [-14.75_32, -16.04_24, -15.60_87]],
														        [[-17.16_51, -15.87_25, -12.96_53], [-17.25_80, -17.37_18, -14.82_23], [-16.60_58, -16.87_83, -16.74_52]],
														        [[-3.64_56, -3.02_09, -1.42_03], [-3.07_97, -3.19_59, -2.00_00], [-1.87_57, -1.92_17, -1.69_97]],
														    ])
							elif model_name == "segformer.b2.1024x1024.city.160k":
														snake_case_				=							torch.tensor(
														    [
														        [[-16.09_76, -16.48_56, -17.39_62], [-16.62_34, -19.03_42, -19.76_85], [-16.09_00, -18.06_61, -19.11_80]],
														        [[-18.47_50, -18.84_88, -19.50_74], [-19.40_30, -22.15_70, -22.59_77], [-19.11_91, -20.84_86, -22.37_83]],
														        [[-4.51_78, -5.50_37, -6.51_09], [-5.08_84, -7.21_74, -8.03_34], [-4.41_56, -5.81_17, -7.29_70]],
														    ])
							elif model_name == "segformer.b3.1024x1024.city.160k":
														snake_case_				=							torch.tensor(
														    [
														        [[-14.20_81, -14.47_32, -14.19_77], [-14.58_67, -16.44_23, -16.63_56], [-13.44_41, -14.96_85, -16.86_96]],
														        [[-14.45_76, -14.70_73, -15.04_51], [-15.08_16, -17.62_37, -17.98_73], [-14.42_13, -16.01_99, -18.59_92]],
														        [[-4.73_49, -4.95_88, -5.09_66], [-4.32_10, -6.93_25, -7.25_91], [-3.43_12, -4.74_84, -7.19_17]],
														    ])
							elif model_name == "segformer.b4.1024x1024.city.160k":
														snake_case_				=							torch.tensor(
														    [
														        [[-11.77_37, -11.95_26, -11.32_73], [-13.66_92, -14.45_74, -13.88_78], [-13.89_37, -14.69_24, -15.93_45]],
														        [[-14.67_06, -14.53_30, -14.13_06], [-16.15_02, -16.81_80, -16.42_69], [-16.83_38, -17.89_39, -20.17_46]],
														        [[1.04_91, 0.82_89, 1.03_10], [1.10_44, 0.52_19, 0.80_55], [1.08_99, 0.69_26, 0.55_90]],
														    ])
							elif model_name == "segformer.b5.1024x1024.city.160k":
														snake_case_				=							torch.tensor(
														    [
														        [[-12.56_41, -13.47_77, -13.06_84], [-13.95_87, -15.89_83, -16.65_57], [-13.31_09, -15.73_50, -16.31_41]],
														        [[-14.70_74, -15.43_52, -14.59_44], [-16.63_53, -18.16_63, -18.61_20], [-15.17_02, -18.03_29, -18.15_47]],
														        [[-1.79_90, -2.09_51, -1.77_84], [-2.63_97, -3.82_45, -3.96_86], [-1.52_64, -2.81_26, -2.93_16]],
														    ])
							else:
														snake_case_				=							logits.argmax(-1).item()
														print('Predicted class:'							,     model.config.idalabel[predicted_class_idx])

							# verify logits
							if not encoder_only:
														assert logits.shape == expected_shape
														assert torch.allclose(logits[0, :3, :3, :3]							,     a_							,     atol=1E-2)

							# finally, save model and image processor
							logger.info(f'''Saving PyTorch model and image processor to {pytorch_dump_folder_path}...''')
							Path(a_).mkdir(exist_ok=a_)
							model.save_pretrained(a_)
							image_processor.save_pretrained(a_)


if __name__ == "__main__":
	lowercase  							=					argparse.ArgumentParser()

	parser.add_argument(
	    "--model_name",
	    default="segformer.b0.512x512.ade.160k",
	    type=str,
	    help="Name of the model you'd like to convert.",
	)
	parser.add_argument(
	    "--checkpoint_path", default=None, type=str, help="Path to the original PyTorch checkpoint (.pth file)."
	)
	parser.add_argument(
	    "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
	)
	lowercase  							=					parser.parse_args()
	convert_segformer_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path)





