

"""simple docstring"""





import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference



class    _lowerCAmelCase    (				SCREAMING_SNAKE_CASE__  ,				unittest.TestCase		):





		"""simple docstring"""






		_lowerCamelCase			  =       ShapEPipeline
		_lowerCamelCase			  =       ['''prompt''']
		_lowerCamelCase			  =       ['''prompt''']
		_lowerCamelCase			  =       [
		    '''num_images_per_prompt''',
		    '''num_inference_steps''',
		    '''generator''',
		    '''latents''',
		    '''guidance_scale''',
		    '''frame_size''',
		    '''output_type''',
		    '''return_dict''',
		]
		_lowerCamelCase			  =       False





		@property
		def    UpperCAmelCase__							(					self   ) ->						int:







								'''simple docstring'''

								return 3_2





		@property
		def    UpperCAmelCase__							(					self   ) ->						Union[str, Any]:







								'''simple docstring'''

								return 3_2





		@property
		def    UpperCAmelCase__							(					self   ) ->						Any:







								'''simple docstring'''

								return self.time_input_dim * 4





		@property
		def    UpperCAmelCase__							(					self   ) ->						Optional[int]:







								'''simple docstring'''

								return 8





		@property
		def    UpperCAmelCase__							(					self   ) ->						List[str]:







								'''simple docstring'''

								snake_case_   : int						=			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""   )
								return tokenizer





		@property
		def    UpperCAmelCase__							(					self   ) ->						List[str]:







								'''simple docstring'''

								torch.manual_seed(0   )
								snake_case_   : Dict						=			CLIPTextConfig(
								    bos_token_id=0							,    eos_token_id=2							,    hidden_size=self.text_embedder_hidden_size							,    projection_dim=self.text_embedder_hidden_size							,    intermediate_size=3_7							,    layer_norm_eps=1E-05							,    num_attention_heads=4							,    num_hidden_layers=5							,    pad_token_id=1							,    vocab_size=1_0_0_0							,    )
								return CLIPTextModelWithProjection(_lowercase   )





		@property
		def    UpperCAmelCase__							(					self   ) ->						List[Any]:







								'''simple docstring'''

								torch.manual_seed(0   )

								snake_case_   : int						=			{
								    """num_attention_heads""": 2,
								    """attention_head_dim""": 1_6,
								    """embedding_dim""": self.time_input_dim,
								    """num_embeddings""": 3_2,
								    """embedding_proj_dim""": self.text_embedder_hidden_size,
								    """time_embed_dim""": self.time_embed_dim,
								    """num_layers""": 1,
								    """clip_embed_dim""": self.time_input_dim * 2,
								    """additional_embeddings""": 0,
								    """time_embed_act_fn""": """gelu""",
								    """norm_in_type""": """layer""",
								    """encoder_hid_proj_type""": None,
								    """added_emb_type""": None,
								}

								snake_case_   : Any						=			PriorTransformer(**_lowercase   )
								return model





		@property
		def    UpperCAmelCase__							(					self   ) ->						Any:







								'''simple docstring'''

								torch.manual_seed(0   )

								snake_case_   : Union[str, Any]						=			{
								    """param_shapes""": (
								        (self.renderer_dim, 9_3),
								        (self.renderer_dim, 8),
								        (self.renderer_dim, 8),
								        (self.renderer_dim, 8),
								    ),
								    """d_latent""": self.time_input_dim,
								    """d_hidden""": self.renderer_dim,
								    """n_output""": 1_2,
								    """background""": (
								        0.1,
								        0.1,
								        0.1,
								    ),
								}
								snake_case_   : str						=			ShapERenderer(**_lowercase   )
								return model





		def    UpperCAmelCase__							(					self   ) ->						List[Any]:







								'''simple docstring'''

								snake_case_   : str						=			self.dummy_prior
								snake_case_   : List[str]						=			self.dummy_text_encoder
								snake_case_   : Optional[Any]						=			self.dummy_tokenizer
								snake_case_   : Optional[Any]						=			self.dummy_renderer

								snake_case_   : Any						=			HeunDiscreteScheduler(
								    beta_schedule="""exp"""							,    num_train_timesteps=1_0_2_4							,    prediction_type="""sample"""							,    use_karras_sigmas=_lowercase							,    clip_sample=_lowercase							,    clip_sample_range=1.0							,    )
								snake_case_   : Union[str, Any]						=			{
								    """prior""": prior,
								    """text_encoder""": text_encoder,
								    """tokenizer""": tokenizer,
								    """renderer""": renderer,
								    """scheduler""": scheduler,
								}

								return components





		def    UpperCAmelCase__							(					self							,    _lowercase							,    _lowercase=0   ) ->						int:







								'''simple docstring'''

								if str(_lowercase   ).startswith("""mps"""   ):
														snake_case_   : Optional[Any]						=			torch.manual_seed(_lowercase   )
								else:
														snake_case_   : Tuple						=			torch.Generator(device=_lowercase   ).manual_seed(_lowercase   )
								snake_case_   : Optional[Any]						=			{
								    """prompt""": """horse""",
								    """generator""": generator,
								    """num_inference_steps""": 1,
								    """frame_size""": 3_2,
								    """output_type""": """np""",
								}
								return inputs





		def    UpperCAmelCase__							(					self   ) ->						int:







								'''simple docstring'''

								snake_case_   : Any						=			"""cpu"""

								snake_case_   : List[str]						=			self.get_dummy_components()

								snake_case_   : int						=			self.pipeline_class(**_lowercase   )
								snake_case_   : List[str]						=			pipe.to(_lowercase   )

								pipe.set_progress_bar_config(disable=_lowercase   )

								snake_case_   : List[Any]						=			pipe(**self.get_dummy_inputs(_lowercase   )   )
								snake_case_   : List[Any]						=			output.images[0]
								snake_case_   : Union[str, Any]						=			image[0, -3:, -3:, -1]

								assert image.shape == (2_0, 3_2, 3_2, 3)

								snake_case_   : Any						=			np.array(
								    [
								        0.0003_9216,
								        0.0003_9216,
								        0.0003_9216,
								        0.0003_9216,
								        0.0003_9216,
								        0.0003_9216,
								        0.0003_9216,
								        0.0003_9216,
								        0.0003_9216,
								    ]   )

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





		def    UpperCAmelCase__							(					self   ) ->						List[str]:







								'''simple docstring'''

								self._test_inference_batch_consistent(batch_sizes=[1, 2]   )





		def    UpperCAmelCase__							(					self   ) ->						Union[str, Any]:







								'''simple docstring'''

								snake_case_   : Any						=			torch_device == """cpu"""
								snake_case_   : List[Any]						=			True

								self._test_inference_batch_single_identical(
								    batch_size=2							,    test_max_difference=_lowercase							,    relax_max_difference=_lowercase							,    )





		def    UpperCAmelCase__							(					self   ) ->						Dict:







								'''simple docstring'''

								snake_case_   : Any						=			self.get_dummy_components()
								snake_case_   : str						=			self.pipeline_class(**_lowercase   )
								snake_case_   : List[str]						=			pipe.to(_lowercase   )
								pipe.set_progress_bar_config(disable=_lowercase   )

								snake_case_   : Dict						=			1
								snake_case_   : Optional[int]						=			2

								snake_case_   : Tuple						=			self.get_dummy_inputs(_lowercase   )

								for key in inputs.keys():
														if key in self.batch_params:
																				snake_case_   : Dict						=			batch_size * [inputs[key]]

								snake_case_   : Tuple						=			pipe(**_lowercase							,    num_images_per_prompt=_lowercase   )[0]

								assert images.shape[0] == batch_size * num_images_per_prompt




