


import importlib
import inspect
import json
import os
import re
import shutil
import sys
from pathlib import Path
from typing import Dict, Optional, Union
from urllib import request

from huggingface_hub import HfFolder, cached_download, hf_hub_download, model_info
from packaging import version

from .. import __version__
from . import DIFFUSERS_DYNAMIC_MODULE_NAME, HF_MODULES_CACHE, logging


A				:       int      =  (
    'https://raw.githubusercontent.com/huggingface/diffusers/{revision}/examples/community/{pipeline}.py'
)


A				:       str      =  logging.get_logger(__name__)  # pylint: disable=invalid-name





def 						_lowerCAmelCase							(     )    ->						Dict:


						'''simple docstring'''






						__snake_case				=		"https://pypi.org/pypi/diffusers/json"
						__snake_case				=		json.loads(request.urlopen(_lowerCAmelCase			).read()			)["releases"].keys()
						return sorted(_lowerCAmelCase   ,						key=lambda _lowerCAmelCase			: version.Version(_lowerCAmelCase			)			)





def 						_lowerCAmelCase							(     )    ->						Tuple:


						'''simple docstring'''






						if HF_MODULES_CACHE in sys.path:
												return

						sys.path.append(_lowerCAmelCase			)
						os.makedirs(_lowerCAmelCase   ,						exist_ok=_lowerCAmelCase			)
						__snake_case				=		Path(_lowerCAmelCase			) / "__init__.py"
						if not init_path.exists():
												init_path.touch()





def 						_lowerCAmelCase							(     _lowerCAmelCase			)    ->						int:


						'''simple docstring'''






						init_hf_modules()
						__snake_case				=		Path(_lowerCAmelCase			) / name
						# If the parent module does not exist yet, recursively create it.
						if not dynamic_module_path.parent.exists():
												create_dynamic_module(dynamic_module_path.parent			)
						os.makedirs(_lowerCAmelCase   ,						exist_ok=_lowerCAmelCase			)
						__snake_case				=		dynamic_module_path / "__init__.py"
						if not init_path.exists():
												init_path.touch()





def 						_lowerCAmelCase							(     _lowerCAmelCase			)    ->						str:


						'''simple docstring'''






						with open(_lowerCAmelCase   ,						"r"   ,						encoding="utf-8"			) as f:
												__snake_case				=		f.read()

						# Imports of the form `import .xxx`
						__snake_case				=		re.findall("^\s*import\s+\.(\S+)\s*$"   ,						_lowerCAmelCase   ,						flags=re.MULTILINE			)
						# Imports of the form `from .xxx import yyy`
						relative_imports += re.findall("^\s*from\s+\.(\S+)\s+import"   ,						_lowerCAmelCase   ,						flags=re.MULTILINE			)
						# Unique-ify
						return list(set(_lowerCAmelCase			)			)





def 						_lowerCAmelCase							(     _lowerCAmelCase			)    ->						Dict:


						'''simple docstring'''






						__snake_case				=		False
						__snake_case				=		[module_file]
						__snake_case				=		[]

						# Let's recurse through all relative imports
						while not no_change:
												__snake_case				=		[]
												for f in files_to_check:
																		new_imports.extend(get_relative_imports(_lowerCAmelCase			)			)

												__snake_case				=		Path(_lowerCAmelCase			).parent
												__snake_case				=		[str(module_path / m			) for m in new_imports]
												__snake_case				=		[f for f in new_import_files if f not in all_relative_imports]
												__snake_case				=		[F'''{f}.py''' for f in new_import_files]

												__snake_case				=		len(_lowerCAmelCase			) == 0
												all_relative_imports.extend(_lowerCAmelCase			)

						return all_relative_imports





