






import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionAttendAndExcitePipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_numpy, skip_mps, slow
from diffusers.utils.testing_utils import require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


lowerCamelCase_       			=   False







@skip_mps
class     UpperCamelCase_ (lowercase__       ,		lowercase__       ,		lowercase__       ,		unittest.TestCase  ):
	__magic_name__			   =       StableDiffusionAttendAndExcitePipeline
	__magic_name__			   =       False
	__magic_name__			   =       TEXT_TO_IMAGE_PARAMS
	__magic_name__			   =       TEXT_TO_IMAGE_BATCH_PARAMS.union({'''token_indices'''}  )
	__magic_name__			   =       TEXT_TO_IMAGE_IMAGE_PARAMS
	__magic_name__			   =       TEXT_TO_IMAGE_IMAGE_PARAMS





	@classmethod
	def   _SCREAMING_SNAKE_CASE   (   cls      :				Any		)		->    Any:
				super().setUpClass()
				torch.use_deterministic_algorithms(UpperCAmelCase__		)





	@classmethod
	def   _SCREAMING_SNAKE_CASE   (   cls      :				Optional[int]		)		->    Optional[int]:
				super().tearDownClass()
				torch.use_deterministic_algorithms(UpperCAmelCase__		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				List[str]		)		->    Optional[Any]:
				torch.manual_seed(0		)
				UpperCAmelCase_  :				Union[str, Any]      =      UNetaDConditionModel(
				    block_out_channels=(32, 64)			,			layers_per_block=1			,			sample_size=32			,			in_channels=4			,			out_channels=4			,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")			,			up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")			,			cross_attention_dim=32			,			attention_head_dim=(2, 4)			,			use_linear_projection=UpperCAmelCase__			,			)
				UpperCAmelCase_  :				List[str]      =      DDIMScheduler(
				    beta_start=0.0_0_0_8_5			,			beta_end=0.0_1_2			,			beta_schedule="scaled_linear"			,			clip_sample=UpperCAmelCase__			,			set_alpha_to_one=UpperCAmelCase__			,			)
				torch.manual_seed(0		)
				UpperCAmelCase_  :				List[str]      =      AutoencoderKL(
				    block_out_channels=[32, 64]			,			in_channels=3			,			out_channels=3			,			down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]			,			up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]			,			latent_channels=4			,			sample_size=128			,			)
				torch.manual_seed(0		)
				UpperCAmelCase_  :				int      =      CLIPTextConfig(
				    bos_token_id=0			,			eos_token_id=2			,			hidden_size=32			,			intermediate_size=37			,			layer_norm_eps=1e-05			,			num_attention_heads=4			,			num_hidden_layers=5			,			pad_token_id=1			,			vocab_size=1_000			,			hidden_act="gelu"			,			projection_dim=512			,			)
				UpperCAmelCase_  :				int      =      CLIPTextModel(UpperCAmelCase__		)
				UpperCAmelCase_  :				List[Any]      =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

				UpperCAmelCase_  :				Optional[Any]      =      {
				    '''unet''': unet,
				    '''scheduler''': scheduler,
				    '''vae''': vae,
				    '''text_encoder''': text_encoder,
				    '''tokenizer''': tokenizer,
				    '''safety_checker''': None,
				    '''feature_extractor''': None,
				}

				return components





	def   _SCREAMING_SNAKE_CASE   (   self      :				int			,			lowerCAmelCase_      :				List[Any]			,			lowerCAmelCase_      :				Optional[int]=0		)		->    List[Any]:
				if str(UpperCAmelCase__		).startswith("mps"		):
							UpperCAmelCase_  :				List[str]      =      torch.manual_seed(UpperCAmelCase__		)
				else:
							UpperCAmelCase_  :				Dict      =      torch.Generator(device=UpperCAmelCase__		).manual_seed(UpperCAmelCase__		)
				UpperCAmelCase_  :				Optional[int]      =      {
				    '''prompt''': '''a cat and a frog''',
				    '''token_indices''': [2, 5],
				    '''generator''': generator,
				    '''num_inference_steps''': 1,
				    '''guidance_scale''': 6.0,
				    '''output_type''': '''numpy''',
				    '''max_iter_to_alter''': 2,
				    '''thresholds''': {0: 0.7},
				}
				return inputs





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[Any]		)		->    Dict:
				UpperCAmelCase_  :				int      =      '''cpu'''

				UpperCAmelCase_  :				Tuple      =      self.get_dummy_components()
				UpperCAmelCase_  :				Optional[Any]      =      self.pipeline_class(**UpperCAmelCase__		)
				pipe.to(UpperCAmelCase__		)
				pipe.set_progress_bar_config(disable=UpperCAmelCase__		)

				UpperCAmelCase_  :				int      =      self.get_dummy_inputs(UpperCAmelCase__		)
				UpperCAmelCase_  :				Union[str, Any]      =      pipe(**UpperCAmelCase__		).images
				UpperCAmelCase_  :				Optional[Any]      =      image[0, -3:, -3:, -1]

				self.assertEqual(image.shape			,			(1, 64, 64, 3)		)
				UpperCAmelCase_  :				Optional[Any]      =      np.array(
				    [0.6_3_9_0_5_3_6_4, 0.6_2_8_9_7_3_0_7, 0.4_8_5_9_9_0_1_7, 0.5_1_3_3_6_2_4, 0.5_5_5_0_0_4_8, 0.4_5_7_6_9_5_1_6, 0.5_0_3_2_6_9_7_3, 0.5_0_2_3_1_3_9, 0.4_5_3_8_4_4_9_6]		)
				UpperCAmelCase_  :				List[Any]      =      np.abs(image_slice.flatten() - expected_slice		).max()
				self.assertLessEqual(UpperCAmelCase__			,			1e-3		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				List[str]		)		->    int:
				super().test_cpu_offload_forward_pass(expected_max_diff=5e-4		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				List[str]		)		->    List[Any]:
				self._test_inference_batch_consistent(batch_sizes=[1, 2]		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				Any		)		->    List[Any]:
				self._test_inference_batch_single_identical(batch_size=2			,			expected_max_diff=7e-4		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[int]		)		->    Optional[int]:
				super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				str		)		->    int:
				super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5e-4		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				List[str]		)		->    List[str]:
				super().test_save_load_local(expected_max_difference=5e-4		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				List[str]		)		->    Union[str, Any]:
				super().test_save_load_optional_components(expected_max_difference=4e-4		)







