





"""simple docstring"""
from __future__ import annotations

import os
import tempfile
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import is_tensorflow_text_available, is_tf_available
from transformers.testing_utils import require_tensorflow_text, require_tf, slow

from ..test_modeling_tf_common import floats_tensor
from .test_framework_agnostic import GenerationIntegrationTestsMixin


if is_tf_available():
       import tensorflow as tf

       from transformers import (
           AutoTokenizer,
           TFAutoModelForCausalLM,
           TFAutoModelForSeqaSeqLM,
           TFAutoModelForSpeechSeqaSeq,
           TFAutoModelForVisionaSeq,
           TFBartForConditionalGeneration,
           TFLogitsProcessorList,
           TFMinLengthLogitsProcessor,
           tf_top_k_top_p_filtering,
       )

if is_tensorflow_text_available():
       import tensorflow_text as text







@require_tf
class       a							(   unittest.TestCase  ):



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= tf.convert_to_tensor(
      [
          [
              8.2_220_991,  # 3rd highest value; idx. 0
              -0.5_620_044,
              5.23_229_752,
              4.0_386_393,
              -6.8_798_378,
              -0.54_785_802,
              -3.2_012_153,
              2.92_777_176,
              1.88_171_953,
              7.35_341_276,  # 5th highest value; idx. 9
              8.43_207_833,  # 2nd highest value; idx. 10
              -9.85_711_836,
              -5.96_209_236,
              -1.13_039_161,
              -7.1_115_294,
              -0.8_369_633,
              -5.3_186_408,
              7.06_427_407,
              0.81_369_344,
              -0.82_023_817,
              -5.9_179_796,
              0.58_813_443,
              -6.99_778_438,
              4.71_551_189,
              -0.18_771_637,
              7.44_020_759,  # 4th highest value; idx. 25
              9.38_450_987,  # 1st highest value; idx. 26
              2.12_662_941,
              -9.32_562_038,
              2.35_652_522,
          ],  # cummulative prob of 5 highest values <= 0.6
          [
              0.58_425_518,
              4.53_139_238,
              -5.57_510_464,
              -6.28_030_699,
              -7.19_529_503,
              -4.02_122_551,
              1.39_337_037,
              -6.06_707_057,
              1.59_480_517,
              -9.643_119,
              0.03_907_799,
              0.67_231_762,
              -8.88_206_726,
              6.27_115_922,  # 4th highest value; idx. 13
              2.28_520_723,
              4.82_767_506,
              4.30_421_368,
              8.8_275_313,  # 2nd highest value; idx. 17
              5.44_029_958,  # 5th highest value; idx. 18
              -4.4_735_794,
              7.38_579_536,  # 3rd highest value; idx. 20
              -2.91_051_663,
              2.61_946_077,
              -2.5_674_762,
              -9.48_959_302,
              -4.02_922_645,
              -1.35_416_918,
              9.67_702_323,  # 1st highest value; idx. 27
              -5.89_478_553,
              1.85_370_467,
          ],  # cummulative prob of 5 highest values <= 0.6
      ]				,       dtype=tf.floataa				,       )

  lowerCAmelCase    				= tf.convert_to_tensor(
      [[0, 0], [0, 9], [0, 10], [0, 25], [0, 26], [1, 13], [1, 17], [1, 18], [1, 20], [1, 27]]				,       dtype=tf.intaa				,       )  # expected non filtered idx as noted above

  lowerCAmelCase    				= tf.convert_to_tensor(
      [8.222_099, 7.3_534_126, 8.432_078, 7.4_402_075, 9.38_451, 6.271_159, 8.827_531, 5.4_402_995, 7.3_857_956, 9.677_023]				,       dtype=tf.floataa				,       )  # expected non filtered values as noted above

  lowerCAmelCase    				= tf_top_k_top_p_filtering(_snake_case				,       top_k=10				,       top_p=0.6				,       min_tokens_to_keep=4     )

  lowerCAmelCase    				= output[output != -float('inf'     )]
  lowerCAmelCase    				= tf.cast(
      tf.where(tf.not_equal(_snake_case				,       tf.constant(-float('inf'     )				,       dtype=tf.floataa     )     )     )				,       dtype=tf.intaa				,       )

  tf.debugging.assert_near(_snake_case				,       _snake_case				,       rtol=1E-12     )
  tf.debugging.assert_equal(_snake_case				,       _snake_case     )







