





import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
     from transformers import WhisperFeatureExtractor

if is_torch_available():
     import torch

__lowercase	: Optional[Any]       =random.Random()




def      a__			(       lowercase__	,     lowercase__=1.0	,     lowercase__=None	,     lowercase__=None	):






   '''simple docstring'''






   if rng is None:
      UpperCAmelCase_           =global_rng

   UpperCAmelCase_           =[]
   for batch_idx in range(shape[0]	):
      values.append([]	)
      for _ in range(shape[1]	):
         values[-1].append(rng.random() * scale	)

   return values

@require_torch
@require_torchaudio
class   A							(						unittest.TestCase    ):



      def __init__(						self:    Optional[int]		,	_lowerCAmelCase:    Any		,	_lowerCAmelCase:    str=7		,	_lowerCAmelCase:    Optional[int]=400		,	_lowerCAmelCase:    Optional[Any]=2000		,	_lowerCAmelCase:    Union[str, Any]=10		,	_lowerCAmelCase:    Optional[int]=160		,	_lowerCAmelCase:    Any=8		,	_lowerCAmelCase:    str=0.0		,	_lowerCAmelCase:    Optional[int]=4000		,	_lowerCAmelCase:    Dict=False		,	_lowerCAmelCase:    Optional[Any]=True		,	)    ->     Dict:
         '''simple docstring'''

         UpperCAmelCase_           =parent
         UpperCAmelCase_           =batch_size
         UpperCAmelCase_           =min_seq_length
         UpperCAmelCase_           =max_seq_length
         UpperCAmelCase_           =(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
         UpperCAmelCase_           =padding_value
         UpperCAmelCase_           =sampling_rate
         UpperCAmelCase_           =return_attention_mask
         UpperCAmelCase_           =do_normalize
         UpperCAmelCase_           =feature_size
         UpperCAmelCase_           =chunk_length
         UpperCAmelCase_           =hop_length



      def 						lowerCAmelCase__     (						self:    Union[str, Any]							)    ->     int:
         '''simple docstring'''

         return {
             "feature_size": self.feature_size,
             "hop_length": self.hop_length,
             "chunk_length": self.chunk_length,
             "padding_value": self.padding_value,
             "sampling_rate": self.sampling_rate,
             "return_attention_mask": self.return_attention_mask,
             "do_normalize": self.do_normalize,
         }






      def 						lowerCAmelCase__     (						self:    List[Any]		,	_lowerCAmelCase:    List[str]=False		,	_lowerCAmelCase:    Dict=False							)    ->     str:
         '''simple docstring'''

         def _flatten(_lowerCAmelCase:    List[str]							):
            return list(itertools.chain(*_lowerCAmelCase							)							)

         if equal_length:
            UpperCAmelCase_           =[floats_list((self.max_seq_length, self.feature_size)							) for _ in range(self.batch_size							)]
         else:
            # make sure that inputs increase in size
            UpperCAmelCase_           =[
                floats_list((x, self.feature_size)							)
                for x in range(self.min_seq_length		,	self.max_seq_length		,	self.seq_length_diff							)
            ]
         if numpify:
            UpperCAmelCase_           =[np.asarray(_lowerCAmelCase							) for x in speech_inputs]
         return speech_inputs

@require_torch
@require_torchaudio
class   A							(						__lowercase		,   unittest.TestCase    ):
      _snake_case				  =WhisperFeatureExtractor if is_speech_available() else None



      def 						lowerCAmelCase__     (						self:    Union[str, Any]							)    ->     Optional[int]:
         '''simple docstring'''

         UpperCAmelCase_           =WhisperFeatureExtractionTester(self							)



      def 						lowerCAmelCase__     (						self:    Union[str, Any]							)    ->     Tuple:
         '''simple docstring'''

         UpperCAmelCase_           =self.feature_extraction_class(**self.feat_extract_dict							)

         with tempfile.TemporaryDirectory() as tmpdirname:
            UpperCAmelCase_           =feat_extract_first.save_pretrained(_lowerCAmelCase							)[0]
            check_json_file_has_correct_format(_lowerCAmelCase							)
            UpperCAmelCase_           =self.feature_extraction_class.from_pretrained(_lowerCAmelCase							)

         UpperCAmelCase_           =feat_extract_first.to_dict()
         UpperCAmelCase_           =feat_extract_second.to_dict()
         UpperCAmelCase_           =feat_extract_first.mel_filters
         UpperCAmelCase_           =feat_extract_second.mel_filters
         self.assertTrue(np.allclose(_lowerCAmelCase		,	_lowerCAmelCase							)							)
         self.assertEqual(_lowerCAmelCase		,	_lowerCAmelCase							)



      def 						lowerCAmelCase__     (						self:    Union[str, Any]							)    ->     Dict:
         '''simple docstring'''

         UpperCAmelCase_           =self.feature_extraction_class(**self.feat_extract_dict							)

         with tempfile.TemporaryDirectory() as tmpdirname:
            UpperCAmelCase_           =os.path.join(_lowerCAmelCase		,	"feat_extract.json"							)
            feat_extract_first.to_json_file(_lowerCAmelCase							)
            UpperCAmelCase_           =self.feature_extraction_class.from_json_file(_lowerCAmelCase							)

         UpperCAmelCase_           =feat_extract_first.to_dict()
         UpperCAmelCase_           =feat_extract_second.to_dict()
         UpperCAmelCase_           =feat_extract_first.mel_filters
         UpperCAmelCase_           =feat_extract_second.mel_filters
         self.assertTrue(np.allclose(_lowerCAmelCase		,	_lowerCAmelCase							)							)
         self.assertEqual(_lowerCAmelCase		,	_lowerCAmelCase							)



      def 						lowerCAmelCase__     (						self:    List[str]							)    ->     Optional[Any]:
         '''simple docstring'''

         UpperCAmelCase_           =self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
         # create three inputs of length 800, 1000, and 1200
         UpperCAmelCase_           =[floats_list((1, x)							)[0] for x in range(800		,	1400		,	200							)]
         UpperCAmelCase_           =[np.asarray(_lowerCAmelCase							) for speech_input in speech_inputs]

         # Test feature size
         UpperCAmelCase_           =feature_extractor(_lowerCAmelCase		,	padding="max_length"		,	return_tensors="np"							).input_features
         self.assertTrue(input_features.ndim == 3							)
         self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames							)
         self.assertTrue(input_features.shape[-2] == feature_extractor.feature_size							)

         # Test not batched input
         UpperCAmelCase_           =feature_extractor(speech_inputs[0]		,	return_tensors="np"							).input_features
         UpperCAmelCase_           =feature_extractor(np_speech_inputs[0]		,	return_tensors="np"							).input_features
         self.assertTrue(np.allclose(_lowerCAmelCase		,	_lowerCAmelCase		,	atol=1e-3							)							)

         # Test batched
         UpperCAmelCase_           =feature_extractor(_lowerCAmelCase		,	return_tensors="np"							).input_features
         UpperCAmelCase_           =feature_extractor(_lowerCAmelCase		,	return_tensors="np"							).input_features
         for enc_seq_a, enc_seq_a in zip(_lowerCAmelCase		,	_lowerCAmelCase							):
            self.assertTrue(np.allclose(_lowerCAmelCase		,	_lowerCAmelCase		,	atol=1e-3							)							)

         # Test 2-D numpy arrays are batched.
         UpperCAmelCase_           =[floats_list((1, x)							)[0] for x in (800, 800, 800)]
         UpperCAmelCase_           =np.asarray(_lowerCAmelCase							)
         UpperCAmelCase_           =feature_extractor(_lowerCAmelCase		,	return_tensors="np"							).input_features
         UpperCAmelCase_           =feature_extractor(_lowerCAmelCase		,	return_tensors="np"							).input_features
         for enc_seq_a, enc_seq_a in zip(_lowerCAmelCase		,	_lowerCAmelCase							):
            self.assertTrue(np.allclose(_lowerCAmelCase		,	_lowerCAmelCase		,	atol=1e-3							)							)

         # Test truncation required
         UpperCAmelCase_           =[floats_list((1, x)							)[0] for x in range(200		,	(feature_extractor.n_samples + 500)		,	200							)]
         UpperCAmelCase_           =[np.asarray(_lowerCAmelCase							) for speech_input in speech_inputs]

         UpperCAmelCase_           =[x[: feature_extractor.n_samples] for x in speech_inputs]
         UpperCAmelCase_           =[np.asarray(_lowerCAmelCase							) for speech_input in speech_inputs_truncated]

         UpperCAmelCase_           =feature_extractor(_lowerCAmelCase		,	return_tensors="np"							).input_features
         UpperCAmelCase_           =feature_extractor(_lowerCAmelCase		,	return_tensors="np"							).input_features
         for enc_seq_a, enc_seq_a in zip(_lowerCAmelCase		,	_lowerCAmelCase							):
            self.assertTrue(np.allclose(_lowerCAmelCase		,	_lowerCAmelCase		,	atol=1e-3							)							)



      def 						lowerCAmelCase__     (						self:    str							)    ->     Tuple:
         '''simple docstring'''

         import torch

         UpperCAmelCase_           =self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
         UpperCAmelCase_           =np.random.rand(100		,	32							).astype(np.floataa							)
         UpperCAmelCase_           =np_speech_inputs.tolist()

         for inputs in [py_speech_inputs, np_speech_inputs]:
            UpperCAmelCase_           =feature_extractor.pad([{"input_features": inputs}]		,	return_tensors="np"							)
            self.assertTrue(np_processed.input_features.dtype == np.floataa							)
            UpperCAmelCase_           =feature_extractor.pad([{"input_features": inputs}]		,	return_tensors="pt"							)
            self.assertTrue(pt_processed.input_features.dtype == torch.floataa							)



      def 						lowerCAmelCase__     (						self:    int		,	_lowerCAmelCase:    Optional[int]							)    ->     List[Any]:
         '''simple docstring'''

         UpperCAmelCase_           =load_dataset("hf-internal-testing/librispeech_asr_dummy"		,	"clean"		,	split="validation"							)
         # automatic decoding with librispeech
         UpperCAmelCase_           =ds.sort("id"							).select(range(_lowerCAmelCase							)							)[:num_samples]["audio"]

         return [x["array"] for x in speech_samples]



      def 						lowerCAmelCase__     (						self:    str							)    ->     Tuple:
         '''simple docstring'''

         UpperCAmelCase_           =torch.tensor(
             [
                 0.11_93, -0.09_46, -0.10_98, -0.01_96, 0.02_25, -0.06_90, -0.17_36, 0.09_51,
                 0.09_71, -0.08_17, -0.07_02, 0.01_62, 0.02_60, 0.00_17, -0.01_92, -0.16_78,
                 0.07_09, -0.18_67, -0.06_55, -0.02_74, -0.02_34, -0.18_84, -0.05_16, -0.05_54,
                 -0.02_74, -0.14_25, -0.14_23, 0.08_37, 0.03_77, -0.08_54
             ]							)
         # fmt: on

         UpperCAmelCase_           =self._load_datasamples(1							)
         UpperCAmelCase_           =WhisperFeatureExtractor()
         UpperCAmelCase_           =feature_extractor(_lowerCAmelCase		,	return_tensors="pt"							).input_features
         self.assertEqual(input_features.shape		,	(1, 80, 3000)							)
         self.assertTrue(torch.allclose(input_features[0, 0, :30]		,	_lowerCAmelCase		,	atol=1e-4							)							)






      def 						lowerCAmelCase__     (						self:    str							)    ->     Any:
         '''simple docstring'''

         UpperCAmelCase_           =self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
         UpperCAmelCase_           =self._load_datasamples(1							)[0]
         UpperCAmelCase_           =((audio - audio.min()) / (audio.max() - audio.min())) * 6_5535  # Rescale to [0, 65535] to show issue
         UpperCAmelCase_           =feat_extract.zero_mean_unit_var_norm([audio]		,	attention_mask=_lowerCAmelCase							)[0]

         self.assertTrue(np.all(np.mean(_lowerCAmelCase							) < 1e-3							)							)
         self.assertTrue(np.all(np.abs(np.var(_lowerCAmelCase							) - 1							) < 1e-3							)							)









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

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

   return config





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

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

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

   return name