@require_torch_gpu
@slow
class     UpperCamelCase_ (unittest.TestCase  ):





	@classmethod
	def   _SCREAMING_SNAKE_CASE   (   cls      :				Tuple		)		->    Optional[Any]:
				super().setUpClass()
				torch.use_deterministic_algorithms(UpperCAmelCase__		)





	@classmethod
	def   _SCREAMING_SNAKE_CASE   (   cls      :				Optional[Any]		)		->    int:
				super().tearDownClass()
				torch.use_deterministic_algorithms(UpperCAmelCase__		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[Any]		)		->    List[Any]:
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()





	def   _SCREAMING_SNAKE_CASE   (   self      :				Union[str, Any]		)		->    Any:
				UpperCAmelCase_  :				Dict      =      torch.manual_seed(51		)

				UpperCAmelCase_  :				Any      =      StableDiffusionAttendAndExcitePipeline.from_pretrained(
				    "CompVis/stable-diffusion-v1-4"			,			safety_checker=UpperCAmelCase__			,			torch_dtype=torch.floataa		)
				pipe.to("cuda"		)

				UpperCAmelCase_  :				Dict      =      '''a painting of an elephant with glasses'''
				UpperCAmelCase_  :				Optional[int]      =      [5, 7]

				UpperCAmelCase_  :				Union[str, Any]      =      pipe(
				    prompt=UpperCAmelCase__			,			token_indices=UpperCAmelCase__			,			guidance_scale=7.5			,			generator=UpperCAmelCase__			,			num_inference_steps=5			,			max_iter_to_alter=5			,			output_type="numpy"			,			).images[0]

				UpperCAmelCase_  :				Optional[int]      =      load_numpy(
				    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/attend-and-excite/elephant_glasses.npy"		)
				assert np.abs((expected_image - image).max()		) < 5e-1











"""simple docstring"""
import argparse
import json
from collections import OrderedDict

import torch
from huggingface_hub import cached_download, hf_hub_url

from transformers import AutoImageProcessor, CvtConfig, CvtForImageClassification

def        snake_case   (     A__		):
			UpperCAmelCase_  :				Tuple      =      []
			embed.append(
			    (
			        F"""cvt.encoder.stages.{idx}.embedding.convolution_embeddings.projection.weight""",
			        F"""stage{idx}.patch_embed.proj.weight""",
			    )		)
			embed.append(
			    (
			        F"""cvt.encoder.stages.{idx}.embedding.convolution_embeddings.projection.bias""",
			        F"""stage{idx}.patch_embed.proj.bias""",
			    )		)
			embed.append(
			    (
			        F"""cvt.encoder.stages.{idx}.embedding.convolution_embeddings.normalization.weight""",
			        F"""stage{idx}.patch_embed.norm.weight""",
			    )		)
			embed.append(
			    (
			        F"""cvt.encoder.stages.{idx}.embedding.convolution_embeddings.normalization.bias""",
			        F"""stage{idx}.patch_embed.norm.bias""",
			    )		)
			return embed

def        snake_case   (     A__			,A__		):
			UpperCAmelCase_  :				List[Any]      =      []
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.convolution.weight""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_q.conv.weight""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.weight""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.weight""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.bias""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.bias""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.running_mean""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.running_mean""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.running_var""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.running_var""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.num_batches_tracked""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.num_batches_tracked""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.convolution.weight""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_k.conv.weight""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.weight""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.weight""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.bias""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.bias""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.running_mean""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.running_mean""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.running_var""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.running_var""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.num_batches_tracked""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.num_batches_tracked""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.convolution.weight""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_v.conv.weight""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.weight""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.weight""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.bias""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.bias""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.running_mean""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.running_mean""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.running_var""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.running_var""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.num_batches_tracked""",
			        F"""stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.num_batches_tracked""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_query.weight""",
			        F"""stage{idx}.blocks.{cnt}.attn.proj_q.weight""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_query.bias""",
			        F"""stage{idx}.blocks.{cnt}.attn.proj_q.bias""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_key.weight""",
			        F"""stage{idx}.blocks.{cnt}.attn.proj_k.weight""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_key.bias""",
			        F"""stage{idx}.blocks.{cnt}.attn.proj_k.bias""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_value.weight""",
			        F"""stage{idx}.blocks.{cnt}.attn.proj_v.weight""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_value.bias""",
			        F"""stage{idx}.blocks.{cnt}.attn.proj_v.bias""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.output.dense.weight""",
			        F"""stage{idx}.blocks.{cnt}.attn.proj.weight""",
			    )		)
			attention_weights.append(
			    (
			        F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.output.dense.bias""",
			        F"""stage{idx}.blocks.{cnt}.attn.proj.bias""",
			    )		)
			attention_weights.append(
			    (F"""cvt.encoder.stages.{idx}.layers.{cnt}.intermediate.dense.weight""", F"""stage{idx}.blocks.{cnt}.mlp.fc1.weight""")		)
			attention_weights.append(
			    (F"""cvt.encoder.stages.{idx}.layers.{cnt}.intermediate.dense.bias""", F"""stage{idx}.blocks.{cnt}.mlp.fc1.bias""")		)
			attention_weights.append(
			    (F"""cvt.encoder.stages.{idx}.layers.{cnt}.output.dense.weight""", F"""stage{idx}.blocks.{cnt}.mlp.fc2.weight""")		)
			attention_weights.append(
			    (F"""cvt.encoder.stages.{idx}.layers.{cnt}.output.dense.bias""", F"""stage{idx}.blocks.{cnt}.mlp.fc2.bias""")		)
			attention_weights.append(
			    (F"""cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_before.weight""", F"""stage{idx}.blocks.{cnt}.norm1.weight""")		)
			attention_weights.append(
			    (F"""cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_before.bias""", F"""stage{idx}.blocks.{cnt}.norm1.bias""")		)
			attention_weights.append(
			    (F"""cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_after.weight""", F"""stage{idx}.blocks.{cnt}.norm2.weight""")		)
			attention_weights.append(
			    (F"""cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_after.bias""", F"""stage{idx}.blocks.{cnt}.norm2.bias""")		)
			return attention_weights

