
"""simple docstring"""






import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


# =================#
# UNet Conversion #
# =================#

lowerCamelCase__ :       Optional[int]						     =			[
    # (stable-diffusion, HF Diffusers)
    ("time_embed.0.weight", "time_embedding.linear_1.weight"),
    ("time_embed.0.bias", "time_embedding.linear_1.bias"),
    ("time_embed.2.weight", "time_embedding.linear_2.weight"),
    ("time_embed.2.bias", "time_embedding.linear_2.bias"),
    ("input_blocks.0.0.weight", "conv_in.weight"),
    ("input_blocks.0.0.bias", "conv_in.bias"),
    ("out.0.weight", "conv_norm_out.weight"),
    ("out.0.bias", "conv_norm_out.bias"),
    ("out.2.weight", "conv_out.weight"),
    ("out.2.bias", "conv_out.bias"),
]

lowerCamelCase__ :       Optional[Any]						     =			[
    # (stable-diffusion, HF Diffusers)
    ("in_layers.0", "norm1"),
    ("in_layers.2", "conv1"),
    ("out_layers.0", "norm2"),
    ("out_layers.3", "conv2"),
    ("emb_layers.1", "time_emb_proj"),
    ("skip_connection", "conv_shortcut"),
]

lowerCamelCase__ :       Union[str, Any]						     =			[]
# hardcoded number of downblocks and resnets/attentions...
# would need smarter logic for other networks.
for i in range(4):
					# loop over downblocks/upblocks

					for j in range(2):
										# loop over resnets/attentions for downblocks
										lowerCamelCase__ :       Union[str, Any]						     =			f'''down_blocks.{i}.resnets.{j}.'''
										lowerCamelCase__ :       Optional[int]						     =			f'''input_blocks.{3*i + j + 1}.0.'''
										unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))

										if i < 3:
															# no attention layers in down_blocks.3
															lowerCamelCase__ :       int						     =			f'''down_blocks.{i}.attentions.{j}.'''
															lowerCamelCase__ :       Optional[Any]						     =			f'''input_blocks.{3*i + j + 1}.1.'''
															unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))

					for j in range(3):
										# loop over resnets/attentions for upblocks
										lowerCamelCase__ :       Tuple						     =			f'''up_blocks.{i}.resnets.{j}.'''
										lowerCamelCase__ :       str						     =			f'''output_blocks.{3*i + j}.0.'''
										unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))

										if i > 0:
															# no attention layers in up_blocks.0
															lowerCamelCase__ :       str						     =			f'''up_blocks.{i}.attentions.{j}.'''
															lowerCamelCase__ :       Any						     =			f'''output_blocks.{3*i + j}.1.'''
															unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))

					if i < 3:
										# no downsample in down_blocks.3
										lowerCamelCase__ :       List[str]						     =			f'''down_blocks.{i}.downsamplers.0.conv.'''
										lowerCamelCase__ :       List[Any]						     =			f'''input_blocks.{3*(i+1)}.0.op.'''
										unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))

										# no upsample in up_blocks.3
										lowerCamelCase__ :       Any						     =			f'''up_blocks.{i}.upsamplers.0.'''
										lowerCamelCase__ :       List[str]						     =			f'''output_blocks.{3*i + 2}.{1 if i == 0 else 2}.'''
										unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

lowerCamelCase__ :       Union[str, Any]						     =			"mid_block.attentions.0."
lowerCamelCase__ :       Union[str, Any]						     =			"middle_block.1."
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

for j in range(2):
					lowerCamelCase__ :       Optional[Any]						     =			f'''mid_block.resnets.{j}.'''
					lowerCamelCase__ :       Optional[int]						     =			f'''middle_block.{2*j}.'''
					unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))





def __A						(   a_ :							Dict     )->							str:




		'''simple docstring'''






		SCREAMING_SNAKE_CASE       :     Tuple				= {k: k for k in unet_state_dict.keys()}
		for sd_name, hf_name in unet_conversion_map:
				SCREAMING_SNAKE_CASE       :     int				= sd_name
		for k, v in mapping.items():
				if "resnets" in k:
						for sd_part, hf_part in unet_conversion_map_resnet:
								SCREAMING_SNAKE_CASE       :     Optional[Any]				= v.replace(__snake_case		,		__snake_case     )
						SCREAMING_SNAKE_CASE       :     Any				= v
		for k, v in mapping.items():
				for sd_part, hf_part in unet_conversion_map_layer:
						SCREAMING_SNAKE_CASE       :     Optional[Any]				= v.replace(__snake_case		,		__snake_case     )
				SCREAMING_SNAKE_CASE       :     Tuple				= v
		SCREAMING_SNAKE_CASE       :     Union[str, Any]				= {v: unet_state_dict[k] for k, v in mapping.items()}
		return new_state_dict


