





from __future__ import annotations

import os
import tempfile
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import is_tensorflow_text_available, is_tf_available
from transformers.testing_utils import require_tensorflow_text, require_tf, slow

from ..test_modeling_tf_common import floats_tensor
from .test_framework_agnostic import GenerationIntegrationTestsMixin


if is_tf_available():
			import tensorflow as tf

			from transformers import (
			    AutoTokenizer,
			    TFAutoModelForCausalLM,
			    TFAutoModelForSeqaSeqLM,
			    TFAutoModelForSpeechSeqaSeq,
			    TFAutoModelForVisionaSeq,
			    TFBartForConditionalGeneration,
			    TFLogitsProcessorList,
			    TFMinLengthLogitsProcessor,
			    tf_top_k_top_p_filtering,
			)

if is_tensorflow_text_available():
			import tensorflow_text as text


@require_tf
class       __lowerCAmelCase					(       unittest.TestCase      ):


	'''simple docstring'''





	def    snake_case__(						self:						Optional[int]			):
			lowercase__					:      str      =		tf.convert_to_tensor(
			    [
			        [
			            8.2_2_2_0_9_9_1,  # 3rd highest value; idx. 0
			            -0.5_6_2_0_0_4_4,
			            5.2_3_2_2_9_7_5_2,
			            4.0_3_8_6_3_9_3,
			            -6.8_7_9_8_3_7_8,
			            -0.5_4_7_8_5_8_0_2,
			            -3.2_0_1_2_1_5_3,
			            2.9_2_7_7_7_1_7_6,
			            1.8_8_1_7_1_9_5_3,
			            7.3_5_3_4_1_2_7_6,  # 5th highest value; idx. 9
			            8.4_3_2_0_7_8_3_3,  # 2nd highest value; idx. 10
			            -9.8_5_7_1_1_8_3_6,
			            -5.9_6_2_0_9_2_3_6,
			            -1.1_3_0_3_9_1_6_1,
			            -7.1_1_1_5_2_9_4,
			            -0.8_3_6_9_6_3_3,
			            -5.3_1_8_6_4_0_8,
			            7.0_6_4_2_7_4_0_7,
			            0.8_1_3_6_9_3_4_4,
			            -0.8_2_0_2_3_8_1_7,
			            -5.9_1_7_9_7_9_6,
			            0.5_8_8_1_3_4_4_3,
			            -6.9_9_7_7_8_4_3_8,
			            4.7_1_5_5_1_1_8_9,
			            -0.1_8_7_7_1_6_3_7,
			            7.4_4_0_2_0_7_5_9,  # 4th highest value; idx. 25
			            9.3_8_4_5_0_9_8_7,  # 1st highest value; idx. 26
			            2.1_2_6_6_2_9_4_1,
			            -9.3_2_5_6_2_0_3_8,
			            2.3_5_6_5_2_5_2_2,
			        ],  # cummulative prob of 5 highest values <= 0.6
			        [
			            0.5_8_4_2_5_5_1_8,
			            4.5_3_1_3_9_2_3_8,
			            -5.5_7_5_1_0_4_6_4,
			            -6.2_8_0_3_0_6_9_9,
			            -7.1_9_5_2_9_5_0_3,
			            -4.0_2_1_2_2_5_5_1,
			            1.3_9_3_3_7_0_3_7,
			            -6.0_6_7_0_7_0_5_7,
			            1.5_9_4_8_0_5_1_7,
			            -9.6_4_3_1_1_9,
			            0.0_3_9_0_7_7_9_9,
			            0.6_7_2_3_1_7_6_2,
			            -8.8_8_2_0_6_7_2_6,
			            6.2_7_1_1_5_9_2_2,  # 4th highest value; idx. 13
			            2.2_8_5_2_0_7_2_3,
			            4.8_2_7_6_7_5_0_6,
			            4.3_0_4_2_1_3_6_8,
			            8.8_2_7_5_3_1_3,  # 2nd highest value; idx. 17
			            5.4_4_0_2_9_9_5_8,  # 5th highest value; idx. 18
			            -4.4_7_3_5_7_9_4,
			            7.3_8_5_7_9_5_3_6,  # 3rd highest value; idx. 20
			            -2.9_1_0_5_1_6_6_3,
			            2.6_1_9_4_6_0_7_7,
			            -2.5_6_7_4_7_6_2,
			            -9.4_8_9_5_9_3_0_2,
			            -4.0_2_9_2_2_6_4_5,
			            -1.3_5_4_1_6_9_1_8,
			            9.6_7_7_0_2_3_2_3,  # 1st highest value; idx. 27
			            -5.8_9_4_7_8_5_5_3,
			            1.8_5_3_7_0_4_6_7,
			        ],  # cummulative prob of 5 highest values <= 0.6
			    ],					dtype=tf.floataa,					)

			lowercase__					:      Any      =		tf.convert_to_tensor(
			    [[0, 0], [0, 9], [0, 10], [0, 25], [0, 26], [1, 13], [1, 17], [1, 18], [1, 20], [1, 27]],					dtype=tf.intaa,					)  # expected non filtered idx as noted above

			lowercase__					:      Tuple      =		tf.convert_to_tensor(
			    [8.2_2_2_0_9_9, 7.3_5_3_4_1_2_6, 8.4_3_2_0_7_8, 7.4_4_0_2_0_7_5, 9.3_8_4_5_1, 6.2_7_1_1_5_9, 8.8_2_7_5_3_1, 5.4_4_0_2_9_9_5, 7.3_8_5_7_9_5_6, 9.6_7_7_0_2_3],					dtype=tf.floataa,					)  # expected non filtered values as noted above

			lowercase__					:      Any      =		tf_top_k_top_p_filtering(lowerCamelCase_,					top_k=10,					top_p=0.6,					min_tokens_to_keep=4			)

			lowercase__					:      Union[str, Any]      =		output[output != -float('inf'			)]
			lowercase__					:      Optional[Any]      =		tf.cast(
			    tf.where(tf.not_equal(lowerCamelCase_,					tf.constant(-float('inf'			),					dtype=tf.floataa			)			)			),					dtype=tf.intaa,					)

			tf.debugging.assert_near(lowerCamelCase_,					lowerCamelCase_,					rtol=1E-12			)
			tf.debugging.assert_equal(lowerCamelCase_,					lowerCamelCase_			)


