






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      _lowerCamelCase	(						)      ->					str:
    _UpperCAmelCase				:      List[str]			  =				'''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
    _UpperCAmelCase				:      Dict			  =				Image.open(requests.get(__A		,		stream=__A ).raw ).convert('''RGB''' )

    return image

def      _lowerCamelCase	(						__A    :							Any )      ->					Optional[int]:
    _UpperCAmelCase				:      Optional[int]			  =				[]
    # 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      _lowerCamelCase	(						__A    :							Any		,		__A    :							Optional[Any]		,		__A    :							int )      ->					List[Any]:
    _UpperCAmelCase				:      List[Any]			  =				dct.pop(__A )
    _UpperCAmelCase				:      Optional[int]			  =				val

def      _lowerCamelCase	(						__A    :							Optional[int]		,		__A    :							Tuple )      ->					Optional[Any]:
    for i in range(config.vision_config.num_hidden_layers ):
        # read in original q and v biases
        _UpperCAmelCase				:      int			  =				state_dict.pop(f'''visual_encoder.blocks.{i}.attn.q_bias''' )
        _UpperCAmelCase				:      List[str]			  =				state_dict.pop(f'''visual_encoder.blocks.{i}.attn.v_bias''' )

        # next, set bias in the state dict
        _UpperCAmelCase				:      str			  =				torch.cat((q_bias, torch.zeros_like(__A		,		requires_grad=__A ), v_bias) )
        _UpperCAmelCase				:      int			  =				qkv_bias

def      _lowerCamelCase	(						__A    :							Optional[Any]		,		__A    :							List[str] )      ->					Any:
    _UpperCAmelCase				:      Tuple			  =				364 if '''coco''' in model_name else 224
    _UpperCAmelCase				:      Any			  =				BlipaVisionConfig(image_size=__A ).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:
        _UpperCAmelCase				:      Tuple			  =				OPTConfig.from_pretrained('''facebook/opt-2.7b'''		,		eos_token_id=__A ).to_dict()
    elif "opt-6.7b" in model_name:
        _UpperCAmelCase				:      Optional[int]			  =				OPTConfig.from_pretrained('''facebook/opt-6.7b'''		,		eos_token_id=__A ).to_dict()
    elif "t5-xl" in model_name:
        _UpperCAmelCase				:      Dict			  =				TaConfig.from_pretrained('''google/flan-t5-xl'''		,		dense_act_fn='''gelu'''		,		bos_token_id=1 ).to_dict()
    elif "t5-xxl" in model_name:
        _UpperCAmelCase				:      Tuple			  =				TaConfig.from_pretrained('''google/flan-t5-xxl'''		,		dense_act_fn='''gelu'''		,		bos_token_id=1 ).to_dict()

    _UpperCAmelCase				:      str			  =				BlipaConfig(vision_config=__A		,		text_config=__A )

    return config, image_size

@torch.no_grad()
def      _lowerCamelCase	(						__A    :							Dict		,		__A    :							int=None		,		__A    :							Dict=False )      ->					Optional[int]:
    _UpperCAmelCase				:      Tuple			  =				(
        AutoTokenizer.from_pretrained('''facebook/opt-2.7b''' )
        if '''opt''' in model_name
        else AutoTokenizer.from_pretrained('''google/flan-t5-xl''' )
    )
    _UpperCAmelCase				:      Any			  =				tokenizer('''\n'''		,		add_special_tokens=__A ).input_ids[0]
    _UpperCAmelCase  , _UpperCAmelCase				:      Dict			  =				get_blipa_config(__A		,		eos_token_id=__A )

    _UpperCAmelCase				:      str			  =				BlipaForConditionalGeneration(__A ).eval()

    _UpperCAmelCase				:      Optional[int]			  =				{
        '''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'''),
    }

    _UpperCAmelCase  , _UpperCAmelCase				:      Tuple			  =				model_name_to_original[model_name]

    # load original model
    print('''Loading original model...''' )
    _UpperCAmelCase				:      Dict			  =				'''cuda''' if torch.cuda.is_available() else '''cpu'''
    _UpperCAmelCase  , _UpperCAmelCase  , _UpperCAmelCase				:      Any			  =				load_model_and_preprocess(
        name=__A		,		model_type=__A		,		is_eval=__A		,		device=__A )
    original_model.eval()
    print('''Done!''' )

    # update state dict keys
    _UpperCAmelCase				:      Any			  =				original_model.state_dict()
    _UpperCAmelCase				:      Tuple			  =				create_rename_keys(__A )
    for src, dest in rename_keys:
        rename_key(__A		,		__A		,		__A )

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

    # read in qv biases
    read_in_q_v_bias(__A		,		__A )

    _UpperCAmelCase  , _UpperCAmelCase				:      str			  =				hf_model.load_state_dict(__A		,		strict=__A )
    assert len(__A ) == 0
    assert unexpected_keys == ["qformer.embeddings.position_ids"]

    _UpperCAmelCase				:      str			  =				load_demo_image()
    _UpperCAmelCase				:      Optional[int]			  =				vis_processors['''eval'''](__A ).unsqueeze(0 ).to(__A )
    _UpperCAmelCase				:      Optional[int]			  =				tokenizer(['''\n''']		,		return_tensors='''pt''' ).input_ids.to(__A )

    # create processor
    _UpperCAmelCase				:      str			  =				BlipImageProcessor(
        size={'''height''': image_size, '''width''': image_size}		,		image_mean=__A		,		image_std=__A )
    _UpperCAmelCase				:      Dict			  =				BlipaProcessor(image_processor=__A		,		tokenizer=__A )
    _UpperCAmelCase				:      Tuple			  =				processor(images=__A		,		return_tensors='''pt''' ).pixel_values.to(__A )

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

    original_model.to(__A )
    hf_model.to(__A )
    with torch.no_grad():
        if "opt" in model_name:
            _UpperCAmelCase				:      str			  =				original_model({'''image''': original_pixel_values, '''text_input''': ['''''']} ).logits
            _UpperCAmelCase				:      Tuple			  =				hf_model(__A		,		__A ).logits
        else:
            _UpperCAmelCase				:      str			  =				original_model(
                {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']} ).logits
            _UpperCAmelCase				:      Dict			  =				input_ids.masked_fill(input_ids == tokenizer.pad_token_id		,		-100 )
            _UpperCAmelCase				:      Optional[int]			  =				hf_model(__A		,		__A		,		labels=__A ).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":
        _UpperCAmelCase				:      List[Any]			  =				torch.tensor(
            [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]		,		device=__A )
        assert torch.allclose(logits[0, :3, :3]		,		__A		,		atol=1e-4 )
    elif model_name == "blip2-flan-t5-xl-coco":
        _UpperCAmelCase				:      Optional[int]			  =				torch.tensor(
            [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]		,		device=__A )
    else:
        # cast to same type
        _UpperCAmelCase				:      Optional[int]			  =				logits.dtype
        assert torch.allclose(original_logits.to(__A )		,		__A		,		atol=1e-2 )
    print('''Looks ok!''' )

    print('''Generating a caption...''' )
    _UpperCAmelCase				:      Tuple			  =				''''''
    _UpperCAmelCase				:      Any			  =				tokenizer(__A		,		return_tensors='''pt''' ).input_ids.to(__A )

    _UpperCAmelCase				:      List[Any]			  =				original_model.generate({'''image''': original_pixel_values} )
    _UpperCAmelCase				:      Optional[Any]			  =				hf_model.generate(
        __A		,		__A		,		do_sample=__A		,		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:'''		,		__A )
    _UpperCAmelCase				:      List[str]			  =				input_ids.shape[1]
    _UpperCAmelCase				:      List[Any]			  =				processor.batch_decode(outputs[:, prompt_length:]		,		skip_special_tokens=__A )
    _UpperCAmelCase				:      List[Any]			  =				[text.strip() for text in output_text]
    print('''HF generation:'''		,		__A )

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

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


