



import argparse
import json
import os

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

from transformers import ConvNextImageProcessor, ConvNextVaConfig, ConvNextVaForImageClassification
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


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




def       UpperCAmelCase__	(							__snake_case							)				->	List[str]:
  _A          =      ConvNextVaConfig()

  if "atto" in checkpoint_url:
    _A          =      [2, 2, 6, 2]
    _A          =      [40, 80, 160, 320]
  if "femto" in checkpoint_url:
    _A          =      [2, 2, 6, 2]
    _A          =      [48, 96, 192, 384]
  if "pico" in checkpoint_url:
    _A          =      [2, 2, 6, 2]
    _A          =      [64, 128, 256, 512]
  if "nano" in checkpoint_url:
    _A          =      [2, 2, 8, 2]
    _A          =      [80, 160, 320, 640]
  if "tiny" in checkpoint_url:
    _A          =      [3, 3, 9, 3]
    _A          =      [96, 192, 384, 768]
  if "base" in checkpoint_url:
    _A          =      [3, 3, 27, 3]
    _A          =      [128, 256, 512, 1_024]
  if "large" in checkpoint_url:
    _A          =      [3, 3, 27, 3]
    _A          =      [192, 384, 768, 1_536]
  if "huge" in checkpoint_url:
    _A          =      [3, 3, 27, 3]
    _A          =      [352, 704, 1_408, 2_816]

  _A          =      1_000
  _A          =      '''imagenet-1k-id2label.json'''
  _A          =      (1, 1_000)

  _A          =      '''huggingface/label-files'''
  _A          =      num_labels
  _A          =      json.load(open(hf_hub_download(__snake_case  ,			__snake_case  ,			repo_type='''dataset'''							)  ,			'''r'''							)							)
  _A          =      {int(__snake_case							): v for k, v in idalabel.items()}

  _A          =      idalabel
  _A          =      {v: k for k, v in idalabel.items()}
  _A          =      hidden_sizes
  _A          =      depths

  return config, expected_shape




def       UpperCAmelCase__	(							__snake_case							)				->	str:
  if "downsample_layers.0.0" in name:
    _A          =      name.replace('''downsample_layers.0.0'''  ,			'''embeddings.patch_embeddings'''							)
  if "downsample_layers.0.1" in name:
    _A          =      name.replace('''downsample_layers.0.1'''  ,			'''embeddings.norm'''							)  # we rename to layernorm later on
  if "downsample_layers.1.0" in name:
    _A          =      name.replace('''downsample_layers.1.0'''  ,			'''stages.1.downsampling_layer.0'''							)
  if "downsample_layers.1.1" in name:
    _A          =      name.replace('''downsample_layers.1.1'''  ,			'''stages.1.downsampling_layer.1'''							)
  if "downsample_layers.2.0" in name:
    _A          =      name.replace('''downsample_layers.2.0'''  ,			'''stages.2.downsampling_layer.0'''							)
  if "downsample_layers.2.1" in name:
    _A          =      name.replace('''downsample_layers.2.1'''  ,			'''stages.2.downsampling_layer.1'''							)
  if "downsample_layers.3.0" in name:
    _A          =      name.replace('''downsample_layers.3.0'''  ,			'''stages.3.downsampling_layer.0'''							)
  if "downsample_layers.3.1" in name:
    _A          =      name.replace('''downsample_layers.3.1'''  ,			'''stages.3.downsampling_layer.1'''							)
  if "stages" in name and "downsampling_layer" not in name:
    # stages.0.0. for instance should be renamed to stages.0.layers.0.
    _A          =      name[: len('''stages.0'''							)] + '''.layers''' + name[len('''stages.0'''							) :]
  if "gamma" in name:
    _A          =      name.replace('''gamma'''  ,			'''weight'''							)
  if "beta" in name:
    _A          =      name.replace('''beta'''  ,			'''bias'''							)
  if "stages" in name:
    _A          =      name.replace('''stages'''  ,			'''encoder.stages'''							)
  if "norm" in name:
    _A          =      name.replace('''norm'''  ,			'''layernorm'''							)
  if "head" in name:
    _A          =      name.replace('''head'''  ,			'''classifier'''							)

  return name




def       UpperCAmelCase__	(							)				->	List[Any]:
  _A          =      '''http://images.cocodataset.org/val2017/000000039769.jpg'''
  _A          =      Image.open(requests.get(__snake_case  ,			stream=__snake_case							).raw							)
  return im




