






'''simple docstring'''







import argparse
import json

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

from transformers import ConvNextConfig, SegformerImageProcessor, UperNetConfig, UperNetForSemanticSegmentation
def     A__      (					__lowerCAmelCase		:      List[str]    ):
						lowerCamelCase__								=	384
						if "tiny" in model_name:
												lowerCamelCase__								=	[3, 3, 9, 3]
												lowerCamelCase__								=	[96, 192, 384, 768]
						if "small" in model_name:
												lowerCamelCase__								=	[3, 3, 27, 3]
												lowerCamelCase__								=	[96, 192, 384, 768]
						if "base" in model_name:
												lowerCamelCase__								=	[3, 3, 27, 3]
												lowerCamelCase__								=	[128, 256, 512, 1024]
												lowerCamelCase__								=	512
						if "large" in model_name:
												lowerCamelCase__								=	[3, 3, 27, 3]
												lowerCamelCase__								=	[192, 384, 768, 1536]
												lowerCamelCase__								=	768
						if "xlarge" in model_name:
												lowerCamelCase__								=	[3, 3, 27, 3]
												lowerCamelCase__								=	[256, 512, 1024, 2048]
												lowerCamelCase__								=	1024

						# set label information
						lowerCamelCase__								=	150
						lowerCamelCase__								=	"""huggingface/label-files"""
						lowerCamelCase__								=	"""ade20k-id2label.json"""
						lowerCamelCase__								=	json.load(open(hf_hub_download(__lowerCAmelCase ,     __lowerCAmelCase ,     repo_type="""dataset"""    ) ,     """r"""    )    )
						lowerCamelCase__								=	{int(__lowerCAmelCase    ): v for k, v in idalabel.items()}
						lowerCamelCase__								=	{v: k for k, v in idalabel.items()}

						lowerCamelCase__								=	ConvNextConfig(
						    depths=__lowerCAmelCase ,     hidden_sizes=__lowerCAmelCase ,     out_features=["""stage1""", """stage2""", """stage3""", """stage4"""]    )
						lowerCamelCase__								=	UperNetConfig(
						    backbone_config=__lowerCAmelCase ,     auxiliary_in_channels=__lowerCAmelCase ,     num_labels=__lowerCAmelCase ,     idalabel=__lowerCAmelCase ,     labelaid=__lowerCAmelCase ,     )

						return config
