

"""simple docstring"""





import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision import transforms

from transformers import BitImageProcessor, FocalNetConfig, FocalNetForImageClassification
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling

def     __SCREAMING_SNAKE_CASE    ( A_					):
	lowerCAmelCase__      :		Tuple    							=      [2, 2, 6, 2] if '''tiny''' in model_name else [2, 2, 18, 2]
	lowerCAmelCase__      :		Optional[Any]    							=      True if '''large''' in model_name or '''huge''' in model_name else False
	lowerCAmelCase__      :		Union[str, Any]    							=      True if '''large''' in model_name or '''huge''' in model_name else False
	lowerCAmelCase__      :		List[Any]    							=      True if '''large''' in model_name or '''huge''' in model_name else False

	if "large" in model_name or "xlarge" in model_name or "huge" in model_name:
		if "fl3" in model_name:
			lowerCAmelCase__      :		Optional[Any]    							=      [3, 3, 3, 3]
			lowerCAmelCase__      :		Optional[Any]    							=      [5, 5, 5, 5]
		elif "fl4" in model_name:
			lowerCAmelCase__      :		Any    							=      [4, 4, 4, 4]
			lowerCAmelCase__      :		Optional[int]    							=      [3, 3, 3, 3]

	if "tiny" in model_name or "small" in model_name or "base" in model_name:
		lowerCAmelCase__      :		Union[str, Any]    							=      [3, 3, 3, 3]
		if "lrf" in model_name:
			lowerCAmelCase__      :		Optional[Any]    							=      [3, 3, 3, 3]
		else:
			lowerCAmelCase__      :		Union[str, Any]    							=      [2, 2, 2, 2]

	if "tiny" in model_name:
		lowerCAmelCase__      :		List[Any]    							=      96
	elif "small" in model_name:
		lowerCAmelCase__      :		Any    							=      96
	elif "base" in model_name:
		lowerCAmelCase__      :		List[str]    							=      1_28
	elif "large" in model_name:
		lowerCAmelCase__      :		Optional[Any]    							=      1_92
	elif "xlarge" in model_name:
		lowerCAmelCase__      :		Any    							=      2_56
	elif "huge" in model_name:
		lowerCAmelCase__      :		Optional[Any]    							=      3_52

	# set label information
	lowerCAmelCase__      :		List[Any]    							=      '''huggingface/label-files'''
	if "large" in model_name or "huge" in model_name:
		lowerCAmelCase__      :		Dict    							=      '''imagenet-22k-id2label.json'''
	else:
		lowerCAmelCase__      :		Optional[int]    							=      '''imagenet-1k-id2label.json'''

	lowerCAmelCase__      :		Union[str, Any]    							=      json.load(open(hf_hub_download(A_		,	A_		,	repo_type='''dataset'''					)		,	'''r'''					)					)
	lowerCAmelCase__      :		List[Any]    							=      {int(A_					): v for k, v in idalabel.items()}
	lowerCAmelCase__      :		Any    							=      {v: k for k, v in idalabel.items()}

	lowerCAmelCase__      :		List[str]    							=      FocalNetConfig(
	    embed_dim=A_		,	depths=A_		,	focal_levels=A_		,	focal_windows=A_		,	use_conv_embed=A_		,	idalabel=A_		,	labelaid=A_		,	use_post_layernorm=A_		,	use_layerscale=A_		,	)

	return config

