



'''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   a_   (		_UpperCAmelCase    :      Optional[Any]				)    ->  List[Any]:
      __snake_case :   int													=    SwinaSRConfig()

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

      return config





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

      if name == "norm.weight":
            __snake_case :   str													=    'layernorm.weight'
      if name == "norm.bias":
            __snake_case :   Tuple													=    'layernorm.bias'

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

      return name





def   a_   (		_UpperCAmelCase    :      Tuple ,_UpperCAmelCase    :      str				)    ->  List[str]:
      for key in orig_state_dict.copy().keys():
            __snake_case :   Dict													=    orig_state_dict.pop(_UpperCAmelCase				)

            if "qkv" in key:
                  __snake_case :   Union[str, Any]													=    key.split('.'				)
                  __snake_case :   str													=    int(key_split[1]				)
                  __snake_case :   str													=    int(key_split[4]				)
                  __snake_case :   Any													=    config.embed_dim

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

      return orig_state_dict





def   a_   (		_UpperCAmelCase    :      Optional[int] ,_UpperCAmelCase    :      Dict ,_UpperCAmelCase    :      List[Any]				)    ->  List[Any]:
      __snake_case :   Dict													=    get_config(_UpperCAmelCase				)
      __snake_case :   str													=    SwinaSRForImageSuperResolution(_UpperCAmelCase				)
      model.eval()

      __snake_case :   int													=    torch.hub.load_state_dict_from_url(_UpperCAmelCase ,map_location='cpu'				)
      __snake_case :   Union[str, Any]													=    convert_state_dict(_UpperCAmelCase ,_UpperCAmelCase				)
      __snake_case							,     __snake_case :   int													=    model.load_state_dict(_UpperCAmelCase ,strict=_UpperCAmelCase				)

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

      __snake_case :   Dict													=    1_26 if 'Jpeg' in checkpoint_url else 2_56
      __snake_case :   Any													=    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]				),
          ]				)
      __snake_case :   Tuple													=    transforms(_UpperCAmelCase				).unsqueeze(0				)

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

      __snake_case :   Any													=    model(_UpperCAmelCase				)

      # assert values
      if "Swin2SR_ClassicalSR_X2_64" in checkpoint_url:
            __snake_case :   List[Any]													=    torch.Size([1, 3, 5_12, 5_12]				)
            __snake_case :   Optional[int]													=    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:
            __snake_case :   Dict													=    torch.Size([1, 3, 10_24, 10_24]				)
            __snake_case :   Any													=    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
            __snake_case :   int													=    torch.Size([1, 3, 10_24, 10_24]				)
            __snake_case :   str													=    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:
            __snake_case :   Tuple													=    torch.Size([1, 3, 5_12, 5_12]				)
            __snake_case :   Any													=    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:
            __snake_case :   Dict													=    torch.Size([1, 3, 10_24, 10_24]				)
            __snake_case :   str													=    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] ,_UpperCAmelCase ,atol=1E-3				)
      print('Looks ok!'				)

      __snake_case :   Tuple													=    {
          '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 :   int													=    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(_UpperCAmelCase				)
            print(f'''Saving image processor to {pytorch_dump_folder_path}'''				)
            processor.save_pretrained(_UpperCAmelCase				)

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


if __name__ == "__main__":
    A__    :   Optional[int]           =					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.''')

    A__    :   Any           =					parser.parse_args()
    convert_swinasr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)








'''simple docstring'''





import gc
import unittest

from parameterized import parameterized

from diffusers import FlaxUNetaDConditionModel
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import load_hf_numpy, require_flax, slow


if is_flax_available():
    import jax
    import jax.numpy as jnp



