






'''simple docstring'''





import argparse
import json
from pathlib import Path

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

from transformers import ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
__A				:				Tuple					=		logging.get_logger(__name__)






def  lowerCamelCase_  (     lowercase__):
 lowerCamelCase__	=  ConvNextConfig()

 if "tiny" in checkpoint_url:
  lowerCamelCase__	=  [3, 3, 9, 3]
  lowerCamelCase__	=  [96, 192, 384, 768]
 if "small" in checkpoint_url:
  lowerCamelCase__	=  [3, 3, 27, 3]
  lowerCamelCase__	=  [96, 192, 384, 768]
 if "base" in checkpoint_url:
  lowerCamelCase__	=  [3, 3, 27, 3]
  lowerCamelCase__	=  [128, 256, 512, 1024]
 if "large" in checkpoint_url:
  lowerCamelCase__	=  [3, 3, 27, 3]
  lowerCamelCase__	=  [192, 384, 768, 1536]
 if "xlarge" in checkpoint_url:
  lowerCamelCase__	=  [3, 3, 27, 3]
  lowerCamelCase__	=  [256, 512, 1024, 2048]

 if "1k" in checkpoint_url:
  lowerCamelCase__	=  1000
  lowerCamelCase__	=  '''imagenet-1k-id2label.json'''
  lowerCamelCase__	=  (1, 1000)
 else:
  lowerCamelCase__	=  21841
  lowerCamelCase__	=  '''imagenet-22k-id2label.json'''
  lowerCamelCase__	=  (1, 21841)

 lowerCamelCase__	=  '''huggingface/label-files'''
 lowerCamelCase__	=  num_labels
 lowerCamelCase__	=  json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  repo_type="dataset")					,  "r"))
 lowerCamelCase__	=  {int(SCREAMING_SNAKE_CASE_): v for k, v in idalabel.items()}
 if "1k" not in checkpoint_url:
  # this dataset contains 21843 labels but the model only has 21841
  # we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
  del idalabel[9205]
  del idalabel[15027]
 lowerCamelCase__	=  idalabel
 lowerCamelCase__	=  {v: k for k, v in idalabel.items()}
 lowerCamelCase__	=  hidden_sizes
 lowerCamelCase__	=  depths

 return config, expected_shape






def  lowerCamelCase_  (     lowercase__):





 if "downsample_layers.0.0" in name:
  lowerCamelCase__	=  name.replace("downsample_layers.0.0"					,  "embeddings.patch_embeddings")
 if "downsample_layers.0.1" in name:
  lowerCamelCase__	=  name.replace("downsample_layers.0.1"					,  "embeddings.norm")  # we rename to layernorm later on
 if "downsample_layers.1.0" in name:
  lowerCamelCase__	=  name.replace("downsample_layers.1.0"					,  "stages.1.downsampling_layer.0")
 if "downsample_layers.1.1" in name:
  lowerCamelCase__	=  name.replace("downsample_layers.1.1"					,  "stages.1.downsampling_layer.1")
 if "downsample_layers.2.0" in name:
  lowerCamelCase__	=  name.replace("downsample_layers.2.0"					,  "stages.2.downsampling_layer.0")
 if "downsample_layers.2.1" in name:
  lowerCamelCase__	=  name.replace("downsample_layers.2.1"					,  "stages.2.downsampling_layer.1")
 if "downsample_layers.3.0" in name:
  lowerCamelCase__	=  name.replace("downsample_layers.3.0"					,  "stages.3.downsampling_layer.0")
 if "downsample_layers.3.1" in name:
  lowerCamelCase__	=  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.
  lowerCamelCase__	=  name[: len("stages.0")] + '''.layers''' + name[len("stages.0") :]
 if "stages" in name:
  lowerCamelCase__	=  name.replace("stages"					,  "encoder.stages")
 if "norm" in name:
  lowerCamelCase__	=  name.replace("norm"					,  "layernorm")
 if "gamma" in name:
  lowerCamelCase__	=  name.replace("gamma"					,  "layer_scale_parameter")
 if "head" in name:
  lowerCamelCase__	=  name.replace("head"					,  "classifier")

 return name






