

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__      ( ):


       """simple docstring"""




       __snake_case       	=							"""https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"""
       __snake_case       	=							Image.open(requests.get(__A	,stream=__A		).raw		).convert("""RGB"""		)

       return image







def 						lowerCamelCase__      ( __A    :int		):


       """simple docstring"""




       __snake_case       	=							[]
       # 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    :str	,__A    :List[Any]	,__A    :str		):


       """simple docstring"""




       __snake_case       	=							dct.pop(__A		)
       __snake_case       	=							val







def 						lowerCamelCase__      ( __A    :List[str]	,__A    :Dict		):


       """simple docstring"""




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

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







def 						lowerCamelCase__      ( __A    :Optional[int]	,__A    :Union[str, Any]		):


       """simple docstring"""




       __snake_case       	=							3_6_4 if """coco""" in model_name else 2_2_4
       __snake_case       	=							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:
              __snake_case       	=							OPTConfig.from_pretrained("""facebook/opt-2.7b"""	,eos_token_id=__A		).to_dict()
       elif "opt-6.7b" in model_name:
              __snake_case       	=							OPTConfig.from_pretrained("""facebook/opt-6.7b"""	,eos_token_id=__A		).to_dict()
       elif "t5-xl" in model_name:
              __snake_case       	=							TaConfig.from_pretrained("""google/flan-t5-xl"""	,dense_act_fn="""gelu"""	,bos_token_id=1		).to_dict()
       elif "t5-xxl" in model_name:
              __snake_case       	=							TaConfig.from_pretrained("""google/flan-t5-xxl"""	,dense_act_fn="""gelu"""	,bos_token_id=1		).to_dict()

       __snake_case       	=							BlipaConfig(vision_config=__A	,text_config=__A		)

       return config, image_size