def 		a__      (		_UpperCamelCase      :			List[Any]			,_UpperCamelCase      :			Union[str, Any]				):
   for key in orig_state_dict.copy().keys():
      __lowerCamelCase    				=     orig_state_dict.pop(_UpperCamelCase				)

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

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

   return orig_state_dict





def 		a__      (		_UpperCamelCase      :			str			,_UpperCamelCase      :			int			,_UpperCamelCase      :			Any				):
   __lowerCamelCase    				=     get_config(_UpperCamelCase				)
   __lowerCamelCase    				=     SwinaSRForImageSuperResolution(_UpperCamelCase				)
   model.eval()

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

   __lowerCamelCase    				=     1_26 if '''Jpeg''' in checkpoint_url else 2_56
   __lowerCamelCase    				=     Compose(
       [
           Resize((image_size, image_size)				),
           ToTensor(),
           Normalize(mean=[0.485, 0.456, 0.406]			,std=[0.229, 0.224, 0.225]				),
       ]				)
   __lowerCamelCase    				=     transforms(_UpperCamelCase				).unsqueeze(0				)

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

   __lowerCamelCase    				=     model(_UpperCamelCase				)

   # assert values
   if "Swin2SR_ClassicalSR_X2_64" in checkpoint_url:
      __lowerCamelCase    				=     torch.Size([1, 3, 5_12, 5_12]				)
      __lowerCamelCase    				=     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:
      __lowerCamelCase    				=     torch.Size([1, 3, 10_24, 10_24]				)
      __lowerCamelCase    				=     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
      __lowerCamelCase    				=     torch.Size([1, 3, 10_24, 10_24]				)
      __lowerCamelCase    				=     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:
      __lowerCamelCase    				=     torch.Size([1, 3, 5_12, 5_12]				)
      __lowerCamelCase    				=     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:
      __lowerCamelCase    				=     torch.Size([1, 3, 10_24, 10_24]				)
      __lowerCamelCase    				=     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]			,_UpperCamelCase			,atol=1e-3				)
   print('''Looks ok!'''				)

   __lowerCamelCase    				=     {
       '''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'''
       ),
   }
   __lowerCamelCase    				=     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_								=       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_								=       parser.parse_args()
  convert_swinasr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)






