





import gc
import unittest

import numpy as np
import torch
import torch.nn.functional as F
from transformers import (
    ClapTextConfig,
    ClapTextModelWithProjection,
    RobertaTokenizer,
    SpeechTaHifiGan,
    SpeechTaHifiGanConfig,
)

from diffusers import (
    AudioLDMPipeline,
    AutoencoderKL,
    DDIMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class       __lowerCAmelCase					(       _UpperCamelCase   ,    unittest.TestCase      ):


	'''simple docstring'''





	_A													=				AudioLDMPipeline
	_A													=				TEXT_TO_AUDIO_PARAMS
	_A													=				TEXT_TO_AUDIO_BATCH_PARAMS
	_A													=				frozenset(
	    [
	        "num_inference_steps",
	        "num_waveforms_per_prompt",
	        "generator",
	        "latents",
	        "output_type",
	        "return_dict",
	        "callback",
	        "callback_steps",
	    ]      )

	def    snake_case__(						self:						Optional[Any]			):
			torch.manual_seed(0			)
			lowercase__					:      Optional[int]      =		UNetaDConditionModel(
			    block_out_channels=(32, 64),					layers_per_block=2,					sample_size=32,					in_channels=4,					out_channels=4,					down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D'),					up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D'),					cross_attention_dim=(32, 64),					class_embed_type='simple_projection',					projection_class_embeddings_input_dim=32,					class_embeddings_concat=lowerCamelCase_,					)
			lowercase__					:      Dict      =		DDIMScheduler(
			    beta_start=0.0_0_0_8_5,					beta_end=0.0_1_2,					beta_schedule='scaled_linear',					clip_sample=lowerCamelCase_,					set_alpha_to_one=lowerCamelCase_,					)
			torch.manual_seed(0			)
			lowercase__					:      str      =		AutoencoderKL(
			    block_out_channels=[32, 64],					in_channels=1,					out_channels=1,					down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D'],					up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D'],					latent_channels=4,					)
			torch.manual_seed(0			)
			lowercase__					:      str      =		ClapTextConfig(
			    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=1000,					projection_dim=32,					)
			lowercase__					:      Optional[Any]      =		ClapTextModelWithProjection(lowerCamelCase_			)
			lowercase__					:      Dict      =		RobertaTokenizer.from_pretrained('hf-internal-testing/tiny-random-roberta',					model_max_length=77			)

			lowercase__					:      List[str]      =		SpeechTaHifiGanConfig(
			    model_in_dim=8,					sampling_rate=16000,					upsample_initial_channel=16,					upsample_rates=[2, 2],					upsample_kernel_sizes=[4, 4],					resblock_kernel_sizes=[3, 7],					resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]],					normalize_before=lowerCamelCase_,					)

			lowercase__					:      Tuple      =		SpeechTaHifiGan(lowerCamelCase_			)

			lowercase__					:      Optional[int]      =		{
			    'unet': unet,
			    'scheduler': scheduler,
			    'vae': vae,
			    'text_encoder': text_encoder,
			    'tokenizer': tokenizer,
			    'vocoder': vocoder,
			}
			return components

	def    snake_case__(						self:						Any,					lowerCamelCase_:						Optional[int],					lowerCamelCase_:						List[Any]=0			):
			if str(lowerCamelCase_			).startswith('mps'			):
					lowercase__					:      int      =		torch.manual_seed(lowerCamelCase_			)
			else:
					lowercase__					:      List[Any]      =		torch.Generator(device=lowerCamelCase_			).manual_seed(lowerCamelCase_			)
			lowercase__					:      Optional[int]      =		{
			    'prompt': 'A hammer hitting a wooden surface',
			    'generator': generator,
			    'num_inference_steps': 2,
			    'guidance_scale': 6.0,
			}
			return inputs

	def    snake_case__(						self:						List[Any]			):
			lowercase__					:      int      =		'cpu'  # ensure determinism for the device-dependent torch.Generator

			lowercase__					:      Optional[int]      =		self.get_dummy_components()
			lowercase__					:      str      =		AudioLDMPipeline(**lowerCamelCase_			)
			lowercase__					:      Optional[int]      =		audioldm_pipe.to(lowerCamelCase_			)
			audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_			)

			lowercase__					:      int      =		self.get_dummy_inputs(lowerCamelCase_			)
			lowercase__					:      List[Any]      =		audioldm_pipe(**lowerCamelCase_			)
			lowercase__					:      Tuple      =		output.audios[0]

			assert audio.ndim == 1
			assert len(lowerCamelCase_			) == 256

			lowercase__					:      Optional[Any]      =		audio[:10]
			lowercase__					:      List[str]      =		np.array(
			    [-0.0_0_5_0, 0.0_0_5_0, -0.0_0_6_0, 0.0_0_3_3, -0.0_0_2_6, 0.0_0_3_3, -0.0_0_2_7, 0.0_0_3_3, -0.0_0_2_8, 0.0_0_3_3]			)

			assert np.abs(audio_slice - expected_slice			).max() < 1E-2

	def    snake_case__(						self:						List[Any]			):
			lowercase__					:      Dict      =		self.get_dummy_components()
			lowercase__					:      Tuple      =		AudioLDMPipeline(**lowerCamelCase_			)
			lowercase__					:      Union[str, Any]      =		audioldm_pipe.to(lowerCamelCase_			)
			lowercase__					:      List[Any]      =		audioldm_pipe.to(lowerCamelCase_			)
			audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_			)

			lowercase__					:      Optional[Any]      =		self.get_dummy_inputs(lowerCamelCase_			)
			lowercase__					:      Dict      =		3 * [inputs['prompt']]

			# forward
			lowercase__					:      List[Any]      =		audioldm_pipe(**lowerCamelCase_			)
			lowercase__					:      Optional[Any]      =		output.audios[0]

			lowercase__					:      Tuple      =		self.get_dummy_inputs(lowerCamelCase_			)
			lowercase__					:      Optional[Any]      =		3 * [inputs.pop('prompt'			)]

			lowercase__					:      Any      =		audioldm_pipe.tokenizer(
			    lowerCamelCase_,					padding='max_length',					max_length=audioldm_pipe.tokenizer.model_max_length,					truncation=lowerCamelCase_,					return_tensors='pt',					)
			lowercase__					:      Union[str, Any]      =		text_inputs['input_ids'].to(lowerCamelCase_			)

			lowercase__					:      List[str]      =		audioldm_pipe.text_encoder(
			    lowerCamelCase_,					)
			lowercase__					:      Optional[int]      =		prompt_embeds.text_embeds
			# additional L_2 normalization over each hidden-state
			lowercase__					:      Dict      =		F.normalize(lowerCamelCase_,					dim=-1			)

			lowercase__					:      Tuple      =		prompt_embeds

			# forward
			lowercase__					:      str      =		audioldm_pipe(**lowerCamelCase_			)
			lowercase__					:      Any      =		output.audios[0]

			assert np.abs(audio_a - audio_a			).max() < 1E-2

	def    snake_case__(						self:						Optional[Any]			):
			lowercase__					:      Any      =		self.get_dummy_components()
			lowercase__					:      Optional[int]      =		AudioLDMPipeline(**lowerCamelCase_			)
			lowercase__					:      Any      =		audioldm_pipe.to(lowerCamelCase_			)
			lowercase__					:      List[str]      =		audioldm_pipe.to(lowerCamelCase_			)
			audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_			)

			lowercase__					:      Any      =		self.get_dummy_inputs(lowerCamelCase_			)
			lowercase__					:      Tuple      =		3 * ['this is a negative prompt']
			lowercase__					:      Optional[int]      =		negative_prompt
			lowercase__					:      Optional[int]      =		3 * [inputs['prompt']]

			# forward
			lowercase__					:      Optional[int]      =		audioldm_pipe(**lowerCamelCase_			)
			lowercase__					:      Union[str, Any]      =		output.audios[0]

			lowercase__					:      str      =		self.get_dummy_inputs(lowerCamelCase_			)
			lowercase__					:      Tuple      =		3 * [inputs.pop('prompt'			)]

			lowercase__					:      List[str]      =		[]
			for p in [prompt, negative_prompt]:
					lowercase__					:      List[str]      =		audioldm_pipe.tokenizer(
					    lowerCamelCase_,					padding='max_length',					max_length=audioldm_pipe.tokenizer.model_max_length,					truncation=lowerCamelCase_,					return_tensors='pt',					)
					lowercase__					:      Any      =		text_inputs['input_ids'].to(lowerCamelCase_			)

					lowercase__					:      Union[str, Any]      =		audioldm_pipe.text_encoder(
					    lowerCamelCase_,					)
					lowercase__					:      Union[str, Any]      =		text_embeds.text_embeds
					# additional L_2 normalization over each hidden-state
					lowercase__					:      Optional[int]      =		F.normalize(lowerCamelCase_,					dim=-1			)

					embeds.append(lowerCamelCase_			)

			lowercase__   ,				lowercase__					:      Optional[int]      =		embeds

			# forward
			lowercase__					:      int      =		audioldm_pipe(**lowerCamelCase_			)
			lowercase__					:      str      =		output.audios[0]

			assert np.abs(audio_a - audio_a			).max() < 1E-2

	def    snake_case__(						self:						Dict			):
			lowercase__					:      Dict      =		'cpu'  # ensure determinism for the device-dependent torch.Generator
			lowercase__					:      Optional[int]      =		self.get_dummy_components()
			lowercase__					:      List[Any]      =		PNDMScheduler(skip_prk_steps=lowerCamelCase_			)
			lowercase__					:      int      =		AudioLDMPipeline(**lowerCamelCase_			)
			lowercase__					:      Union[str, Any]      =		audioldm_pipe.to(lowerCamelCase_			)
			audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_			)

			lowercase__					:      Any      =		self.get_dummy_inputs(lowerCamelCase_			)
			lowercase__					:      List[str]      =		'egg cracking'
			lowercase__					:      str      =		audioldm_pipe(**lowerCamelCase_,					negative_prompt=lowerCamelCase_			)
			lowercase__					:      str      =		output.audios[0]

			assert audio.ndim == 1
			assert len(lowerCamelCase_			) == 256

			lowercase__					:      Tuple      =		audio[:10]
			lowercase__					:      Any      =		np.array(
			    [-0.0_0_5_1, 0.0_0_5_0, -0.0_0_6_0, 0.0_0_3_4, -0.0_0_2_6, 0.0_0_3_3, -0.0_0_2_7, 0.0_0_3_3, -0.0_0_2_8, 0.0_0_3_2]			)

			assert np.abs(audio_slice - expected_slice			).max() < 1E-2

	def    snake_case__(						self:						Dict			):
			lowercase__					:      List[Any]      =		'cpu'  # ensure determinism for the device-dependent torch.Generator
			lowercase__					:      List[Any]      =		self.get_dummy_components()
			lowercase__					:      Optional[Any]      =		PNDMScheduler(skip_prk_steps=lowerCamelCase_			)
			lowercase__					:      Dict      =		AudioLDMPipeline(**lowerCamelCase_			)
			lowercase__					:      Dict      =		audioldm_pipe.to(lowerCamelCase_			)
			audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_			)

			lowercase__					:      Dict      =		'A hammer hitting a wooden surface'

			# test num_waveforms_per_prompt=1 (default)
			lowercase__					:      Tuple      =		audioldm_pipe(lowerCamelCase_,					num_inference_steps=2			).audios

			assert audios.shape == (1, 256)

			# test num_waveforms_per_prompt=1 (default) for batch of prompts
			lowercase__					:      List[str]      =		2
			lowercase__					:      Optional[int]      =		audioldm_pipe([prompt] * batch_size,					num_inference_steps=2			).audios

			assert audios.shape == (batch_size, 256)

			# test num_waveforms_per_prompt for single prompt
			lowercase__					:      Union[str, Any]      =		2
			lowercase__					:      List[Any]      =		audioldm_pipe(lowerCamelCase_,					num_inference_steps=2,					num_waveforms_per_prompt=lowerCamelCase_			).audios

			assert audios.shape == (num_waveforms_per_prompt, 256)

			# test num_waveforms_per_prompt for batch of prompts
			lowercase__					:      Tuple      =		2
			lowercase__					:      Any      =		audioldm_pipe(
			    [prompt] * batch_size,					num_inference_steps=2,					num_waveforms_per_prompt=lowerCamelCase_			).audios

			assert audios.shape == (batch_size * num_waveforms_per_prompt, 256)

	def    snake_case__(						self:						str			):
			lowercase__					:      Tuple      =		'cpu'  # ensure determinism for the device-dependent torch.Generator
			lowercase__					:      Optional[Any]      =		self.get_dummy_components()
			lowercase__					:      int      =		AudioLDMPipeline(**lowerCamelCase_			)
			lowercase__					:      Optional[int]      =		audioldm_pipe.to(lowerCamelCase_			)
			audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_			)
			lowercase__					:      Optional[Any]      =		audioldm_pipe.vocoder.config.sampling_rate

			lowercase__					:      List[str]      =		self.get_dummy_inputs(lowerCamelCase_			)
			lowercase__					:      str      =		audioldm_pipe(audio_length_in_s=0.0_1_6,					**lowerCamelCase_			)
			lowercase__					:      List[Any]      =		output.audios[0]

			assert audio.ndim == 1
			assert len(lowerCamelCase_			) / vocoder_sampling_rate == 0.0_1_6

			lowercase__					:      int      =		audioldm_pipe(audio_length_in_s=0.0_3_2,					**lowerCamelCase_			)
			lowercase__					:      List[str]      =		output.audios[0]

			assert audio.ndim == 1
			assert len(lowerCamelCase_			) / vocoder_sampling_rate == 0.0_3_2

	def    snake_case__(						self:						Any			):
			lowercase__					:      Any      =		self.get_dummy_components()
			lowercase__					:      List[Any]      =		AudioLDMPipeline(**lowerCamelCase_			)
			lowercase__					:      Any      =		audioldm_pipe.to(lowerCamelCase_			)
			audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_			)

			lowercase__					:      Optional[Any]      =		['hey']

			lowercase__					:      List[str]      =		audioldm_pipe(lowerCamelCase_,					num_inference_steps=1			)
			lowercase__					:      Tuple      =		output.audios.shape
			assert audio_shape == (1, 256)

			lowercase__					:      Optional[Any]      =		audioldm_pipe.vocoder.config
			config.model_in_dim *= 2
			lowercase__					:      Tuple      =		SpeechTaHifiGan(lowerCamelCase_			).to(lowerCamelCase_			)
			lowercase__					:      Optional[int]      =		audioldm_pipe(lowerCamelCase_,					num_inference_steps=1			)
			lowercase__					:      Dict      =		output.audios.shape
			# waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
			assert audio_shape == (1, 256)

	def    snake_case__(						self:						List[Any]			):
			self._test_attention_slicing_forward_pass(test_mean_pixel_difference=lowerCamelCase_			)

	def    snake_case__(						self:						Tuple			):
			self._test_inference_batch_single_identical(test_mean_pixel_difference=lowerCamelCase_			)





	@unittest.skipIf(
	    torch_device != 'cuda' or not is_xformers_available(),					reason='XFormers attention is only available with CUDA and `xformers` installed',					)
	def    snake_case__(						self:						Union[str, Any]			):
			self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=lowerCamelCase_			)