def     A__      (					__lowerCAmelCase		:      Optional[Any]    ):
						lowerCamelCase__								=	[]

						# fmt: off
						# stem
						rename_keys.append(("""backbone.downsample_layers.0.0.weight""", """backbone.embeddings.patch_embeddings.weight""")    )
						rename_keys.append(("""backbone.downsample_layers.0.0.bias""", """backbone.embeddings.patch_embeddings.bias""")    )
						rename_keys.append(("""backbone.downsample_layers.0.1.weight""", """backbone.embeddings.layernorm.weight""")    )
						rename_keys.append(("""backbone.downsample_layers.0.1.bias""", """backbone.embeddings.layernorm.bias""")    )
						# stages
						for i in range(len(config.backbone_config.depths    )    ):
												for j in range(config.backbone_config.depths[i]    ):
																		rename_keys.append((F'''backbone.stages.{i}.{j}.gamma''', F'''backbone.encoder.stages.{i}.layers.{j}.layer_scale_parameter''')    )
																		rename_keys.append((F'''backbone.stages.{i}.{j}.depthwise_conv.weight''', F'''backbone.encoder.stages.{i}.layers.{j}.dwconv.weight''')    )
																		rename_keys.append((F'''backbone.stages.{i}.{j}.depthwise_conv.bias''', F'''backbone.encoder.stages.{i}.layers.{j}.dwconv.bias''')    )
																		rename_keys.append((F'''backbone.stages.{i}.{j}.norm.weight''', F'''backbone.encoder.stages.{i}.layers.{j}.layernorm.weight''')    )
																		rename_keys.append((F'''backbone.stages.{i}.{j}.norm.bias''', F'''backbone.encoder.stages.{i}.layers.{j}.layernorm.bias''')    )
																		rename_keys.append((F'''backbone.stages.{i}.{j}.pointwise_conv1.weight''', F'''backbone.encoder.stages.{i}.layers.{j}.pwconv1.weight''')    )
																		rename_keys.append((F'''backbone.stages.{i}.{j}.pointwise_conv1.bias''', F'''backbone.encoder.stages.{i}.layers.{j}.pwconv1.bias''')    )
																		rename_keys.append((F'''backbone.stages.{i}.{j}.pointwise_conv2.weight''', F'''backbone.encoder.stages.{i}.layers.{j}.pwconv2.weight''')    )
																		rename_keys.append((F'''backbone.stages.{i}.{j}.pointwise_conv2.bias''', F'''backbone.encoder.stages.{i}.layers.{j}.pwconv2.bias''')    )
												if i > 0:
																		rename_keys.append((F'''backbone.downsample_layers.{i}.0.weight''', F'''backbone.encoder.stages.{i}.downsampling_layer.0.weight''')    )
																		rename_keys.append((F'''backbone.downsample_layers.{i}.0.bias''', F'''backbone.encoder.stages.{i}.downsampling_layer.0.bias''')    )
																		rename_keys.append((F'''backbone.downsample_layers.{i}.1.weight''', F'''backbone.encoder.stages.{i}.downsampling_layer.1.weight''')    )
																		rename_keys.append((F'''backbone.downsample_layers.{i}.1.bias''', F'''backbone.encoder.stages.{i}.downsampling_layer.1.bias''')    )

												rename_keys.append((F'''backbone.norm{i}.weight''', F'''backbone.hidden_states_norms.stage{i+1}.weight''')    )
												rename_keys.append((F'''backbone.norm{i}.bias''', F'''backbone.hidden_states_norms.stage{i+1}.bias''')    )

						# decode head
						rename_keys.extend(
						    [
						        ("""decode_head.conv_seg.weight""", """decode_head.classifier.weight"""),
						        ("""decode_head.conv_seg.bias""", """decode_head.classifier.bias"""),
						        ("""auxiliary_head.conv_seg.weight""", """auxiliary_head.classifier.weight"""),
						        ("""auxiliary_head.conv_seg.bias""", """auxiliary_head.classifier.bias"""),
						    ]    )
						# fmt: on

						return rename_keys
def     A__      (					__lowerCAmelCase		:      Optional[Any] ,     __lowerCAmelCase		:      Union[str, Any] ,     __lowerCAmelCase		:      Tuple    ):
						lowerCamelCase__								=	dct.pop(__lowerCAmelCase    )
						lowerCamelCase__								=	val
