





'''simple docstring'''




import math
import unittest

from transformers import BioGptConfig, 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 (
      BioGptForCausalLM,
      BioGptForSequenceClassification,
      BioGptForTokenClassification,
      BioGptModel,
      BioGptTokenizer,
  )
  from transformers.models.biogpt.modeling_biogpt import BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST


class  SCREAMING_SNAKE_CASE    :


     def __init__(						self						,	_UpperCAmelCase						,	_UpperCAmelCase=13						,	_UpperCAmelCase=7						,	_UpperCAmelCase=True						,	_UpperCAmelCase=True						,	_UpperCAmelCase=False						,	_UpperCAmelCase=True						,	_UpperCAmelCase=99						,	_UpperCAmelCase=32						,	_UpperCAmelCase=5						,	_UpperCAmelCase=4						,	_UpperCAmelCase=37						,	_UpperCAmelCase="gelu"						,	_UpperCAmelCase=0.1						,	_UpperCAmelCase=0.1						,	_UpperCAmelCase=512						,	_UpperCAmelCase=16						,	_UpperCAmelCase=2						,	_UpperCAmelCase=0.02						,	_UpperCAmelCase=3						,	_UpperCAmelCase=4						,	_UpperCAmelCase=None						,	):

          '''simple docstring'''






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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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

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

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

          __A						: Optional[int]      =       None
          __A						: Any      =       None
          __A						: str      =       None
          if self.use_labels:
               __A						: int      =       ids_tensor([self.batch_size]						,	self.type_sequence_label_size)
               __A						: Tuple      =       ids_tensor([self.batch_size, self.seq_length]						,	self.num_labels)
               __A						: List[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 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase):

          '''simple docstring'''






          __A						: str      =       BioGptModel(config=_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()
          __A						: Optional[int]      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase)
          __A						: Tuple      =       model(_UpperCAmelCase)
          self.parent.assertEqual(result.last_hidden_state.shape						,	(self.batch_size, self.seq_length, self.hidden_size))


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	):

          '''simple docstring'''






          __A						: List[str]      =       BioGptForCausalLM(config=_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()
          __A						: Optional[int]      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase						,	token_type_ids=_UpperCAmelCase						,	labels=_UpperCAmelCase)
          self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.seq_length, self.vocab_size))


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	*_UpperCAmelCase):

          '''simple docstring'''






          __A						: Union[str, Any]      =       BioGptModel(config=_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()

          # create attention mask
          __A						: Tuple      =       torch.ones(input_ids.shape						,	dtype=torch.long						,	device=_UpperCAmelCase)
          __A						: List[str]      =       self.seq_length // 2
          __A						: str      =       0

          # first forward pass
          __A ,__A						: Union[str, Any]      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase).to_tuple()

          # create hypothetical next token and extent to next_input_ids
          __A						: Optional[int]      =       ids_tensor((self.batch_size, 1)						,	config.vocab_size)

          # change a random masked slice from input_ids
          __A						: Dict      =       ids_tensor((1,)						,	_UpperCAmelCase).item() + 1
          __A						: Any      =       ids_tensor((self.batch_size, 1)						,	config.vocab_size).squeeze(-1)
          __A						: Dict      =       random_other_next_tokens

          # append to next input_ids and attn_mask
          __A						: Any      =       torch.cat([input_ids, next_tokens]						,	dim=-1)
          __A						: Dict      =       torch.cat(
              [attn_mask, torch.ones((attn_mask.shape[0], 1)						,	dtype=torch.long						,	device=_UpperCAmelCase)]						,	dim=1						,	)

          # get two different outputs
          __A						: Dict      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase)['last_hidden_state']
          __A						: List[Any]      =       model(_UpperCAmelCase						,	past_key_values=_UpperCAmelCase						,	attention_mask=_UpperCAmelCase)['last_hidden_state']

          # select random slice
          __A						: str      =       ids_tensor((1,)						,	output_from_past.shape[-1]).item()
          __A						: Optional[Any]      =       output_from_no_past[:, -1, random_slice_idx].detach()
          __A						: Any      =       output_from_past[:, 0, random_slice_idx].detach()

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


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	*_UpperCAmelCase):

          '''simple docstring'''






          __A						: str      =       BioGptModel(config=_UpperCAmelCase).to(_UpperCAmelCase).eval()

          __A						: Dict      =       torch.ones(input_ids.shape						,	dtype=torch.long						,	device=_UpperCAmelCase)

          # first forward pass
          __A						: List[Any]      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase						,	use_cache=_UpperCAmelCase)

          __A ,__A						: List[Any]      =       outputs.to_tuple()

          # create hypothetical multiple next token and extent to next_input_ids
          __A						: Union[str, Any]      =       ids_tensor((self.batch_size, 3)						,	config.vocab_size)
          __A						: Optional[int]      =       ids_tensor((self.batch_size, 3)						,	2)

          # append to next input_ids and
          __A						: Optional[int]      =       torch.cat([input_ids, next_tokens]						,	dim=-1)
          __A						: Optional[Any]      =       torch.cat([attention_mask, next_attn_mask]						,	dim=-1)

          __A						: Dict      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase)['last_hidden_state']
          __A						: Any      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase						,	past_key_values=_UpperCAmelCase)[
              'last_hidden_state'
          ]

          # select random slice
          __A						: Dict      =       ids_tensor((1,)						,	output_from_past.shape[-1]).item()
          __A						: Union[str, Any]      =       output_from_no_past[:, -3:, random_slice_idx].detach()
          __A						: Union[str, 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(_UpperCAmelCase						,	_UpperCAmelCase						,	atol=1e-3))


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	*_UpperCAmelCase						,	_UpperCAmelCase=False):

          '''simple docstring'''






          __A						: str      =       BioGptForCausalLM(_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          if gradient_checkpointing:
               model.gradient_checkpointing_enable()

          __A						: str      =       model(_UpperCAmelCase						,	labels=_UpperCAmelCase)
          self.parent.assertEqual(result.loss.shape						,	())
          self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.seq_length, self.vocab_size))
          result.loss.backward()


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	*_UpperCAmelCase):

          '''simple docstring'''






          __A						: List[str]      =       BioGptModel(_UpperCAmelCase)
          __A						: Optional[int]      =       model.config.initializer_range / math.sqrt(2 * model.config.num_hidden_layers)
          for key in model.state_dict().keys():
               if "c_proj" in key and "weight" in key:
                    self.parent.assertLessEqual(abs(torch.std(model.state_dict()[key]) - model_std)						,	0.001)
                    self.parent.assertLessEqual(abs(torch.mean(model.state_dict()[key]) - 0.0)						,	0.01)


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	*_UpperCAmelCase):

          '''simple docstring'''






          __A						: List[str]      =       self.num_labels
          __A						: List[Any]      =       BioGptForTokenClassification(_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()
          __A						: Dict      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase						,	token_type_ids=_UpperCAmelCase)
          self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.seq_length, self.num_labels))


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Dict      =       self.prepare_config_and_inputs()
          (
              (
              __A
          ) ,(
              __A
          ) ,(
              __A
          ) ,(
              __A
          ) ,(
              __A
          ) ,(
              __A
          ) ,(
              __A
          ) ,
          )						: str      =       config_and_inputs
          __A						: List[str]      =       {'input_ids': input_ids, 'attention_mask': input_mask}
          return config, inputs_dict




