



import argparse

import requests
import torch
from PIL import Image
from torchvision.transforms import Compose, Normalize, Resize, ToTensor

from transformers import SwinaSRConfig, SwinaSRForImageSuperResolution, SwinaSRImageProcessor

def      lowerCAmelCase_					(  _lowercase    :       int)							->	Tuple:







			"""simple docstring"""
			a__       :						str							=  SwinaSRConfig()

			if "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
						a__       :						Optional[int]							=  4
			elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
						a__       :						Tuple							=  4
						a__       :						str							=  48
						a__       :						int							=  """pixelshuffle_aux"""
			elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
						a__       :						Optional[Any]							=  [6, 6, 6, 6]
						a__       :						List[str]							=  60
						a__       :						List[Any]							=  [6, 6, 6, 6]
						a__       :						Any							=  """pixelshuffledirect"""
			elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
						a__       :						Dict							=  4
						a__       :						Tuple							=  """nearest+conv"""
			elif "Swin2SR_Jpeg_dynamic" in checkpoint_url:
						a__       :						Dict							=  1
						a__       :						str							=  1
						a__       :						List[str]							=  126
						a__       :						Optional[Any]							=  7
						a__       :						Tuple							=  255.0
						a__       :						Dict							=  """"""

			return config

def      lowerCAmelCase_					(  _lowercase    :       str						,       _lowercase    :       Optional[int])							->	List[Any]:







			"""simple docstring"""
			if "patch_embed.proj" in name and "layers" not in name:
						a__       :						str							=  name.replace("""patch_embed.proj"""						,       """embeddings.patch_embeddings.projection""")
			if "patch_embed.norm" in name:
						a__       :						str							=  name.replace("""patch_embed.norm"""						,       """embeddings.patch_embeddings.layernorm""")
			if "layers" in name:
						a__       :						Optional[Any]							=  name.replace("""layers"""						,       """encoder.stages""")
			if "residual_group.blocks" in name:
						a__       :						Union[str, Any]							=  name.replace("""residual_group.blocks"""						,       """layers""")
			if "attn.proj" in name:
						a__       :						int							=  name.replace("""attn.proj"""						,       """attention.output.dense""")
			if "attn" in name:
						a__       :						Any							=  name.replace("""attn"""						,       """attention.self""")
			if "norm1" in name:
						a__       :						int							=  name.replace("""norm1"""						,       """layernorm_before""")
			if "norm2" in name:
						a__       :						Dict							=  name.replace("""norm2"""						,       """layernorm_after""")
			if "mlp.fc1" in name:
						a__       :						Any							=  name.replace("""mlp.fc1"""						,       """intermediate.dense""")
			if "mlp.fc2" in name:
						a__       :						Dict							=  name.replace("""mlp.fc2"""						,       """output.dense""")
			if "q_bias" in name:
						a__       :						str							=  name.replace("""q_bias"""						,       """query.bias""")
			if "k_bias" in name:
						a__       :						Optional[int]							=  name.replace("""k_bias"""						,       """key.bias""")
			if "v_bias" in name:
						a__       :						str							=  name.replace("""v_bias"""						,       """value.bias""")
			if "cpb_mlp" in name:
						a__       :						Optional[int]							=  name.replace("""cpb_mlp"""						,       """continuous_position_bias_mlp""")
			if "patch_embed.proj" in name:
						a__       :						Any							=  name.replace("""patch_embed.proj"""						,       """patch_embed.projection""")

			if name == "norm.weight":
						a__       :						Optional[int]							=  """layernorm.weight"""
			if name == "norm.bias":
						a__       :						List[str]							=  """layernorm.bias"""

			if "conv_first" in name:
						a__       :						List[Any]							=  name.replace("""conv_first"""						,       """first_convolution""")

			if (
			    "upsample" in name
			    or "conv_before_upsample" in name
			    or "conv_bicubic" in name
			    or "conv_up" in name
			    or "conv_hr" in name
			    or "conv_last" in name
			    or "aux" in name
			):
						# heads
						if "conv_last" in name:
									a__       :						Dict							=  name.replace("""conv_last"""						,       """final_convolution""")
						if config.upsampler in ["pixelshuffle", "pixelshuffle_aux", "nearest+conv"]:
									if "conv_before_upsample.0" in name:
												a__       :						Any							=  name.replace("""conv_before_upsample.0"""						,       """conv_before_upsample""")
									if "upsample.0" in name:
												a__       :						Any							=  name.replace("""upsample.0"""						,       """upsample.convolution_0""")
									if "upsample.2" in name:
												a__       :						Optional[Any]							=  name.replace("""upsample.2"""						,       """upsample.convolution_1""")
									a__       :						List[str]							=  """upsample.""" + name
						elif config.upsampler == "pixelshuffledirect":
									a__       :						List[Any]							=  name.replace("""upsample.0.weight"""						,       """upsample.conv.weight""")
									a__       :						List[Any]							=  name.replace("""upsample.0.bias"""						,       """upsample.conv.bias""")
						else:
									pass
			else:
						a__       :						Tuple							=  """swin2sr.""" + name

			return name