def     A__      (					__lowerCAmelCase		:      Union[str, Any] ,     __lowerCAmelCase		:      Union[str, Any] ,     __lowerCAmelCase		:      int    ):
						lowerCamelCase__								=	{
						    """upernet-convnext-tiny""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_tiny_fp16_512x512_160k_ade20k/upernet_convnext_tiny_fp16_512x512_160k_ade20k_20220227_124553-cad485de.pth""",
						    """upernet-convnext-small""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_small_fp16_512x512_160k_ade20k/upernet_convnext_small_fp16_512x512_160k_ade20k_20220227_131208-1b1e394f.pth""",
						    """upernet-convnext-base""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_base_fp16_512x512_160k_ade20k/upernet_convnext_base_fp16_512x512_160k_ade20k_20220227_181227-02a24fc6.pth""",
						    """upernet-convnext-large""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_large_fp16_640x640_160k_ade20k/upernet_convnext_large_fp16_640x640_160k_ade20k_20220226_040532-e57aa54d.pth""",
						    """upernet-convnext-xlarge""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_xlarge_fp16_640x640_160k_ade20k/upernet_convnext_xlarge_fp16_640x640_160k_ade20k_20220226_080344-95fc38c2.pth""",
						}
						lowerCamelCase__								=	model_name_to_url[model_name]
						lowerCamelCase__								=	torch.hub.load_state_dict_from_url(__lowerCAmelCase ,     map_location="""cpu"""    )["""state_dict"""]

						lowerCamelCase__								=	get_upernet_config(__lowerCAmelCase    )
						lowerCamelCase__								=	UperNetForSemanticSegmentation(__lowerCAmelCase    )
						model.eval()

						# replace "bn" => "batch_norm"
						for key in state_dict.copy().keys():
												lowerCamelCase__								=	state_dict.pop(__lowerCAmelCase    )
												if "bn" in key:
																		lowerCamelCase__								=	key.replace("""bn""" ,     """batch_norm"""    )
												lowerCamelCase__								=	val

						# rename keys
						lowerCamelCase__								=	create_rename_keys(__lowerCAmelCase    )
						for src, dest in rename_keys:
												rename_key(__lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase    )

						model.load_state_dict(__lowerCAmelCase    )

						# verify on image
						lowerCamelCase__								=	"""https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"""
						lowerCamelCase__								=	Image.open(requests.get(__lowerCAmelCase ,     stream=__lowerCAmelCase    ).raw    ).convert("""RGB"""    )

						lowerCamelCase__								=	SegformerImageProcessor()
						lowerCamelCase__								=	processor(__lowerCAmelCase ,     return_tensors="""pt"""    ).pixel_values

						with torch.no_grad():
												lowerCamelCase__								=	model(__lowerCAmelCase    )

						if model_name == "upernet-convnext-tiny":
												lowerCamelCase__								=	torch.tensor(
												    [[-8.8110, -8.8110, -8.6521], [-8.8110, -8.8110, -8.6521], [-8.7746, -8.7746, -8.6130]]    )
						elif model_name == "upernet-convnext-small":
												lowerCamelCase__								=	torch.tensor(
												    [[-8.8236, -8.8236, -8.6771], [-8.8236, -8.8236, -8.6771], [-8.7638, -8.7638, -8.6240]]    )
						elif model_name == "upernet-convnext-base":
												lowerCamelCase__								=	torch.tensor(
												    [[-8.8558, -8.8558, -8.6905], [-8.8558, -8.8558, -8.6905], [-8.7669, -8.7669, -8.6021]]    )
						elif model_name == "upernet-convnext-large":
												lowerCamelCase__								=	torch.tensor(
												    [[-8.6660, -8.6660, -8.6210], [-8.6660, -8.6660, -8.6210], [-8.6310, -8.6310, -8.5964]]    )
						elif model_name == "upernet-convnext-xlarge":
												lowerCamelCase__								=	torch.tensor(
												    [[-8.4980, -8.4980, -8.3977], [-8.4980, -8.4980, -8.3977], [-8.4379, -8.4379, -8.3412]]    )
						print("""Logits:""" ,     outputs.logits[0, 0, :3, :3]    )
						assert torch.allclose(outputs.logits[0, 0, :3, :3] ,     __lowerCAmelCase ,     atol=1e-4    )
						print("""Looks ok!"""    )

						if pytorch_dump_folder_path is not None:
												print(F'''Saving model {model_name} to {pytorch_dump_folder_path}'''    )
												model.save_pretrained(__lowerCAmelCase    )
												print(F'''Saving processor to {pytorch_dump_folder_path}'''    )
												processor.save_pretrained(__lowerCAmelCase    )

						if push_to_hub:
												print(F'''Pushing model and processor for {model_name} to hub'''    )
												model.push_to_hub(F'''openmmlab/{model_name}'''    )
												processor.push_to_hub(F'''openmmlab/{model_name}'''    )


if __name__ == "__main__":
							UpperCamelCase     :   Tuple	      =		argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '--model_name',
							    default='upernet-convnext-tiny',
							    type=str,
							    choices=[F'upernet-convnext-{size}' for size in ['tiny', 'small', 'base', 'large', 'xlarge']],
							    help='Name of the ConvNext UperNet model you\'d like to convert.',
							)
							parser.add_argument(
							    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
							)
							parser.add_argument(
							    '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
							)

							UpperCamelCase     :   Any	      =		parser.parse_args()
							convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