if __name__ == "__main__":
  SCREAMING_SNAKE_CASE 							=						argparse.ArgumentParser()
  SCREAMING_SNAKE_CASE 							=						[
      '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',
  )

  SCREAMING_SNAKE_CASE 							=						parser.parse_args()

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












from argparse import ArgumentParser, Namespace
from typing import Any, List, Optional

from ..pipelines import Pipeline, get_supported_tasks, pipeline
from ..utils import logging
from . import BaseTransformersCLICommand


try:
  from fastapi import Body, FastAPI, HTTPException
  from fastapi.routing import APIRoute
  from pydantic import BaseModel
  from starlette.responses import JSONResponse
  from uvicorn import run

  SCREAMING_SNAKE_CASE 							=						True
except (ImportError, AttributeError):
  SCREAMING_SNAKE_CASE 							=						object

  def      _lowerCamelCase	(						*__A    :							List[Any]		,		**__A    :							Any )      ->					Any:
      pass


  SCREAMING_SNAKE_CASE 							=						False


SCREAMING_SNAKE_CASE 							=						logging.get_logger('transformers-cli/serving')

def      _lowerCamelCase	(						__A    :							Namespace )      ->					Tuple:
    _UpperCAmelCase				:      Any			  =				pipeline(
        task=args.task		,		model=args.model if args.model else None		,		config=args.config		,		tokenizer=args.tokenizer		,		device=args.device		,		)
    return ServeCommand(__A		,		args.host		,		args.port		,		args.workers )
class 							A_    (						__lowercase			):






 '''simple docstring'''




 _SCREAMING_SNAKE_CASE							:			dict



class 							A_    (						__lowercase			):






 '''simple docstring'''




 _SCREAMING_SNAKE_CASE							:			List[str]
 _SCREAMING_SNAKE_CASE							:			Optional[List[int]]



class 							A_    (						__lowercase			):






 '''simple docstring'''




 _SCREAMING_SNAKE_CASE							:			str



class 							A_    (						__lowercase			):






 '''simple docstring'''




 _SCREAMING_SNAKE_CASE							:			Any