def     __SCREAMING_SNAKE_CASE    ( A_					):
	if "patch_embed.proj" in name:
		lowerCAmelCase__      :		List[str]    							=      name.replace('''patch_embed.proj'''		,	'''embeddings.patch_embeddings.projection'''					)
	if "patch_embed.norm" in name:
		lowerCAmelCase__      :		List[str]    							=      name.replace('''patch_embed.norm'''		,	'''embeddings.norm'''					)
	if "layers" in name:
		lowerCAmelCase__      :		List[Any]    							=      '''encoder.''' + name
	if "encoder.layers" in name:
		lowerCAmelCase__      :		List[Any]    							=      name.replace('''encoder.layers'''		,	'''encoder.stages'''					)
	if "downsample.proj" in name:
		lowerCAmelCase__      :		List[Any]    							=      name.replace('''downsample.proj'''		,	'''downsample.projection'''					)
	if "blocks" in name:
		lowerCAmelCase__      :		List[str]    							=      name.replace('''blocks'''		,	'''layers'''					)
	if "modulation.f.weight" in name or "modulation.f.bias" in name:
		lowerCAmelCase__      :		Optional[int]    							=      name.replace('''modulation.f'''		,	'''modulation.projection_in'''					)
	if "modulation.h.weight" in name or "modulation.h.bias" in name:
		lowerCAmelCase__      :		List[str]    							=      name.replace('''modulation.h'''		,	'''modulation.projection_context'''					)
	if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
		lowerCAmelCase__      :		Tuple    							=      name.replace('''modulation.proj'''		,	'''modulation.projection_out'''					)

	if name == "norm.weight":
		lowerCAmelCase__      :		Dict    							=      '''layernorm.weight'''
	if name == "norm.bias":
		lowerCAmelCase__      :		List[str]    							=      '''layernorm.bias'''

	if "head" in name:
		lowerCAmelCase__      :		Dict    							=      name.replace('''head'''		,	'''classifier'''					)
	else:
		lowerCAmelCase__      :		Dict    							=      '''focalnet.''' + name

	return name

def     __SCREAMING_SNAKE_CASE    ( A_		,	A_		,	A_=False					):
	# fmt: off
	lowerCAmelCase__      :		Any    							=      {
	    '''focalnet-tiny''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_srf.pth''',
	    '''focalnet-tiny-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_lrf.pth''',
	    '''focalnet-small''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_srf.pth''',
	    '''focalnet-small-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_lrf.pth''',
	    '''focalnet-base''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_srf.pth''',
	    '''focalnet-base-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_lrf.pth''',
	    '''focalnet-large-lrf-fl3''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384.pth''',
	    '''focalnet-large-lrf-fl4''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384_fl4.pth''',
	    '''focalnet-xlarge-lrf-fl3''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384.pth''',
	    '''focalnet-xlarge-lrf-fl4''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384_fl4.pth''',
	}
	# fmt: on

	lowerCAmelCase__      :		Optional[Any]    							=      model_name_to_url[model_name]
	print('''Checkpoint URL: '''		,	A_					)
	lowerCAmelCase__      :		Tuple    							=      torch.hub.load_state_dict_from_url(A_		,	map_location='''cpu'''					)['''model''']

	# rename keys
	for key in state_dict.copy().keys():
		lowerCAmelCase__      :		Optional[Any]    							=      state_dict.pop(A_					)
		lowerCAmelCase__      :		List[str]    							=      val

	lowerCAmelCase__      :		Tuple    							=      get_focalnet_config(A_					)
	lowerCAmelCase__      :		Union[str, Any]    							=      FocalNetForImageClassification(A_					)
	model.eval()

	# load state dict
	model.load_state_dict(A_					)

	# verify conversion
	lowerCAmelCase__      :		int    							=      '''http://images.cocodataset.org/val2017/000000039769.jpg'''

	lowerCAmelCase__      :		int    							=      BitImageProcessor(
	    do_resize=A_		,	size={'''shortest_edge''': 2_56}		,	resample=PILImageResampling.BILINEAR		,	do_center_crop=A_		,	crop_size=2_24		,	do_normalize=A_		,	image_mean=A_		,	image_std=A_		,	)
	lowerCAmelCase__      :		Optional[Any]    							=      Image.open(requests.get(A_		,	stream=A_					).raw					)
	lowerCAmelCase__      :		Dict    							=      processor(images=A_		,	return_tensors='''pt'''					)

	lowerCAmelCase__      :		List[str]    							=      transforms.Compose(
	    [
	        transforms.Resize(2_56					),
	        transforms.CenterCrop(2_24					),
	        transforms.ToTensor(),
	        transforms.Normalize(mean=[0.485, 0.456, 0.406]		,	std=[0.229, 0.224, 0.225]					),
	    ]					)

	lowerCAmelCase__      :		Optional[int]    							=      image_transforms(A_					).unsqueeze(0					)

	# verify pixel_values
	assert torch.allclose(inputs.pixel_values		,	A_		,	atol=1e-4					)

	lowerCAmelCase__      :		str    							=      model(**A_					)

	lowerCAmelCase__      :		Any    							=      outputs.logits.argmax(-1					).item()
	print('''Predicted class:'''		,	model.config.idalabel[predicted_class_idx]					)

	print('''First values of logits:'''		,	outputs.logits[0, :3]					)

	if model_name == "focalnet-tiny":
		lowerCAmelCase__      :		Optional[int]    							=      torch.tensor([0.2_166, -0.4_368, 0.2_191]					)
	elif model_name == "focalnet-tiny-lrf":
		lowerCAmelCase__      :		Optional[Any]    							=      torch.tensor([1.1_669, 0.0_125, -0.1_695]					)
	elif model_name == "focalnet-small":
		lowerCAmelCase__      :		Dict    							=      torch.tensor([0.4_917, -0.0_430, 0.1_341]					)
	elif model_name == "focalnet-small-lrf":
		lowerCAmelCase__      :		Union[str, Any]    							=      torch.tensor([-0.2_588, -0.5_342, -0.2_331]					)
	elif model_name == "focalnet-base":
		lowerCAmelCase__      :		List[Any]    							=      torch.tensor([-0.1_655, -0.4_090, -0.1_730]					)
	elif model_name == "focalnet-base-lrf":
		lowerCAmelCase__      :		int    							=      torch.tensor([0.5_306, -0.0_483, -0.3_928]					)
	assert torch.allclose(outputs.logits[0, :3]		,	A_		,	atol=1e-4					)
	print('''Looks ok!'''					)

	if pytorch_dump_folder_path is not None:
		print(f'Saving model and processor of {model_name} to {pytorch_dump_folder_path}'					)
		model.save_pretrained(A_					)
		processor.save_pretrained(A_					)

	if push_to_hub:
		print(f'Pushing model and processor of {model_name} to the hub...'					)
		model.push_to_hub(f'{model_name}'					)
		processor.push_to_hub(f'{model_name}'					)


