
"""simple docstring"""




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()
snake_case					  =						logging.get_logger(__name__)

def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_     ):
					SCREAMING_SNAKE_CASE           =      ConvNextVaConfig()

					if "atto" in checkpoint_url:
										SCREAMING_SNAKE_CASE           =      [2, 2, 6, 2]
										SCREAMING_SNAKE_CASE           =      [4_0, 8_0, 1_6_0, 3_2_0]
					if "femto" in checkpoint_url:
										SCREAMING_SNAKE_CASE           =      [2, 2, 6, 2]
										SCREAMING_SNAKE_CASE           =      [4_8, 9_6, 1_9_2, 3_8_4]
					if "pico" in checkpoint_url:
										SCREAMING_SNAKE_CASE           =      [2, 2, 6, 2]
										SCREAMING_SNAKE_CASE           =      [6_4, 1_2_8, 2_5_6, 5_1_2]
					if "nano" in checkpoint_url:
										SCREAMING_SNAKE_CASE           =      [2, 2, 8, 2]
										SCREAMING_SNAKE_CASE           =      [8_0, 1_6_0, 3_2_0, 6_4_0]
					if "tiny" in checkpoint_url:
										SCREAMING_SNAKE_CASE           =      [3, 3, 9, 3]
										SCREAMING_SNAKE_CASE           =      [9_6, 1_9_2, 3_8_4, 7_6_8]
					if "base" in checkpoint_url:
										SCREAMING_SNAKE_CASE           =      [3, 3, 2_7, 3]
										SCREAMING_SNAKE_CASE           =      [1_2_8, 2_5_6, 5_1_2, 1_0_2_4]
					if "large" in checkpoint_url:
										SCREAMING_SNAKE_CASE           =      [3, 3, 2_7, 3]
										SCREAMING_SNAKE_CASE           =      [1_9_2, 3_8_4, 7_6_8, 1_5_3_6]
					if "huge" in checkpoint_url:
										SCREAMING_SNAKE_CASE           =      [3, 3, 2_7, 3]
										SCREAMING_SNAKE_CASE           =      [3_5_2, 7_0_4, 1_4_0_8, 2_8_1_6]

					SCREAMING_SNAKE_CASE           =      1_0_0_0
					SCREAMING_SNAKE_CASE           =      'imagenet-1k-id2label.json'
					SCREAMING_SNAKE_CASE           =      (1, 1_0_0_0)

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

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

					return config, expected_shape

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

					return name

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

def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_     ):
					if "224" in checkpoint_url:
										SCREAMING_SNAKE_CASE           =      2_2_4
										SCREAMING_SNAKE_CASE           =      2_2_4 / 2_5_6
					elif "384" in checkpoint_url:
										SCREAMING_SNAKE_CASE           =      3_8_4
										SCREAMING_SNAKE_CASE           =      None
					else:
										SCREAMING_SNAKE_CASE           =      5_1_2
										SCREAMING_SNAKE_CASE           =      None

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

@torch.no_grad()
def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_     ):
					print('Downloading original model from checkpoint...'     )
					# define ConvNeXTV2 configuration based on URL
					SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE           =      get_convnextva_config(SCREAMING_SNAKE_CASE_     )
					# load original state_dict from URL
					SCREAMING_SNAKE_CASE           =      torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_     )['model']

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

					# load HuggingFace model
					SCREAMING_SNAKE_CASE           =      ConvNextVaForImageClassification(SCREAMING_SNAKE_CASE_     )
					model.load_state_dict(SCREAMING_SNAKE_CASE_     )
					model.eval()

					# Check outputs on an image, prepared by ConvNextImageProcessor
					SCREAMING_SNAKE_CASE           =      convert_preprocessor(SCREAMING_SNAKE_CASE_     )
					SCREAMING_SNAKE_CASE           =      preprocessor(images=prepare_img(),	return_tensors='pt'     )
					SCREAMING_SNAKE_CASE           =      model(**SCREAMING_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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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":
										SCREAMING_SNAKE_CASE           =      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],	SCREAMING_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(SCREAMING_SNAKE_CASE_     ):
															os.mkdir(SCREAMING_SNAKE_CASE_     )

										model.save_pretrained(SCREAMING_SNAKE_CASE_     )
										preprocessor.save_pretrained(SCREAMING_SNAKE_CASE_     )

					SCREAMING_SNAKE_CASE           =      '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(SCREAMING_SNAKE_CASE_     )
										preprocessor.push_to_hub(SCREAMING_SNAKE_CASE_     )


if __name__ == "__main__":
	snake_case					  =						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')

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


'''simple docstring'''





import gc
import unittest

import numpy as np
import torch

from diffusers import RePaintPipeline, RePaintScheduler, UNetaDModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_image,
    load_numpy,
    nightly,
    require_torch_gpu,
    skip_mps,
    torch_device,
)

