



"""simple docstring"""







import math

def       _lowerCamelCase				(							UpperCAmelCase_ :       int     )	-> bool:






	"""simple docstring"""





	if 1 < number < 4:
		# 2 and 3 are primes
		return True
	elif number < 2 or number % 2 == 0 or number % 3 == 0:
		# Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
		return False

	# All primes number are in format of 6k +/- 1
	for i in range(5,	int(math.sqrt(UpperCAmelCase_     ) + 1     ),	6     ):
		if number % i == 0 or number % (i + 2) == 0:
			return False
	return True

def       _lowerCamelCase				(							UpperCAmelCase_ :       int = 10001     )	-> int:






	"""simple docstring"""





	try:
		A__											=					int(UpperCAmelCase_     )
	except (TypeError, ValueError):
		raise TypeError("Parameter nth must be int or castable to int."     ) from None
	if nth <= 0:
		raise ValueError("Parameter nth must be greater than or equal to one."     )
	A__											=					[]
	A__											=					2
	while len(UpperCAmelCase_     ) < nth:
		if is_prime(UpperCAmelCase_     ):
			primes.append(UpperCAmelCase_     )
			num += 1
		else:
			num += 1
	return primes[len(UpperCAmelCase_     ) - 1]


if __name__ == "__main__":
			print(f'{solution() = }')








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					( UpperCamelCase__:  Any    )			-> Optional[int]:





	"""simple docstring"""







	A             =					SwinaSRConfig()

	if "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
		A             =					4
	elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
		A             =					4
		A             =					48
		A             =					"""pixelshuffle_aux"""
	elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
		A             =					[6, 6, 6, 6]
		A             =					60
		A             =					[6, 6, 6, 6]
		A             =					"""pixelshuffledirect"""
	elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
		A             =					4
		A             =					"""nearest+conv"""
	elif "Swin2SR_Jpeg_dynamic" in checkpoint_url:
		A             =					1
		A             =					1
		A             =					1_26
		A             =					7
		A             =					2_55.0
		A             =					""""""

	return config





def 	_lowerCAmelCase					( UpperCamelCase__:  Optional[Any]				,       UpperCamelCase__:  Dict    )			-> Optional[Any]:





	"""simple docstring"""







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

	if name == "norm.weight":
		A             =					"""layernorm.weight"""
	if name == "norm.bias":
		A             =					"""layernorm.bias"""

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

	return name





def 	_lowerCAmelCase					( UpperCamelCase__:  Optional[int]				,       UpperCamelCase__:  Union[str, Any]    )			-> Tuple:





	"""simple docstring"""







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

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

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

	return orig_state_dict





def 	_lowerCAmelCase					( UpperCamelCase__:  Optional[Any]				,       UpperCamelCase__:  Any				,       UpperCamelCase__:  List[Any]    )			-> str:





	"""simple docstring"""







	A             =					get_config(UpperCamelCase__    )
	A             =					SwinaSRForImageSuperResolution(UpperCamelCase__    )
	model.eval()

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

	A             =					1_26 if """Jpeg""" in checkpoint_url else 2_56
	A             =					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]    ),
	    ]    )
	A             =					transforms(UpperCamelCase__    ).unsqueeze(0    )

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

	A             =					model(UpperCamelCase__    )

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

	A             =					{
	    """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"""
	    ),
	}
	A             =					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__":
				_lowercase							:   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.")

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



