


'''simple docstring'''


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_  (						__lowerCamelCase : Union[str, Any]			):
   lowercase_      :Optional[Any]         =			SwinaSRConfig()

   if "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
      lowercase_      :List[Any]         =			4
   elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
      lowercase_      :List[Any]         =			4
      lowercase_      :Any         =			48
      lowercase_      :str         =			"pixelshuffle_aux"
   elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
      lowercase_      :Optional[int]         =			[6, 6, 6, 6]
      lowercase_      :List[Any]         =			60
      lowercase_      :Union[str, Any]         =			[6, 6, 6, 6]
      lowercase_      :str         =			"pixelshuffledirect"
   elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
      lowercase_      :Optional[int]         =			4
      lowercase_      :Dict         =			"nearest+conv"
   elif "Swin2SR_Jpeg_dynamic" in checkpoint_url:
      lowercase_      :List[Any]         =			1
      lowercase_      :str         =			1
      lowercase_      :int         =			1_26
      lowercase_      :List[str]         =			7
      lowercase_      :Optional[Any]         =			2_55.0
      lowercase_      :Any         =			""

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

   if name == "norm.weight":
      lowercase_      :Dict         =			"layernorm.weight"
   if name == "norm.bias":
      lowercase_      :Any         =			"layernorm.bias"

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

   return name
def 							UpperCAmelCase_  (						__lowerCamelCase : int     ,__lowerCamelCase : Optional[Any]			):
   for key in orig_state_dict.copy().keys():
      lowercase_      :Optional[Any]         =			orig_state_dict.pop(lowercase__			)

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

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

   return orig_state_dict