class 							A_    (						__lowercase			):






 '''simple docstring'''




 @staticmethod
 def  snake_case__				(	_A) ->		List[Any]:

     """simple docstring"""



     _UpperCAmelCase				:      Union[str, Any]			  =				parser.add_parser(
         '''serve'''							,  help='''CLI tool to run inference requests through REST and GraphQL endpoints.''')
     serve_parser.add_argument(
         '''--task'''							,  type=_A							,  choices=get_supported_tasks()							,  help='''The task to run the pipeline on'''							,  )
     serve_parser.add_argument('''--host'''							,  type=_A							,  default='''localhost'''							,  help='''Interface the server will listen on.''')
     serve_parser.add_argument('''--port'''							,  type=_A							,  default=8888							,  help='''Port the serving will listen to.''')
     serve_parser.add_argument('''--workers'''							,  type=_A							,  default=1							,  help='''Number of http workers''')
     serve_parser.add_argument('''--model'''							,  type=_A							,  help='''Model\'s name or path to stored model.''')
     serve_parser.add_argument('''--config'''							,  type=_A							,  help='''Model\'s config name or path to stored model.''')
     serve_parser.add_argument('''--tokenizer'''							,  type=_A							,  help='''Tokenizer name to use.''')
     serve_parser.add_argument(
         '''--device'''							,  type=_A							,  default=-1							,  help='''Indicate the device to run onto, -1 indicates CPU, >= 0 indicates GPU (default: -1)'''							,  )
     serve_parser.set_defaults(func=_A)


 def __init__(	self							,  _A							,  _A							,  _A							,  _A) ->		List[Any]:

     """simple docstring"""



     _UpperCAmelCase				:      List[Any]			  =				pipeline

     _UpperCAmelCase				:      List[Any]			  =				host
     _UpperCAmelCase				:      Tuple			  =				port
     _UpperCAmelCase				:      Union[str, Any]			  =				workers

     if not _serve_dependencies_installed:
         raise RuntimeError(
             '''Using serve command requires FastAPI and uvicorn. '''
             '''Please install transformers with [serving]: pip install "transformers[serving]".'''
             '''Or install FastAPI and uvicorn separately.''')
     else:
         logger.info(f'''Serving model over {host}:{port}''')
         _UpperCAmelCase				:      str			  =				FastAPI(
             routes=[
                 APIRoute(
                     '''/'''							,  self.model_info							,  response_model=_A							,  response_class=_A							,  methods=['''GET''']							,  ),
                 APIRoute(
                     '''/tokenize'''							,  self.tokenize							,  response_model=_A							,  response_class=_A							,  methods=['''POST''']							,  ),
                 APIRoute(
                     '''/detokenize'''							,  self.detokenize							,  response_model=_A							,  response_class=_A							,  methods=['''POST''']							,  ),
                 APIRoute(
                     '''/forward'''							,  self.forward							,  response_model=_A							,  response_class=_A							,  methods=['''POST''']							,  ),
             ]							,  timeout=600							,  )


 def  snake_case__				(	self) ->		Optional[Any]:

     """simple docstring"""



     run(self._app							,  host=self.host							,  port=self.port							,  workers=self.workers)


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

     """simple docstring"""



     return ServeModelInfoResult(infos=vars(self._pipeline.model.config))


 def  snake_case__				(	self							,  _A = Body(_A							,  embed=_A)							,  _A = Body(_A							,  embed=_A)) ->		List[Any]:

     """simple docstring"""



     try:
         _UpperCAmelCase				:      int			  =				self._pipeline.tokenizer.tokenize(_A)

         if return_ids:
             _UpperCAmelCase				:      List[Any]			  =				self._pipeline.tokenizer.convert_tokens_to_ids(_A)
             return ServeTokenizeResult(tokens=_A							,  tokens_ids=_A)
         else:
             return ServeTokenizeResult(tokens=_A)

     except Exception as e:
         raise HTTPException(status_code=500							,  detail={'''model''': '''''', '''error''': str(_A)})


 def  snake_case__				(	self							,  _A = Body(_A							,  embed=_A)							,  _A = Body(_A							,  embed=_A)							,  _A = Body(_A							,  embed=_A)							,  ) ->		str:

     """simple docstring"""



     try:
         _UpperCAmelCase				:      Optional[Any]			  =				self._pipeline.tokenizer.decode(_A							,  _A							,  _A)
         return ServeDeTokenizeResult(model=''''''							,  text=_A)
     except Exception as e:
         raise HTTPException(status_code=500							,  detail={'''model''': '''''', '''error''': str(_A)})


 async def  snake_case__				(	self							,  _A=Body(_A							,  embed=_A)) ->		str:

     """simple docstring"""



     if len(_A) == 0:
         return ServeForwardResult(output=[]							,  attention=[])

     try:
         # Forward through the model
         _UpperCAmelCase				:      Any			  =				self._pipeline(_A)
         return ServeForwardResult(output=_A)
     except Exception as e:
         raise HTTPException(500							,  {'''error''': str(_A)})




