




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							(					lowerCAmelCase )->	Optional[int]:
		SCREAMING_SNAKE_CASE_					      =							SwinaSRConfig()

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

		return config


def 					_a							(					lowerCAmelCase    ,						lowerCAmelCase )->	Optional[Any]:







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

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

		if "conv_first" in name:
				SCREAMING_SNAKE_CASE_					      =							name.replace('conv_first'    ,						'first_convolution' )

		if (
		    "upsample" in name
		    or "conv_before_upsample" in name
		    or "conv_bicubic" in name
		    or "conv_up" in name
		    or "conv_hr" in name
		    or "conv_last" in name
		    or "aux" in name
		):
				# heads
				if "conv_last" in name:
						SCREAMING_SNAKE_CASE_					      =							name.replace('conv_last'    ,						'final_convolution' )
				if config.upsampler in ["pixelshuffle", "pixelshuffle_aux", "nearest+conv"]:
						if "conv_before_upsample.0" in name:
								SCREAMING_SNAKE_CASE_					      =							name.replace('conv_before_upsample.0'    ,						'conv_before_upsample' )
						if "upsample.0" in name:
								SCREAMING_SNAKE_CASE_					      =							name.replace('upsample.0'    ,						'upsample.convolution_0' )
						if "upsample.2" in name:
								SCREAMING_SNAKE_CASE_					      =							name.replace('upsample.2'    ,						'upsample.convolution_1' )
						SCREAMING_SNAKE_CASE_					      =							'upsample.' + name
				elif config.upsampler == "pixelshuffledirect":
						SCREAMING_SNAKE_CASE_					      =							name.replace('upsample.0.weight'    ,						'upsample.conv.weight' )
						SCREAMING_SNAKE_CASE_					      =							name.replace('upsample.0.bias'    ,						'upsample.conv.bias' )
				else:
						pass
		else:
				SCREAMING_SNAKE_CASE_					      =							'swin2sr.' + name

		return name


def 					_a							(					lowerCAmelCase    ,						lowerCAmelCase )->	Tuple:







		for key in orig_state_dict.copy().keys():
				SCREAMING_SNAKE_CASE_					      =							orig_state_dict.pop(UpperCamelCase__ )

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

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

		return orig_state_dict


