
'''simple docstring'''






import argparse

import requests
import torch

# pip3 install salesforce-lavis
# I'm actually installing a slightly modified version: pip3 install git+https://github.com/nielsrogge/LAVIS.git@fix_lavis
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipaConfig,
    BlipaForConditionalGeneration,
    BlipaProcessor,
    BlipaVisionConfig,
    BlipImageProcessor,
    OPTConfig,
    TaConfig,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD





def    UpperCAmelCase							(  ):
				lowerCamelCase					:			Dict									= 'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
				lowerCamelCase					:			Tuple									= Image.open(requests.get(_lowercase       ,							stream=_lowercase).raw).convert('RGB')

				return image





def    UpperCAmelCase							(  UpperCAmelCase__							:							Tuple):
				lowerCamelCase					:			str									= []
				# fmt: off

				# vision encoder
				rename_keys.append(('visual_encoder.cls_token', 'vision_model.embeddings.class_embedding'))
				rename_keys.append(('visual_encoder.pos_embed', 'vision_model.embeddings.position_embedding'))
				rename_keys.append(('visual_encoder.patch_embed.proj.weight', 'vision_model.embeddings.patch_embedding.weight'))
				rename_keys.append(('visual_encoder.patch_embed.proj.bias', 'vision_model.embeddings.patch_embedding.bias'))
				rename_keys.append(('ln_vision.weight', 'vision_model.post_layernorm.weight'))
				rename_keys.append(('ln_vision.bias', 'vision_model.post_layernorm.bias'))

				for i in range(config.vision_config.num_hidden_layers):
								rename_keys.append((F'''visual_encoder.blocks.{i}.norm1.weight''', F'''vision_model.encoder.layers.{i}.layer_norm1.weight'''))
								rename_keys.append((F'''visual_encoder.blocks.{i}.norm1.bias''', F'''vision_model.encoder.layers.{i}.layer_norm1.bias'''))
								rename_keys.append((F'''visual_encoder.blocks.{i}.norm2.weight''', F'''vision_model.encoder.layers.{i}.layer_norm2.weight'''))
								rename_keys.append((F'''visual_encoder.blocks.{i}.norm2.bias''', F'''vision_model.encoder.layers.{i}.layer_norm2.bias'''))
								rename_keys.append((F'''visual_encoder.blocks.{i}.attn.qkv.weight''', F'''vision_model.encoder.layers.{i}.self_attn.qkv.weight'''))
								rename_keys.append((F'''visual_encoder.blocks.{i}.attn.proj.weight''', F'''vision_model.encoder.layers.{i}.self_attn.projection.weight''',))
								rename_keys.append((F'''visual_encoder.blocks.{i}.attn.proj.bias''', F'''vision_model.encoder.layers.{i}.self_attn.projection.bias'''))
								rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc1.weight''', F'''vision_model.encoder.layers.{i}.mlp.fc1.weight'''))
								rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc1.bias''', F'''vision_model.encoder.layers.{i}.mlp.fc1.bias'''))
								rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc2.weight''', F'''vision_model.encoder.layers.{i}.mlp.fc2.weight'''))
								rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc2.bias''', F'''vision_model.encoder.layers.{i}.mlp.fc2.bias'''))

				# QFormer
				rename_keys.append(('Qformer.bert.embeddings.LayerNorm.weight', 'qformer.layernorm.weight'))
				rename_keys.append(('Qformer.bert.embeddings.LayerNorm.bias', 'qformer.layernorm.bias'))

				# fmt: on
				return rename_keys





def    UpperCAmelCase							(  UpperCAmelCase__							:							List[str]       ,							UpperCAmelCase__							:							Dict       ,							UpperCAmelCase__							:							Union[str, Any]):
				lowerCamelCase					:			Any									= dct.pop(_lowercase)
				lowerCamelCase					:			Dict									= val





def    UpperCAmelCase							(  UpperCAmelCase__							:							Optional[Any]       ,							UpperCAmelCase__							:							Dict):







				for i in range(config.vision_config.num_hidden_layers):
								# read in original q and v biases
								lowerCamelCase					:			Dict									= state_dict.pop(F'''visual_encoder.blocks.{i}.attn.q_bias''')
								lowerCamelCase					:			str									= state_dict.pop(F'''visual_encoder.blocks.{i}.attn.v_bias''')

								# next, set bias in the state dict
								lowerCamelCase					:			List[Any]									= torch.cat((q_bias, torch.zeros_like(_lowercase       ,							requires_grad=_lowercase), v_bias))
								lowerCamelCase					:			Any									= qkv_bias





