

import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    BertTokenizer,
    ViltConfig,
    ViltForImageAndTextRetrieval,
    ViltForImagesAndTextClassification,
    ViltForMaskedLM,
    ViltForQuestionAnswering,
    ViltImageProcessor,
    ViltProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
SCREAMING_SNAKE_CASE_            = logging.get_logger(__name__)

def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					Optional[Any]					, lowerCAmelCase:					Any=False					, lowerCAmelCase:					List[str]=False					, lowerCAmelCase:					Optional[int]=False	)				->    Dict:
				_UpperCAmelCase    :						Optional[int]							=      []
				for i in range(config.num_hidden_layers	):
								# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
								rename_keys.append((F'transformer.blocks.{i}.norm1.weight', F'vilt.encoder.layer.{i}.layernorm_before.weight')	)
								rename_keys.append((F'transformer.blocks.{i}.norm1.bias', F'vilt.encoder.layer.{i}.layernorm_before.bias')	)
								rename_keys.append(
								    (F'transformer.blocks.{i}.attn.proj.weight', F'vilt.encoder.layer.{i}.attention.output.dense.weight')	)
								rename_keys.append(
								    (F'transformer.blocks.{i}.attn.proj.bias', F'vilt.encoder.layer.{i}.attention.output.dense.bias')	)
								rename_keys.append((F'transformer.blocks.{i}.norm2.weight', F'vilt.encoder.layer.{i}.layernorm_after.weight')	)
								rename_keys.append((F'transformer.blocks.{i}.norm2.bias', F'vilt.encoder.layer.{i}.layernorm_after.bias')	)
								rename_keys.append(
								    (F'transformer.blocks.{i}.mlp.fc1.weight', F'vilt.encoder.layer.{i}.intermediate.dense.weight')	)
								rename_keys.append((F'transformer.blocks.{i}.mlp.fc1.bias', F'vilt.encoder.layer.{i}.intermediate.dense.bias')	)
								rename_keys.append((F'transformer.blocks.{i}.mlp.fc2.weight', F'vilt.encoder.layer.{i}.output.dense.weight')	)
								rename_keys.append((F'transformer.blocks.{i}.mlp.fc2.bias', F'vilt.encoder.layer.{i}.output.dense.bias')	)

				# embeddings
				rename_keys.extend(
				    [
				        # text embeddings
				        ("text_embeddings.word_embeddings.weight", "vilt.embeddings.text_embeddings.word_embeddings.weight"),
				        (
				            "text_embeddings.position_embeddings.weight",
				            "vilt.embeddings.text_embeddings.position_embeddings.weight",
				        ),
				        ("text_embeddings.position_ids", "vilt.embeddings.text_embeddings.position_ids"),
				        (
				            "text_embeddings.token_type_embeddings.weight",
				            "vilt.embeddings.text_embeddings.token_type_embeddings.weight",
				        ),
				        ("text_embeddings.LayerNorm.weight", "vilt.embeddings.text_embeddings.LayerNorm.weight"),
				        ("text_embeddings.LayerNorm.bias", "vilt.embeddings.text_embeddings.LayerNorm.bias"),
				        # patch embeddings
				        ("transformer.cls_token", "vilt.embeddings.cls_token"),
				        ("transformer.patch_embed.proj.weight", "vilt.embeddings.patch_embeddings.projection.weight"),
				        ("transformer.patch_embed.proj.bias", "vilt.embeddings.patch_embeddings.projection.bias"),
				        ("transformer.pos_embed", "vilt.embeddings.position_embeddings"),
				        # token type embeddings
				        ("token_type_embeddings.weight", "vilt.embeddings.token_type_embeddings.weight"),
				    ]	)

				# final layernorm + pooler
				rename_keys.extend(
				    [
				        ("transformer.norm.weight", "vilt.layernorm.weight"),
				        ("transformer.norm.bias", "vilt.layernorm.bias"),
				        ("pooler.dense.weight", "vilt.pooler.dense.weight"),
				        ("pooler.dense.bias", "vilt.pooler.dense.bias"),
				    ]	)

				# classifier head(s)
				if vqa_model:
								# classification head
								rename_keys.extend(
								    [
								        ("vqa_classifier.0.weight", "classifier.0.weight"),
								        ("vqa_classifier.0.bias", "classifier.0.bias"),
								        ("vqa_classifier.1.weight", "classifier.1.weight"),
								        ("vqa_classifier.1.bias", "classifier.1.bias"),
								        ("vqa_classifier.3.weight", "classifier.3.weight"),
								        ("vqa_classifier.3.bias", "classifier.3.bias"),
								    ]	)
				elif nlvr_model:
								# classification head
								rename_keys.extend(
								    [
								        ("nlvr2_classifier.0.weight", "classifier.0.weight"),
								        ("nlvr2_classifier.0.bias", "classifier.0.bias"),
								        ("nlvr2_classifier.1.weight", "classifier.1.weight"),
								        ("nlvr2_classifier.1.bias", "classifier.1.bias"),
								        ("nlvr2_classifier.3.weight", "classifier.3.weight"),
								        ("nlvr2_classifier.3.bias", "classifier.3.bias"),
								    ]	)
				else:
								pass

				return rename_keys