def 						_lowerCAmelCase							(     _lowerCAmelCase			)    ->						Any:


						'''simple docstring'''






						with open(_lowerCAmelCase   ,						"r"   ,						encoding="utf-8"			) as f:
												__snake_case				=		f.read()

						# Imports of the form `import xxx`
						__snake_case				=		re.findall("^\s*import\s+(\S+)\s*$"   ,						_lowerCAmelCase   ,						flags=re.MULTILINE			)
						# Imports of the form `from xxx import yyy`
						imports += re.findall("^\s*from\s+(\S+)\s+import"   ,						_lowerCAmelCase   ,						flags=re.MULTILINE			)
						# Only keep the top-level module
						__snake_case				=		[imp.split("."			)[0] for imp in imports if not imp.startswith("."			)]

						# Unique-ify and test we got them all
						__snake_case				=		list(set(_lowerCAmelCase			)			)
						__snake_case				=		[]
						for imp in imports:
												try:
																		importlib.import_module(_lowerCAmelCase			)
												except ImportError:
																		missing_packages.append(_lowerCAmelCase			)

						if len(_lowerCAmelCase			) > 0:
												raise ImportError(
												    "This modeling file requires the following packages that were not found in your environment: "
												    F'''{', '.join(_lowerCAmelCase			)}. Run `pip install {' '.join(_lowerCAmelCase			)}`'''			)

						return get_relative_imports(_lowerCAmelCase			)





def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase			)    ->						Union[str, Any]:


						'''simple docstring'''






						__snake_case				=		module_path.replace(os.path.sep   ,						"."			)
						__snake_case				=		importlib.import_module(_lowerCAmelCase			)

						if class_name is None:
												return find_pipeline_class(_lowerCAmelCase			)
						return getattr(_lowerCAmelCase   ,						_lowerCAmelCase			)





def 						_lowerCAmelCase							(     _lowerCAmelCase			)    ->						Optional[int]:


						'''simple docstring'''






						from ..pipelines import DiffusionPipeline

						__snake_case				=		dict(inspect.getmembers(_lowerCAmelCase   ,						inspect.isclass			)			)

						__snake_case				=		None
						for cls_name, cls in cls_members.items():
												if (
												    cls_name != DiffusionPipeline.__name__
												    and issubclass(cls   ,						_lowerCAmelCase			)
												    and cls.__module__.split("."			)[0] != "diffusers"
												):
																		if pipeline_class is not None:
																								raise ValueError(
																								    F'''Multiple classes that inherit from {DiffusionPipeline.__name__} have been found:'''
																								    F''' {pipeline_class.__name__}, and {cls_name}. Please make sure to define only one in'''
																								    F''' {loaded_module}.'''			)
																		__snake_case				=		cls

						return pipeline_class





