


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_float32 (there's also the fix_lavis branch)
# also note: to convert Vicuna checkpoints, we had to include /home/niels/python_projects/checkpoints/FastChat/vicuna-7b in lavis/configs/models/blip2/blip2_instruct_vicuna7b.yaml
# same for Vicuna-13b
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipImageProcessor,
    InstructBlipConfig,
    InstructBlipForConditionalGeneration,
    InstructBlipProcessor,
    InstructBlipQFormerConfig,
    InstructBlipVisionConfig,
    LlamaConfig,
    LlamaTokenizerFast,
    TaConfig,
    TaTokenizerFast,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD






def 							__lowercase     (				):
							a__    =  'https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'
							a__    =  Image.open(requests.get(__lowerCAmelCase      ,					stream=__lowerCAmelCase     ).raw     ).convert('RGB'     )

							return image






def 							__lowercase     (				__lowerCAmelCase							:       Optional[int]     ):
							a__    =  []
							# 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.embeddings.layernorm.weight')     )
							rename_keys.append(('Qformer.bert.embeddings.LayerNorm.bias', 'qformer.embeddings.layernorm.bias')     )

							# fmt: on
							return rename_keys






def 							__lowercase     (				__lowerCAmelCase							:       Tuple      ,					__lowerCAmelCase							:       Dict      ,					__lowerCAmelCase							:       List[Any]     ):
							a__    =  dct.pop(__lowerCAmelCase     )
							a__    =  val






def 							__lowercase     (				__lowerCAmelCase							:       int      ,					__lowerCAmelCase							:       Union[str, Any]     ):
							for i in range(config.vision_config.num_hidden_layers     ):
														# read in original q and v biases
														a__    =  state_dict.pop(F'visual_encoder.blocks.{i}.attn.q_bias'     )
														a__    =  state_dict.pop(F'visual_encoder.blocks.{i}.attn.v_bias'     )

														# next, set bias in the state dict
														a__    =  torch.cat((q_bias, torch.zeros_like(__lowerCAmelCase      ,					requires_grad=__lowerCAmelCase     ), v_bias)     )
														a__    =  qkv_bias






def 							__lowercase     (				__lowerCAmelCase							:       Tuple     ):
							a__    =  3_6_4 if 'coco' in model_name else 2_2_4
							a__    =  InstructBlipVisionConfig(image_size=__lowerCAmelCase     ).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 "t5-xl" in model_name:
														a__    =  TaConfig.from_pretrained('google/flan-t5-xl'      ,					dense_act_fn='gelu'      ,					bos_token_id=1     ).to_dict()
							elif "t5-xxl" in model_name:
														a__    =  TaConfig.from_pretrained('google/flan-t5-xxl'      ,					dense_act_fn='gelu'      ,					bos_token_id=1     ).to_dict()
							elif "vicuna-7b" in model_name:
														a__    =  LlamaConfig.from_pretrained('decapoda-research/llama-7b-hf'      ,					vocab_size=3_2_0_0_1     ).to_dict()
							elif "vicuna-13b" in model_name:
														a__    =  LlamaConfig.from_pretrained('decapoda-research/llama-13b-hf'      ,					vocab_size=3_2_0_0_1     ).to_dict()
							else:
														raise ValueError('Model name not supported'     )

							# the authors add one special "[DEC]" token to the vocab of Q-Former, hence vocab size = 30522 + 1
							a__    =  InstructBlipQFormerConfig(vocab_size=3_0_5_2_3     ).to_dict()
							a__    =  InstructBlipConfig(vision_config=__lowerCAmelCase      ,					text_config=__lowerCAmelCase      ,					qformer_config=__lowerCAmelCase     )

							return config, image_size






@torch.no_grad()
def 							__lowercase     (				__lowerCAmelCase							:       Optional[Any]      ,					__lowerCAmelCase							:       Tuple=None      ,					__lowerCAmelCase							:       Optional[Any]=False     ):
							a__    =  AutoTokenizer.from_pretrained('bert-base-uncased'      ,					truncation_side='left'     )
							qformer_tokenizer.add_special_tokens({'bos_token': '[DEC]'}     )

							if "t5" in model_name:
														a__    =  TaTokenizerFast.from_pretrained('google/flan-t5-xl'      ,					truncation_side='left'     )
							elif "vicuna" in model_name:
														# the following was used in the original implementation:
														# tokenizer = LlamaTokenizer.from_pretrained("huggyllama/llama-7b", use_fast=False, truncation_side="left")
														# tokenizer.add_special_tokens({"pad_token": "[PAD]"})
														# tokenizer.add_special_tokens({"bos_token": "</s>"})
														# tokenizer.add_special_tokens({"eos_token": "</s>"})
														# tokenizer.add_special_tokens({"unk_token": "</s>"})
														a__    =  LlamaTokenizerFast.from_pretrained(
														    'huggyllama/llama-7b'      ,					truncation_side='left'      ,					bos_token='</s>'      ,					unk_token='</s>'     )
														tokenizer.add_special_tokens({'pad_token': '[PAD]'}     )

							a__   ,      a__    =  get_blipa_config(__lowerCAmelCase     )
							a__    =  InstructBlipForConditionalGeneration(__lowerCAmelCase     ).eval()

							a__    =  {
							    'instructblip-vicuna-7b': ('blip2_vicuna_instruct', 'vicuna7b'),
							    'instructblip-vicuna-13b': ('blip2_vicuna_instruct', 'vicuna13b'),
							    'instructblip-flan-t5-xl': ('blip2_t5_instruct', 'flant5xl'),
							    'instructblip-flan-t5-xxl': ('blip2_t5_instruct', 'flant5xxl'),
							}

							a__   ,      a__    =  model_name_to_original[model_name]

							# load original model
							print('Loading original model...'     )
							a__    =  'cuda:1' if torch.cuda.is_available() else 'cpu'
							a__    =  'cuda:2' if torch.cuda.is_available() else 'cpu'
							a__   ,      a__   ,      a__    =  load_model_and_preprocess(
							    name=__lowerCAmelCase      ,					model_type=__lowerCAmelCase      ,					is_eval=__lowerCAmelCase      ,					device=__lowerCAmelCase     )
							original_model.eval()
							print('Done!'     )

							# update state dict keys
							a__    =  original_model.state_dict()
							a__    =  create_rename_keys(__lowerCAmelCase     )
							for src, dest in rename_keys:
														rename_key(__lowerCAmelCase      ,					__lowerCAmelCase      ,					__lowerCAmelCase     )

							# some keys can be renamed efficiently
							for key, val in state_dict.copy().items():
														a__    =  state_dict.pop(__lowerCAmelCase     )
														if key.startswith('Qformer.bert'     ):
																					a__    =  key.replace('Qformer.bert'      ,					'qformer'     )
														if "attention.self" in key:
																					a__    =  key.replace('self'      ,					'attention'     )
														if "llm_proj" in key:
																					a__    =  key.replace('llm_proj'      ,					'language_projection'     )
														if "t5_proj" in key:
																					a__    =  key.replace('t5_proj'      ,					'language_projection'     )
														if key.startswith('llm_model'     ):
																					a__    =  key.replace('llm_model'      ,					'language_model'     )
														if key.startswith('t5'     ):
																					a__    =  key.replace('t5'      ,					'language'     )
														a__    =  val

							# read in qv biases
							read_in_q_v_bias(__lowerCAmelCase      ,					__lowerCAmelCase     )

							# note: weights get loaded in torch.float32 by default
							hf_model.load_state_dict(__lowerCAmelCase      ,					strict=__lowerCAmelCase     )

							a__    =  load_demo_image()
							a__    =  'What is unusual about this image?'

							# create processor
							a__    =  BlipImageProcessor(
							    size={'height': image_size, 'width': image_size}      ,					image_mean=__lowerCAmelCase      ,					image_std=__lowerCAmelCase     )
							a__    =  InstructBlipProcessor(
							    image_processor=__lowerCAmelCase      ,					tokenizer=__lowerCAmelCase      ,					qformer_tokenizer=__lowerCAmelCase      ,					)
							a__    =  processor(images=__lowerCAmelCase      ,					text=__lowerCAmelCase      ,					return_tensors='pt'     ).to(__lowerCAmelCase     )

							# make sure processor creates exact same pixel values
							a__    =  vis_processors['eval'](__lowerCAmelCase     ).unsqueeze(0     ).to(__lowerCAmelCase     )
							a__    =  inputs.pixel_values
							assert torch.allclose(original_pixel_values.to(pixel_values.device     )      ,					__lowerCAmelCase     )

							original_model.to(__lowerCAmelCase     )
							hf_model.to(__lowerCAmelCase     )
							with torch.no_grad():
														if "vicuna" in model_name:
																					a__    =  original_model({'image': original_pixel_values, 'text_input': [prompt]}     ).logits
																					a__    =  hf_model(**__lowerCAmelCase     ).logits
														else:
																					a__    =  original_model(
																					    {'image': original_pixel_values, 'text_input': [prompt], 'text_output': ['\n']}     ).logits
																					a__    =  tokenizer('\n'      ,					return_tensors='pt'     ).input_ids.to(__lowerCAmelCase     )
																					a__    =  label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id      ,					-1_0_0     )
																					a__    =  hf_model(**__lowerCAmelCase      ,					labels=__lowerCAmelCase     ).logits

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

							# assert values
							assert original_logits.shape == logits.shape
							a__    =  1E-4 if 'vicuna' in model_name else 1E-5
							assert torch.allclose(original_logits.to(logits.device     )      ,					__lowerCAmelCase      ,					atol=__lowerCAmelCase     )
							print('Looks ok!'     )

							print('Generating with original model...'     )
							a__    =  original_model.generate({'image': original_pixel_values, 'prompt': prompt}      ,					num_beams=5     )

							# important: we need to cast the weights of the HF model to the appropriate type
							print('Generating with HF model...'     )
							a__    =  hf_model.generate(
							    **__lowerCAmelCase      ,					do_sample=__lowerCAmelCase      ,					num_beams=5      ,					max_length=2_5_6      ,					min_length=1      ,					top_p=0.9      ,					repetition_penalty=1.5      ,					length_penalty=1.0      ,					temperature=1      ,					)
							if "vicuna" in model_name:
														# convert output id 0 to 2 (eos_token_id)
														# TODO add this in the generate method?
														a__    =  2
							print('Original generation:'      ,					__lowerCAmelCase     )
							a__    =  processor.batch_decode(__lowerCAmelCase      ,					skip_special_tokens=__lowerCAmelCase     )
							a__    =  [text.strip() for text in output_text]
							print('HF generation:'      ,					__lowerCAmelCase     )

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

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


