


"""simple docstring"""


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 _lowerCamelCase			(	):

			'''simple docstring'''




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


def _lowerCamelCase			(	_UpperCamelCase					,				_UpperCamelCase					,				_UpperCamelCase   ):

			'''simple docstring'''




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

			__lowerCAmelCase					,			__lowerCAmelCase       					=    imgs[0].size
			__lowerCAmelCase       					=    Image.new("RGB"					,				size=(cols * w, rows * h)   )
			__lowerCAmelCase					,			__lowerCAmelCase       					=    grid.size

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


def _lowerCamelCase			(	_UpperCamelCase					,				_UpperCamelCase="robotic cat with wings"					,				_UpperCamelCase=7.5					,				_UpperCamelCase=50					,				_UpperCamelCase=1					,				_UpperCamelCase=42					,				):

			'''simple docstring'''




			__lowerCAmelCase       					=    torch.Generator(pipeline.device   ).manual_seed(__UpperCamelCase   )
			__lowerCAmelCase       					=    pipeline(
			    __UpperCamelCase					,				guidance_scale=__UpperCamelCase					,				num_inference_steps=__UpperCamelCase					,				generator=__UpperCamelCase					,				num_images_per_prompt=__UpperCamelCase					,				).images
			__lowerCAmelCase       					=    int(math.sqrt(__UpperCamelCase   )   )
			__lowerCAmelCase       					=    image_grid(__UpperCamelCase					,				rows=_rows					,				cols=num_images_per_prompt // _rows   )
			return grid, images


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

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








"""simple docstring"""


import os
import tempfile
import unittest
from pathlib import Path

from transformers import AutoConfig, is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
	import tensorflow as tf

	from transformers import TensorFlowBenchmark, TensorFlowBenchmarkArguments
@require_tf
class 						_UpperCamelCase						(     unittest.TestCase	):
							'''simple docstring'''







							def     snake_case       (  self					,     __a			):
										for model_result in results.values():
													for batch_size, sequence_length in zip(model_result["bs"]					,     model_result["ss"]			):
																__lowerCAmelCase       					=    model_result["result"][batch_size][sequence_length]
																self.assertIsNotNone(__a			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    "sshleifer/tiny-gpt2"
										__lowerCAmelCase       					=    TensorFlowBenchmarkArguments(
										    models=[MODEL_ID]					,     training=__a					,     inference=__a					,     sequence_lengths=[8]					,     batch_sizes=[1]					,     eager_mode=__a					,     multi_process=__a					,     )
										__lowerCAmelCase       					=    TensorFlowBenchmark(__a			)
										__lowerCAmelCase       					=    benchmark.run()
										self.check_results_dict_not_empty(results.time_inference_result			)
										self.check_results_dict_not_empty(results.memory_inference_result			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    "sgugger/tiny-distilbert-classification"
										__lowerCAmelCase       					=    TensorFlowBenchmarkArguments(
										    models=[MODEL_ID]					,     training=__a					,     inference=__a					,     sequence_lengths=[8]					,     batch_sizes=[1]					,     multi_process=__a					,     only_pretrain_model=__a					,     )
										__lowerCAmelCase       					=    TensorFlowBenchmark(__a			)
										__lowerCAmelCase       					=    benchmark.run()
										self.check_results_dict_not_empty(results.time_inference_result			)
										self.check_results_dict_not_empty(results.memory_inference_result			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    "sshleifer/tiny-gpt2"
										__lowerCAmelCase       					=    TensorFlowBenchmarkArguments(
										    models=[MODEL_ID]					,     training=__a					,     inference=__a					,     sequence_lengths=[8]					,     batch_sizes=[1]					,     multi_process=__a					,     )
										__lowerCAmelCase       					=    TensorFlowBenchmark(__a			)
										__lowerCAmelCase       					=    benchmark.run()
										self.check_results_dict_not_empty(results.time_inference_result			)
										self.check_results_dict_not_empty(results.memory_inference_result			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    "sshleifer/tiny-gpt2"
										__lowerCAmelCase       					=    AutoConfig.from_pretrained(__a			)
										__lowerCAmelCase       					=    TensorFlowBenchmarkArguments(
										    models=[MODEL_ID]					,     training=__a					,     inference=__a					,     sequence_lengths=[8]					,     batch_sizes=[1]					,     eager_mode=__a					,     multi_process=__a					,     )
										__lowerCAmelCase       					=    TensorFlowBenchmark(__a					,     [config]			)
										__lowerCAmelCase       					=    benchmark.run()
										self.check_results_dict_not_empty(results.time_inference_result			)
										self.check_results_dict_not_empty(results.memory_inference_result			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    "sshleifer/tiny-gpt2"
										__lowerCAmelCase       					=    AutoConfig.from_pretrained(__a			)
										__lowerCAmelCase       					=    TensorFlowBenchmarkArguments(
										    models=[MODEL_ID]					,     training=__a					,     inference=__a					,     sequence_lengths=[8]					,     batch_sizes=[1]					,     multi_process=__a					,     )
										__lowerCAmelCase       					=    TensorFlowBenchmark(__a					,     [config]			)
										__lowerCAmelCase       					=    benchmark.run()
										self.check_results_dict_not_empty(results.time_inference_result			)
										self.check_results_dict_not_empty(results.memory_inference_result			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    "sshleifer/tiny-gpt2"
										__lowerCAmelCase       					=    TensorFlowBenchmarkArguments(
										    models=[MODEL_ID]					,     training=__a					,     inference=__a					,     sequence_lengths=[8]					,     batch_sizes=[1]					,     multi_process=__a					,     )
										__lowerCAmelCase       					=    TensorFlowBenchmark(__a			)
										__lowerCAmelCase       					=    benchmark.run()
										self.check_results_dict_not_empty(results.time_train_result			)
										self.check_results_dict_not_empty(results.memory_train_result			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    "sshleifer/tiny-gpt2"
										__lowerCAmelCase       					=    AutoConfig.from_pretrained(__a			)
										__lowerCAmelCase       					=    TensorFlowBenchmarkArguments(
										    models=[MODEL_ID]					,     training=__a					,     inference=__a					,     sequence_lengths=[8]					,     batch_sizes=[1]					,     multi_process=__a					,     )
										__lowerCAmelCase       					=    TensorFlowBenchmark(__a					,     [config]			)
										__lowerCAmelCase       					=    benchmark.run()
										self.check_results_dict_not_empty(results.time_train_result			)
										self.check_results_dict_not_empty(results.memory_train_result			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    "patrickvonplaten/t5-tiny-random"
										__lowerCAmelCase       					=    AutoConfig.from_pretrained(__a			)
										__lowerCAmelCase       					=    TensorFlowBenchmarkArguments(
										    models=[MODEL_ID]					,     training=__a					,     inference=__a					,     sequence_lengths=[8]					,     batch_sizes=[1]					,     multi_process=__a					,     )
										__lowerCAmelCase       					=    TensorFlowBenchmark(__a					,     configs=[config]			)
										__lowerCAmelCase       					=    benchmark.run()
										self.check_results_dict_not_empty(results.time_inference_result			)
										self.check_results_dict_not_empty(results.memory_inference_result			)

							@unittest.skipIf(is_tf_available() and len(tf.config.list_physical_devices("GPU"			)			) == 0					,     "Cannot do xla on CPU."			)
							def     snake_case       (  self			):
										__lowerCAmelCase       					=    "sshleifer/tiny-gpt2"
										__lowerCAmelCase       					=    TensorFlowBenchmarkArguments(
										    models=[MODEL_ID]					,     training=__a					,     inference=__a					,     sequence_lengths=[8]					,     batch_sizes=[1]					,     use_xla=__a					,     multi_process=__a					,     )
										__lowerCAmelCase       					=    TensorFlowBenchmark(__a			)
										__lowerCAmelCase       					=    benchmark.run()
										self.check_results_dict_not_empty(results.time_inference_result			)
										self.check_results_dict_not_empty(results.memory_inference_result			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    "sshleifer/tiny-gpt2"
										with tempfile.TemporaryDirectory() as tmp_dir:
													__lowerCAmelCase       					=    TensorFlowBenchmarkArguments(
													    models=[MODEL_ID]					,     inference=__a					,     save_to_csv=__a					,     sequence_lengths=[8]					,     batch_sizes=[1]					,     inference_time_csv_file=os.path.join(__a					,     "inf_time.csv"			)					,     inference_memory_csv_file=os.path.join(__a					,     "inf_mem.csv"			)					,     env_info_csv_file=os.path.join(__a					,     "env.csv"			)					,     multi_process=__a					,     )
													__lowerCAmelCase       					=    TensorFlowBenchmark(__a			)
													benchmark.run()
													self.assertTrue(Path(os.path.join(__a					,     "inf_time.csv"			)			).exists()			)
													self.assertTrue(Path(os.path.join(__a					,     "inf_mem.csv"			)			).exists()			)
													self.assertTrue(Path(os.path.join(__a					,     "env.csv"			)			).exists()			)



							def     snake_case       (  self			):
										__lowerCAmelCase       					=    "sshleifer/tiny-gpt2"

										def _check_summary_is_not_empty(__a			):
													self.assertTrue(hasattr(__a					,     "sequential"			)			)
													self.assertTrue(hasattr(__a					,     "cumulative"			)			)
													self.assertTrue(hasattr(__a					,     "current"			)			)
													self.assertTrue(hasattr(__a					,     "total"			)			)

										with tempfile.TemporaryDirectory() as tmp_dir:
													__lowerCAmelCase       					=    TensorFlowBenchmarkArguments(
													    models=[MODEL_ID]					,     inference=__a					,     sequence_lengths=[8]					,     batch_sizes=[1]					,     log_filename=os.path.join(__a					,     "log.txt"			)					,     log_print=__a					,     trace_memory_line_by_line=__a					,     eager_mode=__a					,     multi_process=__a					,     )
													__lowerCAmelCase       					=    TensorFlowBenchmark(__a			)
													__lowerCAmelCase       					=    benchmark.run()
													_check_summary_is_not_empty(result.inference_summary			)
													self.assertTrue(Path(os.path.join(__a					,     "log.txt"			)			).exists()			)




