
import os
from glob import glob

import imageio
import torch
import torchvision
import wandb
from img_processing import custom_to_pil, loop_post_process, preprocess, preprocess_vqgan
from loaders import load_vqgan
from PIL import Image
from torch import nn

from transformers import CLIPModel, CLIPTokenizerFast
from utils import get_device, get_timestamp, show_pil






class   UpperCAmelCase		:



   def __init__(self		:					List[str]    ,						snake_case__		:					str = "cpu"    ,						snake_case__		:					str = "openai/clip-vit-large-patch14" )		->	None:


          '''simple docstring'''


          snake_case       :   int							=      device
          snake_case       :   str							=      CLIPTokenizerFast.from_pretrained(snake_case__ )
          snake_case       :   List[Any]							=      [0.48145466, 0.4578275, 0.40821073]
          snake_case       :   Optional[int]							=      [0.26862954, 0.26130258, 0.27577711]
          snake_case       :   List[Any]							=      torchvision.transforms.Normalize(self.image_mean    ,						self.image_std )
          snake_case       :   List[str]							=      torchvision.transforms.Resize(2_24 )
          snake_case       :   List[Any]							=      torchvision.transforms.CenterCrop(2_24 )



   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[int]    ,						snake_case__		:					Union[str, Any] )		->	Optional[Any]:


          '''simple docstring'''


          snake_case       :   List[str]							=      self.resize(snake_case__ )
          snake_case       :   Union[str, Any]							=      self.center_crop(snake_case__ )
          snake_case       :   str							=      self.normalize(snake_case__ )
          return images



   def __call__(self		:					Optional[int]    ,						snake_case__		:					Optional[int]=None    ,						snake_case__		:					List[Any]=None    ,						**snake_case__		:					int )		->	List[Any]:


          '''simple docstring'''


          snake_case       :   Dict							=      self.tokenizer(text=snake_case__    ,						**snake_case__ )
          snake_case       :   str							=      self.preprocess_img(snake_case__ )
          snake_case       :   Any							=      {key: value.to(self.device ) for (key, value) in encoding.items()}
          return encoding