def    UpperCAmelCase							(  UpperCAmelCase__							:							Tuple       ,							UpperCAmelCase__							:							Union[str, Any]):
				lowerCamelCase					:			Tuple									= 3_64 if 'coco' in model_name else 2_24
				lowerCamelCase					:			Union[str, Any]									= BlipaVisionConfig(image_size=_lowercase).to_dict()

				# make sure the models have proper bos_token_id and eos_token_id set (important for generation)
				# seems like flan-T5 models don't have bos_token_id properly set?
				if "opt-2.7b" in model_name:
								lowerCamelCase					:			Tuple									= OPTConfig.from_pretrained('facebook/opt-2.7b'       ,							eos_token_id=_lowercase).to_dict()
				elif "opt-6.7b" in model_name:
								lowerCamelCase					:			Any									= OPTConfig.from_pretrained('facebook/opt-6.7b'       ,							eos_token_id=_lowercase).to_dict()
				elif "t5-xl" in model_name:
								lowerCamelCase					:			Dict									= TaConfig.from_pretrained('google/flan-t5-xl'       ,							dense_act_fn='gelu'       ,							bos_token_id=1).to_dict()
				elif "t5-xxl" in model_name:
								lowerCamelCase					:			Union[str, Any]									= TaConfig.from_pretrained('google/flan-t5-xxl'       ,							dense_act_fn='gelu'       ,							bos_token_id=1).to_dict()

				lowerCamelCase					:			Tuple									= BlipaConfig(vision_config=_lowercase       ,							text_config=_lowercase)

				return config, image_size





