





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       UpperCamelCase__(			UpperCamelCase__     :			Tuple    )->Tuple:
       A__		   =		3_84
       if "tiny" in model_name:
              A__		   =		[3, 3, 9, 3]
              A__		   =		[96, 1_92, 3_84, 7_68]
       if "small" in model_name:
              A__		   =		[3, 3, 27, 3]
              A__		   =		[96, 1_92, 3_84, 7_68]
       if "base" in model_name:
              A__		   =		[3, 3, 27, 3]
              A__		   =		[1_28, 2_56, 5_12, 10_24]
              A__		   =		5_12
       if "large" in model_name:
              A__		   =		[3, 3, 27, 3]
              A__		   =		[1_92, 3_84, 7_68, 15_36]
              A__		   =		7_68
       if "xlarge" in model_name:
              A__		   =		[3, 3, 27, 3]
              A__		   =		[2_56, 5_12, 10_24, 20_48]
              A__		   =		10_24

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

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

       return config


def       UpperCamelCase__(			UpperCamelCase__     :			Union[str, Any]    )->Optional[Any]:
       A__		   =		[]

       # 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       UpperCamelCase__(			UpperCamelCase__     :			Optional[int]					,	UpperCamelCase__     :			List[Any]					,	UpperCamelCase__     :			int    )->Tuple:
       A__		   =		dct.pop(UpperCamelCase__    )
       A__		   =		val


def       UpperCamelCase__(			UpperCamelCase__     :			int					,	UpperCamelCase__     :			Any					,	UpperCamelCase__     :			int    )->Tuple:
       A__		   =		{
           '''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''',
       }
       A__		   =		model_name_to_url[model_name]
       A__		   =		torch.hub.load_state_dict_from_url(UpperCamelCase__					,	map_location='''cpu'''    )['''state_dict''']

       A__		   =		get_upernet_config(UpperCamelCase__    )
       A__		   =		UperNetForSemanticSegmentation(UpperCamelCase__    )
       model.eval()

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

       # rename keys
       A__		   =		create_rename_keys(UpperCamelCase__    )
       for src, dest in rename_keys:
              rename_key(UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__    )

       model.load_state_dict(UpperCamelCase__    )

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

       A__		   =		SegformerImageProcessor()
       A__		   =		processor(UpperCamelCase__					,	return_tensors='''pt'''    ).pixel_values

       with torch.no_grad():
              A__		   =		model(UpperCamelCase__    )

       if model_name == "upernet-convnext-tiny":
              A__		   =		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":
              A__		   =		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":
              A__		   =		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":
              A__		   =		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":
              A__		   =		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]					,	UpperCamelCase__					,	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(UpperCamelCase__    )
              print(f"Saving processor to {pytorch_dump_folder_path}"    )
              processor.save_pretrained(UpperCamelCase__    )

       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__":
       a__: Any							   =						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.'
       )

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












import math


def       UpperCamelCase__(			UpperCamelCase__     :			int    )->list:
       A__		   =		[True] * n
       A__		   =		False
       A__		   =		False
       A__		   =		True

       for i in range(3					,	int(n**0.5 + 1    )					,	2    ):
              A__		   =		i * 2
              while index < n:
                     A__		   =		False
                     A__		   =		index + i

       A__		   =		[2]

       for i in range(3					,	UpperCamelCase__					,	2    ):
              if is_prime[i]:
                     primes.append(UpperCamelCase__    )

       return primes


def       UpperCamelCase__(			UpperCamelCase__     :			int = 99_99_66_66_33_33    )->int:
       A__		   =		math.floor(math.sqrt(UpperCamelCase__    )    ) + 1_00
       A__		   =		prime_sieve(UpperCamelCase__    )

       A__		   =		0
       A__		   =		0
       A__		   =		primes[prime_index]

       while (last_prime**2) <= limit:
              A__		   =		primes[prime_index + 1]

              A__		   =		last_prime**2
              A__		   =		next_prime**2

              # Get numbers divisible by lps(current)
              A__		   =		lower_bound + last_prime
              while upper_bound > current <= limit:
                     matches_sum += current
                     current += last_prime

              # Reset the upper_bound
              while (upper_bound - next_prime) > limit:
                     upper_bound -= next_prime

              # Add the numbers divisible by ups(current)
              A__		   =		upper_bound - next_prime
              while current > lower_bound:
                     matches_sum += current
                     current -= next_prime

              # Remove the numbers divisible by both ups and lps
              A__		   =		0
              while upper_bound > current <= limit:
                     if current <= lower_bound:
                            # Increment the current number
                            current += last_prime * next_prime
                            continue

                     if current > limit:
                            break

                     # Remove twice since it was added by both ups and lps
                     matches_sum -= current * 2

                     # Increment the current number
                     current += last_prime * next_prime

              # Setup for next pair
              A__		   =		next_prime
              prime_index += 1

       return matches_sum


if __name__ == "__main__":
       print(solution())