def 					_a							(					lowerCAmelCase    ,						lowerCAmelCase    ,						lowerCAmelCase )->	str:
		SCREAMING_SNAKE_CASE_					      =							get_config(UpperCamelCase__ )
		SCREAMING_SNAKE_CASE_					      =							SwinaSRForImageSuperResolution(UpperCamelCase__ )
		model.eval()

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

		SCREAMING_SNAKE_CASE_					      =							126 if 'Jpeg' in checkpoint_url else 256
		SCREAMING_SNAKE_CASE_					      =							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] ),
		    ] )
		SCREAMING_SNAKE_CASE_					      =							transforms(UpperCamelCase__ ).unsqueeze(0 )

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

		SCREAMING_SNAKE_CASE_					      =							model(UpperCamelCase__ )

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

		SCREAMING_SNAKE_CASE_					      =							{
		    'https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_ClassicalSR_X2_64.pth': (
		        'swin2SR-classical-sr-x2-64'
		    ),
		    'https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_ClassicalSR_X4_64.pth': (
		        'swin2SR-classical-sr-x4-64'
		    ),
		    'https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_CompressedSR_X4_48.pth': (
		        'swin2SR-compressed-sr-x4-48'
		    ),
		    'https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_Lightweight_X2_64.pth': (
		        'swin2SR-lightweight-x2-64'
		    ),
		    'https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR.pth': (
		        'swin2SR-realworld-sr-x4-64-bsrgan-psnr'
		    ),
		}
		SCREAMING_SNAKE_CASE_					      =							url_to_name[checkpoint_url]

		if pytorch_dump_folder_path is not None:
				print(F'''Saving model {model_name} to {pytorch_dump_folder_path}''' )
				model.save_pretrained(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__":
						SCREAMING_SNAKE_CASE:     List[str]									= 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.''')

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



import unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
				import jax

				from transformers.models.big_bird.modeling_flax_big_bird import (
				    FlaxBigBirdForCausalLM,
				    FlaxBigBirdForMaskedLM,
				    FlaxBigBirdForMultipleChoice,
				    FlaxBigBirdForPreTraining,
				    FlaxBigBirdForQuestionAnswering,
				    FlaxBigBirdForSequenceClassification,
				    FlaxBigBirdForTokenClassification,
				    FlaxBigBirdModel,
				)






class    _UpperCamelCase     (       unittest.TestCase					):

				"""simple docstring"""







				def __init__(    self					, a__					, a__=2					, a__=56					, a__=True					, a__=True					, a__=True					, a__=True					, a__=99					, a__=32					, a__=2					, a__=2					, a__=7					, a__="gelu_new"					, a__=0.1					, a__=0.1					, a__=512					, a__=16					, a__=2					, a__=0.02					, a__=4					, a__="block_sparse"					, a__=True					, a__=False					, a__=2					, a__=3					, )						->					List[Any]:
					A             =					parent
					A             =					batch_size
					A             =					seq_length
					A             =					is_training
					A             =					use_attention_mask
					A             =					use_token_type_ids
					A             =					use_labels
					A             =					vocab_size
					A             =					hidden_size
					A             =					num_hidden_layers
					A             =					num_attention_heads
					A             =					intermediate_size
					A             =					hidden_act
					A             =					hidden_dropout_prob
					A             =					attention_probs_dropout_prob
					A             =					max_position_embeddings
					A             =					type_vocab_size
					A             =					type_sequence_label_size
					A             =					initializer_range
					A             =					num_choices

					A             =					rescale_embeddings
					A             =					attention_type
					A             =					use_bias
					A             =					block_size
					A             =					num_random_blocks

				def   _UpperCAmelCase       (    self )						->					List[Any]:
					A             =					ids_tensor([self.batch_size, self.seq_length]					, self.vocab_size )

					A             =					None
					if self.use_attention_mask:
						A             =					random_attention_mask([self.batch_size, self.seq_length] )

					A             =					None
					if self.use_token_type_ids:
						A             =					ids_tensor([self.batch_size, self.seq_length]					, self.type_vocab_size )

					A             =					BigBirdConfig(
					    vocab_size=self.vocab_size					, hidden_size=self.hidden_size					, num_hidden_layers=self.num_hidden_layers					, num_attention_heads=self.num_attention_heads					, intermediate_size=self.intermediate_size					, hidden_act=self.hidden_act					, hidden_dropout_prob=self.hidden_dropout_prob					, attention_probs_dropout_prob=self.attention_probs_dropout_prob					, max_position_embeddings=self.max_position_embeddings					, type_vocab_size=self.type_vocab_size					, is_decoder=a__					, initializer_range=self.initializer_range					, attention_type=self.attention_type					, block_size=self.block_size					, num_random_blocks=self.num_random_blocks					, use_bias=self.use_bias					, rescale_embeddings=self.rescale_embeddings					, )

					return config, input_ids, token_type_ids, attention_mask

				def   _UpperCAmelCase       (    self )						->					Optional[int]:
					A             =					self.prepare_config_and_inputs()
					A       ,      A       ,      A       ,      A             =					config_and_inputs
					A             =					{
					    """input_ids""": input_ids,
					    """token_type_ids""": token_type_ids,
					    """attention_mask""": attention_mask,
					}
					return config, inputs_dict






@require_flax
class    _UpperCamelCase     (       __snake_case     ,		unittest.TestCase					):

				"""simple docstring"""







				lowerCAmelCase					=			(
				    (
				        FlaxBigBirdForCausalLM,
				        FlaxBigBirdModel,
				        FlaxBigBirdForPreTraining,
				        FlaxBigBirdForMaskedLM,
				        FlaxBigBirdForMultipleChoice,
				        FlaxBigBirdForQuestionAnswering,
				        FlaxBigBirdForSequenceClassification,
				        FlaxBigBirdForTokenClassification,
				    )
				    if is_flax_available()
				    else ()
				)

				lowerCAmelCase					=			False
				lowerCAmelCase					=			False

				def   _UpperCAmelCase       (    self )						->					Optional[int]:
					A             =					FlaxBigBirdModelTester(self )

				@slow
				# copied from `test_modeling_flax_common` because it takes much longer than other models
				def   _UpperCAmelCase       (    self )						->					Tuple:
					super().test_from_pretrained_save_pretrained()

				@slow
				# copied from `test_modeling_flax_common` because it takes much longer than other models
				def   _UpperCAmelCase       (    self )						->					Dict:
					super().test_from_pretrained_with_no_automatic_init()

				@slow
				# copied from `test_modeling_flax_common` because it takes much longer than other models
				def   _UpperCAmelCase       (    self )						->					str:
					super().test_no_automatic_init()

				@slow
				# copied from `test_modeling_flax_common` because it takes much longer than other models
				def   _UpperCAmelCase       (    self )						->					List[Any]:
					super().test_hidden_states_output()

				@slow
				def   _UpperCAmelCase       (    self )						->					str:
					for model_class_name in self.all_model_classes:
						A             =					model_class_name.from_pretrained("""google/bigbird-roberta-base""" )
						self.assertIsNotNone(a__ )

				def   _UpperCAmelCase       (    self )						->					Optional[int]:
					if self.test_attn_probs:
						super().test_attention_outputs()

				@slow
				# copied from `test_modeling_flax_common` because it takes much longer than other models
				def   _UpperCAmelCase       (    self )						->					List[Any]:
					A       ,      A             =					self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
						with self.subTest(model_class.__name__ ):
							A             =					self._prepare_for_class(a__					, a__ )
							A             =					model_class(a__ )

							@jax.jit
							def model_jitted(a__					, a__=None					, **a__ ):
								return model(input_ids=a__					, attention_mask=a__					, **a__ )

							with self.subTest("""JIT Enabled""" ):
								A             =					model_jitted(**a__ ).to_tuple()

							with self.subTest("""JIT Disabled""" ):
								with jax.disable_jit():
									A             =					model_jitted(**a__ ).to_tuple()

							self.assertEqual(len(a__ )					, len(a__ ) )
							for jitted_output, output in zip(a__					, a__ ):
								self.assertEqual(jitted_output.shape					, output.shape )

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__=1e-5					, a__="outputs"					, a__=None )						->					Tuple:
					# `bigbird_block_sparse_attention` in `FlaxBigBird` returns `attention_probs = None`, while in PyTorch version,
					# an effort was done to return `attention_probs` (yet to be verified).
					if name.startswith("""outputs.attentions""" ):
						return
					else:
						super().check_pt_flax_outputs(a__					, a__					, a__					, a__					, a__					, a__ )



