


import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin







@skip_mps
class        snake_case_   (_UpperCAmelCase	,						_UpperCAmelCase	,						unittest.TestCase      ):
					UpperCAmelCase__					:    List[str]          =			IFPipeline
					UpperCAmelCase__					:    str          =			TEXT_TO_IMAGE_PARAMS - {"""width""", """height""", """latents"""}
					UpperCAmelCase__					:    Optional[int]          =			TEXT_TO_IMAGE_BATCH_PARAMS
					UpperCAmelCase__					:    Dict          =			PipelineTesterMixin.required_optional_params - {"""latents"""}
					def  lowerCamelCase__(      self   :Dict      )					->    Union[str, Any]:
												return self._get_dummy_components()
					def  lowerCamelCase__(      self   :Optional[Any]			,__snake_case   :int			,__snake_case   :Optional[int]=0      )					->    List[Any]:
												if str(UpperCAmelCase_      ).startswith('mps'      ):
																			a__             =     torch.manual_seed(UpperCAmelCase_      )
												else:
																			a__             =     torch.Generator(device=UpperCAmelCase_      ).manual_seed(UpperCAmelCase_      )

												a__             =     {
												    'prompt': 'A painting of a squirrel eating a burger',
												    'generator': generator,
												    'num_inference_steps': 2,
												    'output_type': 'numpy',
												}

												return inputs
					def  lowerCamelCase__(      self   :Tuple      )					->    List[str]:
												self._test_save_load_optional_components()
					@unittest.skipIf(torch_device != 'cuda'			,reason='float16 requires CUDA'      )
					def  lowerCamelCase__(      self   :str      )					->    str:
												# Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
												super().test_save_load_floataa(expected_max_diff=1E-1      )
					def  lowerCamelCase__(      self   :Dict      )					->    List[str]:
												self._test_attention_slicing_forward_pass(expected_max_diff=1E-2      )
					def  lowerCamelCase__(      self   :List[str]      )					->    List[str]:
												self._test_save_load_local()
					def  lowerCamelCase__(      self   :Dict      )					->    str:
												self._test_inference_batch_single_identical(
												    expected_max_diff=1E-2			,)



					@unittest.skipIf(
					    torch_device != 'cuda' or not is_xformers_available()			,reason='XFormers attention is only available with CUDA and `xformers` installed'			,)
					def  lowerCamelCase__(      self   :List[Any]      )					->    int:
												self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3      )