@slow
@require_torch_gpu
class    _lowerCAmelCase    (				unittest.TestCase		):





		"""simple docstring"""






		def    UpperCAmelCase__							(					self   ) ->						Optional[Any]:







								'''simple docstring'''

								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()





		def    UpperCAmelCase__							(					self   ) ->						Dict:







								'''simple docstring'''

								snake_case_   : Tuple						=			load_numpy(
								    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
								    """/shap_e/test_shap_e_np_out.npy"""   )
								snake_case_   : Any						=			ShapEPipeline.from_pretrained("""openai/shap-e"""   )
								snake_case_   : Any						=			pipe.to(_lowercase   )
								pipe.set_progress_bar_config(disable=_lowercase   )

								snake_case_   : Tuple						=			torch.Generator(device=_lowercase   ).manual_seed(0   )

								snake_case_   : int						=			pipe(
								    """a shark"""							,    generator=_lowercase							,    guidance_scale=15.0							,    num_inference_steps=6_4							,    frame_size=6_4							,    output_type="""np"""							,    ).images[0]

								assert images.shape == (2_0, 6_4, 6_4, 3)

								assert_mean_pixel_difference(_lowercase							,    _lowercase   )






"""simple docstring"""





import unittest

import numpy as np

from transformers import RobertaPreLayerNormConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
					import jax.numpy as jnp

					from transformers.models.roberta_prelayernorm.modeling_flax_roberta_prelayernorm import (
					    FlaxRobertaPreLayerNormForCausalLM,
					    FlaxRobertaPreLayerNormForMaskedLM,
					    FlaxRobertaPreLayerNormForMultipleChoice,
					    FlaxRobertaPreLayerNormForQuestionAnswering,
					    FlaxRobertaPreLayerNormForSequenceClassification,
					    FlaxRobertaPreLayerNormForTokenClassification,
					    FlaxRobertaPreLayerNormModel,
					)