@slow
@require_flax
class  snake_case__						(    unittest.TestCase				):
 def A_   (   self						:							int      ,     __a						:							Optional[int]      ,     __a						:							Union[str, Any]  )      ->   str:







       '''simple docstring'''
       return f'''gaussian_noise_s={seed}_shape={"_".join([str(__a  ) for s in shape]  )}.npy'''
 def A_   (   self						:							Any  )      ->   List[str]:







       '''simple docstring'''
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
 def A_   (   self						:							List[Any]      ,     __a						:							Union[str, Any]=0      ,     __a						:							List[str]=(4, 4, 64, 64)      ,     __a						:							Optional[Any]=False  )      ->   List[str]:







       '''simple docstring'''
       __snake_case :   Optional[Any]													=    jnp.bfloataa if fpaa else jnp.floataa
       __snake_case :   List[Any]													=    jnp.array(load_hf_numpy(self.get_file_format(__a      ,     __a  )  )      ,     dtype=__a  )
       return image
 def A_   (   self						:							Any      ,     __a						:							Any=False      ,     __a						:							Dict="CompVis/stable-diffusion-v1-4"  )      ->   Optional[int]:







       '''simple docstring'''
       __snake_case :   Tuple													=    jnp.bfloataa if fpaa else jnp.floataa
       __snake_case :   int													=    'bf16' if fpaa else None

       __snake_case							,     __snake_case :   Union[str, Any]													=    FlaxUNetaDConditionModel.from_pretrained(
           __a      ,     subfolder='unet'      ,     dtype=__a      ,     revision=__a  )
       return model, params
 def A_   (   self						:							Any      ,     __a						:							Dict=0      ,     __a						:							Dict=(4, 77, 768)      ,     __a						:							List[str]=False  )      ->   List[Any]:







       '''simple docstring'''
       __snake_case :   List[Any]													=    jnp.bfloataa if fpaa else jnp.floataa
       __snake_case :   Any													=    jnp.array(load_hf_numpy(self.get_file_format(__a      ,     __a  )  )      ,     dtype=__a  )
       return hidden_states
 @parameterized.expand(
     [
         # fmt: off
         [83, 4, [-0.2_3_2_3, -0.1_3_0_4, 0.0_8_1_3, -0.3_0_9_3, -0.0_9_1_9, -0.1_5_7_1, -0.1_1_2_5, -0.5_8_0_6]],
         [17, 0.5_5, [-0.0_8_3_1, -0.2_4_4_3, 0.0_9_0_1, -0.0_9_1_9, 0.3_3_9_6, 0.0_1_0_3, -0.3_7_4_3, 0.0_7_0_1]],
         [8, 0.8_9, [-0.4_8_6_3, 0.0_8_5_9, 0.0_8_7_5, -0.1_6_5_8, 0.9_1_9_9, -0.0_1_1_4, 0.4_8_3_9, 0.4_6_3_9]],
         [3, 1000, [-0.5_6_4_9, 0.2_4_0_2, -0.5_5_1_8, 0.1_2_4_8, 1.1_3_2_8, -0.2_4_4_3, -0.0_3_2_5, -1.0_0_7_8]],
         # fmt: on
     ]  )
 def A_   (   self						:							List[Any]      ,     __a						:							List[str]      ,     __a						:							Union[str, Any]      ,     __a						:							Any  )      ->   Dict:







       '''simple docstring'''
       __snake_case							,     __snake_case :   Tuple													=    self.get_unet_model(model_id='CompVis/stable-diffusion-v1-4'      ,     fpaa=__a  )
       __snake_case :   Tuple													=    self.get_latents(__a      ,     fpaa=__a  )
       __snake_case :   int													=    self.get_encoder_hidden_states(__a      ,     fpaa=__a  )

       __snake_case :   List[str]													=    model.apply(
           {'params': params}      ,     __a      ,     jnp.array(__a      ,     dtype=jnp.intaa  )      ,     encoder_hidden_states=__a      ,     ).sample

       assert sample.shape == latents.shape

       __snake_case :   List[Any]													=    jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())  )      ,     dtype=jnp.floataa  )
       __snake_case :   str													=    jnp.array(__a      ,     dtype=jnp.floataa  )

       # Found torch (float16) and flax (bfloat16) outputs to be within this tolerance, in the same hardware
       assert jnp.allclose(__a      ,     __a      ,     atol=1e-2  )





 @parameterized.expand(
     [
         # fmt: off
         [83, 4, [0.1_5_1_4, 0.0_8_0_7, 0.1_6_2_4, 0.1_0_1_6, -0.1_8_9_6, 0.0_2_6_3, 0.0_6_7_7, 0.2_3_1_0]],
         [17, 0.5_5, [0.1_1_6_4, -0.0_2_1_6, 0.0_1_7_0, 0.1_5_8_9, -0.3_1_2_0, 0.1_0_0_5, -0.0_5_8_1, -0.1_4_5_8]],
         [8, 0.8_9, [-0.1_7_5_8, -0.0_1_6_9, 0.1_0_0_4, -0.1_4_1_1, 0.1_3_1_2, 0.1_1_0_3, -0.1_9_9_6, 0.2_1_3_9]],
         [3, 1000, [0.1_2_1_4, 0.0_3_5_2, -0.0_7_3_1, -0.1_5_6_2, -0.0_9_9_4, -0.0_9_0_6, -0.2_3_4_0, -0.0_5_3_9]],
         # fmt: on
     ]  )
 def A_   (   self						:							str      ,     __a						:							Optional[int]      ,     __a						:							Union[str, Any]      ,     __a						:							Tuple  )      ->   Optional[Any]:







       '''simple docstring'''
       __snake_case							,     __snake_case :   int													=    self.get_unet_model(model_id='stabilityai/stable-diffusion-2'      ,     fpaa=__a  )
       __snake_case :   int													=    self.get_latents(__a      ,     shape=(4, 4, 96, 96)      ,     fpaa=__a  )
       __snake_case :   Optional[Any]													=    self.get_encoder_hidden_states(__a      ,     shape=(4, 77, 1024)      ,     fpaa=__a  )

       __snake_case :   List[str]													=    model.apply(
           {'params': params}      ,     __a      ,     jnp.array(__a      ,     dtype=jnp.intaa  )      ,     encoder_hidden_states=__a      ,     ).sample

       assert sample.shape == latents.shape

       __snake_case :   List[Any]													=    jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())  )      ,     dtype=jnp.floataa  )
       __snake_case :   int													=    jnp.array(__a      ,     dtype=jnp.floataa  )

       # Found torch (float16) and flax (bfloat16) outputs to be within this tolerance, on the same hardware
       assert jnp.allclose(__a      ,     __a      ,     atol=1e-2  )



