


'''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    lowerCamelCase__   ( A     :    str  ):




						'''simple docstring'''






						UpperCAmelCase         =			3_84
						if "tiny" in model_name:
												UpperCAmelCase         =			[3, 3, 9, 3]
												UpperCAmelCase         =			[96, 1_92, 3_84, 7_68]
						if "small" in model_name:
												UpperCAmelCase         =			[3, 3, 27, 3]
												UpperCAmelCase         =			[96, 1_92, 3_84, 7_68]
						if "base" in model_name:
												UpperCAmelCase         =			[3, 3, 27, 3]
												UpperCAmelCase         =			[1_28, 2_56, 5_12, 10_24]
												UpperCAmelCase         =			5_12
						if "large" in model_name:
												UpperCAmelCase         =			[3, 3, 27, 3]
												UpperCAmelCase         =			[1_92, 3_84, 7_68, 15_36]
												UpperCAmelCase         =			7_68
						if "xlarge" in model_name:
												UpperCAmelCase         =			[3, 3, 27, 3]
												UpperCAmelCase         =			[2_56, 5_12, 10_24, 20_48]
												UpperCAmelCase         =			10_24

						# set label information
						UpperCAmelCase         =			1_50
						UpperCAmelCase         =			'''huggingface/label-files'''
						UpperCAmelCase         =			'''ade20k-id2label.json'''
						UpperCAmelCase         =			json.load(open(hf_hub_download(A   ,     A   ,     repo_type='''dataset'''  )   ,     '''r'''  )  )
						UpperCAmelCase         =			{int(A  ): v for k, v in idalabel.items()}
						UpperCAmelCase         =			{v: k for k, v in idalabel.items()}

						UpperCAmelCase         =			ConvNextConfig(
						    depths=A   ,     hidden_sizes=A   ,     out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4''']  )
						UpperCAmelCase         =			UperNetConfig(
						    backbone_config=A   ,     auxiliary_in_channels=A   ,     num_labels=A   ,     idalabel=A   ,     labelaid=A   ,     )

						return config


def    lowerCamelCase__   ( A     :    Tuple  ):




						'''simple docstring'''






						UpperCAmelCase         =			[]

						# 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    lowerCamelCase__   ( A     :    Optional[Any]   ,     A     :    List[Any]   ,     A     :    Any  ):




						'''simple docstring'''






						UpperCAmelCase         =			dct.pop(A  )
						UpperCAmelCase         =			val


def    lowerCamelCase__   ( A     :    Union[str, Any]   ,     A     :    Union[str, Any]   ,     A     :    str  ):




						'''simple docstring'''






						UpperCAmelCase         =			{
						    '''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''',
						}
						UpperCAmelCase         =			model_name_to_url[model_name]
						UpperCAmelCase         =			torch.hub.load_state_dict_from_url(A   ,     map_location='''cpu'''  )['''state_dict''']

						UpperCAmelCase         =			get_upernet_config(A  )
						UpperCAmelCase         =			UperNetForSemanticSegmentation(A  )
						model.eval()

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

						# rename keys
						UpperCAmelCase         =			create_rename_keys(A  )
						for src, dest in rename_keys:
												rename_key(A   ,     A   ,     A  )

						model.load_state_dict(A  )

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

						UpperCAmelCase         =			SegformerImageProcessor()
						UpperCAmelCase         =			processor(A   ,     return_tensors='''pt'''  ).pixel_values

						with torch.no_grad():
												UpperCAmelCase         =			model(A  )

						if model_name == "upernet-convnext-tiny":
												UpperCAmelCase         =			torch.tensor(
												    [[-8.8_110, -8.8_110, -8.6_521], [-8.8_110, -8.8_110, -8.6_521], [-8.7_746, -8.7_746, -8.6_130]]  )
						elif model_name == "upernet-convnext-small":
												UpperCAmelCase         =			torch.tensor(
												    [[-8.8_236, -8.8_236, -8.6_771], [-8.8_236, -8.8_236, -8.6_771], [-8.7_638, -8.7_638, -8.6_240]]  )
						elif model_name == "upernet-convnext-base":
												UpperCAmelCase         =			torch.tensor(
												    [[-8.8_558, -8.8_558, -8.6_905], [-8.8_558, -8.8_558, -8.6_905], [-8.7_669, -8.7_669, -8.6_021]]  )
						elif model_name == "upernet-convnext-large":
												UpperCAmelCase         =			torch.tensor(
												    [[-8.6_660, -8.6_660, -8.6_210], [-8.6_660, -8.6_660, -8.6_210], [-8.6_310, -8.6_310, -8.5_964]]  )
						elif model_name == "upernet-convnext-xlarge":
												UpperCAmelCase         =			torch.tensor(
												    [[-8.4_980, -8.4_980, -8.3_977], [-8.4_980, -8.4_980, -8.3_977], [-8.4_379, -8.4_379, -8.3_412]]  )
						print('''Logits:'''   ,     outputs.logits[0, 0, :3, :3]  )
						assert torch.allclose(outputs.logits[0, 0, :3, :3]   ,     A   ,     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(A  )
												print(f"""Saving processor to {pytorch_dump_folder_path}"""  )
												processor.save_pretrained(A  )

						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__":
	_lowercase    :    List[str]			     =  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."""
	)

	_lowercase    :    int			     =  parser.parse_args()
	convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









'''simple docstring'''



import math
from typing import Optional

import numpy as np

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowercase    :    Union[str, Any]			     =  logging.get_logger(__name__)

_lowercase    :    List[str]			     =  {
    """facebook/encodec_24khz""": """https://huggingface.co/facebook/encodec_24khz/resolve/main/config.json""",
    """facebook/encodec_48khz""": """https://huggingface.co/facebook/encodec_48khz/resolve/main/config.json""",
}






