





'''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 				_lowerCAmelCase  (			lowerCamelCase_     :					Tuple			):
			__lowercase			  =							SwinaSRConfig()

			if "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
						__lowercase			  =							4
			elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
						__lowercase			  =							4
						__lowercase			  =							4_8
						__lowercase			  =							'''pixelshuffle_aux'''
			elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
						__lowercase			  =							[6, 6, 6, 6]
						__lowercase			  =							6_0
						__lowercase			  =							[6, 6, 6, 6]
						__lowercase			  =							'''pixelshuffledirect'''
			elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
						__lowercase			  =							4
						__lowercase			  =							'''nearest+conv'''
			elif "Swin2SR_Jpeg_dynamic" in checkpoint_url:
						__lowercase			  =							1
						__lowercase			  =							1
						__lowercase			  =							1_2_6
						__lowercase			  =							7
						__lowercase			  =							2_5_5.0
						__lowercase			  =							''''''

			return config




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

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

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

			return name




def 				_lowerCAmelCase  (			lowerCamelCase_     :					Tuple						,    lowerCamelCase_     :					List[str]			):
			for key in orig_state_dict.copy().keys():
						__lowercase			  =							orig_state_dict.pop(lowerCamelCase_			)

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

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

			return orig_state_dict




def 				_lowerCAmelCase  (			lowerCamelCase_     :					Any						,    lowerCamelCase_     :					int						,    lowerCamelCase_     :					int			):
			__lowercase			  =							get_config(lowerCamelCase_			)
			__lowercase			  =							SwinaSRForImageSuperResolution(lowerCamelCase_			)
			model.eval()

			__lowercase			  =							torch.hub.load_state_dict_from_url(lowerCamelCase_						,    map_location='''cpu'''			)
			__lowercase			  =							convert_state_dict(lowerCamelCase_						,    lowerCamelCase_			)
			__lowercase  ,						__lowercase			  =							model.load_state_dict(lowerCamelCase_						,    strict=lowerCamelCase_			)

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

			__lowercase			  =							1_2_6 if '''Jpeg''' in checkpoint_url else 2_5_6
			__lowercase			  =							Compose(
			    [
			        Resize((image_size, image_size)			),
			        ToTensor(),
			        Normalize(mean=[0.4_85, 0.4_56, 0.4_06]						,    std=[0.2_29, 0.2_24, 0.2_25]			),
			    ]			)
			__lowercase			  =							transforms(lowerCamelCase_			).unsqueeze(0			)

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

			__lowercase			  =							model(lowerCamelCase_			)

			# assert values
			if "Swin2SR_ClassicalSR_X2_64" in checkpoint_url:
						__lowercase			  =							torch.Size([1, 3, 5_1_2, 5_1_2]			)
						__lowercase			  =							torch.tensor(
						    [[-0.70_87, -0.71_38, -0.67_21], [-0.83_40, -0.80_95, -0.72_98], [-0.91_49, -0.84_14, -0.79_40]]			)
			elif "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
						__lowercase			  =							torch.Size([1, 3, 1_0_2_4, 1_0_2_4]			)
						__lowercase			  =							torch.tensor(
						    [[-0.77_75, -0.81_05, -0.89_33], [-0.77_64, -0.83_56, -0.92_25], [-0.79_76, -0.86_86, -0.95_79]]			)
			elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
						# TODO values didn't match exactly here
						__lowercase			  =							torch.Size([1, 3, 1_0_2_4, 1_0_2_4]			)
						__lowercase			  =							torch.tensor(
						    [[-0.80_35, -0.75_04, -0.74_91], [-0.85_38, -0.81_24, -0.77_82], [-0.88_04, -0.86_51, -0.84_93]]			)
			elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
						__lowercase			  =							torch.Size([1, 3, 5_1_2, 5_1_2]			)
						__lowercase			  =							torch.tensor(
						    [[-0.76_69, -0.86_62, -0.87_67], [-0.88_10, -0.99_62, -0.98_20], [-0.93_40, -1.03_22, -1.11_49]]			)
			elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
						__lowercase			  =							torch.Size([1, 3, 1_0_2_4, 1_0_2_4]			)
						__lowercase			  =							torch.tensor(
						    [[-0.52_38, -0.55_57, -0.63_21], [-0.60_16, -0.59_03, -0.63_91], [-0.62_44, -0.63_34, -0.68_89]]			)

			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]						,    lowerCamelCase_						,    atol=1E-3			)
			print('''Looks ok!'''			)

			__lowercase			  =							{
			    '''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			  =							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(lowerCamelCase_			)
						print(f"Saving image processor to {pytorch_dump_folder_path}"			)
						processor.save_pretrained(lowerCamelCase_			)

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












