

'''simple docstring'''

import gc
import unittest

import numpy as np
import torch

from diffusers import RePaintPipeline, RePaintScheduler, UNetaDModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_image,
    load_numpy,
    nightly,
    require_torch_gpu,
    skip_mps,
    torch_device,
)

from ..pipeline_params import IMAGE_INPAINTING_BATCH_PARAMS, IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()





class 							_snake_case	(    _a		,					unittest.TestCase  ):
   _A    :       int      				=   RePaintPipeline
   _A    :       Optional[Any]      				=   IMAGE_INPAINTING_PARAMS - {'''width''', '''height''', '''guidance_scale'''}
   _A    :       Optional[int]      				=   PipelineTesterMixin.required_optional_params - {
       '''latents''',
       '''num_images_per_prompt''',
       '''callback''',
       '''callback_steps''',
   }
   _A    :       List[Any]      				=   IMAGE_INPAINTING_BATCH_PARAMS
   _A    :       str      				=   False
   def       __UpperCamelCase  (	self     : Tuple							):
     torch.manual_seed(0							)
     torch.manual_seed(0							)
     SCREAMING_SNAKE_CASE:Any   =   UNetaDModel(
         block_out_channels=(32, 64)       ,layers_per_block=2       ,sample_size=32       ,in_channels=3       ,out_channels=3       ,down_block_types=("DownBlock2D", "AttnDownBlock2D")       ,up_block_types=("AttnUpBlock2D", "UpBlock2D")       ,)
     SCREAMING_SNAKE_CASE:int   =   RePaintScheduler()
     SCREAMING_SNAKE_CASE:int   =   {"unet": unet, "scheduler": scheduler}
     return components
   def       __UpperCamelCase  (	self     : Optional[Any]       ,SCREAMING_SNAKE_CASE__     : int       ,SCREAMING_SNAKE_CASE__     : List[str]=0							):
     if str(SCREAMING_SNAKE_CASE__							).startswith("mps"							):
       SCREAMING_SNAKE_CASE:List[Any]   =   torch.manual_seed(SCREAMING_SNAKE_CASE__							)
     else:
       SCREAMING_SNAKE_CASE:Dict   =   torch.Generator(device=SCREAMING_SNAKE_CASE__							).manual_seed(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:int   =   np.random.RandomState(SCREAMING_SNAKE_CASE__							).standard_normal((1, 3, 32, 32)							)
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   torch.from_numpy(SCREAMING_SNAKE_CASE__							).to(device=SCREAMING_SNAKE_CASE__       ,dtype=torch.floataa							)
     SCREAMING_SNAKE_CASE:str   =   (image > 0).to(device=SCREAMING_SNAKE_CASE__       ,dtype=torch.floataa							)
     SCREAMING_SNAKE_CASE:List[Any]   =   {
         "image": image,
         "mask_image": mask,
         "generator": generator,
         "num_inference_steps": 5,
         "eta": 0.0,
         "jump_length": 2,
         "jump_n_sample": 2,
         "output_type": "numpy",
     }
     return inputs
   def       __UpperCamelCase  (	self     : int							):
     SCREAMING_SNAKE_CASE:List[Any]   =   "cpu"  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE:Dict   =   self.get_dummy_components()
     SCREAMING_SNAKE_CASE:int   =   RePaintPipeline(**SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Any   =   sd_pipe.to(SCREAMING_SNAKE_CASE__							)
     sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Any   =   self.get_dummy_inputs(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:int   =   sd_pipe(**SCREAMING_SNAKE_CASE__							).images
     SCREAMING_SNAKE_CASE:List[Any]   =   image[0, -3:, -3:, -1]

     assert image.shape == (1, 32, 32, 3)
     SCREAMING_SNAKE_CASE:Tuple   =   np.array([1.0_000, 0.5_426, 0.5_497, 0.2_200, 1.0_000, 1.0_000, 0.5_623, 1.0_000, 0.6_274]							)

     assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-3
   @skip_mps
   def       __UpperCamelCase  (	self     : List[str]							):
     return super().test_save_load_local()
   @unittest.skip("non-deterministic pipeline"							)
   def       __UpperCamelCase  (	self     : str							):
     return super().test_inference_batch_single_identical()
   @skip_mps
   def       __UpperCamelCase  (	self     : Tuple							):
     return super().test_dict_tuple_outputs_equivalent()
   @skip_mps
   def       __UpperCamelCase  (	self     : Union[str, Any]							):
     return super().test_save_load_optional_components()







   @skip_mps
   def       __UpperCamelCase  (	self     : Tuple							):
     return super().test_attention_slicing_forward_pass()





@nightly
@require_torch_gpu
class 							_snake_case	(    unittest.TestCase  ):
   def       __UpperCamelCase  (	self     : str							):
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()







   def       __UpperCamelCase  (	self     : Any							):
     SCREAMING_SNAKE_CASE:Optional[Any]   =   load_image(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/"
         "repaint/celeba_hq_256.png"							)
     SCREAMING_SNAKE_CASE:Tuple   =   load_image(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/mask_256.png"							)
     SCREAMING_SNAKE_CASE:Tuple   =   load_numpy(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/"
         "repaint/celeba_hq_256_result.npy"							)

     SCREAMING_SNAKE_CASE:List[str]   =   "google/ddpm-ema-celebahq-256"
     SCREAMING_SNAKE_CASE:Dict   =   UNetaDModel.from_pretrained(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Tuple   =   RePaintScheduler.from_pretrained(SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:str   =   RePaintPipeline(unet=SCREAMING_SNAKE_CASE__       ,scheduler=SCREAMING_SNAKE_CASE__							).to(SCREAMING_SNAKE_CASE__							)
     repaint.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)
     repaint.enable_attention_slicing()

     SCREAMING_SNAKE_CASE:List[str]   =   torch.manual_seed(0							)
     SCREAMING_SNAKE_CASE:List[Any]   =   repaint(
         SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,num_inference_steps=250       ,eta=0.0       ,jump_length=10       ,jump_n_sample=10       ,generator=SCREAMING_SNAKE_CASE__       ,output_type="np"       ,)
     SCREAMING_SNAKE_CASE:Tuple   =   output.images[0]

     assert image.shape == (256, 256, 3)
     assert np.abs(expected_image - image							).mean() < 1e-2









'''simple docstring'''

import shutil
import tempfile
import unittest

import pytest

from transformers import WhisperTokenizer, is_speech_available
from transformers.testing_utils import require_sentencepiece, require_torch, require_torchaudio

from .test_feature_extraction_whisper import floats_list


if is_speech_available():
    from transformers import WhisperFeatureExtractor, WhisperProcessor


A_      =							5_03_58
A_      =							5_03_62





@require_torch
@require_torchaudio
@require_sentencepiece
class 							_snake_case	(    unittest.TestCase  ):
   def       __UpperCamelCase  (	self     : Any							):
     SCREAMING_SNAKE_CASE:int   =   "openai/whisper-small.en"
     SCREAMING_SNAKE_CASE:Optional[Any]   =   tempfile.mkdtemp()
   def       __UpperCamelCase  (	self     : Any       ,**SCREAMING_SNAKE_CASE__     : List[str]							):
     return WhisperTokenizer.from_pretrained(self.checkpoint       ,**SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : List[str]       ,**SCREAMING_SNAKE_CASE__     : Dict							):
     return WhisperFeatureExtractor.from_pretrained(self.checkpoint       ,**SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : str							):
     shutil.rmtree(self.tmpdirname							)
   def       __UpperCamelCase  (	self     : Any							):
     SCREAMING_SNAKE_CASE:Dict   =   self.get_tokenizer()
     SCREAMING_SNAKE_CASE:Optional[int]   =   self.get_feature_extractor()

     SCREAMING_SNAKE_CASE:Dict   =   WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE__       ,feature_extractor=SCREAMING_SNAKE_CASE__							)

     processor.save_pretrained(self.tmpdirname							)
     SCREAMING_SNAKE_CASE:List[Any]   =   WhisperProcessor.from_pretrained(self.tmpdirname							)

     self.assertEqual(processor.tokenizer.get_vocab()       ,tokenizer.get_vocab()							)
     self.assertIsInstance(processor.tokenizer       ,SCREAMING_SNAKE_CASE__							)

     self.assertEqual(processor.feature_extractor.to_json_string()       ,feature_extractor.to_json_string()							)
     self.assertIsInstance(processor.feature_extractor       ,SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : Union[str, Any]							):
     SCREAMING_SNAKE_CASE:int   =   WhisperProcessor(tokenizer=self.get_tokenizer()       ,feature_extractor=self.get_feature_extractor()							)
     processor.save_pretrained(self.tmpdirname							)

     SCREAMING_SNAKE_CASE:List[str]   =   self.get_tokenizer(bos_token="(BOS)"       ,eos_token="(EOS)"							)
     SCREAMING_SNAKE_CASE:Tuple   =   self.get_feature_extractor(do_normalize=SCREAMING_SNAKE_CASE__       ,padding_value=1.0							)

     SCREAMING_SNAKE_CASE:Any   =   WhisperProcessor.from_pretrained(
         self.tmpdirname       ,bos_token="(BOS)"       ,eos_token="(EOS)"       ,do_normalize=SCREAMING_SNAKE_CASE__       ,padding_value=1.0							)

     self.assertEqual(processor.tokenizer.get_vocab()       ,tokenizer_add_kwargs.get_vocab()							)
     self.assertIsInstance(processor.tokenizer       ,SCREAMING_SNAKE_CASE__							)

     self.assertEqual(processor.feature_extractor.to_json_string()       ,feature_extractor_add_kwargs.to_json_string()							)
     self.assertIsInstance(processor.feature_extractor       ,SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : List[Any]							):
     SCREAMING_SNAKE_CASE:Any   =   self.get_feature_extractor()
     SCREAMING_SNAKE_CASE:List[Any]   =   self.get_tokenizer()

     SCREAMING_SNAKE_CASE:Optional[Any]   =   WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE__       ,feature_extractor=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Any   =   floats_list((3, 1_000)							)

     SCREAMING_SNAKE_CASE:Any   =   feature_extractor(SCREAMING_SNAKE_CASE__       ,return_tensors="np"							)
     SCREAMING_SNAKE_CASE:List[str]   =   processor(SCREAMING_SNAKE_CASE__       ,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       __UpperCamelCase  (	self     : List[Any]							):
     SCREAMING_SNAKE_CASE:int   =   self.get_feature_extractor()
     SCREAMING_SNAKE_CASE:Tuple   =   self.get_tokenizer()

     SCREAMING_SNAKE_CASE:Union[str, Any]   =   WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE__       ,feature_extractor=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:str   =   "This is a test string"

     SCREAMING_SNAKE_CASE:int   =   processor(text=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:str   =   tokenizer(SCREAMING_SNAKE_CASE__							)

     for key in encoded_tok.keys():
       self.assertListEqual(encoded_tok[key]       ,encoded_processor[key]							)
   def       __UpperCamelCase  (	self     : str							):
     SCREAMING_SNAKE_CASE:List[Any]   =   self.get_feature_extractor()
     SCREAMING_SNAKE_CASE:List[str]   =   self.get_tokenizer()

     SCREAMING_SNAKE_CASE:Union[str, Any]   =   WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE__       ,feature_extractor=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Optional[Any]   =   [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

     SCREAMING_SNAKE_CASE:Any   =   processor.batch_decode(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   tokenizer.batch_decode(SCREAMING_SNAKE_CASE__							)

     self.assertListEqual(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : int							):
     SCREAMING_SNAKE_CASE:Tuple   =   self.get_feature_extractor()
     SCREAMING_SNAKE_CASE:Optional[int]   =   self.get_tokenizer()

     SCREAMING_SNAKE_CASE:Union[str, Any]   =   WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE__       ,feature_extractor=SCREAMING_SNAKE_CASE__							)

     self.assertListEqual(
         processor.model_input_names       ,feature_extractor.model_input_names       ,msg="`processor` and `feature_extractor` model input names do not match"       ,)
   def       __UpperCamelCase  (	self     : int							):
     SCREAMING_SNAKE_CASE:Tuple   =   self.get_feature_extractor()
     SCREAMING_SNAKE_CASE:Any   =   self.get_tokenizer()

     SCREAMING_SNAKE_CASE:Any   =   WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE__       ,feature_extractor=SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Optional[int]   =   processor.get_decoder_prompt_ids(task="transcribe"       ,no_timestamps=SCREAMING_SNAKE_CASE__							)

     self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__							)
     for ids in forced_decoder_ids:
       self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,(list, tuple)							)

     SCREAMING_SNAKE_CASE:Any   =   [TRANSCRIBE, NOTIMESTAMPS]
     self.assertListEqual([ids[-1] for ids in forced_decoder_ids]       ,SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : Dict							):
     SCREAMING_SNAKE_CASE:Tuple   =   WhisperProcessor(tokenizer=self.get_tokenizer()       ,feature_extractor=self.get_feature_extractor()							)
     SCREAMING_SNAKE_CASE:Dict   =   processor.get_prompt_ids("Mr. Quilter"							)
     SCREAMING_SNAKE_CASE:int   =   processor.tokenizer.decode(SCREAMING_SNAKE_CASE__							)

     self.assertListEqual(prompt_ids.tolist()       ,[50_360, 1_770, 13, 2_264, 346, 353]							)
     self.assertEqual(SCREAMING_SNAKE_CASE__       ,"<|startofprev|> Mr. Quilter"							)
   def       __UpperCamelCase  (	self     : Union[str, Any]							):
     SCREAMING_SNAKE_CASE:Dict   =   WhisperProcessor(tokenizer=self.get_tokenizer()       ,feature_extractor=self.get_feature_extractor()							)
     SCREAMING_SNAKE_CASE:List[Any]   =   processor.get_prompt_ids(""							)
     SCREAMING_SNAKE_CASE:Tuple   =   processor.tokenizer.decode(SCREAMING_SNAKE_CASE__							)

     self.assertListEqual(prompt_ids.tolist()       ,[50_360, 220]							)
     self.assertEqual(SCREAMING_SNAKE_CASE__       ,"<|startofprev|> "							)







   def       __UpperCamelCase  (	self     : Union[str, Any]							):
     SCREAMING_SNAKE_CASE:Tuple   =   WhisperProcessor(tokenizer=self.get_tokenizer()       ,feature_extractor=self.get_feature_extractor()							)

     def _test_prompt_error_raised_helper(SCREAMING_SNAKE_CASE__     : int       ,SCREAMING_SNAKE_CASE__     : List[str]							):
       with pytest.raises(SCREAMING_SNAKE_CASE__							) as excinfo:
         processor.get_prompt_ids(SCREAMING_SNAKE_CASE__							)
       SCREAMING_SNAKE_CASE:List[Any]   =   F'''Encountered text in the prompt corresponding to disallowed special token: {special_token}.'''
       self.assertEqual(SCREAMING_SNAKE_CASE__       ,str(excinfo.value							)							)

     _test_prompt_error_raised_helper("<|startofprev|> test"       ,"<|startofprev|>"							)
     _test_prompt_error_raised_helper("test <|notimestamps|>"       ,"<|notimestamps|>"							)
     _test_prompt_error_raised_helper("test <|zh|> test <|transcribe|>"       ,"<|zh|>"							)