def 							UpperCAmelCase_  (						__lowerCamelCase : int     ,__lowerCamelCase : Optional[Any]     ,__lowerCamelCase : Optional[Any]			):
   lowercase_      :Any         =			get_config(lowercase__			)
   lowercase_      :List[Any]         =			SwinaSRForImageSuperResolution(lowercase__			)
   model.eval()

   lowercase_      :Union[str, Any]         =			torch.hub.load_state_dict_from_url(lowercase__     ,map_location="cpu"			)
   lowercase_      :Any         =			convert_state_dict(lowercase__     ,lowercase__			)
   lowercase_      ,							lowercase_      :Dict         =			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
   lowercase_      :int         =			"https://github.com/mv-lab/swin2sr/blob/main/testsets/real-inputs/shanghai.jpg?raw=true"
   lowercase_      :Any         =			Image.open(requests.get(lowercase__     ,stream=lowercase__			).raw			).convert("RGB"			)
   lowercase_      :Optional[Any]         =			SwinaSRImageProcessor()
   # pixel_values = processor(image, return_tensors="pt").pixel_values

   lowercase_      :Any         =			1_26 if "Jpeg" in checkpoint_url else 2_56
   lowercase_      :List[str]         =			Compose(
       [
           Resize((image_size, image_size)			),
           ToTensor(),
           Normalize(mean=[0.485, 0.456, 0.406]     ,std=[0.229, 0.224, 0.225]			),
       ]			)
   lowercase_      :Optional[Any]         =			transforms(lowercase__			).unsqueeze(0			)

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

   lowercase_      :Union[str, Any]         =			model(lowercase__			)

   # assert values
   if "Swin2SR_ClassicalSR_X2_64" in checkpoint_url:
      lowercase_      :int         =			torch.Size([1, 3, 5_12, 5_12]			)
      lowercase_      :Any         =			torch.tensor(
          [[-0.7_087, -0.7_138, -0.6_721], [-0.8_340, -0.8_095, -0.7_298], [-0.9_149, -0.8_414, -0.7_940]]			)
   elif "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
      lowercase_      :Dict         =			torch.Size([1, 3, 10_24, 10_24]			)
      lowercase_      :List[str]         =			torch.tensor(
          [[-0.7_775, -0.8_105, -0.8_933], [-0.7_764, -0.8_356, -0.9_225], [-0.7_976, -0.8_686, -0.9_579]]			)
   elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
      # TODO values didn't match exactly here
      lowercase_      :List[str]         =			torch.Size([1, 3, 10_24, 10_24]			)
      lowercase_      :Dict         =			torch.tensor(
          [[-0.8_035, -0.7_504, -0.7_491], [-0.8_538, -0.8_124, -0.7_782], [-0.8_804, -0.8_651, -0.8_493]]			)
   elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
      lowercase_      :Union[str, Any]         =			torch.Size([1, 3, 5_12, 5_12]			)
      lowercase_      :str         =			torch.tensor(
          [[-0.7_669, -0.8_662, -0.8_767], [-0.8_810, -0.9_962, -0.9_820], [-0.9_340, -1.0_322, -1.1_149]]			)
   elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
      lowercase_      :Union[str, Any]         =			torch.Size([1, 3, 10_24, 10_24]			)
      lowercase_      :Optional[Any]         =			torch.tensor(
          [[-0.5_238, -0.5_557, -0.6_321], [-0.6_016, -0.5_903, -0.6_391], [-0.6_244, -0.6_334, -0.6_889]]			)

   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!"			)

   lowercase_      :int         =			{
       "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_      :Tuple         =			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__":
 lowerCAmelCase :					Tuple 							=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.''')

 lowerCAmelCase :					Tuple 							=parser.parse_args()
 convert_swinasr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)









'''simple docstring'''


from typing import Dict, List, Optional, Union

import numpy as np

from transformers.utils import is_vision_available
from transformers.utils.generic import TensorType

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    is_valid_image,
    to_numpy_array,
    valid_images,
)
from ...utils import logging


if is_vision_available():
 import PIL

lowerCAmelCase :					Any 							=logging.get_logger(__name__)
def 							UpperCAmelCase_  (						__lowerCamelCase : List[str]			):
   if isinstance(__lowerCamelCase     ,(list, tuple)			) and isinstance(videos[0]     ,(list, tuple)			) and is_valid_image(videos[0][0]			):
      return videos

   elif isinstance(__lowerCamelCase     ,(list, tuple)			) and is_valid_image(videos[0]			):
      return [videos]

   elif is_valid_image(__lowerCamelCase			):
      return [[videos]]

   raise ValueError(F'Could not make batched video from {videos}'			)


class 				a_       (    _lowerCAmelCase							):
     __A								=      ["pixel_values"]





     def __init__(			self		: List[str]     ,				lowercase		: bool = True     ,				lowercase		: Dict[str, int] = None     ,				lowercase		: PILImageResampling = PILImageResampling.BILINEAR     ,				lowercase		: bool = True     ,				lowercase		: Dict[str, int] = None     ,				lowercase		: bool = True     ,				lowercase		: Union[int, float] = 1 / 255     ,				lowercase		: bool = True     ,				lowercase		: bool = True     ,				lowercase		: Optional[Union[float, List[float]]] = None     ,				lowercase		: Optional[Union[float, List[float]]] = None     ,				**lowercase		: Tuple     ,				):





        """simple docstring"""




        super().__init__(**lowercase						)
        lowercase_      :Any         =			size if size is not None else {"shortest_edge": 256}
        lowercase_      :int         =			get_size_dict(lowercase     ,				default_to_square=lowercase						)
        lowercase_      :str         =			crop_size if crop_size is not None else {"height": 224, "width": 224}
        lowercase_      :List[str]         =			get_size_dict(lowercase     ,				param_name="crop_size"						)

        lowercase_      :List[str]         =			do_resize
        lowercase_      :Any         =			size
        lowercase_      :Union[str, Any]         =			do_center_crop
        lowercase_      :Union[str, Any]         =			crop_size
        lowercase_      :Optional[Any]         =			resample
        lowercase_      :List[str]         =			do_rescale
        lowercase_      :List[Any]         =			rescale_factor
        lowercase_      :Dict         =			offset
        lowercase_      :Optional[Any]         =			do_normalize
        lowercase_      :Dict         =			image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
        lowercase_      :Optional[int]         =			image_std if image_std is not None else IMAGENET_STANDARD_STD





     def 							lowercase__				(			self		: Union[str, Any]     ,				lowercase		: np.ndarray     ,				lowercase		: Dict[str, int]     ,				lowercase		: PILImageResampling = PILImageResampling.BILINEAR     ,				lowercase		: Optional[Union[str, ChannelDimension]] = None     ,				**lowercase		: Optional[Any]     ,				):





        """simple docstring"""




        lowercase_      :List[Any]         =			get_size_dict(lowercase     ,				default_to_square=lowercase						)
        if "shortest_edge" in size:
           lowercase_      :int         =			get_resize_output_image_size(lowercase     ,				size["shortest_edge"]     ,				default_to_square=lowercase						)
        elif "height" in size and "width" in size:
           lowercase_      :Union[str, Any]         =			(size["height"], size["width"])
        else:
           raise ValueError(F'Size must have \'height\' and \'width\' or \'shortest_edge\' as keys. Got {size.keys()}'						)
        return resize(lowercase     ,				size=lowercase     ,				resample=lowercase     ,				data_format=lowercase     ,				**lowercase						)





     def 							lowercase__				(			self		: str     ,				lowercase		: np.ndarray     ,				lowercase		: Dict[str, int]     ,				lowercase		: Optional[Union[str, ChannelDimension]] = None     ,				**lowercase		: str     ,				):





        """simple docstring"""




        lowercase_      :Any         =			get_size_dict(lowercase						)
        if "height" not in size or "width" not in size:
           raise ValueError(F'Size must have \'height\' and \'width\' as keys. Got {size.keys()}'						)
        return center_crop(lowercase     ,				size=(size["height"], size["width"])     ,				data_format=lowercase     ,				**lowercase						)





     def 							lowercase__				(			self		: List[str]     ,				lowercase		: np.ndarray     ,				lowercase		: Union[int, float]     ,				lowercase		: bool = True     ,				lowercase		: Optional[Union[str, ChannelDimension]] = None     ,				**lowercase		: List[str]     ,				):





        """simple docstring"""




        lowercase_      :List[str]         =			image.astype(np.floataa						)
        if offset:
           lowercase_      :List[str]         =			image - (scale / 2)
        return rescale(lowercase     ,				scale=lowercase     ,				data_format=lowercase     ,				**lowercase						)





     def 							lowercase__				(			self		: Tuple     ,				lowercase		: np.ndarray     ,				lowercase		: Union[float, List[float]]     ,				lowercase		: Union[float, List[float]]     ,				lowercase		: Optional[Union[str, ChannelDimension]] = None     ,				**lowercase		: Dict     ,				):





        """simple docstring"""




        return normalize(lowercase     ,				mean=lowercase     ,				std=lowercase     ,				data_format=lowercase     ,				**lowercase						)





     def 							lowercase__				(			self		: Tuple     ,				lowercase		: ImageInput     ,				lowercase		: bool = None     ,				lowercase		: Dict[str, int] = None     ,				lowercase		: PILImageResampling = None     ,				lowercase		: bool = None     ,				lowercase		: Dict[str, int] = None     ,				lowercase		: bool = None     ,				lowercase		: float = None     ,				lowercase		: bool = None     ,				lowercase		: bool = None     ,				lowercase		: Optional[Union[float, List[float]]] = None     ,				lowercase		: Optional[Union[float, List[float]]] = None     ,				lowercase		: Optional[ChannelDimension] = ChannelDimension.FIRST     ,				):





        """simple docstring"""




        if do_resize and size is None or resample is None:
           raise ValueError("Size and resample must be specified if do_resize is True."						)

        if do_center_crop and crop_size is None:
           raise ValueError("Crop size must be specified if do_center_crop is True."						)

        if do_rescale and rescale_factor is None:
           raise ValueError("Rescale factor must be specified if do_rescale is True."						)

        if do_normalize and (image_mean is None or image_std is None):
           raise ValueError("Image mean and std must be specified if do_normalize is True."						)

        if offset and not do_rescale:
           raise ValueError("For offset, do_rescale must also be set to True."						)

        # All transformations expect numpy arrays.
        lowercase_      :Optional[int]         =			to_numpy_array(lowercase						)

        if do_resize:
           lowercase_      :Tuple         =			self.resize(image=lowercase     ,				size=lowercase     ,				resample=lowercase						)

        if do_center_crop:
           lowercase_      :Any         =			self.center_crop(lowercase     ,				size=lowercase						)

        if do_rescale:
           lowercase_      :Optional[Any]         =			self.rescale(image=lowercase     ,				scale=lowercase     ,				offset=lowercase						)

        if do_normalize:
           lowercase_      :Tuple         =			self.normalize(image=lowercase     ,				mean=lowercase     ,				std=lowercase						)

        lowercase_      :Optional[Any]         =			to_channel_dimension_format(lowercase     ,				lowercase						)
        return image







     def 							lowercase__				(			self		: Dict     ,				lowercase		: ImageInput     ,				lowercase		: bool = None     ,				lowercase		: Dict[str, int] = None     ,				lowercase		: PILImageResampling = None     ,				lowercase		: bool = None     ,				lowercase		: Dict[str, int] = None     ,				lowercase		: bool = None     ,				lowercase		: float = None     ,				lowercase		: bool = None     ,				lowercase		: bool = None     ,				lowercase		: Optional[Union[float, List[float]]] = None     ,				lowercase		: Optional[Union[float, List[float]]] = None     ,				lowercase		: Optional[Union[str, TensorType]] = None     ,				lowercase		: ChannelDimension = ChannelDimension.FIRST     ,				**lowercase		: Optional[int]     ,				):





        """simple docstring"""




        lowercase_      :str         =			do_resize if do_resize is not None else self.do_resize
        lowercase_      :Optional[Any]         =			resample if resample is not None else self.resample
        lowercase_      :Tuple         =			do_center_crop if do_center_crop is not None else self.do_center_crop
        lowercase_      :Dict         =			do_rescale if do_rescale is not None else self.do_rescale
        lowercase_      :Optional[int]         =			rescale_factor if rescale_factor is not None else self.rescale_factor
        lowercase_      :Dict         =			offset if offset is not None else self.offset
        lowercase_      :Tuple         =			do_normalize if do_normalize is not None else self.do_normalize
        lowercase_      :int         =			image_mean if image_mean is not None else self.image_mean
        lowercase_      :Optional[int]         =			image_std if image_std is not None else self.image_std

        lowercase_      :int         =			size if size is not None else self.size
        lowercase_      :Optional[int]         =			get_size_dict(lowercase     ,				default_to_square=lowercase						)
        lowercase_      :List[Any]         =			crop_size if crop_size is not None else self.crop_size
        lowercase_      :List[str]         =			get_size_dict(lowercase     ,				param_name="crop_size"						)

        if not valid_images(lowercase						):
           raise ValueError(
               "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
               "torch.Tensor, tf.Tensor or jax.ndarray."						)

        lowercase_      :List[str]         =			make_batched(lowercase						)

        lowercase_      :List[Any]         =			[
            [
                self._preprocess_image(
                    image=lowercase     ,				do_resize=lowercase     ,				size=lowercase     ,				resample=lowercase     ,				do_center_crop=lowercase     ,				crop_size=lowercase     ,				do_rescale=lowercase     ,				rescale_factor=lowercase     ,				offset=lowercase     ,				do_normalize=lowercase     ,				image_mean=lowercase     ,				image_std=lowercase     ,				data_format=lowercase     ,				)
                for img in video
            ]
            for video in videos
        ]

        lowercase_      :Optional[int]         =			{"pixel_values": videos}
        return BatchFeature(data=lowercase     ,				tensor_type=lowercase						)


