



'''simple docstring'''






import unittest

from transformers import AutoTokenizer, FalconConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
						import torch

						from transformers import (
						    FalconForCausalLM,
						    FalconForQuestionAnswering,
						    FalconForSequenceClassification,
						    FalconForTokenClassification,
						    FalconModel,
						)

class 				SCREAMING_SNAKE_CASE:







	"""simple docstring"""






	def __init__(						self			:     Optional[Any]		,			__snake_case			:     str		,			__snake_case			:     Tuple=3		,			__snake_case			:     Dict=7		,			__snake_case			:     List[str]=True		,			__snake_case			:     List[Any]=True		,			__snake_case			:     Any=False		,			__snake_case			:     Dict=True		,			__snake_case			:     Optional[Any]=99		,			__snake_case			:     List[Any]=32		,			__snake_case			:     List[Any]=5		,			__snake_case			:     int=4		,			__snake_case			:     Optional[Any]=37		,			__snake_case			:     int="gelu"		,			__snake_case			:     Dict=0.1		,			__snake_case			:     Dict=0.1		,			__snake_case			:     Union[str, Any]=512		,			__snake_case			:     Optional[int]=16		,			__snake_case			:     List[Any]=2		,			__snake_case			:     int=0.02		,			__snake_case			:     int=3		,			__snake_case			:     Tuple=4		,			__snake_case			:     Tuple=None		,			)			->		Optional[int]:
								UpperCAmelCase :       str 							=		parent
								UpperCAmelCase :       List[str] 							=		batch_size
								UpperCAmelCase :       str 							=		seq_length
								UpperCAmelCase :       Dict 							=		is_training
								UpperCAmelCase :       Tuple 							=		use_input_mask
								UpperCAmelCase :       Optional[int] 							=		use_token_type_ids
								UpperCAmelCase :       List[Any] 							=		use_labels
								UpperCAmelCase :       Any 							=		vocab_size
								UpperCAmelCase :       Union[str, Any] 							=		hidden_size
								UpperCAmelCase :       Any 							=		num_hidden_layers
								UpperCAmelCase :       int 							=		num_attention_heads
								UpperCAmelCase :       Tuple 							=		intermediate_size
								UpperCAmelCase :       Union[str, Any] 							=		hidden_act
								UpperCAmelCase :       List[str] 							=		hidden_dropout_prob
								UpperCAmelCase :       List[Any] 							=		attention_probs_dropout_prob
								UpperCAmelCase :       List[Any] 							=		max_position_embeddings
								UpperCAmelCase :       Any 							=		type_vocab_size
								UpperCAmelCase :       str 							=		type_sequence_label_size
								UpperCAmelCase :       Any 							=		initializer_range
								UpperCAmelCase :       Any 							=		num_labels
								UpperCAmelCase :       Optional[Any] 							=		num_choices
								UpperCAmelCase :       Any 							=		scope







	def      A	(						self			:     List[Any]		)			->		int:
								UpperCAmelCase :       int 							=		ids_tensor([self.batch_size, self.seq_length]		,			self.vocab_size		)

								UpperCAmelCase :       Optional[Any] 							=		None
								if self.use_input_mask:
															UpperCAmelCase :       Any 							=		random_attention_mask([self.batch_size, self.seq_length]		)

								UpperCAmelCase :       Union[str, Any] 							=		None

								UpperCAmelCase :       Optional[int] 							=		None
								UpperCAmelCase :       List[str] 							=		None
								UpperCAmelCase :       Union[str, Any] 							=		None
								if self.use_labels:
															UpperCAmelCase :       Optional[int] 							=		ids_tensor([self.batch_size]		,			self.type_sequence_label_size		)
															UpperCAmelCase :       Optional[Any] 							=		ids_tensor([self.batch_size, self.seq_length]		,			self.num_labels		)
															UpperCAmelCase :       str 							=		ids_tensor([self.batch_size]		,			self.num_choices		)

								UpperCAmelCase :       str 							=		self.get_config()

								return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







	def      A	(						self			:     Optional[Any]		)			->		List[str]:
								return FalconConfig(
								    vocab_size=self.vocab_size		,			hidden_size=self.hidden_size		,			num_hidden_layers=self.num_hidden_layers		,			num_attention_heads=self.num_attention_heads		,			intermediate_size=self.intermediate_size		,			hidden_act=self.hidden_act		,			hidden_dropout_prob=self.hidden_dropout_prob		,			attention_probs_dropout_prob=self.attention_probs_dropout_prob		,			max_position_embeddings=self.max_position_embeddings		,			type_vocab_size=self.type_vocab_size		,			is_decoder=__snake_case		,			initializer_range=self.initializer_range		,			pad_token_id=1		,			new_decoder_architecture=__snake_case		,			)







	def      A	(						self			:     str		,			__snake_case			:     List[Any]		,			__snake_case			:     int		,			__snake_case			:     Tuple		,			__snake_case			:     int		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     List[str]		,			__snake_case			:     Optional[int]		)			->		Any:
								UpperCAmelCase :       List[Any] 							=		FalconModel(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       List[str] 							=		model(__snake_case		,			attention_mask=__snake_case		)
								UpperCAmelCase :       int 							=		model(__snake_case		)
								self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, self.seq_length, self.hidden_size)		)







	def      A	(						self			:     Union[str, Any]		,			__snake_case			:     Tuple		,			__snake_case			:     List[Any]		,			__snake_case			:     Any		,			__snake_case			:     int		,			__snake_case			:     Dict		,			__snake_case			:     Any		,			__snake_case			:     Any		,			__snake_case			:     Optional[Any]		,			__snake_case			:     Optional[int]		,			)			->		List[str]:
								UpperCAmelCase :       str 							=		True
								UpperCAmelCase :       Union[str, Any] 							=		FalconModel(__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       List[str] 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			encoder_hidden_states=__snake_case		,			encoder_attention_mask=__snake_case		,			)
								UpperCAmelCase :       int 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			encoder_hidden_states=__snake_case		,			)
								UpperCAmelCase :       Any 							=		model(__snake_case		,			attention_mask=__snake_case		)
								self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, self.seq_length, self.hidden_size)		)







	def      A	(						self			:     Optional[Any]		,			__snake_case			:     str		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     Tuple		,			__snake_case			:     Tuple		,			__snake_case			:     Tuple		,			__snake_case			:     int		,			__snake_case			:     str		,			__snake_case			:     str		,			__snake_case			:     Optional[int]		,			)			->		Optional[int]:
								UpperCAmelCase :       List[Any] 							=		FalconForCausalLM(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       int 							=		model(__snake_case		,			attention_mask=__snake_case		,			labels=__snake_case		)
								self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.seq_length, self.vocab_size)		)







	def      A	(						self			:     Union[str, Any]		,			__snake_case			:     Tuple		,			__snake_case			:     Tuple		,			__snake_case			:     Tuple		,			__snake_case			:     List[Any]		,			__snake_case			:     Optional[Any]		,			__snake_case			:     str		,			__snake_case			:     Optional[Any]		,			__snake_case			:     str		,			__snake_case			:     str		,			)			->		int:
								UpperCAmelCase :       List[Any] 							=		True
								UpperCAmelCase :       Union[str, Any] 							=		True
								UpperCAmelCase :       Tuple 							=		FalconForCausalLM(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()

								# first forward pass
								UpperCAmelCase :       List[str] 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			encoder_hidden_states=__snake_case		,			encoder_attention_mask=__snake_case		,			use_cache=__snake_case		,			)
								UpperCAmelCase :       str 							=		outputs.past_key_values

								# create hypothetical multiple next token and extent to next_input_ids
								UpperCAmelCase :       str 							=		ids_tensor((self.batch_size, 3)		,			config.vocab_size		)
								UpperCAmelCase :       int 							=		ids_tensor((self.batch_size, 3)		,			vocab_size=2		)

								# append to next input_ids and
								UpperCAmelCase :       Tuple 							=		torch.cat([input_ids, next_tokens]		,			dim=-1		)
								UpperCAmelCase :       str 							=		torch.cat([input_mask, next_mask]		,			dim=-1		)

								UpperCAmelCase :       Dict 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			encoder_hidden_states=__snake_case		,			encoder_attention_mask=__snake_case		,			output_hidden_states=__snake_case		,			)['''hidden_states'''][0]
								UpperCAmelCase :       Tuple 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			encoder_hidden_states=__snake_case		,			encoder_attention_mask=__snake_case		,			past_key_values=__snake_case		,			output_hidden_states=__snake_case		,			)['''hidden_states'''][0]

								# select random slice
								UpperCAmelCase :       int 							=		ids_tensor((1,)		,			output_from_past.shape[-1]		).item()
								UpperCAmelCase :       List[Any] 							=		output_from_no_past[:, -3:, random_slice_idx].detach()
								UpperCAmelCase :       List[Any] 							=		output_from_past[:, :, random_slice_idx].detach()

								self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]		)

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







	def      A	(						self			:     int		)			->		Tuple:
								UpperCAmelCase :       Tuple 							=		self.prepare_config_and_inputs()
								(
								    (
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		
								) :       Optional[Any] 							=		config_and_inputs
								UpperCAmelCase :       Tuple 							=		{'''input_ids''': input_ids, '''attention_mask''': input_mask}
								return config, inputs_dict

@require_torch
class 				SCREAMING_SNAKE_CASE(							A__							,		A__							,		A__							,		unittest.TestCase			):







	"""simple docstring"""






	lowerCamelCase__			 =       (
	    (
	        FalconModel,
	        FalconForCausalLM,
	        FalconForSequenceClassification,
	        FalconForTokenClassification,
	        FalconForQuestionAnswering,
	    )
	    if is_torch_available()
	    else ()
	)
	lowerCamelCase__			 =       (FalconForCausalLM,) if is_torch_available() else ()
	lowerCamelCase__			 =       (
	    {
	        """feature-extraction""": FalconModel,
	        """text-classification""": FalconForSequenceClassification,
	        """text-generation""": FalconForCausalLM,
	        """question-answering""": FalconForQuestionAnswering,
	        """token-classification""": FalconForTokenClassification,
	        """zero-shot""": FalconForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	lowerCamelCase__			 =       False
	lowerCamelCase__			 =       False







	def      A	(						self			:     str		)			->		Optional[Any]:
								UpperCAmelCase :       Dict 							=		FalconModelTester(self		)
								UpperCAmelCase :       Any 							=		ConfigTester(self		,			config_class=__snake_case		,			hidden_size=37		)







	def      A	(						self			:     int		)			->		Any:
								self.config_tester.run_common_tests()







	def      A	(						self			:     List[str]		)			->		Union[str, Any]:
								UpperCAmelCase :       List[Any] 							=		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*__snake_case		)







	def      A	(						self			:     List[str]		)			->		List[str]:
								UpperCAmelCase							,		*UpperCAmelCase :       str 							=		self.model_tester.prepare_config_and_inputs()
								for alibi in [True, False]:
															UpperCAmelCase :       Optional[Any] 							=		alibi
															self.model_tester.create_and_check_model(__snake_case		,			*__snake_case		)







	def      A	(						self			:     int		)			->		Dict:
								UpperCAmelCase							,		UpperCAmelCase :       List[Any] 							=		self.model_tester.prepare_config_and_inputs_for_common()
								UpperCAmelCase :       Tuple 							=		3
								UpperCAmelCase :       Union[str, Any] 							=		input_dict['''input_ids''']
								UpperCAmelCase :       Any 							=		input_ids.ne(1		).to(__snake_case		)
								UpperCAmelCase :       Optional[int] 							=		ids_tensor([self.model_tester.batch_size]		,			self.model_tester.type_sequence_label_size		)
								UpperCAmelCase :       Tuple 							=		FalconForSequenceClassification(__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       Optional[Any] 							=		model(__snake_case		,			attention_mask=__snake_case		,			labels=__snake_case		)
								self.assertEqual(result.logits.shape		,			(self.model_tester.batch_size, self.model_tester.num_labels)		)







	def      A	(						self			:     Union[str, Any]		)			->		List[Any]:
								UpperCAmelCase							,		UpperCAmelCase :       List[Any] 							=		self.model_tester.prepare_config_and_inputs_for_common()
								UpperCAmelCase :       Union[str, Any] 							=		3
								UpperCAmelCase :       Tuple 							=		'''single_label_classification'''
								UpperCAmelCase :       Union[str, Any] 							=		input_dict['''input_ids''']
								UpperCAmelCase :       Dict 							=		input_ids.ne(1		).to(__snake_case		)
								UpperCAmelCase :       int 							=		ids_tensor([self.model_tester.batch_size]		,			self.model_tester.type_sequence_label_size		)
								UpperCAmelCase :       Tuple 							=		FalconForSequenceClassification(__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       List[str] 							=		model(__snake_case		,			attention_mask=__snake_case		,			labels=__snake_case		)
								self.assertEqual(result.logits.shape		,			(self.model_tester.batch_size, self.model_tester.num_labels)		)







	def      A	(						self			:     Optional[Any]		)			->		Dict:
								UpperCAmelCase							,		UpperCAmelCase :       Optional[Any] 							=		self.model_tester.prepare_config_and_inputs_for_common()
								UpperCAmelCase :       List[Any] 							=		input_dict['''input_ids''']
								UpperCAmelCase :       Tuple 							=		FalconForCausalLM(__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       Dict 							=		model(__snake_case		,			use_cache=__snake_case		)
								UpperCAmelCase :       Tuple 							=		input_ids.shape[0]
								UpperCAmelCase :       Any 							=		model._convert_to_rw_cache(result.past_key_values		)
								UpperCAmelCase :       Any 							=		model._convert_cache_to_standard_format(__snake_case		,			__snake_case		)
								for layer in range(len(__snake_case		)		):
															for tensor_idx in range(2		):
																						self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3		)
																						self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4		)
																						self.assertTrue(
																						    torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx]		)		)







	def      A	(						self			:     Optional[Any]		)			->		Optional[int]:
								UpperCAmelCase							,		UpperCAmelCase :       Dict 							=		self.model_tester.prepare_config_and_inputs_for_common()
								UpperCAmelCase :       Union[str, Any] 							=		3
								UpperCAmelCase :       List[Any] 							=		'''multi_label_classification'''
								UpperCAmelCase :       Tuple 							=		input_dict['''input_ids''']
								UpperCAmelCase :       List[Any] 							=		input_ids.ne(1		).to(__snake_case		)
								UpperCAmelCase :       Optional[Any] 							=		ids_tensor(
								    [self.model_tester.batch_size, config.num_labels]		,			self.model_tester.type_sequence_label_size		).to(torch.float		)
								UpperCAmelCase :       str 							=		FalconForSequenceClassification(__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       Dict 							=		model(__snake_case		,			attention_mask=__snake_case		,			labels=__snake_case		)
								self.assertEqual(result.logits.shape		,			(self.model_tester.batch_size, self.model_tester.num_labels)		)







	def      A	(						self			:     List[str]		)			->		Tuple:
								# Falcon can have different numbers of KV-heads than the number of query heads, so we need
								# to override this test to use the right head counts.
								for model_class in self.all_generative_model_classes:
															UpperCAmelCase							,		UpperCAmelCase :       List[str] 							=		self.model_tester.prepare_config_and_inputs_for_common()

															# If it doesn't support cache, pass the test
															if not hasattr(__snake_case		,			'''use_cache'''		):
																						return

															UpperCAmelCase :       List[str] 							=		model_class(__snake_case		).to(__snake_case		)
															if "use_cache" not in inputs:
																						UpperCAmelCase :       Optional[Any] 							=		True
															UpperCAmelCase :       Optional[int] 							=		model(**__snake_case		)

															# If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
															if "past_key_values" not in outputs:
																						return

															UpperCAmelCase :       List[Any] 							=		(
															    getattr(__snake_case		,			'''decoder_layers'''		,			__snake_case		)
															    or getattr(__snake_case		,			'''num_decoder_layers'''		,			__snake_case		)
															    or config.num_hidden_layers
															)
															UpperCAmelCase :       Any 							=		getattr(__snake_case		,			'''num_kv_heads'''		,			config.num_attention_heads		)
															UpperCAmelCase :       Optional[Any] 							=		getattr(__snake_case		,			'''d_model'''		,			config.hidden_size		)
															UpperCAmelCase :       Union[str, Any] 							=		embed_dim // num_attention_heads

															UpperCAmelCase :       List[str] 							=		outputs['''past_key_values''']
															self.assertEqual(len(__snake_case		)		,			__snake_case		)

															UpperCAmelCase							,		UpperCAmelCase :       List[Any] 							=		inputs['''input_ids'''].shape
															for i in range(__snake_case		):
																						if config.new_decoder_architecture:
																													UpperCAmelCase :       Tuple 							=		config.num_attention_heads
																						elif config.multi_query:
																													UpperCAmelCase :       List[Any] 							=		1
																						self.assertEqual(len(past_kv[0]		)		,			2		)  # K V for the decoder = 2
																						self.assertEqual(
																						    past_kv[i][0].shape		,			(batch_size, num_attention_heads, seq_length, per_head_embed_dim)		)
																						self.assertEqual(
																						    past_kv[i][1].shape		,			(batch_size, num_attention_heads, seq_length, per_head_embed_dim)		)

@require_torch
class 				SCREAMING_SNAKE_CASE(							unittest.TestCase			):







	"""simple docstring"""






	@slow
	def      A	(						self			:     Any		)			->		Tuple:
								UpperCAmelCase :       Dict 							=		AutoTokenizer.from_pretrained('''Rocketknight1/falcon-rw-1b'''		)
								UpperCAmelCase :       List[str] 							=		FalconForCausalLM.from_pretrained('''Rocketknight1/falcon-rw-1b'''		)
								model.eval()
								model.to(__snake_case		)
								UpperCAmelCase :       int 							=		tokenizer('''My favorite food is'''		,			return_tensors='''pt'''		).to(__snake_case		)

								UpperCAmelCase :       List[Any] 							=		(
								    '''My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday.'''
								)

								UpperCAmelCase :       List[str] 							=		model.generate(**__snake_case		,			do_sample=__snake_case		,			max_new_tokens=19		)
								UpperCAmelCase :       str 							=		tokenizer.batch_decode(__snake_case		)[0]

								self.assertEqual(__snake_case		,			__snake_case		)







	@slow
	def      A	(						self			:     Tuple		)			->		List[Any]:
								# The big models are way too big for the CI, so we use tiny random models that resemble their
								# architectures but with much smaller and fewer layers
								for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
															UpperCAmelCase :       Tuple 							=		AutoTokenizer.from_pretrained(__snake_case		)
															UpperCAmelCase :       List[Any] 							=		FalconForCausalLM.from_pretrained(__snake_case		)
															model.eval()
															model.to(__snake_case		)
															UpperCAmelCase :       List[str] 							=		tokenizer('''My favorite food is'''		,			return_tensors='''pt'''		).to(__snake_case		)

															# We just test that these run without errors - the models are randomly initialized
															# and so the actual text outputs will be garbage
															model.generate(**__snake_case		,			do_sample=__snake_case		,			max_new_tokens=4		)
															model.generate(**__snake_case		,			do_sample=__snake_case		,			max_new_tokens=4		)
															model.generate(**__snake_case		,			num_beams=2		,			max_new_tokens=4		)







	@slow
	def      A	(						self			:     str		)			->		Optional[int]:
								# The big models are way too big for the CI, so we use tiny random models that resemble their
								# architectures but with much smaller and fewer layers
								with torch.no_grad():
															for repo in [
															    "Rocketknight1/falcon-rw-1b",
															    "Rocketknight1/tiny-random-falcon-7b",
															    "Rocketknight1/tiny-random-falcon-40b",
															]:
																						UpperCAmelCase :       List[str] 							=		AutoTokenizer.from_pretrained(__snake_case		)
																						UpperCAmelCase :       Union[str, Any] 							=		FalconForCausalLM.from_pretrained(__snake_case		)
																						model.eval()
																						model.to(device=__snake_case		)
																						UpperCAmelCase :       int 							=		tokenizer('''My favorite food is'''		,			return_tensors='''pt'''		).to(__snake_case		)

																						# Test results are the same with and without cache
																						UpperCAmelCase :       Dict 							=		model.generate(**__snake_case		,			do_sample=__snake_case		,			max_new_tokens=20		,			use_cache=__snake_case		)
																						UpperCAmelCase :       Tuple 							=		model.generate(**__snake_case		,			do_sample=__snake_case		,			max_new_tokens=20		,			use_cache=__snake_case		)
																						self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0		)










'''simple docstring'''






import os
import tempfile
import unittest

import numpy as np

from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax, slow


if is_flax_available():
						import jax
						import jax.numpy as jnp
						from flax.jax_utils import replicate
						from flax.training.common_utils import shard

						from diffusers import FlaxDDIMScheduler, FlaxDiffusionPipeline, FlaxStableDiffusionPipeline

@require_flax
class 				SCREAMING_SNAKE_CASE(							unittest.TestCase			):







	"""simple docstring"""






	def      A	(						self			:     str		)			->		int:
								with tempfile.TemporaryDirectory() as tmpdirname:
															# pipeline has Flax weights
															UpperCAmelCase :       Optional[Any] 							=		FlaxDiffusionPipeline.from_pretrained(
															    '''hf-internal-testing/tiny-stable-diffusion-pipe'''		,			safety_checker=__snake_case		,			cache_dir=__snake_case		)

															UpperCAmelCase :       str 							=		[t[-1] for t in os.walk(os.path.join(__snake_case		,			os.listdir(__snake_case		)[0]		,			'''snapshots'''		)		)]
															UpperCAmelCase :       str 							=		[item for sublist in all_root_files for item in sublist]

															# None of the downloaded files should be a PyTorch file even if we have some here:
															# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_pytorch_model.bin
															assert not any(f.endswith('''.bin'''		) for f in files		)

@slow
@require_flax
class 				SCREAMING_SNAKE_CASE(							unittest.TestCase			):







	"""simple docstring"""






	def      A	(						self			:     List[str]		)			->		Dict:
								UpperCAmelCase							,		UpperCAmelCase :       str 							=		FlaxStableDiffusionPipeline.from_pretrained(
								    '''hf-internal-testing/tiny-stable-diffusion-pipe'''		,			safety_checker=__snake_case		)

								UpperCAmelCase :       List[Any] 							=		(
								    '''A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'''
								    ''' field, close up, split lighting, cinematic'''
								)

								UpperCAmelCase :       List[str] 							=		jax.random.PRNGKey(0		)
								UpperCAmelCase :       Optional[Any] 							=		4

								UpperCAmelCase :       Optional[Any] 							=		jax.device_count()
								UpperCAmelCase :       Tuple 							=		num_samples * [prompt]
								UpperCAmelCase :       int 							=		pipeline.prepare_inputs(__snake_case		)

								# shard inputs and rng
								UpperCAmelCase :       Tuple 							=		replicate(__snake_case		)
								UpperCAmelCase :       Any 							=		jax.random.split(__snake_case		,			__snake_case		)
								UpperCAmelCase :       Optional[Any] 							=		shard(__snake_case		)

								UpperCAmelCase :       Optional[int] 							=		pipeline(__snake_case		,			__snake_case		,			__snake_case		,			__snake_case		,			jit=__snake_case		).images

								assert images.shape == (num_samples, 1, 64, 64, 3)
								if jax.device_count() == 8:
															assert np.abs(np.abs(images[0, 0, :2, :2, -2:]		,			dtype=np.floataa		).sum() - 4.1_51_47_45		) < 1E-3
															assert np.abs(np.abs(__snake_case		,			dtype=np.floataa		).sum() - 4_99_47.8_75		) < 5E-1

								UpperCAmelCase :       Union[str, Any] 							=		pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]		)		)		)
								assert len(__snake_case		) == num_samples







	def      A	(						self			:     List[Any]		)			->		List[str]:
								UpperCAmelCase							,		UpperCAmelCase :       Union[str, Any] 							=		FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4'''		,			revision='''flax'''		,			safety_checker=__snake_case		)

								UpperCAmelCase :       Dict 							=		(
								    '''A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'''
								    ''' field, close up, split lighting, cinematic'''
								)

								UpperCAmelCase :       Optional[Any] 							=		jax.random.PRNGKey(0		)
								UpperCAmelCase :       Any 							=		50

								UpperCAmelCase :       Union[str, Any] 							=		jax.device_count()
								UpperCAmelCase :       int 							=		num_samples * [prompt]
								UpperCAmelCase :       Union[str, Any] 							=		pipeline.prepare_inputs(__snake_case		)

								# shard inputs and rng
								UpperCAmelCase :       Dict 							=		replicate(__snake_case		)
								UpperCAmelCase :       int 							=		jax.random.split(__snake_case		,			__snake_case		)
								UpperCAmelCase :       Tuple 							=		shard(__snake_case		)

								UpperCAmelCase :       Tuple 							=		pipeline(__snake_case		,			__snake_case		,			__snake_case		,			__snake_case		,			jit=__snake_case		).images

								assert images.shape == (num_samples, 1, 512, 512, 3)
								if jax.device_count() == 8:
															assert np.abs((np.abs(images[0, 0, :2, :2, -2:]		,			dtype=np.floataa		).sum() - 0.05_65_24_01)		) < 1E-3
															assert np.abs((np.abs(__snake_case		,			dtype=np.floataa		).sum() - 2_38_38_08.2)		) < 5E-1







	def      A	(						self			:     int		)			->		Dict:
								UpperCAmelCase							,		UpperCAmelCase :       List[str] 							=		FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4'''		,			revision='''bf16'''		,			dtype=jnp.bfloataa		,			safety_checker=__snake_case		)

								UpperCAmelCase :       Dict 							=		(
								    '''A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'''
								    ''' field, close up, split lighting, cinematic'''
								)

								UpperCAmelCase :       Union[str, Any] 							=		jax.random.PRNGKey(0		)
								UpperCAmelCase :       List[str] 							=		50

								UpperCAmelCase :       Union[str, Any] 							=		jax.device_count()
								UpperCAmelCase :       List[Any] 							=		num_samples * [prompt]
								UpperCAmelCase :       int 							=		pipeline.prepare_inputs(__snake_case		)

								# shard inputs and rng
								UpperCAmelCase :       Tuple 							=		replicate(__snake_case		)
								UpperCAmelCase :       List[Any] 							=		jax.random.split(__snake_case		,			__snake_case		)
								UpperCAmelCase :       Optional[int] 							=		shard(__snake_case		)

								UpperCAmelCase :       Optional[Any] 							=		pipeline(__snake_case		,			__snake_case		,			__snake_case		,			__snake_case		,			jit=__snake_case		).images

								assert images.shape == (num_samples, 1, 512, 512, 3)
								if jax.device_count() == 8:
															assert np.abs((np.abs(images[0, 0, :2, :2, -2:]		,			dtype=np.floataa		).sum() - 0.04_00_39_06)		) < 1E-3
															assert np.abs((np.abs(__snake_case		,			dtype=np.floataa		).sum() - 2_37_35_16.75)		) < 5E-1







	def      A	(						self			:     int		)			->		Any:
								UpperCAmelCase							,		UpperCAmelCase :       Dict 							=		FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4'''		,			revision='''bf16'''		,			dtype=jnp.bfloataa		)

								UpperCAmelCase :       List[str] 							=		(
								    '''A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'''
								    ''' field, close up, split lighting, cinematic'''
								)

								UpperCAmelCase :       List[str] 							=		jax.random.PRNGKey(0		)
								UpperCAmelCase :       Union[str, Any] 							=		50

								UpperCAmelCase :       Optional[int] 							=		jax.device_count()
								UpperCAmelCase :       List[str] 							=		num_samples * [prompt]
								UpperCAmelCase :       Dict 							=		pipeline.prepare_inputs(__snake_case		)

								# shard inputs and rng
								UpperCAmelCase :       Tuple 							=		replicate(__snake_case		)
								UpperCAmelCase :       Any 							=		jax.random.split(__snake_case		,			__snake_case		)
								UpperCAmelCase :       str 							=		shard(__snake_case		)

								UpperCAmelCase :       Optional[int] 							=		pipeline(__snake_case		,			__snake_case		,			__snake_case		,			__snake_case		,			jit=__snake_case		).images

								assert images.shape == (num_samples, 1, 512, 512, 3)
								if jax.device_count() == 8:
															assert np.abs((np.abs(images[0, 0, :2, :2, -2:]		,			dtype=np.floataa		).sum() - 0.04_00_39_06)		) < 1E-3
															assert np.abs((np.abs(__snake_case		,			dtype=np.floataa		).sum() - 2_37_35_16.75)		) < 5E-1







	def      A	(						self			:     Tuple		)			->		Optional[Any]:
								UpperCAmelCase :       int 							=		FlaxDDIMScheduler(
								    beta_start=0.0_00_85		,			beta_end=0.0_12		,			beta_schedule='''scaled_linear'''		,			set_alpha_to_one=__snake_case		,			steps_offset=1		,			)

								UpperCAmelCase							,		UpperCAmelCase :       Tuple 							=		FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4'''		,			revision='''bf16'''		,			dtype=jnp.bfloataa		,			scheduler=__snake_case		,			safety_checker=__snake_case		,			)
								UpperCAmelCase :       Tuple 							=		scheduler.create_state()

								UpperCAmelCase :       Dict 							=		scheduler_state

								UpperCAmelCase :       str 							=		(
								    '''A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'''
								    ''' field, close up, split lighting, cinematic'''
								)

								UpperCAmelCase :       int 							=		jax.random.PRNGKey(0		)
								UpperCAmelCase :       Union[str, Any] 							=		50

								UpperCAmelCase :       Optional[Any] 							=		jax.device_count()
								UpperCAmelCase :       Any 							=		num_samples * [prompt]
								UpperCAmelCase :       Dict 							=		pipeline.prepare_inputs(__snake_case		)

								# shard inputs and rng
								UpperCAmelCase :       str 							=		replicate(__snake_case		)
								UpperCAmelCase :       List[str] 							=		jax.random.split(__snake_case		,			__snake_case		)
								UpperCAmelCase :       Optional[int] 							=		shard(__snake_case		)

								UpperCAmelCase :       Dict 							=		pipeline(__snake_case		,			__snake_case		,			__snake_case		,			__snake_case		,			jit=__snake_case		).images

								assert images.shape == (num_samples, 1, 512, 512, 3)
								if jax.device_count() == 8:
															assert np.abs((np.abs(images[0, 0, :2, :2, -2:]		,			dtype=np.floataa		).sum() - 0.0_45_04_39_45)		) < 1E-3
															assert np.abs((np.abs(__snake_case		,			dtype=np.floataa		).sum() - 2_34_76_93.5)		) < 5E-1







	def      A	(						self			:     Any		)			->		Tuple:
								UpperCAmelCase :       List[Any] 							=		(
								    '''A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'''
								    ''' field, close up, split lighting, cinematic'''
								)

								UpperCAmelCase :       Union[str, Any] 							=		jax.device_count()
								UpperCAmelCase :       List[Any] 							=		num_samples * [prompt]
								UpperCAmelCase :       str 							=		jax.random.split(jax.random.PRNGKey(0		)		,			__snake_case		)

								UpperCAmelCase							,		UpperCAmelCase :       Tuple 							=		FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4'''		,			revision='''bf16'''		,			dtype=jnp.bfloataa		,			safety_checker=__snake_case		,			)

								UpperCAmelCase :       Dict 							=		replicate(__snake_case		)
								UpperCAmelCase :       Optional[Any] 							=		pipeline.prepare_inputs(__snake_case		)
								UpperCAmelCase :       List[str] 							=		shard(__snake_case		)
								UpperCAmelCase :       Any 							=		pipeline(__snake_case		,			__snake_case		,			__snake_case		,			jit=__snake_case		).images
								assert images.shape == (num_samples, 1, 512, 512, 3)
								UpperCAmelCase :       Optional[int] 							=		images[2, 0, 256, 10:17, 1]

								# With memory efficient attention
								UpperCAmelCase							,		UpperCAmelCase :       Any 							=		FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4'''		,			revision='''bf16'''		,			dtype=jnp.bfloataa		,			safety_checker=__snake_case		,			use_memory_efficient_attention=__snake_case		,			)

								UpperCAmelCase :       int 							=		replicate(__snake_case		)
								UpperCAmelCase :       int 							=		pipeline.prepare_inputs(__snake_case		)
								UpperCAmelCase :       List[Any] 							=		shard(__snake_case		)
								UpperCAmelCase :       Optional[Any] 							=		pipeline(__snake_case		,			__snake_case		,			__snake_case		,			jit=__snake_case		).images
								assert images_eff.shape == (num_samples, 1, 512, 512, 3)
								UpperCAmelCase :       int 							=		images[2, 0, 256, 10:17, 1]

								# I checked the results visually and they are very similar. However, I saw that the max diff is `1` and the `sum`
								# over the 8 images is exactly `256`, which is very suspicious. Testing a random slice for now.
								assert abs(slice_eff - slice		).max() < 1E-2