# ================#
# VAE Conversion #
# ================#

lowerCamelCase__ :       Optional[Any]						     =			[
    # (stable-diffusion, HF Diffusers)
    ("nin_shortcut", "conv_shortcut"),
    ("norm_out", "conv_norm_out"),
    ("mid.attn_1.", "mid_block.attentions.0."),
]

for i in range(4):
					# down_blocks have two resnets
					for j in range(2):
										lowerCamelCase__ :       Any						     =			f'''encoder.down_blocks.{i}.resnets.{j}.'''
										lowerCamelCase__ :       Optional[Any]						     =			f'''encoder.down.{i}.block.{j}.'''
										vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

					if i < 3:
										lowerCamelCase__ :       List[Any]						     =			f'''down_blocks.{i}.downsamplers.0.'''
										lowerCamelCase__ :       str						     =			f'''down.{i}.downsample.'''
										vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

										lowerCamelCase__ :       int						     =			f'''up_blocks.{i}.upsamplers.0.'''
										lowerCamelCase__ :       Optional[Any]						     =			f'''up.{3-i}.upsample.'''
										vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))

					# up_blocks have three resnets
					# also, up blocks in hf are numbered in reverse from sd
					for j in range(3):
										lowerCamelCase__ :       Any						     =			f'''decoder.up_blocks.{i}.resnets.{j}.'''
										lowerCamelCase__ :       Tuple						     =			f'''decoder.up.{3-i}.block.{j}.'''
										vae_conversion_map.append((sd_up_prefix, hf_up_prefix))

# this part accounts for mid blocks in both the encoder and the decoder
for i in range(2):
					lowerCamelCase__ :       str						     =			f'''mid_block.resnets.{i}.'''
					lowerCamelCase__ :       Union[str, Any]						     =			f'''mid.block_{i+1}.'''
					vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


lowerCamelCase__ :       List[Any]						     =			[
    # (stable-diffusion, HF Diffusers)
    ("norm.", "group_norm."),
    ("q.", "query."),
    ("k.", "key."),
    ("v.", "value."),
    ("proj_out.", "proj_attn."),
]





def __A						(   a_ :							Optional[Any]     )->							List[str]:




		'''simple docstring'''






		return w.reshape(*w.shape		,		1		,		1     )





def __A						(   a_ :							Union[str, Any]     )->							Optional[Any]:




		'''simple docstring'''






		SCREAMING_SNAKE_CASE       :     Union[str, Any]				= {k: k for k in vae_state_dict.keys()}
		for k, v in mapping.items():
				for sd_part, hf_part in vae_conversion_map:
						SCREAMING_SNAKE_CASE       :     str				= v.replace(__snake_case		,		__snake_case     )
				SCREAMING_SNAKE_CASE       :     Any				= v
		for k, v in mapping.items():
				if "attentions" in k:
						for sd_part, hf_part in vae_conversion_map_attn:
								SCREAMING_SNAKE_CASE       :     Dict				= v.replace(__snake_case		,		__snake_case     )
						SCREAMING_SNAKE_CASE       :     Tuple				= v
		SCREAMING_SNAKE_CASE       :     Dict				= {v: vae_state_dict[k] for k, v in mapping.items()}
		SCREAMING_SNAKE_CASE       :     Dict				= ['''q''', '''k''', '''v''', '''proj_out''']
		for k, v in new_state_dict.items():
				for weight_name in weights_to_convert:
						if F"mid.attn_1.{weight_name}.weight" in k:
								print(F"Reshaping {k} for SD format"     )
								SCREAMING_SNAKE_CASE       :     int				= reshape_weight_for_sd(__snake_case     )
		return new_state_dict


# =========================#
# Text Encoder Conversion #
# =========================#


lowerCamelCase__ :       Tuple						     =			[
    # (stable-diffusion, HF Diffusers)
    ("resblocks.", "text_model.encoder.layers."),
    ("ln_1", "layer_norm1"),
    ("ln_2", "layer_norm2"),
    (".c_fc.", ".fc1."),
    (".c_proj.", ".fc2."),
    (".attn", ".self_attn"),
    ("ln_final.", "transformer.text_model.final_layer_norm."),
    ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"),
    ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"),
]
lowerCamelCase__ :       int						     =			{re.escape(x[1]): x[0] for x in textenc_conversion_lst}
lowerCamelCase__ :       List[str]						     =			re.compile("|".join(protected.keys()))

# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
lowerCamelCase__ :       Dict						     =			{"q": 0, "k": 1, "v": 2}





