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






     """simple docstring"""







     lowerCAmelCase_							   =			IFPipeline
     lowerCAmelCase_							   =			TEXT_TO_IMAGE_PARAMS - {'''width''', '''height''', '''latents'''}
     lowerCAmelCase_							   =			TEXT_TO_IMAGE_BATCH_PARAMS
     lowerCAmelCase_							   =			PipelineTesterMixin.required_optional_params - {'''latents'''}



     def  UpperCAmelCase__							(				self							:							Union[str, Any] ):







       """simple docstring"""




       return self._get_dummy_components()



     def  UpperCAmelCase__							(				self							:							Optional[int]     ,						_A							:							Dict     ,						_A							:							Dict=0 ):







       """simple docstring"""




       if str(_A ).startswith('''mps''' ):
         __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		torch.manual_seed(_A )
       else:
         __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		torch.Generator(device=_A ).manual_seed(_A )

       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		{
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''output_type''': '''numpy''',
       }

       return inputs



     def  UpperCAmelCase__							(				self							:							List[Any] ):







       """simple docstring"""




       self._test_save_load_optional_components()



     @unittest.skipIf(torch_device != '''cuda'''     ,						reason='''float16 requires CUDA''' )
     def  UpperCAmelCase__							(				self							:							Union[str, Any] ):







       """simple docstring"""




       super().test_save_load_floataa(expected_max_diff=1e-1 )



     def  UpperCAmelCase__							(				self							:							Tuple ):







       """simple docstring"""




       self._test_attention_slicing_forward_pass(expected_max_diff=1e-2 )



     def  UpperCAmelCase__							(				self							:							List[str] ):







       """simple docstring"""




       self._test_save_load_local()



     def  UpperCAmelCase__							(				self							:							Tuple ):







       """simple docstring"""




       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  UpperCAmelCase__							(				self							:							Union[str, Any] ):







       """simple docstring"""




       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3 )






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






     """simple docstring"""







     def  UpperCAmelCase__							(				self							:							Any ):







       """simple docstring"""




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



     def  UpperCAmelCase__							(				self							:							Union[str, Any] ):







       """simple docstring"""




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

       __SCREAMING_SNAKE_CASE      :	int					=		IFSuperResolutionPipeline.from_pretrained(
           '''DeepFloyd/IF-II-L-v1.0'''     ,						variant='''fp16'''     ,						torch_dtype=torch.floataa     ,						text_encoder=_A     ,						tokenizer=_A )

       # pre compute text embeddings and remove T5 to save memory

       pipe_a.text_encoder.to('''cuda''' )

       __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE      :	Optional[int]					=		pipe_a.encode_prompt('''anime turtle'''     ,						device='''cuda''' )

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

       __SCREAMING_SNAKE_CASE      :	List[Any]					=		None
       __SCREAMING_SNAKE_CASE      :	str					=		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(_A     ,						_A     ,						_A     ,						_A )

       pipe_a.remove_all_hooks()
       pipe_a.remove_all_hooks()

       # img2img

       __SCREAMING_SNAKE_CASE      :	Dict					=		IFImgaImgPipeline(**pipe_a.components )
       __SCREAMING_SNAKE_CASE      :	int					=		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(_A     ,						_A     ,						_A     ,						_A )

       pipe_a.remove_all_hooks()
       pipe_a.remove_all_hooks()

       # inpainting

       __SCREAMING_SNAKE_CASE      :	int					=		IFInpaintingPipeline(**pipe_a.components )
       __SCREAMING_SNAKE_CASE      :	Dict					=		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(_A     ,						_A     ,						_A     ,						_A )



     def  UpperCAmelCase__							(				self							:							Optional[Any]     ,						_A							:							str     ,						_A							:							Optional[Any]     ,						_A							:							Tuple     ,						_A							:							List[str] ):







       """simple docstring"""




       _start_torch_memory_measurement()

       __SCREAMING_SNAKE_CASE      :	List[Any]					=		torch.Generator(device='''cpu''' ).manual_seed(0 )
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		pipe_a(
           prompt_embeds=_A     ,						negative_prompt_embeds=_A     ,						num_inference_steps=2     ,						generator=_A     ,						output_type='''np'''     ,						)

       __SCREAMING_SNAKE_CASE      :	int					=		output.images[0]

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

       __SCREAMING_SNAKE_CASE      :	Dict					=		torch.cuda.max_memory_allocated()
       assert mem_bytes < 13 * 10**9

       __SCREAMING_SNAKE_CASE      :	List[str]					=		load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy''' )
       assert_mean_pixel_difference(_A     ,						_A )

       # pipeline 2

       _start_torch_memory_measurement()

       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		torch.Generator(device='''cpu''' ).manual_seed(0 )

       __SCREAMING_SNAKE_CASE      :	Dict					=		floats_tensor((1, 3, 64, 64)     ,						rng=random.Random(0 ) ).to(_A )

       __SCREAMING_SNAKE_CASE      :	int					=		pipe_a(
           prompt_embeds=_A     ,						negative_prompt_embeds=_A     ,						image=_A     ,						generator=_A     ,						num_inference_steps=2     ,						output_type='''np'''     ,						)

       __SCREAMING_SNAKE_CASE      :	int					=		output.images[0]

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

       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 10**9

       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy''' )
       assert_mean_pixel_difference(_A     ,						_A )



     def  UpperCAmelCase__							(				self							:							Dict     ,						_A							:							Optional[Any]     ,						_A							:							List[Any]     ,						_A							:							Optional[int]     ,						_A							:							Union[str, Any] ):







       """simple docstring"""




       _start_torch_memory_measurement()

       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		floats_tensor((1, 3, 64, 64)     ,						rng=random.Random(0 ) ).to(_A )

       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		torch.Generator(device='''cpu''' ).manual_seed(0 )

       __SCREAMING_SNAKE_CASE      :	List[Any]					=		pipe_a(
           prompt_embeds=_A     ,						negative_prompt_embeds=_A     ,						image=_A     ,						num_inference_steps=2     ,						generator=_A     ,						output_type='''np'''     ,						)

       __SCREAMING_SNAKE_CASE      :	Any					=		output.images[0]

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

       __SCREAMING_SNAKE_CASE      :	List[str]					=		torch.cuda.max_memory_allocated()
       assert mem_bytes < 10 * 10**9

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

       # pipeline 2

       _start_torch_memory_measurement()

       __SCREAMING_SNAKE_CASE      :	Tuple					=		torch.Generator(device='''cpu''' ).manual_seed(0 )

       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		floats_tensor((1, 3, 256, 256)     ,						rng=random.Random(0 ) ).to(_A )
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		floats_tensor((1, 3, 64, 64)     ,						rng=random.Random(0 ) ).to(_A )

       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		pipe_a(
           prompt_embeds=_A     ,						negative_prompt_embeds=_A     ,						image=_A     ,						original_image=_A     ,						generator=_A     ,						num_inference_steps=2     ,						output_type='''np'''     ,						)

       __SCREAMING_SNAKE_CASE      :	List[Any]					=		output.images[0]

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

       __SCREAMING_SNAKE_CASE      :	Dict					=		torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 10**9

       __SCREAMING_SNAKE_CASE      :	List[Any]					=		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(_A     ,						_A )




     def  UpperCAmelCase__							(				self							:							Optional[int]     ,						_A							:							List[str]     ,						_A							:							List[str]     ,						_A							:							Any     ,						_A							:							Dict ):







       """simple docstring"""




       _start_torch_memory_measurement()

       __SCREAMING_SNAKE_CASE      :	Tuple					=		floats_tensor((1, 3, 64, 64)     ,						rng=random.Random(0 ) ).to(_A )
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		floats_tensor((1, 3, 64, 64)     ,						rng=random.Random(1 ) ).to(_A )

       __SCREAMING_SNAKE_CASE      :	int					=		torch.Generator(device='''cpu''' ).manual_seed(0 )
       __SCREAMING_SNAKE_CASE      :	Tuple					=		pipe_a(
           prompt_embeds=_A     ,						negative_prompt_embeds=_A     ,						image=_A     ,						mask_image=_A     ,						num_inference_steps=2     ,						generator=_A     ,						output_type='''np'''     ,						)

       __SCREAMING_SNAKE_CASE      :	List[Any]					=		output.images[0]

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

       __SCREAMING_SNAKE_CASE      :	int					=		torch.cuda.max_memory_allocated()
       assert mem_bytes < 10 * 10**9

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

       # pipeline 2

       _start_torch_memory_measurement()

       __SCREAMING_SNAKE_CASE      :	Dict					=		torch.Generator(device='''cpu''' ).manual_seed(0 )

       __SCREAMING_SNAKE_CASE      :	int					=		floats_tensor((1, 3, 64, 64)     ,						rng=random.Random(0 ) ).to(_A )
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		floats_tensor((1, 3, 256, 256)     ,						rng=random.Random(0 ) ).to(_A )
       __SCREAMING_SNAKE_CASE      :	Tuple					=		floats_tensor((1, 3, 256, 256)     ,						rng=random.Random(1 ) ).to(_A )

       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		pipe_a(
           prompt_embeds=_A     ,						negative_prompt_embeds=_A     ,						image=_A     ,						mask_image=_A     ,						original_image=_A     ,						generator=_A     ,						num_inference_steps=2     ,						output_type='''np'''     ,						)

       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		output.images[0]

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

       __SCREAMING_SNAKE_CASE      :	List[str]					=		torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 10**9

       __SCREAMING_SNAKE_CASE      :	List[Any]					=		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(_A     ,						_A )






def 							a__					(							):





  """simple docstring"""
  torch.cuda.empty_cache()
  torch.cuda.reset_max_memory_allocated()
  torch.cuda.reset_peak_memory_stats()







from typing import List, Optional

import numpy as np

from ...processing_utils import ProcessorMixin
from ...utils import to_numpy

class     snake_case_		(   a							):






 '''simple docstring'''





 __UpperCamelCase       =      'EncodecFeatureExtractor'
 __UpperCamelCase       =      ('T5Tokenizer', 'T5TokenizerFast')

 def __init__(  self,		A_,		A_		)	->       Optional[int]:
    super().__init__(A_,		A_		)
    UpperCAmelCase__						=self.feature_extractor
    UpperCAmelCase__						=False

 def  __UpperCAmelCase						(  self,		A_=None,		A_=None,		A_=True		)	->       Union[str, Any]:
    return self.tokenizer.get_decoder_prompt_ids(task=A_,		language=A_,		no_timestamps=A_		)

 def __call__(  self,		*A_,		**A_		)	->       Any:
    # For backward compatibility
    if self._in_target_context_manager:
       return self.current_processor(*A_,		**A_		)

    UpperCAmelCase__						=kwargs.pop("audio",		A_		)
    UpperCAmelCase__						=kwargs.pop("sampling_rate",		A_		)
    UpperCAmelCase__						=kwargs.pop("text",		A_		)
    if len(A_		) > 0:
       UpperCAmelCase__						=args[0]
       UpperCAmelCase__						=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 text is not None:
       UpperCAmelCase__						=self.tokenizer(A_,		**A_		)

    if audio is not None:
       UpperCAmelCase__						=self.feature_extractor(A_,		*A_,		sampling_rate=A_,		**A_		)

    if audio is None:
       return inputs

    elif text is None:
       return audio_inputs

    else:
       UpperCAmelCase__						=audio_inputs["input_values"]
       if "padding_mask" in audio_inputs:
          UpperCAmelCase__						=audio_inputs["padding_mask"]
       return inputs

 def  __UpperCAmelCase						(  self,		*A_,		**A_		)	->       Dict:
    UpperCAmelCase__						=kwargs.pop("audio",		A_		)
    UpperCAmelCase__						=kwargs.pop("padding_mask",		A_		)

    if len(A_		) > 0:
       UpperCAmelCase__						=args[0]
       UpperCAmelCase__						=args[1:]

    if audio_values is not None:
       return self._decode_audio(A_,		padding_mask=A_		)
    else:
       return self.tokenizer.batch_decode(*A_,		**A_		)

 def  __UpperCAmelCase						(  self,		*A_,		**A_		)	->       int:
    return self.tokenizer.decode(*A_,		**A_		)






 def  __UpperCAmelCase						(  self,		A_,		A_ = None		)	->       List[np.ndarray]:
    UpperCAmelCase__						=to_numpy(A_		)
    UpperCAmelCase__			,							UpperCAmelCase__			,							UpperCAmelCase__						=audio_values.shape

    if padding_mask is None:
       return list(A_		)

    UpperCAmelCase__						=to_numpy(A_		)

    # match the sequence length of the padding mask to the generated audio arrays by padding with the **non-padding**
    # token (so that the generated audio values are **not** treated as padded tokens)
    UpperCAmelCase__						=seq_len - padding_mask.shape[-1]
    UpperCAmelCase__						=1 - self.feature_extractor.padding_value
    UpperCAmelCase__						=np.pad(A_,		((0, 0), (0, difference)),		"constant",		constant_values=A_		)

    UpperCAmelCase__						=audio_values.tolist()
    for i in range(A_		):
       UpperCAmelCase__						=np.asarray(audio_values[i]		)[
           padding_mask[i][None, :] != self.feature_extractor.padding_value
       ]
       UpperCAmelCase__						=sliced_audio.reshape(A_,		-1		)

    return audio_values



