

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 							__lowercase     (	UpperCAmelCase__   ):

					"""simple docstring"""






					__lowerCAmelCase       =				ConvNextVaConfig()

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

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

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

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

					return config, expected_shape


def 							__lowercase     (	UpperCAmelCase__   ):

					"""simple docstring"""






					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 "gamma" in name:
										__lowerCAmelCase       =				name.replace('gamma'		,       'weight'   )
					if "beta" in name:
										__lowerCAmelCase       =				name.replace('beta'		,       'bias'   )
					if "stages" in name:
										__lowerCAmelCase       =				name.replace('stages'		,       'encoder.stages'   )
					if "norm" in name:
										__lowerCAmelCase       =				name.replace('norm'		,       'layernorm'   )
					if "head" in name:
										__lowerCAmelCase       =				name.replace('head'		,       'classifier'   )

					return name


def 							__lowercase     (	):

					"""simple docstring"""






					__lowerCAmelCase       =				'http://images.cocodataset.org/val2017/000000039769.jpg'
					__lowerCAmelCase       =				Image.open(requests.get(UpperCAmelCase__		,       stream=UpperCAmelCase__   ).raw   )
					return im


def 							__lowercase     (	UpperCAmelCase__   ):

					"""simple docstring"""






					if "224" in checkpoint_url:
										__lowerCAmelCase       =				224
										__lowerCAmelCase       =				224 / 256
					elif "384" in checkpoint_url:
										__lowerCAmelCase       =				384
										__lowerCAmelCase       =				None
					else:
										__lowerCAmelCase       =				512
										__lowerCAmelCase       =				None

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