@slow
@require_torch_gpu
class        snake_case_   (unittest.TestCase      ):
					def  lowerCamelCase__(      self   :List[str]      )					->    int:
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()
					def  lowerCamelCase__(      self   :Dict      )					->    Optional[int]:
												# if

												a__             =     IFPipeline.from_pretrained('DeepFloyd/IF-I-XL-v1.0'			,variant='fp16'			,torch_dtype=torch.floataa      )

												a__             =     IFSuperResolutionPipeline.from_pretrained(
												    'DeepFloyd/IF-II-L-v1.0'			,variant='fp16'			,torch_dtype=torch.floataa			,text_encoder=UpperCAmelCase_			,tokenizer=UpperCAmelCase_      )

												# pre compute text embeddings and remove T5 to save memory

												pipe_a.text_encoder.to('cuda'      )

												a__      ,					a__             =     pipe_a.encode_prompt('anime turtle'			,device='cuda'      )

												del pipe_a.tokenizer
												del pipe_a.text_encoder
												gc.collect()

												a__             =     None
												a__             =     None

												pipe_a.enable_model_cpu_offload()
												pipe_a.enable_model_cpu_offload()

												pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )
												pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )

												self._test_if(UpperCAmelCase_			,UpperCAmelCase_			,UpperCAmelCase_			,UpperCAmelCase_      )

												pipe_a.remove_all_hooks()
												pipe_a.remove_all_hooks()

												# img2img

												a__             =     IFImgaImgPipeline(**pipe_a.components      )
												a__             =     IFImgaImgSuperResolutionPipeline(**pipe_a.components      )

												pipe_a.enable_model_cpu_offload()
												pipe_a.enable_model_cpu_offload()

												pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )
												pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )

												self._test_if_imgaimg(UpperCAmelCase_			,UpperCAmelCase_			,UpperCAmelCase_			,UpperCAmelCase_      )

												pipe_a.remove_all_hooks()
												pipe_a.remove_all_hooks()

												# inpainting

												a__             =     IFInpaintingPipeline(**pipe_a.components      )
												a__             =     IFInpaintingSuperResolutionPipeline(**pipe_a.components      )

												pipe_a.enable_model_cpu_offload()
												pipe_a.enable_model_cpu_offload()

												pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )
												pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )

												self._test_if_inpainting(UpperCAmelCase_			,UpperCAmelCase_			,UpperCAmelCase_			,UpperCAmelCase_      )
					def  lowerCamelCase__(      self   :Optional[int]			,__snake_case   :Optional[int]			,__snake_case   :Union[str, Any]			,__snake_case   :str			,__snake_case   :Union[str, Any]      )					->    Any:
												# pipeline 1

												_start_torch_memory_measurement()

												a__             =     torch.Generator(device='cpu'      ).manual_seed(0      )
												a__             =     pipe_a(
												    prompt_embeds=UpperCAmelCase_			,negative_prompt_embeds=UpperCAmelCase_			,num_inference_steps=2			,generator=UpperCAmelCase_			,output_type='np'			,)

												a__             =     output.images[0]

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

												a__             =     torch.cuda.max_memory_allocated()
												assert mem_bytes < 13 * 10**9

												a__             =     load_numpy(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy'      )
												assert_mean_pixel_difference(UpperCAmelCase_			,UpperCAmelCase_      )

												# pipeline 2

												_start_torch_memory_measurement()

												a__             =     torch.Generator(device='cpu'      ).manual_seed(0      )

												a__             =     floats_tensor((1, 3, 64, 64)			,rng=random.Random(0      )      ).to(UpperCAmelCase_      )

												a__             =     pipe_a(
												    prompt_embeds=UpperCAmelCase_			,negative_prompt_embeds=UpperCAmelCase_			,image=UpperCAmelCase_			,generator=UpperCAmelCase_			,num_inference_steps=2			,output_type='np'			,)

												a__             =     output.images[0]

												assert image.shape == (2_56, 2_56, 3)

												a__             =     torch.cuda.max_memory_allocated()
												assert mem_bytes < 4 * 10**9

												a__             =     load_numpy(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy'      )
												assert_mean_pixel_difference(UpperCAmelCase_			,UpperCAmelCase_      )
					def  lowerCamelCase__(      self   :str			,__snake_case   :List[Any]			,__snake_case   :int			,__snake_case   :Optional[Any]			,__snake_case   :Union[str, Any]      )					->    Tuple:
												# pipeline 1

												_start_torch_memory_measurement()

												a__             =     floats_tensor((1, 3, 64, 64)			,rng=random.Random(0      )      ).to(UpperCAmelCase_      )

												a__             =     torch.Generator(device='cpu'      ).manual_seed(0      )

												a__             =     pipe_a(
												    prompt_embeds=UpperCAmelCase_			,negative_prompt_embeds=UpperCAmelCase_			,image=UpperCAmelCase_			,num_inference_steps=2			,generator=UpperCAmelCase_			,output_type='np'			,)

												a__             =     output.images[0]

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

												a__             =     torch.cuda.max_memory_allocated()
												assert mem_bytes < 10 * 10**9

												a__             =     load_numpy(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy'      )
												assert_mean_pixel_difference(UpperCAmelCase_			,UpperCAmelCase_      )

												# pipeline 2

												_start_torch_memory_measurement()

												a__             =     torch.Generator(device='cpu'      ).manual_seed(0      )

												a__             =     floats_tensor((1, 3, 2_56, 2_56)			,rng=random.Random(0      )      ).to(UpperCAmelCase_      )
												a__             =     floats_tensor((1, 3, 64, 64)			,rng=random.Random(0      )      ).to(UpperCAmelCase_      )

												a__             =     pipe_a(
												    prompt_embeds=UpperCAmelCase_			,negative_prompt_embeds=UpperCAmelCase_			,image=UpperCAmelCase_			,original_image=UpperCAmelCase_			,generator=UpperCAmelCase_			,num_inference_steps=2			,output_type='np'			,)

												a__             =     output.images[0]

												assert image.shape == (2_56, 2_56, 3)

												a__             =     torch.cuda.max_memory_allocated()
												assert mem_bytes < 4 * 10**9

												a__             =     load_numpy(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy'      )
												assert_mean_pixel_difference(UpperCAmelCase_			,UpperCAmelCase_      )



					def  lowerCamelCase__(      self   :List[str]			,__snake_case   :Dict			,__snake_case   :List[str]			,__snake_case   :Dict			,__snake_case   :Optional[Any]      )					->    Dict:
												# pipeline 1

												_start_torch_memory_measurement()

												a__             =     floats_tensor((1, 3, 64, 64)			,rng=random.Random(0      )      ).to(UpperCAmelCase_      )
												a__             =     floats_tensor((1, 3, 64, 64)			,rng=random.Random(1      )      ).to(UpperCAmelCase_      )

												a__             =     torch.Generator(device='cpu'      ).manual_seed(0      )
												a__             =     pipe_a(
												    prompt_embeds=UpperCAmelCase_			,negative_prompt_embeds=UpperCAmelCase_			,image=UpperCAmelCase_			,mask_image=UpperCAmelCase_			,num_inference_steps=2			,generator=UpperCAmelCase_			,output_type='np'			,)

												a__             =     output.images[0]

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

												a__             =     torch.cuda.max_memory_allocated()
												assert mem_bytes < 10 * 10**9

												a__             =     load_numpy(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy'      )
												assert_mean_pixel_difference(UpperCAmelCase_			,UpperCAmelCase_      )

												# pipeline 2

												_start_torch_memory_measurement()

												a__             =     torch.Generator(device='cpu'      ).manual_seed(0      )

												a__             =     floats_tensor((1, 3, 64, 64)			,rng=random.Random(0      )      ).to(UpperCAmelCase_      )
												a__             =     floats_tensor((1, 3, 2_56, 2_56)			,rng=random.Random(0      )      ).to(UpperCAmelCase_      )
												a__             =     floats_tensor((1, 3, 2_56, 2_56)			,rng=random.Random(1      )      ).to(UpperCAmelCase_      )

												a__             =     pipe_a(
												    prompt_embeds=UpperCAmelCase_			,negative_prompt_embeds=UpperCAmelCase_			,image=UpperCAmelCase_			,mask_image=UpperCAmelCase_			,original_image=UpperCAmelCase_			,generator=UpperCAmelCase_			,num_inference_steps=2			,output_type='np'			,)

												a__             =     output.images[0]

												assert image.shape == (2_56, 2_56, 3)

												a__             =     torch.cuda.max_memory_allocated()
												assert mem_bytes < 4 * 10**9

												a__             =     load_numpy(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy'      )
												assert_mean_pixel_difference(UpperCAmelCase_			,UpperCAmelCase_      )






def 							__lowercase   ( ):
							torch.cuda.empty_cache()
							torch.cuda.reset_max_memory_allocated()
							torch.cuda.reset_peak_memory_stats()






import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin
from .feature_extraction_wavaveca import WavaVecaFeatureExtractor
from .tokenization_wavaveca import WavaVecaCTCTokenizer




class 						lowercase__  (						_UpperCAmelCase  ):
	A__    :   Union[str, Any]				       ="""Wav2Vec2FeatureExtractor"""
	A__    :   Any				       ="""AutoTokenizer"""
	def __init__(			self							:   List[Any]       ,   UpperCAmelCase_							:   List[str]       ,   UpperCAmelCase_							:   str							):
						super().__init__(UpperCAmelCase_       ,   UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= self.feature_extractor
						SCREAMING_SNAKE_CASE__													= False
	@classmethod
	def      A_				(			cls							:   Union[str, Any]       ,   UpperCAmelCase_							:   Any       ,   **UpperCAmelCase_							:   Optional[Any]							):
						try:
											return super().from_pretrained(UpperCAmelCase_       ,   **UpperCAmelCase_							)
						except OSError:
											warnings.warn(
											    F'Loading a tokenizer inside {cls.__name__} from a config that does not'
											    ' include a `tokenizer_class` attribute is deprecated and will be '
											    'removed in v5. Please add `\'tokenizer_class\': \'Wav2Vec2CTCTokenizer\'`'
											    ' attribute to either your `config.json` or `tokenizer_config.json` '
											    'file to suppress this warning: '       ,   UpperCAmelCase_       ,   )

											SCREAMING_SNAKE_CASE__													= WavaVecaFeatureExtractor.from_pretrained(UpperCAmelCase_       ,   **UpperCAmelCase_							)
											SCREAMING_SNAKE_CASE__													= WavaVecaCTCTokenizer.from_pretrained(UpperCAmelCase_       ,   **UpperCAmelCase_							)

											return cls(feature_extractor=UpperCAmelCase_       ,   tokenizer=UpperCAmelCase_							)
	def __call__(			self							:   Any       ,   *UpperCAmelCase_							:   Dict       ,   **UpperCAmelCase_							:   Any							):
						# For backward compatibility
						if self._in_target_context_manager:
											return self.current_processor(*UpperCAmelCase_       ,   **UpperCAmelCase_							)

						if "raw_speech" in kwargs:
											warnings.warn('Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'							)
											SCREAMING_SNAKE_CASE__													= kwargs.pop('raw_speech'							)
						else:
											SCREAMING_SNAKE_CASE__													= kwargs.pop('audio'       ,   UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= kwargs.pop('sampling_rate'       ,   UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= kwargs.pop('text'       ,   UpperCAmelCase_							)
						if len(UpperCAmelCase_							) > 0:
											SCREAMING_SNAKE_CASE__													= args[0]
											SCREAMING_SNAKE_CASE__													= args[1:]

						if audio is None and text is None:
											raise ValueError('You need to specify either an `audio` or `text` input to process.'							)

						if audio is not None:
											SCREAMING_SNAKE_CASE__													= self.feature_extractor(UpperCAmelCase_       ,   *UpperCAmelCase_       ,   sampling_rate=UpperCAmelCase_       ,   **UpperCAmelCase_							)
						if text is not None:
											SCREAMING_SNAKE_CASE__													= self.tokenizer(UpperCAmelCase_       ,   **UpperCAmelCase_							)

						if text is None:
											return inputs
						elif audio is None:
											return encodings
						else:
											SCREAMING_SNAKE_CASE__													= encodings['input_ids']
											return inputs
	def      A_				(			self							:   List[str]       ,   *UpperCAmelCase_							:   List[Any]       ,   **UpperCAmelCase_							:   Optional[int]							):
						# For backward compatibility
						if self._in_target_context_manager:
											return self.current_processor.pad(*UpperCAmelCase_       ,   **UpperCAmelCase_							)

						SCREAMING_SNAKE_CASE__													= kwargs.pop('input_features'       ,   UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= kwargs.pop('labels'       ,   UpperCAmelCase_							)
						if len(UpperCAmelCase_							) > 0:
											SCREAMING_SNAKE_CASE__													= args[0]
											SCREAMING_SNAKE_CASE__													= args[1:]

						if input_features is not None:
											SCREAMING_SNAKE_CASE__													= self.feature_extractor.pad(UpperCAmelCase_       ,   *UpperCAmelCase_       ,   **UpperCAmelCase_							)
						if labels is not None:
											SCREAMING_SNAKE_CASE__													= self.tokenizer.pad(UpperCAmelCase_       ,   **UpperCAmelCase_							)

						if labels is None:
											return input_features
						elif input_features is None:
											return labels
						else:
											SCREAMING_SNAKE_CASE__													= labels['input_ids']
											return input_features
	def      A_				(			self							:   List[Any]       ,   *UpperCAmelCase_							:   List[str]       ,   **UpperCAmelCase_							:   List[str]							):
						return self.tokenizer.batch_decode(*UpperCAmelCase_       ,   **UpperCAmelCase_							)
	def      A_				(			self							:   Tuple       ,   *UpperCAmelCase_							:   int       ,   **UpperCAmelCase_							:   List[str]							):
						return self.tokenizer.decode(*UpperCAmelCase_       ,   **UpperCAmelCase_							)





	@contextmanager
	def      A_				(			self							:   str							):
						warnings.warn(
						    '`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '
						    'labels by using the argument `text` of the regular `__call__` method (either in the same call as '
						    'your audio inputs, or in a separate call.'							)
						SCREAMING_SNAKE_CASE__													= True
						SCREAMING_SNAKE_CASE__													= self.tokenizer
						yield
						SCREAMING_SNAKE_CASE__													= self.feature_extractor
						SCREAMING_SNAKE_CASE__													= False