@require_tf
class       __lowerCAmelCase					(       unittest.TestCase   ,    _UpperCamelCase      ):


	'''simple docstring'''





	if is_tf_available():
		_A													=				{
		    "AutoModelForCausalLM": TFAutoModelForCausalLM,
		    "AutoModelForSpeechSeq2Seq": TFAutoModelForSpeechSeqaSeq,
		    "AutoModelForSeq2SeqLM": TFAutoModelForSeqaSeqLM,
		    "AutoModelForVision2Seq": TFAutoModelForVisionaSeq,
		    "LogitsProcessorList": TFLogitsProcessorList,
		    "MinLengthLogitsProcessor": TFMinLengthLogitsProcessor,
		    "create_tensor_fn": tf.convert_to_tensor,
		    "floats_tensor": floats_tensor,
		    "return_tensors": "tf",
		}

	@slow
	def    snake_case__(						self:						List[Any]			):
			# TF-only test: tf.saved_model export
			lowercase__					:      int      =		TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'			)
			lowercase__					:      Dict      =		2
			lowercase__					:      Dict      =		2


			class       __lowerCAmelCase					(       tf.Module      ):


				'''simple docstring'''





				def __init__(						self:						int,					lowerCamelCase_:						Any			):
						super(lowerCamelCase_,					self			).__init__()
						lowercase__					:      List[Any]      =		model





				@tf.function(
				    input_signature=(
				        tf.TensorSpec((None, input_length),					tf.intaa,					name='input_ids'			),
				        tf.TensorSpec((None, input_length),					tf.intaa,					name='attention_mask'			),
				    ),					jit_compile=lowerCamelCase_,					)
				def    snake_case__(						self:						Dict,					lowerCamelCase_:						List[Any],					lowerCamelCase_:						str			):
						lowercase__					:      Optional[int]      =		self.model.generate(
						    input_ids=lowerCamelCase_,					attention_mask=lowerCamelCase_,					max_new_tokens=lowerCamelCase_,					return_dict_in_generate=lowerCamelCase_,					)
						return {"sequences": outputs["sequences"]}
			lowercase__					:      int      =		[[2, 0], [102, 103]]
			lowercase__					:      str      =		[[1, 0], [1, 1]]
			lowercase__					:      Tuple      =		DummyModel(model=lowerCamelCase_			)
			with tempfile.TemporaryDirectory() as tmp_dir:
					tf.saved_model.save(lowerCamelCase_,					lowerCamelCase_,					signatures={'serving_default': dummy_model.serving}			)
					lowercase__					:      List[str]      =		tf.saved_model.load(lowerCamelCase_			).signatures['serving_default']
					for batch_size in range(1,					len(lowerCamelCase_			) + 1			):
							lowercase__					:      Any      =		{
							    'input_ids': tf.constant(dummy_input_ids[:batch_size]			),
							    'attention_mask': tf.constant(dummy_attention_masks[:batch_size]			),
							}
							lowercase__					:      List[Any]      =		serving_func(**lowerCamelCase_			)['sequences']
							lowercase__					:      str      =		test_model.generate(**lowerCamelCase_,					max_new_tokens=lowerCamelCase_			)
							tf.debugging.assert_equal(lowerCamelCase_,					lowerCamelCase_			)

	@slow
	def    snake_case__(						self:						Union[str, Any]			):
			# TF-only test: tf.saved_model export
			lowercase__					:      Dict      =		TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'			)
			lowercase__					:      Any      =		1
			lowercase__					:      Optional[int]      =		2


			class       __lowerCAmelCase					(       tf.Module      ):


				'''simple docstring'''





				def __init__(						self:						Optional[Any],					lowerCamelCase_:						Optional[int]			):
						super(lowerCamelCase_,					self			).__init__()
						lowercase__					:      Tuple      =		model





				@tf.function(
				    input_signature=(
				        tf.TensorSpec((batch_size, None),					tf.intaa,					name='input_ids'			),
				        tf.TensorSpec((batch_size, None),					tf.intaa,					name='attention_mask'			),
				    ),					jit_compile=lowerCamelCase_,					)
				def    snake_case__(						self:						Optional[int],					lowerCamelCase_:						str,					lowerCamelCase_:						str			):
						lowercase__					:      int      =		self.model.generate(
						    input_ids=lowerCamelCase_,					attention_mask=lowerCamelCase_,					max_new_tokens=lowerCamelCase_,					return_dict_in_generate=lowerCamelCase_,					)
						return {"sequences": outputs["sequences"]}
			lowercase__					:      str      =		[[2], [102, 103]]
			lowercase__					:      List[str]      =		[[1], [1, 1]]
			lowercase__					:      List[str]      =		DummyModel(model=lowerCamelCase_			)
			with tempfile.TemporaryDirectory() as tmp_dir:
					tf.saved_model.save(lowerCamelCase_,					lowerCamelCase_,					signatures={'serving_default': dummy_model.serving}			)
					lowercase__					:      str      =		tf.saved_model.load(lowerCamelCase_			).signatures['serving_default']
					for input_row in range(len(lowerCamelCase_			)			):
							lowercase__					:      int      =		{
							    'input_ids': tf.constant([dummy_input_ids[input_row]]			),
							    'attention_mask': tf.constant([dummy_attention_masks[input_row]]			),
							}
							lowercase__					:      Tuple      =		serving_func(**lowerCamelCase_			)['sequences']
							lowercase__					:      Union[str, Any]      =		test_model.generate(**lowerCamelCase_,					max_new_tokens=lowerCamelCase_			)
							tf.debugging.assert_equal(lowerCamelCase_,					lowerCamelCase_			)

	@slow
	@require_tensorflow_text
	def    snake_case__(						self:						List[str]			):
			# TF-only test: tf.saved_model export
			with tempfile.TemporaryDirectory() as tmp_dir:
					# file needed to load the TF tokenizer
					hf_hub_download(repo_id='google/flan-t5-small',					filename='spiece.model',					local_dir=lowerCamelCase_			)


					class       __lowerCAmelCase					(       tf.keras.layers.Layer      ):


						'''simple docstring'''





						def __init__(						self:						Optional[Any]			):
								super().__init__()
								lowercase__					:      Any      =		text.SentencepieceTokenizer(
								    model=tf.io.gfile.GFile(os.path.join(lowerCamelCase_,					'spiece.model'			),					'rb'			).read()			)
								lowercase__					:      Any      =		TFAutoModelForSeqaSeqLM.from_pretrained('hf-internal-testing/tiny-random-t5'			)





						def    snake_case__(						self:						Union[str, Any],					lowerCamelCase_:						Optional[Any],					*lowerCamelCase_:						List[str],					**lowerCamelCase_:						Optional[int]			):
								lowercase__					:      Any      =		self.tokenizer.tokenize(lowerCamelCase_			)
								lowercase__   ,				lowercase__					:      Tuple      =		text.pad_model_inputs(
								    lowerCamelCase_,					max_seq_length=64,					pad_value=self.model.config.pad_token_id			)
								lowercase__					:      Dict      =		self.model.generate(input_ids=lowerCamelCase_,					attention_mask=lowerCamelCase_			)
								return self.tokenizer.detokenize(lowerCamelCase_			)
					lowercase__					:      Union[str, Any]      =		CompleteSentenceTransformer()
					lowercase__					:      Optional[Any]      =		tf.keras.layers.Input(shape=(1,),					dtype=tf.string,					name='inputs'			)
					lowercase__					:      Optional[int]      =		complete_model(lowerCamelCase_			)
					lowercase__					:      Dict      =		tf.keras.Model(lowerCamelCase_,					lowerCamelCase_			)
					keras_model.save(lowerCamelCase_			)

	def    snake_case__(						self:						str			):
			# Has PT equivalent: this test relies on random sampling
			lowercase__					:      Any      =		{
			    'do_sample': True,
			    'num_beams': 1,
			    'top_p': 0.7,
			    'top_k': 10,
			    'temperature': 0.7,
			}
			lowercase__					:      int      =		14

			lowercase__					:      str      =		AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'			)
			lowercase__					:      Any      =		'Hello, my dog is cute and'
			lowercase__					:      List[Any]      =		tokenizer(lowerCamelCase_,					return_tensors='tf'			)
			lowercase__					:      List[Any]      =		TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'			)

			lowercase__					:      int      =		638
			# forces the generation to happen on CPU, to avoid GPU-related quirks
			with tf.device(':/CPU:0'			):
					tf.random.set_seed(0			)
					lowercase__					:      Union[str, Any]      =		model.generate(**lowerCamelCase_,					eos_token_id=lowerCamelCase_,					**lowerCamelCase_			)
			self.assertTrue(expectation == len(generated_tokens[0]			)			)

			lowercase__					:      Union[str, Any]      =		[638, 198]
			with tf.device(':/CPU:0'			):
					tf.random.set_seed(0			)
					lowercase__					:      Any      =		model.generate(**lowerCamelCase_,					eos_token_id=lowerCamelCase_,					**lowerCamelCase_			)
			self.assertTrue(expectation == len(generated_tokens[0]			)			)





	def    snake_case__(						self:						Union[str, Any]			):
			# Has PT equivalent: ample use of framework-specific code
			lowercase__					:      Dict      =		AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-bart'			)
			lowercase__					:      Any      =		'Hugging Face is a technology company based in New York and Paris.'
			lowercase__					:      int      =		bart_tokenizer(lowerCamelCase_,					return_tensors='tf'			).input_ids
			lowercase__					:      Tuple      =		TFBartForConditionalGeneration.from_pretrained('hf-internal-testing/tiny-random-bart'			)
			lowercase__					:      Optional[int]      =		bart_model.generate(lowerCamelCase_			).numpy()


			class       __lowerCAmelCase					(       _UpperCamelCase      ):


				'''simple docstring'''





				def    snake_case__(						self:						Dict,					lowerCamelCase_:						Tuple,					lowerCamelCase_:						str=None,					**lowerCamelCase_:						Union[str, Any]			):
						return super().call(lowerCamelCase_,					**lowerCamelCase_			)
			lowercase__					:      int      =		FakeBart.from_pretrained('hf-internal-testing/tiny-random-bart'			)
			lowercase__					:      Any      =		bart_model.generate(lowerCamelCase_,					foo='bar'			).numpy()
			self.assertTrue(np.array_equal(lowerCamelCase_,					lowerCamelCase_			)			)


			class       __lowerCAmelCase					(       bart_model.model.encoder.__class__      ):


				'''simple docstring'''





				def    snake_case__(						self:						str,					lowerCamelCase_:						int,					**lowerCamelCase_:						List[str]			):
						return super().call(lowerCamelCase_,					**lowerCamelCase_			)
			lowercase__					:      Optional[Any]      =		FakeEncoder(bart_model.config,					bart_model.model.shared			)
			lowercase__					:      List[Any]      =		fake_encoder

			# Normal generation still works (the output will be different because the encoder weights are different)
			lowercase__					:      Optional[Any]      =		bart_model.generate(lowerCamelCase_			).numpy()
			with self.assertRaises(lowerCamelCase_			):
					# FakeEncoder.call() accepts **kwargs -> no filtering -> value error due to unexpected input "foo"
					bart_model.generate(lowerCamelCase_,					foo='bar'			)









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)))