from ..pipeline_params import IMAGE_INPAINTING_BATCH_PARAMS, IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class 			UpperCamelCase__	(     SCREAMING_SNAKE_CASE   ,  unittest.TestCase):
 UpperCAmelCase__     :     Union[str, Any]    					=       RePaintPipeline
 UpperCAmelCase__     :     Any    					=       IMAGE_INPAINTING_PARAMS - {'width', 'height', 'guidance_scale'}
 UpperCAmelCase__     :     Optional[int]    					=       PipelineTesterMixin.required_optional_params - {
     'latents',
     'num_images_per_prompt',
     'callback',
     'callback_steps',
 }
 UpperCAmelCase__     :     Optional[int]    					=       IMAGE_INPAINTING_BATCH_PARAMS
 UpperCAmelCase__     :     Any    					=       False






 def 						lowercase_    (      self	:Any      )					->      int:







       '''simple docstring'''



       torch.manual_seed(0      )
       torch.manual_seed(0      )
       __A													=      UNetaDModel(
           block_out_channels=(32, 64)						,					layers_per_block=2						,					sample_size=32						,					in_channels=3						,					out_channels=3						,					down_block_types=('DownBlock2D', 'AttnDownBlock2D')						,					up_block_types=('AttnUpBlock2D', 'UpBlock2D')						,					)
       __A													=      RePaintScheduler()
       __A													=      {'unet': unet, 'scheduler': scheduler}
       return components






 def 						lowercase_    (      self	:Any						,					_A	:List[Any]						,					_A	:Tuple=0      )					->      List[str]:







       '''simple docstring'''



       if str(_A      ).startswith('mps'      ):
             __A													=      torch.manual_seed(_A      )
       else:
             __A													=      torch.Generator(device=_A      ).manual_seed(_A      )
       __A													=      np.random.RandomState(_A      ).standard_normal((1, 3, 32, 32)      )
       __A													=      torch.from_numpy(_A      ).to(device=_A						,					dtype=torch.floataa      )
       __A													=      (image > 0).to(device=_A						,					dtype=torch.floataa      )
       __A													=      {
           'image': image,
           'mask_image': mask,
           'generator': generator,
           'num_inference_steps': 5,
           'eta': 0.0,
           'jump_length': 2,
           'jump_n_sample': 2,
           'output_type': 'numpy',
       }
       return inputs






 def 						lowercase_    (      self	:Union[str, Any]      )					->      Optional[int]:







       '''simple docstring'''



       __A													=      'cpu'  # ensure determinism for the device-dependent torch.Generator
       __A													=      self.get_dummy_components()
       __A													=      RePaintPipeline(**_A      )
       __A													=      sd_pipe.to(_A      )
       sd_pipe.set_progress_bar_config(disable=_A      )

       __A													=      self.get_dummy_inputs(_A      )
       __A													=      sd_pipe(**_A      ).images
       __A													=      image[0, -3:, -3:, -1]

       assert image.shape == (1, 32, 32, 3)
       __A													=      np.array([1.0_000, 0.5_426, 0.5_497, 0.2_200, 1.0_000, 1.0_000, 0.5_623, 1.0_000, 0.6_274]      )

       assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-3






 @skip_mps
 def 						lowercase_    (      self	:Dict      )					->      Optional[int]:







       '''simple docstring'''



       return super().test_save_load_local()






 @unittest.skip('non-deterministic pipeline'      )
 def 						lowercase_    (      self	:Union[str, Any]      )					->      int:







       '''simple docstring'''



       return super().test_inference_batch_single_identical()






 @skip_mps
 def 						lowercase_    (      self	:Dict      )					->      Union[str, Any]:







       '''simple docstring'''



       return super().test_dict_tuple_outputs_equivalent()






 @skip_mps
 def 						lowercase_    (      self	:List[str]      )					->      int:







       '''simple docstring'''



       return super().test_save_load_optional_components()






 @skip_mps
 def 						lowercase_    (      self	:Optional[int]      )					->      Optional[Any]:







       '''simple docstring'''



       return super().test_attention_slicing_forward_pass()


@nightly
@require_torch_gpu
class 			UpperCamelCase__	(     unittest.TestCase):






 def 						lowercase_    (      self	:Optional[int]      )					->      int:







       '''simple docstring'''



       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






 def 						lowercase_    (      self	:Dict      )					->      Optional[Any]:







       '''simple docstring'''



       __A													=      load_image(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/'
           'repaint/celeba_hq_256.png'      )
       __A													=      load_image(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/mask_256.png'      )
       __A													=      load_numpy(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/'
           'repaint/celeba_hq_256_result.npy'      )

       __A													=      'google/ddpm-ema-celebahq-256'
       __A													=      UNetaDModel.from_pretrained(_A      )
       __A													=      RePaintScheduler.from_pretrained(_A      )

       __A													=      RePaintPipeline(unet=_A						,					scheduler=_A      ).to(_A      )
       repaint.set_progress_bar_config(disable=_A      )
       repaint.enable_attention_slicing()

       __A													=      torch.manual_seed(0      )
       __A													=      repaint(
           _A						,					_A						,					num_inference_steps=250						,					eta=0.0						,					jump_length=10						,					jump_n_sample=10						,					generator=_A						,					output_type='np'						,					)
       __A													=      output.images[0]

       assert image.shape == (256, 256, 3)
       assert np.abs(expected_image - image      ).mean() < 1E-2