def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					Dict					, lowerCAmelCase:					Any	)				->    Dict:
				for i in range(config.num_hidden_layers	):
								_UpperCAmelCase    :						str							=      "vilt."
								# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
								_UpperCAmelCase    :						List[Any]							=      state_dict.pop(F'transformer.blocks.{i}.attn.qkv.weight'	)
								_UpperCAmelCase    :						Optional[int]							=      state_dict.pop(F'transformer.blocks.{i}.attn.qkv.bias'	)
								# next, add query, keys and values (in that order) to the state dict
								_UpperCAmelCase    :						Union[str, Any]							=      in_proj_weight[
								    : config.hidden_size, :
								]
								_UpperCAmelCase    :						List[Any]							=      in_proj_bias[: config.hidden_size]
								_UpperCAmelCase    :						Dict							=      in_proj_weight[
								    config.hidden_size : config.hidden_size * 2, :
								]
								_UpperCAmelCase    :						List[str]							=      in_proj_bias[
								    config.hidden_size : config.hidden_size * 2
								]
								_UpperCAmelCase    :						Optional[Any]							=      in_proj_weight[
								    -config.hidden_size :, :
								]
								_UpperCAmelCase    :						Any							=      in_proj_bias[-config.hidden_size :]

def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					int	)				->    Tuple:
				_UpperCAmelCase    :						Any							=      ["head.weight", "head.bias"]
				for k in ignore_keys:
								state_dict.pop(lowerCAmelCase					, lowerCAmelCase	)

def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					List[Any]					, lowerCAmelCase:					Any					, lowerCAmelCase:					Optional[Any]	)				->    Dict:
				_UpperCAmelCase    :						Tuple							=      dct.pop(lowerCAmelCase	)
				_UpperCAmelCase    :						Dict							=      val