@torch.no_grad()
def    UpperCAmelCase							(  UpperCAmelCase__							:							Tuple       ,							UpperCAmelCase__							:							List[str]=None       ,							UpperCAmelCase__							:							Tuple=False):
				lowerCamelCase					:			Optional[Any]									= (
				    AutoTokenizer.from_pretrained('facebook/opt-2.7b')
				    if 'opt' in model_name
				    else AutoTokenizer.from_pretrained('google/flan-t5-xl')
				)
				lowerCamelCase					:			Union[str, Any]									= tokenizer('\n'       ,							add_special_tokens=_lowercase).input_ids[0]
				lowerCamelCase			, lowerCamelCase					:			List[str]									= get_blipa_config(_lowercase       ,							eos_token_id=_lowercase)

				lowerCamelCase					:			Union[str, Any]									= BlipaForConditionalGeneration(_lowercase).eval()

				lowerCamelCase					:			Union[str, Any]									= {
				    'blip2-opt-2.7b': ('blip2_opt', 'pretrain_opt2.7b'),
				    'blip2-opt-6.7b': ('blip2_opt', 'pretrain_opt6.7b'),
				    'blip2-opt-2.7b-coco': ('blip2_opt', 'caption_coco_opt2.7b'),
				    'blip2-opt-6.7b-coco': ('blip2_opt', 'caption_coco_opt6.7b'),
				    'blip2-flan-t5-xl': ('blip2_t5', 'pretrain_flant5xl'),
				    'blip2-flan-t5-xl-coco': ('blip2_t5', 'caption_coco_flant5xl'),
				    'blip2-flan-t5-xxl': ('blip2_t5', 'pretrain_flant5xxl'),
				}

				lowerCamelCase			, lowerCamelCase					:			Any									= model_name_to_original[model_name]

				# load original model
				print('Loading original model...')
				lowerCamelCase					:			Any									= 'cuda' if torch.cuda.is_available() else 'cpu'
				lowerCamelCase			, lowerCamelCase			, lowerCamelCase					:			Optional[int]									= load_model_and_preprocess(
				    name=_lowercase       ,							model_type=_lowercase       ,							is_eval=_lowercase       ,							device=_lowercase)
				original_model.eval()
				print('Done!')

				# update state dict keys
				lowerCamelCase					:			str									= original_model.state_dict()
				lowerCamelCase					:			Dict									= create_rename_keys(_lowercase)
				for src, dest in rename_keys:
								rename_key(_lowercase       ,							_lowercase       ,							_lowercase)

				# some keys can be renamed efficiently
				for key, val in state_dict.copy().items():
								lowerCamelCase					:			Tuple									= state_dict.pop(_lowercase)
								if key.startswith('Qformer.bert'):
												lowerCamelCase					:			Tuple									= key.replace('Qformer.bert'       ,							'qformer')
								if "attention.self" in key:
												lowerCamelCase					:			Union[str, Any]									= key.replace('self'       ,							'attention')
								if "opt_proj" in key:
												lowerCamelCase					:			int									= key.replace('opt_proj'       ,							'language_projection')
								if "t5_proj" in key:
												lowerCamelCase					:			int									= key.replace('t5_proj'       ,							'language_projection')
								if key.startswith('opt'):
												lowerCamelCase					:			Any									= key.replace('opt'       ,							'language')
								if key.startswith('t5'):
												lowerCamelCase					:			str									= key.replace('t5'       ,							'language')
								lowerCamelCase					:			List[Any]									= val

				# read in qv biases
				read_in_q_v_bias(_lowercase       ,							_lowercase)

				lowerCamelCase			, lowerCamelCase					:			Optional[int]									= hf_model.load_state_dict(_lowercase       ,							strict=_lowercase)
				assert len(_lowercase) == 0
				assert unexpected_keys == ["qformer.embeddings.position_ids"]

				lowerCamelCase					:			Optional[Any]									= load_demo_image()
				lowerCamelCase					:			Dict									= vis_processors['eval'](_lowercase).unsqueeze(0).to(_lowercase)
				lowerCamelCase					:			Optional[int]									= tokenizer(['\n']       ,							return_tensors='pt').input_ids.to(_lowercase)

				# create processor
				lowerCamelCase					:			Optional[int]									= BlipImageProcessor(
				    size={'height': image_size, 'width': image_size}       ,							image_mean=_lowercase       ,							image_std=_lowercase)
				lowerCamelCase					:			Optional[int]									= BlipaProcessor(image_processor=_lowercase       ,							tokenizer=_lowercase)
				lowerCamelCase					:			Optional[int]									= processor(images=_lowercase       ,							return_tensors='pt').pixel_values.to(_lowercase)

				# make sure processor creates exact same pixel values
				assert torch.allclose(_lowercase       ,							_lowercase)

				original_model.to(_lowercase)
				hf_model.to(_lowercase)
				with torch.no_grad():
								if "opt" in model_name:
												lowerCamelCase					:			Optional[int]									= original_model({'image': original_pixel_values, 'text_input': ['']}).logits
												lowerCamelCase					:			Optional[Any]									= hf_model(_lowercase       ,							_lowercase).logits
								else:
												lowerCamelCase					:			Union[str, Any]									= original_model(
												    {'image': original_pixel_values, 'text_input': ['\n'], 'text_output': ['\n']}).logits
												lowerCamelCase					:			int									= input_ids.masked_fill(input_ids == tokenizer.pad_token_id       ,							-1_00)
												lowerCamelCase					:			int									= hf_model(_lowercase       ,							_lowercase       ,							labels=_lowercase).logits

				assert original_logits.shape == logits.shape
				print('First values of original logits:'       ,							original_logits[0, :3, :3])
				print('First values of HF logits:'       ,							logits[0, :3, :3])

				# assert values
				if model_name == "blip2-flan-t5-xl":
								lowerCamelCase					:			Any									= torch.tensor(
								    [[-4_1.5_8_5_0, -4.4_4_4_0, -8.9_9_2_2], [-4_7.4_3_2_2, -5.9_1_4_3, -1.7_3_4_0]]       ,							device=_lowercase)
								assert torch.allclose(logits[0, :3, :3]       ,							_lowercase       ,							atol=1E-4)
				elif model_name == "blip2-flan-t5-xl-coco":
								lowerCamelCase					:			Any									= torch.tensor(
								    [[-5_7.0_1_0_9, -9.8_9_6_7, -1_2.6_2_8_0], [-6_8.6_5_7_8, -1_2.7_1_9_1, -1_0.5_0_6_5]]       ,							device=_lowercase)
				else:
								# cast to same type
								lowerCamelCase					:			Tuple									= logits.dtype
								assert torch.allclose(original_logits.to(_lowercase)       ,							_lowercase       ,							atol=1E-2)
				print('Looks ok!')

				print('Generating a caption...')
				lowerCamelCase					:			Any									= ''
				lowerCamelCase					:			Union[str, Any]									= tokenizer(_lowercase       ,							return_tensors='pt').input_ids.to(_lowercase)

				lowerCamelCase					:			str									= original_model.generate({'image': original_pixel_values})
				lowerCamelCase					:			Optional[int]									= hf_model.generate(
				    _lowercase       ,							_lowercase       ,							do_sample=_lowercase       ,							num_beams=5       ,							max_length=30       ,							min_length=1       ,							top_p=0.9       ,							repetition_penalty=1.0       ,							length_penalty=1.0       ,							temperature=1       ,							)
				print('Original generation:'       ,							_lowercase)
				lowerCamelCase					:			int									= input_ids.shape[1]
				lowerCamelCase					:			Dict									= processor.batch_decode(outputs[:, prompt_length:]       ,							skip_special_tokens=_lowercase)
				lowerCamelCase					:			int									= [text.strip() for text in output_text]
				print('HF generation:'       ,							_lowercase)

				if pytorch_dump_folder_path is not None:
								processor.save_pretrained(_lowercase)
								hf_model.save_pretrained(_lowercase)

				if push_to_hub:
								processor.push_to_hub(F'''nielsr/{model_name}''')
								hf_model.push_to_hub(F'''nielsr/{model_name}''')