if __name__ == "__main__":
						snake_case							: Tuple      =   argparse.ArgumentParser()
						snake_case							: Optional[int]      =   [
						    '''instructblip-vicuna-7b''',
						    '''instructblip-vicuna-13b''',
						    '''instructblip-flan-t5-xl''',
						    '''instructblip-flan-t5-xxl''',
						]
						parser.add_argument(
						    '''--model_name''',
						    default='''instructblip-flan-t5-xl''',
						    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''',
						)

						snake_case							: Any      =   parser.parse_args()

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








from decimal import Decimal, getcontext
from math import ceil, factorial






def 							__lowercase     (				__lowerCAmelCase							:       int     ):

							if not isinstance(__lowerCAmelCase      ,					__lowerCAmelCase     ):
														raise TypeError('Undefined for non-integers'     )
							elif precision < 1:
														raise ValueError('Undefined for non-natural numbers'     )

							a__    =  precision
							a__    =  ceil(precision / 1_4     )
							a__    =  4_2_6_8_8_0 * Decimal(1_0_0_0_5     ).sqrt()
							a__    =  1
							a__    =  1_3_5_9_1_4_0_9
							a__    =  Decimal(__lowerCAmelCase     )
							for k in range(1      ,					__lowerCAmelCase     ):
														a__    =  factorial(6 * k     ) // (factorial(3 * k     ) * factorial(__lowerCAmelCase     ) ** 3)
														linear_term += 5_4_5_1_4_0_1_3_4
														exponential_term *= -2_6_2_5_3_7_4_1_2_6_4_0_7_6_8_0_0_0
														partial_sum += Decimal(multinomial_term * linear_term     ) / exponential_term
							return str(constant_term / partial_sum     )[:-1]


if __name__ == "__main__":
						snake_case							: Tuple      =   50
						print(f"""The first {n} digits of pi is: {pi(n)}""")