@torch.no_grad()
def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					Any					, lowerCAmelCase:					Tuple	)				->    Any:
				_UpperCAmelCase    :						int							=      ViltConfig(image_size=384					, patch_size=32					, tie_word_embeddings=lowerCAmelCase	)
				_UpperCAmelCase    :						Union[str, Any]							=      False
				_UpperCAmelCase    :						Dict							=      False
				_UpperCAmelCase    :						Dict							=      False
				_UpperCAmelCase    :						int							=      False
				if "vqa" in checkpoint_url:
								_UpperCAmelCase    :						Optional[Any]							=      True
								_UpperCAmelCase    :						Dict							=      3129
								_UpperCAmelCase    :						List[Any]							=      "huggingface/label-files"
								_UpperCAmelCase    :						str							=      "vqa2-id2label.json"
								_UpperCAmelCase    :						int							=      json.load(open(hf_hub_download(lowerCAmelCase					, lowerCAmelCase					, repo_type="dataset"	)					, "r"	)	)
								_UpperCAmelCase    :						Any							=      {int(lowerCAmelCase	): v for k, v in idalabel.items()}
								_UpperCAmelCase    :						List[str]							=      idalabel
								_UpperCAmelCase    :						Union[str, Any]							=      {v: k for k, v in idalabel.items()}
								_UpperCAmelCase    :						List[Any]							=      ViltForQuestionAnswering(lowerCAmelCase	)
				elif "nlvr" in checkpoint_url:
								_UpperCAmelCase    :						str							=      True
								_UpperCAmelCase    :						Optional[Any]							=      2
								_UpperCAmelCase    :						str							=      {0: "False", 1: "True"}
								_UpperCAmelCase    :						Any							=      {v: k for k, v in config.idalabel.items()}
								_UpperCAmelCase    :						Optional[Any]							=      3
								_UpperCAmelCase    :						Any							=      ViltForImagesAndTextClassification(lowerCAmelCase	)
				elif "irtr" in checkpoint_url:
								_UpperCAmelCase    :						Any							=      True
								_UpperCAmelCase    :						Optional[Any]							=      ViltForImageAndTextRetrieval(lowerCAmelCase	)
				elif "mlm_itm" in checkpoint_url:
								_UpperCAmelCase    :						Tuple							=      True
								_UpperCAmelCase    :						Dict							=      ViltForMaskedLM(lowerCAmelCase	)
				else:
								raise ValueError("Unknown model type"	)

				# load state_dict of original model, remove and rename some keys
				_UpperCAmelCase    :						Optional[int]							=      torch.hub.load_state_dict_from_url(lowerCAmelCase					, map_location="cpu"	)["state_dict"]
				_UpperCAmelCase    :						int							=      create_rename_keys(lowerCAmelCase					, lowerCAmelCase					, lowerCAmelCase					, lowerCAmelCase	)
				for src, dest in rename_keys:
								rename_key(lowerCAmelCase					, lowerCAmelCase					, lowerCAmelCase	)
				read_in_q_k_v(lowerCAmelCase					, lowerCAmelCase	)
				if mlm_model or irtr_model:
								_UpperCAmelCase    :						Dict							=      ["itm_score.fc.weight", "itm_score.fc.bias"]
								for k in ignore_keys:
												state_dict.pop(lowerCAmelCase					, lowerCAmelCase	)

    # load state dict into HuggingFace model
				model.eval()
				if mlm_model:
								_UpperCAmelCase						,							_UpperCAmelCase    :						Dict							=      model.load_state_dict(lowerCAmelCase					, strict=lowerCAmelCase	)
								assert missing_keys == ["mlm_score.decoder.bias"]
				else:
								model.load_state_dict(lowerCAmelCase	)

				# Define processor
				_UpperCAmelCase    :						Optional[int]							=      ViltImageProcessor(size=384	)
				_UpperCAmelCase    :						Union[str, Any]							=      BertTokenizer.from_pretrained("bert-base-uncased"	)
				_UpperCAmelCase    :						Optional[Any]							=      ViltProcessor(lowerCAmelCase					, lowerCAmelCase	)

				# Forward pass on example inputs (image + text)
				if nlvr_model:
								_UpperCAmelCase    :						List[str]							=      Image.open(requests.get("https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg"					, stream=lowerCAmelCase	).raw	)
								_UpperCAmelCase    :						str							=      Image.open(requests.get("https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg"					, stream=lowerCAmelCase	).raw	)
								_UpperCAmelCase    :						Tuple							=      (
								    "The left image contains twice the number of dogs as the right image, and at least two dogs in total are"
								    " standing."
								)
								_UpperCAmelCase    :						Optional[Any]							=      processor(lowerCAmelCase					, lowerCAmelCase					, return_tensors="pt"	)
								_UpperCAmelCase    :						int							=      processor(lowerCAmelCase					, lowerCAmelCase					, return_tensors="pt"	)
								_UpperCAmelCase    :						List[Any]							=      model(
								    input_ids=encoding_a.input_ids					, pixel_values=encoding_a.pixel_values					, pixel_values_a=encoding_a.pixel_values					, )
				else:
								_UpperCAmelCase    :						Dict							=      Image.open(requests.get("http://images.cocodataset.org/val2017/000000039769.jpg"					, stream=lowerCAmelCase	).raw	)
								if mlm_model:
												_UpperCAmelCase    :						Dict							=      "a bunch of [MASK] laying on a [MASK]."
								else:
												_UpperCAmelCase    :						int							=      "How many cats are there?"
								_UpperCAmelCase    :						Optional[Any]							=      processor(lowerCAmelCase					, lowerCAmelCase					, return_tensors="pt"	)
								_UpperCAmelCase    :						Tuple							=      model(**lowerCAmelCase	)

				# Verify outputs
				if mlm_model:
								_UpperCAmelCase    :						List[str]							=      torch.Size([1, 11, 3_0522]	)
								_UpperCAmelCase    :						Tuple							=      torch.tensor([-12.5_061, -12.5_123, -12.5_174]	)
								assert outputs.logits.shape == expected_shape
								assert torch.allclose(outputs.logits[0, 0, :3]					, lowerCAmelCase					, atol=1E-4	)

								# verify masked token prediction equals "cats"
								_UpperCAmelCase    :						Dict							=      outputs.logits[0, 4, :].argmax(-1	).item()
								assert tokenizer.decode([predicted_id]	) == "cats"
				elif vqa_model:
								_UpperCAmelCase    :						str							=      torch.Size([1, 3129]	)
								_UpperCAmelCase    :						List[Any]							=      torch.tensor([-15.9_495, -18.1_472, -10.3_041]	)
								assert torch.allclose(outputs.logits[0, :3]					, lowerCAmelCase					, atol=1E-4	)
								assert outputs.logits.shape == expected_shape
								assert torch.allclose(outputs.logits[0, 0, :3]					, lowerCAmelCase					, atol=1E-4	)

								# verify vqa prediction equals "2"
								_UpperCAmelCase    :						Optional[int]							=      outputs.logits.argmax(-1	).item()
								assert model.config.idalabel[predicted_idx] == "2"
				elif nlvr_model:
								_UpperCAmelCase    :						Tuple							=      torch.Size([1, 2]	)
								_UpperCAmelCase    :						Dict							=      torch.tensor([-2.8_721, 2.1_291]	)
								assert torch.allclose(outputs.logits[0, :3]					, lowerCAmelCase					, atol=1E-4	)
								assert outputs.logits.shape == expected_shape

				Path(lowerCAmelCase	).mkdir(exist_ok=lowerCAmelCase	)
				print(F'Saving model and processor to {pytorch_dump_folder_path}'	)
				model.save_pretrained(lowerCAmelCase	)
				processor.save_pretrained(lowerCAmelCase	)


