





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     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	)       ->     List[str]:
  lowercase__			    =						SwinaSRConfig()

  if "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
    lowercase__			    =						4
  elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
    lowercase__			    =						4
    lowercase__			    =						48
    lowercase__			    =						'pixelshuffle_aux'
  elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
    lowercase__			    =						[6, 6, 6, 6]
    lowercase__			    =						60
    lowercase__			    =						[6, 6, 6, 6]
    lowercase__			    =						'pixelshuffledirect'
  elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
    lowercase__			    =						4
    lowercase__			    =						'nearest+conv'
  elif "Swin2SR_Jpeg_dynamic" in checkpoint_url:
    lowercase__			    =						1
    lowercase__			    =						1
    lowercase__			    =						126
    lowercase__			    =						7
    lowercase__			    =						2_5_5.0
    lowercase__			    =						''

  return config




def     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)       ->     str:
  if "patch_embed.proj" in name and "layers" not in name:
    lowercase__			    =						name.replace('patch_embed.proj'	,		'embeddings.patch_embeddings.projection'	)
  if "patch_embed.norm" in name:
    lowercase__			    =						name.replace('patch_embed.norm'	,		'embeddings.patch_embeddings.layernorm'	)
  if "layers" in name:
    lowercase__			    =						name.replace('layers'	,		'encoder.stages'	)
  if "residual_group.blocks" in name:
    lowercase__			    =						name.replace('residual_group.blocks'	,		'layers'	)
  if "attn.proj" in name:
    lowercase__			    =						name.replace('attn.proj'	,		'attention.output.dense'	)
  if "attn" in name:
    lowercase__			    =						name.replace('attn'	,		'attention.self'	)
  if "norm1" in name:
    lowercase__			    =						name.replace('norm1'	,		'layernorm_before'	)
  if "norm2" in name:
    lowercase__			    =						name.replace('norm2'	,		'layernorm_after'	)
  if "mlp.fc1" in name:
    lowercase__			    =						name.replace('mlp.fc1'	,		'intermediate.dense'	)
  if "mlp.fc2" in name:
    lowercase__			    =						name.replace('mlp.fc2'	,		'output.dense'	)
  if "q_bias" in name:
    lowercase__			    =						name.replace('q_bias'	,		'query.bias'	)
  if "k_bias" in name:
    lowercase__			    =						name.replace('k_bias'	,		'key.bias'	)
  if "v_bias" in name:
    lowercase__			    =						name.replace('v_bias'	,		'value.bias'	)
  if "cpb_mlp" in name:
    lowercase__			    =						name.replace('cpb_mlp'	,		'continuous_position_bias_mlp'	)
  if "patch_embed.proj" in name:
    lowercase__			    =						name.replace('patch_embed.proj'	,		'patch_embed.projection'	)

  if name == "norm.weight":
    lowercase__			    =						'layernorm.weight'
  if name == "norm.bias":
    lowercase__			    =						'layernorm.bias'

  if "conv_first" in name:
    lowercase__			    =						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:
      lowercase__			    =						name.replace('conv_last'	,		'final_convolution'	)
    if config.upsampler in ["pixelshuffle", "pixelshuffle_aux", "nearest+conv"]:
      if "conv_before_upsample.0" in name:
        lowercase__			    =						name.replace('conv_before_upsample.0'	,		'conv_before_upsample'	)
      if "upsample.0" in name:
        lowercase__			    =						name.replace('upsample.0'	,		'upsample.convolution_0'	)
      if "upsample.2" in name:
        lowercase__			    =						name.replace('upsample.2'	,		'upsample.convolution_1'	)
      lowercase__			    =						'upsample.' + name
    elif config.upsampler == "pixelshuffledirect":
      lowercase__			    =						name.replace('upsample.0.weight'	,		'upsample.conv.weight'	)
      lowercase__			    =						name.replace('upsample.0.bias'	,		'upsample.conv.bias'	)
    else:
      pass
  else:
    lowercase__			    =						'swin2sr.' + name

  return name




def     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)       ->     Union[str, Any]:
  for key in orig_state_dict.copy().keys():
    lowercase__			    =						orig_state_dict.pop(SCREAMING_SNAKE_CASE_	)

    if "qkv" in key:
      lowercase__			    =						key.split('.'	)
      lowercase__			    =						int(key_split[1]	)
      lowercase__			    =						int(key_split[4]	)
      lowercase__			    =						config.embed_dim

      if "weight" in key:
        lowercase__			    =						val[:dim, :]
        lowercase__			    =						val[dim : dim * 2, :]
        lowercase__			    =						val[-dim:, :]
      else:
        lowercase__			    =						val[:dim]
        lowercase__			    =						val[dim : dim * 2]
        lowercase__			    =						val[-dim:]
      pass
    else:
      lowercase__			    =						val

  return orig_state_dict