@torch.no_grad()
def 						lowerCamelCase__      ( __A    :int	,__A    :Any=None	,__A    :List[Any]=False		):


       """simple docstring"""




       __snake_case       	=							(
           AutoTokenizer.from_pretrained("""facebook/opt-2.7b"""		)
           if """opt""" in model_name
           else AutoTokenizer.from_pretrained("""google/flan-t5-xl"""		)
       )
       __snake_case       	=							tokenizer("""\n"""	,add_special_tokens=__A		).input_ids[0]
       __snake_case	,    __snake_case       	=							get_blipa_config(__A	,eos_token_id=__A		)

       __snake_case       	=							BlipaForConditionalGeneration(__A		).eval()

       __snake_case       	=							{
           """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"""),
       }

       __snake_case	,    __snake_case       	=							model_name_to_original[model_name]

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

       # update state dict keys
       __snake_case       	=							original_model.state_dict()
       __snake_case       	=							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():
              __snake_case       	=							state_dict.pop(__A		)
              if key.startswith("""Qformer.bert"""		):
                     __snake_case       	=							key.replace("""Qformer.bert"""	,"""qformer"""		)
              if "attention.self" in key:
                     __snake_case       	=							key.replace("""self"""	,"""attention"""		)
              if "opt_proj" in key:
                     __snake_case       	=							key.replace("""opt_proj"""	,"""language_projection"""		)
              if "t5_proj" in key:
                     __snake_case       	=							key.replace("""t5_proj"""	,"""language_projection"""		)
              if key.startswith("""opt"""		):
                     __snake_case       	=							key.replace("""opt"""	,"""language"""		)
              if key.startswith("""t5"""		):
                     __snake_case       	=							key.replace("""t5"""	,"""language"""		)
              __snake_case       	=							val

       # read in qv biases
       read_in_q_v_bias(__A	,__A		)

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

       __snake_case       	=							load_demo_image()
       __snake_case       	=							vis_processors["""eval"""](__A		).unsqueeze(0		).to(__A		)
       __snake_case       	=							tokenizer(["""\n"""]	,return_tensors="""pt"""		).input_ids.to(__A		)

       # create processor
       __snake_case       	=							BlipImageProcessor(
           size={"""height""": image_size, """width""": image_size}	,image_mean=__A	,image_std=__A		)
       __snake_case       	=							BlipaProcessor(image_processor=__A	,tokenizer=__A		)
       __snake_case       	=							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:
                     __snake_case       	=							original_model({"""image""": original_pixel_values, """text_input""": [""""""]}		).logits
                     __snake_case       	=							hf_model(__A	,__A		).logits
              else:
                     __snake_case       	=							original_model(
                         {"""image""": original_pixel_values, """text_input""": ["""\n"""], """text_output""": ["""\n"""]}		).logits
                     __snake_case       	=							input_ids.masked_fill(input_ids == tokenizer.pad_token_id	,-1_0_0		)
                     __snake_case       	=							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":
              __snake_case       	=							torch.tensor(
                  [[-41.5_850, -4.4_440, -8.9_922], [-47.4_322, -5.9_143, -1.7_340]]	,device=__A		)
              assert torch.allclose(logits[0, :3, :3]	,__A	,atol=1e-4		)
       elif model_name == "blip2-flan-t5-xl-coco":
              __snake_case       	=							torch.tensor(
                  [[-57.0_109, -9.8_967, -12.6_280], [-68.6_578, -12.7_191, -10.5_065]]	,device=__A		)
       else:
              # cast to same type
              __snake_case       	=							logits.dtype
              assert torch.allclose(original_logits.to(__A		)	,__A	,atol=1e-2		)
       print("""Looks ok!"""		)

       print("""Generating a caption..."""		)
       __snake_case       	=							""""""
       __snake_case       	=							tokenizer(__A	,return_tensors="""pt"""		).input_ids.to(__A		)

       __snake_case       	=							original_model.generate({"""image""": original_pixel_values}		)
       __snake_case       	=							hf_model.generate(
           __A	,__A	,do_sample=__A	,num_beams=5	,max_length=3_0	,min_length=1	,top_p=0.9	,repetition_penalty=1.0	,length_penalty=1.0	,temperature=1	,)
       print("""Original generation:"""	,__A		)
       __snake_case       	=							input_ids.shape[1]
       __snake_case       	=							processor.batch_decode(outputs[:, prompt_length:]	,skip_special_tokens=__A		)
       __snake_case       	=							[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__":
    UpperCamelCase__		 =     argparse.ArgumentParser()
    UpperCamelCase__		 =     [
        '''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''',
    )

    UpperCamelCase__		 =     parser.parse_args()

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




from __future__ import annotations

import random

# Maximum size of the population.  Bigger could be faster but is more memory expensive.
UpperCamelCase__		 =     200
# Number of elements selected in every generation of evolution. The selection takes
# place from best to worst of that generation and must be smaller than N_POPULATION.
UpperCamelCase__		 =     50
# Probability that an element of a generation can mutate, changing one of its genes.
# This will guarantee that all genes will be used during evolution.
UpperCamelCase__		 =     0.4
# Just a seed to improve randomness required by the algorithm.
random.seed(random.randint(0, 1000))







def 						lowerCamelCase__      ( __A    :str	,__A    :str		):


       """simple docstring"""




       __snake_case       	=							len([g for position, g in enumerate(__A		) if g == main_target[position]]		)
       return (item, float(__A		))







def 						lowerCamelCase__      ( __A    :str	,__A    :str		):


       """simple docstring"""




       __snake_case       	=							random.randint(0	,len(__A		) - 1		)
       __snake_case       	=							parent_a[:random_slice] + parent_a[random_slice:]
       __snake_case       	=							parent_a[:random_slice] + parent_a[random_slice:]
       return (child_a, child_a)







def 						lowerCamelCase__      ( __A    :str	,__A    :list[str]		):


       """simple docstring"""




       __snake_case       	=							list(__A		)
       if random.uniform(0	,1		) < MUTATION_PROBABILITY:
              __snake_case       	=							random.choice(__A		)
       return "".join(__A		)







def 						lowerCamelCase__      ( __A    :tuple[str, float]	,__A    :list[tuple[str, float]]	,__A    :list[str]	,):


       """simple docstring"""




       __snake_case       	=							[]
       # Generate more children proportionally to the fitness score.
       __snake_case       	=							int(parent_a[1] * 1_0_0		) + 1
       __snake_case       	=							1_0 if child_n >= 1_0 else child_n
       for _ in range(__A		):
              __snake_case       	=							population_score[random.randint(0	,__A		)][0]

              __snake_case	,    __snake_case       	=							crossover(parent_a[0]	,__A		)
              # Append new string to the population list.
              pop.append(mutate(__A	,__A		)		)
              pop.append(mutate(__A	,__A		)		)
       return pop







def 						lowerCamelCase__      ( __A    :str	,__A    :list[str]	,__A    :bool = True		):


       """simple docstring"""




       if N_POPULATION < N_SELECTED:
              __snake_case       	=							F'{N_POPULATION} must be bigger than {N_SELECTED}'
              raise ValueError(__A		)
       # Verify that the target contains no genes besides the ones inside genes variable.
       __snake_case       	=							sorted({c for c in target if c not in genes}		)
       if not_in_genes_list:
              __snake_case       	=							F'{not_in_genes_list} is not in genes list, evolution cannot converge'
              raise ValueError(__A		)

       # Generate random starting population.
       __snake_case       	=							[]
       for _ in range(__A		):
              population.append("""""".join([random.choice(__A		) for i in range(len(__A		)		)]		)		)

       # Just some logs to know what the algorithms is doing.
       __snake_case	,    __snake_case       	=							0, 0

       # This loop will end when we find a perfect match for our target.
       while True:
              generation += 1
              total_population += len(__A		)

              # Random population created. Now it's time to evaluate.

              # Adding a bit of concurrency can make everything faster,
              #
              # import concurrent.futures
              # population_score: list[tuple[str, float]] = []
              # with concurrent.futures.ThreadPoolExecutor(
              #                                   max_workers=NUM_WORKERS) as executor:
              #     futures = {executor.submit(evaluate, item) for item in population}
              #     concurrent.futures.wait(futures)
              #     population_score = [item.result() for item in futures]
              #
              # but with a simple algorithm like this, it will probably be slower.
              # We just need to call evaluate for every item inside the population.
              __snake_case       	=							[evaluate(__A	,__A		) for item in population]

              # Check if there is a matching evolution.
              __snake_case       	=							sorted(__A	,key=lambda __A		: x[1]	,reverse=__A		)
              if population_score[0][0] == target:
                     return (generation, total_population, population_score[0][0])

              # Print the best result every 10 generation.
              # Just to know that the algorithm is working.
              if debug and generation % 1_0 == 0:
                     print(
                         F'\nGeneration: {generation}'
                         F'\nTotal Population:{total_population}'
                         F'\nBest score: {population_score[0][1]}'
                         F'\nBest string: {population_score[0][0]}'		)

              # Flush the old population, keeping some of the best evolutions.
              # Keeping this avoid regression of evolution.
              __snake_case       	=							population[: int(N_POPULATION / 3		)]
              population.clear()
              population.extend(__A		)
              # Normalize population score to be between 0 and 1.
              __snake_case       	=							[
                  (item, score / len(__A		)) for item, score in population_score
              ]

              # This is selection
              for i in range(__A		):
                     population.extend(select(population_score[int(__A		)]	,__A	,__A		)		)
                     # Check if the population has already reached the maximum value and if so,
                     # break the cycle.  If this check is disabled, the algorithm will take
                     # forever to compute large strings, but will also calculate small strings in
                     # a far fewer generations.
                     if len(__A		) > N_POPULATION:
                            break


if __name__ == "__main__":
    UpperCamelCase__		 =     (
        '''This is a genetic algorithm to evaluate, combine, evolve, and mutate a string!'''
    )
    UpperCamelCase__		 =     list(
        ''' ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklm'''
        '''nopqrstuvwxyz.,;!?+-*#@^\'èéòà€ù=)(&%$£/\\'''
    )
    UpperCamelCase__ ,UpperCamelCase__ ,UpperCamelCase__		 =     basic(target_str, genes_list)
    print(
        F'\nGeneration: {generation}\nTotal Population: {population}\nTarget: {target}'
    )

