



"""simple docstring"""






import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTNeoConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
		import jax
		import jax.numpy as jnp

		from transformers.modeling_flax_pytorch_utils import (
		    convert_pytorch_state_dict_to_flax,
		    load_flax_weights_in_pytorch_model,
		)
		from transformers.models.gpt_neo.modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel

if is_torch_available():
		import torch

class 			__SCREAMING_SNAKE_CASE      :


					'''simple docstring'''





					def __init__(			self				:    Optional[int],    lowerCamelCase				:    int,    lowerCamelCase				:    Union[str, Any]=14,    lowerCamelCase				:    List[str]=7,    lowerCamelCase				:    int=True,    lowerCamelCase				:    List[str]=True,    lowerCamelCase				:    Dict=False,    lowerCamelCase				:    List[str]=True,    lowerCamelCase				:    Optional[int]=99,    lowerCamelCase				:    Any=32,    lowerCamelCase				:    List[Any]=4,    lowerCamelCase				:    Any=4,    lowerCamelCase				:    Optional[int]=[[["global", "local"], 2]],    lowerCamelCase				:    Any=37,    lowerCamelCase				:    Optional[Any]="gelu",    lowerCamelCase				:    str=0.1,    lowerCamelCase				:    Optional[Any]=0.1,    lowerCamelCase				:    str=512,    lowerCamelCase				:    Dict=7,    lowerCamelCase				:    Optional[int]=0.02,    )->   List[str]:
								lowerCamelCase__						:						Optional[int]						      =parent
								lowerCamelCase__						:						List[str]						      =batch_size
								lowerCamelCase__						:						List[str]						      =seq_length
								lowerCamelCase__						:						Optional[Any]						      =is_training
								lowerCamelCase__						:						int						      =use_input_mask
								lowerCamelCase__						:						List[str]						      =use_token_type_ids
								lowerCamelCase__						:						Optional[int]						      =use_labels
								lowerCamelCase__						:						Tuple						      =vocab_size
								lowerCamelCase__						:						Optional[Any]						      =hidden_size
								lowerCamelCase__						:						str						      =num_hidden_layers
								lowerCamelCase__						:						List[Any]						      =num_attention_heads
								lowerCamelCase__						:						Optional[Any]						      =attention_types
								lowerCamelCase__						:						Tuple						      =intermediate_size
								lowerCamelCase__						:						List[Any]						      =hidden_act
								lowerCamelCase__						:						str						      =hidden_dropout_prob
								lowerCamelCase__						:						List[str]						      =attention_probs_dropout_prob
								lowerCamelCase__						:						str						      =max_position_embeddings
								lowerCamelCase__						:						Dict						      =window_size
								lowerCamelCase__						:						str						      =initializer_range
								lowerCamelCase__						:						int						      =None
								lowerCamelCase__						:						int						      =vocab_size - 1
								lowerCamelCase__						:						Any						      =vocab_size - 1
								lowerCamelCase__						:						List[str]						      =vocab_size - 1



					def 						snake_case			(			self				:    Union[str, Any]	)->   str:
								lowerCamelCase__						:						Optional[int]						      =ids_tensor([self.batch_size, self.seq_length],    self.vocab_size	)

								lowerCamelCase__						:						int						      =None
								if self.use_input_mask:
											lowerCamelCase__						:						Any						      =random_attention_mask([self.batch_size, self.seq_length]	)

								lowerCamelCase__						:						Optional[int]						      =GPTNeoConfig(
								    vocab_size=self.vocab_size,    hidden_size=self.hidden_size,    num_layers=self.num_hidden_layers,    num_heads=self.num_attention_heads,    max_position_embeddings=self.max_position_embeddings,    use_cache=lowerCamelCase,    bos_token_id=self.bos_token_id,    eos_token_id=self.eos_token_id,    pad_token_id=self.pad_token_id,    window_size=self.window_size,    attention_types=self.attention_types,    )

								return (config, input_ids, input_mask)



					def 						snake_case			(			self				:    Tuple	)->   int:
								lowerCamelCase__						:						List[str]						      =self.prepare_config_and_inputs()
								lowerCamelCase__      ,  lowerCamelCase__      ,  lowerCamelCase__						:						Optional[int]						      =config_and_inputs
								lowerCamelCase__						:						List[str]						      ={'''input_ids''': input_ids, '''attention_mask''': attention_mask}
								return config, inputs_dict



					def 						snake_case			(			self				:    str,    lowerCamelCase				:    Union[str, Any],    lowerCamelCase				:    Any,    lowerCamelCase				:    Dict,    lowerCamelCase				:    Dict	)->   Tuple:
								lowerCamelCase__						:						Any						      =20
								lowerCamelCase__						:						Optional[Any]						      =model_class_name(lowerCamelCase	)

								lowerCamelCase__						:						List[Any]						      =model.init_cache(input_ids.shape[0],    lowerCamelCase	)
								lowerCamelCase__						:						Optional[Any]						      =jnp.ones((input_ids.shape[0], max_decoder_length),    dtype='''i4'''	)

								lowerCamelCase__						:						Any						      =jnp.broadcast_to(
								    jnp.arange(input_ids.shape[-1] - 1	)[None, :],    (input_ids.shape[0], input_ids.shape[-1] - 1)	)
								lowerCamelCase__						:						Union[str, Any]						      =model(
								    input_ids[:, :-1],    attention_mask=lowerCamelCase,    past_key_values=lowerCamelCase,    position_ids=lowerCamelCase,    )

								lowerCamelCase__						:						Any						      =jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]],    dtype='''i4'''	)
								lowerCamelCase__						:						Optional[int]						      =model(
								    input_ids[:, -1:],    attention_mask=lowerCamelCase,    past_key_values=outputs_cache.past_key_values,    position_ids=lowerCamelCase,    )

								lowerCamelCase__						:						List[Any]						      =model(lowerCamelCase	)

								lowerCamelCase__						:						Any						      =np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
								self.parent.assertTrue(diff < 1E-3,    msg=F'''Max diff is {diff}'''	)




					def 						snake_case			(			self				:    str,    lowerCamelCase				:    Optional[int],    lowerCamelCase				:    List[str],    lowerCamelCase				:    Any,    lowerCamelCase				:    Any	)->   List[str]:
								lowerCamelCase__						:						List[Any]						      =20
								lowerCamelCase__						:						str						      =model_class_name(lowerCamelCase	)

								lowerCamelCase__						:						List[str]						      =jnp.concatenate(
								    [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])	)],    axis=-1,    )

								lowerCamelCase__						:						Tuple						      =model.init_cache(input_ids.shape[0],    lowerCamelCase	)
								lowerCamelCase__						:						Union[str, Any]						      =jnp.broadcast_to(
								    jnp.arange(input_ids.shape[-1] - 1	)[None, :],    (input_ids.shape[0], input_ids.shape[-1] - 1)	)

								lowerCamelCase__						:						str						      =model(
								    input_ids[:, :-1],    attention_mask=lowerCamelCase,    past_key_values=lowerCamelCase,    position_ids=lowerCamelCase,    )
								lowerCamelCase__						:						Any						      =jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]],    dtype='''i4'''	)
								lowerCamelCase__						:						Any						      =model(
								    input_ids[:, -1:],    past_key_values=outputs_cache.past_key_values,    attention_mask=lowerCamelCase,    position_ids=lowerCamelCase,    )

								lowerCamelCase__						:						Dict						      =model(lowerCamelCase,    attention_mask=lowerCamelCase	)

								lowerCamelCase__						:						Tuple						      =np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
								self.parent.assertTrue(diff < 1E-3,    msg=F'''Max diff is {diff}'''	)





