





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

     return image






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






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

          # next, set bias in the state dict
          __A						:			Any								=	torch.cat((q_bias, torch.zeros_like(__snake_case			,							requires_grad=__snake_case      ), v_bias)      )
          __A						:			Dict								=	qkv_bias






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

     __A						:			List[Any]								=	BlipaConfig(vision_config=__snake_case			,							text_config=__snake_case      )

     return config, image_size






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

     __A						:			Optional[Any]								=	BlipaForConditionalGeneration(__snake_case      ).eval()

     __A						:			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'),
     }

     __A			,__A						:			Optional[int]								=	model_name_to_original[model_name]

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

     # update state dict keys
     __A						:			List[str]								=	original_model.state_dict()
     __A						:			Dict								=	create_rename_keys(__snake_case      )
     for src, dest in rename_keys:
          rename_key(__snake_case			,							__snake_case			,							__snake_case      )

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

     # read in qv biases
     read_in_q_v_bias(__snake_case			,							__snake_case      )

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

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

     # create processor
     __A						:			Tuple								=	BlipImageProcessor(
         size={'height': image_size, 'width': image_size}			,							image_mean=__snake_case			,							image_std=__snake_case      )
     __A						:			List[Any]								=	BlipaProcessor(image_processor=__snake_case			,							tokenizer=__snake_case      )
     __A						:			int								=	processor(images=__snake_case			,							return_tensors='pt'      ).pixel_values.to(__snake_case      )

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

     original_model.to(__snake_case      )
     hf_model.to(__snake_case      )
     with torch.no_grad():
          if "opt" in model_name:
               __A						:			Dict								=	original_model({'image': original_pixel_values, 'text_input': ['']}      ).logits
               __A						:			int								=	hf_model(__snake_case			,							__snake_case      ).logits
          else:
               __A						:			int								=	original_model(
                   {'image': original_pixel_values, 'text_input': ['\n'], 'text_output': ['\n']}      ).logits
               __A						:			Any								=	input_ids.masked_fill(input_ids == tokenizer.pad_token_id			,							-1_00      )
               __A						:			Optional[Any]								=	hf_model(__snake_case			,							__snake_case			,							labels=__snake_case      ).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":
          __A						:			Tuple								=	torch.tensor(
              [[-41.5_850, -4.4_440, -8.9_922], [-47.4_322, -5.9_143, -1.7_340]]			,							device=__snake_case      )
          assert torch.allclose(logits[0, :3, :3]			,							__snake_case			,							atol=1e-4      )
     elif model_name == "blip2-flan-t5-xl-coco":
          __A						:			str								=	torch.tensor(
              [[-57.0_109, -9.8_967, -12.6_280], [-68.6_578, -12.7_191, -10.5_065]]			,							device=__snake_case      )
     else:
          # cast to same type
          __A						:			str								=	logits.dtype
          assert torch.allclose(original_logits.to(__snake_case      )			,							__snake_case			,							atol=1e-2      )
     print('Looks ok!'      )

     print('Generating a caption...'      )
     __A						:			Dict								=	''
     __A						:			List[Any]								=	tokenizer(__snake_case			,							return_tensors='pt'      ).input_ids.to(__snake_case      )

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

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

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


if __name__ == "__main__":
  lowercase__   :				Optional[Any]      =   argparse.ArgumentParser()
  lowercase__   :				Tuple      =   [
      '''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''',
  )

  lowercase__   :				str      =   parser.parse_args()

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






'''simple docstring'''




from __future__ import annotations

from math import gcd






def 	_lowerCAmelCase      (  __snake_case  :			int			,							__snake_case  :			int = 2			,							__snake_case  :			int = 1			,							__snake_case  :			int = 3			,							)					->					int | None:
     # A value less than 2 can cause an infinite loop in the algorithm.
     if num < 2:
          raise ValueError('The input value cannot be less than 2'      )

     # Because of the relationship between ``f(f(x))`` and ``f(x)``, this
     # algorithm struggles to find factors that are divisible by two.
     # As a workaround, we specifically check for two and even inputs.
     #   See: https://math.stackexchange.com/a/2856214/165820
     if num > 2 and num % 2 == 0:
          return 2

     # Pollard's Rho algorithm requires a function that returns pseudorandom
     # values between 0 <= X < ``num``.  It doesn't need to be random in the
     # sense that the output value is cryptographically secure or difficult
     # to calculate, it only needs to be random in the sense that all output
     # values should be equally likely to appear.
     # For this reason, Pollard suggested using ``f(x) = (x**2 - 1) % num``
     # However, the success of Pollard's algorithm isn't guaranteed and is
     # determined in part by the initial seed and the chosen random function.
     # To make retries easier, we will instead use ``f(x) = (x**2 + C) % num``
     # where ``C`` is a value that we can modify between each attempt.
     def rand_fn(__snake_case  :			int			,							__snake_case  :			int			,							__snake_case  :			int      ) -> int:
          return (pow(__snake_case			,							2      ) + step) % modulus

     for _ in range(__snake_case      ):
          # These track the position within the cycle detection logic.
          __A						:			int								=	seed
          __A						:			Union[str, Any]								=	seed

          while True:
               # At each iteration, the tortoise moves one step and the hare moves two.
               __A						:			List[Any]								=	rand_fn(__snake_case			,							__snake_case			,							__snake_case      )
               __A						:			Optional[Any]								=	rand_fn(__snake_case			,							__snake_case			,							__snake_case      )
               __A						:			Any								=	rand_fn(__snake_case			,							__snake_case			,							__snake_case      )

               # At some point both the tortoise and the hare will enter a cycle whose
               # length ``p`` is a divisor of ``num``.  Once in that cycle, at some point
               # the tortoise and hare will end up on the same value modulo ``p``.
               # We can detect when this happens because the position difference between
               # the tortoise and the hare will share a common divisor with ``num``.
               __A						:			Optional[int]								=	gcd(hare - tortoise			,							__snake_case      )

               if divisor == 1:
                    # No common divisor yet, just keep searching.
                    continue
               else:
                    # We found a common divisor!
                    if divisor == num:
                         # Unfortunately, the divisor is ``num`` itself and is useless.
                         break
                    else:
                         # The divisor is a nontrivial factor of ``num``!
                         return divisor

        # If we made it here, then this attempt failed.
        # We need to pick a new starting seed for the tortoise and hare
        # in addition to a new step value for the random function.
        # To keep this example implementation deterministic, the
        # new values will be generated based on currently available
        # values instead of using something like ``random.randint``.

        # We can use the hare's position as the new seed.
        # This is actually what Richard Brent's the "optimized" variant does.
          __A						:			Union[str, Any]								=	hare

          # The new step value for the random function can just be incremented.
          # At first the results will be similar to what the old function would
          # have produced, but the value will quickly diverge after a bit.
          step += 1

     # We haven't found a divisor within the requested number of attempts.
     # We were unlucky or ``num`` itself is actually prime.
     return None


if __name__ == "__main__":
  import argparse

  lowercase__   :				str      =   argparse.ArgumentParser()
  parser.add_argument(
      '''num''',
      type=int,
      help='''The value to find a divisor of''',
  )
  parser.add_argument(
      '''--attempts''',
      type=int,
      default=3,
      help='''The number of attempts before giving up''',
  )
  lowercase__   :				Optional[int]      =   parser.parse_args()

  lowercase__   :				int      =   pollard_rho(args.num, attempts=args.attempts)
  if divisor is None:
    print(f"""{args.num} is probably prime""")
  else:
    lowercase__   :				List[str]      =   args.num // divisor
    print(f"""{args.num} = {divisor} * {quotient}""")