class   UpperCAmelCase		(						nn.Module ):



   def __init__(self		:					int    ,						snake_case__		:					Tuple=10    ,						snake_case__		:					List[str]=0.01    ,						snake_case__		:					Dict=None    ,						snake_case__		:					List[str]=None    ,						snake_case__		:					Dict=None    ,						snake_case__		:					Optional[Any]=None    ,						snake_case__		:					Union[str, Any]=None    ,						snake_case__		:					Dict=None    ,						snake_case__		:					Any=False    ,						snake_case__		:					Optional[Any]=True    ,						snake_case__		:					Any="image"    ,						snake_case__		:					str=True    ,						snake_case__		:					Any=False    ,						snake_case__		:					List[str]=False    ,						snake_case__		:					Optional[Any]=False    ,						)		->	None:


          '''simple docstring'''


          super().__init__()
          snake_case       :   List[str]							=      None
          snake_case       :   Optional[int]							=      device if device else get_device()
          if vqgan:
                 snake_case       :   List[Any]							=      vqgan
          else:
                 snake_case       :   Tuple							=      load_vqgan(self.device    ,						conf_path=snake_case__    ,						ckpt_path=snake_case__ )
          self.vqgan.eval()
          if clip:
                 snake_case       :   Any							=      clip
          else:
                 snake_case       :   Optional[Any]							=      CLIPModel.from_pretrained("openai/clip-vit-base-patch32" )
          self.clip.to(self.device )
          snake_case       :   Optional[Any]							=      ProcessorGradientFlow(device=self.device )

          snake_case       :   Optional[int]							=      iterations
          snake_case       :   str							=      lr
          snake_case       :   List[Any]							=      log
          snake_case       :   Optional[int]							=      make_grid
          snake_case       :   str							=      return_val
          snake_case       :   List[Any]							=      quantize
          snake_case       :   List[str]							=      self.vqgan.decoder.z_shape



   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[int]    ,						snake_case__		:					Dict=None    ,						snake_case__		:					Tuple=None    ,						snake_case__		:					Any=5    ,						snake_case__		:					Optional[Any]=True )		->	Any:


          '''simple docstring'''


          snake_case       :   List[Any]							=      []
          if output_path is None:
                 snake_case       :   List[Any]							=      "./animation.gif"
          if input_path is None:
                 snake_case       :   Any							=      self.save_path
          snake_case       :   int							=      sorted(glob(input_path + "/*" ) )
          if not len(snake_case__ ):
                 raise ValueError(
                     "No images found in save path, aborting (did you pass save_intermediate=True to the generate"
                     " function?)" )
          if len(snake_case__ ) == 1:
                 print("Only one image found in save path, (did you pass save_intermediate=True to the generate function?)" )
          snake_case       :   List[Any]							=      total_duration / len(snake_case__ )
          snake_case       :   Union[str, Any]							=      [frame_duration] * len(snake_case__ )
          if extend_frames:
                 snake_case       :   Union[str, Any]							=      1.5
                 snake_case       :   int							=      3
          for file_name in paths:
                 if file_name.endswith(".png" ):
                        images.append(imageio.imread(snake_case__ ) )
          imageio.mimsave(snake_case__    ,						snake_case__    ,						duration=snake_case__ )
          print(f"""gif saved to {output_path}""" )



   def 				_SCREAMING_SNAKE_CASE   (self		:					str    ,						snake_case__		:					str=None    ,						snake_case__		:					List[Any]=None )		->	Union[str, Any]:


          '''simple docstring'''


          if not (path or img):
                 raise ValueError("Input either path or tensor" )
          if img is not None:
                 raise NotImplementedError
          snake_case       :   Optional[Any]							=      preprocess(Image.open(snake_case__ )    ,						target_image_size=2_56 ).to(self.device )
          snake_case       :   Any							=      preprocess_vqgan(snake_case__ )
          snake_case				,					*snake_case       :   str							=      self.vqgan.encode(snake_case__ )
          return z



   def 				_SCREAMING_SNAKE_CASE   (self		:					List[str]    ,						snake_case__		:					Tuple )		->	Optional[Any]:


          '''simple docstring'''


          snake_case       :   Dict							=      self.latent.detach().requires_grad_()
          snake_case       :   Optional[int]							=      base_latent + transform_vector
          if self.quantize:
                 snake_case				,					*snake_case       :   Any							=      self.vqgan.quantize(snake_case__ )
          else:
                 snake_case       :   Union[str, Any]							=      trans_latent
          return self.vqgan.decode(snake_case__ )



   def 				_SCREAMING_SNAKE_CASE   (self		:					Dict    ,						snake_case__		:					str    ,						snake_case__		:					Optional[int]    ,						snake_case__		:					Optional[int]=None )		->	str:


          '''simple docstring'''


          snake_case       :   Any							=      self.clip_preprocessor(text=snake_case__    ,						images=snake_case__    ,						return_tensors="pt"    ,						padding=snake_case__ )
          snake_case       :   Optional[int]							=      self.clip(**snake_case__ )
          snake_case       :   List[Any]							=      clip_outputs.logits_per_image
          if weights is not None:
                 snake_case       :   Any							=      similarity_logits * weights
          return similarity_logits.sum()



   def 				_SCREAMING_SNAKE_CASE   (self		:					Union[str, Any]    ,						snake_case__		:					Optional[Any]    ,						snake_case__		:					Optional[Any]    ,						snake_case__		:					Dict )		->	Optional[Any]:


          '''simple docstring'''


          snake_case       :   Any							=      self._get_clip_similarity(pos_prompts["prompts"]    ,						snake_case__    ,						weights=(1 / pos_prompts["weights"]) )
          if neg_prompts:
                 snake_case       :   Union[str, Any]							=      self._get_clip_similarity(neg_prompts["prompts"]    ,						snake_case__    ,						weights=neg_prompts["weights"] )
          else:
                 snake_case       :   Union[str, Any]							=      torch.tensor([1]    ,						device=self.device )
          snake_case       :   List[Any]							=      -torch.log(snake_case__ ) + torch.log(snake_case__ )
          return loss



   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[Any]    ,						snake_case__		:					Union[str, Any]    ,						snake_case__		:					str    ,						snake_case__		:					Optional[Any] )		->	str:


          '''simple docstring'''


          snake_case       :   Union[str, Any]							=      torch.randn_like(self.latent    ,						requires_grad=snake_case__    ,						device=self.device )
          snake_case       :   Optional[int]							=      torch.optim.Adam([vector]    ,						lr=self.lr )

          for i in range(self.iterations ):
                 optim.zero_grad()
                 snake_case       :   Dict							=      self._add_vector(snake_case__ )
                 snake_case       :   Union[str, Any]							=      loop_post_process(snake_case__ )
                 snake_case       :   List[str]							=      self._get_CLIP_loss(snake_case__    ,						snake_case__    ,						snake_case__ )
                 print("CLIP loss"    ,						snake_case__ )
                 if self.log:
                        wandb.log({"CLIP Loss": clip_loss} )
                 clip_loss.backward(retain_graph=snake_case__ )
                 optim.step()
                 if self.return_val == "image":
                        yield custom_to_pil(transformed_img[0] )
                 else:
                        yield vector



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


          '''simple docstring'''


          wandb.init(reinit=snake_case__    ,						project="face-editor" )
          wandb.config.update({"Positive Prompts": positive_prompts} )
          wandb.config.update({"Negative Prompts": negative_prompts} )
          wandb.config.update({"lr": self.lr, "iterations": self.iterations} )
          if image_path:
                 snake_case       :   Optional[Any]							=      Image.open(snake_case__ )
                 snake_case       :   Dict							=      image.resize((2_56, 2_56) )
                 wandb.log("Original Image"    ,						wandb.Image(snake_case__ ) )



   def 				_SCREAMING_SNAKE_CASE   (self		:					Dict    ,						snake_case__		:					Union[str, Any] )		->	Optional[int]:


          '''simple docstring'''


          if not prompts:
                 return []
          snake_case       :   Tuple							=      []
          snake_case       :   List[str]							=      []
          if isinstance(snake_case__    ,						snake_case__ ):
                 snake_case       :   Union[str, Any]							=      [prompt.strip() for prompt in prompts.split("|" )]
          for prompt in prompts:
                 if isinstance(snake_case__    ,						(tuple, list) ):
                        snake_case       :   Tuple							=      prompt[0]
                        snake_case       :   Optional[int]							=      float(prompt[1] )
                 elif ":" in prompt:
                        snake_case				,					snake_case       :   Optional[int]							=      prompt.split(":" )
                        snake_case       :   List[str]							=      float(snake_case__ )
                 else:
                        snake_case       :   Optional[Any]							=      prompt
                        snake_case       :   Tuple							=      1.0
                 processed_prompts.append(snake_case__ )
                 weights.append(snake_case__ )
          return {
              "prompts": processed_prompts,
              "weights": torch.tensor(snake_case__    ,						device=self.device ),
          }



   def 				_SCREAMING_SNAKE_CASE   (self		:					Dict    ,						snake_case__		:					Dict    ,						snake_case__		:					Union[str, Any]=None    ,						snake_case__		:					str=None    ,						snake_case__		:					Union[str, Any]=True    ,						snake_case__		:					Any=False    ,						snake_case__		:					List[Any]=True    ,						snake_case__		:					List[Any]=True    ,						snake_case__		:					Dict=None    ,						)		->	Any:


          '''simple docstring'''


          if image_path:
                 snake_case       :   List[Any]							=      self._get_latent(snake_case__ )
          else:
                 snake_case       :   Optional[int]							=      torch.randn(self.latent_dim    ,						device=self.device )
          if self.log:
                 self._init_logging(snake_case__    ,						snake_case__    ,						snake_case__ )

          assert pos_prompts, "You must provide at least one positive prompt."
          snake_case       :   Optional[Any]							=      self.process_prompts(snake_case__ )
          snake_case       :   List[Any]							=      self.process_prompts(snake_case__ )
          if save_final and save_path is None:
                 snake_case       :   Tuple							=      os.path.join("./outputs/"    ,						"_".join(pos_prompts["prompts"] ) )
          if not os.path.exists(snake_case__ ):
                 os.makedirs(snake_case__ )
          else:
                 snake_case       :   int							=      save_path + "_" + get_timestamp()
                 os.makedirs(snake_case__ )
          snake_case       :   Tuple							=      save_path

          snake_case       :   str							=      self.vqgan.decode(self.latent )[0]
          if show_intermediate:
                 print("Original Image" )
                 show_pil(custom_to_pil(snake_case__ ) )

          snake_case       :   List[str]							=      loop_post_process(snake_case__ )
          for iter, transformed_img in enumerate(self._optimize_CLIP(snake_case__    ,						snake_case__    ,						snake_case__ ) ):
                 if show_intermediate:
                        show_pil(snake_case__ )
                 if save_intermediate:
                        transformed_img.save(os.path.join(self.save_path    ,						f"""iter_{iter:03d}.png""" ) )
                 if self.log:
                        wandb.log({"Image": wandb.Image(snake_case__ )} )
          if show_final:
                 show_pil(snake_case__ )
          if save_final:
                 transformed_img.save(os.path.join(self.save_path    ,						f"""iter_{iter:03d}_final.png""" ) )