def     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)       ->     List[Any]:
  lowercase__			    =						get_config(SCREAMING_SNAKE_CASE_	)
  lowercase__			    =						SwinaSRForImageSuperResolution(SCREAMING_SNAKE_CASE_	)
  model.eval()

  lowercase__			    =						torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_	,		map_location='cpu'	)
  lowercase__			    =						convert_state_dict(SCREAMING_SNAKE_CASE_	,		SCREAMING_SNAKE_CASE_	)
  lowercase__    ,					lowercase__			    =						model.load_state_dict(SCREAMING_SNAKE_CASE_	,		strict=SCREAMING_SNAKE_CASE_	)

  if len(SCREAMING_SNAKE_CASE_	) > 0:
    raise ValueError('Missing keys when converting: {}'.format(SCREAMING_SNAKE_CASE_	)	)
  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
  lowercase__			    =						'https://github.com/mv-lab/swin2sr/blob/main/testsets/real-inputs/shanghai.jpg?raw=true'
  lowercase__			    =						Image.open(requests.get(SCREAMING_SNAKE_CASE_	,		stream=SCREAMING_SNAKE_CASE_	).raw	).convert('RGB'	)
  lowercase__			    =						SwinaSRImageProcessor()
  # pixel_values = processor(image, return_tensors="pt").pixel_values

  lowercase__			    =						126 if 'Jpeg' in checkpoint_url else 256
  lowercase__			    =						Compose(
      [
          Resize((image_size, image_size)	),
          ToTensor(),
          Normalize(mean=[0.4_8_5, 0.4_5_6, 0.4_0_6]	,		std=[0.2_2_9, 0.2_2_4, 0.2_2_5]	),
      ]	)
  lowercase__			    =						transforms(SCREAMING_SNAKE_CASE_	).unsqueeze(0	)

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

  lowercase__			    =						model(SCREAMING_SNAKE_CASE_	)

  # assert values
  if "Swin2SR_ClassicalSR_X2_64" in checkpoint_url:
    lowercase__			    =						torch.Size([1, 3, 512, 512]	)
    lowercase__			    =						torch.tensor(
        [[-0.7_0_8_7, -0.7_1_3_8, -0.6_7_2_1], [-0.8_3_4_0, -0.8_0_9_5, -0.7_2_9_8], [-0.9_1_4_9, -0.8_4_1_4, -0.7_9_4_0]]	)
  elif "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
    lowercase__			    =						torch.Size([1, 3, 1024, 1024]	)
    lowercase__			    =						torch.tensor(
        [[-0.7_7_7_5, -0.8_1_0_5, -0.8_9_3_3], [-0.7_7_6_4, -0.8_3_5_6, -0.9_2_2_5], [-0.7_9_7_6, -0.8_6_8_6, -0.9_5_7_9]]	)
  elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
    # TODO values didn't match exactly here
    lowercase__			    =						torch.Size([1, 3, 1024, 1024]	)
    lowercase__			    =						torch.tensor(
        [[-0.8_0_3_5, -0.7_5_0_4, -0.7_4_9_1], [-0.8_5_3_8, -0.8_1_2_4, -0.7_7_8_2], [-0.8_8_0_4, -0.8_6_5_1, -0.8_4_9_3]]	)
  elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
    lowercase__			    =						torch.Size([1, 3, 512, 512]	)
    lowercase__			    =						torch.tensor(
        [[-0.7_6_6_9, -0.8_6_6_2, -0.8_7_6_7], [-0.8_8_1_0, -0.9_9_6_2, -0.9_8_2_0], [-0.9_3_4_0, -1.0_3_2_2, -1.1_1_4_9]]	)
  elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
    lowercase__			    =						torch.Size([1, 3, 1024, 1024]	)
    lowercase__			    =						torch.tensor(
        [[-0.5_2_3_8, -0.5_5_5_7, -0.6_3_2_1], [-0.6_0_1_6, -0.5_9_0_3, -0.6_3_9_1], [-0.6_2_4_4, -0.6_3_3_4, -0.6_8_8_9]]	)

  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]	,		SCREAMING_SNAKE_CASE_	,		atol=1E-3	)
  print('Looks ok!'	)

  lowercase__			    =						{
      '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'
      ),
  }
  lowercase__			    =						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(SCREAMING_SNAKE_CASE_	)
    print(F"""Saving image processor to {pytorch_dump_folder_path}"""	)
    processor.save_pretrained(SCREAMING_SNAKE_CASE_	)

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


if __name__ == "__main__":
 lowercase_					     =	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_					     =	parser.parse_args()
 convert_swinasr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)









from string import ascii_uppercase

lowercase_					     =	{str(ord(c) - 55): c for c in ascii_uppercase}




def     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)       ->     str:
  if isinstance(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	):
    raise TypeError('int() can\'t convert non-string with explicit base'	)
  if num < 0:
    raise ValueError('parameter must be positive int'	)
  if isinstance(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	):
    raise TypeError('\'str\' object cannot be interpreted as an integer'	)
  if isinstance(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	):
    raise TypeError('\'float\' object cannot be interpreted as an integer'	)
  if base in (0, 1):
    raise ValueError('base must be >= 2'	)
  if base > 36:
    raise ValueError('base must be <= 36'	)
  lowercase__			    =						''
  lowercase__			    =						0
  lowercase__			    =						0
  while div != 1:
    lowercase__    ,					lowercase__			    =						divmod(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)
    if base >= 11 and 9 < mod < 36:
      lowercase__			    =						ALPHABET_VALUES[str(_SCREAMING_SNAKE_CASE	)]
    else:
      lowercase__			    =						str(_SCREAMING_SNAKE_CASE	)
    new_value += actual_value
    lowercase__			    =						num // base
    lowercase__			    =						div
    if div == 0:
      return str(new_value[::-1]	)
    elif div == 1:
      new_value += str(_SCREAMING_SNAKE_CASE	)
      return str(new_value[::-1]	)

  return new_value[::-1]


if __name__ == "__main__":
      import doctest

      doctest.testmod()

      for base in range(2, 37):
            for num in range(1_000):
                  assert int(decimal_to_any(num, base), base) == num, (
                      num,
                      base,
                      decimal_to_any(num, base),
                      int(decimal_to_any(num, base), base),
                  )