@require_flax
class 			__SCREAMING_SNAKE_CASE      ( lowerCAmelCase_						,  lowerCAmelCase_						,  unittest.TestCase  ):


					'''simple docstring'''





					_a       =				(FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
					_a       =				(FlaxGPTNeoForCausalLM,) if is_flax_available() else ()



					def 						snake_case			(			self				:    Optional[Any]	)->   List[Any]:
								lowerCamelCase__						:						Optional[Any]						      =FlaxGPTNeoModelTester(self	)



					def 						snake_case			(			self				:    List[str]	)->   int:
								for model_class_name in self.all_model_classes:
											lowerCamelCase__      ,  lowerCamelCase__      ,  lowerCamelCase__						:						List[Any]						      =self.model_tester.prepare_config_and_inputs()
											self.model_tester.check_use_cache_forward(lowerCamelCase,    lowerCamelCase,    lowerCamelCase,    lowerCamelCase	)



					def 						snake_case			(			self				:    Union[str, Any]	)->   str:
								for model_class_name in self.all_model_classes:
											lowerCamelCase__      ,  lowerCamelCase__      ,  lowerCamelCase__						:						Union[str, Any]						      =self.model_tester.prepare_config_and_inputs()
											self.model_tester.check_use_cache_forward_with_attn_mask(
											    lowerCamelCase,    lowerCamelCase,    lowerCamelCase,    lowerCamelCase	)



					@slow
					def 						snake_case			(			self				:    str	)->   Union[str, Any]:
								lowerCamelCase__						:						str						      =GPTaTokenizer.from_pretrained('''gpt2''',    pad_token='''<|endoftext|>''',    padding_side='''left'''	)
								lowerCamelCase__						:						List[Any]						      =tokenizer(['''Hello this is a long string''', '''Hey'''],    return_tensors='''np''',    padding=lowerCamelCase,    truncation=lowerCamelCase	)

								lowerCamelCase__						:						Any						      =FlaxGPTNeoForCausalLM.from_pretrained('''EleutherAI/gpt-neo-125M'''	)
								lowerCamelCase__						:						Any						      =False
								lowerCamelCase__						:						str						      =model.config.eos_token_id

								lowerCamelCase__						:						Any						      =jax.jit(model.generate	)

								lowerCamelCase__						:						List[Any]						      =jit_generate(
								    inputs['''input_ids'''],    attention_mask=inputs['''attention_mask'''],    pad_token_id=tokenizer.pad_token_id	).sequences

								lowerCamelCase__						:						int						      =tokenizer.batch_decode(lowerCamelCase,    skip_special_tokens=lowerCamelCase	)

								lowerCamelCase__						:						List[str]						      =[
								    '''Hello this is a long string of text.\n\nI\'m trying to get the text of the''',
								    '''Hey, I\'m a little late to the party. I\'m going to''',
								]

								self.assertListEqual(lowerCamelCase,    lowerCamelCase	)



					@is_pt_flax_cross_test
					def 						snake_case			(			self				:    Tuple	)->   Tuple:
								lowerCamelCase__      ,  lowerCamelCase__						:						List[Any]						      =self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
											with self.subTest(model_class.__name__	):
														# prepare inputs
														lowerCamelCase__						:						Optional[int]						      =self._prepare_for_class(lowerCamelCase,    lowerCamelCase	)
														lowerCamelCase__						:						Optional[int]						      ={k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

														# load corresponding PyTorch class
														lowerCamelCase__						:						List[str]						      =model_class.__name__[4:]  # Skip the "Flax" at the beginning
														lowerCamelCase__						:						List[Any]						      =getattr(lowerCamelCase,    lowerCamelCase	)

														lowerCamelCase__      ,  lowerCamelCase__						:						List[str]						      =pt_inputs['''input_ids'''].shape
														lowerCamelCase__						:						Dict						      =np.random.randint(0,    seq_length - 1,    size=(batch_size,)	)
														for batch_idx, start_index in enumerate(lowerCamelCase	):
																	lowerCamelCase__						:						List[str]						      =0
																	lowerCamelCase__						:						Dict						      =1
																	lowerCamelCase__						:						Optional[int]						      =0
																	lowerCamelCase__						:						Optional[int]						      =1
														lowerCamelCase__						:						Dict						      =pt_model_class(lowerCamelCase	).eval()
														lowerCamelCase__						:						Dict						      =model_class(lowerCamelCase,    dtype=jnp.floataa	)

														lowerCamelCase__						:						Union[str, Any]						      =convert_pytorch_state_dict_to_flax(pt_model.state_dict(),    lowerCamelCase	)
														lowerCamelCase__						:						Union[str, Any]						      =fx_state

														with torch.no_grad():
																	lowerCamelCase__						:						Optional[int]						      =pt_model(**lowerCamelCase	).to_tuple()

														lowerCamelCase__						:						Tuple						      =fx_model(**lowerCamelCase	).to_tuple()
														self.assertEqual(len(lowerCamelCase	),    len(lowerCamelCase	),    '''Output lengths differ between Flax and PyTorch'''	)
														for fx_output, pt_output in zip(lowerCamelCase,    lowerCamelCase	):
																	self.assert_almost_equals(fx_output[:, -1],    pt_output[:, -1].numpy(),    4E-2	)

														with tempfile.TemporaryDirectory() as tmpdirname:
																	pt_model.save_pretrained(lowerCamelCase	)
																	lowerCamelCase__						:						Any						      =model_class.from_pretrained(lowerCamelCase,    from_pt=lowerCamelCase	)

														lowerCamelCase__						:						List[str]						      =fx_model_loaded(**lowerCamelCase	).to_tuple()
														self.assertEqual(
														    len(lowerCamelCase	),    len(lowerCamelCase	),    '''Output lengths differ between Flax and PyTorch'''	)
														for fx_output_loaded, pt_output in zip(lowerCamelCase,    lowerCamelCase	):
																	self.assert_almost_equals(fx_output_loaded[:, -1],    pt_output[:, -1].numpy(),    4E-2	)



					@is_pt_flax_cross_test
					def 						snake_case			(			self				:    int	)->   Dict:
								lowerCamelCase__      ,  lowerCamelCase__						:						Any						      =self.model_tester.prepare_config_and_inputs_for_common()
								for model_class in self.all_model_classes:
											with self.subTest(model_class.__name__	):
														# prepare inputs
														lowerCamelCase__						:						str						      =self._prepare_for_class(lowerCamelCase,    lowerCamelCase	)
														lowerCamelCase__						:						Tuple						      ={k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

														# load corresponding PyTorch class
														lowerCamelCase__						:						str						      =model_class.__name__[4:]  # Skip the "Flax" at the beginning
														lowerCamelCase__						:						List[str]						      =getattr(lowerCamelCase,    lowerCamelCase	)

														lowerCamelCase__						:						Optional[Any]						      =pt_model_class(lowerCamelCase	).eval()
														lowerCamelCase__						:						List[Any]						      =model_class(lowerCamelCase,    dtype=jnp.floataa	)

														lowerCamelCase__						:						List[Any]						      =load_flax_weights_in_pytorch_model(lowerCamelCase,    fx_model.params	)
														lowerCamelCase__      ,  lowerCamelCase__						:						int						      =pt_inputs['''input_ids'''].shape
														lowerCamelCase__						:						Union[str, Any]						      =np.random.randint(0,    seq_length - 1,    size=(batch_size,)	)
														for batch_idx, start_index in enumerate(lowerCamelCase	):
																	lowerCamelCase__						:						Union[str, Any]						      =0
																	lowerCamelCase__						:						List[str]						      =1
																	lowerCamelCase__						:						Dict						      =0
																	lowerCamelCase__						:						List[Any]						      =1

														# make sure weights are tied in PyTorch
														pt_model.tie_weights()

														with torch.no_grad():
																	lowerCamelCase__						:						Optional[Any]						      =pt_model(**lowerCamelCase	).to_tuple()

														lowerCamelCase__						:						str						      =fx_model(**lowerCamelCase	).to_tuple()
														self.assertEqual(len(lowerCamelCase	),    len(lowerCamelCase	),    '''Output lengths differ between Flax and PyTorch'''	)
														for fx_output, pt_output in zip(lowerCamelCase,    lowerCamelCase	):
																	self.assert_almost_equals(fx_output[:, -1],    pt_output[:, -1].numpy(),    4E-2	)

														with tempfile.TemporaryDirectory() as tmpdirname:
																	fx_model.save_pretrained(lowerCamelCase	)
																	lowerCamelCase__						:						Tuple						      =pt_model_class.from_pretrained(lowerCamelCase,    from_flax=lowerCamelCase	)

														with torch.no_grad():
																	lowerCamelCase__						:						List[str]						      =pt_model_loaded(**lowerCamelCase	).to_tuple()

														self.assertEqual(
														    len(lowerCamelCase	),    len(lowerCamelCase	),    '''Output lengths differ between Flax and PyTorch'''	)
														for fx_output, pt_output in zip(lowerCamelCase,    lowerCamelCase	):
																	self.assert_almost_equals(fx_output[:, -1],    pt_output[:, -1].numpy(),    4E-2	)




					@slow
					def 						snake_case			(			self				:    Optional[Any]	)->   Optional[Any]:
								for model_class_name in self.all_model_classes:
											lowerCamelCase__						:						int						      =model_class_name.from_pretrained('''EleutherAI/gpt-neo-125M'''	)
											lowerCamelCase__						:						Optional[int]						      =model(np.ones((1, 1)	)	)
											self.assertIsNotNone(lowerCamelCase	)











import argparse
import os

import torch
from torchvision.datasets.utils import download_url

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, TransformeraDModel


__lowercase				    =      {512: """DiT-XL-2-512x512.pt""", 256: """DiT-XL-2-256x256.pt"""}


def        __SCREAMING_SNAKE_CASE	(				__lowercase:		Any):



  '''simple docstring'''






  __lowerCAmelCase								=					F'''pretrained_models/{model_name}'''
  if not os.path.isfile(__lowercase):
    os.makedirs('''pretrained_models'''  ,				exist_ok=__lowercase)
    __lowerCAmelCase								=					F'''https://dl.fbaipublicfiles.com/DiT/models/{model_name}'''
    download_url(__lowercase  ,				'''pretrained_models''')
  __lowerCAmelCase								=					torch.load(__lowercase  ,				map_location=lambda __lowercase  ,				__lowercase: storage)
  return model


def        __SCREAMING_SNAKE_CASE	(				__lowercase:		Tuple):



  '''simple docstring'''






  __lowerCAmelCase								=					download_model(pretrained_models[args.image_size])

  __lowerCAmelCase								=					state_dict['''x_embedder.proj.weight''']
  __lowerCAmelCase								=					state_dict['''x_embedder.proj.bias''']
  state_dict.pop('''x_embedder.proj.weight''')
  state_dict.pop('''x_embedder.proj.bias''')

  for depth in range(28):
    __lowerCAmelCase								=					state_dict[
        '''t_embedder.mlp.0.weight'''
    ]
    __lowerCAmelCase								=					state_dict[
        '''t_embedder.mlp.0.bias'''
    ]
    __lowerCAmelCase								=					state_dict[
        '''t_embedder.mlp.2.weight'''
    ]
    __lowerCAmelCase								=					state_dict[
        '''t_embedder.mlp.2.bias'''
    ]
    __lowerCAmelCase								=					state_dict[
        '''y_embedder.embedding_table.weight'''
    ]

    __lowerCAmelCase								=					state_dict[
        F'''blocks.{depth}.adaLN_modulation.1.weight'''
    ]
    __lowerCAmelCase								=					state_dict[
        F'''blocks.{depth}.adaLN_modulation.1.bias'''
    ]

    __lowerCAmelCase						,    __lowerCAmelCase						,    __lowerCAmelCase								=					torch.chunk(state_dict[F'''blocks.{depth}.attn.qkv.weight''']  ,				3  ,				dim=0)
    __lowerCAmelCase						,    __lowerCAmelCase						,    __lowerCAmelCase								=					torch.chunk(state_dict[F'''blocks.{depth}.attn.qkv.bias''']  ,				3  ,				dim=0)

    __lowerCAmelCase								=					q
    __lowerCAmelCase								=					q_bias
    __lowerCAmelCase								=					k
    __lowerCAmelCase								=					k_bias
    __lowerCAmelCase								=					v
    __lowerCAmelCase								=					v_bias

    __lowerCAmelCase								=					state_dict[
        F'''blocks.{depth}.attn.proj.weight'''
    ]
    __lowerCAmelCase								=					state_dict[F'''blocks.{depth}.attn.proj.bias''']

    __lowerCAmelCase								=					state_dict[F'''blocks.{depth}.mlp.fc1.weight''']
    __lowerCAmelCase								=					state_dict[F'''blocks.{depth}.mlp.fc1.bias''']
    __lowerCAmelCase								=					state_dict[F'''blocks.{depth}.mlp.fc2.weight''']
    __lowerCAmelCase								=					state_dict[F'''blocks.{depth}.mlp.fc2.bias''']

    state_dict.pop(F'''blocks.{depth}.attn.qkv.weight''')
    state_dict.pop(F'''blocks.{depth}.attn.qkv.bias''')
    state_dict.pop(F'''blocks.{depth}.attn.proj.weight''')
    state_dict.pop(F'''blocks.{depth}.attn.proj.bias''')
    state_dict.pop(F'''blocks.{depth}.mlp.fc1.weight''')
    state_dict.pop(F'''blocks.{depth}.mlp.fc1.bias''')
    state_dict.pop(F'''blocks.{depth}.mlp.fc2.weight''')
    state_dict.pop(F'''blocks.{depth}.mlp.fc2.bias''')
    state_dict.pop(F'''blocks.{depth}.adaLN_modulation.1.weight''')
    state_dict.pop(F'''blocks.{depth}.adaLN_modulation.1.bias''')

  state_dict.pop('''t_embedder.mlp.0.weight''')
  state_dict.pop('''t_embedder.mlp.0.bias''')
  state_dict.pop('''t_embedder.mlp.2.weight''')
  state_dict.pop('''t_embedder.mlp.2.bias''')
  state_dict.pop('''y_embedder.embedding_table.weight''')

  __lowerCAmelCase								=					state_dict['''final_layer.adaLN_modulation.1.weight''']
  __lowerCAmelCase								=					state_dict['''final_layer.adaLN_modulation.1.bias''']
  __lowerCAmelCase								=					state_dict['''final_layer.linear.weight''']
  __lowerCAmelCase								=					state_dict['''final_layer.linear.bias''']

  state_dict.pop('''final_layer.linear.weight''')
  state_dict.pop('''final_layer.linear.bias''')
  state_dict.pop('''final_layer.adaLN_modulation.1.weight''')
  state_dict.pop('''final_layer.adaLN_modulation.1.bias''')

  # DiT XL/2
  __lowerCAmelCase								=					TransformeraDModel(
      sample_size=args.image_size // 8  ,				num_layers=28  ,				attention_head_dim=72  ,				in_channels=4  ,				out_channels=8  ,				patch_size=2  ,				attention_bias=__lowercase  ,				num_attention_heads=16  ,				activation_fn='''gelu-approximate'''  ,				num_embeds_ada_norm=10_00  ,				norm_type='''ada_norm_zero'''  ,				norm_elementwise_affine=__lowercase  ,				)
  transformer.load_state_dict(__lowercase  ,				strict=__lowercase)

  __lowerCAmelCase								=					DDIMScheduler(
      num_train_timesteps=10_00  ,				beta_schedule='''linear'''  ,				prediction_type='''epsilon'''  ,				clip_sample=__lowercase  ,				)

  __lowerCAmelCase								=					AutoencoderKL.from_pretrained(args.vae_model)

  __lowerCAmelCase								=					DiTPipeline(transformer=__lowercase  ,				vae=__lowercase  ,				scheduler=__lowercase)

  if args.save:
    pipeline.save_pretrained(args.checkpoint_path)


if __name__ == "__main__":
      __lowercase				    =      argparse.ArgumentParser()

      parser.add_argument(
          """--image_size""",
          default=256,
          type=int,
          required=False,
          help="""Image size of pretrained model, either 256 or 512.""",
      )
      parser.add_argument(
          """--vae_model""",
          default="""stabilityai/sd-vae-ft-ema""",
          type=str,
          required=False,
          help="""Path to pretrained VAE model, either stabilityai/sd-vae-ft-mse or stabilityai/sd-vae-ft-ema.""",
      )
      parser.add_argument(
          """--save""", default=True, type=bool, required=False, help="""Whether to save the converted pipeline or not."""
      )
      parser.add_argument(
          """--checkpoint_path""", default=None, type=str, required=True, help="""Path to the output pipeline."""
      )

      __lowercase				    =      parser.parse_args()
      main(args)