def  lowerCamelCase_  (     ):
 lowerCamelCase__	=  '''http://images.cocodataset.org/val2017/000000039769.jpg'''
 lowerCamelCase__	=  Image.open(requests.get(SCREAMING_SNAKE_CASE_					,  stream=SCREAMING_SNAKE_CASE_).raw)
 return im






@torch.no_grad()
def  lowerCamelCase_  (     lowercase__					,  lowercase__):
 lowerCamelCase__	=  get_convnext_config(SCREAMING_SNAKE_CASE_)
 # load original state_dict from URL
 lowerCamelCase__	=  torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_)['''model''']
 # rename keys
 for key in state_dict.copy().keys():
  lowerCamelCase__	=  state_dict.pop(SCREAMING_SNAKE_CASE_)
  lowerCamelCase__	=  val
 # add prefix to all keys expect classifier head
 for key in state_dict.copy().keys():
  lowerCamelCase__	=  state_dict.pop(SCREAMING_SNAKE_CASE_)
  if not key.startswith("classifier"):
   lowerCamelCase__	=  '''convnext.''' + key
  lowerCamelCase__	=  val

 # load HuggingFace model
 lowerCamelCase__	=  ConvNextForImageClassification(SCREAMING_SNAKE_CASE_)
 model.load_state_dict(SCREAMING_SNAKE_CASE_)
 model.eval()

 # Check outputs on an image, prepared by ConvNextImageProcessor
 lowerCamelCase__	=  224 if '''224''' in checkpoint_url else 384
 lowerCamelCase__	=  ConvNextImageProcessor(size=SCREAMING_SNAKE_CASE_)
 lowerCamelCase__	=  image_processor(images=prepare_img()					,  return_tensors="pt").pixel_values

 lowerCamelCase__	=  model(SCREAMING_SNAKE_CASE_).logits

 # note: the logits below were obtained without center cropping
 if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
  lowerCamelCase__	=  torch.tensor([-0.1_210, -0.6_605, 0.1_918])
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
  lowerCamelCase__	=  torch.tensor([-0.4_473, -0.1_847, -0.6_365])
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
  lowerCamelCase__	=  torch.tensor([0.4_525, 0.7_539, 0.0_308])
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
  lowerCamelCase__	=  torch.tensor([0.3_561, 0.6_350, -0.0_384])
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
  lowerCamelCase__	=  torch.tensor([0.4_174, -0.0_989, 0.1_489])
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
  lowerCamelCase__	=  torch.tensor([0.2_513, -0.1_349, -0.1_613])
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
  lowerCamelCase__	=  torch.tensor([1.2_980, 0.3_631, -0.1_198])
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
  lowerCamelCase__	=  torch.tensor([1.2_963, 0.1_227, 0.1_723])
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
  lowerCamelCase__	=  torch.tensor([1.7_956, 0.8_390, 0.2_820])
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
  lowerCamelCase__	=  torch.tensor([-0.2_822, -0.0_502, -0.0_878])
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
  lowerCamelCase__	=  torch.tensor([-0.5_672, -0.0_730, -0.4_348])
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
  lowerCamelCase__	=  torch.tensor([0.2_681, 0.2_365, 0.6_246])
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
  lowerCamelCase__	=  torch.tensor([-0.2_642, 0.3_931, 0.5_116])
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
  lowerCamelCase__	=  torch.tensor([-0.6_677, -0.1_873, -0.8_379])
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
  lowerCamelCase__	=  torch.tensor([-0.7_749, -0.2_967, -0.6_444])
 else:
  raise ValueError(F'''Unknown URL: {checkpoint_url}''')

 assert torch.allclose(logits[0, :3]					,  SCREAMING_SNAKE_CASE_					,  atol=1E-3)
 assert logits.shape == expected_shape

 Path(SCREAMING_SNAKE_CASE_).mkdir(exist_ok=SCREAMING_SNAKE_CASE_)
 print(F'''Saving model to {pytorch_dump_folder_path}''')
 model.save_pretrained(SCREAMING_SNAKE_CASE_)
 print(F'''Saving image processor to {pytorch_dump_folder_path}''')
 image_processor.save_pretrained(SCREAMING_SNAKE_CASE_)

 print("Pushing model to the hub...")
 lowerCamelCase__	=  '''convnext'''
 if "tiny" in checkpoint_url:
  model_name += "-tiny"
 elif "small" in checkpoint_url:
  model_name += "-small"
 elif "base" in checkpoint_url:
  model_name += "-base"
 elif "xlarge" in checkpoint_url:
  model_name += "-xlarge"
 elif "large" in checkpoint_url:
  model_name += "-large"
 if "224" in checkpoint_url:
  model_name += "-224"
 elif "384" in checkpoint_url:
  model_name += "-384"
 if "22k" in checkpoint_url and "1k" not in checkpoint_url:
  model_name += "-22k"
 if "22k" in checkpoint_url and "1k" in checkpoint_url:
  model_name += "-22k-1k"

 model.push_to_hub(
     repo_path_or_name=Path(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)					,  organization="nielsr"					,  commit_message="Add model"					,  )