@require_torch
class  SCREAMING_SNAKE_CASE    (a__					,		a__					,		a__					,		unittest.TestCase							):
     lowerCAmelCase					=       (
         (BioGptModel, BioGptForCausalLM, BioGptForSequenceClassification, BioGptForTokenClassification)
         if is_torch_available()
         else ()
     )
     lowerCAmelCase					=       (BioGptForCausalLM,) if is_torch_available() else ()
     lowerCAmelCase					=       (
         {
             '''feature-extraction''': BioGptModel,
             '''text-classification''': BioGptForSequenceClassification,
             '''text-generation''': BioGptForCausalLM,
             '''token-classification''': BioGptForTokenClassification,
             '''zero-shot''': BioGptForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     lowerCAmelCase					=       False


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[Any]      =       BioGptModelTester(self)
          __A						: List[str]      =       ConfigTester(self						,	config_class=_UpperCAmelCase						,	hidden_size=37)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          self.config_tester.run_common_tests()


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: str      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[str]      =       self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
               __A						: Optional[int]      =       type
               self.model_tester.create_and_check_model(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Dict      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_biogpt_model_attention_mask_past(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Tuple      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_forward_and_backwards(*_UpperCAmelCase						,	gradient_checkpointing=_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Optional[Any]      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_biogpt_model_past_large_inputs(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: int      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_biogpt_weight_initialization(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Union[str, Any]      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_biogpt_for_token_classification(*_UpperCAmelCase)


     @slow
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Any      =       BioGptForCausalLM.from_pretrained('microsoft/biogpt')
          model.to(_UpperCAmelCase)
          __A						: Any      =       BioGptTokenizer.from_pretrained('microsoft/biogpt')

          __A						: Dict      =       'left'

          # Define PAD Token = EOS Token = 50256
          __A						: Optional[int]      =       tokenizer.eos_token
          __A						: Dict      =       model.config.eos_token_id

          # use different length sentences to test batching
          __A						: Dict      =       [
              'Hello, my dog is a little',
              'Today, I',
          ]

          __A						: Union[str, Any]      =       tokenizer(_UpperCAmelCase						,	return_tensors='pt'						,	padding=_UpperCAmelCase)
          __A						: Union[str, Any]      =       inputs['input_ids'].to(_UpperCAmelCase)

          __A						: Dict      =       model.generate(
              input_ids=_UpperCAmelCase						,	attention_mask=inputs['attention_mask'].to(_UpperCAmelCase)						,	)

          __A						: Dict      =       tokenizer(sentences[0]						,	return_tensors='pt').input_ids.to(_UpperCAmelCase)
          __A						: Optional[int]      =       model.generate(input_ids=_UpperCAmelCase)

          __A						: Optional[int]      =       inputs_non_padded.shape[-1] - inputs['attention_mask'][-1].long().sum().cpu().item()
          __A						: str      =       tokenizer(sentences[1]						,	return_tensors='pt').input_ids.to(_UpperCAmelCase)
          __A						: Any      =       model.generate(input_ids=_UpperCAmelCase						,	max_length=model.config.max_length - num_paddings)

          __A						: Optional[Any]      =       tokenizer.batch_decode(_UpperCAmelCase						,	skip_special_tokens=_UpperCAmelCase)
          __A						: Union[str, Any]      =       tokenizer.decode(output_non_padded[0]						,	skip_special_tokens=_UpperCAmelCase)
          __A						: Dict      =       tokenizer.decode(output_padded[0]						,	skip_special_tokens=_UpperCAmelCase)

          __A						: Optional[int]      =       [
              'Hello, my dog is a little bit bigger than a little bit.',
              'Today, I have a good idea of how to use the information',
          ]
          self.assertListEqual(_UpperCAmelCase						,	_UpperCAmelCase)
          self.assertListEqual(_UpperCAmelCase						,	[non_padded_sentence, padded_sentence])


     @slow
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          for model_name in BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __A						: int      =       BioGptModel.from_pretrained(_UpperCAmelCase)
               self.assertIsNotNone(_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A ,__A						: Dict      =       self.model_tester.prepare_config_and_inputs_for_common()
          __A						: List[str]      =       3
          __A						: int      =       input_dict['input_ids']
          __A						: Any      =       input_ids.ne(1).to(_UpperCAmelCase)
          __A						: Tuple      =       ids_tensor([self.model_tester.batch_size]						,	self.model_tester.type_sequence_label_size)
          __A						: Dict      =       BioGptForSequenceClassification(_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()
          __A						: int      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase						,	labels=_UpperCAmelCase)
          self.assertEqual(result.logits.shape						,	(self.model_tester.batch_size, self.model_tester.num_labels))


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A ,__A						: Any      =       self.model_tester.prepare_config_and_inputs_for_common()
          __A						: Optional[Any]      =       3
          __A						: int      =       'multi_label_classification'
          __A						: int      =       input_dict['input_ids']
          __A						: Optional[int]      =       input_ids.ne(1).to(_UpperCAmelCase)
          __A						: Dict      =       ids_tensor(
              [self.model_tester.batch_size, config.num_labels]						,	self.model_tester.type_sequence_label_size).to(torch.float)
          __A						: Dict      =       BioGptForSequenceClassification(_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()
          __A						: Any      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase						,	labels=_UpperCAmelCase)
          self.assertEqual(result.logits.shape						,	(self.model_tester.batch_size, self.model_tester.num_labels))




@require_torch
class  SCREAMING_SNAKE_CASE    (unittest.TestCase							):


     @slow
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[str]      =       BioGptForCausalLM.from_pretrained('microsoft/biogpt')
          __A						: Dict      =       torch.tensor([[2, 4805, 9, 656, 21]])
          __A						: List[str]      =       model(_UpperCAmelCase)[0]

          __A						: Union[str, Any]      =       4_2384

          __A						: Any      =       torch.Size((1, 5, vocab_size))
          self.assertEqual(output.shape						,	_UpperCAmelCase)

          __A						: str      =       torch.tensor(
              [[[-9.5236, -9.8918, 10.4557], [-11.0469, -9.6423, 8.1022], [-8.8664, -7.8826, 5.5325]]])

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


     @slow
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: int      =       BioGptTokenizer.from_pretrained('microsoft/biogpt')
          __A						: Tuple      =       BioGptForCausalLM.from_pretrained('microsoft/biogpt')
          model.to(_UpperCAmelCase)

          torch.manual_seed(0)
          __A						: Any      =       tokenizer('COVID-19 is'						,	return_tensors='pt').to(_UpperCAmelCase)
          __A						: Tuple      =       model.generate(
              **_UpperCAmelCase						,	min_length=100						,	max_length=1024						,	num_beams=5						,	early_stopping=_UpperCAmelCase						,	)
          __A						: str      =       tokenizer.decode(output_ids[0]						,	skip_special_tokens=_UpperCAmelCase)

          __A						: int      =       (
              'COVID-19 is a global pandemic caused by severe acute respiratory syndrome coronavirus 2 (SARS-CoV-2), the'
              ' causative agent of coronavirus disease 2019 (COVID-19), which has spread to more than 200 countries and'
              ' territories, including the United States (US), Canada, Australia, New Zealand, the United Kingdom (UK),'
              ' and the United States of America (USA), as of March 11, 2020, with more than 800,000 confirmed cases and'
              ' more than 800,000 deaths.'
          )
          self.assertEqual(_UpperCAmelCase						,	_UpperCAmelCase)






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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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

               __A						: Optional[Any]      =       [t[-1] for t in os.walk(os.path.join(_UpperCAmelCase						,	os.listdir(_UpperCAmelCase)[0]						,	'snapshots'))]
               __A						: int      =       [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							):


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A ,__A						: Optional[int]      =       FlaxStableDiffusionPipeline.from_pretrained(
              'hf-internal-testing/tiny-stable-diffusion-pipe'						,	safety_checker=_UpperCAmelCase)

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

          __A						: Optional[Any]      =       jax.random.PRNGKey(0)
          __A						: int      =       4

          __A						: Tuple      =       jax.device_count()
          __A						: Union[str, Any]      =       num_samples * [prompt]
          __A						: Tuple      =       pipeline.prepare_inputs(_UpperCAmelCase)

          # shard inputs and rng
          __A						: str      =       replicate(_UpperCAmelCase)
          __A						: Tuple      =       jax.random.split(_UpperCAmelCase						,	_UpperCAmelCase)
          __A						: Union[str, Any]      =       shard(_UpperCAmelCase)

          __A						: Union[str, Any]      =       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.1514745) < 1e-3
               assert np.abs(np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 49947.875) < 5e-1

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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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

          __A						: 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'
          )

          __A						: Tuple      =       jax.random.PRNGKey(0)
          __A						: Any      =       50

          __A						: str      =       jax.device_count()
          __A						: Union[str, Any]      =       num_samples * [prompt]
          __A						: List[str]      =       pipeline.prepare_inputs(_UpperCAmelCase)

          # shard inputs and rng
          __A						: Dict      =       replicate(_UpperCAmelCase)
          __A						: Optional[Any]      =       jax.random.split(_UpperCAmelCase						,	_UpperCAmelCase)
          __A						: int      =       shard(_UpperCAmelCase)

          __A						: Tuple      =       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.05652401)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 2383808.2)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A ,__A						: List[str]      =       FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,	revision='bf16'						,	dtype=jnp.bfloataa						,	safety_checker=_UpperCAmelCase)

          __A						: 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'
          )

          __A						: str      =       jax.random.PRNGKey(0)
          __A						: Any      =       50

          __A						: Optional[int]      =       jax.device_count()
          __A						: int      =       num_samples * [prompt]
          __A						: Optional[int]      =       pipeline.prepare_inputs(_UpperCAmelCase)

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

          __A						: str      =       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.04003906)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 2373516.75)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A ,__A						: Optional[Any]      =       FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,	revision='bf16'						,	dtype=jnp.bfloataa)

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

          __A						: Any      =       jax.random.PRNGKey(0)
          __A						: List[str]      =       50

          __A						: Optional[int]      =       jax.device_count()
          __A						: List[Any]      =       num_samples * [prompt]
          __A						: List[Any]      =       pipeline.prepare_inputs(_UpperCAmelCase)

          # shard inputs and rng
          __A						: Union[str, Any]      =       replicate(_UpperCAmelCase)
          __A						: Optional[Any]      =       jax.random.split(_UpperCAmelCase						,	_UpperCAmelCase)
          __A						: List[str]      =       shard(_UpperCAmelCase)

          __A						: int      =       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.04003906)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 2373516.75)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[Any]      =       FlaxDDIMScheduler(
              beta_start=0.00085						,	beta_end=0.012						,	beta_schedule='scaled_linear'						,	set_alpha_to_one=_UpperCAmelCase						,	steps_offset=1						,	)

          __A ,__A						: Any      =       FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,	revision='bf16'						,	dtype=jnp.bfloataa						,	scheduler=_UpperCAmelCase						,	safety_checker=_UpperCAmelCase						,	)
          __A						: Optional[Any]      =       scheduler.create_state()

          __A						: Any      =       scheduler_state

          __A						: 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'
          )

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

          __A						: Optional[Any]      =       jax.device_count()
          __A						: Any      =       num_samples * [prompt]
          __A						: Optional[Any]      =       pipeline.prepare_inputs(_UpperCAmelCase)

          # shard inputs and rng
          __A						: int      =       replicate(_UpperCAmelCase)
          __A						: Any      =       jax.random.split(_UpperCAmelCase						,	_UpperCAmelCase)
          __A						: Tuple      =       shard(_UpperCAmelCase)

          __A						: Union[str, Any]      =       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.045043945)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 2347693.5)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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

          __A						: int      =       jax.device_count()
          __A						: List[Any]      =       num_samples * [prompt]
          __A						: List[Any]      =       jax.random.split(jax.random.PRNGKey(0)						,	_UpperCAmelCase)

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

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

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

          __A						: Any      =       replicate(_UpperCAmelCase)
          __A						: List[Any]      =       pipeline.prepare_inputs(_UpperCAmelCase)
          __A						: Optional[Any]      =       shard(_UpperCAmelCase)
          __A						: List[Any]      =       pipeline(_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	jit=_UpperCAmelCase).images
          assert images_eff.shape == (num_samples, 1, 512, 512, 3)
          __A						: List[Any]      =       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