'''simple docstring'''

import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel

from diffusers import DDIMScheduler, LDMPipeline, UNetaDModel, VQModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()

class       __magic_name__		(       unittest.TestCase  ):







							@property
							def 				__lowercase     (				self   :			int						):
									torch.manual_seed(0						)
									_a	: int								=				UNetaDModel(
									    block_out_channels=(32, 64)      ,layers_per_block=2      ,sample_size=32      ,in_channels=3      ,out_channels=3      ,down_block_types=('DownBlock2D', 'AttnDownBlock2D')      ,up_block_types=('AttnUpBlock2D', 'UpBlock2D')      ,)
									return model







							@property
							def 				__lowercase     (				self   :			Any						):
									torch.manual_seed(0						)
									_a	: int								=				VQModel(
									    block_out_channels=[32, 64]      ,in_channels=3      ,out_channels=3      ,down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']      ,up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']      ,latent_channels=3      ,)
									return model







							@property
							def 				__lowercase     (				self   :			Union[str, Any]						):
									torch.manual_seed(0						)
									_a	: Dict								=				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=1000      ,)
									return CLIPTextModel(_UpperCAmelCase						)







							def 				__lowercase     (				self   :			Dict						):
									_a	: Union[str, Any]								=				self.dummy_uncond_unet
									_a	: Union[str, Any]								=				DDIMScheduler()
									_a	: Union[str, Any]								=				self.dummy_vq_model

									_a	: List[Any]								=				LDMPipeline(unet=_UpperCAmelCase      ,vqvae=_UpperCAmelCase      ,scheduler=_UpperCAmelCase						)
									ldm.to(_UpperCAmelCase						)
									ldm.set_progress_bar_config(disable=_UpperCAmelCase						)

									_a	: Dict								=				torch.manual_seed(0						)
									_a	: int								=				ldm(generator=_UpperCAmelCase      ,num_inference_steps=2      ,output_type='numpy'						).images

									_a	: Tuple								=				torch.manual_seed(0						)
									_a	: List[Any]								=				ldm(generator=_UpperCAmelCase      ,num_inference_steps=2      ,output_type='numpy'      ,return_dict=_UpperCAmelCase						)[0]

									_a	: Any								=				image[0, -3:, -3:, -1]
									_a	: List[Any]								=				image_from_tuple[0, -3:, -3:, -1]

									assert image.shape == (1, 64, 64, 3)
									_a	: Union[str, Any]								=				np.array([0.85_12, 0.8_18, 0.64_11, 0.68_08, 0.44_65, 0.56_18, 0.46, 0.62_31, 0.51_72]						)
									_a	: Any								=				1E-2 if torch_device != 'mps' else 3E-2

									assert np.abs(image_slice.flatten() - expected_slice						).max() < tolerance
									assert np.abs(image_from_tuple_slice.flatten() - expected_slice						).max() < tolerance

@slow
@require_torch
class       __magic_name__		(       unittest.TestCase  ):







							def 				__lowercase     (				self   :			Optional[int]						):
									_a	: Dict								=				LDMPipeline.from_pretrained('CompVis/ldm-celebahq-256'						)
									ldm.to(_UpperCAmelCase						)
									ldm.set_progress_bar_config(disable=_UpperCAmelCase						)

									_a	: str								=				torch.manual_seed(0						)
									_a	: Optional[Any]								=				ldm(generator=_UpperCAmelCase      ,num_inference_steps=5      ,output_type='numpy'						).images

									_a	: Optional[Any]								=				image[0, -3:, -3:, -1]

									assert image.shape == (1, 256, 256, 3)
									_a	: Union[str, Any]								=				np.array([0.43_99, 0.4_49_75, 0.4_68_25, 0.4_74, 0.43_59, 0.45_81, 0.4_50_95, 0.43_41, 0.44_47]						)
									_a	: List[str]								=				1E-2 if torch_device != 'mps' else 3E-2

									assert np.abs(image_slice.flatten() - expected_slice						).max() < tolerance