def       UpperCAmelCase__	(							__snake_case							)				->	Tuple:
  if "224" in checkpoint_url:
    _A          =      224
    _A          =      224 / 256
  elif "384" in checkpoint_url:
    _A          =      384
    _A          =      None
  else:
    _A          =      512
    _A          =      None

  return ConvNextImageProcessor(
      size=__snake_case  ,			crop_pct=__snake_case  ,			image_mean=[0.4_85, 0.4_56, 0.4_06]  ,			image_std=[0.2_29, 0.2_24, 0.2_25]  ,			resample=PILImageResampling.BICUBIC  ,			)




@torch.no_grad()
def       UpperCAmelCase__	(							__snake_case  ,			__snake_case  ,			__snake_case  ,			__snake_case							)				->	Tuple:
  print('''Downloading original model from checkpoint...'''							)
  # define ConvNeXTV2 configuration based on URL
  _A      ,				_A          =      get_convnextva_config(__snake_case							)
  # load original state_dict from URL
  _A          =      torch.hub.load_state_dict_from_url(__snake_case							)['''model''']

  print('''Converting model parameters...'''							)
  # rename keys
  for key in state_dict.copy().keys():
    _A          =      state_dict.pop(__snake_case							)
    _A          =      val
  # add prefix to all keys expect classifier head
  for key in state_dict.copy().keys():
    _A          =      state_dict.pop(__snake_case							)
    if not key.startswith('''classifier'''							):
      _A          =      '''convnextv2.''' + key
    _A          =      val

  # load HuggingFace model
  _A          =      ConvNextVaForImageClassification(__snake_case							)
  model.load_state_dict(__snake_case							)
  model.eval()

  # Check outputs on an image, prepared by ConvNextImageProcessor
  _A          =      convert_preprocessor(__snake_case							)
  _A          =      preprocessor(images=prepare_img()  ,			return_tensors='''pt'''							)
  _A          =      model(**__snake_case							).logits

  # note: the logits below were obtained without center cropping
  if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
    _A          =      torch.tensor([-0.39_30, 0.17_47, -0.52_46, 0.41_77, 0.42_95]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
    _A          =      torch.tensor([-0.17_27, -0.53_41, -0.78_18, -0.47_45, -0.65_66]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
    _A          =      torch.tensor([-0.03_33, 0.15_63, -0.91_37, 0.10_54, 0.03_81]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
    _A          =      torch.tensor([-0.17_44, -0.15_55, -0.07_13, 0.09_50, -0.14_31]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
    _A          =      torch.tensor([0.99_96, 0.19_66, -0.43_86, -0.34_72, 0.66_61]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
    _A          =      torch.tensor([-0.25_53, -0.67_08, -0.13_59, 0.25_18, -0.24_88]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
    _A          =      torch.tensor([-0.06_73, -0.56_27, -0.37_53, -0.27_22, 0.01_78]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
    _A          =      torch.tensor([-0.63_77, -0.74_58, -0.21_50, 0.11_84, -0.05_97]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
    _A          =      torch.tensor([1.07_99, 0.23_22, -0.88_60, 1.02_19, 0.62_31]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
    _A          =      torch.tensor([0.37_66, 0.49_17, -1.14_26, 0.99_42, 0.60_24]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
    _A          =      torch.tensor([0.42_20, -0.69_19, -0.43_17, -0.28_81, -0.66_09]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
    _A          =      torch.tensor([0.10_82, -0.82_86, -0.50_95, 0.46_81, -0.80_85]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
    _A          =      torch.tensor([-0.24_19, -0.62_21, 0.21_76, -0.09_80, -0.75_27]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
    _A          =      torch.tensor([0.03_91, -0.43_71, 0.37_86, 0.12_51, -0.27_84]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
    _A          =      torch.tensor([-0.05_04, 0.56_36, -0.17_29, -0.65_07, -0.39_49]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
    _A          =      torch.tensor([0.35_60, 0.94_86, 0.31_49, -0.26_67, -0.51_38]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
    _A          =      torch.tensor([-0.24_69, -0.45_50, -0.58_53, -0.08_10, 0.03_09]							)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
    _A          =      torch.tensor([-0.30_90, 0.08_02, -0.06_82, -0.19_79, -0.28_26]							)
  else:
    raise ValueError(F'''Unknown URL: {checkpoint_url}'''							)

  assert torch.allclose(logits[0, :5]  ,			__snake_case  ,			atol=1e-3							)
  assert logits.shape == expected_shape
  print('''Model outputs match the original results!'''							)

  if save_model:
    print('''Saving model to local...'''							)
    # Create folder to save model
    if not os.path.isdir(__snake_case							):
      os.mkdir(__snake_case							)

    model.save_pretrained(__snake_case							)
    preprocessor.save_pretrained(__snake_case							)

  _A          =      '''convnextv2'''
  if "atto" in checkpoint_url:
    model_name += "-atto"
  if "femto" in checkpoint_url:
    model_name += "-femto"
  if "pico" in checkpoint_url:
    model_name += "-pico"
  if "nano" in checkpoint_url:
    model_name += "-nano"
  elif "tiny" in checkpoint_url:
    model_name += "-tiny"
  elif "base" in checkpoint_url:
    model_name += "-base"
  elif "large" in checkpoint_url:
    model_name += "-large"
  elif "huge" in checkpoint_url:
    model_name += "-huge"
  if "22k" in checkpoint_url and "1k" not in checkpoint_url:
    model_name += "-22k"
  elif "22k" in checkpoint_url and "1k" in checkpoint_url:
    model_name += "-22k-1k"
  elif "1k" in checkpoint_url:
    model_name += "-1k"
  if "224" in checkpoint_url:
    model_name += "-224"
  elif "384" in checkpoint_url:
    model_name += "-384"
  elif "512" in checkpoint_url:
    model_name += "-512"

  if push_to_hub:
    print(F'''Pushing {model_name} to the hub...'''							)
    model.push_to_hub(__snake_case							)
    preprocessor.push_to_hub(__snake_case							)


if __name__ == "__main__":
    __lowerCamelCase								=   argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        """--checkpoint_url""",
        default="""https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt""",
        type=str,
        help="""URL of the original ConvNeXTV2 checkpoint you'd like to convert.""",
    )
    parser.add_argument(
        """--pytorch_dump_folder_path""",
        default="""model""",
        type=str,
        help="""Path to the output PyTorch model directory.""",
    )
    parser.add_argument("""--save_model""", action="""store_true""", help="""Save model to local""")
    parser.add_argument("""--push_to_hub""", action="""store_true""", help="""Push model and image preprocessor to the hub""")

    __lowerCamelCase								=   parser.parse_args()
    convert_convnextva_checkpoint(
        args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
    )




from __future__ import annotations




def       UpperCAmelCase__	(							__snake_case							)				->	list[int]:
  _A          =      [0] * len(__snake_case							)
  for i in range(1  ,			len(__snake_case							)							):
    _A          =      duration_times[i - 1] + waiting_times[i - 1]
  return waiting_times




def       UpperCAmelCase__	(							__snake_case  ,			__snake_case							)				->	list[int]:
  return [
      duration_time + waiting_times[i]
      for i, duration_time in enumerate(__snake_case							)
  ]




def       UpperCAmelCase__	(							__snake_case							)				->	float:
  return sum(__snake_case							) / len(__snake_case							)




def       UpperCAmelCase__	(							__snake_case							)				->	float:
  return sum(__snake_case							) / len(__snake_case							)


if __name__ == "__main__":
    # process id's
    __lowerCamelCase								=   [1, 2, 3]

    # ensure that we actually have processes
    if len(processes) == 0:
        print("""Zero amount of processes""")
        raise SystemExit(0)

    # duration time of all processes
    __lowerCamelCase								=   [1_9, 8, 9]

    # ensure we can match each id to a duration time
    if len(duration_times) != len(processes):
        print("""Unable to match all id's with their duration time""")
        raise SystemExit(0)

    # get the waiting times and the turnaround times
    __lowerCamelCase								=   calculate_waiting_times(duration_times)
    __lowerCamelCase								=   calculate_turnaround_times(duration_times, waiting_times)

    # get the average times
    __lowerCamelCase								=   calculate_average_waiting_time(waiting_times)
    __lowerCamelCase								=   calculate_average_turnaround_time(turnaround_times)

    # print all the results
    print("""Process ID\tDuration Time\tWaiting Time\tTurnaround Time""")
    for i, process in enumerate(processes):
        print(
            f'''{process}\t\t{duration_times[i]}\t\t{waiting_times[i]}\t\t'''
            f'''{turnaround_times[i]}'''
        )
    print(f'''Average waiting time = {average_waiting_time}''')
    print(f'''Average turn around time = {average_turnaround_time}''')