"""simple docstring"""




import argparse
import collections

import numpy as np
import torch
from flax import traverse_util
from tax import checkpoints

from transformers import MTaConfig, UMTaEncoderModel, UMTaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()



def     a__ (				snake_case__       , snake_case__       , snake_case__				)						->					List[Any]:
			return params[F'{prefix}/{prefix}/relpos_bias/rel_embedding'][:, i, :]



def     a__ (				snake_case__       , snake_case__       , snake_case__       , snake_case__="attention"				)						->					List[Any]:
			lowerCamelCase			      =					lowerCamelCase			      =					np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/key/kernel'][:, i, :, :]				)
			lowerCamelCase			      =					k_tmp.reshape(k_tmp.shape[0]       , k_tmp.shape[1] * k_tmp.shape[2]				)
			lowerCamelCase			      =					np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/out/kernel'][:, i, :, :]				)
			lowerCamelCase			      =					o_tmp.reshape(o_tmp.shape[0] * o_tmp.shape[1]       , o_tmp.shape[2]				)
			lowerCamelCase			      =					np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/query/kernel'][:, i, :, :]				)
			lowerCamelCase			      =					q_tmp.reshape(q_tmp.shape[0]       , q_tmp.shape[1] * q_tmp.shape[2]				)
			lowerCamelCase			      =					np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/value/kernel'][:, i, :, :]				)
			lowerCamelCase			      =					v_tmp.reshape(v_tmp.shape[0]       , v_tmp.shape[1] * v_tmp.shape[2]				)
			return k, o, q, v