@slow
class       __lowerCAmelCase					(       unittest.TestCase      ):


	'''simple docstring'''





	def    snake_case__(						self:						Dict			):
			super().tearDown()
			gc.collect()
			torch.cuda.empty_cache()

	def    snake_case__(						self:						str,					lowerCamelCase_:						List[Any],					lowerCamelCase_:						Optional[Any]="cpu",					lowerCamelCase_:						Tuple=torch.floataa,					lowerCamelCase_:						Optional[int]=0			):
			lowercase__					:      Optional[int]      =		torch.Generator(device=lowerCamelCase_			).manual_seed(lowerCamelCase_			)
			lowercase__					:      List[Any]      =		np.random.RandomState(lowerCamelCase_			).standard_normal((1, 8, 128, 16)			)
			lowercase__					:      Optional[int]      =		torch.from_numpy(lowerCamelCase_			).to(device=lowerCamelCase_,					dtype=lowerCamelCase_			)
			lowercase__					:      int      =		{
			    'prompt': 'A hammer hitting a wooden surface',
			    'latents': latents,
			    'generator': generator,
			    'num_inference_steps': 3,
			    'guidance_scale': 2.5,
			}
			return inputs

	def    snake_case__(						self:						List[str]			):
			lowercase__					:      str      =		AudioLDMPipeline.from_pretrained('cvssp/audioldm'			)
			lowercase__					:      Tuple      =		audioldm_pipe.to(lowerCamelCase_			)
			audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_			)

			lowercase__					:      List[str]      =		self.get_inputs(lowerCamelCase_			)
			lowercase__					:      Union[str, Any]      =		25
			lowercase__					:      str      =		audioldm_pipe(**lowerCamelCase_			).audios[0]

			assert audio.ndim == 1
			assert len(lowerCamelCase_			) == 81920

			lowercase__					:      Optional[int]      =		audio[77230:77240]
			lowercase__					:      List[Any]      =		np.array(
			    [-0.4_8_8_4, -0.4_6_0_7, 0.0_0_2_3, 0.5_0_0_7, 0.5_8_9_6, 0.5_1_5_1, 0.3_8_1_3, -0.0_2_0_8, -0.3_6_8_7, -0.4_3_1_5]			)
			lowercase__					:      Optional[Any]      =		np.abs(expected_slice - audio_slice			).max()
			assert max_diff < 1E-2





	def    snake_case__(						self:						Optional[int]			):
			lowercase__					:      str      =		AudioLDMPipeline.from_pretrained('cvssp/audioldm'			)
			lowercase__					:      Dict      =		LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config			)
			lowercase__					:      str      =		audioldm_pipe.to(lowerCamelCase_			)
			audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_			)

			lowercase__					:      int      =		self.get_inputs(lowerCamelCase_			)
			lowercase__					:      Optional[int]      =		audioldm_pipe(**lowerCamelCase_			).audios[0]

			assert audio.ndim == 1
			assert len(lowerCamelCase_			) == 81920

			lowercase__					:      Any      =		audio[27780:27790]
			lowercase__					:      Dict      =		np.array([-0.2_1_3_1, -0.0_8_7_3, -0.0_1_2_4, -0.0_1_8_9, 0.0_5_6_9, 0.1_3_7_3, 0.1_8_8_3, 0.2_8_8_6, 0.3_2_9_7, 0.2_2_1_2]			)
			lowercase__					:      Tuple      =		np.abs(expected_slice - audio_slice			).max()
			assert max_diff < 3E-2









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       SCREAMING_SNAKE_CASE__  (       )   ->							int:




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