class    _lowerCAmelCase    (				unittest.TestCase		):





		"""simple docstring"""






		def __init__(					self							,    _lowercase							,    _lowercase=1_3							,    _lowercase=7							,    _lowercase=True							,    _lowercase=True							,    _lowercase=True							,    _lowercase=True							,    _lowercase=9_9							,    _lowercase=3_2							,    _lowercase=5							,    _lowercase=4							,    _lowercase=3_7							,    _lowercase="gelu"							,    _lowercase=0.1							,    _lowercase=0.1							,    _lowercase=5_1_2							,    _lowercase=1_6							,    _lowercase=2							,    _lowercase=0.02							,    _lowercase=4							,    ) ->						Union[str, Any]:







								'''simple docstring'''

								snake_case_   : Optional[int]						=			parent
								snake_case_   : Dict						=			batch_size
								snake_case_   : Any						=			seq_length
								snake_case_   : Tuple						=			is_training
								snake_case_   : Dict						=			use_attention_mask
								snake_case_   : int						=			use_token_type_ids
								snake_case_   : List[Any]						=			use_labels
								snake_case_   : List[str]						=			vocab_size
								snake_case_   : str						=			hidden_size
								snake_case_   : Optional[Any]						=			num_hidden_layers
								snake_case_   : List[Any]						=			num_attention_heads
								snake_case_   : Any						=			intermediate_size
								snake_case_   : Optional[Any]						=			hidden_act
								snake_case_   : Tuple						=			hidden_dropout_prob
								snake_case_   : List[Any]						=			attention_probs_dropout_prob
								snake_case_   : int						=			max_position_embeddings
								snake_case_   : Dict						=			type_vocab_size
								snake_case_   : Dict						=			type_sequence_label_size
								snake_case_   : Optional[Any]						=			initializer_range
								snake_case_   : Tuple						=			num_choices





		def    UpperCAmelCase__							(					self   ) ->						Union[str, Any]:







								'''simple docstring'''

								snake_case_   : int						=			ids_tensor([self.batch_size, self.seq_length]							,    self.vocab_size   )

								snake_case_   : Tuple						=			None
								if self.use_attention_mask:
														snake_case_   : Optional[Any]						=			random_attention_mask([self.batch_size, self.seq_length]   )

								snake_case_   : List[str]						=			None
								if self.use_token_type_ids:
														snake_case_   : int						=			ids_tensor([self.batch_size, self.seq_length]							,    self.type_vocab_size   )

								snake_case_   : Tuple						=			RobertaPreLayerNormConfig(
								    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=_lowercase							,    initializer_range=self.initializer_range							,    )

								return config, input_ids, token_type_ids, attention_mask





		def    UpperCAmelCase__							(					self   ) ->						Any:







								'''simple docstring'''

								snake_case_   : Any						=			self.prepare_config_and_inputs()
								snake_case_			,		snake_case_			,		snake_case_			,		snake_case_   : List[Any]						=			config_and_inputs
								snake_case_   : int						=			{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
								return config, inputs_dict





		def    UpperCAmelCase__							(					self   ) ->						str:







								'''simple docstring'''

								snake_case_   : Tuple						=			self.prepare_config_and_inputs()
								snake_case_			,		snake_case_			,		snake_case_			,		snake_case_   : Any						=			config_and_inputs

								snake_case_   : Union[str, Any]						=			True
								snake_case_   : List[Any]						=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
								snake_case_   : Optional[int]						=			ids_tensor([self.batch_size, self.seq_length]							,    vocab_size=2   )

								return (
								    config,
								    input_ids,
								    token_type_ids,
								    encoder_hidden_states,
								    encoder_attention_mask,
								)




@require_flax
# Copied from tests.models.roberta.test_modelling_flax_roberta.FlaxRobertaPreLayerNormModelTest with ROBERTA->ROBERTA_PRELAYERNORM,Roberta->RobertaPreLayerNorm,roberta-base->andreasmadsen/efficient_mlm_m0.40
class    _lowerCAmelCase    (				SCREAMING_SNAKE_CASE__  ,				unittest.TestCase		):





		"""simple docstring"""






		_lowerCamelCase			  =       True

		_lowerCamelCase			  =       (
		    (
		        FlaxRobertaPreLayerNormModel,
		        FlaxRobertaPreLayerNormForCausalLM,
		        FlaxRobertaPreLayerNormForMaskedLM,
		        FlaxRobertaPreLayerNormForSequenceClassification,
		        FlaxRobertaPreLayerNormForTokenClassification,
		        FlaxRobertaPreLayerNormForMultipleChoice,
		        FlaxRobertaPreLayerNormForQuestionAnswering,
		    )
		    if is_flax_available()
		    else ()
		)





		def    UpperCAmelCase__							(					self   ) ->						int:







								'''simple docstring'''

								snake_case_   : Any						=			FlaxRobertaPreLayerNormModelTester(self   )





		@slow
		def    UpperCAmelCase__							(					self   ) ->						Tuple:







								'''simple docstring'''

								for model_class_name in self.all_model_classes:
														snake_case_   : Tuple						=			model_class_name.from_pretrained("""andreasmadsen/efficient_mlm_m0.40"""							,    from_pt=_lowercase   )
														snake_case_   : Tuple						=			model(np.ones((1, 1)   )   )
														self.assertIsNotNone(_lowercase   )




@require_flax
class    _lowerCAmelCase    (				unittest.TestCase		):





		"""simple docstring"""






		@slow
		def    UpperCAmelCase__							(					self   ) ->						Union[str, Any]:







								'''simple docstring'''

								snake_case_   : Union[str, Any]						=			FlaxRobertaPreLayerNormForMaskedLM.from_pretrained("""andreasmadsen/efficient_mlm_m0.40"""							,    from_pt=_lowercase   )

								snake_case_   : List[str]						=			np.array([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]]							,    dtype=jnp.intaa   )
								snake_case_   : str						=			model(_lowercase   )[0]
								snake_case_   : int						=			[1, 1_1, 5_0_2_6_5]
								self.assertEqual(list(output.shape   )							,    _lowercase   )
								# compare the actual values for a slice.
								snake_case_   : Tuple						=			np.array(
								    [[[40.4880, 18.0199, -5.2367], [-1.8877, -4.0885, 10.7085], [-2.2613, -5.6110, 7.2665]]]							,    dtype=np.floataa   )
								self.assertTrue(np.allclose(output[:, :3, :3]							,    _lowercase							,    atol=1E-4   )   )





		@slow
		def    UpperCAmelCase__							(					self   ) ->						int:







								'''simple docstring'''

								snake_case_   : Any						=			FlaxRobertaPreLayerNormModel.from_pretrained("""andreasmadsen/efficient_mlm_m0.40"""							,    from_pt=_lowercase   )

								snake_case_   : Dict						=			np.array([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]]							,    dtype=jnp.intaa   )
								snake_case_   : Any						=			model(_lowercase   )[0]
								# compare the actual values for a slice.
								snake_case_   : Optional[Any]						=			np.array(
								    [[[0.0208, -0.0356, 0.0237], [-0.1569, -0.0411, -0.2626], [0.1879, 0.0125, -0.0089]]]							,    dtype=np.floataa   )
								self.assertTrue(np.allclose(output[:, :3, :3]							,    _lowercase							,    atol=1E-4   )   )