def __A						(   a_ :							Optional[int]     )->							Any:




		'''simple docstring'''






		SCREAMING_SNAKE_CASE       :     str				= {}
		SCREAMING_SNAKE_CASE       :     Union[str, Any]				= {}
		SCREAMING_SNAKE_CASE       :     Optional[Any]				= {}
		for k, v in text_enc_dict.items():
				if (
				    k.endswith('''.self_attn.q_proj.weight'''     )
				    or k.endswith('''.self_attn.k_proj.weight'''     )
				    or k.endswith('''.self_attn.v_proj.weight'''     )
				):
						SCREAMING_SNAKE_CASE       :     Dict				= k[: -len('''.q_proj.weight'''     )]
						SCREAMING_SNAKE_CASE       :     List[Any]				= k[-len('''q_proj.weight'''     )]
						if k_pre not in capture_qkv_weight:
								SCREAMING_SNAKE_CASE       :     List[str]				= [None, None, None]
						SCREAMING_SNAKE_CASE       :     int				= v
						continue

				if (
				    k.endswith('''.self_attn.q_proj.bias'''     )
				    or k.endswith('''.self_attn.k_proj.bias'''     )
				    or k.endswith('''.self_attn.v_proj.bias'''     )
				):
						SCREAMING_SNAKE_CASE       :     Optional[Any]				= k[: -len('''.q_proj.bias'''     )]
						SCREAMING_SNAKE_CASE       :     Dict				= k[-len('''q_proj.bias'''     )]
						if k_pre not in capture_qkv_bias:
								SCREAMING_SNAKE_CASE       :     List[str]				= [None, None, None]
						SCREAMING_SNAKE_CASE       :     Optional[int]				= v
						continue

				SCREAMING_SNAKE_CASE       :     Any				= textenc_pattern.sub(lambda a_     : protected[re.escape(m.group(0     )     )]		,		__snake_case     )
				SCREAMING_SNAKE_CASE       :     List[str]				= v

		for k_pre, tensors in capture_qkv_weight.items():
				if None in tensors:
						raise Exception('''CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing'''     )
				SCREAMING_SNAKE_CASE       :     Dict				= textenc_pattern.sub(lambda a_     : protected[re.escape(m.group(0     )     )]		,		__snake_case     )
				SCREAMING_SNAKE_CASE       :     Optional[Any]				= torch.cat(__snake_case     )

		for k_pre, tensors in capture_qkv_bias.items():
				if None in tensors:
						raise Exception('''CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing'''     )
				SCREAMING_SNAKE_CASE       :     Optional[Any]				= textenc_pattern.sub(lambda a_     : protected[re.escape(m.group(0     )     )]		,		__snake_case     )
				SCREAMING_SNAKE_CASE       :     Dict				= torch.cat(__snake_case     )

		return new_state_dict





def __A						(   a_ :							Optional[Any]     )->							Tuple:




		'''simple docstring'''






		return text_enc_dict


