



from typing import Optional

import torch
from PIL import Image
from tqdm.auto import tqdm
from transformers import CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, DiffusionPipeline, UNetaDConditionModel
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils import (
    deprecate,
)





class 							lowerCAmelCase_		(		lowercase_     ):







		def __init__(    self  :   int       ,		UpperCAmelCase_  :   AutoencoderKL       ,		UpperCAmelCase_  :   CLIPTextModel       ,		UpperCAmelCase_  :   CLIPTokenizer       ,		UpperCAmelCase_  :   UNetaDConditionModel       ,		UpperCAmelCase_  :   DDIMScheduler       ,		UpperCAmelCase_  :   float = 0.93       ,		UpperCAmelCase_  :   bool = True       ,		)       ->		List[Any]:


						'''simple docstring'''







						_UpperCAmelCase						:  List[str]     					=	mixing_coeff
						_UpperCAmelCase						:  Optional[int]     					=	leapfrog_steps

						super().__init__()
						self.register_modules(
						    vae=UpperCAmelCase_       ,		text_encoder=UpperCAmelCase_       ,		tokenizer=UpperCAmelCase_       ,		unet=UpperCAmelCase_       ,		scheduler=UpperCAmelCase_       ,		)

						_UpperCAmelCase						:  Optional[int]     					=	2 ** (len(self.vae.config.block_out_channels					) - 1)
						_UpperCAmelCase						:  Any     					=	VaeImageProcessor(vae_scale_factor=self.vae_scale_factor					)







		def 					a_  (    self  :   List[Any]       ,		UpperCAmelCase_  :   str       ,		UpperCAmelCase_  :   Optional[str] = None       ,		UpperCAmelCase_  :   bool = False					)       ->		List[str]:


						'''simple docstring'''







						_UpperCAmelCase						:  Any     					=	self.tokenizer(
						    UpperCAmelCase_       ,		padding='''max_length'''       ,		max_length=self.tokenizer.model_max_length       ,		truncation=UpperCAmelCase_       ,		return_tensors='''pt'''       ,		)

						_UpperCAmelCase						:  Optional[int]     					=	self.text_encoder(text_inputs.input_ids.to(self.device					)					).last_hidden_state

						_UpperCAmelCase						:  Any     					=	prompt_embeds.to(dtype=self.text_encoder.dtype       ,		device=self.device					)

						if do_classifier_free_guidance:
										_UpperCAmelCase						:  Union[str, Any]     					=	'''''' if negative_prompt is None else negative_prompt

										_UpperCAmelCase						:  int     					=	self.tokenizer(
										    UpperCAmelCase_       ,		padding='''max_length'''       ,		max_length=self.tokenizer.model_max_length       ,		truncation=UpperCAmelCase_       ,		return_tensors='''pt'''       ,		)

										_UpperCAmelCase						:  int     					=	self.text_encoder(uncond_input.input_ids.to(self.device					)					).last_hidden_state

										_UpperCAmelCase						:  Optional[Any]     					=	torch.cat([negative_prompt_embeds, prompt_embeds]					)

						return prompt_embeds







		def 					a_  (    self  :   Union[str, Any]       ,		UpperCAmelCase_  :   torch.Tensor       ,		UpperCAmelCase_  :   torch.Tensor					)       ->		Any:


						'''simple docstring'''







						_UpperCAmelCase						:  str     					=	self.mixing_coeff * x + (1 - self.mixing_coeff) * y
						_UpperCAmelCase						:  Tuple     					=	self.mixing_coeff * y + (1 - self.mixing_coeff) * x

						return [x, y]







		def 					a_  (    self  :   Optional[Any]       ,		UpperCAmelCase_  :   torch.Tensor       ,		UpperCAmelCase_  :   torch.Tensor					)       ->		Union[str, Any]:


						'''simple docstring'''







						_UpperCAmelCase						:  int     					=	(y - (1 - self.mixing_coeff) * x) / self.mixing_coeff
						_UpperCAmelCase						:  Any     					=	(x - (1 - self.mixing_coeff) * y) / self.mixing_coeff

						return [x, y]







		def 					a_  (    self  :   List[str]       ,		UpperCAmelCase_  :   torch.Tensor					)       ->		Any:


						'''simple docstring'''







						_UpperCAmelCase						:  List[str]     					=	int(UpperCAmelCase_					)

						_UpperCAmelCase						:  Optional[int]     					=	self.scheduler.alphas_cumprod[t] if t >= 0 else self.scheduler.final_alpha_cumprod

						return alpha_prod, 1 - alpha_prod







		def 					a_  (    self  :   Dict       ,		UpperCAmelCase_  :   torch.Tensor       ,		UpperCAmelCase_  :   torch.Tensor       ,		UpperCAmelCase_  :   torch.Tensor       ,		UpperCAmelCase_  :   torch.Tensor       ,		)       ->		Dict:


						'''simple docstring'''







						_UpperCAmelCase						:  Tuple     					=	timestep - self.scheduler.config.num_train_timesteps / self.scheduler.num_inference_steps

						_UpperCAmelCase					, _UpperCAmelCase						:  List[str]     					=	self._get_alpha_and_beta(UpperCAmelCase_					)
						_UpperCAmelCase					, _UpperCAmelCase						:  Tuple     					=	self._get_alpha_and_beta(UpperCAmelCase_					)

						_UpperCAmelCase						:  Any     					=	(alpha_prod_t_prev / alpha_prod_t) ** 0.5
						_UpperCAmelCase						:  List[Any]     					=	-a_t * (beta_prod_t**0.5) + beta_prod_t_prev**0.5

						_UpperCAmelCase						:  Tuple     					=	(base - b_t * model_output) / a_t

						return model_input, next_model_input.to(base.dtype					)







		def 					a_  (    self  :   Dict       ,		UpperCAmelCase_  :   torch.Tensor       ,		UpperCAmelCase_  :   torch.Tensor       ,		UpperCAmelCase_  :   torch.Tensor       ,		UpperCAmelCase_  :   torch.Tensor       ,		)       ->		List[Any]:


						'''simple docstring'''







						_UpperCAmelCase						:  List[str]     					=	timestep - self.scheduler.config.num_train_timesteps / self.scheduler.num_inference_steps

						_UpperCAmelCase					, _UpperCAmelCase						:  Optional[Any]     					=	self._get_alpha_and_beta(UpperCAmelCase_					)
						_UpperCAmelCase					, _UpperCAmelCase						:  List[str]     					=	self._get_alpha_and_beta(UpperCAmelCase_					)

						_UpperCAmelCase						:  List[str]     					=	(alpha_prod_t_prev / alpha_prod_t) ** 0.5
						_UpperCAmelCase						:  Optional[Any]     					=	-a_t * (beta_prod_t**0.5) + beta_prod_t_prev**0.5
						_UpperCAmelCase						:  str     					=	a_t * base + b_t * model_output

						return model_input, next_model_input.to(base.dtype					)







		@torch.no_grad()
		def 					a_  (    self  :   int       ,		UpperCAmelCase_  :   torch.Tensor					)       ->		List[Any]:


						'''simple docstring'''







						_UpperCAmelCase						:  Optional[int]     					=	1 / self.vae.config.scaling_factor * latents
						_UpperCAmelCase						:  Any     					=	self.vae.decode(UpperCAmelCase_					).sample
						_UpperCAmelCase						:  Dict     					=	(image / 2 + 0.5).clamp(0       ,		1					)
						return image







		@torch.no_grad()
		def 					a_  (    self  :   Any       ,		UpperCAmelCase_  :   Image.Image       ,		UpperCAmelCase_  :   torch.Tensor       ,		UpperCAmelCase_  :   torch.Tensor       ,		UpperCAmelCase_  :   float       ,		UpperCAmelCase_  :   Optional[torch.Generator] = None       ,		)       ->		List[str]:


						'''simple docstring'''







						_UpperCAmelCase						:  Any     					=	guidance_scale > 1.0

						_UpperCAmelCase						:  Dict     					=	image.to(device=self.device       ,		dtype=text_embeds.dtype					)
						_UpperCAmelCase						:  str     					=	self.vae.encode(UpperCAmelCase_					).latent_dist.sample(UpperCAmelCase_					)

						_UpperCAmelCase						:  Optional[int]     					=	self.vae.config.scaling_factor * latent

						_UpperCAmelCase						:  Tuple     					=	[latent.clone(), latent.clone()]

						for i, t in tqdm(enumerate(UpperCAmelCase_					)       ,		total=len(UpperCAmelCase_					)					):
										_UpperCAmelCase						:  Dict     					=	self.noise_mixing_layer(x=coupled_latents[0]       ,		y=coupled_latents[1]					)

										# j - model_input index, k - base index
										for j in range(2					):
														_UpperCAmelCase						:  List[str]     					=	j ^ 1

														if self.leapfrog_steps:
																		if i % 2 == 0:
																						_UpperCAmelCase					, _UpperCAmelCase						:  List[str]     					=	j, k

														_UpperCAmelCase						:  List[Any]     					=	coupled_latents[j]
														_UpperCAmelCase						:  str     					=	coupled_latents[k]

														_UpperCAmelCase						:  List[str]     					=	torch.cat([model_input] * 2					) if do_classifier_free_guidance else model_input

														_UpperCAmelCase						:  Union[str, Any]     					=	self.unet(UpperCAmelCase_       ,		UpperCAmelCase_       ,		encoder_hidden_states=UpperCAmelCase_					).sample

														if do_classifier_free_guidance:
																		_UpperCAmelCase					, _UpperCAmelCase						:  Union[str, Any]     					=	noise_pred.chunk(2					)
																		_UpperCAmelCase						:  str     					=	noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

														_UpperCAmelCase					, _UpperCAmelCase						:  int     					=	self.noise_step(
														    base=UpperCAmelCase_       ,		model_input=UpperCAmelCase_       ,		model_output=UpperCAmelCase_       ,		timestep=UpperCAmelCase_       ,		)

														_UpperCAmelCase						:  List[Any]     					=	model_input

						return coupled_latents







		@torch.no_grad()
		def __call__(    self  :   Optional[int]       ,		UpperCAmelCase_  :   str       ,		UpperCAmelCase_  :   str       ,		UpperCAmelCase_  :   Image.Image       ,		UpperCAmelCase_  :   float = 3.0       ,		UpperCAmelCase_  :   int = 50       ,		UpperCAmelCase_  :   float = 0.8       ,		UpperCAmelCase_  :   Optional[str] = None       ,		UpperCAmelCase_  :   Optional[torch.Generator] = None       ,		UpperCAmelCase_  :   Optional[str] = "pil"       ,		)       ->		str:


						'''simple docstring'''







						_UpperCAmelCase						:  Union[str, Any]     					=	guidance_scale > 1.0

						_UpperCAmelCase						:  str     					=	self.image_processor.preprocess(UpperCAmelCase_					)

						_UpperCAmelCase						:  Optional[Any]     					=	self._encode_prompt(UpperCAmelCase_       ,		UpperCAmelCase_       ,		UpperCAmelCase_					)
						_UpperCAmelCase						:  List[str]     					=	self._encode_prompt(UpperCAmelCase_       ,		UpperCAmelCase_       ,		UpperCAmelCase_					)

						self.scheduler.set_timesteps(UpperCAmelCase_       ,		self.device					)

						_UpperCAmelCase						:  Dict     					=	num_inference_steps - int(num_inference_steps * strength					)
						_UpperCAmelCase						:  Optional[Any]     					=	self.scheduler.timesteps[t_limit:]
						_UpperCAmelCase						:  Optional[Any]     					=	fwd_timesteps.flip(0					)

						_UpperCAmelCase						:  Dict     					=	self.prepare_latents(UpperCAmelCase_       ,		UpperCAmelCase_       ,		UpperCAmelCase_       ,		UpperCAmelCase_       ,		UpperCAmelCase_					)

						for i, t in tqdm(enumerate(UpperCAmelCase_					)       ,		total=len(UpperCAmelCase_					)					):
										# j - model_input index, k - base index
										for k in range(2					):
														_UpperCAmelCase						:  Any     					=	k ^ 1

														if self.leapfrog_steps:
																		if i % 2 == 1:
																						_UpperCAmelCase					, _UpperCAmelCase						:  List[Any]     					=	j, k

														_UpperCAmelCase						:  Dict     					=	coupled_latents[j]
														_UpperCAmelCase						:  str     					=	coupled_latents[k]

														_UpperCAmelCase						:  str     					=	torch.cat([model_input] * 2					) if do_classifier_free_guidance else model_input

														_UpperCAmelCase						:  Optional[Any]     					=	self.unet(UpperCAmelCase_       ,		UpperCAmelCase_       ,		encoder_hidden_states=UpperCAmelCase_					).sample

														if do_classifier_free_guidance:
																		_UpperCAmelCase					, _UpperCAmelCase						:  Any     					=	noise_pred.chunk(2					)
																		_UpperCAmelCase						:  List[str]     					=	noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

														_UpperCAmelCase					, _UpperCAmelCase						:  str     					=	self.denoise_step(
														    base=UpperCAmelCase_       ,		model_input=UpperCAmelCase_       ,		model_output=UpperCAmelCase_       ,		timestep=UpperCAmelCase_       ,		)

														_UpperCAmelCase						:  str     					=	model_input

										_UpperCAmelCase						:  Optional[int]     					=	self.denoise_mixing_layer(x=coupled_latents[0]       ,		y=coupled_latents[1]					)

						# either one is fine
						_UpperCAmelCase						:  Union[str, Any]     					=	coupled_latents[0]

						if output_type not in ["latent", "pt", "np", "pil"]:
										_UpperCAmelCase						:  Dict     					=	(
										    F'''the output_type {output_type} is outdated. Please make sure to set it to one of these instead: '''
										    '''`pil`, `np`, `pt`, `latent`'''
										)
										deprecate('''Unsupported output_type'''       ,		'''1.0.0'''       ,		UpperCAmelCase_       ,		standard_warn=UpperCAmelCase_					)
										_UpperCAmelCase						:  str     					=	'''np'''

						if output_type == "latent":
										_UpperCAmelCase						:  List[Any]     					=	final_latent
						else:
										_UpperCAmelCase						:  int     					=	self.decode_latents(UpperCAmelCase_					)
										_UpperCAmelCase						:  str     					=	self.image_processor.postprocess(UpperCAmelCase_       ,		output_type=UpperCAmelCase_					)

						return image





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()
UpperCAmelCase__			:   Optional[Any]      			=    logging.get_logger(__name__)


