
"""simple docstring"""







import argparse
import json
from pathlib import Path

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

from transformers import DeiTImageProcessor, ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
_SCREAMING_SNAKE_CASE	       =      logging.get_logger(__name__)
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE=False )  ->     Tuple:







							"""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"),
														        ("pre_logits.fc.weight", "pooler.dense.weight"),
														        ("pre_logits.fc.bias", "pooler.dense.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   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE=False )  ->     str:







							"""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   __UpperCamelCase   (					SCREAMING_SNAKE_CASE )  ->     List[Any]:







							"""simple docstring"""


							__snake_case     			=       ["head.weight", "head.bias"]
							for k in ignore_keys:
														state_dict.pop(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )  ->     Dict:







							"""simple docstring"""


							__snake_case     			=       dct.pop(SCREAMING_SNAKE_CASE )
							__snake_case     			=       val
def   __UpperCamelCase   (					)  ->     str:







							"""simple docstring"""


							__snake_case     			=       "http://images.cocodataset.org/val2017/000000039769.jpg"
							__snake_case     			=       Image.open(requests.get(SCREAMING_SNAKE_CASE							,    stream=SCREAMING_SNAKE_CASE ).raw )
							return im
@torch.no_grad()
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )  ->     Dict:







							"""simple docstring"""


							__snake_case     			=       ViTConfig()
							__snake_case     			=       False
							# dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
							if vit_name[-5:] == "in21k":
														__snake_case     			=       True
														__snake_case     			=       int(vit_name[-12:-10] )
														__snake_case     			=       int(vit_name[-9:-6] )
							else:
														__snake_case     			=       10_00
														__snake_case     			=       "huggingface/label-files"
														__snake_case     			=       "imagenet-1k-id2label.json"
														__snake_case     			=       json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    repo_type="dataset" )							,    "r" ) )
														__snake_case     			=       {int(SCREAMING_SNAKE_CASE ): v for k, v in idalabel.items()}
														__snake_case     			=       idalabel
														__snake_case     			=       {v: k for k, v in idalabel.items()}
														__snake_case     			=       int(vit_name[-6:-4] )
														__snake_case     			=       int(vit_name[-3:] )
							# size of the architecture
							if "deit" in vit_name:
														if vit_name[9:].startswith("tiny" ):
																					__snake_case     			=       1_92
																					__snake_case     			=       7_68
																					__snake_case     			=       12
																					__snake_case     			=       3
														elif vit_name[9:].startswith("small" ):
																					__snake_case     			=       3_84
																					__snake_case     			=       15_36
																					__snake_case     			=       12
																					__snake_case     			=       6
														else:
																					pass
							else:
														if vit_name[4:].startswith("small" ):
																					__snake_case     			=       7_68
																					__snake_case     			=       23_04
																					__snake_case     			=       8
																					__snake_case     			=       8
														elif vit_name[4:].startswith("base" ):
																					pass
														elif vit_name[4:].startswith("large" ):
																					__snake_case     			=       10_24
																					__snake_case     			=       40_96
																					__snake_case     			=       24
																					__snake_case     			=       16
														elif vit_name[4:].startswith("huge" ):
																					__snake_case     			=       12_80
																					__snake_case     			=       51_20
																					__snake_case     			=       32
																					__snake_case     			=       16

    # load original model from timm
							__snake_case     			=       timm.create_model(SCREAMING_SNAKE_CASE							,    pretrained=SCREAMING_SNAKE_CASE )
							timm_model.eval()

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

							# load HuggingFace model
							if vit_name[-5:] == "in21k":
														__snake_case     			=       ViTModel(SCREAMING_SNAKE_CASE ).eval()
							else:
														__snake_case     			=       ViTForImageClassification(SCREAMING_SNAKE_CASE ).eval()
							model.load_state_dict(SCREAMING_SNAKE_CASE )

							# Check outputs on an image, prepared by ViTImageProcessor/DeiTImageProcessor
							if "deit" in vit_name:
														__snake_case     			=       DeiTImageProcessor(size=config.image_size )
							else:
														__snake_case     			=       ViTImageProcessor(size=config.image_size )
							__snake_case     			=       image_processor(images=prepare_img()							,    return_tensors="pt" )
							__snake_case     			=       encoding["pixel_values"]
							__snake_case     			=       model(SCREAMING_SNAKE_CASE )

							if base_model:
														__snake_case     			=       timm_model.forward_features(SCREAMING_SNAKE_CASE )
														assert timm_pooled_output.shape == outputs.pooler_output.shape
														assert torch.allclose(SCREAMING_SNAKE_CASE							,    outputs.pooler_output							,    atol=1e-3 )
							else:
														__snake_case     			=       timm_model(SCREAMING_SNAKE_CASE )
														assert timm_logits.shape == outputs.logits.shape
														assert torch.allclose(SCREAMING_SNAKE_CASE							,    outputs.logits							,    atol=1e-3 )

							Path(SCREAMING_SNAKE_CASE ).mkdir(exist_ok=SCREAMING_SNAKE_CASE )
							print(F'''Saving model {vit_name} 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 )


if __name__ == "__main__":
						_SCREAMING_SNAKE_CASE	       =      argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    """--vit_name""",
						    default="""vit_base_patch16_224""",
						    type=str,
						    help="""Name of the ViT timm model you'd like to convert.""",
						)
						parser.add_argument(
						    """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
						)

						_SCREAMING_SNAKE_CASE	       =      parser.parse_args()
						convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)


"""simple docstring"""







from collections import UserDict
from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
		from PIL import Image

		from ..image_utils import load_image

if is_torch_available():
		from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING

if is_tf_available():
		from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
		from ..tf_utils import stable_softmax

_SCREAMING_SNAKE_CASE	       =      logging.get_logger(__name__)







@add_end_docstrings(lowercase__   )
class 							__magic_name__     (		lowercase__   ):





					def __init__(				self   :							Tuple							,						**snake_case_   :							str    ):
												super().__init__(**snake_case_    )

												requires_backends(self							,						"vision"    )
												self.check_model_type(
												    TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
												    if self.framework == "tf"
												    else MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING    )





					def __call__(				self   :							Optional[Any]							,						snake_case_   :							Union[str, List[str], "Image", List["Image"]]							,						**snake_case_   :							Optional[int]    ):
												return super().__call__(snake_case_							,						**snake_case_    )





					def 							lowerCAmelCase      (				self   :							Tuple							,						**snake_case_   :							Any    ):
												__snake_case     			=       {}
												if "candidate_labels" in kwargs:
																			__snake_case     			=       kwargs["candidate_labels"]
												if "hypothesis_template" in kwargs:
																			__snake_case     			=       kwargs["hypothesis_template"]

												return preprocess_params, {}, {}





					def 							lowerCAmelCase      (				self   :							int							,						snake_case_   :							List[str]							,						snake_case_   :							List[str]=None							,						snake_case_   :							int="This is a photo of {}."    ):
												__snake_case     			=       load_image(snake_case_    )
												__snake_case     			=       self.image_processor(images=[image]							,						return_tensors=self.framework    )
												__snake_case     			=       candidate_labels
												__snake_case     			=       [hypothesis_template.format(snake_case_    ) for x in candidate_labels]
												__snake_case     			=       self.tokenizer(snake_case_							,						return_tensors=self.framework							,						padding=snake_case_    )
												__snake_case     			=       [text_inputs]
												return inputs





					def 							lowerCAmelCase      (				self   :							Any							,						snake_case_   :							Dict    ):
												__snake_case     			=       model_inputs.pop("candidate_labels"    )
												__snake_case     			=       model_inputs.pop("text_inputs"    )
												if isinstance(text_inputs[0]							,						snake_case_    ):
																			__snake_case     			=       text_inputs[0]
												else:
																			# Batching case.
																			__snake_case     			=       text_inputs[0][0]

												__snake_case     			=       self.model(**snake_case_							,						**snake_case_    )

												__snake_case     			=       {
												    "candidate_labels": candidate_labels,
												    "logits": outputs.logits_per_image,
												}
												return model_outputs





					def 							lowerCAmelCase      (				self   :							Optional[int]							,						snake_case_   :							Any    ):
												__snake_case     			=       model_outputs.pop("candidate_labels"    )
												__snake_case     			=       model_outputs["logits"][0]
												if self.framework == "pt":
																			__snake_case     			=       logits.softmax(dim=-1    ).squeeze(-1    )
																			__snake_case     			=       probs.tolist()
																			if not isinstance(snake_case_							,						snake_case_    ):
																										__snake_case     			=       [scores]
												elif self.framework == "tf":
																			__snake_case     			=       stable_softmax(snake_case_							,						axis=-1    )
																			__snake_case     			=       probs.numpy().tolist()
												else:
																			raise ValueError(F'''Unsupported framework: {self.framework}'''    )

												__snake_case     			=       [
												    {"score": score, "label": candidate_label}
												    for score, candidate_label in sorted(zip(snake_case_							,						snake_case_    )							,						key=lambda snake_case_    : -x[0]    )
												]
												return result