'''simple docstring'''



from ...configuration_utils import PretrainedConfig
from ...utils import logging


_SCREAMING_SNAKE_CASE   =      logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE   =      {
    '''s-JoL/Open-Llama-V1''': '''https://huggingface.co/s-JoL/Open-Llama-V1/blob/main/config.json''',
}






class 							__lowercase    ( lowerCAmelCase__  ):







	'''simple docstring'''

	a      :			Dict     =    "open-llama"


	def __init__(self       ,_lowerCamelCase=100000       ,_lowerCamelCase=4096       ,_lowerCamelCase=11008       ,_lowerCamelCase=32       ,_lowerCamelCase=32       ,_lowerCamelCase="silu"       ,_lowerCamelCase=2048       ,_lowerCamelCase=0.0_2       ,_lowerCamelCase=1E-6       ,_lowerCamelCase=True       ,_lowerCamelCase=0       ,_lowerCamelCase=1       ,_lowerCamelCase=2       ,_lowerCamelCase=False       ,_lowerCamelCase=True       ,_lowerCamelCase=0.1       ,_lowerCamelCase=0.1       ,_lowerCamelCase=True       ,_lowerCamelCase=True       ,_lowerCamelCase=None       ,**_lowerCamelCase       ,)							->							Union[str, Any]:







				'''simple docstring'''





				__lowercase			  =							vocab_size
				__lowercase			  =							max_position_embeddings
				__lowercase			  =							hidden_size
				__lowercase			  =							intermediate_size
				__lowercase			  =							num_hidden_layers
				__lowercase			  =							num_attention_heads
				__lowercase			  =							hidden_act
				__lowercase			  =							initializer_range
				__lowercase			  =							rms_norm_eps
				__lowercase			  =							use_cache
				__lowercase			  =							kwargs.pop(
				    '''use_memorry_efficient_attention'''       ,_lowerCamelCase	)
				__lowercase			  =							hidden_dropout_prob
				__lowercase			  =							attention_dropout_prob
				__lowercase			  =							use_stable_embedding
				__lowercase			  =							shared_input_output_embedding
				__lowercase			  =							rope_scaling
				self._rope_scaling_validation()

				super().__init__(
				    pad_token_id=_lowerCamelCase       ,bos_token_id=_lowerCamelCase       ,eos_token_id=_lowerCamelCase       ,tie_word_embeddings=_lowerCamelCase       ,**_lowerCamelCase       ,)


	def   _UpperCAmelCase       (self	)							->							List[Any]:







				'''simple docstring'''





				if self.rope_scaling is None:
							return

				if not isinstance(self.rope_scaling       ,_lowerCamelCase	) or len(self.rope_scaling	) != 2:
							raise ValueError(
							    '''`rope_scaling` must be a dictionary with with two fields, `name` and `factor`, '''
							    f"got {self.rope_scaling}"	)
				__lowercase			  =							self.rope_scaling.get('''type'''       ,_lowerCamelCase	)
				__lowercase			  =							self.rope_scaling.get('''factor'''       ,_lowerCamelCase	)
				if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
							raise ValueError(
							    f"`rope_scaling`'s name field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"	)
				if rope_scaling_factor is None or not isinstance(_lowerCamelCase       ,_lowerCamelCase	) or rope_scaling_factor <= 1.0:
							raise ValueError(f"`rope_scaling`'s factor field must be an float > 1, got {rope_scaling_factor}"	)