class  UpperCamelCase__(     lowerCAmelCase ):
				__magic_name__       :   List[str]							= "encodec"




				def __init__(	self    :				List[str]   ,					lowerCAmelCase    :				int=[1.5, 3.0, 6.0, 12.0, 24.0]   ,					lowerCAmelCase    :				Tuple=24000   ,					lowerCAmelCase    :				List[Any]=1   ,					lowerCAmelCase    :				Union[str, Any]=False   ,					lowerCAmelCase    :				str=None   ,					lowerCAmelCase    :				Dict=None   ,					lowerCAmelCase    :				str=128   ,					lowerCAmelCase    :				Any=32   ,					lowerCAmelCase    :				Any=1   ,					lowerCAmelCase    :				List[Any]=[8, 5, 4, 2]   ,					lowerCAmelCase    :				Union[str, Any]="weight_norm"   ,					lowerCAmelCase    :				str=7   ,					lowerCAmelCase    :				Optional[int]=7   ,					lowerCAmelCase    :				Any=3   ,					lowerCAmelCase    :				Tuple=2   ,					lowerCAmelCase    :				Optional[Any]=True   ,					lowerCAmelCase    :				List[str]="reflect"   ,					lowerCAmelCase    :				Optional[int]=2   ,					lowerCAmelCase    :				List[Any]=2   ,					lowerCAmelCase    :				Union[str, Any]=1.0   ,					lowerCAmelCase    :				Optional[Any]=1024   ,					lowerCAmelCase    :				Optional[Any]=None   ,					lowerCAmelCase    :				str=True   ,					**lowerCAmelCase    :				str   ,					)->  List[str]:

										"""simple docstring"""




										UpperCAmelCase         =			target_bandwidths
										UpperCAmelCase         =			sampling_rate
										UpperCAmelCase         =			audio_channels
										UpperCAmelCase         =			normalize
										UpperCAmelCase         =			chunk_length_s
										UpperCAmelCase         =			overlap
										UpperCAmelCase         =			hidden_size
										UpperCAmelCase         =			num_filters
										UpperCAmelCase         =			num_residual_layers
										UpperCAmelCase         =			upsampling_ratios
										UpperCAmelCase         =			norm_type
										UpperCAmelCase         =			kernel_size
										UpperCAmelCase         =			last_kernel_size
										UpperCAmelCase         =			residual_kernel_size
										UpperCAmelCase         =			dilation_growth_rate
										UpperCAmelCase         =			use_causal_conv
										UpperCAmelCase         =			pad_mode
										UpperCAmelCase         =			compress
										UpperCAmelCase         =			num_lstm_layers
										UpperCAmelCase         =			trim_right_ratio
										UpperCAmelCase         =			codebook_size
										UpperCAmelCase         =			codebook_dim if codebook_dim is not None else hidden_size
										UpperCAmelCase         =			use_conv_shortcut

										if self.norm_type not in ["weight_norm", "time_group_norm"]:
																raise ValueError(
																    F"""self.norm_type must be one of `\"weight_norm\"`, `\"time_group_norm\"`), got {self.norm_type}"""   )

										super().__init__(**lowerCAmelCase   )




				@property
				def   a__(	self    :				str   )->  Optional[int]:

										"""simple docstring"""




										if self.chunk_length_s is None:
																return None
										else:
																return int(self.chunk_length_s * self.sampling_rate   )




				@property
				def   a__(	self    :				List[str]   )->  Optional[int]:

										"""simple docstring"""




										if self.chunk_length_s is None or self.overlap is None:
																return None
										else:
																return max(1   ,					int((1.0 - self.overlap) * self.chunk_length   )   )




				@property
				def   a__(	self    :				List[Any]   )->  int:

										"""simple docstring"""




										UpperCAmelCase         =			np.prod(self.upsampling_ratios   )
										return math.ceil(self.sampling_rate / hop_length   )





				@property
				def   a__(	self    :				int   )->  int:

										"""simple docstring"""




										return int(1000 * self.target_bandwidths[-1] // (self.frame_rate * 10)   )





