
'''simple docstring'''




import unittest

from transformers import AlbertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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 (
          MODEL_FOR_PRETRAINING_MAPPING,
          AlbertForMaskedLM,
          AlbertForMultipleChoice,
          AlbertForPreTraining,
          AlbertForQuestionAnswering,
          AlbertForSequenceClassification,
          AlbertForTokenClassification,
          AlbertModel,
      )
      from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST





class  A				:





     def __init__(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=13						,			SCREAMING_SNAKE_CASE=7						,			SCREAMING_SNAKE_CASE=True						,			SCREAMING_SNAKE_CASE=True						,			SCREAMING_SNAKE_CASE=True						,			SCREAMING_SNAKE_CASE=True						,			SCREAMING_SNAKE_CASE=99						,			SCREAMING_SNAKE_CASE=16						,			SCREAMING_SNAKE_CASE=36						,			SCREAMING_SNAKE_CASE=6						,			SCREAMING_SNAKE_CASE=6						,			SCREAMING_SNAKE_CASE=6						,			SCREAMING_SNAKE_CASE=37						,			SCREAMING_SNAKE_CASE="gelu"						,			SCREAMING_SNAKE_CASE=0.1						,			SCREAMING_SNAKE_CASE=0.1						,			SCREAMING_SNAKE_CASE=512						,			SCREAMING_SNAKE_CASE=16						,			SCREAMING_SNAKE_CASE=2						,			SCREAMING_SNAKE_CASE=0.02						,			SCREAMING_SNAKE_CASE=3						,			SCREAMING_SNAKE_CASE=4						,			SCREAMING_SNAKE_CASE=None						,			) ->				Dict:






         """simple docstring"""






         A       :   Optional[int]					=						parent
         A       :   Optional[int]					=						batch_size
         A       :   Union[str, Any]					=						seq_length
         A       :   Union[str, Any]					=						is_training
         A       :   Optional[Any]					=						use_input_mask
         A       :   Tuple					=						use_token_type_ids
         A       :   Optional[int]					=						use_labels
         A       :   Tuple					=						vocab_size
         A       :   Optional[Any]					=						embedding_size
         A       :   str					=						hidden_size
         A       :   Dict					=						num_hidden_layers
         A       :   str					=						num_hidden_groups
         A       :   Dict					=						num_attention_heads
         A       :   List[Any]					=						intermediate_size
         A       :   List[str]					=						hidden_act
         A       :   Optional[Any]					=						hidden_dropout_prob
         A       :   Any					=						attention_probs_dropout_prob
         A       :   Any					=						max_position_embeddings
         A       :   Any					=						type_vocab_size
         A       :   Dict					=						type_sequence_label_size
         A       :   Optional[int]					=						initializer_range
         A       :   List[Any]					=						num_labels
         A       :   int					=						num_choices
         A       :   Optional[int]					=						scope





     def 				__lowerCAmelCase							(							self	) ->				int:






         """simple docstring"""






         A       :   Any					=						ids_tensor([self.batch_size, self.seq_length]						,			self.vocab_size	)

         A       :   str					=						None
         if self.use_input_mask:
             A       :   int					=						random_attention_mask([self.batch_size, self.seq_length]	)

         A       :   Tuple					=						None
         if self.use_token_type_ids:
             A       :   Any					=						ids_tensor([self.batch_size, self.seq_length]						,			self.type_vocab_size	)

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

         A       :   Tuple					=						self.get_config()

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





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






         """simple docstring"""






         return AlbertConfig(
             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						,			initializer_range=self.initializer_range						,			num_hidden_groups=self.num_hidden_groups						,			)





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				List[str]:






         """simple docstring"""






         A       :   Any					=						AlbertModel(config=SCREAMING_SNAKE_CASE	)
         model.to(SCREAMING_SNAKE_CASE	)
         model.eval()
         A       :   Dict					=						model(SCREAMING_SNAKE_CASE						,			attention_mask=SCREAMING_SNAKE_CASE						,			token_type_ids=SCREAMING_SNAKE_CASE	)
         A       :   str					=						model(SCREAMING_SNAKE_CASE						,			token_type_ids=SCREAMING_SNAKE_CASE	)
         A       :   Optional[int]					=						model(SCREAMING_SNAKE_CASE	)
         self.parent.assertEqual(result.last_hidden_state.shape						,			(self.batch_size, self.seq_length, self.hidden_size)	)
         self.parent.assertEqual(result.pooler_output.shape						,			(self.batch_size, self.hidden_size)	)





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Optional[int]:






         """simple docstring"""






         A       :   int					=						AlbertForPreTraining(config=SCREAMING_SNAKE_CASE	)
         model.to(SCREAMING_SNAKE_CASE	)
         model.eval()
         A       :   Optional[Any]					=						model(
             SCREAMING_SNAKE_CASE						,			attention_mask=SCREAMING_SNAKE_CASE						,			token_type_ids=SCREAMING_SNAKE_CASE						,			labels=SCREAMING_SNAKE_CASE						,			sentence_order_label=SCREAMING_SNAKE_CASE						,			)
         self.parent.assertEqual(result.prediction_logits.shape						,			(self.batch_size, self.seq_length, self.vocab_size)	)
         self.parent.assertEqual(result.sop_logits.shape						,			(self.batch_size, config.num_labels)	)





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				str:






         """simple docstring"""






         A       :   Any					=						AlbertForMaskedLM(config=SCREAMING_SNAKE_CASE	)
         model.to(SCREAMING_SNAKE_CASE	)
         model.eval()
         A       :   Tuple					=						model(SCREAMING_SNAKE_CASE						,			attention_mask=SCREAMING_SNAKE_CASE						,			token_type_ids=SCREAMING_SNAKE_CASE						,			labels=SCREAMING_SNAKE_CASE	)
         self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.seq_length, self.vocab_size)	)





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Tuple:






         """simple docstring"""






         A       :   Any					=						AlbertForQuestionAnswering(config=SCREAMING_SNAKE_CASE	)
         model.to(SCREAMING_SNAKE_CASE	)
         model.eval()
         A       :   str					=						model(
             SCREAMING_SNAKE_CASE						,			attention_mask=SCREAMING_SNAKE_CASE						,			token_type_ids=SCREAMING_SNAKE_CASE						,			start_positions=SCREAMING_SNAKE_CASE						,			end_positions=SCREAMING_SNAKE_CASE						,			)
         self.parent.assertEqual(result.start_logits.shape						,			(self.batch_size, self.seq_length)	)
         self.parent.assertEqual(result.end_logits.shape						,			(self.batch_size, self.seq_length)	)





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				List[str]:






         """simple docstring"""






         A       :   List[Any]					=						self.num_labels
         A       :   str					=						AlbertForSequenceClassification(SCREAMING_SNAKE_CASE	)
         model.to(SCREAMING_SNAKE_CASE	)
         model.eval()
         A       :   int					=						model(SCREAMING_SNAKE_CASE						,			attention_mask=SCREAMING_SNAKE_CASE						,			token_type_ids=SCREAMING_SNAKE_CASE						,			labels=SCREAMING_SNAKE_CASE	)
         self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.num_labels)	)





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






         """simple docstring"""






         A       :   Optional[Any]					=						self.num_labels
         A       :   str					=						AlbertForTokenClassification(config=SCREAMING_SNAKE_CASE	)
         model.to(SCREAMING_SNAKE_CASE	)
         model.eval()
         A       :   Union[str, Any]					=						model(SCREAMING_SNAKE_CASE						,			attention_mask=SCREAMING_SNAKE_CASE						,			token_type_ids=SCREAMING_SNAKE_CASE						,			labels=SCREAMING_SNAKE_CASE	)
         self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.seq_length, self.num_labels)	)





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				List[str]:






         """simple docstring"""






         A       :   str					=						self.num_choices
         A       :   List[str]					=						AlbertForMultipleChoice(config=SCREAMING_SNAKE_CASE	)
         model.to(SCREAMING_SNAKE_CASE	)
         model.eval()
         A       :   Union[str, Any]					=						input_ids.unsqueeze(1	).expand(-1						,			self.num_choices						,			-1	).contiguous()
         A       :   Tuple					=						token_type_ids.unsqueeze(1	).expand(-1						,			self.num_choices						,			-1	).contiguous()
         A       :   Optional[Any]					=						input_mask.unsqueeze(1	).expand(-1						,			self.num_choices						,			-1	).contiguous()
         A       :   List[str]					=						model(
             SCREAMING_SNAKE_CASE						,			attention_mask=SCREAMING_SNAKE_CASE						,			token_type_ids=SCREAMING_SNAKE_CASE						,			labels=SCREAMING_SNAKE_CASE						,			)
         self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.num_choices)	)





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






         """simple docstring"""






         A       :   Union[str, Any]					=						self.prepare_config_and_inputs()
         (
             (
             A
         ),  (
             A
         ),  (
             A
         ),  (
             A
         ),  (
             A
         ),  (
             A
         ),  (
             A
         ),  
         )       :   str					=						config_and_inputs
         A       :   Union[str, Any]					=						{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
         return config, inputs_dict





@require_torch
class  A				(							__snake_case    ,			__snake_case    ,			unittest.TestCase						):
     __magic_name__						=		(
         (
             AlbertModel,
             AlbertForPreTraining,
             AlbertForMaskedLM,
             AlbertForMultipleChoice,
             AlbertForSequenceClassification,
             AlbertForTokenClassification,
             AlbertForQuestionAnswering,
         )
         if is_torch_available()
         else ()
     )
     __magic_name__						=		(
         {
             '''feature-extraction''': AlbertModel,
             '''fill-mask''': AlbertForMaskedLM,
             '''question-answering''': AlbertForQuestionAnswering,
             '''text-classification''': AlbertForSequenceClassification,
             '''token-classification''': AlbertForTokenClassification,
             '''zero-shot''': AlbertForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     __magic_name__						=		True





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=False	) ->				Tuple:






         """simple docstring"""






         A       :   Dict					=						super()._prepare_for_class(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			return_labels=SCREAMING_SNAKE_CASE	)

         if return_labels:
             if model_class in get_values(SCREAMING_SNAKE_CASE	):
                 A       :   Any					=						torch.zeros(
                     (self.model_tester.batch_size, self.model_tester.seq_length)						,			dtype=torch.long						,			device=SCREAMING_SNAKE_CASE	)
                 A       :   str					=						torch.zeros(
                     self.model_tester.batch_size						,			dtype=torch.long						,			device=SCREAMING_SNAKE_CASE	)
         return inputs_dict





     def 				__lowerCAmelCase							(							self	) ->				Dict:






         """simple docstring"""






         A       :   Any					=						AlbertModelTester(self	)
         A       :   int					=						ConfigTester(self						,			config_class=SCREAMING_SNAKE_CASE						,			hidden_size=37	)





     def 				__lowerCAmelCase							(							self	) ->				int:






         """simple docstring"""






         self.config_tester.run_common_tests()





     def 				__lowerCAmelCase							(							self	) ->				Optional[int]:






         """simple docstring"""






         A       :   Union[str, Any]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Optional[int]:






         """simple docstring"""






         A       :   List[str]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_pretraining(*SCREAMING_SNAKE_CASE	)





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






         """simple docstring"""






         A       :   Optional[Any]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE	)





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






         """simple docstring"""






         A       :   str					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE	)





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






         """simple docstring"""






         A       :   List[str]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Dict:






         """simple docstring"""






         A       :   Tuple					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE	)





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






         """simple docstring"""






         A       :   Tuple					=						self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
             A       :   str					=						type
             self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE	)





     @slow
     def 				__lowerCAmelCase							(							self	) ->				Tuple:






         """simple docstring"""






         for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             A       :   int					=						AlbertModel.from_pretrained(SCREAMING_SNAKE_CASE	)
             self.assertIsNotNone(SCREAMING_SNAKE_CASE	)





