





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     SCREAMING_SNAKE_CASE			(      ):
       _UpperCamelCase              =						'''https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'''
       _UpperCamelCase              =						Image.open(requests.get(lowerCAmelCase    ,       stream=lowerCAmelCase					).raw					).convert('''RGB'''					)

       return image


def     SCREAMING_SNAKE_CASE			(      lowerCAmelCase					):
       _UpperCamelCase              =						[]
       # 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     SCREAMING_SNAKE_CASE			(      lowerCAmelCase    ,       lowerCAmelCase    ,       lowerCAmelCase					):
       _UpperCamelCase              =						dct.pop(lowerCAmelCase					)
       _UpperCamelCase              =						val


def     SCREAMING_SNAKE_CASE			(      lowerCAmelCase    ,       lowerCAmelCase					):
       for i in range(config.vision_config.num_hidden_layers					):
              # read in original q and v biases
              _UpperCamelCase              =						state_dict.pop(f'''visual_encoder.blocks.{i}.attn.q_bias'''					)
              _UpperCamelCase              =						state_dict.pop(f'''visual_encoder.blocks.{i}.attn.v_bias'''					)

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


def     SCREAMING_SNAKE_CASE			(      lowerCAmelCase					):
       _UpperCamelCase              =						364 if '''coco''' in model_name else 224
       _UpperCamelCase              =						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:
              _UpperCamelCase              =						TaConfig.from_pretrained('''google/flan-t5-xl'''    ,       dense_act_fn='''gelu'''    ,       bos_token_id=1					).to_dict()
       elif "t5-xxl" in model_name:
              _UpperCamelCase              =						TaConfig.from_pretrained('''google/flan-t5-xxl'''    ,       dense_act_fn='''gelu'''    ,       bos_token_id=1					).to_dict()
       elif "vicuna-7b" in model_name:
              _UpperCamelCase              =						LlamaConfig.from_pretrained('''decapoda-research/llama-7b-hf'''    ,       vocab_size=32_001					).to_dict()
       elif "vicuna-13b" in model_name:
              _UpperCamelCase              =						LlamaConfig.from_pretrained('''decapoda-research/llama-13b-hf'''    ,       vocab_size=32_001					).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
       _UpperCamelCase              =						InstructBlipQFormerConfig(vocab_size=30_523					).to_dict()
       _UpperCamelCase              =						InstructBlipConfig(vision_config=lowerCAmelCase    ,       text_config=lowerCAmelCase    ,       qformer_config=lowerCAmelCase					)

       return config, image_size