if __name__ == "__main__":
					A			      =      argparse.ArgumentParser()
					A			      =      [
					    'blip2-opt-2.7b',
					    'blip2-opt-6.7b',
					    'blip2-opt-2.7b-coco',
					    'blip2-opt-6.7b-coco',
					    'blip2-flan-t5-xl',
					    'blip2-flan-t5-xl-coco',
					    'blip2-flan-t5-xxl',
					]
					parser.add_argument(
					    '--model_name',
					    default='blip2-opt-2.7b',
					    choices=choices,
					    type=str,
					    help='Path to hf config.json of model to convert',
					)
					parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
					parser.add_argument(
					    '--push_to_hub',
					    action='store_true',
					    help='Whether to push the model and processor to the hub after converting',
					)

					A			      =      parser.parse_args()

					convert_blipa_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)


'''simple docstring'''






import warnings
from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


A			      =      logging.get_logger(__name__)

A			      =      {
    'nvidia/segformer-b0-finetuned-ade-512-512': (
        'https://huggingface.co/nvidia/segformer-b0-finetuned-ade-512-512/resolve/main/config.json'
    ),
    # See all SegFormer models at https://huggingface.co/models?filter=segformer
}




class     __snake_case		(   a__):
					_lowerCAmelCase			     =	'''segformer'''
					def __init__(     self,    A=3,    A=4,    A=[2, 2, 2, 2],    A=[8, 4, 2, 1],    A=[32, 64, 160, 256],    A=[7, 3, 3, 3],    A=[4, 2, 2, 2],    A=[1, 2, 5, 8],    A=[4, 4, 4, 4],    A="gelu",    A=0.0,    A=0.0,    A=0.1,    A=0.02,    A=0.1,    A=1e-6,    A=256,    A=255,    **A,    ):

									"""simple docstring"""


									super().__init__(**A      )

									if "reshape_last_stage" in kwargs and kwargs["reshape_last_stage"] is False:
													warnings.warn(
													    'Reshape_last_stage is set to False in this config. This argument is deprecated and will soon be'
													    ' removed, as the behaviour will default to that of reshape_last_stage = True.',    A,    )

									lowerCamelCase					:			Tuple									= num_channels
									lowerCamelCase					:			Dict									= num_encoder_blocks
									lowerCamelCase					:			Dict									= depths
									lowerCamelCase					:			Any									= sr_ratios
									lowerCamelCase					:			List[Any]									= hidden_sizes
									lowerCamelCase					:			List[Any]									= patch_sizes
									lowerCamelCase					:			Tuple									= strides
									lowerCamelCase					:			int									= mlp_ratios
									lowerCamelCase					:			Dict									= num_attention_heads
									lowerCamelCase					:			Union[str, Any]									= hidden_act
									lowerCamelCase					:			Tuple									= hidden_dropout_prob
									lowerCamelCase					:			Optional[int]									= attention_probs_dropout_prob
									lowerCamelCase					:			List[str]									= classifier_dropout_prob
									lowerCamelCase					:			Union[str, Any]									= initializer_range
									lowerCamelCase					:			Union[str, Any]									= drop_path_rate
									lowerCamelCase					:			int									= layer_norm_eps
									lowerCamelCase					:			Any									= decoder_hidden_size
									lowerCamelCase					:			Optional[int]									= kwargs.get('reshape_last_stage',    A      )
									lowerCamelCase					:			Any									= semantic_loss_ignore_index





class     __snake_case		(   a__):
					_lowerCAmelCase			     =	version.parse('''1.11''')
					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									return OrderedDict(
									    [
									        ('pixel_values', {0: 'batch', 1: 'num_channels', 2: 'height', 3: 'width'}),
									    ]      )
					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									return 1e-4




					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									return 12