def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase = None   ,						_lowerCAmelCase = False   ,						_lowerCAmelCase = False   ,						_lowerCAmelCase = None   ,						_lowerCAmelCase = None   ,						_lowerCAmelCase = None   ,						_lowerCAmelCase = False   ,						)    ->						List[Any]:


						'''simple docstring'''






						__snake_case				=		str(_lowerCAmelCase			)

						__snake_case				=		os.path.join(_lowerCAmelCase   ,						_lowerCAmelCase			)

						if os.path.isfile(_lowerCAmelCase			):
												__snake_case				=		module_file_or_url
												__snake_case				=		"local"
						elif pretrained_model_name_or_path.count("/"			) == 0:
												__snake_case				=		get_diffusers_versions()
												# cut ".dev0"
												__snake_case				=		"v" + ".".join(__version__.split("."			)[:3]			)

												# retrieve github version that matches
												if revision is None:
																		__snake_case				=		latest_version if latest_version[1:] in available_versions else "main"
																		logger.info(F'''Defaulting to latest_version: {revision}.'''			)
												elif revision in available_versions:
																		__snake_case				=		F'''v{revision}'''
												elif revision == "main":
																		__snake_case				=		revision
												else:
																		raise ValueError(
																		    F'''`custom_revision`: {revision} does not exist. Please make sure to choose one of'''
																		    F''' {', '.join(available_versions + ['main']			)}.'''			)

												# community pipeline on GitHub
												__snake_case				=		COMMUNITY_PIPELINES_URL.format(revision=_lowerCAmelCase   ,						pipeline=_lowerCAmelCase			)
												try:
																		__snake_case				=		cached_download(
																		    _lowerCAmelCase   ,						cache_dir=_lowerCAmelCase   ,						force_download=_lowerCAmelCase   ,						proxies=_lowerCAmelCase   ,						resume_download=_lowerCAmelCase   ,						local_files_only=_lowerCAmelCase   ,						use_auth_token=_lowerCAmelCase   ,						)
																		__snake_case				=		"git"
																		__snake_case				=		pretrained_model_name_or_path + ".py"
												except EnvironmentError:
																		logger.error(F'''Could not locate the {module_file} inside {pretrained_model_name_or_path}.'''			)
																		raise
						else:
												try:
																		# Load from URL or cache if already cached
																		__snake_case				=		hf_hub_download(
																		    _lowerCAmelCase   ,						_lowerCAmelCase   ,						cache_dir=_lowerCAmelCase   ,						force_download=_lowerCAmelCase   ,						proxies=_lowerCAmelCase   ,						resume_download=_lowerCAmelCase   ,						local_files_only=_lowerCAmelCase   ,						use_auth_token=_lowerCAmelCase   ,						)
																		__snake_case				=		os.path.join("local"   ,						"--".join(pretrained_model_name_or_path.split("/"			)			)			)
												except EnvironmentError:
																		logger.error(F'''Could not locate the {module_file} inside {pretrained_model_name_or_path}.'''			)
																		raise

    # Check we have all the requirements in our environment
						__snake_case				=		check_imports(_lowerCAmelCase			)

						# Now we move the module inside our cached dynamic modules.
						__snake_case				=		DIFFUSERS_DYNAMIC_MODULE_NAME + os.path.sep + submodule
						create_dynamic_module(_lowerCAmelCase			)
						__snake_case				=		Path(_lowerCAmelCase			) / full_submodule
						if submodule == "local" or submodule == "git":
												# We always copy local files (we could hash the file to see if there was a change, and give them the name of
												# that hash, to only copy when there is a modification but it seems overkill for now).
												# The only reason we do the copy is to avoid putting too many folders in sys.path.
												shutil.copy(_lowerCAmelCase   ,						submodule_path / module_file			)
												for module_needed in modules_needed:
																		__snake_case				=		F'''{module_needed}.py'''
																		shutil.copy(os.path.join(_lowerCAmelCase   ,						_lowerCAmelCase			)   ,						submodule_path / module_needed			)
						else:
												# Get the commit hash
												# TODO: we will get this info in the etag soon, so retrieve it from there and not here.
												if isinstance(_lowerCAmelCase   ,						_lowerCAmelCase			):
																		__snake_case				=		use_auth_token
												elif use_auth_token is True:
																		__snake_case				=		HfFolder.get_token()
												else:
																		__snake_case				=		None

												__snake_case				=		model_info(_lowerCAmelCase   ,						revision=_lowerCAmelCase   ,						token=_lowerCAmelCase			).sha

												# The module file will end up being placed in a subfolder with the git hash of the repo. This way we get the
												# benefit of versioning.
												__snake_case				=		submodule_path / commit_hash
												__snake_case				=		full_submodule + os.path.sep + commit_hash
												create_dynamic_module(_lowerCAmelCase			)

												if not (submodule_path / module_file).exists():
																		shutil.copy(_lowerCAmelCase   ,						submodule_path / module_file			)
												# Make sure we also have every file with relative
												for module_needed in modules_needed:
																		if not (submodule_path / module_needed).exists():
																								get_cached_module_file(
																								    _lowerCAmelCase   ,						F'''{module_needed}.py'''   ,						cache_dir=_lowerCAmelCase   ,						force_download=_lowerCAmelCase   ,						resume_download=_lowerCAmelCase   ,						proxies=_lowerCAmelCase   ,						use_auth_token=_lowerCAmelCase   ,						revision=_lowerCAmelCase   ,						local_files_only=_lowerCAmelCase   ,						)
						return os.path.join(_lowerCAmelCase   ,						_lowerCAmelCase			)