def     a__ (				snake_case__       , snake_case__       , snake_case__       , snake_case__=False				)						->					List[str]:
			if split_mlp_wi:
						lowerCamelCase			      =					params[F'{prefix}/{prefix}/mlp/wi_0/kernel'][:, i, :]
						lowerCamelCase			      =					params[F'{prefix}/{prefix}/mlp/wi_1/kernel'][:, i, :]
						lowerCamelCase			      =					(wi_a, wi_a)
			else:
						lowerCamelCase			      =					params[F'{prefix}/{prefix}/mlp/wi/kernel'][:, i, :]

			lowerCamelCase			      =					params[F'{prefix}/{prefix}/mlp/wo/kernel'][:, i, :]
			return wi, wo



def     a__ (				snake_case__       , snake_case__       , snake_case__       , snake_case__				)						->					Tuple:
			return params[F'{prefix}/{prefix}/{layer_name}/scale'][:, i]



def     a__ (				snake_case__       , *, snake_case__       , snake_case__       , snake_case__ = False				)						->					Dict:
			lowerCamelCase			      =					traverse_util.flatten_dict(variables["""target"""]				)
			lowerCamelCase			      =					{"""/""".join(snake_case__				): v for k, v in old.items()}

			# v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
			lowerCamelCase			      =					"""encoder/encoder/mlp/wi_0/kernel""" in old
			print("""Split MLP:"""       , snake_case__				)

			lowerCamelCase			      =					collections.OrderedDict()

			# Shared embeddings.
			lowerCamelCase			      =					old["""token_embedder/embedding"""]

			# Encoder.
			for i in range(snake_case__				):
						# Block i, layer 0 (Self Attention).
						lowerCamelCase			      =					tax_layer_norm_lookup(snake_case__       , snake_case__       , """encoder"""       , """pre_attention_layer_norm"""				)
						lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase			      =					tax_attention_lookup(snake_case__       , snake_case__       , """encoder"""       , """attention"""				)
						lowerCamelCase			      =					layer_norm
						lowerCamelCase			      =					k.T
						lowerCamelCase			      =					o.T
						lowerCamelCase			      =					q.T
						lowerCamelCase			      =					v.T

						# Block i, layer 1 (MLP).
						lowerCamelCase			      =					tax_layer_norm_lookup(snake_case__       , snake_case__       , """encoder"""       , """pre_mlp_layer_norm"""				)
						lowerCamelCase      ,   lowerCamelCase			      =					tax_mlp_lookup(snake_case__       , snake_case__       , """encoder"""       , snake_case__				)
						lowerCamelCase			      =					layer_norm
						if split_mlp_wi:
									lowerCamelCase			      =					wi[0].T
									lowerCamelCase			      =					wi[1].T
						else:
									lowerCamelCase			      =					wi.T
						lowerCamelCase			      =					wo.T
						if scalable_attention:
									# convert the rel_embedding of each layer
									lowerCamelCase			      =					tax_relpos_bias_lookup(
									    snake_case__       , snake_case__       , """encoder"""				).T

			lowerCamelCase			      =					old["""encoder/encoder_norm/scale"""]

			if not scalable_attention:
						lowerCamelCase			      =					tax_relpos_bias_lookup(
						    snake_case__       , 0       , """encoder"""				).T
						lowerCamelCase			      =					tax_relpos_bias_lookup(
						    snake_case__       , 0       , """decoder"""				).T

			if not is_encoder_only:
						# Decoder.
						for i in range(snake_case__				):
									# Block i, layer 0 (Self Attention).
									lowerCamelCase			      =					tax_layer_norm_lookup(snake_case__       , snake_case__       , """decoder"""       , """pre_self_attention_layer_norm"""				)
									lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase			      =					tax_attention_lookup(snake_case__       , snake_case__       , """decoder"""       , """self_attention"""				)
									lowerCamelCase			      =					layer_norm
									lowerCamelCase			      =					k.T
									lowerCamelCase			      =					o.T
									lowerCamelCase			      =					q.T
									lowerCamelCase			      =					v.T

									# Block i, layer 1 (Cross Attention).
									lowerCamelCase			      =					tax_layer_norm_lookup(snake_case__       , snake_case__       , """decoder"""       , """pre_cross_attention_layer_norm"""				)
									lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase			      =					tax_attention_lookup(snake_case__       , snake_case__       , """decoder"""       , """encoder_decoder_attention"""				)
									lowerCamelCase			      =					layer_norm
									lowerCamelCase			      =					k.T
									lowerCamelCase			      =					o.T
									lowerCamelCase			      =					q.T
									lowerCamelCase			      =					v.T

									# Block i, layer 2 (MLP).
									lowerCamelCase			      =					tax_layer_norm_lookup(snake_case__       , snake_case__       , """decoder"""       , """pre_mlp_layer_norm"""				)
									lowerCamelCase      ,   lowerCamelCase			      =					tax_mlp_lookup(snake_case__       , snake_case__       , """decoder"""       , snake_case__				)
									lowerCamelCase			      =					layer_norm
									if split_mlp_wi:
												lowerCamelCase			      =					wi[0].T
												lowerCamelCase			      =					wi[1].T
									else:
												lowerCamelCase			      =					wi.T
									lowerCamelCase			      =					wo.T

									if scalable_attention:
												# convert the rel_embedding of each layer
												lowerCamelCase			      =					tax_relpos_bias_lookup(snake_case__       , snake_case__       , """decoder"""				).T

						lowerCamelCase			      =					old["""decoder/decoder_norm/scale"""]

						# LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
						if "decoder/logits_dense/kernel" in old:
									lowerCamelCase			      =					old["""decoder/logits_dense/kernel"""].T

			return new