def        snake_case   (     A__		):
			UpperCAmelCase_  :				Union[str, Any]      =      []
			token.append((F"""cvt.encoder.stages.{idx}.cls_token""", "stage2.cls_token")		)
			return token

def        snake_case   (     ):
			UpperCAmelCase_  :				str      =      []
			head.append(("layernorm.weight", "norm.weight")		)
			head.append(("layernorm.bias", "norm.bias")		)
			head.append(("classifier.weight", "head.weight")		)
			head.append(("classifier.bias", "head.bias")		)
			return head

def        snake_case   (     A__			,A__			,A__			,A__		):
			UpperCAmelCase_  :				int      =      "imagenet-1k-id2label.json"
			UpperCAmelCase_  :				Union[str, Any]      =      10_00

			UpperCAmelCase_  :				Any      =      "huggingface/label-files"
			UpperCAmelCase_  :				Optional[int]      =      num_labels
			UpperCAmelCase_  :				Union[str, Any]      =      json.load(open(cached_download(hf_hub_url(A__			,A__			,repo_type="dataset"		)		)			,"r"		)		)
			UpperCAmelCase_  :				Optional[int]      =      {int(A__		): v for k, v in idalabel.items()}

			UpperCAmelCase_  :				Optional[int]      =      idalabel
			UpperCAmelCase_  :				Optional[Any]      =      {v: k for k, v in idalabel.items()}

			UpperCAmelCase_  :				str      =      CvtConfig(num_labels=A__			,idalabel=A__			,labelaid=A__		)

			# For depth size 13 (13 = 1+2+10)
			if cvt_model.rsplit("/"			,1		)[-1][4:6] == "13":
						UpperCAmelCase_  :				Optional[int]      =      [1, 2, 10]

			# For depth size 21 (21 = 1+4+16)
			elif cvt_model.rsplit("/"			,1		)[-1][4:6] == "21":
						UpperCAmelCase_  :				str      =      [1, 4, 16]

			# For wide cvt (similar to wide-resnet) depth size 24 (w24 = 2 + 2 20)
			else:
						UpperCAmelCase_  :				Optional[int]      =      [2, 2, 20]
						UpperCAmelCase_  :				List[str]      =      [3, 12, 16]
						UpperCAmelCase_  :				Optional[Any]      =      [1_92, 7_68, 10_24]

			UpperCAmelCase_  :				Tuple      =      CvtForImageClassification(A__		)
			UpperCAmelCase_  :				Tuple      =      AutoImageProcessor.from_pretrained("facebook/convnext-base-224-22k-1k"		)
			UpperCAmelCase_  :				int      =      image_size
			UpperCAmelCase_  :				List[str]      =      torch.load(A__			,map_location=torch.device("cpu"		)		)

			UpperCAmelCase_  :				int      =      OrderedDict()
			UpperCAmelCase_  :				Optional[Any]      =      []

			for idx in range(len(config.depth		)		):
						if config.cls_token[idx]:
									UpperCAmelCase_  :				Optional[Any]      =      list_of_state_dict + cls_token(A__		)
						UpperCAmelCase_  :				Any      =      list_of_state_dict + embeddings(A__		)
						for cnt in range(config.depth[idx]		):
									UpperCAmelCase_  :				Dict      =      list_of_state_dict + attention(A__			,A__		)

			UpperCAmelCase_  :				Union[str, Any]      =      list_of_state_dict + final()
			for gg in list_of_state_dict:
						print(A__		)
			for i in range(len(A__		)		):
						UpperCAmelCase_  :				Any      =      original_weights[list_of_state_dict[i][1]]

			model.load_state_dict(A__		)
			model.save_pretrained(A__		)
			image_processor.save_pretrained(A__		)


# Download the weights from zoo: https://1drv.ms/u/s!AhIXJn_J-blW9RzF3rMW7SsLHa8h?e=blQ0Al

if __name__ == "__main__":
			lowerCamelCase_       			=   argparse.ArgumentParser()
			parser.add_argument(
			    '''--cvt_model''',
			    default='''cvt-w24''',
			    type=str,
			    help='''Name of the cvt model you\'d like to convert.''',
			)
			parser.add_argument(
			    '''--image_size''',
			    default=384,
			    type=int,
			    help='''Input Image Size''',
			)
			parser.add_argument(
			    '''--cvt_file_name''',
			    default=r'''cvtmodels\CvT-w24-384x384-IN-22k.pth''',
			    type=str,
			    help='''Input Image Size''',
			)
			parser.add_argument(
			    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
			)

			lowerCamelCase_       			=   parser.parse_args()
			convert_cvt_checkpoint(args.cvt_model, args.image_size, args.cvt_file_name, args.pytorch_dump_folder_path)