def 	_A	(   _UpperCamelCase			):
				_UpperCAmelCase						:  Optional[Any]     					=	ConvNextConfig()

				if "tiny" in checkpoint_url:
								_UpperCAmelCase						:  Any     					=	[3, 3, 9, 3]
								_UpperCAmelCase						:  int     					=	[96, 192, 384, 768]
				if "small" in checkpoint_url:
								_UpperCAmelCase						:  str     					=	[3, 3, 27, 3]
								_UpperCAmelCase						:  List[Any]     					=	[96, 192, 384, 768]
				if "base" in checkpoint_url:
								_UpperCAmelCase						:  Dict     					=	[3, 3, 27, 3]
								_UpperCAmelCase						:  Tuple     					=	[128, 256, 512, 1_024]
				if "large" in checkpoint_url:
								_UpperCAmelCase						:  Optional[int]     					=	[3, 3, 27, 3]
								_UpperCAmelCase						:  Optional[int]     					=	[192, 384, 768, 1_536]
				if "xlarge" in checkpoint_url:
								_UpperCAmelCase						:  Tuple     					=	[3, 3, 27, 3]
								_UpperCAmelCase						:  int     					=	[256, 512, 1_024, 2_048]

				if "1k" in checkpoint_url:
								_UpperCAmelCase						:  int     					=	1_000
								_UpperCAmelCase						:  Union[str, Any]     					=	'''imagenet-1k-id2label.json'''
								_UpperCAmelCase						:  Any     					=	(1, 1_000)
				else:
								_UpperCAmelCase						:  Union[str, Any]     					=	21_841
								_UpperCAmelCase						:  int     					=	'''imagenet-22k-id2label.json'''
								_UpperCAmelCase						:  List[str]     					=	(1, 21_841)

				_UpperCAmelCase						:  List[str]     					=	'''huggingface/label-files'''
				_UpperCAmelCase						:  Any     					=	num_labels
				_UpperCAmelCase						:  str     					=	json.load(open(hf_hub_download(_UpperCamelCase      ,					_UpperCamelCase      ,					repo_type='''dataset'''			)      ,					'''r'''			)			)
				_UpperCAmelCase						:  List[Any]     					=	{int(_UpperCamelCase			): 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[9_205]
								del idalabel[15_027]
				_UpperCAmelCase						:  List[str]     					=	idalabel
				_UpperCAmelCase						:  Any     					=	{v: k for k, v in idalabel.items()}
				_UpperCAmelCase						:  Union[str, Any]     					=	hidden_sizes
				_UpperCAmelCase						:  int     					=	depths

				return config, expected_shape


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

				return name


def 	_A	(   ):
				_UpperCAmelCase						:  Tuple     					=	'''http://images.cocodataset.org/val2017/000000039769.jpg'''
				_UpperCAmelCase						:  Union[str, Any]     					=	Image.open(requests.get(_UpperCamelCase      ,					stream=_UpperCamelCase			).raw			)
				return im


@torch.no_grad()
def 	_A	(   _UpperCamelCase      ,					_UpperCamelCase			):
				_UpperCAmelCase					, _UpperCAmelCase						:  Union[str, Any]     					=	get_convnext_config(_UpperCamelCase			)
				# load original state_dict from URL
				_UpperCAmelCase						:  List[Any]     					=	torch.hub.load_state_dict_from_url(_UpperCamelCase			)['''model''']
				# rename keys
				for key in state_dict.copy().keys():
								_UpperCAmelCase						:  str     					=	state_dict.pop(_UpperCamelCase			)
								_UpperCAmelCase						:  int     					=	val
				# add prefix to all keys expect classifier head
				for key in state_dict.copy().keys():
								_UpperCAmelCase						:  Any     					=	state_dict.pop(_UpperCamelCase			)
								if not key.startswith('''classifier'''			):
												_UpperCAmelCase						:  Dict     					=	'''convnext.''' + key
								_UpperCAmelCase						:  List[Any]     					=	val

				# load HuggingFace model
				_UpperCAmelCase						:  Optional[int]     					=	ConvNextForImageClassification(_UpperCamelCase			)
				model.load_state_dict(_UpperCamelCase			)
				model.eval()

				# Check outputs on an image, prepared by ConvNextImageProcessor
				_UpperCAmelCase						:  Optional[Any]     					=	224 if '''224''' in checkpoint_url else 384
				_UpperCAmelCase						:  int     					=	ConvNextImageProcessor(size=_UpperCamelCase			)
				_UpperCAmelCase						:  Union[str, Any]     					=	image_processor(images=prepare_img()      ,					return_tensors='''pt'''			).pixel_values

				_UpperCAmelCase						:  Any     					=	model(_UpperCamelCase			).logits

				# note: the logits below were obtained without center cropping
				if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
								_UpperCAmelCase						:  Optional[int]     					=	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":
								_UpperCAmelCase						:  Optional[Any]     					=	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":
								_UpperCAmelCase						:  List[str]     					=	torch.tensor([0.4_525, 0.7_539, 0.0_308]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
								_UpperCAmelCase						:  Any     					=	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":
								_UpperCAmelCase						:  str     					=	torch.tensor([0.4_174, -0.0_989, 0.1_489]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
								_UpperCAmelCase						:  Tuple     					=	torch.tensor([0.2_513, -0.1_349, -0.1_613]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
								_UpperCAmelCase						:  Dict     					=	torch.tensor([1.2_980, 0.3_631, -0.1_198]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
								_UpperCAmelCase						:  Optional[int]     					=	torch.tensor([1.2_963, 0.1_227, 0.1_723]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
								_UpperCAmelCase						:  str     					=	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":
								_UpperCAmelCase						:  List[str]     					=	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":
								_UpperCAmelCase						:  List[Any]     					=	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":
								_UpperCAmelCase						:  str     					=	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":
								_UpperCAmelCase						:  List[str]     					=	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":
								_UpperCAmelCase						:  Tuple     					=	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":
								_UpperCAmelCase						:  Any     					=	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]      ,					_UpperCamelCase      ,					atol=1e-3			)
				assert logits.shape == expected_shape

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

				print('''Pushing model to the hub...'''			)
				_UpperCAmelCase						:  Union[str, Any]     					=	'''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(_UpperCamelCase      ,					_UpperCamelCase			)      ,					organization='''nielsr'''      ,					commit_message='''Add model'''      ,					)


if __name__ == "__main__":
					UpperCAmelCase__			:   List[Any]      			=    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.',
					)

					UpperCAmelCase__			:   int      			=    parser.parse_args()
					convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)