def      lowerCAmelCase_					(  _lowercase    :       List[Any]						,       _lowercase    :       List[Any])							->	Tuple:







			"""simple docstring"""
			for key in orig_state_dict.copy().keys():
						a__       :						int							=  orig_state_dict.pop(_lowercase)

						if "qkv" in key:
									a__       :						Tuple							=  key.split(""".""")
									a__       :						Any							=  int(key_split[1])
									a__       :						Tuple							=  int(key_split[4])
									a__       :						List[Any]							=  config.embed_dim

									if "weight" in key:
												a__       :						int							=  val[:dim, :]
												a__       :						str							=  val[dim : dim * 2, :]
												a__       :						Tuple							=  val[-dim:, :]
									else:
												a__       :						Dict							=  val[:dim]
												a__       :						Optional[Any]							=  val[dim : dim * 2]
												a__       :						Optional[Any]							=  val[-dim:]
									pass
						else:
									a__       :						Optional[Any]							=  val

			return orig_state_dict

def      lowerCAmelCase_					(  _lowercase    :       Any						,       _lowercase    :       List[Any]						,       _lowercase    :       Union[str, Any])							->	Union[str, Any]:







			"""simple docstring"""
			a__       :						List[Any]							=  get_config(_lowercase)
			a__       :						Optional[Any]							=  SwinaSRForImageSuperResolution(_lowercase)
			model.eval()

			a__       :						List[Any]							=  torch.hub.load_state_dict_from_url(_lowercase						,       map_location="""cpu""")
			a__       :						Any							=  convert_state_dict(_lowercase						,       _lowercase)
			a__	,						a__       :						Tuple							=  model.load_state_dict(_lowercase						,       strict=_lowercase)

			if len(_lowercase) > 0:
						raise ValueError("""Missing keys when converting: {}""".format(_lowercase))
			for key in unexpected_keys:
						if not ("relative_position_index" in key or "relative_coords_table" in key or "self_mask" in key):
									raise ValueError(F'''Unexpected key {key} in state_dict''')

    # verify values
			a__       :						int							=  """https://github.com/mv-lab/swin2sr/blob/main/testsets/real-inputs/shanghai.jpg?raw=true"""
			a__       :						str							=  Image.open(requests.get(_lowercase						,       stream=_lowercase).raw).convert("""RGB""")
			a__       :						Any							=  SwinaSRImageProcessor()
			# pixel_values = processor(image, return_tensors="pt").pixel_values

			a__       :						str							=  126 if """Jpeg""" in checkpoint_url else 256
			a__       :						int							=  Compose(
			    [
			        Resize((image_size, image_size)),
			        ToTensor(),
			        Normalize(mean=[0.485, 0.456, 0.406]						,       std=[0.229, 0.224, 0.225]),
			    ])
			a__       :						Optional[Any]							=  transforms(_lowercase).unsqueeze(0)

			if config.num_channels == 1:
						a__       :						Tuple							=  pixel_values[:, 0, :, :].unsqueeze(1)

			a__       :						List[str]							=  model(_lowercase)

			# assert values
			if "Swin2SR_ClassicalSR_X2_64" in checkpoint_url:
						a__       :						Optional[Any]							=  torch.Size([1, 3, 512, 512])
						a__       :						str							=  torch.tensor(
						    [[-0.7087, -0.7138, -0.6721], [-0.8340, -0.8095, -0.7298], [-0.9149, -0.8414, -0.7940]])
			elif "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
						a__       :						Tuple							=  torch.Size([1, 3, 1024, 1024])
						a__       :						Any							=  torch.tensor(
						    [[-0.7775, -0.8105, -0.8933], [-0.7764, -0.8356, -0.9225], [-0.7976, -0.8686, -0.9579]])
			elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
						# TODO values didn't match exactly here
						a__       :						str							=  torch.Size([1, 3, 1024, 1024])
						a__       :						str							=  torch.tensor(
						    [[-0.8035, -0.7504, -0.7491], [-0.8538, -0.8124, -0.7782], [-0.8804, -0.8651, -0.8493]])
			elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
						a__       :						List[Any]							=  torch.Size([1, 3, 512, 512])
						a__       :						Any							=  torch.tensor(
						    [[-0.7669, -0.8662, -0.8767], [-0.8810, -0.9962, -0.9820], [-0.9340, -1.0322, -1.1149]])
			elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
						a__       :						Any							=  torch.Size([1, 3, 1024, 1024])
						a__       :						Tuple							=  torch.tensor(
						    [[-0.5238, -0.5557, -0.6321], [-0.6016, -0.5903, -0.6391], [-0.6244, -0.6334, -0.6889]])

			assert (
			    outputs.reconstruction.shape == expected_shape
			), F'''Shape of reconstruction should be {expected_shape}, but is {outputs.reconstruction.shape}'''
			assert torch.allclose(outputs.reconstruction[0, 0, :3, :3]						,       _lowercase						,       atol=1e-3)
			print("""Looks ok!""")

			a__       :						List[Any]							=  {
			    """https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_ClassicalSR_X2_64.pth""": (
			        """swin2SR-classical-sr-x2-64"""
			    ),
			    """https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_ClassicalSR_X4_64.pth""": (
			        """swin2SR-classical-sr-x4-64"""
			    ),
			    """https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_CompressedSR_X4_48.pth""": (
			        """swin2SR-compressed-sr-x4-48"""
			    ),
			    """https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_Lightweight_X2_64.pth""": (
			        """swin2SR-lightweight-x2-64"""
			    ),
			    """https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR.pth""": (
			        """swin2SR-realworld-sr-x4-64-bsrgan-psnr"""
			    ),
			}
			a__       :						Dict							=  url_to_name[checkpoint_url]

			if pytorch_dump_folder_path is not None:
						print(F'''Saving model {model_name} to {pytorch_dump_folder_path}''')
						model.save_pretrained(_lowercase)
						print(F'''Saving image processor to {pytorch_dump_folder_path}''')
						processor.save_pretrained(_lowercase)

			if push_to_hub:
						model.push_to_hub(F'''caidas/{model_name}''')
						processor.push_to_hub(F'''caidas/{model_name}''')