def     a__ (				snake_case__       , snake_case__				)						->					Optional[int]:
			lowerCamelCase			      =					collections.OrderedDict([(k, torch.from_numpy(v.copy()				)) for (k, v) in converted_params.items()]				)

			# Add what is missing.
			if "encoder.embed_tokens.weight" not in state_dict:
						lowerCamelCase			      =					state_dict["""shared.weight"""]

			if not is_encoder_only:
						if "decoder.embed_tokens.weight" not in state_dict:
									lowerCamelCase			      =					state_dict["""shared.weight"""]

						if "lm_head.weight" not in state_dict:  # For old 1.0 models.
									print("""Using shared word embeddings as lm_head."""				)
									lowerCamelCase			      =					state_dict["""shared.weight"""]

			return state_dict



def     a__ (				snake_case__       , snake_case__       , snake_case__       , snake_case__       , snake_case__				)						->					List[Any]:
			lowerCamelCase			      =					checkpoints.load_tax_checkpoint(snake_case__				)
			lowerCamelCase			      =					convert_tax_to_pytorch(
			    snake_case__       , num_layers=config.num_layers       , is_encoder_only=snake_case__       , scalable_attention=snake_case__				)
			lowerCamelCase			      =					make_state_dict(snake_case__       , snake_case__				)
			model.load_state_dict(snake_case__       , strict=snake_case__				)



