





'''simple docstring'''




import json
import os
import shutil
import tempfile
import unittest
from multiprocessing import get_context
from pathlib import Path

import datasets
import numpy as np
from datasets import load_dataset
from parameterized import parameterized

from transformers import AutoProcessor
from transformers.models.wavaveca import WavaVecaCTCTokenizer, WavaVecaFeatureExtractor
from transformers.models.wavaveca.tokenization_wavaveca import VOCAB_FILES_NAMES
from transformers.testing_utils import require_pyctcdecode, require_torch, require_torchaudio, slow
from transformers.utils import FEATURE_EXTRACTOR_NAME, is_pyctcdecode_available, is_torch_available

from ..wavaveca.test_feature_extraction_wavaveca import floats_list


if is_pyctcdecode_available():
  from huggingface_hub import snapshot_download
  from pyctcdecode import BeamSearchDecoderCTC

  from transformers.models.wavaveca_with_lm import WavaVecaProcessorWithLM
  from transformers.models.wavaveca_with_lm.processing_wavaveca_with_lm import WavaVecaDecoderWithLMOutput

if is_torch_available():
  from transformers import WavaVecaForCTC


@require_pyctcdecode
class  SCREAMING_SNAKE_CASE    (unittest.TestCase							):


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: str      =       '| <pad> <unk> <s> </s> a b c d e f g h i j k'.split()
          __A						: List[Any]      =       dict(zip(_UpperCAmelCase						,	range(len(_UpperCAmelCase))))

          __A						: List[str]      =       {
              'unk_token': '<unk>',
              'bos_token': '<s>',
              'eos_token': '</s>',
          }
          __A						: str      =       {
              'feature_size': 1,
              'padding_value': 0.0,
              'sampling_rate': 1_6000,
              'return_attention_mask': False,
              'do_normalize': True,
          }

          __A						: List[str]      =       tempfile.mkdtemp()
          __A						: List[Any]      =       os.path.join(self.tmpdirname						,	VOCAB_FILES_NAMES['vocab_file'])
          __A						: Dict      =       os.path.join(self.tmpdirname						,	_UpperCAmelCase)
          with open(self.vocab_file						,	'w'						,	encoding='utf-8') as fp:
               fp.write(json.dumps(_UpperCAmelCase) + '\n')

          with open(self.feature_extraction_file						,	'w'						,	encoding='utf-8') as fp:
               fp.write(json.dumps(_UpperCAmelCase) + '\n')

          # load decoder from hub
          __A						: Tuple      =       'hf-internal-testing/ngram-beam-search-decoder'


     def 							SCREAMING_SNAKE_CASE  (						self						,	**_UpperCAmelCase):

          '''simple docstring'''






          __A						: Tuple      =       self.add_kwargs_tokens_map.copy()
          kwargs.update(_UpperCAmelCase)
          return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname						,	**_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self						,	**_UpperCAmelCase):

          '''simple docstring'''






          return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname						,	**_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self						,	**_UpperCAmelCase):

          '''simple docstring'''






          return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name						,	**_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          shutil.rmtree(self.tmpdirname)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: int      =       self.get_tokenizer()
          __A						: Any      =       self.get_feature_extractor()
          __A						: int      =       self.get_decoder()

          __A						: Optional[Any]      =       WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase						,	feature_extractor=_UpperCAmelCase						,	decoder=_UpperCAmelCase)

          processor.save_pretrained(self.tmpdirname)
          __A						: Dict      =       WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname)

          # tokenizer
          self.assertEqual(processor.tokenizer.get_vocab()						,	tokenizer.get_vocab())
          self.assertIsInstance(processor.tokenizer						,	_UpperCAmelCase)

          # feature extractor
          self.assertEqual(processor.feature_extractor.to_json_string()						,	feature_extractor.to_json_string())
          self.assertIsInstance(processor.feature_extractor						,	_UpperCAmelCase)

          # decoder
          self.assertEqual(processor.decoder._alphabet.labels						,	decoder._alphabet.labels)
          self.assertEqual(
              processor.decoder.model_container[decoder._model_key]._unigram_set						,	decoder.model_container[decoder._model_key]._unigram_set						,	)
          self.assertIsInstance(processor.decoder						,	_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Optional[Any]      =       WavaVecaProcessorWithLM(
              tokenizer=self.get_tokenizer()						,	feature_extractor=self.get_feature_extractor()						,	decoder=self.get_decoder())
          processor.save_pretrained(self.tmpdirname)

          # make sure that error is thrown when decoder alphabet doesn't match
          __A						: Any      =       WavaVecaProcessorWithLM.from_pretrained(
              self.tmpdirname						,	alpha=5.0						,	beta=3.0						,	score_boundary=-7.0						,	unk_score_offset=3)

          # decoder
          self.assertEqual(processor.language_model.alpha						,	5.0)
          self.assertEqual(processor.language_model.beta						,	3.0)
          self.assertEqual(processor.language_model.score_boundary						,	-7.0)
          self.assertEqual(processor.language_model.unk_score_offset						,	3)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Any      =       self.get_tokenizer()
          # add token to trigger raise
          tokenizer.add_tokens(['xx'])
          with self.assertRaisesRegex(_UpperCAmelCase						,	'include'):
               WavaVecaProcessorWithLM(
                   tokenizer=_UpperCAmelCase						,	feature_extractor=self.get_feature_extractor()						,	decoder=self.get_decoder())


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Union[str, Any]      =       self.get_feature_extractor()
          __A						: Union[str, Any]      =       self.get_tokenizer()
          __A						: str      =       self.get_decoder()

          __A						: Tuple      =       WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase						,	feature_extractor=_UpperCAmelCase						,	decoder=_UpperCAmelCase)

          __A						: Optional[int]      =       floats_list((3, 1000))

          __A						: int      =       feature_extractor(_UpperCAmelCase						,	return_tensors='np')
          __A						: Optional[Any]      =       processor(_UpperCAmelCase						,	return_tensors='np')

          for key in input_feat_extract.keys():
               self.assertAlmostEqual(input_feat_extract[key].sum()						,	input_processor[key].sum()						,	delta=1e-2)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: str      =       self.get_feature_extractor()
          __A						: int      =       self.get_tokenizer()
          __A						: Union[str, Any]      =       self.get_decoder()

          __A						: Dict      =       WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase						,	feature_extractor=_UpperCAmelCase						,	decoder=_UpperCAmelCase)

          __A						: Union[str, Any]      =       'This is a test string'

          __A						: Optional[Any]      =       processor(text=_UpperCAmelCase)

          __A						: Any      =       tokenizer(_UpperCAmelCase)

          for key in encoded_tok.keys():
               self.assertListEqual(encoded_tok[key]						,	encoded_processor[key])


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase=(2, 10, 16)						,	_UpperCAmelCase=77):

          '''simple docstring'''






          np.random.seed(_UpperCAmelCase)
          return np.random.rand(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: int      =       self.get_feature_extractor()
          __A						: Union[str, Any]      =       self.get_tokenizer()
          __A						: Union[str, Any]      =       self.get_decoder()

          __A						: Optional[int]      =       WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase						,	feature_extractor=_UpperCAmelCase						,	decoder=_UpperCAmelCase)

          __A						: Union[str, Any]      =       self._get_dummy_logits(shape=(10, 16)						,	seed=13)

          __A						: List[str]      =       processor.decode(_UpperCAmelCase)

          __A						: Union[str, Any]      =       decoder.decode_beams(_UpperCAmelCase)[0]

          self.assertEqual(decoded_decoder[0]						,	decoded_processor.text)
          self.assertEqual('</s> <s> </s>'						,	decoded_processor.text)
          self.assertEqual(decoded_decoder[-2]						,	decoded_processor.logit_score)
          self.assertEqual(decoded_decoder[-1]						,	decoded_processor.lm_score)


     @parameterized.expand([[None], ['fork'], ['spawn']])
     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase):

          '''simple docstring'''






          __A						: Any      =       self.get_feature_extractor()
          __A						: List[str]      =       self.get_tokenizer()
          __A						: List[str]      =       self.get_decoder()

          __A						: int      =       WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase						,	feature_extractor=_UpperCAmelCase						,	decoder=_UpperCAmelCase)

          __A						: List[Any]      =       self._get_dummy_logits()

          # note: pool should be instantiated *after* Wav2Vec2ProcessorWithLM.
          #       otherwise, the LM won't be available to the pool's sub-processes.
          # manual logic used to allow parameterized test for both pool=None and pool=Pool(...)
          if pool_context is None:
               __A						: Union[str, Any]      =       processor.batch_decode(_UpperCAmelCase)
          else:
               with get_context(_UpperCAmelCase).Pool() as pool:
                    __A						: List[Any]      =       processor.batch_decode(_UpperCAmelCase						,	_UpperCAmelCase)

          __A						: List[str]      =       list(_UpperCAmelCase)

          with get_context('fork').Pool() as p:
               __A						: List[Any]      =       decoder.decode_beams_batch(_UpperCAmelCase						,	_UpperCAmelCase)

          __A ,__A ,__A						: List[str]      =       [], [], []
          for beams in decoded_beams:
               texts_decoder.append(beams[0][0])
               logit_scores_decoder.append(beams[0][-2])
               lm_scores_decoder.append(beams[0][-1])

          self.assertListEqual(_UpperCAmelCase						,	decoded_processor.text)
          self.assertListEqual(['<s> <s> </s>', '<s> <s> <s>']						,	decoded_processor.text)
          self.assertListEqual(_UpperCAmelCase						,	decoded_processor.logit_score)
          self.assertListEqual(_UpperCAmelCase						,	decoded_processor.lm_score)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[str]      =       self.get_feature_extractor()
          __A						: List[str]      =       self.get_tokenizer()
          __A						: Optional[int]      =       self.get_decoder()

          __A						: Optional[int]      =       WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase						,	feature_extractor=_UpperCAmelCase						,	decoder=_UpperCAmelCase)

          __A						: List[Any]      =       self._get_dummy_logits()

          __A						: Dict      =       15
          __A						: str      =       -20.0
          __A						: Optional[Any]      =       -4.0

          __A						: int      =       processor.batch_decode(
              _UpperCAmelCase						,	beam_width=_UpperCAmelCase						,	beam_prune_logp=_UpperCAmelCase						,	token_min_logp=_UpperCAmelCase						,	)
          __A						: Any      =       decoded_processor_out.text

          __A						: Any      =       list(_UpperCAmelCase)

          with get_context('fork').Pool() as pool:
               __A						: Tuple      =       decoder.decode_beams_batch(
                   _UpperCAmelCase						,	_UpperCAmelCase						,	beam_width=_UpperCAmelCase						,	beam_prune_logp=_UpperCAmelCase						,	token_min_logp=_UpperCAmelCase						,	)

          __A						: List[str]      =       [d[0][0] for d in decoded_decoder_out]
          __A						: List[str]      =       [d[0][2] for d in decoded_decoder_out]
          __A						: Dict      =       [d[0][3] for d in decoded_decoder_out]

          self.assertListEqual(_UpperCAmelCase						,	_UpperCAmelCase)
          self.assertListEqual(['</s> <s> <s>', '<s> <s> <s>']						,	_UpperCAmelCase)

          self.assertTrue(np.array_equal(_UpperCAmelCase						,	decoded_processor_out.logit_score))
          self.assertTrue(np.allclose([-20.054, -18.447]						,	_UpperCAmelCase						,	atol=1e-3))

          self.assertTrue(np.array_equal(_UpperCAmelCase						,	decoded_processor_out.lm_score))
          self.assertTrue(np.allclose([-15.554, -13.9474]						,	_UpperCAmelCase						,	atol=1e-3))


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Optional[Any]      =       self.get_feature_extractor()
          __A						: Union[str, Any]      =       self.get_tokenizer()
          __A						: int      =       self.get_decoder()

          __A						: Optional[int]      =       WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase						,	feature_extractor=_UpperCAmelCase						,	decoder=_UpperCAmelCase)

          __A						: List[Any]      =       self._get_dummy_logits()

          __A						: List[str]      =       2.0
          __A						: Union[str, Any]      =       5.0
          __A						: Optional[Any]      =       -20.0
          __A						: str      =       True

          __A						: str      =       processor.batch_decode(
              _UpperCAmelCase						,	alpha=_UpperCAmelCase						,	beta=_UpperCAmelCase						,	unk_score_offset=_UpperCAmelCase						,	lm_score_boundary=_UpperCAmelCase						,	)
          __A						: str      =       decoded_processor_out.text

          __A						: str      =       list(_UpperCAmelCase)
          decoder.reset_params(
              alpha=_UpperCAmelCase						,	beta=_UpperCAmelCase						,	unk_score_offset=_UpperCAmelCase						,	lm_score_boundary=_UpperCAmelCase						,	)

          with get_context('fork').Pool() as pool:
               __A						: Dict      =       decoder.decode_beams_batch(
                   _UpperCAmelCase						,	_UpperCAmelCase						,	)

          __A						: int      =       [d[0][0] for d in decoded_decoder_out]

          self.assertListEqual(_UpperCAmelCase						,	_UpperCAmelCase)
          self.assertListEqual(['<s> </s> <s> </s> </s>', '</s> </s> <s> </s> </s>']						,	_UpperCAmelCase)
          __A						: Tuple      =       processor.decoder.model_container[processor.decoder._model_key]
          self.assertEqual(lm_model.alpha						,	2.0)
          self.assertEqual(lm_model.beta						,	5.0)
          self.assertEqual(lm_model.unk_score_offset						,	-20.0)
          self.assertEqual(lm_model.score_boundary						,	_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[str]      =       WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm')

          __A						: str      =       processor.decoder.model_container[processor.decoder._model_key]
          __A						: Union[str, Any]      =       Path(language_model._kenlm_model.path.decode('utf-8')).parent.parent.absolute()

          __A						: Any      =       os.listdir(_UpperCAmelCase)
          __A						: List[str]      =       ['alphabet.json', 'language_model']

          downloaded_decoder_files.sort()
          expected_decoder_files.sort()

          # test that only decoder relevant files from
          # https://huggingface.co/hf-internal-testing/processor_with_lm/tree/main
          # are downloaded and none of the rest (e.g. README.md, ...)
          self.assertListEqual(_UpperCAmelCase						,	_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Tuple      =       snapshot_download('hf-internal-testing/processor_with_lm')

          __A						: Optional[int]      =       WavaVecaProcessorWithLM.from_pretrained(_UpperCAmelCase)

          __A						: Optional[Any]      =       processor.decoder.model_container[processor.decoder._model_key]
          __A						: Tuple      =       Path(language_model._kenlm_model.path.decode('utf-8')).parent.parent.absolute()

          __A						: Tuple      =       os.listdir(_UpperCAmelCase)
          __A						: str      =       os.listdir(_UpperCAmelCase)

          local_decoder_files.sort()
          expected_decoder_files.sort()

          # test that both decoder form hub and local files in cache are the same
          self.assertListEqual(_UpperCAmelCase						,	_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[str]      =       WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm')
          __A						: Optional[Any]      =       AutoProcessor.from_pretrained('hf-internal-testing/processor_with_lm')

          __A						: str      =       floats_list((3, 1000))

          __A						: Dict      =       processor_wavaveca(_UpperCAmelCase						,	return_tensors='np')
          __A						: Optional[Any]      =       processor_auto(_UpperCAmelCase						,	return_tensors='np')

          for key in input_wavaveca.keys():
               self.assertAlmostEqual(input_wavaveca[key].sum()						,	input_auto[key].sum()						,	delta=1e-2)

          __A						: str      =       self._get_dummy_logits()

          __A						: Optional[int]      =       processor_wavaveca.batch_decode(_UpperCAmelCase)
          __A						: int      =       processor_auto.batch_decode(_UpperCAmelCase)

          self.assertListEqual(decoded_wavaveca.text						,	decoded_auto.text)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[Any]      =       self.get_feature_extractor()
          __A						: Tuple      =       self.get_tokenizer()
          __A						: Tuple      =       self.get_decoder()

          __A						: Optional[int]      =       WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase						,	feature_extractor=_UpperCAmelCase						,	decoder=_UpperCAmelCase)

          self.assertListEqual(
              processor.model_input_names						,	feature_extractor.model_input_names						,	msg='`processor` and `feature_extractor` model input names do not match'						,	)


     @staticmethod
     def 							SCREAMING_SNAKE_CASE  (						_UpperCAmelCase						,	_UpperCAmelCase):

          '''simple docstring'''






          __A						: List[str]      =       [d[key] for d in offsets]
          return retrieved_list


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[Any]      =       WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm')
          __A						: Dict      =       self._get_dummy_logits()[0]

          __A						: Optional[int]      =       processor.decode(_UpperCAmelCase						,	output_word_offsets=_UpperCAmelCase)
          # check Wav2Vec2CTCTokenizerOutput keys for word
          self.assertEqual(len(outputs.keys())						,	4)
          self.assertTrue('text' in outputs)
          self.assertTrue('word_offsets' in outputs)
          self.assertTrue(isinstance(_UpperCAmelCase						,	_UpperCAmelCase))

          self.assertEqual(' '.join(self.get_from_offsets(outputs['word_offsets']						,	'word'))						,	outputs.text)
          self.assertListEqual(self.get_from_offsets(outputs['word_offsets']						,	'word')						,	['<s>', '<s>', '</s>'])
          self.assertListEqual(self.get_from_offsets(outputs['word_offsets']						,	'start_offset')						,	[0, 2, 4])
          self.assertListEqual(self.get_from_offsets(outputs['word_offsets']						,	'end_offset')						,	[1, 3, 5])


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[Any]      =       WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm')
          __A						: Optional[Any]      =       self._get_dummy_logits()

          __A						: Tuple      =       processor.batch_decode(_UpperCAmelCase						,	output_word_offsets=_UpperCAmelCase)

          # check Wav2Vec2CTCTokenizerOutput keys for word
          self.assertEqual(len(outputs.keys())						,	4)
          self.assertTrue('text' in outputs)
          self.assertTrue('word_offsets' in outputs)
          self.assertTrue(isinstance(_UpperCAmelCase						,	_UpperCAmelCase))

          self.assertListEqual(
              [' '.join(self.get_from_offsets(_UpperCAmelCase						,	'word')) for o in outputs['word_offsets']]						,	outputs.text)
          self.assertListEqual(self.get_from_offsets(outputs['word_offsets'][0]						,	'word')						,	['<s>', '<s>', '</s>'])
          self.assertListEqual(self.get_from_offsets(outputs['word_offsets'][0]						,	'start_offset')						,	[0, 2, 4])
          self.assertListEqual(self.get_from_offsets(outputs['word_offsets'][0]						,	'end_offset')						,	[1, 3, 5])


     @slow
     @require_torch
     @require_torchaudio
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          import torch

          __A						: Dict      =       load_dataset('common_voice'						,	'en'						,	split='train'						,	streaming=_UpperCAmelCase)
          __A						: Any      =       ds.cast_column('audio'						,	datasets.Audio(sampling_rate=1_6000))
          __A						: List[str]      =       iter(_UpperCAmelCase)
          __A						: List[str]      =       next(_UpperCAmelCase)

          __A						: str      =       AutoProcessor.from_pretrained('patrickvonplaten/wav2vec2-base-100h-with-lm')
          __A						: Any      =       WavaVecaForCTC.from_pretrained('patrickvonplaten/wav2vec2-base-100h-with-lm')

          # compare to filename `common_voice_en_100038.mp3` of dataset viewer on https://huggingface.co/datasets/common_voice/viewer/en/train
          __A						: str      =       processor(sample['audio']['array']						,	return_tensors='pt').input_values

          with torch.no_grad():
               __A						: Union[str, Any]      =       model(_UpperCAmelCase).logits.cpu().numpy()

          __A						: List[str]      =       processor.decode(logits[0]						,	output_word_offsets=_UpperCAmelCase)

          __A						: List[Any]      =       model.config.inputs_to_logits_ratio / processor.feature_extractor.sampling_rate
          __A						: Union[str, Any]      =       [
              {
                  'start_time': d['start_offset'] * time_offset,
                  'end_time': d['end_offset'] * time_offset,
                  'word': d['word'],
              }
              for d in output['word_offsets']
          ]

          __A						: List[str]      =       'WHY DOES MILISANDRA LOOK LIKE SHE WANTS TO CONSUME JOHN SNOW ON THE RIVER AT THE WALL'

          # output words
          self.assertEqual(' '.join(self.get_from_offsets(_UpperCAmelCase						,	'word'))						,	_UpperCAmelCase)
          self.assertEqual(' '.join(self.get_from_offsets(_UpperCAmelCase						,	'word'))						,	output.text)

          # output times
          __A						: List[Any]      =       torch.tensor(self.get_from_offsets(_UpperCAmelCase						,	'start_time'))
          __A						: Any      =       torch.tensor(self.get_from_offsets(_UpperCAmelCase						,	'end_time'))

          # fmt: off
          __A						: Optional[int]      =       torch.tensor([1.4199, 1.6599, 2.2599, 3.0, 3.24, 3.5999, 3.7999, 4.0999, 4.26, 4.94, 5.28, 5.6599, 5.78, 5.94, 6.32, 6.5399, 6.6599])
          __A						: str      =       torch.tensor([1.5399, 1.8999, 2.9, 3.16, 3.5399, 3.72, 4.0199, 4.1799, 4.76, 5.1599, 5.5599, 5.6999, 5.86, 6.1999, 6.38, 6.6199, 6.94])
          # fmt: on

          self.assertTrue(torch.allclose(_UpperCAmelCase						,	_UpperCAmelCase						,	atol=0.01))
          self.assertTrue(torch.allclose(_UpperCAmelCase						,	_UpperCAmelCase						,	atol=0.01))






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