if __name__ == "__main__":
		_lowercase		:				Any      		=argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    "--checkpoint_url",
		    default="https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_ClassicalSR_X2_64.pth",
		    type=str,
		    help="URL of the original Swin2SR 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."
		)
		parser.add_argument("--push_to_hub", action="store_true", help="Whether to push the converted model to the hub.")

		_lowercase		:				Dict      		=parser.parse_args()
		convert_swinasr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)









import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
_lowercase		:				Union[str, Any]      		=data_utils.TransfoXLTokenizer
_lowercase		:				Tuple      		=data_utils.TransfoXLCorpus
_lowercase		:				Optional[int]      		=data_utils
_lowercase		:				int      		=data_utils

def      lowerCAmelCase_					(  _lowercase    :       Optional[Any]						,       _lowercase    :       List[Any]						,       _lowercase    :       Optional[Any]						,       _lowercase    :       Optional[Any])							->	List[Any]:







			"""simple docstring"""
			if transfo_xl_dataset_file:
						# Convert a pre-processed corpus (see original TensorFlow repo)
						with open(_lowercase						,       """rb""") as fp:
									a__       :						Optional[Any]							=  pickle.load(_lowercase						,       encoding="""latin1""")
						# Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
						a__       :						Dict							=  pytorch_dump_folder_path + """/""" + VOCAB_FILES_NAMES["""pretrained_vocab_file"""]
						print(F'''Save vocabulary to {pytorch_vocab_dump_path}''')
						a__       :						List[str]							=  corpus.vocab.__dict__
						torch.save(_lowercase						,       _lowercase)

						a__       :						Optional[int]							=  corpus.__dict__
						corpus_dict_no_vocab.pop("""vocab"""						,       _lowercase)
						a__       :						Tuple							=  pytorch_dump_folder_path + """/""" + CORPUS_NAME
						print(F'''Save dataset to {pytorch_dataset_dump_path}''')
						torch.save(_lowercase						,       _lowercase)

			if tf_checkpoint_path:
						# Convert a pre-trained TensorFlow model
						a__       :						Optional[Any]							=  os.path.abspath(_lowercase)
						a__       :						str							=  os.path.abspath(_lowercase)

						print(F'''Converting Transformer XL checkpoint from {tf_path} with config at {config_path}.''')
						# Initialise PyTorch model
						if transfo_xl_config_file == "":
									a__       :						List[str]							=  TransfoXLConfig()
						else:
									a__       :						Any							=  TransfoXLConfig.from_json_file(_lowercase)
						print(F'''Building PyTorch model from configuration: {config}''')
						a__       :						Any							=  TransfoXLLMHeadModel(_lowercase)

						a__       :						Optional[Any]							=  load_tf_weights_in_transfo_xl(_lowercase						,       _lowercase						,       _lowercase)
						# Save pytorch-model
						a__       :						Tuple							=  os.path.join(_lowercase						,       _lowercase)
						a__       :						Optional[int]							=  os.path.join(_lowercase						,       _lowercase)
						print(F'''Save PyTorch model to {os.path.abspath(_lowercase)}''')
						torch.save(model.state_dict()						,       _lowercase)
						print(F'''Save configuration file to {os.path.abspath(_lowercase)}''')
						with open(_lowercase						,       """w"""						,       encoding="""utf-8""") as f:
									f.write(config.to_json_string())


if __name__ == "__main__":
		_lowercase		:				Tuple      		=argparse.ArgumentParser()
		parser.add_argument(
		    "--pytorch_dump_folder_path",
		    default=None,
		    type=str,
		    required=True,
		    help="Path to the folder to store the PyTorch model or dataset/vocab.",
		)
		parser.add_argument(
		    "--tf_checkpoint_path",
		    default="",
		    type=str,
		    help="An optional path to a TensorFlow checkpoint path to be converted.",
		)
		parser.add_argument(
		    "--transfo_xl_config_file",
		    default="",
		    type=str,
		    help=(
		        "An optional config json file corresponding to the pre-trained BERT model. \n"
		        "This specifies the model architecture."
		    ),
		)
		parser.add_argument(
		    "--transfo_xl_dataset_file",
		    default="",
		    type=str,
		    help="An optional dataset file to be converted in a vocabulary.",
		)
		_lowercase		:				Dict      		=parser.parse_args()
		convert_transfo_xl_checkpoint_to_pytorch(
		    args.tf_checkpoint_path,
		    args.transfo_xl_config_file,
		    args.pytorch_dump_folder_path,
		    args.transfo_xl_dataset_file,
		)