@require_torch
class  A				(							unittest.TestCase						):





     @slow
     def 				__lowerCAmelCase							(							self	) ->				List[str]:






         """simple docstring"""






         A       :   Optional[Any]					=						AlbertModel.from_pretrained('''albert-base-v2'''	)
         A       :   Any					=						torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]	)
         A       :   Union[str, Any]					=						torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]	)
         with torch.no_grad():
             A       :   Optional[int]					=						model(SCREAMING_SNAKE_CASE						,			attention_mask=SCREAMING_SNAKE_CASE	)[0]
         A       :   str					=						torch.Size((1, 11, 768)	)
         self.assertEqual(output.shape						,			SCREAMING_SNAKE_CASE	)
         A       :   Optional[int]					=						torch.tensor(
             [[[-0.6_513, 1.5_035, -0.2_766], [-0.6_515, 1.5_046, -0.2_780], [-0.6_512, 1.5_049, -0.2_784]]]	)

         self.assertTrue(torch.allclose(output[:, 1:4, 1:4]						,			SCREAMING_SNAKE_CASE						,			atol=1e-4	)	)












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 __lowerCamelCase (      unittest.TestCase):





	"""simple docstring"""






	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							with tempfile.TemporaryDirectory() as tmpdirname:
													# pipeline has Flax weights
													_UpperCAmelCase									=			FlaxDiffusionPipeline.from_pretrained(
													    'hf-internal-testing/tiny-stable-diffusion-pipe'       ,    safety_checker=UpperCAmelCase       ,    cache_dir=UpperCAmelCase    )

													_UpperCAmelCase									=			[t[-1] for t in os.walk(os.path.join(UpperCAmelCase       ,    os.listdir(UpperCAmelCase    )[0]       ,    'snapshots'    )    )]
													_UpperCAmelCase									=			[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 __lowerCamelCase (      unittest.TestCase):





	"""simple docstring"""






	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase					, _UpperCAmelCase									=			FlaxStableDiffusionPipeline.from_pretrained(
							    'hf-internal-testing/tiny-stable-diffusion-pipe'       ,    safety_checker=UpperCAmelCase    )

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

							_UpperCAmelCase									=			jax.random.PRNGKey(0    )
							_UpperCAmelCase									=			4

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).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(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 4_99_47.8_75    ) < 5e-1

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

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase					, _UpperCAmelCase									=			FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'       ,    revision='flax'       ,    safety_checker=UpperCAmelCase    )

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

							_UpperCAmelCase									=			jax.random.PRNGKey(0    )
							_UpperCAmelCase									=			50

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).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(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_38_38_08.2)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

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

							_UpperCAmelCase									=			jax.random.PRNGKey(0    )
							_UpperCAmelCase									=			50

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).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(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_37_35_16.75)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase					, _UpperCAmelCase									=			FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'       ,    revision='bf16'       ,    dtype=jnp.bfloataa    )

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

							_UpperCAmelCase									=			jax.random.PRNGKey(0    )
							_UpperCAmelCase									=			50

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).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(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_37_35_16.75)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			FlaxDDIMScheduler(
							    beta_start=0.0_00_85       ,    beta_end=0.0_12       ,    beta_schedule='scaled_linear'       ,    set_alpha_to_one=UpperCAmelCase       ,    steps_offset=1       ,    )

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

							_UpperCAmelCase									=			scheduler_state

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

							_UpperCAmelCase									=			jax.random.PRNGKey(0    )
							_UpperCAmelCase									=			50

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).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(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_34_76_93.5)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			jax.random.split(jax.random.PRNGKey(0    )       ,    UpperCAmelCase    )

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

							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images
							assert images.shape == (num_samples, 1, 512, 512, 3)
							_UpperCAmelCase									=			images[2, 0, 256, 10:17, 1]

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

							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images
							assert images_eff.shape == (num_samples, 1, 512, 512, 3)
							_UpperCAmelCase									=			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