@torch.no_grad()
def 							__lowercase     (	UpperCAmelCase__		,       UpperCAmelCase__		,       UpperCAmelCase__		,       UpperCAmelCase__   ):

					"""simple docstring"""






					print('Downloading original model from checkpoint...'   )
					# define ConvNeXTV2 configuration based on URL
					__lowerCAmelCase,						__lowerCAmelCase       =				get_convnextva_config(UpperCAmelCase__   )
					# load original state_dict from URL
					__lowerCAmelCase       =				torch.hub.load_state_dict_from_url(UpperCAmelCase__   )['model']

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

					# load HuggingFace model
					__lowerCAmelCase       =				ConvNextVaForImageClassification(UpperCAmelCase__   )
					model.load_state_dict(UpperCAmelCase__   )
					model.eval()

					# Check outputs on an image, prepared by ConvNextImageProcessor
					__lowerCAmelCase       =				convert_preprocessor(UpperCAmelCase__   )
					__lowerCAmelCase       =				preprocessor(images=prepare_img()		,       return_tensors='pt'   )
					__lowerCAmelCase       =				model(**UpperCAmelCase__   ).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":
										__lowerCAmelCase       =				torch.tensor([-0.3930, 0.1747, -0.5246, 0.4177, 0.4295]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
										__lowerCAmelCase       =				torch.tensor([-0.1727, -0.5341, -0.7818, -0.4745, -0.6566]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
										__lowerCAmelCase       =				torch.tensor([-0.0333, 0.1563, -0.9137, 0.1054, 0.0381]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
										__lowerCAmelCase       =				torch.tensor([-0.1744, -0.1555, -0.0713, 0.0950, -0.1431]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
										__lowerCAmelCase       =				torch.tensor([0.9996, 0.1966, -0.4386, -0.3472, 0.6661]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
										__lowerCAmelCase       =				torch.tensor([-0.2553, -0.6708, -0.1359, 0.2518, -0.2488]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
										__lowerCAmelCase       =				torch.tensor([-0.0673, -0.5627, -0.3753, -0.2722, 0.0178]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
										__lowerCAmelCase       =				torch.tensor([-0.6377, -0.7458, -0.2150, 0.1184, -0.0597]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
										__lowerCAmelCase       =				torch.tensor([1.0799, 0.2322, -0.8860, 1.0219, 0.6231]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
										__lowerCAmelCase       =				torch.tensor([0.3766, 0.4917, -1.1426, 0.9942, 0.6024]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
										__lowerCAmelCase       =				torch.tensor([0.4220, -0.6919, -0.4317, -0.2881, -0.6609]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
										__lowerCAmelCase       =				torch.tensor([0.1082, -0.8286, -0.5095, 0.4681, -0.8085]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
										__lowerCAmelCase       =				torch.tensor([-0.2419, -0.6221, 0.2176, -0.0980, -0.7527]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
										__lowerCAmelCase       =				torch.tensor([0.0391, -0.4371, 0.3786, 0.1251, -0.2784]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
										__lowerCAmelCase       =				torch.tensor([-0.0504, 0.5636, -0.1729, -0.6507, -0.3949]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
										__lowerCAmelCase       =				torch.tensor([0.3560, 0.9486, 0.3149, -0.2667, -0.5138]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
										__lowerCAmelCase       =				torch.tensor([-0.2469, -0.4550, -0.5853, -0.0810, 0.0309]   )
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
										__lowerCAmelCase       =				torch.tensor([-0.3090, 0.0802, -0.0682, -0.1979, -0.2826]   )
					else:
										raise ValueError(F"""Unknown URL: {checkpoint_url}"""   )

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

										model.save_pretrained(UpperCAmelCase__   )
										preprocessor.save_pretrained(UpperCAmelCase__   )

					__lowerCAmelCase       =				'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(UpperCAmelCase__   )
										preprocessor.push_to_hub(UpperCAmelCase__   )


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
		)









import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class      snake_case_      (					_a      ,			_a      ,			unittest.TestCase				):


							"""simple docstring"""





							__UpperCAmelCase      =StableDiffusionParadigmsPipeline
							__UpperCAmelCase      =TEXT_TO_IMAGE_PARAMS
							__UpperCAmelCase      =TEXT_TO_IMAGE_BATCH_PARAMS
							__UpperCAmelCase      =TEXT_TO_IMAGE_IMAGE_PARAMS
							__UpperCAmelCase      =TEXT_TO_IMAGE_IMAGE_PARAMS




							def 							A__	(			self				):
												torch.manual_seed(0				)
												__lowerCAmelCase       =				UNetaDConditionModel(
												    block_out_channels=(3_2, 6_4)       ,    layers_per_block=2       ,    sample_size=3_2       ,    in_channels=4       ,    out_channels=4       ,    down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')       ,    up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')       ,    cross_attention_dim=3_2       ,    attention_head_dim=(2, 4)       ,    use_linear_projection=_A       ,    )
												__lowerCAmelCase       =				DDIMParallelScheduler(
												    beta_start=0.0_0085       ,    beta_end=0.012       ,    beta_schedule='scaled_linear'       ,    clip_sample=_A       ,    set_alpha_to_one=_A       ,    )
												torch.manual_seed(0				)
												__lowerCAmelCase       =				AutoencoderKL(
												    block_out_channels=[3_2, 6_4]       ,    in_channels=3       ,    out_channels=3       ,    down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']       ,    up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']       ,    latent_channels=4       ,    sample_size=1_2_8       ,    )
												torch.manual_seed(0				)
												__lowerCAmelCase       =				CLIPTextConfig(
												    bos_token_id=0       ,    eos_token_id=2       ,    hidden_size=3_2       ,    intermediate_size=3_7       ,    layer_norm_eps=1e-0_5       ,    num_attention_heads=4       ,    num_hidden_layers=5       ,    pad_token_id=1       ,    vocab_size=1_0_0_0       ,    hidden_act='gelu'       ,    projection_dim=5_1_2       ,    )
												__lowerCAmelCase       =				CLIPTextModel(_A				)
												__lowerCAmelCase       =				CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'				)

												__lowerCAmelCase       =				{
												    'unet': unet,
												    'scheduler': scheduler,
												    'vae': vae,
												    'text_encoder': text_encoder,
												    'tokenizer': tokenizer,
												    'safety_checker': None,
												    'feature_extractor': None,
												}
												return components




							def 							A__	(			self       ,    _A       ,    _A=0				):
												if str(_A				).startswith('mps'				):
																	__lowerCAmelCase       =				torch.manual_seed(_A				)
												else:
																	__lowerCAmelCase       =				torch.Generator(device=_A				).manual_seed(_A				)
												__lowerCAmelCase       =				{
												    'prompt': 'a photograph of an astronaut riding a horse',
												    'generator': generator,
												    'num_inference_steps': 1_0,
												    'guidance_scale': 6.0,
												    'output_type': 'numpy',
												    'parallel': 3,
												    'debug': True,
												}
												return inputs




							def 							A__	(			self				):
												__lowerCAmelCase       =				'cpu'  # ensure determinism for the device-dependent torch.Generator
												__lowerCAmelCase       =				self.get_dummy_components()
												__lowerCAmelCase       =				StableDiffusionParadigmsPipeline(**_A				)
												__lowerCAmelCase       =				sd_pipe.to(_A				)
												sd_pipe.set_progress_bar_config(disable=_A				)

												__lowerCAmelCase       =				self.get_dummy_inputs(_A				)
												__lowerCAmelCase       =				sd_pipe(**_A				).images
												__lowerCAmelCase       =				image[0, -3:, -3:, -1]
												assert image.shape == (1, 6_4, 6_4, 3)

												__lowerCAmelCase       =				np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023]				)

												assert np.abs(image_slice.flatten() - expected_slice				).max() < 1e-2




							def 							A__	(			self				):
												__lowerCAmelCase       =				'cpu'  # ensure determinism for the device-dependent torch.Generator
												__lowerCAmelCase       =				self.get_dummy_components()
												torch.manual_seed(0				)
												__lowerCAmelCase       =				DDPMParallelScheduler()
												torch.manual_seed(0				)
												__lowerCAmelCase       =				StableDiffusionParadigmsPipeline(**_A				)
												__lowerCAmelCase       =				sd_pipe.to(_A				)
												sd_pipe.set_progress_bar_config(disable=_A				)

												__lowerCAmelCase       =				self.get_dummy_inputs(_A				)
												__lowerCAmelCase       =				sd_pipe(**_A				).images
												__lowerCAmelCase       =				image[0, -3:, -3:, -1]
												assert image.shape == (1, 6_4, 6_4, 3)

												__lowerCAmelCase       =				np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468]				)

												assert np.abs(image_slice.flatten() - expected_slice				).max() < 1e-2




							def 							A__	(			self				):
												super().test_inference_batch_consistent(batch_sizes=[1, 2]				)




							def 							A__	(			self				):
												super().test_inference_batch_single_identical(batch_size=2       ,    expected_max_diff=3e-3				)




							def 							A__	(			self				):
												__lowerCAmelCase       =				'cpu'  # ensure determinism for the device-dependent torch.Generator
												__lowerCAmelCase       =				self.get_dummy_components()
												__lowerCAmelCase       =				StableDiffusionParadigmsPipeline(**_A				)
												__lowerCAmelCase       =				sd_pipe.to(_A				)
												sd_pipe.set_progress_bar_config(disable=_A				)

												__lowerCAmelCase       =				self.get_dummy_inputs(_A				)
												__lowerCAmelCase       =				'french fries'
												__lowerCAmelCase       =				sd_pipe(**_A       ,    negative_prompt=_A				)
												__lowerCAmelCase       =				output.images
												__lowerCAmelCase       =				image[0, -3:, -3:, -1]

												assert image.shape == (1, 6_4, 6_4, 3)

												__lowerCAmelCase       =				np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015]				)

												assert np.abs(image_slice.flatten() - expected_slice				).max() < 1e-2





@slow
@require_torch_gpu
class      snake_case_      (					unittest.TestCase				):


							"""simple docstring"""





							def 							A__	(			self				):
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()




							def 							A__	(			self       ,    _A=0				):
												__lowerCAmelCase       =				torch.Generator(device=_A				).manual_seed(_A				)
												__lowerCAmelCase       =				{
												    'prompt': 'a photograph of an astronaut riding a horse',
												    'generator': generator,
												    'num_inference_steps': 1_0,
												    'guidance_scale': 7.5,
												    'output_type': 'numpy',
												    'parallel': 3,
												    'debug': True,
												}
												return inputs




							def 							A__	(			self				):
												__lowerCAmelCase       =				'stabilityai/stable-diffusion-2-base'
												__lowerCAmelCase       =				DDIMParallelScheduler.from_pretrained(_A       ,    subfolder='scheduler'				)
												__lowerCAmelCase       =				StableDiffusionParadigmsPipeline.from_pretrained(_A       ,    scheduler=_A       ,    safety_checker=_A				)
												pipe.to(_A				)
												pipe.set_progress_bar_config(disable=_A				)
												pipe.enable_attention_slicing()

												__lowerCAmelCase       =				self.get_inputs()
												__lowerCAmelCase       =				pipe(**_A				).images
												__lowerCAmelCase       =				image[0, -3:, -3:, -1].flatten()

												assert image.shape == (1, 5_1_2, 5_1_2, 3)

												__lowerCAmelCase       =				np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741]				)

												assert np.abs(expected_slice - image_slice				).max() < 1e-2