@require_tf
class       a							(   unittest.TestCase   , a__  ):
 # setting framework_dependent_parameters needs to be gated, just like its contents' imports
 if is_tf_available():
  snake_case__             =       {
      '''AutoModelForCausalLM''': TFAutoModelForCausalLM,
      '''AutoModelForSpeechSeq2Seq''': TFAutoModelForSpeechSeqaSeq,
      '''AutoModelForSeq2SeqLM''': TFAutoModelForSeqaSeqLM,
      '''AutoModelForVision2Seq''': TFAutoModelForVisionaSeq,
      '''LogitsProcessorList''': TFLogitsProcessorList,
      '''MinLengthLogitsProcessor''': TFMinLengthLogitsProcessor,
      '''create_tensor_fn''': tf.convert_to_tensor,
      '''floats_tensor''': floats_tensor,
      '''return_tensors''': '''tf''',
  }



 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'     )
  lowerCAmelCase    				= 2
  lowerCAmelCase    				= 2







  class       a							(   tf.Module  ):



   def __init__(     self				,       _snake_case     ):







    """simple docstring"""






    super(_snake_case				,       self     ).__init__()
    lowerCAmelCase    				= model



   @tf.function(
       input_signature=(
           tf.TensorSpec((None, input_length)				,       tf.intaa				,       name='input_ids'     ),
           tf.TensorSpec((None, input_length)				,       tf.intaa				,       name='attention_mask'     ),
       )				,       jit_compile=_snake_case				,       )
   def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case     ):







    """simple docstring"""






    lowerCAmelCase    				= self.model.generate(
        input_ids=_snake_case				,       attention_mask=_snake_case				,       max_new_tokens=_snake_case				,       return_dict_in_generate=_snake_case				,       )
    return {"sequences": outputs["sequences"]}


  lowerCAmelCase    				= [[2, 0], [1_02, 1_03]]
  lowerCAmelCase    				= [[1, 0], [1, 1]]
  lowerCAmelCase    				= DummyModel(model=_snake_case     )
  with tempfile.TemporaryDirectory() as tmp_dir:
   tf.saved_model.save(_snake_case				,       _snake_case				,       signatures={'serving_default': dummy_model.serving}     )
   lowerCAmelCase    				= tf.saved_model.load(_snake_case     ).signatures['serving_default']
   for batch_size in range(1				,       len(_snake_case     ) + 1     ):
    lowerCAmelCase    				= {
        'input_ids': tf.constant(dummy_input_ids[:batch_size]     ),
        'attention_mask': tf.constant(dummy_attention_masks[:batch_size]     ),
    }
    lowerCAmelCase    				= serving_func(**_snake_case     )['sequences']
    lowerCAmelCase    				= test_model.generate(**_snake_case				,       max_new_tokens=_snake_case     )
    tf.debugging.assert_equal(_snake_case				,       _snake_case     )



 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'     )
  lowerCAmelCase    				= 1
  lowerCAmelCase    				= 2







  class       a							(   tf.Module  ):



   def __init__(     self				,       _snake_case     ):







    """simple docstring"""






    super(_snake_case				,       self     ).__init__()
    lowerCAmelCase    				= model



   @tf.function(
       input_signature=(
           tf.TensorSpec((batch_size, None)				,       tf.intaa				,       name='input_ids'     ),
           tf.TensorSpec((batch_size, None)				,       tf.intaa				,       name='attention_mask'     ),
       )				,       jit_compile=_snake_case				,       )
   def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case     ):







    """simple docstring"""






    lowerCAmelCase    				= self.model.generate(
        input_ids=_snake_case				,       attention_mask=_snake_case				,       max_new_tokens=_snake_case				,       return_dict_in_generate=_snake_case				,       )
    return {"sequences": outputs["sequences"]}


  lowerCAmelCase    				= [[2], [1_02, 1_03]]
  lowerCAmelCase    				= [[1], [1, 1]]
  lowerCAmelCase    				= DummyModel(model=_snake_case     )
  with tempfile.TemporaryDirectory() as tmp_dir:
   tf.saved_model.save(_snake_case				,       _snake_case				,       signatures={'serving_default': dummy_model.serving}     )
   lowerCAmelCase    				= tf.saved_model.load(_snake_case     ).signatures['serving_default']
   for input_row in range(len(_snake_case     )     ):
    lowerCAmelCase    				= {
        'input_ids': tf.constant([dummy_input_ids[input_row]]     ),
        'attention_mask': tf.constant([dummy_attention_masks[input_row]]     ),
    }
    lowerCAmelCase    				= serving_func(**_snake_case     )['sequences']
    lowerCAmelCase    				= test_model.generate(**_snake_case				,       max_new_tokens=_snake_case     )
    tf.debugging.assert_equal(_snake_case				,       _snake_case     )



 @slow
 @require_tensorflow_text
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  with tempfile.TemporaryDirectory() as tmp_dir:
   # file needed to load the TF tokenizer
   hf_hub_download(repo_id='google/flan-t5-small'				,       filename='spiece.model'				,       local_dir=_snake_case     )







   class       a							(   tf.keras.layers.Layer  ):



    def __init__(     self     ):







     """simple docstring"""






     super().__init__()
     lowerCAmelCase    				= text.SentencepieceTokenizer(
         model=tf.io.gfile.GFile(os.path.join(_snake_case				,       'spiece.model'     )				,       'rb'     ).read()     )
     lowerCAmelCase    				= TFAutoModelForSeqaSeqLM.from_pretrained('hf-internal-testing/tiny-random-t5'     )



    def    UpperCamelCase__  (     self				,       _snake_case				,       *_snake_case				,       **_snake_case     ):







     """simple docstring"""






     lowerCAmelCase    				= self.tokenizer.tokenize(_snake_case     )
     lowerCAmelCase				,lowerCAmelCase    				= text.pad_model_inputs(
         _snake_case				,       max_seq_length=64				,       pad_value=self.model.config.pad_token_id     )
     lowerCAmelCase    				= self.model.generate(input_ids=_snake_case				,       attention_mask=_snake_case     )
     return self.tokenizer.detokenize(_snake_case     )


   lowerCAmelCase    				= CompleteSentenceTransformer()
   lowerCAmelCase    				= tf.keras.layers.Input(shape=(1,)				,       dtype=tf.string				,       name='inputs'     )
   lowerCAmelCase    				= complete_model(_snake_case     )
   lowerCAmelCase    				= tf.keras.Model(_snake_case				,       _snake_case     )
   keras_model.save(_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= {
      'do_sample': True,
      'num_beams': 1,
      'top_p': 0.7,
      'top_k': 10,
      'temperature': 0.7,
  }
  lowerCAmelCase    				= 14

  lowerCAmelCase    				= AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'     )
  lowerCAmelCase    				= 'Hello, my dog is cute and'
  lowerCAmelCase    				= tokenizer(_snake_case				,       return_tensors='tf'     )
  lowerCAmelCase    				= TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'     )

  lowerCAmelCase    				= 6_38
  # forces the generation to happen on CPU, to avoid GPU-related quirks
  with tf.device(':/CPU:0'     ):
   tf.random.set_seed(0     )
   lowerCAmelCase    				= model.generate(**_snake_case				,       eos_token_id=_snake_case				,       **_snake_case     )
  self.assertTrue(expectation == len(generated_tokens[0]     )     )

  lowerCAmelCase    				= [6_38, 1_98]
  with tf.device(':/CPU:0'     ):
   tf.random.set_seed(0     )
   lowerCAmelCase    				= model.generate(**_snake_case				,       eos_token_id=_snake_case				,       **_snake_case     )
  self.assertTrue(expectation == len(generated_tokens[0]     )     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-bart'     )
  lowerCAmelCase    				= 'Hugging Face is a technology company based in New York and Paris.'
  lowerCAmelCase    				= bart_tokenizer(_snake_case				,       return_tensors='tf'     ).input_ids
  lowerCAmelCase    				= TFBartForConditionalGeneration.from_pretrained('hf-internal-testing/tiny-random-bart'     )
  lowerCAmelCase    				= bart_model.generate(_snake_case     ).numpy()







  class       a							(   a__  ):



   def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case=None				,       **_snake_case     ):







    """simple docstring"""






    return super().call(_snake_case				,       **_snake_case     )


  lowerCAmelCase    				= FakeBart.from_pretrained('hf-internal-testing/tiny-random-bart'     )
  lowerCAmelCase    				= bart_model.generate(_snake_case				,       foo='bar'     ).numpy()
  self.assertTrue(np.array_equal(_snake_case				,       _snake_case     )     )







  class       a							(   bart_model.model.encoder.__class__  ):



   def    UpperCamelCase__  (     self				,       _snake_case				,       **_snake_case     ):







    """simple docstring"""






    return super().call(_snake_case				,       **_snake_case     )


  lowerCAmelCase    				= FakeEncoder(bart_model.config				,       bart_model.model.shared     )
  lowerCAmelCase    				= fake_encoder

  # Normal generation still works (the output will be different because the encoder weights are different)
  lowerCAmelCase    				= bart_model.generate(_snake_case     ).numpy()
  with self.assertRaises(_snake_case     ):
   # FakeEncoder.call() accepts **kwargs -> no filtering -> value error due to unexpected input "foo"
   bart_model.generate(_snake_case				,       foo='bar'     )












"""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 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Any			):
 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    				= 126
  lowerCAmelCase    				= 7
  lowerCAmelCase    				= 255.0
  lowerCAmelCase    				= ''

 return config

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Tuple  ,     _UpperCAmelCase					:    Union[str, Any]			):
 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 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Any  ,     _UpperCAmelCase					:    Dict			):
 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 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Tuple  ,     _UpperCAmelCase					:    Optional[Any]  ,     _UpperCAmelCase					:    Tuple			):
 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    				= 126 if 'Jpeg' in checkpoint_url else 256
 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, 512, 512]			)
  lowerCAmelCase    				= 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:
  lowerCAmelCase    				= torch.Size([1, 3, 1024, 1024]			)
  lowerCAmelCase    				= 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
  lowerCAmelCase    				= torch.Size([1, 3, 1024, 1024]			)
  lowerCAmelCase    				= 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:
  lowerCAmelCase    				= torch.Size([1, 3, 512, 512]			)
  lowerCAmelCase    				= 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:
  lowerCAmelCase    				= torch.Size([1, 3, 1024, 1024]			)
  lowerCAmelCase    				= 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]  ,     _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__":
       __UpperCamelCase     :						Any						       =  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.''')

       __UpperCamelCase     :						Optional[int]						       =  parser.parse_args()
       convert_swinasr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)