if __name__ == "__main__":
					lowerCamelCase__ :       Dict						     =			argparse.ArgumentParser()

					parser.add_argument("--model_path", default=None, type=str, required=True, help="Path to the model to convert.")
					parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.")
					parser.add_argument("--half", action="store_true", help="Save weights in half precision.")
					parser.add_argument(
					    "--use_safetensors", action="store_true", help="Save weights use safetensors, default is ckpt."
					)

					lowerCamelCase__ :       Tuple						     =			parser.parse_args()

					assert args.model_path is not None, "Must provide a model path!"

					assert args.checkpoint_path is not None, "Must provide a checkpoint path!"

					# Path for safetensors
					lowerCamelCase__ :       Dict						     =			osp.join(args.model_path, "unet", "diffusion_pytorch_model.safetensors")
					lowerCamelCase__ :       List[Any]						     =			osp.join(args.model_path, "vae", "diffusion_pytorch_model.safetensors")
					lowerCamelCase__ :       Optional[Any]						     =			osp.join(args.model_path, "text_encoder", "model.safetensors")

					# Load models from safetensors if it exists, if it doesn't pytorch
					if osp.exists(unet_path):
										lowerCamelCase__ :       Optional[Any]						     =			load_file(unet_path, device="cpu")
					else:
										lowerCamelCase__ :       Any						     =			osp.join(args.model_path, "unet", "diffusion_pytorch_model.bin")
										lowerCamelCase__ :       Optional[int]						     =			torch.load(unet_path, map_location="cpu")

					if osp.exists(vae_path):
										lowerCamelCase__ :       Tuple						     =			load_file(vae_path, device="cpu")
					else:
										lowerCamelCase__ :       Dict						     =			osp.join(args.model_path, "vae", "diffusion_pytorch_model.bin")
										lowerCamelCase__ :       str						     =			torch.load(vae_path, map_location="cpu")

					if osp.exists(text_enc_path):
										lowerCamelCase__ :       int						     =			load_file(text_enc_path, device="cpu")
					else:
										lowerCamelCase__ :       Dict						     =			osp.join(args.model_path, "text_encoder", "pytorch_model.bin")
										lowerCamelCase__ :       int						     =			torch.load(text_enc_path, map_location="cpu")

					# Convert the UNet model
					lowerCamelCase__ :       Optional[Any]						     =			convert_unet_state_dict(unet_state_dict)
					lowerCamelCase__ :       List[Any]						     =			{"model.diffusion_model." + k: v for k, v in unet_state_dict.items()}

					# Convert the VAE model
					lowerCamelCase__ :       str						     =			convert_vae_state_dict(vae_state_dict)
					lowerCamelCase__ :       Optional[Any]						     =			{"first_stage_model." + k: v for k, v in vae_state_dict.items()}

					# Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
					lowerCamelCase__ :       Tuple						     =			"text_model.encoder.layers.22.layer_norm2.bias" in text_enc_dict

					if is_vaa_model:
										# Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
										lowerCamelCase__ :       Optional[int]						     =			{"transformer." + k: v for k, v in text_enc_dict.items()}
										lowerCamelCase__ :       int						     =			convert_text_enc_state_dict_vaa(text_enc_dict)
										lowerCamelCase__ :       List[Any]						     =			{"cond_stage_model.model." + k: v for k, v in text_enc_dict.items()}
					else:
										lowerCamelCase__ :       Optional[Any]						     =			convert_text_enc_state_dict(text_enc_dict)
										lowerCamelCase__ :       Optional[int]						     =			{"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()}

					# Put together new checkpoint
					lowerCamelCase__ :       List[str]						     =			{**unet_state_dict, **vae_state_dict, **text_enc_dict}
					if args.half:
										lowerCamelCase__ :       Tuple						     =			{k: v.half() for k, v in state_dict.items()}

					if args.use_safetensors:
										save_file(state_dict, args.checkpoint_path)
					else:
										lowerCamelCase__ :       Optional[int]						     =			{"state_dict": state_dict}
										torch.save(state_dict, args.checkpoint_path)







"""simple docstring"""






import collections
import inspect
import unittest

from transformers import SwinvaConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch
					from torch import nn

					from transformers import SwinvaForImageClassification, SwinvaForMaskedImageModeling, SwinvaModel
					from transformers.models.swinva.modeling_swinva import SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
					from PIL import Image

					from transformers import AutoImageProcessor




class 				lowercase__:
	'''simple docstring'''

	def __init__(						self    :str		, lowerCamelCase_    :List[Any]		, lowerCamelCase_    :Any=13		, lowerCamelCase_    :Any=32		, lowerCamelCase_    :Union[str, Any]=2		, lowerCamelCase_    :Any=3		, lowerCamelCase_    :Union[str, Any]=16		, lowerCamelCase_    :int=[1, 2, 1]		, lowerCamelCase_    :str=[2, 2, 4]		, lowerCamelCase_    :str=2		, lowerCamelCase_    :Tuple=2.0		, lowerCamelCase_    :Union[str, Any]=True		, lowerCamelCase_    :str=0.0		, lowerCamelCase_    :Optional[int]=0.0		, lowerCamelCase_    :Dict=0.1		, lowerCamelCase_    :Union[str, Any]="gelu"		, lowerCamelCase_    :str=False		, lowerCamelCase_    :List[Any]=True		, lowerCamelCase_    :int=0.0_2		, lowerCamelCase_    :List[Any]=1E-5		, lowerCamelCase_    :int=True		, lowerCamelCase_    :str=None		, lowerCamelCase_    :List[Any]=True		, lowerCamelCase_    :Union[str, Any]=10		, lowerCamelCase_    :List[Any]=8		, )							->			int:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE       :     Optional[Any]				= parent
			SCREAMING_SNAKE_CASE       :     int				= batch_size
			SCREAMING_SNAKE_CASE       :     int				= image_size
			SCREAMING_SNAKE_CASE       :     Union[str, Any]				= patch_size
			SCREAMING_SNAKE_CASE       :     Tuple				= num_channels
			SCREAMING_SNAKE_CASE       :     Any				= embed_dim
			SCREAMING_SNAKE_CASE       :     int				= depths
			SCREAMING_SNAKE_CASE       :     List[str]				= num_heads
			SCREAMING_SNAKE_CASE       :     Union[str, Any]				= window_size
			SCREAMING_SNAKE_CASE       :     Optional[Any]				= mlp_ratio
			SCREAMING_SNAKE_CASE       :     List[Any]				= qkv_bias
			SCREAMING_SNAKE_CASE       :     Union[str, Any]				= hidden_dropout_prob
			SCREAMING_SNAKE_CASE       :     Dict				= attention_probs_dropout_prob
			SCREAMING_SNAKE_CASE       :     List[str]				= drop_path_rate
			SCREAMING_SNAKE_CASE       :     List[Any]				= hidden_act
			SCREAMING_SNAKE_CASE       :     Optional[Any]				= use_absolute_embeddings
			SCREAMING_SNAKE_CASE       :     Any				= patch_norm
			SCREAMING_SNAKE_CASE       :     Union[str, Any]				= layer_norm_eps
			SCREAMING_SNAKE_CASE       :     List[Any]				= initializer_range
			SCREAMING_SNAKE_CASE       :     Any				= is_training
			SCREAMING_SNAKE_CASE       :     List[Any]				= scope
			SCREAMING_SNAKE_CASE       :     Optional[Any]				= use_labels
			SCREAMING_SNAKE_CASE       :     Optional[Any]				= type_sequence_label_size
			SCREAMING_SNAKE_CASE       :     Optional[Any]				= encoder_stride





	def   __lowerCAmelCase (						self    :Optional[int]					)							->			List[str]:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE       :     Optional[int]				= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

			SCREAMING_SNAKE_CASE       :     int				= None
			if self.use_labels:
					SCREAMING_SNAKE_CASE       :     Dict				= ids_tensor([self.batch_size]		, self.type_sequence_label_size					)

			SCREAMING_SNAKE_CASE       :     Tuple				= self.get_config()

			return config, pixel_values, labels





	def   __lowerCAmelCase (						self    :int					)							->			int:







			'''simple docstring'''




			return SwinvaConfig(
			    image_size=self.image_size		, patch_size=self.patch_size		, num_channels=self.num_channels		, embed_dim=self.embed_dim		, depths=self.depths		, num_heads=self.num_heads		, window_size=self.window_size		, mlp_ratio=self.mlp_ratio		, qkv_bias=self.qkv_bias		, hidden_dropout_prob=self.hidden_dropout_prob		, attention_probs_dropout_prob=self.attention_probs_dropout_prob		, drop_path_rate=self.drop_path_rate		, hidden_act=self.hidden_act		, use_absolute_embeddings=self.use_absolute_embeddings		, path_norm=self.patch_norm		, layer_norm_eps=self.layer_norm_eps		, initializer_range=self.initializer_range		, encoder_stride=self.encoder_stride		, )





	def   __lowerCAmelCase (						self    :List[Any]		, lowerCamelCase_    :int		, lowerCamelCase_    :Optional[int]		, lowerCamelCase_    :Optional[Any]					)							->			Union[str, Any]:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE       :     str				= SwinvaModel(config=lowerCamelCase_					)
			model.to(lowerCamelCase_					)
			model.eval()
			SCREAMING_SNAKE_CASE       :     int				= model(lowerCamelCase_					)

			SCREAMING_SNAKE_CASE       :     List[Any]				= ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths					) - 1))
			SCREAMING_SNAKE_CASE       :     Dict				= int(config.embed_dim * 2 ** (len(config.depths					) - 1)					)

			self.parent.assertEqual(result.last_hidden_state.shape		, (self.batch_size, expected_seq_len, expected_dim)					)





	def   __lowerCAmelCase (						self    :str		, lowerCamelCase_    :str		, lowerCamelCase_    :Tuple		, lowerCamelCase_    :Any					)							->			Union[str, Any]:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE       :     List[str]				= SwinvaForMaskedImageModeling(config=lowerCamelCase_					)
			model.to(lowerCamelCase_					)
			model.eval()
			SCREAMING_SNAKE_CASE       :     Dict				= model(lowerCamelCase_					)
			self.parent.assertEqual(
			    result.logits.shape		, (self.batch_size, self.num_channels, self.image_size, self.image_size)					)

			# test greyscale images
			SCREAMING_SNAKE_CASE       :     Tuple				= 1
			SCREAMING_SNAKE_CASE       :     List[Any]				= SwinvaForMaskedImageModeling(lowerCamelCase_					)
			model.to(lowerCamelCase_					)
			model.eval()

			SCREAMING_SNAKE_CASE       :     List[str]				= floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
			SCREAMING_SNAKE_CASE       :     Tuple				= model(lowerCamelCase_					)
			self.parent.assertEqual(result.logits.shape		, (self.batch_size, 1, self.image_size, self.image_size)					)





	def   __lowerCAmelCase (						self    :Tuple		, lowerCamelCase_    :Union[str, Any]		, lowerCamelCase_    :Any		, lowerCamelCase_    :int					)							->			Tuple:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE       :     Tuple				= self.type_sequence_label_size
			SCREAMING_SNAKE_CASE       :     Optional[Any]				= SwinvaForImageClassification(lowerCamelCase_					)
			model.to(lowerCamelCase_					)
			model.eval()
			SCREAMING_SNAKE_CASE       :     Optional[int]				= model(lowerCamelCase_		, labels=lowerCamelCase_					)
			self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.type_sequence_label_size)					)







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







			'''simple docstring'''




			SCREAMING_SNAKE_CASE       :     List[str]				= self.prepare_config_and_inputs()
			SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE       :     Optional[int]				= config_and_inputs
			SCREAMING_SNAKE_CASE       :     int				= {'''pixel_values''': pixel_values}
			return config, inputs_dict