if __name__ == "__main__":
					SCREAMING_SNAKE_CASE_            = argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '--checkpoint_url',
					    default='https://github.com/dandelin/ViLT/releases/download/200k/vilt_200k_mlm_itm.ckpt',
					    type=str,
					    help='URL of the checkpoint you\'d like to convert.',
					)
					parser.add_argument(
					    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
					)

					SCREAMING_SNAKE_CASE_            = parser.parse_args()
					convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)





import argparse

import torch

from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_from_original_stable_diffusion_ckpt


if __name__ == "__main__":
					SCREAMING_SNAKE_CASE_            = argparse.ArgumentParser()

					parser.add_argument(
					    '--checkpoint_path', default=None, type=str, required=True, help='Path to the checkpoint to convert.'
					)
					# !wget https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml
					parser.add_argument(
					    '--original_config_file',
					    default=None,
					    type=str,
					    help='The YAML config file corresponding to the original architecture.',
					)
					parser.add_argument(
					    '--num_in_channels',
					    default=None,
					    type=int,
					    help='The number of input channels. If `None` number of input channels will be automatically inferred.',
					)
					parser.add_argument(
					    '--scheduler_type',
					    default='pndm',
					    type=str,
					    help='Type of scheduler to use. Should be one of [\'pndm\', \'lms\', \'ddim\', \'euler\', \'euler-ancestral\', \'dpm\']',
					)
					parser.add_argument(
					    '--pipeline_type',
					    default=None,
					    type=str,
					    help=(
					        'The pipeline type. One of \'FrozenOpenCLIPEmbedder\', \'FrozenCLIPEmbedder\', \'PaintByExample\''
					        '. If `None` pipeline will be automatically inferred.'
					    ),
					)
					parser.add_argument(
					    '--image_size',
					    default=None,
					    type=int,
					    help=(
					        'The image size that the model was trained on. Use 512 for Stable Diffusion v1.X and Stable Siffusion v2'
					        ' Base. Use 768 for Stable Diffusion v2.'
					    ),
					)
					parser.add_argument(
					    '--prediction_type',
					    default=None,
					    type=str,
					    help=(
					        'The prediction type that the model was trained on. Use \'epsilon\' for Stable Diffusion v1.X and Stable'
					        ' Diffusion v2 Base. Use \'v_prediction\' for Stable Diffusion v2.'
					    ),
					)
					parser.add_argument(
					    '--extract_ema',
					    action='store_true',
					    help=(
					        'Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights'
					        ' or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield'
					        ' higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning.'
					    ),
					)
					parser.add_argument(
					    '--upcast_attention',
					    action='store_true',
					    help=(
					        'Whether the attention computation should always be upcasted. This is necessary when running stable'
					        ' diffusion 2.1.'
					    ),
					)
					parser.add_argument(
					    '--from_safetensors',
					    action='store_true',
					    help='If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.',
					)
					parser.add_argument(
					    '--to_safetensors',
					    action='store_true',
					    help='Whether to store pipeline in safetensors format or not.',
					)
					parser.add_argument('--dump_path', default=None, type=str, required=True, help='Path to the output model.')
					parser.add_argument('--device', type=str, help='Device to use (e.g. cpu, cuda:0, cuda:1, etc.)')
					parser.add_argument(
					    '--stable_unclip',
					    type=str,
					    default=None,
					    required=False,
					    help='Set if this is a stable unCLIP model. One of \'txt2img\' or \'img2img\'.',
					)
					parser.add_argument(
					    '--stable_unclip_prior',
					    type=str,
					    default=None,
					    required=False,
					    help='Set if this is a stable unCLIP txt2img model. Selects which prior to use. If `--stable_unclip` is set to `txt2img`, the karlo prior (https://huggingface.co/kakaobrain/karlo-v1-alpha/tree/main/prior) is selected by default.',
					)
					parser.add_argument(
					    '--clip_stats_path',
					    type=str,
					    help='Path to the clip stats file. Only required if the stable unclip model\'s config specifies `model.params.noise_aug_config.params.clip_stats_path`.',
					    required=False,
					)
					parser.add_argument(
					    '--controlnet', action='store_true', default=None, help='Set flag if this is a controlnet checkpoint.'
					)
					parser.add_argument('--half', action='store_true', help='Save weights in half precision.')
					parser.add_argument(
					    '--vae_path',
					    type=str,
					    default=None,
					    required=False,
					    help='Set to a path, hub id to an already converted vae to not convert it again.',
					)
					SCREAMING_SNAKE_CASE_            = parser.parse_args()

					SCREAMING_SNAKE_CASE_            = download_from_original_stable_diffusion_ckpt(
					    checkpoint_path=args.checkpoint_path,
					    original_config_file=args.original_config_file,
					    image_size=args.image_size,
					    prediction_type=args.prediction_type,
					    model_type=args.pipeline_type,
					    extract_ema=args.extract_ema,
					    scheduler_type=args.scheduler_type,
					    num_in_channels=args.num_in_channels,
					    upcast_attention=args.upcast_attention,
					    from_safetensors=args.from_safetensors,
					    device=args.device,
					    stable_unclip=args.stable_unclip,
					    stable_unclip_prior=args.stable_unclip_prior,
					    clip_stats_path=args.clip_stats_path,
					    controlnet=args.controlnet,
					    vae_path=args.vae_path,
					)

					if args.half:
										pipe.to(torch_dtype=torch.floataa)

					if args.controlnet:
										# only save the controlnet model
										pipe.controlnet.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)
					else:
										pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)


