



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 						UpperCAmelCase__     (     __UpperCamelCase		,__UpperCamelCase		,__UpperCamelCase		,unittest.TestCase  ):




			'''simple docstring'''






			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 					snake_case__	(  self				:					Tuple					):
					'''simple docstring'''





					torch.manual_seed(0					)
					__UpperCAmelCase       :	str						=      UNetaDConditionModel(
					    block_out_channels=(32, 64)			,			layers_per_block=2			,			sample_size=32			,			in_channels=4			,			out_channels=4			,			down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')			,			up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')			,			cross_attention_dim=32			,			)
					__UpperCAmelCase       :	Dict						=      DDIMScheduler()
					torch.manual_seed(0					)
					__UpperCAmelCase       :	Any						=      AutoencoderKL(
					    block_out_channels=[32, 64]			,			in_channels=3			,			out_channels=3			,			down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']			,			up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']			,			latent_channels=4			,			)
					torch.manual_seed(0					)
					__UpperCAmelCase       :	Optional[Any]						=      CLIPTextConfig(
					    bos_token_id=0			,			eos_token_id=2			,			hidden_size=32			,			intermediate_size=37			,			layer_norm_eps=1e-05			,			num_attention_heads=4			,			num_hidden_layers=5			,			pad_token_id=1			,			vocab_size=10_00			,			)
					__UpperCAmelCase       :	Any						=      CLIPTextModel(a_					)
					__UpperCAmelCase       :	Optional[int]						=      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

					__UpperCAmelCase       :	Optional[Any]						=      {
					    '''unet''': unet,
					    '''scheduler''': scheduler,
					    '''vae''': vae,
					    '''text_encoder''': text_encoder,
					    '''tokenizer''': tokenizer,
					    '''safety_checker''': None,
					    '''feature_extractor''': None,
					}
					return components




			def 					snake_case__	(  self				:					Optional[Any]			,			a_				:					Dict			,			a_				:					Tuple=0					):
					'''simple docstring'''





					__UpperCAmelCase       :	str						=      torch.manual_seed(a_					)
					__UpperCAmelCase       :	int						=      {
					    '''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 					snake_case__	(  self				:					int					):
					'''simple docstring'''





					__UpperCAmelCase       :	Union[str, Any]						=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
					__UpperCAmelCase       :	Any						=      self.get_dummy_components()
					__UpperCAmelCase       :	str						=      StableDiffusionModelEditingPipeline(**a_					)
					__UpperCAmelCase       :	Optional[int]						=      sd_pipe.to(a_					)
					sd_pipe.set_progress_bar_config(disable=a_					)

					__UpperCAmelCase       :	Tuple						=      self.get_dummy_inputs(a_					)
					__UpperCAmelCase       :	List[str]						=      sd_pipe(**a_					).images
					__UpperCAmelCase       :	Union[str, Any]						=      image[0, -3:, -3:, -1]
					assert image.shape == (1, 64, 64, 3)

					__UpperCAmelCase       :	str						=      np.array([0.4_7_5_5, 0.5_1_3_2, 0.4_9_7_6, 0.3_9_0_4, 0.3_5_5_4, 0.4_7_6_5, 0.5_1_3_9, 0.5_1_5_8, 0.4_8_8_9]					)

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




			def 					snake_case__	(  self				:					Optional[Any]					):
					'''simple docstring'''





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

					__UpperCAmelCase       :	Union[str, Any]						=      self.get_dummy_inputs(a_					)
					__UpperCAmelCase       :	Optional[int]						=      '''french fries'''
					__UpperCAmelCase       :	Optional[int]						=      sd_pipe(**a_			,			negative_prompt=a_					)
					__UpperCAmelCase       :	Dict						=      output.images
					__UpperCAmelCase       :	Tuple						=      image[0, -3:, -3:, -1]

					assert image.shape == (1, 64, 64, 3)

					__UpperCAmelCase       :	Union[str, Any]						=      np.array([0.4_9_9_2, 0.5_1_0_1, 0.5_0_0_4, 0.3_9_4_9, 0.3_6_0_4, 0.4_7_3_5, 0.5_2_1_6, 0.5_2_0_4, 0.4_9_1_3]					)

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




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





					__UpperCAmelCase       :	str						=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
					__UpperCAmelCase       :	Tuple						=      self.get_dummy_components()
					__UpperCAmelCase       :	Optional[int]						=      EulerAncestralDiscreteScheduler(
					    beta_start=0.0_0_0_8_5			,			beta_end=0.0_1_2			,			beta_schedule='''scaled_linear'''					)
					__UpperCAmelCase       :	int						=      StableDiffusionModelEditingPipeline(**a_					)
					__UpperCAmelCase       :	Dict						=      sd_pipe.to(a_					)
					sd_pipe.set_progress_bar_config(disable=a_					)

					__UpperCAmelCase       :	Optional[Any]						=      self.get_dummy_inputs(a_					)
					__UpperCAmelCase       :	Tuple						=      sd_pipe(**a_					).images
					__UpperCAmelCase       :	Any						=      image[0, -3:, -3:, -1]

					assert image.shape == (1, 64, 64, 3)

					__UpperCAmelCase       :	Union[str, Any]						=      np.array([0.4_7_4_7, 0.5_3_7_2, 0.4_7_7_9, 0.4_9_8_2, 0.5_5_4_3, 0.4_8_1_6, 0.5_2_3_8, 0.4_9_0_4, 0.5_0_2_7]					)

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




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





					__UpperCAmelCase       :	Optional[int]						=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
					__UpperCAmelCase       :	int						=      self.get_dummy_components()
					__UpperCAmelCase       :	Union[str, Any]						=      PNDMScheduler()
					__UpperCAmelCase       :	Optional[int]						=      StableDiffusionModelEditingPipeline(**a_					)
					__UpperCAmelCase       :	List[str]						=      sd_pipe.to(a_					)
					sd_pipe.set_progress_bar_config(disable=a_					)

					__UpperCAmelCase       :	Union[str, Any]						=      self.get_dummy_inputs(a_					)
					# the pipeline does not expect pndm so test if it raises error.
					with self.assertRaises(a_					):
							__UpperCAmelCase       :	List[str]						=      sd_pipe(**a_					).images




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





					super().test_inference_batch_single_identical(expected_max_diff=5e-3					)




			def 					snake_case__	(  self				:					Optional[int]					):
					'''simple docstring'''





					super().test_attention_slicing_forward_pass(expected_max_diff=5e-3					)

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




			'''simple docstring'''






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





					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()




			def 					snake_case__	(  self				:					int			,			a_				:					int=0					):
					'''simple docstring'''





					__UpperCAmelCase       :	int						=      torch.manual_seed(a_					)
					__UpperCAmelCase       :	List[Any]						=      {
					    '''prompt''': '''A field of roses''',
					    '''generator''': generator,
					    '''num_inference_steps''': 3,
					    '''guidance_scale''': 7.5,
					    '''output_type''': '''numpy''',
					}
					return inputs




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





					__UpperCAmelCase       :	Union[str, Any]						=      '''CompVis/stable-diffusion-v1-4'''
					__UpperCAmelCase       :	List[str]						=      StableDiffusionModelEditingPipeline.from_pretrained(a_			,			safety_checker=a_					)
					pipe.to(a_					)
					pipe.set_progress_bar_config(disable=a_					)
					pipe.enable_attention_slicing()

					__UpperCAmelCase       :	str						=      self.get_inputs()
					__UpperCAmelCase       :	Dict						=      pipe(**a_					).images
					__UpperCAmelCase       :	List[Any]						=      image[0, -3:, -3:, -1].flatten()

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

					__UpperCAmelCase       :	List[str]						=      np.array(
					    [0.6_7_4_9_4_9_6, 0.6_3_8_6_4_5_3, 0.5_1_4_4_3_2_6_7, 0.6_6_0_9_4_9_0_5, 0.6_1_9_2_1_2_1_5, 0.5_4_9_1_3_3_2, 0.5_7_4_4_4_1_7, 0.5_8_0_7_5_1_0_6, 0.5_1_7_4_6_5_8]					)

					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       :	Optional[int]						=      pipe(**a_					).images
					__UpperCAmelCase       :	Optional[Any]						=      image[0, -3:, -3:, -1].flatten()

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

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




			def 					snake_case__	(  self				:					Optional[Any]					):
					'''simple docstring'''





					torch.cuda.empty_cache()
					torch.cuda.reset_max_memory_allocated()
					torch.cuda.reset_peak_memory_stats()

					__UpperCAmelCase       :	Optional[Any]						=      '''CompVis/stable-diffusion-v1-4'''
					__UpperCAmelCase       :	List[str]						=      DDIMScheduler.from_pretrained(a_			,			subfolder='''scheduler'''					)
					__UpperCAmelCase       :	Tuple						=      StableDiffusionModelEditingPipeline.from_pretrained(
					    a_			,			scheduler=a_			,			safety_checker=a_					)
					__UpperCAmelCase       :	Tuple						=      pipe.to(a_					)
					pipe.set_progress_bar_config(disable=a_					)
					pipe.enable_attention_slicing(1					)
					pipe.enable_sequential_cpu_offload()

					__UpperCAmelCase       :	str						=      self.get_inputs()
					__UpperCAmelCase       :	int						=      pipe(**a_					)

					__UpperCAmelCase       :	int						=      torch.cuda.max_memory_allocated()
					# make sure that less than 4.4 GB is allocated
					assert mem_bytes < 4.4 * 10**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'''			,			)





