






import tempfile
import unittest

from transformers import TaConfig, is_torch_available
from transformers.testing_utils import (
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
    torch_device,
)

from ...generation.test_utils import GenerationTesterMixin
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import AutoTokenizer, UMTaForConditionalGeneration, UMTaForQuestionAnswering, UMTaModel
class 						lowerCamelCase_				:







				def __init__(						self       ,					__lowerCAmelCase       ,					__lowerCAmelCase=9_9       ,					__lowerCAmelCase=1_3       ,					__lowerCAmelCase=7       ,					__lowerCAmelCase=9       ,					__lowerCAmelCase=True       ,					__lowerCAmelCase=True       ,					__lowerCAmelCase=False       ,					__lowerCAmelCase=3_2       ,					__lowerCAmelCase=5       ,					__lowerCAmelCase=4       ,					__lowerCAmelCase=3_7       ,					__lowerCAmelCase=8       ,					__lowerCAmelCase=0.1       ,					__lowerCAmelCase=0.002       ,					__lowerCAmelCase=1       ,					__lowerCAmelCase=0       ,					__lowerCAmelCase=0       ,					__lowerCAmelCase=None       ,					__lowerCAmelCase=None       ,					):
								"""simple docstring"""
								__magic_name__		:Optional[int]        =      parent
								__magic_name__		:int        =      batch_size
								__magic_name__		:str        =      encoder_seq_length
								__magic_name__		:str        =      decoder_seq_length
								# For common tests
								__magic_name__		:str        =      self.decoder_seq_length
								__magic_name__		:Optional[int]        =      is_training
								__magic_name__		:str        =      use_attention_mask
								__magic_name__		:Optional[int]        =      use_labels
								__magic_name__		:Tuple        =      vocab_size
								__magic_name__		:Any        =      hidden_size
								__magic_name__		:List[str]        =      num_hidden_layers
								__magic_name__		:Dict        =      num_attention_heads
								__magic_name__		:Union[str, Any]        =      d_ff
								__magic_name__		:Optional[int]        =      relative_attention_num_buckets
								__magic_name__		:Dict        =      dropout_rate
								__magic_name__		:Tuple        =      initializer_factor
								__magic_name__		:Dict        =      eos_token_id
								__magic_name__		:int        =      pad_token_id
								__magic_name__		:int        =      decoder_start_token_id
								__magic_name__		:Dict        =      None
								__magic_name__		:Dict        =      decoder_layers







				def 			A	(						self     ):
								"""simple docstring"""
								return TaConfig.from_pretrained('''google/umt5-base'''     )







				def 			A	(						self       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase=None       ,					__lowerCAmelCase=None       ,					__lowerCAmelCase=None       ,					__lowerCAmelCase=None       ,					__lowerCAmelCase=None       ,					):
								"""simple docstring"""






								if attention_mask is None:
												__magic_name__		:Tuple        =      input_ids.ne(config.pad_token_id     )
								if decoder_attention_mask is None:
												__magic_name__		:Union[str, Any]        =      decoder_input_ids.ne(config.pad_token_id     )
								if head_mask is None:
												__magic_name__		:Dict        =      torch.ones(config.num_hidden_layers       ,					config.num_attention_heads       ,					device=lowercase__     )
								if decoder_head_mask is None:
												__magic_name__		:List[str]        =      torch.ones(config.num_decoder_layers       ,					config.num_attention_heads       ,					device=lowercase__     )
								if cross_attn_head_mask is None:
												__magic_name__		:Optional[int]        =      torch.ones(
												    config.num_decoder_layers       ,					config.num_attention_heads       ,					device=lowercase__     )
								return {
								    "input_ids": input_ids,
								    "decoder_input_ids": decoder_input_ids,
								    "attention_mask": attention_mask,
								    "decoder_attention_mask": decoder_attention_mask,
								    "head_mask": head_mask,
								    "decoder_head_mask": decoder_head_mask,
								    "cross_attn_head_mask": cross_attn_head_mask,
								}







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:str        =      ids_tensor([self.batch_size, self.encoder_seq_length]       ,					self.vocab_size     )
								__magic_name__		:Dict        =      ids_tensor([self.batch_size, self.decoder_seq_length]       ,					self.vocab_size     )

								# we need to clamp the input ids here to avoid having pad token in between
								# this is because for NllbMoe the position_ids are prepared such that
								# all pad tokens have pos id = 2 and rest are between 2..seq_length
								# and the seq_length here is seq_length - num_pad_tokens
								# but when using past, there is no way of knowing if the past input ids had
								# pad tokens in them, which results in incorrect seq_lenth and which in turn results in
								# position_ids being off by num_pad_tokens in past input
								__magic_name__		:Optional[Any]        =      input_ids.clamp(self.pad_token_id + 1     )
								__magic_name__		:List[Any]        =      decoder_input_ids.clamp(self.pad_token_id + 1     )

								__magic_name__		:Any        =      self.get_config()
								__magic_name__		:int        =      config.num_attention_heads
								__magic_name__		:int        =      self.prepare_inputs_dict(lowercase__       ,					lowercase__       ,					lowercase__     )
								return config, input_dict







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:str        =      self.prepare_config_and_inputs()
								return config, inputs_dict







				def 			A	(						self     ):
								"""simple docstring"""
								return TaConfig(
								    vocab_size=1_6_6       ,					d_model=self.hidden_size       ,					d_ff=self.d_ff       ,					d_kv=self.hidden_size // self.num_attention_heads       ,					num_layers=self.num_hidden_layers       ,					num_decoder_layers=self.decoder_layers       ,					num_heads=self.num_attention_heads       ,					relative_attention_num_buckets=self.relative_attention_num_buckets       ,					dropout_rate=self.dropout_rate       ,					initializer_factor=self.initializer_factor       ,					eos_token_id=self.eos_token_id       ,					bos_token_id=self.pad_token_id       ,					pad_token_id=self.pad_token_id       ,					decoder_start_token_id=self.decoder_start_token_id       ,					)







				def 			A	(						self     ):
								"""simple docstring"""
								return TaConfig(
								    vocab_size=self.vocab_size       ,					d_model=self.hidden_size       ,					d_ff=self.d_ff       ,					d_kv=self.hidden_size // self.num_attention_heads       ,					num_layers=self.num_hidden_layers       ,					num_decoder_layers=self.decoder_layers       ,					num_heads=self.num_attention_heads       ,					relative_attention_num_buckets=self.relative_attention_num_buckets       ,					dropout_rate=self.dropout_rate       ,					initializer_factor=self.initializer_factor       ,					eos_token_id=self.eos_token_id       ,					bos_token_id=self.pad_token_id       ,					pad_token_id=self.pad_token_id       ,					decoder_start_token_id=self.decoder_start_token_id       ,					)







				def 			A	(						self       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					):
								"""simple docstring"""
								__magic_name__		:List[Any]        =      UMTaModel(config=lowercase__     )
								model.to(lowercase__     )
								model.eval()
								__magic_name__		:Optional[Any]        =      model(
								    input_ids=lowercase__       ,					decoder_input_ids=lowercase__       ,					attention_mask=lowercase__       ,					decoder_attention_mask=lowercase__       ,					)
								__magic_name__		:Optional[int]        =      model(input_ids=lowercase__       ,					decoder_input_ids=lowercase__     )
								__magic_name__		:int        =      result.last_hidden_state
								__magic_name__		:Any        =      result.past_key_values
								__magic_name__		:int        =      result.encoder_last_hidden_state

								self.parent.assertEqual(encoder_output.size()       ,					(self.batch_size, self.encoder_seq_length, self.hidden_size)     )
								self.parent.assertEqual(decoder_output.size()       ,					(self.batch_size, self.decoder_seq_length, self.hidden_size)     )
								# There should be `num_layers` key value embeddings stored in decoder_past
								self.parent.assertEqual(len(lowercase__     )       ,					config.num_layers     )
								# There should be a self attn key, a self attn value, a cross attn key and a cross attn value stored in each decoder_past tuple
								self.parent.assertEqual(len(decoder_past[0]     )       ,					4     )







				def 			A	(						self       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					):
								"""simple docstring"""
								__magic_name__		:List[str]        =      UMTaModel(config=lowercase__     ).get_decoder().to(lowercase__     ).eval()
								# first forward pass
								__magic_name__		:Dict        =      model(lowercase__       ,					use_cache=lowercase__     )
								__magic_name__		:Optional[int]        =      model(lowercase__     )
								__magic_name__		:Dict        =      model(lowercase__       ,					use_cache=lowercase__     )

								self.parent.assertTrue(len(lowercase__     ) == len(lowercase__     )     )
								self.parent.assertTrue(len(lowercase__     ) == len(lowercase__     ) + 1     )

								__magic_name__		:Optional[Any]        =      outputs.to_tuple()

								# create hypothetical next token and extent to next_input_ids
								__magic_name__		:Dict        =      ids_tensor((self.batch_size, 1)       ,					config.vocab_size     )

								# append to next input_ids and
								__magic_name__		:int        =      torch.cat([input_ids, next_tokens]       ,					dim=-1     )

								__magic_name__		:Optional[int]        =      model(lowercase__     )["last_hidden_state"]
								__magic_name__		:Tuple        =      model(lowercase__       ,					past_key_values=lowercase__     )["last_hidden_state"]

								# select random slice
								__magic_name__		:int        =      ids_tensor((1,)       ,					output_from_past.shape[-1]     ).item()
								__magic_name__		:Dict        =      output_from_no_past[:, -1, random_slice_idx].detach()
								__magic_name__		:Tuple        =      output_from_past[:, 0, random_slice_idx].detach()

								# test that outputs are equal for slice
								self.parent.assertTrue(torch.allclose(lowercase__       ,					lowercase__       ,					atol=1E-3     )     )







				def 			A	(						self       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					):
								"""simple docstring"""
								__magic_name__		:Dict        =      UMTaModel(config=lowercase__     ).to(lowercase__     ).half().eval()
								__magic_name__		:Union[str, Any]        =      model(**lowercase__     )["last_hidden_state"]
								self.parent.assertFalse(torch.isnan(lowercase__     ).any().item()     )







@require_torch
class 						lowerCamelCase_				(      _UpperCAmelCase			,				_UpperCAmelCase			,				_UpperCAmelCase			,				unittest.TestCase ):
				a__					       =  (
				    (UMTaModel, UMTaForConditionalGeneration, UMTaForQuestionAnswering) if is_torch_available() else ()
				)
				a__					       =  (UMTaForConditionalGeneration,) if is_torch_available() else ()
				a__					       =  (
				    {
				        '''conversational''': UMTaForConditionalGeneration,
				        '''feature-extraction''': UMTaModel,
				        '''summarization''': UMTaForConditionalGeneration,
				        '''text2text-generation''': UMTaForConditionalGeneration,
				        '''translation''': UMTaForConditionalGeneration,
				        '''question-answering''': UMTaForQuestionAnswering,
				    }
				    if is_torch_available()
				    else {}
				)
				a__					       =  True
				a__					       =  False
				a__					       =  False
				a__					       =  True
				a__					       =  True
				# The small UMT5 model needs higher percentages for CPU/MP tests
				a__					       =  [0.8, 0.9]







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:Optional[int]        =      UMTaModelTester(self     )







				@unittest.skip('''Test has a segmentation fault on torch 1.8.0'''     )
				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:Optional[Any]        =      self.model_tester.prepare_config_and_inputs()
								__magic_name__		:Tuple        =      UMTaModel(config_and_inputs[0]     ).to(lowercase__     )
								with tempfile.TemporaryDirectory() as tmpdirname:
												torch.onnx.export(
												    lowercase__       ,					(config_and_inputs[1], config_and_inputs[3], config_and_inputs[2])       ,					F'''{tmpdirname}/t5_test.onnx'''       ,					export_params=lowercase__       ,					opset_version=9       ,					input_names=['''input_ids''', '''decoder_input_ids''']       ,					)







				@unittest.skipIf(torch_device == '''cpu'''       ,					'''Cant do half precision'''     )
				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:int        =      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model_fpaa_forward(*lowercase__     )







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:Optional[int]        =      ["encoder_attentions", "decoder_attentions", "cross_attentions"]
								__magic_name__		:Dict        =      self.model_tester.prepare_config_and_inputs()
								__magic_name__		:Tuple        =      config_and_inputs[0]
								__magic_name__		:str        =      UMTaForConditionalGeneration(lowercase__     ).eval()
								model.to(lowercase__     )

								__magic_name__		:Dict        =      {
								    "head_mask": torch.zeros(config.num_layers       ,					config.num_heads       ,					device=lowercase__     ),
								    "decoder_head_mask": torch.zeros(config.num_decoder_layers       ,					config.num_heads       ,					device=lowercase__     ),
								    "cross_attn_head_mask": torch.zeros(config.num_decoder_layers       ,					config.num_heads       ,					device=lowercase__     ),
								}

								for attn_name, (name, mask) in zip(lowercase__       ,					head_masking.items()     ):
												__magic_name__		:Tuple        =      {name: mask}
												# Explicitly pass decoder_head_mask as it is required from T5 model when head_mask specified
												if name == "head_mask":
																__magic_name__		:str        =      torch.ones(
																    config.num_decoder_layers       ,					config.num_heads       ,					device=lowercase__     )

												__magic_name__		:Any        =      model.generate(
												    config_and_inputs[1]['''input_ids''']       ,					num_beams=1       ,					max_length=3       ,					output_attentions=lowercase__       ,					return_dict_in_generate=lowercase__       ,					**lowercase__       ,					)
												# We check the state of decoder_attentions and cross_attentions just from the last step
												__magic_name__		:Optional[int]        =      out[attn_name] if attn_name == attention_names[0] else out[attn_name][-1]
												self.assertEqual(sum([w.sum().item() for w in attn_weights]     )       ,					0.0     )







				@unittest.skip('''Does not work on the tiny model as we keep hitting edge cases.'''     )
				def 			A	(						self     ):
								"""simple docstring"""
								pass







@require_torch
@require_sentencepiece
@require_tokenizers
class 						lowerCamelCase_				(      unittest.TestCase ):







				@slow
				@unittest.skip(
				    '''Unless we stop stripping left and right by default for all special tokens, the expected ids obtained here will not match the original ones. Wait for https://github.com/huggingface/transformers/pull/23909 to be merged'''     )
				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:Tuple        =      UMTaForConditionalGeneration.from_pretrained('''google/umt5-small'''       ,					return_dict=lowercase__     ).to(lowercase__     )
								__magic_name__		:List[Any]        =      AutoTokenizer.from_pretrained('''google/umt5-small'''       ,					use_fast=lowercase__       ,					legacy=lowercase__     )
								__magic_name__		:Optional[Any]        =      [
								    "Bonjour monsieur <extra_id_0> bien <extra_id_1>.",
								    "No se como puedo <extra_id_0>.",
								    "This is the reason why we <extra_id_0> them.",
								    "The <extra_id_0> walks in <extra_id_1>, seats",
								    "A <extra_id_0> walks into a bar and orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>.",
								]
								__magic_name__		:Dict        =      tokenizer(lowercase__       ,					return_tensors='''pt'''       ,					padding=lowercase__     ).input_ids
								# fmt: off
								__magic_name__		:Optional[Any]        =      torch.tensor(
								    [
								        [ 3_8_5_3_0, 2_1_0_7_0_3, 2_5_6_2_9_9, 1_4_1_0, 2_5_6_2_9_8, 2_7_4, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
								        [   8_2_6, 3_2_1, 6_7_1, 2_5_9_2_2, 2_5_6_2_9_9, 2_7_4, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
								        [  1_4_6_0, 3_3_9, 3_1_2, 1_9_0_1_4, 1_0_6_2_0, 7_5_8, 2_5_6_2_9_9, 2_3_5_5,2_7_4, 1, 0, 0, 0, 0, 0, 0,0, 0],
								        [   5_1_7, 2_5_6_2_9_9, 1_4_8_6_9, 2_8_1, 3_0_1, 2_5_6_2_9_8, 2_7_5, 1_1_9_9_8_3,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
								        [   3_2_0, 2_5_6_2_9_9, 1_4_8_6_9, 2_8_1, 2_2_3_4, 2_8_9, 2_2_7_5, 3_3_3,6_1_3_9_1, 2_8_9, 2_5_6_2_9_8, 5_4_3, 2_5_6_2_9_7, 1_6_8_7_1_4, 3_2_9, 2_5_6_2_9_6,2_7_4, 1],
								    ]     )
								# fmt: on
								torch.testing.assert_allclose(lowercase__       ,					lowercase__     )

								__magic_name__		:List[Any]        =      model.generate(input_ids.to(lowercase__     )     )
								__magic_name__		:Dict        =      [
								    "<pad><extra_id_0> et<extra_id_1> [eod] <extra_id_2><extra_id_55>.. [eod] 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 <extra_id_56>ajšietosto<extra_id_56>lleux<extra_id_19><extra_id_6>ajšie</s>",
								    "<pad><extra_id_0>.<extra_id_1>.,<0x0A>...spech <0x0A><extra_id_20> <extra_id_21></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
								    "<pad><extra_id_0> are not going to be a part of the world. We are not going to be a part of<extra_id_1> and<extra_id_2><0x0A><extra_id_48>.<extra_id_48></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
								    "<pad><extra_id_0> door<extra_id_1>, the door<extra_id_2> 피해[/</s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
								    "<pad><extra_id_0>nyone who<extra_id_1> drink<extra_id_2> a<extra_id_3> alcohol<extra_id_4> A<extra_id_5> A. This<extra_id_6> I<extra_id_7><extra_id_52><extra_id_53></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
								]
								__magic_name__		:List[Any]        =      tokenizer.batch_decode(lowercase__     )
								self.assertEqual(lowercase__       ,					lowercase__     )








import argparse
import math
import os

import torch
from neural_compressor.utils.pytorch import load
from PIL import Image
from transformers import CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, StableDiffusionPipeline, UNetaDConditionModel

def      __lowercase					( ):







				"""simple docstring"""






				__magic_name__		:Optional[int]        =      argparse.ArgumentParser()
				parser.add_argument(
				    '''-m''',  '''--pretrained_model_name_or_path''',  type=snake_case,  default=snake_case,  required=snake_case,  help='''Path to pretrained model or model identifier from huggingface.co/models.''',  )
				parser.add_argument(
				    '''-c''',  '''--caption''',  type=snake_case,  default='''robotic cat with wings''',  help='''Text used to generate images.''',  )
				parser.add_argument(
				    '''-n''',  '''--images_num''',  type=snake_case,  default=4,  help='''How much images to generate.''',  )
				parser.add_argument(
				    '''-s''',  '''--seed''',  type=snake_case,  default=4_2,  help='''Seed for random process.''',  )
				parser.add_argument(
				    '''-ci''',  '''--cuda_id''',  type=snake_case,  default=0,  help='''cuda_id.''',  )
				__magic_name__		:Optional[Any]        =      parser.parse_args()
				return args

def      __lowercase					( snake_case,  snake_case,  snake_case ):







				"""simple docstring"""






				if not len(snake_case ) == rows * cols:
								raise ValueError('''The specified number of rows and columns are not correct.''' )

				__magic_name__     ,				__magic_name__		:Tuple        =      imgs[0].size
				__magic_name__		:List[Any]        =      Image.new('''RGB''',  size=(cols * w, rows * h) )
				__magic_name__     ,				__magic_name__		:List[Any]        =      grid.size

				for i, img in enumerate(snake_case ):
								grid.paste(snake_case,  box=(i % cols * w, i // cols * h) )
				return grid

def      __lowercase					( snake_case,  snake_case="robotic cat with wings",  snake_case=7.5,  snake_case=5_0,  snake_case=1,  snake_case=4_2,  ):







				"""simple docstring"""






				__magic_name__		:List[Any]        =      torch.Generator(pipeline.device ).manual_seed(snake_case )
				__magic_name__		:str        =      pipeline(
				    snake_case,  guidance_scale=snake_case,  num_inference_steps=snake_case,  generator=snake_case,  num_images_per_prompt=snake_case,  ).images
				__magic_name__		:Tuple        =      int(math.sqrt(snake_case ) )
				__magic_name__		:Union[str, Any]        =      image_grid(snake_case,  rows=_rows,  cols=num_images_per_prompt // _rows )
				return grid, images


SCREAMING_SNAKE_CASE__      :			Any							=		parse_args()
# Load models and create wrapper for stable diffusion
SCREAMING_SNAKE_CASE__      :			int							=		CLIPTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="""tokenizer""")
SCREAMING_SNAKE_CASE__      :			str							=		CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="""text_encoder""")
SCREAMING_SNAKE_CASE__      :			Optional[Any]							=		AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="""vae""")
SCREAMING_SNAKE_CASE__      :			int							=		UNetaDConditionModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="""unet""")

SCREAMING_SNAKE_CASE__      :			Union[str, Any]							=		StableDiffusionPipeline.from_pretrained(
    args.pretrained_model_name_or_path, text_encoder=text_encoder, vae=vae, unet=unet, tokenizer=tokenizer
)
SCREAMING_SNAKE_CASE__      :			str							=		lambda images, clip_input: (images, False)
if os.path.exists(os.path.join(args.pretrained_model_name_or_path, """best_model.pt""")):
					SCREAMING_SNAKE_CASE__      :			str							=		load(args.pretrained_model_name_or_path, model=unet)
					unet.eval()
					setattr(pipeline, """unet""", unet)
else:
					SCREAMING_SNAKE_CASE__      :			Dict							=		unet.to(torch.device("""cuda""", args.cuda_id))
SCREAMING_SNAKE_CASE__      :			Any							=		pipeline.to(unet.device)
SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__      :			Dict							=		generate_images(pipeline, prompt=args.caption, num_images_per_prompt=args.images_num, seed=args.seed)
grid.save(os.path.join(args.pretrained_model_name_or_path, """{}.png""".format("""_""".join(args.caption.split()))))
SCREAMING_SNAKE_CASE__      :			int							=		os.path.join(args.pretrained_model_name_or_path, """_""".join(args.caption.split()))
os.makedirs(dirname, exist_ok=True)
for idx, image in enumerate(images):
					image.save(os.path.join(dirname, """{}.png""".format(idx + 1)))