if __name__ == "__main__":
							__UpperCamelCase    : List[str]				 =						argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '''--model_name''',
							    default='''focalnet-tiny''',
							    type=str,
							    help='''Name of the FocalNet 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.'''
							)
							parser.add_argument(
							    '''--push_to_hub''',
							    action='''store_true''',
							    help='''Whether to push the model and processor to the hub.''',
							)

							__UpperCamelCase    : int				 =						parser.parse_args()
							convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)







"""simple docstring"""





# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

from ..models.auto import AutoModelForVisionaSeq
from ..utils import requires_backends
from .base import PipelineTool


if TYPE_CHECKING:
							from PIL import Image







class 							SCREAMING_SNAKE_CASE   (						a_			):






					"""simple docstring"""



					lowercase__												=			"Salesforce/blip-image-captioning-base"
					lowercase__												=			(
					    "This is a tool that generates a description of an image. It takes an input named `image` which should be the "
					    "image to caption, and returns a text that contains the description in English."
					)
					lowercase__												=			"image_captioner"
					lowercase__												=			AutoModelForVisionaSeq

					lowercase__												=			["image"]
					lowercase__												=			["text"]







					def __init__(			self     :							int  ,*lowercase_     :							int  ,**lowercase_     :							List[str]						):
						requires_backends(self  ,['''vision''']						)
						super().__init__(*lowercase_  ,**lowercase_						)







					def 			__lowerCAmelCase					(			self     :							List[str]  ,lowercase_     :							"Image"						):
						return self.pre_processor(images=lowercase_  ,return_tensors='''pt'''						)







					def 			__lowerCAmelCase					(			self     :							List[Any]  ,lowercase_     :							Any						):
						return self.model.generate(**lowercase_						)







					def 			__lowerCAmelCase					(			self     :							Dict  ,lowercase_     :							Union[str, Any]						):
						return self.pre_processor.batch_decode(lowercase_  ,skip_special_tokens=lowercase_						)[0].strip()