@require_torch
class 				lowercase__(				_UpperCAmelCase		,						_UpperCAmelCase		,						unittest.TestCase   ):
	'''simple docstring'''

	UpperCamelCase				=  (
	    (SwinvaModel, SwinvaForImageClassification, SwinvaForMaskedImageModeling) if is_torch_available() else ()
	)
	UpperCamelCase				=  (
	    {"""feature-extraction""": SwinvaModel, """image-classification""": SwinvaForImageClassification}
	    if is_torch_available()
	    else {}
	)

	UpperCamelCase				=  False
	UpperCamelCase				=  False
	UpperCamelCase				=  False
	UpperCamelCase				=  False





	def   __lowerCAmelCase (						self    :Union[str, Any]					)							->			Union[str, Any]:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE       :     int				= SwinvaModelTester(self					)
			SCREAMING_SNAKE_CASE       :     Dict				= ConfigTester(self		, config_class=lowerCamelCase_		, embed_dim=37					)





	def   __lowerCAmelCase (						self    :Dict					)							->			List[str]:







			'''simple docstring'''




			self.config_tester.create_and_test_config_to_json_string()
			self.config_tester.create_and_test_config_to_json_file()
			self.config_tester.create_and_test_config_from_and_save_pretrained()
			self.config_tester.create_and_test_config_with_num_labels()
			self.config_tester.check_config_can_be_init_without_params()
			self.config_tester.check_config_arguments_init()





	def   __lowerCAmelCase (						self    :List[Any]					)							->			List[str]:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE       :     Union[str, Any]				= self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*lowerCamelCase_					)





	@unittest.skip(reason='''Got `CUDA error: misaligned address` with PyTorch 2.0.0.'''					)
	def   __lowerCAmelCase (						self    :str					)							->			List[Any]:







			'''simple docstring'''




			pass





	@unittest.skip(reason='''Swinv2 does not use inputs_embeds'''					)
	def   __lowerCAmelCase (						self    :List[Any]					)							->			Any:







			'''simple docstring'''




			pass





	def   __lowerCAmelCase (						self    :int					)							->			Union[str, Any]:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE       :     Optional[int]				= self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					SCREAMING_SNAKE_CASE       :     List[Any]				= model_class(lowerCamelCase_					)
					self.assertIsInstance(model.get_input_embeddings()		, (nn.Module)					)
					SCREAMING_SNAKE_CASE       :     str				= model.get_output_embeddings()
					self.assertTrue(x is None or isinstance(lowerCamelCase_		, nn.Linear					)					)





	def   __lowerCAmelCase (						self    :int					)							->			List[Any]:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE       :     Tuple				= self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					SCREAMING_SNAKE_CASE       :     Optional[Any]				= model_class(lowerCamelCase_					)
					SCREAMING_SNAKE_CASE       :     List[str]				= inspect.signature(model.forward					)
					# signature.parameters is an OrderedDict => so arg_names order is deterministic
					SCREAMING_SNAKE_CASE       :     Optional[int]				= [*signature.parameters.keys()]

					SCREAMING_SNAKE_CASE       :     Union[str, Any]				= ['''pixel_values''']
					self.assertListEqual(arg_names[:1]		, lowerCamelCase_					)





	def   __lowerCAmelCase (						self    :int					)							->			Optional[Any]:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE       :     str				= self.model_tester.prepare_config_and_inputs_for_common()
			SCREAMING_SNAKE_CASE       :     Tuple				= True

			for model_class in self.all_model_classes:
					SCREAMING_SNAKE_CASE       :     List[str]				= True
					SCREAMING_SNAKE_CASE       :     List[Any]				= False
					SCREAMING_SNAKE_CASE       :     Any				= True
					SCREAMING_SNAKE_CASE       :     Any				= model_class(lowerCamelCase_					)
					model.to(lowerCamelCase_					)
					model.eval()
					with torch.no_grad():
							SCREAMING_SNAKE_CASE       :     str				= model(**self._prepare_for_class(lowerCamelCase_		, lowerCamelCase_					)					)
					SCREAMING_SNAKE_CASE       :     Tuple				= outputs.attentions
					SCREAMING_SNAKE_CASE       :     Tuple				= len(self.model_tester.depths					)
					self.assertEqual(len(lowerCamelCase_					)		, lowerCamelCase_					)

					# check that output_attentions also work using config
					del inputs_dict["output_attentions"]
					SCREAMING_SNAKE_CASE       :     Union[str, Any]				= True
					SCREAMING_SNAKE_CASE       :     Optional[int]				= config.window_size**2
					SCREAMING_SNAKE_CASE       :     str				= model_class(lowerCamelCase_					)
					model.to(lowerCamelCase_					)
					model.eval()
					with torch.no_grad():
							SCREAMING_SNAKE_CASE       :     Union[str, Any]				= model(**self._prepare_for_class(lowerCamelCase_		, lowerCamelCase_					)					)
					SCREAMING_SNAKE_CASE       :     Dict				= outputs.attentions
					self.assertEqual(len(lowerCamelCase_					)		, lowerCamelCase_					)

					self.assertListEqual(
					    list(attentions[0].shape[-3:]					)		, [self.model_tester.num_heads[0], window_size_squared, window_size_squared]		, )
					SCREAMING_SNAKE_CASE       :     Dict				= len(lowerCamelCase_					)

					# Check attention is always last and order is fine
					SCREAMING_SNAKE_CASE       :     int				= True
					SCREAMING_SNAKE_CASE       :     Dict				= True
					SCREAMING_SNAKE_CASE       :     int				= model_class(lowerCamelCase_					)
					model.to(lowerCamelCase_					)
					model.eval()
					with torch.no_grad():
							SCREAMING_SNAKE_CASE       :     List[Any]				= model(**self._prepare_for_class(lowerCamelCase_		, lowerCamelCase_					)					)

					if hasattr(self.model_tester		, '''num_hidden_states_types'''					):
							SCREAMING_SNAKE_CASE       :     Any				= self.model_tester.num_hidden_states_types
					else:
							# also another +1 for reshaped_hidden_states
							SCREAMING_SNAKE_CASE       :     Optional[Any]				= 2
					self.assertEqual(out_len + added_hidden_states		, len(lowerCamelCase_					)					)

					SCREAMING_SNAKE_CASE       :     Tuple				= outputs.attentions

					self.assertEqual(len(lowerCamelCase_					)		, lowerCamelCase_					)

					self.assertListEqual(
					    list(self_attentions[0].shape[-3:]					)		, [self.model_tester.num_heads[0], window_size_squared, window_size_squared]		, )





	def   __lowerCAmelCase (						self    :Dict		, lowerCamelCase_    :int		, lowerCamelCase_    :Any		, lowerCamelCase_    :Tuple		, lowerCamelCase_    :Any					)							->			Optional[Any]:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE       :     Union[str, Any]				= model_class(lowerCamelCase_					)
			model.to(lowerCamelCase_					)
			model.eval()

			with torch.no_grad():
					SCREAMING_SNAKE_CASE       :     Optional[int]				= model(**self._prepare_for_class(lowerCamelCase_		, lowerCamelCase_					)					)

			SCREAMING_SNAKE_CASE       :     Dict				= outputs.hidden_states

			SCREAMING_SNAKE_CASE       :     str				= getattr(
			    self.model_tester		, '''expected_num_hidden_layers'''		, len(self.model_tester.depths					) + 1					)
			self.assertEqual(len(lowerCamelCase_					)		, lowerCamelCase_					)

			# Swinv2 has a different seq_length
			SCREAMING_SNAKE_CASE       :     List[str]				= (
			    config.patch_size
			    if isinstance(config.patch_size		, collections.abc.Iterable					)
			    else (config.patch_size, config.patch_size)
			)

			SCREAMING_SNAKE_CASE       :     Dict				= (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

			self.assertListEqual(
			    list(hidden_states[0].shape[-2:]					)		, [num_patches, self.model_tester.embed_dim]		, )

			SCREAMING_SNAKE_CASE       :     Any				= outputs.reshaped_hidden_states
			self.assertEqual(len(lowerCamelCase_					)		, lowerCamelCase_					)

			SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE       :     Any				= reshaped_hidden_states[0].shape
			SCREAMING_SNAKE_CASE       :     Optional[int]				= (
			    reshaped_hidden_states[0].view(lowerCamelCase_		, lowerCamelCase_		, height * width					).permute(0		, 2		, 1					)
			)
			self.assertListEqual(
			    list(reshaped_hidden_states.shape[-2:]					)		, [num_patches, self.model_tester.embed_dim]		, )





	def   __lowerCAmelCase (						self    :Optional[Any]					)							->			Optional[int]:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE       :     Union[str, Any]				= self.model_tester.prepare_config_and_inputs_for_common()

			SCREAMING_SNAKE_CASE       :     List[str]				= (
			    self.model_tester.image_size
			    if isinstance(self.model_tester.image_size		, collections.abc.Iterable					)
			    else (self.model_tester.image_size, self.model_tester.image_size)
			)

			for model_class in self.all_model_classes:
					SCREAMING_SNAKE_CASE       :     Tuple				= True
					self.check_hidden_states_output(lowerCamelCase_		, lowerCamelCase_		, lowerCamelCase_		, lowerCamelCase_					)

					# check that output_hidden_states also work using config
					del inputs_dict["output_hidden_states"]
					SCREAMING_SNAKE_CASE       :     List[str]				= True

					self.check_hidden_states_output(lowerCamelCase_		, lowerCamelCase_		, lowerCamelCase_		, lowerCamelCase_					)





	def   __lowerCAmelCase (						self    :str					)							->			Optional[int]:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE       :     str				= self.model_tester.prepare_config_and_inputs_for_common()
			SCREAMING_SNAKE_CASE       :     List[Any]				= 3

			SCREAMING_SNAKE_CASE       :     int				= (
			    self.model_tester.image_size
			    if isinstance(self.model_tester.image_size		, collections.abc.Iterable					)
			    else (self.model_tester.image_size, self.model_tester.image_size)
			)
			SCREAMING_SNAKE_CASE       :     Optional[Any]				= (
			    config.patch_size
			    if isinstance(config.patch_size		, collections.abc.Iterable					)
			    else (config.patch_size, config.patch_size)
			)

			SCREAMING_SNAKE_CASE       :     Optional[int]				= image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
			SCREAMING_SNAKE_CASE       :     Optional[int]				= image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

			for model_class in self.all_model_classes:
					SCREAMING_SNAKE_CASE       :     Optional[Any]				= True
					self.check_hidden_states_output(lowerCamelCase_		, lowerCamelCase_		, lowerCamelCase_		, (padded_height, padded_width)					)

					# check that output_hidden_states also work using config
					del inputs_dict["output_hidden_states"]
					SCREAMING_SNAKE_CASE       :     Optional[Any]				= True
					self.check_hidden_states_output(lowerCamelCase_		, lowerCamelCase_		, lowerCamelCase_		, (padded_height, padded_width)					)





	def   __lowerCAmelCase (						self    :str					)							->			List[Any]:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE       :     Tuple				= self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_masked_image_modeling(*lowerCamelCase_					)





	def   __lowerCAmelCase (						self    :Tuple					)							->			Optional[Any]:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE       :     Tuple				= self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_image_classification(*lowerCamelCase_					)





	@slow
	def   __lowerCAmelCase (						self    :Tuple					)							->			List[str]:







			'''simple docstring'''




			for model_name in SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					SCREAMING_SNAKE_CASE       :     Tuple				= SwinvaModel.from_pretrained(lowerCamelCase_					)
					self.assertIsNotNone(lowerCamelCase_					)







	def   __lowerCAmelCase (						self    :int					)							->			Any:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE       :     List[str]				= self.model_tester.prepare_config_and_inputs_for_common()

			SCREAMING_SNAKE_CASE       :     Tuple				= _config_zero_init(lowerCamelCase_					)
			for model_class in self.all_model_classes:
					SCREAMING_SNAKE_CASE       :     Dict				= model_class(config=lowerCamelCase_					)
					for name, param in model.named_parameters():
							if "embeddings" not in name and "logit_scale" not in name and param.requires_grad:
									self.assertIn(
									    ((param.data.mean() * 1E9).round() / 1E9).item()		, [0.0, 1.0]		, msg=f"Parameter {name} of model {model_class} seems not properly initialized"		, )





@require_vision
@require_torch
class 				lowercase__(				unittest.TestCase   ):
	'''simple docstring'''

	@cached_property
	def   __lowerCAmelCase (						self    :Dict					)							->			List[Any]:







			'''simple docstring'''




			return (
			    AutoImageProcessor.from_pretrained('''microsoft/swinv2-tiny-patch4-window8-256'''					)
			    if is_vision_available()
			    else None
			)







	@slow
	def   __lowerCAmelCase (						self    :Dict					)							->			Tuple:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE       :     Optional[int]				= SwinvaForImageClassification.from_pretrained('''microsoft/swinv2-tiny-patch4-window8-256'''					).to(
			    lowerCamelCase_					)
			SCREAMING_SNAKE_CASE       :     Dict				= self.default_image_processor

			SCREAMING_SNAKE_CASE       :     Optional[Any]				= Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''					)
			SCREAMING_SNAKE_CASE       :     List[str]				= image_processor(images=lowerCamelCase_		, return_tensors='''pt'''					).to(lowerCamelCase_					)

			# forward pass
			with torch.no_grad():
					SCREAMING_SNAKE_CASE       :     Tuple				= model(**lowerCamelCase_					)

			# verify the logits
			SCREAMING_SNAKE_CASE       :     List[str]				= torch.Size((1, 10_00)					)
			self.assertEqual(outputs.logits.shape		, lowerCamelCase_					)
			SCREAMING_SNAKE_CASE       :     int				= torch.tensor([-0.3_9_4_7, -0.4_3_0_6, 0.0_0_2_6]					).to(lowerCamelCase_					)
			self.assertTrue(torch.allclose(outputs.logits[0, :3]		, lowerCamelCase_		, atol=1E-4					)					)