if __name__ == "__main__":
 __A				:				Dict					=		argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     """--checkpoint_url""",
     default="""https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth""",
     type=str,
     help="""URL of the original ConvNeXT checkpoint you'd like to convert.""",
 )
 parser.add_argument(
     """--pytorch_dump_folder_path""",
     default=None,
     type=str,
     required=True,
     help="""Path to the output PyTorch model directory.""",
 )

 __A				:				Optional[Any]					=		parser.parse_args()
 convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)





import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()



class  UpperCAmelCase_						(				unittest.TestCase):
       '''simple docstring'''
       __UpperCamelCase   :  Tuple				=       "hf-internal-testing/tiny-stable-diffusion-pipe"
       __UpperCamelCase   :  Optional[Any]				=       1
       __UpperCamelCase   :  str				=       77
       __UpperCamelCase   :  Tuple				=       32
       __UpperCamelCase   :  List[str]				=       4
       __UpperCamelCase   :  Optional[int]				=       64
       __UpperCamelCase   :  Union[str, Any]				=       torch.manual_seed(0)






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE=0.9_999    ):




              """simple docstring"""




              UpperCamelCase     :						Optional[int] 						=		UNetaDConditionModel.from_pretrained(self.model_id			,       subfolder='''unet'''    )
              UpperCamelCase     :						Any 						=		unet.to(__SCREAMING_SNAKE_CASE    )
              UpperCamelCase     :						Dict 						=		EMAModel(unet.parameters()			,       decay=__SCREAMING_SNAKE_CASE			,       model_cls=__SCREAMING_SNAKE_CASE			,       model_config=unet.config    )
              return unet, ema_unet






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						Tuple 						=		torch.randn(
                  self.batch_size			,       self.num_in_channels			,       self.latent_height			,       self.latent_width			,       generator=self.generator    ).to(__SCREAMING_SNAKE_CASE    )
              UpperCamelCase     :						Optional[int] 						=		torch.randint(0			,       1_000			,       size=(self.batch_size,)			,       generator=self.generator    ).to(__SCREAMING_SNAKE_CASE    )
              UpperCamelCase     :						List[Any] 						=		torch.randn(
                  self.batch_size			,       self.prompt_length			,       self.text_encoder_hidden_dim			,       generator=self.generator    ).to(__SCREAMING_SNAKE_CASE    )
              return noisy_latents, timesteps, encoder_hidden_states






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE    ):




              """simple docstring"""




              UpperCamelCase     :						Dict 						=		{}
              for k, param in unet.state_dict().items():
                     UpperCamelCase     :						int 						=		torch.randn_like(__SCREAMING_SNAKE_CASE    ) + (param * torch.randn_like(__SCREAMING_SNAKE_CASE    ))
                     updated_state_dict.update({k: updated_param}    )
              unet.load_state_dict(__SCREAMING_SNAKE_CASE    )
              return unet






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						Optional[Any] 						=		self.get_models()
              # Take the first (hypothetical) EMA step.
              ema_unet.step(unet.parameters()    )
              assert ema_unet.optimization_step == 1

              # Take two more.
              for _ in range(2    ):
                     ema_unet.step(unet.parameters()    )
              assert ema_unet.optimization_step == 3






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						List[Any] 						=		self.get_models()
              # Since the `unet` is not being updated (i.e., backprop'd)
              # there won't be any difference between the `params` of `unet`
              # and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
              ema_unet.step(unet.parameters()    )
              UpperCamelCase     :						int 						=		list(unet.parameters()    )
              for s_param, param in zip(ema_unet.shadow_params			,       __SCREAMING_SNAKE_CASE    ):
                     assert torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )

              # The above holds true even if we call `ema.step()` multiple times since
              # `unet` params are still not being updated.
              for _ in range(4    ):
                     ema_unet.step(unet.parameters()    )
              for s_param, param in zip(ema_unet.shadow_params			,       __SCREAMING_SNAKE_CASE    ):
                     assert torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						Any 						=		self.get_models()
              # Here we simulate the parameter updates for `unet`. Since there might
              # be some parameters which are initialized to zero we take extra care to
              # initialize their values to something non-zero before the multiplication.
              UpperCamelCase     :						Union[str, Any] 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )

              # Take the EMA step.
              ema_unet.step(unet_pseudo_updated_step_one.parameters()    )

              # Now the EMA'd parameters won't be equal to the original model parameters.
              UpperCamelCase     :						Any 						=		list(unet_pseudo_updated_step_one.parameters()    )
              for s_param, param in zip(ema_unet.shadow_params			,       __SCREAMING_SNAKE_CASE    ):
                     assert ~torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )

              # Ensure this is the case when we take multiple EMA steps.
              for _ in range(4    ):
                     ema_unet.step(unet.parameters()    )
              for s_param, param in zip(ema_unet.shadow_params			,       __SCREAMING_SNAKE_CASE    ):
                     assert ~torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						Optional[int] 						=		self.get_models()

              # First backprop + EMA
              UpperCamelCase     :						Dict 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )
              ema_unet.step(unet_step_one.parameters()    )
              UpperCamelCase     :						List[str] 						=		ema_unet.shadow_params

              # Second backprop + EMA
              UpperCamelCase     :						str 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )
              ema_unet.step(unet_step_two.parameters()    )
              UpperCamelCase     :						List[Any] 						=		ema_unet.shadow_params

              for step_one, step_two in zip(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ):
                     assert ~torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						int 						=		self.get_models(decay=0.0    )
              UpperCamelCase     :						Optional[int] 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )
              ema_unet.step(unet_step_one.parameters()    )
              UpperCamelCase     :						Union[str, Any] 						=		ema_unet.shadow_params

              UpperCamelCase     :						Any 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )
              ema_unet.step(unet_step_two.parameters()    )
              UpperCamelCase     :						int 						=		ema_unet.shadow_params

              for step_one, step_two in zip(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ):
                     assert torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )






       @skip_mps
       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						str 						=		self.get_models()
              UpperCamelCase						,			UpperCamelCase						,			UpperCamelCase     :						Union[str, Any] 						=		self.get_dummy_inputs()

              with tempfile.TemporaryDirectory() as tmpdir:
                     ema_unet.save_pretrained(__SCREAMING_SNAKE_CASE    )
                     UpperCamelCase     :						Optional[int] 						=		UNetaDConditionModel.from_pretrained(__SCREAMING_SNAKE_CASE			,       model_cls=__SCREAMING_SNAKE_CASE    )
                     UpperCamelCase     :						List[str] 						=		loaded_unet.to(unet.device    )

              # Since no EMA step has been performed the outputs should match.
              UpperCamelCase     :						Optional[Any] 						=		unet(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ).sample
              UpperCamelCase     :						List[str] 						=		loaded_unet(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ).sample

              assert torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       atol=1e-4    )

