

import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






@skip_mps
class     __magic_name__     (							__UpperCamelCase						,    __UpperCamelCase						,    __UpperCamelCase						,    unittest.TestCase     ):
 UpperCAmelCase				     =    StableDiffusionModelEditingPipeline
 UpperCAmelCase				     =    TEXT_TO_IMAGE_PARAMS
 UpperCAmelCase				     =    TEXT_TO_IMAGE_BATCH_PARAMS
 UpperCAmelCase				     =    TEXT_TO_IMAGE_IMAGE_PARAMS
 UpperCAmelCase				     =    TEXT_TO_IMAGE_IMAGE_PARAMS


 def _UpperCamelCase						(    self		:			Tuple     )     ->   Tuple:
   torch.manual_seed(0     )
   UpperCAmelCase          =	UNetaDConditionModel(
       block_out_channels=(3_2, 6_4)      ,    layers_per_block=2      ,    sample_size=3_2      ,    in_channels=4      ,    out_channels=4      ,    down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")      ,    up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")      ,    cross_attention_dim=3_2      ,    )
   UpperCAmelCase          =	DDIMScheduler()
   torch.manual_seed(0     )
   UpperCAmelCase          =	AutoencoderKL(
       block_out_channels=[3_2, 6_4]      ,    in_channels=3      ,    out_channels=3      ,    down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]      ,    up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]      ,    latent_channels=4      ,    )
   torch.manual_seed(0     )
   UpperCAmelCase          =	CLIPTextConfig(
       bos_token_id=0      ,    eos_token_id=2      ,    hidden_size=3_2      ,    intermediate_size=3_7      ,    layer_norm_eps=1e-0_5      ,    num_attention_heads=4      ,    num_hidden_layers=5      ,    pad_token_id=1      ,    vocab_size=1_0_0_0      ,    )
   UpperCAmelCase          =	CLIPTextModel(a_     )
   UpperCAmelCase          =	CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

   UpperCAmelCase          =	{
       '''unet''': unet,
       '''scheduler''': scheduler,
       '''vae''': vae,
       '''text_encoder''': text_encoder,
       '''tokenizer''': tokenizer,
       '''safety_checker''': None,
       '''feature_extractor''': None,
   }
   return components


 def _UpperCamelCase						(    self		:			Optional[Any]      ,    lowerCAmelCase__		:			Dict      ,    lowerCAmelCase__		:			Tuple=0     )     ->   Optional[Any]:
   UpperCAmelCase          =	torch.manual_seed(a_     )
   UpperCAmelCase          =	{
       '''prompt''': '''A field of roses''',
       '''generator''': generator,
       # Setting height and width to None to prevent OOMs on CPU.
       '''height''': None,
       '''width''': None,
       '''num_inference_steps''': 2,
       '''guidance_scale''': 6.0,
       '''output_type''': '''numpy''',
   }
   return inputs


 def _UpperCamelCase						(    self		:			int     )     ->   List[str]:
   UpperCAmelCase          =	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
   UpperCAmelCase          =	self.get_dummy_components()
   UpperCAmelCase          =	StableDiffusionModelEditingPipeline(**a_     )
   UpperCAmelCase          =	sd_pipe.to(a_     )
   sd_pipe.set_progress_bar_config(disable=a_     )

   UpperCAmelCase          =	self.get_dummy_inputs(a_     )
   UpperCAmelCase          =	sd_pipe(**a_     ).images
   UpperCAmelCase          =	image[0, -3:, -3:, -1]
   assert image.shape == (1, 6_4, 6_4, 3)

   UpperCAmelCase          =	np.array([0.4_755, 0.5_132, 0.4_976, 0.3_904, 0.3_554, 0.4_765, 0.5_139, 0.5_158, 0.4_889]     )

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


 def _UpperCamelCase						(    self		:			Optional[Any]     )     ->   Optional[int]:
   UpperCAmelCase          =	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
   UpperCAmelCase          =	self.get_dummy_components()
   UpperCAmelCase          =	StableDiffusionModelEditingPipeline(**a_     )
   UpperCAmelCase          =	sd_pipe.to(a_     )
   sd_pipe.set_progress_bar_config(disable=a_     )

   UpperCAmelCase          =	self.get_dummy_inputs(a_     )
   UpperCAmelCase          =	'''french fries'''
   UpperCAmelCase          =	sd_pipe(**a_      ,    negative_prompt=a_     )
   UpperCAmelCase          =	output.images
   UpperCAmelCase          =	image[0, -3:, -3:, -1]

   assert image.shape == (1, 6_4, 6_4, 3)

   UpperCAmelCase          =	np.array([0.4_992, 0.5_101, 0.5_004, 0.3_949, 0.3_604, 0.4_735, 0.5_216, 0.5_204, 0.4_913]     )

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


 def _UpperCamelCase						(    self		:			List[str]     )     ->   Optional[Any]:
   UpperCAmelCase          =	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
   UpperCAmelCase          =	self.get_dummy_components()
   UpperCAmelCase          =	EulerAncestralDiscreteScheduler(
       beta_start=0.00_085      ,    beta_end=0.012      ,    beta_schedule="scaled_linear"     )
   UpperCAmelCase          =	StableDiffusionModelEditingPipeline(**a_     )
   UpperCAmelCase          =	sd_pipe.to(a_     )
   sd_pipe.set_progress_bar_config(disable=a_     )

   UpperCAmelCase          =	self.get_dummy_inputs(a_     )
   UpperCAmelCase          =	sd_pipe(**a_     ).images
   UpperCAmelCase          =	image[0, -3:, -3:, -1]

   assert image.shape == (1, 6_4, 6_4, 3)

   UpperCAmelCase          =	np.array([0.4_747, 0.5_372, 0.4_779, 0.4_982, 0.5_543, 0.4_816, 0.5_238, 0.4_904, 0.5_027]     )

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


 def _UpperCamelCase						(    self		:			Any     )     ->   str:
   UpperCAmelCase          =	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
   UpperCAmelCase          =	self.get_dummy_components()
   UpperCAmelCase          =	PNDMScheduler()
   UpperCAmelCase          =	StableDiffusionModelEditingPipeline(**a_     )
   UpperCAmelCase          =	sd_pipe.to(a_     )
   sd_pipe.set_progress_bar_config(disable=a_     )

   UpperCAmelCase          =	self.get_dummy_inputs(a_     )
   # the pipeline does not expect pndm so test if it raises error.
   with self.assertRaises(a_     ):
     UpperCAmelCase          =	sd_pipe(**a_     ).images


 def _UpperCamelCase						(    self		:			int     )     ->   Any:
   super().test_inference_batch_single_identical(expected_max_diff=5e-3     )


 def _UpperCamelCase						(    self		:			Optional[int]     )     ->   List[Any]:
   super().test_attention_slicing_forward_pass(expected_max_diff=5e-3     )






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


 def _UpperCamelCase						(    self		:			List[Any]     )     ->   Optional[int]:
   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()


 def _UpperCamelCase						(    self		:			int      ,    lowerCAmelCase__		:			int=0     )     ->   Optional[int]:
   UpperCAmelCase          =	torch.manual_seed(a_     )
   UpperCAmelCase          =	{
       '''prompt''': '''A field of roses''',
       '''generator''': generator,
       '''num_inference_steps''': 3,
       '''guidance_scale''': 7.5,
       '''output_type''': '''numpy''',
   }
   return inputs


 def _UpperCamelCase						(    self		:			int     )     ->   List[str]:
   UpperCAmelCase          =	'''CompVis/stable-diffusion-v1-4'''
   UpperCAmelCase          =	StableDiffusionModelEditingPipeline.from_pretrained(a_      ,    safety_checker=a_     )
   pipe.to(a_     )
   pipe.set_progress_bar_config(disable=a_     )
   pipe.enable_attention_slicing()

   UpperCAmelCase          =	self.get_inputs()
   UpperCAmelCase          =	pipe(**a_     ).images
   UpperCAmelCase          =	image[0, -3:, -3:, -1].flatten()

   assert image.shape == (1, 5_1_2, 5_1_2, 3)

   UpperCAmelCase          =	np.array(
       [0.6_749_496, 0.6_386_453, 0.51_443_267, 0.66_094_905, 0.61_921_215, 0.5_491_332, 0.5_744_417, 0.58_075_106, 0.5_174_658]     )

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

   # make sure image changes after editing
   pipe.edit_model("A pack of roses"      ,    "A pack of blue roses"     )

   UpperCAmelCase          =	pipe(**a_     ).images
   UpperCAmelCase          =	image[0, -3:, -3:, -1].flatten()

   assert image.shape == (1, 5_1_2, 5_1_2, 3)

   assert np.abs(expected_slice - image_slice     ).max() > 1e-1


 def _UpperCamelCase						(    self		:			Optional[Any]     )     ->   Optional[int]:
   torch.cuda.empty_cache()
   torch.cuda.reset_max_memory_allocated()
   torch.cuda.reset_peak_memory_stats()

   UpperCAmelCase          =	'''CompVis/stable-diffusion-v1-4'''
   UpperCAmelCase          =	DDIMScheduler.from_pretrained(a_      ,    subfolder="scheduler"     )
   UpperCAmelCase          =	StableDiffusionModelEditingPipeline.from_pretrained(
       a_      ,    scheduler=a_      ,    safety_checker=a_     )
   UpperCAmelCase          =	pipe.to(a_     )
   pipe.set_progress_bar_config(disable=a_     )
   pipe.enable_attention_slicing(1     )
   pipe.enable_sequential_cpu_offload()

   UpperCAmelCase          =	self.get_inputs()
   UpperCAmelCase          =	pipe(**a_     )

   UpperCAmelCase          =	torch.cuda.max_memory_allocated()
   # make sure that less than 4.4 GB is allocated
   assert mem_bytes < 4.4 * 1_0**9








