




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()
_UpperCAmelCase						:    Union[str, Any]      =  logging.get_logger(__name__)



def    UpperCAmelCase__	(  lowerCamelCase ):
  lowercase     :Tuple   						=    ConvNextVaConfig()

  if "atto" in checkpoint_url:
    lowercase     :Optional[int]   						=    [2, 2, 6, 2]
    lowercase     :Union[str, Any]   						=    [40, 80, 160, 320]
  if "femto" in checkpoint_url:
    lowercase     :Dict   						=    [2, 2, 6, 2]
    lowercase     :Any   						=    [48, 96, 192, 384]
  if "pico" in checkpoint_url:
    lowercase     :int   						=    [2, 2, 6, 2]
    lowercase     :str   						=    [64, 128, 256, 512]
  if "nano" in checkpoint_url:
    lowercase     :Dict   						=    [2, 2, 8, 2]
    lowercase     :Any   						=    [80, 160, 320, 640]
  if "tiny" in checkpoint_url:
    lowercase     :Optional[int]   						=    [3, 3, 9, 3]
    lowercase     :int   						=    [96, 192, 384, 768]
  if "base" in checkpoint_url:
    lowercase     :Any   						=    [3, 3, 27, 3]
    lowercase     :List[Any]   						=    [128, 256, 512, 1024]
  if "large" in checkpoint_url:
    lowercase     :str   						=    [3, 3, 27, 3]
    lowercase     :str   						=    [192, 384, 768, 1536]
  if "huge" in checkpoint_url:
    lowercase     :Union[str, Any]   						=    [3, 3, 27, 3]
    lowercase     :Union[str, Any]   						=    [352, 704, 1408, 2816]

  lowercase     :Union[str, Any]   						=    1000
  lowercase     :Union[str, Any]   						=    "imagenet-1k-id2label.json"
  lowercase     :str   						=    (1, 1000)

  lowercase     :Union[str, Any]   						=    "huggingface/label-files"
  lowercase     :Optional[int]   						=    num_labels
  lowercase     :Optional[Any]   						=    json.load(open(hf_hub_download(lowerCamelCase,      lowerCamelCase,      repo_type="dataset" ),      "r" ) )
  lowercase     :int   						=    {int(lowerCamelCase ): v for k, v in idalabel.items()}

  lowercase     :Dict   						=    idalabel
  lowercase     :int   						=    {v: k for k, v in idalabel.items()}
  lowercase     :List[Any]   						=    hidden_sizes
  lowercase     :Optional[Any]   						=    depths

  return config, expected_shape



def    UpperCAmelCase__	(  lowerCamelCase ):
  if "downsample_layers.0.0" in name:
    lowercase     :Optional[Any]   						=    name.replace("downsample_layers.0.0",      "embeddings.patch_embeddings" )
  if "downsample_layers.0.1" in name:
    lowercase     :Dict   						=    name.replace("downsample_layers.0.1",      "embeddings.norm" )  # we rename to layernorm later on
  if "downsample_layers.1.0" in name:
    lowercase     :Optional[Any]   						=    name.replace("downsample_layers.1.0",      "stages.1.downsampling_layer.0" )
  if "downsample_layers.1.1" in name:
    lowercase     :int   						=    name.replace("downsample_layers.1.1",      "stages.1.downsampling_layer.1" )
  if "downsample_layers.2.0" in name:
    lowercase     :Union[str, Any]   						=    name.replace("downsample_layers.2.0",      "stages.2.downsampling_layer.0" )
  if "downsample_layers.2.1" in name:
    lowercase     :str   						=    name.replace("downsample_layers.2.1",      "stages.2.downsampling_layer.1" )
  if "downsample_layers.3.0" in name:
    lowercase     :Optional[int]   						=    name.replace("downsample_layers.3.0",      "stages.3.downsampling_layer.0" )
  if "downsample_layers.3.1" in name:
    lowercase     :Dict   						=    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.
    lowercase     :Dict   						=    name[: len("stages.0" )] + ".layers" + name[len("stages.0" ) :]
  if "gamma" in name:
    lowercase     :Optional[int]   						=    name.replace("gamma",      "weight" )
  if "beta" in name:
    lowercase     :str   						=    name.replace("beta",      "bias" )
  if "stages" in name:
    lowercase     :Tuple   						=    name.replace("stages",      "encoder.stages" )
  if "norm" in name:
    lowercase     :Tuple   						=    name.replace("norm",      "layernorm" )
  if "head" in name:
    lowercase     :int   						=    name.replace("head",      "classifier" )

  return name