def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase = None   ,						_lowerCAmelCase = None   ,						_lowerCAmelCase = False   ,						_lowerCAmelCase = False   ,						_lowerCAmelCase = None   ,						_lowerCAmelCase = None   ,						_lowerCAmelCase = None   ,						_lowerCAmelCase = False   ,						**_lowerCAmelCase   ,						)    ->						str:


						'''simple docstring'''






						__snake_case				=		get_cached_module_file(
						    _lowerCAmelCase   ,						_lowerCAmelCase   ,						cache_dir=_lowerCAmelCase   ,						force_download=_lowerCAmelCase   ,						resume_download=_lowerCAmelCase   ,						proxies=_lowerCAmelCase   ,						use_auth_token=_lowerCAmelCase   ,						revision=_lowerCAmelCase   ,						local_files_only=_lowerCAmelCase   ,						)
						return get_class_in_module(_lowerCAmelCase   ,						final_module.replace(".py"   ,						""			)			)









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 ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
A				:       int      =  logging.get_logger(__name__)





def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase=False			)    ->						Optional[int]:


						'''simple docstring'''






						__snake_case				=		[]
						for i in range(config.num_hidden_layers			):
												# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
												rename_keys.append((F'''blocks.{i}.norm1.weight''', F'''vit.encoder.layer.{i}.layernorm_before.weight''')			)
												rename_keys.append((F'''blocks.{i}.norm1.bias''', F'''vit.encoder.layer.{i}.layernorm_before.bias''')			)
												rename_keys.append((F'''blocks.{i}.attn.proj.weight''', F'''vit.encoder.layer.{i}.attention.output.dense.weight''')			)
												rename_keys.append((F'''blocks.{i}.attn.proj.bias''', F'''vit.encoder.layer.{i}.attention.output.dense.bias''')			)
												rename_keys.append((F'''blocks.{i}.norm2.weight''', F'''vit.encoder.layer.{i}.layernorm_after.weight''')			)
												rename_keys.append((F'''blocks.{i}.norm2.bias''', F'''vit.encoder.layer.{i}.layernorm_after.bias''')			)
												rename_keys.append((F'''blocks.{i}.mlp.fc1.weight''', F'''vit.encoder.layer.{i}.intermediate.dense.weight''')			)
												rename_keys.append((F'''blocks.{i}.mlp.fc1.bias''', F'''vit.encoder.layer.{i}.intermediate.dense.bias''')			)
												rename_keys.append((F'''blocks.{i}.mlp.fc2.weight''', F'''vit.encoder.layer.{i}.output.dense.weight''')			)
												rename_keys.append((F'''blocks.{i}.mlp.fc2.bias''', F'''vit.encoder.layer.{i}.output.dense.bias''')			)

						# projection layer + position embeddings
						rename_keys.extend(
						    [
						        ("cls_token", "vit.embeddings.cls_token"),
						        ("patch_embed.proj.weight", "vit.embeddings.patch_embeddings.projection.weight"),
						        ("patch_embed.proj.bias", "vit.embeddings.patch_embeddings.projection.bias"),
						        ("pos_embed", "vit.embeddings.position_embeddings"),
						    ]			)

						if base_model:
												# layernorm + pooler
												rename_keys.extend(
												    [
												        ("norm.weight", "layernorm.weight"),
												        ("norm.bias", "layernorm.bias"),
												    ]			)

												# if just the base model, we should remove "vit" from all keys that start with "vit"
												__snake_case				=		[(pair[0], pair[1][4:]) if pair[1].startswith("vit"			) else pair for pair in rename_keys]
						else:
												# layernorm + classification head
												rename_keys.extend(
												    [
												        ("norm.weight", "vit.layernorm.weight"),
												        ("norm.bias", "vit.layernorm.bias"),
												        ("head.weight", "classifier.weight"),
												        ("head.bias", "classifier.bias"),
												    ]			)

						return rename_keys





def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase=False			)    ->						int:


						'''simple docstring'''






						for i in range(config.num_hidden_layers			):
												if base_model:
																		__snake_case				=		""
												else:
																		__snake_case				=		"vit."
												# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
												__snake_case				=		state_dict.pop(F'''blocks.{i}.attn.qkv.weight'''			)
												__snake_case				=		state_dict.pop(F'''blocks.{i}.attn.qkv.bias'''			)
												# next, add query, keys and values (in that order) to the state dict
												__snake_case				=		in_proj_weight[
												    : config.hidden_size, :
												]
												__snake_case				=		in_proj_bias[: config.hidden_size]
												__snake_case				=		in_proj_weight[
												    config.hidden_size : config.hidden_size * 2, :
												]
												__snake_case				=		in_proj_bias[
												    config.hidden_size : config.hidden_size * 2
												]
												__snake_case				=		in_proj_weight[
												    -config.hidden_size :, :
												]
												__snake_case				=		in_proj_bias[-config.hidden_size :]