import json
import os
import shutil
import tempfile
import unittest

from transformers import is_speech_available, is_torch_available
from transformers.models.speechta import SpeechTaTokenizer
from transformers.testing_utils import get_tests_dir, require_torch
from transformers.utils import FEATURE_EXTRACTOR_NAME


if is_speech_available() and is_torch_available():
				from transformers import SpeechTaFeatureExtractor, SpeechTaProcessor

				from .test_feature_extraction_speechta import floats_list


__A       =get_tests_dir("fixtures/test_sentencepiece_bpe_char.model")

@require_torch
class 						UpperCAmelCase__     (     unittest.TestCase  ):




			'''simple docstring'''






			def 					snake_case__	(  self				:					List[str]					):
					'''simple docstring'''





					__UpperCAmelCase       :	str						=      tempfile.mkdtemp()

					__UpperCAmelCase       :	Optional[int]						=      SpeechTaTokenizer(a_					)
					tokenizer.save_pretrained(self.tmpdirname					)

					__UpperCAmelCase       :	Union[str, Any]						=      {
					    '''feature_size''': 1,
					    '''padding_value''': 0.0,
					    '''sampling_rate''': 1_60_00,
					    '''do_normalize''': False,
					    '''num_mel_bins''': 80,
					    '''hop_length''': 16,
					    '''win_length''': 64,
					    '''win_function''': '''hann_window''',
					    '''fmin''': 80,
					    '''fmax''': 76_00,
					    '''mel_floor''': 1e-10,
					    '''reduction_factor''': 2,
					    '''return_attention_mask''': True,
					}

					__UpperCAmelCase       :	Tuple						=      os.path.join(self.tmpdirname			,			a_					)
					with open(self.feature_extraction_file			,			'''w'''			,			encoding='''utf-8'''					) as fp:
							fp.write(json.dumps(a_					) + '''\n'''					)




			def 					snake_case__	(  self				:					Any			,			**a_				:					Tuple					):
					'''simple docstring'''





					return SpeechTaTokenizer.from_pretrained(self.tmpdirname			,			**a_					)




			def 					snake_case__	(  self				:					Dict			,			**a_				:					str					):
					'''simple docstring'''





					return SpeechTaFeatureExtractor.from_pretrained(self.tmpdirname			,			**a_					)




			def 					snake_case__	(  self				:					List[str]					):
					'''simple docstring'''





					shutil.rmtree(self.tmpdirname					)




			def 					snake_case__	(  self				:					Dict					):
					'''simple docstring'''





					__UpperCAmelCase       :	int						=      self.get_tokenizer()
					__UpperCAmelCase       :	str						=      self.get_feature_extractor()

					__UpperCAmelCase       :	List[Any]						=      SpeechTaProcessor(tokenizer=a_			,			feature_extractor=a_					)

					processor.save_pretrained(self.tmpdirname					)
					__UpperCAmelCase       :	List[str]						=      SpeechTaProcessor.from_pretrained(self.tmpdirname					)

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

					self.assertEqual(processor.feature_extractor.to_json_string()			,			feature_extractor.to_json_string()					)
					self.assertIsInstance(processor.feature_extractor			,			a_					)




			def 					snake_case__	(  self				:					int					):
					'''simple docstring'''





					__UpperCAmelCase       :	Optional[int]						=      SpeechTaProcessor(tokenizer=self.get_tokenizer()			,			feature_extractor=self.get_feature_extractor()					)
					processor.save_pretrained(self.tmpdirname					)

					__UpperCAmelCase       :	Optional[Any]						=      self.get_tokenizer(bos_token='''(BOS)'''			,			eos_token='''(EOS)'''					)
					__UpperCAmelCase       :	List[Any]						=      self.get_feature_extractor(do_normalize=a_			,			padding_value=1.0					)

					__UpperCAmelCase       :	List[str]						=      SpeechTaProcessor.from_pretrained(
					    self.tmpdirname			,			bos_token='''(BOS)'''			,			eos_token='''(EOS)'''			,			do_normalize=a_			,			padding_value=1.0					)

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

					self.assertEqual(processor.feature_extractor.to_json_string()			,			feature_extractor_add_kwargs.to_json_string()					)
					self.assertIsInstance(processor.feature_extractor			,			a_					)




			def 					snake_case__	(  self				:					Union[str, Any]					):
					'''simple docstring'''





					__UpperCAmelCase       :	Optional[Any]						=      self.get_feature_extractor()
					__UpperCAmelCase       :	Union[str, Any]						=      self.get_tokenizer()

					__UpperCAmelCase       :	List[Any]						=      SpeechTaProcessor(tokenizer=a_			,			feature_extractor=a_					)

					__UpperCAmelCase       :	List[Any]						=      floats_list((3, 10_00)					)

					__UpperCAmelCase       :	Tuple						=      feature_extractor(audio=a_			,			return_tensors='''np'''					)
					__UpperCAmelCase       :	List[str]						=      processor(audio=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 					snake_case__	(  self				:					Optional[Any]					):
					'''simple docstring'''





					__UpperCAmelCase       :	List[Any]						=      self.get_feature_extractor()
					__UpperCAmelCase       :	Dict						=      self.get_tokenizer()

					__UpperCAmelCase       :	Optional[Any]						=      SpeechTaProcessor(tokenizer=a_			,			feature_extractor=a_					)

					__UpperCAmelCase       :	Union[str, Any]						=      floats_list((3, 10_00)					)

					__UpperCAmelCase       :	Dict						=      feature_extractor(audio_target=a_			,			return_tensors='''np'''					)
					__UpperCAmelCase       :	str						=      processor(audio_target=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 					snake_case__	(  self				:					Union[str, Any]					):
					'''simple docstring'''





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

					__UpperCAmelCase       :	Optional[Any]						=      SpeechTaProcessor(tokenizer=a_			,			feature_extractor=a_					)

					__UpperCAmelCase       :	List[str]						=      '''This is a test string'''

					__UpperCAmelCase       :	List[Any]						=      processor(text=a_					)
					__UpperCAmelCase       :	Optional[int]						=      tokenizer(a_					)

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




			def 					snake_case__	(  self				:					List[str]					):
					'''simple docstring'''





					__UpperCAmelCase       :	Optional[Any]						=      self.get_feature_extractor()
					__UpperCAmelCase       :	Optional[Any]						=      self.get_tokenizer()

					__UpperCAmelCase       :	List[Any]						=      SpeechTaProcessor(tokenizer=a_			,			feature_extractor=a_					)

					__UpperCAmelCase       :	List[str]						=      '''This is a test string'''

					__UpperCAmelCase       :	Optional[Any]						=      processor(text_target=a_					)
					__UpperCAmelCase       :	Optional[int]						=      tokenizer(a_					)

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




			def 					snake_case__	(  self				:					List[str]					):
					'''simple docstring'''





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

					__UpperCAmelCase       :	int						=      SpeechTaProcessor(tokenizer=a_			,			feature_extractor=a_					)

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

					__UpperCAmelCase       :	Optional[Any]						=      processor.batch_decode(a_					)
					__UpperCAmelCase       :	Any						=      tokenizer.batch_decode(a_					)

					self.assertListEqual(a_			,			a_					)




			def 					snake_case__	(  self				:					str					):
					'''simple docstring'''





					__UpperCAmelCase       :	str						=      self.get_feature_extractor()
					__UpperCAmelCase       :	str						=      self.get_tokenizer()

					__UpperCAmelCase       :	Optional[Any]						=      SpeechTaProcessor(tokenizer=a_			,			feature_extractor=a_					)

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