def    UpperCAmelCase__	(  ):
  lowercase     :List[str]   						=    "http://images.cocodataset.org/val2017/000000039769.jpg"
  lowercase     :Optional[Any]   						=    Image.open(requests.get(lowerCamelCase,      stream=lowerCamelCase ).raw )
  return im



def    UpperCAmelCase__	(  lowerCamelCase ):
  if "224" in checkpoint_url:
    lowercase     :Any   						=    224
    lowercase     :Tuple   						=    224 / 256
  elif "384" in checkpoint_url:
    lowercase     :Dict   						=    384
    lowercase     :List[Any]   						=    None
  else:
    lowercase     :Optional[int]   						=    512
    lowercase     :Optional[int]   						=    None

  return ConvNextImageProcessor(
      size=lowerCamelCase,      crop_pct=lowerCamelCase,      image_mean=[0.485, 0.456, 0.406],      image_std=[0.229, 0.224, 0.225],      resample=PILImageResampling.BICUBIC,      )



@torch.no_grad()
def    UpperCAmelCase__	(  lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase ):
  print("Downloading original model from checkpoint..." )
  # define ConvNeXTV2 configuration based on URL
  lowercase						,     lowercase     :Optional[int]   						=    get_convnextva_config(lowerCamelCase )
  # load original state_dict from URL
  lowercase     :Any   						=    torch.hub.load_state_dict_from_url(lowerCamelCase )["model"]

  print("Converting model parameters..." )
  # rename keys
  for key in state_dict.copy().keys():
    lowercase     :Optional[int]   						=    state_dict.pop(lowerCamelCase )
    lowercase     :Any   						=    val
  # add prefix to all keys expect classifier head
  for key in state_dict.copy().keys():
    lowercase     :int   						=    state_dict.pop(lowerCamelCase )
    if not key.startswith("classifier" ):
      lowercase     :Any   						=    "convnextv2." + key
    lowercase     :Any   						=    val

  # load HuggingFace model
  lowercase     :Tuple   						=    ConvNextVaForImageClassification(lowerCamelCase )
  model.load_state_dict(lowerCamelCase )
  model.eval()

  # Check outputs on an image, prepared by ConvNextImageProcessor
  lowercase     :str   						=    convert_preprocessor(lowerCamelCase )
  lowercase     :Optional[int]   						=    preprocessor(images=prepare_img(),      return_tensors="pt" )
  lowercase     :Tuple   						=    model(**lowerCamelCase ).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":
    lowercase     :Tuple   						=    torch.tensor([-0.3_930, 0.1_747, -0.5_246, 0.4_177, 0.4_295] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
    lowercase     :Tuple   						=    torch.tensor([-0.1_727, -0.5_341, -0.7_818, -0.4_745, -0.6_566] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
    lowercase     :Optional[Any]   						=    torch.tensor([-0.0_333, 0.1_563, -0.9_137, 0.1_054, 0.0_381] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
    lowercase     :Union[str, Any]   						=    torch.tensor([-0.1_744, -0.1_555, -0.0_713, 0.0_950, -0.1_431] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
    lowercase     :Union[str, Any]   						=    torch.tensor([0.9_996, 0.1_966, -0.4_386, -0.3_472, 0.6_661] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
    lowercase     :Optional[Any]   						=    torch.tensor([-0.2_553, -0.6_708, -0.1_359, 0.2_518, -0.2_488] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
    lowercase     :List[Any]   						=    torch.tensor([-0.0_673, -0.5_627, -0.3_753, -0.2_722, 0.0_178] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
    lowercase     :Union[str, Any]   						=    torch.tensor([-0.6_377, -0.7_458, -0.2_150, 0.1_184, -0.0_597] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
    lowercase     :Dict   						=    torch.tensor([1.0_799, 0.2_322, -0.8_860, 1.0_219, 0.6_231] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
    lowercase     :Optional[int]   						=    torch.tensor([0.3_766, 0.4_917, -1.1_426, 0.9_942, 0.6_024] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
    lowercase     :int   						=    torch.tensor([0.4_220, -0.6_919, -0.4_317, -0.2_881, -0.6_609] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
    lowercase     :List[str]   						=    torch.tensor([0.1_082, -0.8_286, -0.5_095, 0.4_681, -0.8_085] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
    lowercase     :str   						=    torch.tensor([-0.2_419, -0.6_221, 0.2_176, -0.0_980, -0.7_527] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
    lowercase     :List[Any]   						=    torch.tensor([0.0_391, -0.4_371, 0.3_786, 0.1_251, -0.2_784] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
    lowercase     :Union[str, Any]   						=    torch.tensor([-0.0_504, 0.5_636, -0.1_729, -0.6_507, -0.3_949] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
    lowercase     :Optional[Any]   						=    torch.tensor([0.3_560, 0.9_486, 0.3_149, -0.2_667, -0.5_138] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
    lowercase     :int   						=    torch.tensor([-0.2_469, -0.4_550, -0.5_853, -0.0_810, 0.0_309] )
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
    lowercase     :Dict   						=    torch.tensor([-0.3_090, 0.0_802, -0.0_682, -0.1_979, -0.2_826] )
  else:
    raise ValueError(F"Unknown URL: {checkpoint_url}" )

  assert torch.allclose(logits[0, :5],      lowerCamelCase,      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(lowerCamelCase ):
      os.mkdir(lowerCamelCase )

    model.save_pretrained(lowerCamelCase )
    preprocessor.save_pretrained(lowerCamelCase )

  lowercase     :Any   						=    "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(lowerCamelCase )
    preprocessor.push_to_hub(lowerCamelCase )


if __name__ == "__main__":
     _UpperCAmelCase						:    Optional[int]      =  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")

     _UpperCAmelCase						:    str      =  parser.parse_args()
     convert_convnextva_checkpoint(
         args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
     )









import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only

from utils import save_json



def    UpperCAmelCase__	(  lowerCamelCase ):
  lowercase     :Optional[Any]   						=    filter(lambda lowerCamelCase : p.requires_grad,      model.parameters() )
  lowercase     :Optional[int]   						=    sum([np.prod(p.size() ) for p in model_parameters] )
  return params


_UpperCAmelCase						:    Optional[Any]      =  logging.getLogger(__name__)




class  __lowerCAmelCase							(				pl.Callback):






       def        SCREAMING_SNAKE_CASE     ( self:       Union[str, Any]							,       _lowerCAmelCase:       List[Any]							,       _lowerCAmelCase:       Union[str, Any]					):
         lowercase     :Dict   						=    {F"lr_group_{i}": param["lr"] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups					)}
         pl_module.logger.log_metrics(_lowerCAmelCase					)






       @rank_zero_only
       def        SCREAMING_SNAKE_CASE     ( self:       Optional[Any]							,       _lowerCAmelCase:       pl.Trainer							,       _lowerCAmelCase:       pl.LightningModule							,       _lowerCAmelCase:       str							,       _lowerCAmelCase:       List[str]=True					):
         logger.info(F"***** {type_path} results at step {trainer.global_step:05d} *****"					)
         lowercase     :Tuple   						=    trainer.callback_metrics
         trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ["log", "progress_bar", "preds"]}					)
         # Log results
         lowercase     :Optional[int]   						=    Path(pl_module.hparams.output_dir					)
         if type_path == "test":
           lowercase     :List[Any]   						=    od / "test_results.txt"
           lowercase     :Optional[int]   						=    od / "test_generations.txt"
         else:
           # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
           # If people want this it will be easy enough to add back.
           lowercase     :int   						=    od / F"{type_path}_results/{trainer.global_step:05d}.txt"
           lowercase     :Tuple   						=    od / F"{type_path}_generations/{trainer.global_step:05d}.txt"
           results_file.parent.mkdir(exist_ok=_lowerCAmelCase					)
           generations_file.parent.mkdir(exist_ok=_lowerCAmelCase					)
         with open(_lowerCAmelCase							,       "a+"					) as writer:
           for key in sorted(_lowerCAmelCase					):
             if key in ["log", "progress_bar", "preds"]:
               continue
             lowercase     :int   						=    metrics[key]
             if isinstance(_lowerCAmelCase							,       torch.Tensor					):
               lowercase     :int   						=    val.item()
             lowercase     :Optional[Any]   						=    F"{key}: {val:.6f}\n"
             writer.write(_lowerCAmelCase					)

         if not save_generations:
           return

         if "preds" in metrics:
           lowercase     :str   						=    "\n".join(metrics["preds"]					)
           generations_file.open("w+"					).write(_lowerCAmelCase					)






       @rank_zero_only
       def        SCREAMING_SNAKE_CASE     ( self:       Optional[int]							,       _lowerCAmelCase:       Tuple							,       _lowerCAmelCase:       Tuple					):
         try:
           lowercase     :Optional[int]   						=    pl_module.model.model.num_parameters()
         except AttributeError:
           lowercase     :Dict   						=    pl_module.model.num_parameters()

         lowercase     :Any   						=    count_trainable_parameters(_lowerCAmelCase					)
         # mp stands for million parameters
         trainer.logger.log_metrics({"n_params": npars, "mp": npars / 1e6, "grad_mp": n_trainable_pars / 1e6}					)






       @rank_zero_only
       def        SCREAMING_SNAKE_CASE     ( self:       Union[str, Any]							,       _lowerCAmelCase:       pl.Trainer							,       _lowerCAmelCase:       pl.LightningModule					):
         save_json(pl_module.metrics							,       pl_module.metrics_save_path					)
         return self._write_logs(_lowerCAmelCase							,       _lowerCAmelCase							,       "test"					)






       @rank_zero_only
       def        SCREAMING_SNAKE_CASE     ( self:       List[str]							,       _lowerCAmelCase:       pl.Trainer							,       _lowerCAmelCase:       Dict					):
         save_json(pl_module.metrics							,       pl_module.metrics_save_path					)
         # Uncommenting this will save val generations
         # return self._write_logs(trainer, pl_module, "valid")







def    UpperCAmelCase__	(  lowerCamelCase,      lowerCamelCase,      lowerCamelCase=1,      lowerCamelCase=False ):
  if metric == "rouge2":
    lowercase     :str   						=    "{val_avg_rouge2:.4f}-{step_count}"
  elif metric == "bleu":
    lowercase     :int   						=    "{val_avg_bleu:.4f}-{step_count}"
  elif metric == "loss":
    lowercase     :List[str]   						=    "{val_avg_loss:.4f}-{step_count}"
  else:
    raise NotImplementedError(
        F"seq2seq callbacks only support rouge2, bleu and loss, got {metric}, You can make your own by adding to"
        " this function." )

  lowercase     :int   						=    ModelCheckpoint(
      dirpath=lowerCamelCase,      filename=lowerCamelCase,      monitor=F"val_{metric}",      mode="min" if "loss" in metric else "max",      save_top_k=lowerCamelCase,      )
  return checkpoint_callback







def    UpperCAmelCase__	(  lowerCamelCase,      lowerCamelCase ):
  return EarlyStopping(
      monitor=F"val_{metric}",      mode="min" if "loss" in metric else "max",      patience=lowerCamelCase,      verbose=lowerCamelCase,      )