def 						_lowerCAmelCase							(     _lowerCAmelCase			)    ->						List[str]:


						'''simple docstring'''






						__snake_case				=		["head.weight", "head.bias"]
						for k in ignore_keys:
												state_dict.pop(_lowerCAmelCase   ,						_lowerCAmelCase			)





def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase			)    ->						Dict:


						'''simple docstring'''






						__snake_case				=		dct.pop(_lowerCAmelCase			)
						__snake_case				=		val





def 						_lowerCAmelCase							(     )    ->						Any:


						'''simple docstring'''






						__snake_case				=		"http://images.cocodataset.org/val2017/000000039769.jpg"
						__snake_case				=		Image.open(requests.get(_lowerCAmelCase   ,						stream=_lowerCAmelCase			).raw			)
						return im





@torch.no_grad()
def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase=True			)    ->						int:


						'''simple docstring'''






						__snake_case				=		ViTConfig()
						# patch_size
						if model_name[-1] == "8":
												__snake_case				=		8
						# set labels if required
						if not base_model:
												__snake_case				=		1000
												__snake_case				=		"huggingface/label-files"
												__snake_case				=		"imagenet-1k-id2label.json"
												__snake_case				=		json.load(open(hf_hub_download(_lowerCAmelCase   ,						_lowerCAmelCase   ,						repo_type="dataset"			)   ,						"r"			)			)
												__snake_case				=		{int(_lowerCAmelCase			): v for k, v in idalabel.items()}
												__snake_case				=		idalabel
												__snake_case				=		{v: k for k, v in idalabel.items()}
						# size of the architecture
						if model_name in ["dino_vits8", "dino_vits16"]:
												__snake_case				=		384
												__snake_case				=		1536
												__snake_case				=		12
												__snake_case				=		6

						# load original model from torch hub
						__snake_case				=		torch.hub.load("facebookresearch/dino:main"   ,						_lowerCAmelCase			)
						original_model.eval()

						# load state_dict of original model, remove and rename some keys
						__snake_case				=		original_model.state_dict()
						if base_model:
												remove_classification_head_(_lowerCAmelCase			)
						__snake_case				=		create_rename_keys(_lowerCAmelCase   ,						base_model=_lowerCAmelCase			)
						for src, dest in rename_keys:
												rename_key(_lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase			)
						read_in_q_k_v(_lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase			)

						# load HuggingFace model
						if base_model:
												__snake_case				=		ViTModel(_lowerCAmelCase   ,						add_pooling_layer=_lowerCAmelCase			).eval()
						else:
												__snake_case				=		ViTForImageClassification(_lowerCAmelCase			).eval()
						model.load_state_dict(_lowerCAmelCase			)

						# Check outputs on an image, prepared by ViTImageProcessor
						__snake_case				=		ViTImageProcessor()
						__snake_case				=		image_processor(images=prepare_img()   ,						return_tensors="pt"			)
						__snake_case				=		encoding["pixel_values"]
						__snake_case				=		model(_lowerCAmelCase			)

						if base_model:
												__snake_case				=		original_model(_lowerCAmelCase			)
												assert torch.allclose(_lowerCAmelCase   ,						outputs.last_hidden_state[:, 0, :]   ,						atol=1E-1			)
						else:
												__snake_case				=		original_model(_lowerCAmelCase			)
												assert logits.shape == outputs.logits.shape
												assert torch.allclose(_lowerCAmelCase   ,						outputs.logits   ,						atol=1E-3			)

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


if __name__ == "__main__":
	A				:       str      =  argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    '--model_name',
	    default='dino_vitb16',
	    type=str,
	    help='Name of the model trained with DINO you\'d like to convert.',
	)
	parser.add_argument(
	    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
	)
	parser.add_argument(
	    '--base_model',
	    action='store_true',
	    help='Whether to only convert the base model (no projection head weights).',
	)

	parser.set_defaults(base_model=True)
	A				:       Any      =  parser.parse_args()
	convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)





