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




  '''simple docstring'''



  def  lowerCAmelCase__ (      self			:       Any  )				->int:
        UpperCAmelCase_            =  '''| <pad> <unk> <s> </s> a b c d e f g h i j k'''.split()
        UpperCAmelCase_            =  dict(zip(_a  ,							range(len(_a  )  )  )  )

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

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

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

        # load decoder from hub
        UpperCAmelCase_            =  '''hf-internal-testing/ngram-beam-search-decoder'''






  def  lowerCAmelCase__ (      self			:       Any  ,							**UpperCAmelCase__			:       Optional[Any]  )				->Dict:
        UpperCAmelCase_            =  self.add_kwargs_tokens_map.copy()
        kwargs.update(_a  )
        return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname  ,							**_a  )






  def  lowerCAmelCase__ (      self			:       int  ,							**UpperCAmelCase__			:       List[Any]  )				->Optional[Any]:
        return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname  ,							**_a  )






  def  lowerCAmelCase__ (      self			:       Optional[int]  ,							**UpperCAmelCase__			:       Union[str, Any]  )				->List[str]:
        return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name  ,							**_a  )






  def  lowerCAmelCase__ (      self			:       Union[str, Any]  )				->List[Any]:
        shutil.rmtree(self.tmpdirname  )






  def  lowerCAmelCase__ (      self			:       int  )				->Optional[Any]:
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=_a  ,							feature_extractor=_a  ,							decoder=_a  )

        processor.save_pretrained(self.tmpdirname  )
        UpperCAmelCase_            =  WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname  )

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

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

        # 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  ,							_a  )






  def  lowerCAmelCase__ (      self			:       Optional[int]  )				->str:
        UpperCAmelCase_            =  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
        UpperCAmelCase_            =  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  lowerCAmelCase__ (      self			:       Optional[int]  )				->int:
        UpperCAmelCase_            =  self.get_tokenizer()
        # add token to trigger raise
        tokenizer.add_tokens(['''xx''']  )
        with self.assertRaisesRegex(_a  ,							'''include'''  ):
              WavaVecaProcessorWithLM(
                  tokenizer=_a  ,							feature_extractor=self.get_feature_extractor()  ,							decoder=self.get_decoder()  )






  def  lowerCAmelCase__ (      self			:       Optional[int]  )				->int:
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=_a  ,							feature_extractor=_a  ,							decoder=_a  )

        UpperCAmelCase_            =  floats_list((3, 1000)  )

        UpperCAmelCase_            =  feature_extractor(_a  ,							return_tensors='''np'''  )
        UpperCAmelCase_            =  processor(_a  ,							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  lowerCAmelCase__ (      self			:       List[str]  )				->List[Any]:
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=_a  ,							feature_extractor=_a  ,							decoder=_a  )

        UpperCAmelCase_            =  '''This is a test string'''

        UpperCAmelCase_            =  processor(text=_a  )

        UpperCAmelCase_            =  tokenizer(_a  )

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






  def  lowerCAmelCase__ (      self			:       Union[str, Any]  ,							UpperCAmelCase__			:       str=(2, 10, 16)  ,							UpperCAmelCase__			:       Any=77  )				->List[Any]:
        np.random.seed(_a  )
        return np.random.rand(*_a  )






  def  lowerCAmelCase__ (      self			:       Union[str, Any]  )				->List[str]:
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=_a  ,							feature_extractor=_a  ,							decoder=_a  )

        UpperCAmelCase_            =  self._get_dummy_logits(shape=(10, 16)  ,							seed=13  )

        UpperCAmelCase_            =  processor.decode(_a  )

        UpperCAmelCase_            =  decoder.decode_beams(_a  )[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  lowerCAmelCase__ (      self			:       Any  ,							UpperCAmelCase__			:       Optional[Any]  )				->List[str]:
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=_a  ,							feature_extractor=_a  ,							decoder=_a  )

        UpperCAmelCase_            =  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:
              UpperCAmelCase_            =  processor.batch_decode(_a  )
        else:
              with get_context(_a  ).Pool() as pool:
                    UpperCAmelCase_            =  processor.batch_decode(_a  ,							_a  )

        UpperCAmelCase_            =  list(_a  )

        with get_context('''fork'''  ).Pool() as p:
              UpperCAmelCase_            =  decoder.decode_beams_batch(_a  ,							_a  )

        UpperCAmelCase_      ,		UpperCAmelCase_      ,		UpperCAmelCase_            =  [], [], []
        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(_a  ,							decoded_processor.text  )
        self.assertListEqual(['''<s> <s> </s>''', '''<s> <s> <s>''']  ,							decoded_processor.text  )
        self.assertListEqual(_a  ,							decoded_processor.logit_score  )
        self.assertListEqual(_a  ,							decoded_processor.lm_score  )






  def  lowerCAmelCase__ (      self			:       Tuple  )				->Tuple:
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=_a  ,							feature_extractor=_a  ,							decoder=_a  )

        UpperCAmelCase_            =  self._get_dummy_logits()

        UpperCAmelCase_            =  15
        UpperCAmelCase_            =  -20.0
        UpperCAmelCase_            =  -4.0

        UpperCAmelCase_            =  processor.batch_decode(
            _a  ,							beam_width=_a  ,							beam_prune_logp=_a  ,							token_min_logp=_a  ,							)
        UpperCAmelCase_            =  decoded_processor_out.text

        UpperCAmelCase_            =  list(_a  )

        with get_context('''fork'''  ).Pool() as pool:
              UpperCAmelCase_            =  decoder.decode_beams_batch(
                  _a  ,							_a  ,							beam_width=_a  ,							beam_prune_logp=_a  ,							token_min_logp=_a  ,							)

        UpperCAmelCase_            =  [d[0][0] for d in decoded_decoder_out]
        UpperCAmelCase_            =  [d[0][2] for d in decoded_decoder_out]
        UpperCAmelCase_            =  [d[0][3] for d in decoded_decoder_out]

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

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

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






  def  lowerCAmelCase__ (      self			:       Tuple  )				->List[Any]:
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=_a  ,							feature_extractor=_a  ,							decoder=_a  )

        UpperCAmelCase_            =  self._get_dummy_logits()

        UpperCAmelCase_            =  2.0
        UpperCAmelCase_            =  5.0
        UpperCAmelCase_            =  -20.0
        UpperCAmelCase_            =  True

        UpperCAmelCase_            =  processor.batch_decode(
            _a  ,							alpha=_a  ,							beta=_a  ,							unk_score_offset=_a  ,							lm_score_boundary=_a  ,							)
        UpperCAmelCase_            =  decoded_processor_out.text

        UpperCAmelCase_            =  list(_a  )
        decoder.reset_params(
            alpha=_a  ,							beta=_a  ,							unk_score_offset=_a  ,							lm_score_boundary=_a  ,							)

        with get_context('''fork'''  ).Pool() as pool:
              UpperCAmelCase_            =  decoder.decode_beams_batch(
                  _a  ,							_a  ,							)

        UpperCAmelCase_            =  [d[0][0] for d in decoded_decoder_out]

        self.assertListEqual(_a  ,							_a  )
        self.assertListEqual(['''<s> </s> <s> </s> </s>''', '''</s> </s> <s> </s> </s>''']  ,							_a  )
        UpperCAmelCase_            =  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  ,							_a  )






  def  lowerCAmelCase__ (      self			:       str  )				->Tuple:
        UpperCAmelCase_            =  WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''  )

        UpperCAmelCase_            =  processor.decoder.model_container[processor.decoder._model_key]
        UpperCAmelCase_            =  Path(language_model._kenlm_model.path.decode('''utf-8'''  )  ).parent.parent.absolute()

        UpperCAmelCase_            =  os.listdir(_a  )
        UpperCAmelCase_            =  ['''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(_a  ,							_a  )






  def  lowerCAmelCase__ (      self			:       Tuple  )				->int:
        UpperCAmelCase_            =  snapshot_download('''hf-internal-testing/processor_with_lm'''  )

        UpperCAmelCase_            =  WavaVecaProcessorWithLM.from_pretrained(_a  )

        UpperCAmelCase_            =  processor.decoder.model_container[processor.decoder._model_key]
        UpperCAmelCase_            =  Path(language_model._kenlm_model.path.decode('''utf-8'''  )  ).parent.parent.absolute()

        UpperCAmelCase_            =  os.listdir(_a  )
        UpperCAmelCase_            =  os.listdir(_a  )

        local_decoder_files.sort()
        expected_decoder_files.sort()

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






  def  lowerCAmelCase__ (      self			:       Optional[Any]  )				->Dict:
        UpperCAmelCase_            =  WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''  )
        UpperCAmelCase_            =  AutoProcessor.from_pretrained('''hf-internal-testing/processor_with_lm'''  )

        UpperCAmelCase_            =  floats_list((3, 1000)  )

        UpperCAmelCase_            =  processor_wavaveca(_a  ,							return_tensors='''np'''  )
        UpperCAmelCase_            =  processor_auto(_a  ,							return_tensors='''np'''  )

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

        UpperCAmelCase_            =  self._get_dummy_logits()

        UpperCAmelCase_            =  processor_wavaveca.batch_decode(_a  )
        UpperCAmelCase_            =  processor_auto.batch_decode(_a  )

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






  def  lowerCAmelCase__ (      self			:       List[str]  )				->int:
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=_a  ,							feature_extractor=_a  ,							decoder=_a  )

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






  @staticmethod
  def  lowerCAmelCase__ (      UpperCAmelCase__			:       Optional[Any]  ,							UpperCAmelCase__			:       Dict  )				->List[str]:
        UpperCAmelCase_            =  [d[key] for d in offsets]
        return retrieved_list






  def  lowerCAmelCase__ (      self			:       Any  )				->List[str]:
        UpperCAmelCase_            =  WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''  )
        UpperCAmelCase_            =  self._get_dummy_logits()[0]

        UpperCAmelCase_            =  processor.decode(_a  ,							output_word_offsets=_a  )
        # 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(_a  ,							_a  )  )

        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  lowerCAmelCase__ (      self			:       Any  )				->Optional[Any]:
        UpperCAmelCase_            =  WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''  )
        UpperCAmelCase_            =  self._get_dummy_logits()

        UpperCAmelCase_            =  processor.batch_decode(_a  ,							output_word_offsets=_a  )

        # 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(_a  ,							_a  )  )

        self.assertListEqual(
            [''' '''.join(self.get_from_offsets(_a  ,							'''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  lowerCAmelCase__ (      self			:       Optional[int]  )				->Tuple:
        import torch

        UpperCAmelCase_            =  load_dataset('''common_voice'''  ,							'''en'''  ,							split='''train'''  ,							streaming=_a  )
        UpperCAmelCase_            =  ds.cast_column('''audio'''  ,							datasets.Audio(sampling_rate=1_6000  )  )
        UpperCAmelCase_            =  iter(_a  )
        UpperCAmelCase_            =  next(_a  )

        UpperCAmelCase_            =  AutoProcessor.from_pretrained('''patrickvonplaten/wav2vec2-base-100h-with-lm'''  )
        UpperCAmelCase_            =  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
        UpperCAmelCase_            =  processor(sample['''audio''']['''array''']  ,							return_tensors='''pt'''  ).input_values

        with torch.no_grad():
              UpperCAmelCase_            =  model(_a  ).logits.cpu().numpy()

        UpperCAmelCase_            =  processor.decode(logits[0]  ,							output_word_offsets=_a  )

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

        UpperCAmelCase_            =  '''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(_a  ,							'''word'''  )  )  ,							_a  )
        self.assertEqual(''' '''.join(self.get_from_offsets(_a  ,							'''word'''  )  )  ,							output.text  )

        # output times
        UpperCAmelCase_            =  torch.tensor(self.get_from_offsets(_a  ,							'''start_time'''  )  )
        UpperCAmelCase_            =  torch.tensor(self.get_from_offsets(_a  ,							'''end_time'''  )  )

        # fmt: off
        UpperCAmelCase_            =  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]  )
        UpperCAmelCase_            =  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(_a  ,							_a  ,							atol=0.01  )  )
        self.assertTrue(torch.allclose(_a  ,							_a  ,							atol=0.01  )  )






'''simple docstring'''
import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







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




  '''simple docstring'''



  def  lowerCAmelCase__ (      self			:       Any  )				->Optional[int]:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()






  def  lowerCAmelCase__ (      self			:       str  )				->List[str]:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_euler'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe([prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=9.0  ,							num_inference_steps=20  ,							output_type='''np'''  )

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array([0.0447, 0.0492, 0.0468, 0.0408, 0.0383, 0.0408, 0.0354, 0.0380, 0.0339]  )

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






  def  lowerCAmelCase__ (      self			:       List[str]  )				->int:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_euler'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe([prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=9.0  ,							num_inference_steps=20  ,							output_type='''np'''  )

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array([0.1237, 0.1320, 0.1438, 0.1359, 0.1390, 0.1132, 0.1277, 0.1175, 0.1112]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() < 5e-1






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->Optional[int]:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_dpmpp_2m'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe(
            [prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=7.5  ,							num_inference_steps=15  ,							output_type='''np'''  ,							use_karras_sigmas=UpperCAmelCase__  ,							)

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array(
            [0.1138_1689, 0.1211_2921, 0.138_9457, 0.1254_9606, 0.124_4964, 0.1083_1517, 0.1156_2866, 0.1086_7816, 0.1049_9048]  )

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





