





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	(	a    : str     ) ->List[str]:
 snake_case				=			SwinaSRConfig()

 if "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
  snake_case				=			4
 elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
  snake_case				=			4
  snake_case				=			48
  snake_case				=			"""pixelshuffle_aux"""
 elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
  snake_case				=			[6, 6, 6, 6]
  snake_case				=			60
  snake_case				=			[6, 6, 6, 6]
  snake_case				=			"""pixelshuffledirect"""
 elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
  snake_case				=			4
  snake_case				=			"""nearest+conv"""
 elif "Swin2SR_Jpeg_dynamic" in checkpoint_url:
  snake_case				=			1
  snake_case				=			1
  snake_case				=			126
  snake_case				=			7
  snake_case				=			255.0
  snake_case				=			""""""

 return config

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

 if name == "norm.weight":
  snake_case				=			"""layernorm.weight"""
 if name == "norm.bias":
  snake_case				=			"""layernorm.bias"""

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

 return name

def    __UpperCamelCase	(	a    : Optional[int]					,   a    : Any     ) ->Tuple:
 for key in orig_state_dict.copy().keys():
  snake_case				=			orig_state_dict.pop(__lowerCAmelCase     )

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

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

 return orig_state_dict

def    __UpperCamelCase	(	a    : Optional[Any]					,   a    : Optional[Any]					,   a    : List[Any]     ) ->Optional[Any]:
 snake_case				=			get_config(__lowerCAmelCase     )
 snake_case				=			SwinaSRForImageSuperResolution(__lowerCAmelCase     )
 model.eval()

 snake_case				=			torch.hub.load_state_dict_from_url(__lowerCAmelCase					,   map_location='''cpu'''     )
 snake_case				=			convert_state_dict(__lowerCAmelCase					,   __lowerCAmelCase     )
 snake_case				=			model.load_state_dict(__lowerCAmelCase					,   strict=__lowerCAmelCase     )

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

 snake_case				=			126 if """Jpeg""" in checkpoint_url else 256
 snake_case				=			Compose(
     [
         Resize((image_size, image_size)     ),
         ToTensor(),
         Normalize(mean=[0.485, 0.456, 0.406]					,   std=[0.229, 0.224, 0.225]     ),
     ]     )
 snake_case				=			transforms(__lowerCAmelCase     ).unsqueeze(0     )

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

 snake_case				=			model(__lowerCAmelCase     )

 # assert values
 if "Swin2SR_ClassicalSR_X2_64" in checkpoint_url:
  snake_case				=			torch.Size([1, 3, 512, 512]     )
  snake_case				=			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:
  snake_case				=			torch.Size([1, 3, 1024, 1024]     )
  snake_case				=			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
  snake_case				=			torch.Size([1, 3, 1024, 1024]     )
  snake_case				=			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:
  snake_case				=			torch.Size([1, 3, 512, 512]     )
  snake_case				=			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:
  snake_case				=			torch.Size([1, 3, 1024, 1024]     )
  snake_case				=			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]					,   __lowerCAmelCase					,   atol=1e-3     )
 print('''Looks ok!'''     )

 snake_case				=			{
     """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"""
     ),
 }
 snake_case				=			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(__lowerCAmelCase     )
  print(f"""Saving image processor to {pytorch_dump_folder_path}"""     )
  processor.save_pretrained(__lowerCAmelCase     )

 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)








'''simple docstring'''







from ...processing_utils import ProcessorMixin
class 			_lowercase						(       __a					):
     _UpperCAmelCase										=					'''WhisperFeatureExtractor'''
     _UpperCAmelCase										=					'''WhisperTokenizer'''







     def __init__( self    ,   A__    ,   A__    ) ->	Optional[Any]:
      super().__init__(A__    ,   A__    )
      snake_case				=			self.feature_extractor
      snake_case				=			False







     def     UpperCamelCase			( self    ,   A__=None    ,   A__=None    ,   A__=True    ) ->	Union[str, Any]:
      return self.tokenizer.get_decoder_prompt_ids(task=A__    ,   language=A__    ,   no_timestamps=A__    )







     def __call__( self    ,   *A__    ,   **A__    ) ->	Dict:
      # For backward compatibility
      if self._in_target_context_manager:
       return self.current_processor(*A__    ,   **A__    )

      snake_case				=			kwargs.pop('''audio'''    ,   A__    )
      snake_case				=			kwargs.pop('''sampling_rate'''    ,   A__    )
      snake_case				=			kwargs.pop('''text'''    ,   A__    )
      if len(A__    ) > 0:
       snake_case				=			args[0]
       snake_case				=			args[1:]

      if audio is None and text is None:
       raise ValueError('''You need to specify either an `audio` or `text` input to process.'''    )

      if audio is not None:
       snake_case				=			self.feature_extractor(A__    ,   *A__    ,   sampling_rate=A__    ,   **A__    )
      if text is not None:
       snake_case				=			self.tokenizer(A__    ,   **A__    )

      if text is None:
       return inputs

      elif audio is None:
       return encodings
      else:
       snake_case				=			encodings['''input_ids''']
       return inputs







     def     UpperCamelCase			( self    ,   *A__    ,   **A__    ) ->	Optional[Any]:
      return self.tokenizer.batch_decode(*A__    ,   **A__    )







     def     UpperCamelCase			( self    ,   *A__    ,   **A__    ) ->	str:
      return self.tokenizer.decode(*A__    ,   **A__    )







     def     UpperCamelCase			( self    ,   A__    ,   A__="np"    ) ->	Optional[Any]:
      return self.tokenizer.get_prompt_ids(A__    ,   return_tensors=A__    )