def       SCREAMING_SNAKE_CASE__  (       _lowercase     : Union[str, Any]   ,    _lowercase     : Optional[int]   ,    _lowercase     : Union[str, Any]			)   ->							str:




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

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

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


def       SCREAMING_SNAKE_CASE__  (       _lowercase     : Optional[int]   ,    _lowercase     : Tuple="robotic cat with wings"   ,    _lowercase     : Optional[Any]=7.5   ,    _lowercase     : int=50   ,    _lowercase     : str=1   ,    _lowercase     : Any=42   ,    )   ->							str:




		'''simple docstring'''
		lowercase__					:      int      =		torch.Generator(pipeline.device			).manual_seed(_lowercase			)
		lowercase__					:      Any      =		pipeline(
		    _lowercase   ,    guidance_scale=_lowercase   ,    num_inference_steps=_lowercase   ,    generator=_lowercase   ,    num_images_per_prompt=_lowercase   ,    ).images
		lowercase__					:      Dict      =		int(math.sqrt(_lowercase			)			)
		lowercase__					:      int      =		image_grid(_lowercase   ,    rows=_rows   ,    cols=num_images_per_prompt // _rows			)
		return grid, images


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

__UpperCamelCase:	Tuple = StableDiffusionPipeline.from_pretrained(
    args.pretrained_model_name_or_path, text_encoder=text_encoder, vae=vae, unet=unet, tokenizer=tokenizer
)
__UpperCamelCase:	List[str] = lambda images, clip_input: (images, False)
if os.path.exists(os.path.join(args.pretrained_model_name_or_path, """best_model.pt""")):
			__UpperCamelCase:	Dict = load(args.pretrained_model_name_or_path, model=unet)
			unet.eval()
			setattr(pipeline, """unet""", unet)
else:
			__UpperCamelCase:	Dict = unet.to(torch.device("""cuda""", args.cuda_id))
__UpperCamelCase:	Optional[int] = pipeline.to(unet.device)
__UpperCamelCase,	__UpperCamelCase:	Union[str, Any] = 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()))))
__UpperCamelCase:	List[str] = 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)))