@torch.no_grad()
def     SCREAMING_SNAKE_CASE			(      lowerCAmelCase    ,       lowerCAmelCase=None    ,       lowerCAmelCase=False					):
       _UpperCamelCase              =						AutoTokenizer.from_pretrained('''bert-base-uncased'''    ,       truncation_side='''left'''					)
       qformer_tokenizer.add_special_tokens({'''bos_token''': '''[DEC]'''}					)

       if "t5" in model_name:
              _UpperCamelCase              =						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>"})
              _UpperCamelCase              =						LlamaTokenizerFast.from_pretrained(
                  '''huggyllama/llama-7b'''    ,       truncation_side='''left'''    ,       bos_token='''</s>'''    ,       unk_token='''</s>'''					)
              tokenizer.add_special_tokens({'''pad_token''': '''[PAD]'''}					)

       _UpperCamelCase       ,		_UpperCamelCase              =						get_blipa_config(lowerCAmelCase					)
       _UpperCamelCase              =						InstructBlipForConditionalGeneration(lowerCAmelCase					).eval()

       _UpperCamelCase              =						{
           '''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'''),
       }

       _UpperCamelCase       ,		_UpperCamelCase              =						model_name_to_original[model_name]

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

       # update state dict keys
       _UpperCamelCase              =						original_model.state_dict()
       _UpperCamelCase              =						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():
              _UpperCamelCase              =						state_dict.pop(lowerCAmelCase					)
              if key.startswith('''Qformer.bert'''					):
                     _UpperCamelCase              =						key.replace('''Qformer.bert'''    ,       '''qformer'''					)
              if "attention.self" in key:
                     _UpperCamelCase              =						key.replace('''self'''    ,       '''attention'''					)
              if "llm_proj" in key:
                     _UpperCamelCase              =						key.replace('''llm_proj'''    ,       '''language_projection'''					)
              if "t5_proj" in key:
                     _UpperCamelCase              =						key.replace('''t5_proj'''    ,       '''language_projection'''					)
              if key.startswith('''llm_model'''					):
                     _UpperCamelCase              =						key.replace('''llm_model'''    ,       '''language_model'''					)
              if key.startswith('''t5'''					):
                     _UpperCamelCase              =						key.replace('''t5'''    ,       '''language'''					)
              _UpperCamelCase              =						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					)

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

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

       # make sure processor creates exact same pixel values
       _UpperCamelCase              =						vis_processors['''eval'''](lowerCAmelCase					).unsqueeze(0					).to(lowerCAmelCase					)
       _UpperCamelCase              =						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:
                     _UpperCamelCase              =						original_model({'''image''': original_pixel_values, '''text_input''': [prompt]}					).logits
                     _UpperCamelCase              =						hf_model(**lowerCAmelCase					).logits
              else:
                     _UpperCamelCase              =						original_model(
                         {'''image''': original_pixel_values, '''text_input''': [prompt], '''text_output''': ['''\n''']}					).logits
                     _UpperCamelCase              =						tokenizer('''\n'''    ,       return_tensors='''pt'''					).input_ids.to(lowerCAmelCase					)
                     _UpperCamelCase              =						label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id    ,       -100					)
                     _UpperCamelCase              =						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
       _UpperCamelCase              =						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...'''					)
       _UpperCamelCase              =						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...'''					)
       _UpperCamelCase              =						hf_model.generate(
           **lowerCAmelCase    ,       do_sample=lowerCAmelCase    ,       num_beams=5    ,       max_length=256    ,       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?
              _UpperCamelCase              =						2
       print('''Original generation:'''    ,       lowerCAmelCase					)
       _UpperCamelCase              =						processor.batch_decode(lowerCAmelCase    ,       skip_special_tokens=lowerCAmelCase					)
       _UpperCamelCase              =						[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__":
       lowercase    : List[Any]    	=  argparse.ArgumentParser()
       lowercase    : Tuple    	=  [
           """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""",
       )

       lowercase    : List[str]    	=  parser.parse_args()

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













import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
       import jax
       import jax.numpy as jnp
       from flax.jax_utils import replicate
       from flax.training.common_utils import shard

@slow
@require_flax
class  __A(						unittest.TestCase ):

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            super().tearDown()
            gc.collect()

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase       ,		_UpperCamelCase              =						FlaxStableDiffusionPipeline.from_pretrained(
                '''stabilityai/stable-diffusion-2''',       revision='''bf16''',       dtype=jnp.bfloataa,       )

            _UpperCamelCase              =						'''A painting of a squirrel eating a burger'''
            _UpperCamelCase              =						jax.device_count()
            _UpperCamelCase              =						num_samples * [prompt]
            _UpperCamelCase              =						sd_pipe.prepare_inputs(A      )

            _UpperCamelCase              =						replicate(A      )
            _UpperCamelCase              =						shard(A      )

            _UpperCamelCase              =						jax.random.PRNGKey(0      )
            _UpperCamelCase              =						jax.random.split(A,       jax.device_count()      )

            _UpperCamelCase              =						sd_pipe(A,       A,       A,       num_inference_steps=25,       jit=A      )[0]
            assert images.shape == (jax.device_count(), 1, 768, 768, 3)

            _UpperCamelCase              =						images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]      )
            _UpperCamelCase              =						images[0, 253:256, 253:256, -1]

            _UpperCamelCase              =						jnp.asarray(jax.device_get(image_slice.flatten()      )      )
            _UpperCamelCase              =						jnp.array([0.4_238, 0.4_414, 0.4_395, 0.4_453, 0.4_629, 0.4_590, 0.4_531, 0.45_508, 0.4_512]      )
            print(F'''output_slice: {output_slice}'''      )
            assert jnp.abs(output_slice - expected_slice      ).max() < 1E-2



     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						'''stabilityai/stable-diffusion-2'''
            _UpperCamelCase       ,		_UpperCamelCase              =						FlaxDPMSolverMultistepScheduler.from_pretrained(A,       subfolder='''scheduler'''      )
            _UpperCamelCase       ,		_UpperCamelCase              =						FlaxStableDiffusionPipeline.from_pretrained(
                A,       scheduler=A,       revision='''bf16''',       dtype=jnp.bfloataa,       )
            _UpperCamelCase              =						scheduler_params

            _UpperCamelCase              =						'''A painting of a squirrel eating a burger'''
            _UpperCamelCase              =						jax.device_count()
            _UpperCamelCase              =						num_samples * [prompt]
            _UpperCamelCase              =						sd_pipe.prepare_inputs(A      )

            _UpperCamelCase              =						replicate(A      )
            _UpperCamelCase              =						shard(A      )

            _UpperCamelCase              =						jax.random.PRNGKey(0      )
            _UpperCamelCase              =						jax.random.split(A,       jax.device_count()      )

            _UpperCamelCase              =						sd_pipe(A,       A,       A,       num_inference_steps=25,       jit=A      )[0]
            assert images.shape == (jax.device_count(), 1, 768, 768, 3)

            _UpperCamelCase              =						images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]      )
            _UpperCamelCase              =						images[0, 253:256, 253:256, -1]

            _UpperCamelCase              =						jnp.asarray(jax.device_get(image_slice.flatten()      )      )
            _UpperCamelCase              =						jnp.array([0.4_336, 0.42_969, 0.4_453, 0.4_199, 0.4_297, 0.4_531, 0.4_434, 0.4_434, 0.4_297]      )
            print(F'''output_slice: {output_slice}'''      )
            assert jnp.abs(output_slice - expected_slice      ).max() < 1E-2