def     a__ (				snake_case__       , snake_case__       , snake_case__       , snake_case__ = False       , snake_case__ = False       , )						->					str:
			lowerCamelCase			      =					MTaConfig.from_json_file(snake_case__				)
			print(F'Building PyTorch model from configuration: {config}'				)
			# Non-v1.1 checkpoints could also use T5Model, but this works for all.
			# The v1.0 checkpoints will simply have an LM head that is the word embeddings.
			if is_encoder_only:
						lowerCamelCase			      =					UMTaEncoderModel(snake_case__				)
			else:
						lowerCamelCase			      =					UMTaForConditionalGeneration(snake_case__				)

			# Load weights from tf checkpoint
			load_tax_weights_in_ta(snake_case__       , snake_case__       , snake_case__       , snake_case__       , snake_case__				)

			# Save pytorch-model
			print(F'Save PyTorch model to {pytorch_dump_path}'				)
			model.save_pretrained(snake_case__				)

			# Verify that we can load the checkpoint.
			model.from_pretrained(snake_case__				)
			print("""Done"""				)


if __name__ == "__main__":
					lowerCAmelCase							:			Optional[int]     							=							argparse.ArgumentParser(description="""Converts a native T5X checkpoint into a PyTorch checkpoint.""")
					# Required parameters
					parser.add_argument(
					    """--t5x_checkpoint_path""", default=None, type=str, required=True, help="""Path to the T5X checkpoint."""
					)
					parser.add_argument(
					    """--config_file""",
					    default=None,
					    type=str,
					    required=True,
					    help="""The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.""",
					)
					parser.add_argument(
					    """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
					)
					parser.add_argument(
					    """--is_encoder_only""", action="""store_true""", help="""Check if the model is encoder-decoder model""", default=False
					)
					parser.add_argument(
					    """--scalable_attention""",
					    action="""store_true""",
					    help="""Whether the model uses scaled attention (umt5 model)""",
					    default=False,
					)
					lowerCAmelCase							:			int     							=							parser.parse_args()
					convert_tax_checkpoint_to_pytorch(
					    args.tax_checkpoint_path,
					    args.config_file,
					    args.pytorch_dump_path,
					    args.is_encoder_only,
					    args.scalable_attention,
					)



