






"""simple docstring"""



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       ( _A   ,						_A   ,						unittest.TestCase							):
					SCREAMING_SNAKE_CASE_		:				List[Any]							    =   IFPipeline
					SCREAMING_SNAKE_CASE_		:				Optional[int]							    =   TEXT_TO_IMAGE_PARAMS - {"width", "height", "latents"}
					SCREAMING_SNAKE_CASE_		:				int							    =   TEXT_TO_IMAGE_BATCH_PARAMS
					SCREAMING_SNAKE_CASE_		:				Any							    =   PipelineTesterMixin.required_optional_params - {"latents"}




					def   A  (				self     :						Optional[int]   )  ->  Tuple:
							return self._get_dummy_components()




					def   A  (				self     :						Dict    ,						A     :						Union[str, Any]    ,						A     :						Optional[Any]=0   )  ->  Optional[Any]:
							if str(A   ).startswith('''mps'''   ):
									lowercase_ : Optional[Any]							       =   torch.manual_seed(A   )
							else:
									lowercase_ : Optional[int]							       =   torch.Generator(device=A   ).manual_seed(A   )

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

							return inputs




					def   A  (				self     :						int   )  ->  Dict:
							self._test_save_load_optional_components()




					@unittest.skipIf(torch_device != '''cuda'''    ,						reason='''float16 requires CUDA'''   )
					def   A  (				self     :						Union[str, Any]   )  ->  int:
							# 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   A  (				self     :						int   )  ->  Union[str, Any]:
							self._test_attention_slicing_forward_pass(expected_max_diff=1e-2   )




					def   A  (				self     :						int   )  ->  Optional[int]:
							self._test_save_load_local()




					def   A  (				self     :						Tuple   )  ->  Optional[Any]:
							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   A  (				self     :						Tuple   )  ->  Any:
							self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3   )





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




					def   A  (				self     :						Any   )  ->  Optional[Any]:
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()




					def   A  (				self     :						Tuple   )  ->  Optional[Any]:
							# if

							lowercase_ : Union[str, Any]							       =   IFPipeline.from_pretrained('''DeepFloyd/IF-I-XL-v1.0'''    ,						variant='''fp16'''    ,						torch_dtype=torch.floataa   )

							lowercase_ : str							       =   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'''   )

							lowercase_       ,					lowercase_ : List[str]							       =   pipe_a.encode_prompt('''anime turtle'''    ,						device='''cuda'''   )

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

							lowercase_ : Optional[int]							       =   None
							lowercase_ : List[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

							lowercase_ : List[Any]							       =   IFImgaImgPipeline(**pipe_a.components   )
							lowercase_ : Optional[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

							lowercase_ : str							       =   IFInpaintingPipeline(**pipe_a.components   )
							lowercase_ : int							       =   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   A  (				self     :						Optional[Any]    ,						A     :						int    ,						A     :						List[str]    ,						A     :						Optional[int]    ,						A     :						List[Any]   )  ->  str:
							# pipeline 1

							_start_torch_memory_measurement()

							lowercase_ : Optional[int]							       =   torch.Generator(device='''cpu'''   ).manual_seed(0   )
							lowercase_ : Optional[int]							       =   pipe_a(
							    prompt_embeds=A    ,						negative_prompt_embeds=A    ,						num_inference_steps=2    ,						generator=A    ,						output_type='''np'''    ,						)

							lowercase_ : Optional[Any]							       =   output.images[0]

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

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

							lowercase_ : 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()

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

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

							lowercase_ : str							       =   pipe_a(
							    prompt_embeds=A    ,						negative_prompt_embeds=A    ,						image=A    ,						generator=A    ,						num_inference_steps=2    ,						output_type='''np'''    ,						)

							lowercase_ : Optional[int]							       =   output.images[0]

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

							lowercase_ : Optional[Any]							       =   torch.cuda.max_memory_allocated()
							assert mem_bytes < 4 * 10**9

							lowercase_ : 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   A  (				self     :						List[Any]    ,						A     :						str    ,						A     :						Tuple    ,						A     :						Optional[int]    ,						A     :						Optional[Any]   )  ->  int:
							# pipeline 1

							_start_torch_memory_measurement()

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

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

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

							lowercase_ : Any							       =   output.images[0]

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

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

							lowercase_ : Dict							       =   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()

							lowercase_ : str							       =   torch.Generator(device='''cpu'''   ).manual_seed(0   )

							lowercase_ : Any							       =   floats_tensor((1, 3, 2_56, 2_56)    ,						rng=random.Random(0   )   ).to(A   )
							lowercase_ : int							       =   floats_tensor((1, 3, 64, 64)    ,						rng=random.Random(0   )   ).to(A   )

							lowercase_ : Any							       =   pipe_a(
							    prompt_embeds=A    ,						negative_prompt_embeds=A    ,						image=A    ,						original_image=A    ,						generator=A    ,						num_inference_steps=2    ,						output_type='''np'''    ,						)

							lowercase_ : Optional[int]							       =   output.images[0]

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

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

							lowercase_ : Tuple							       =   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   A  (				self     :						str    ,						A     :						int    ,						A     :						Any    ,						A     :						Tuple    ,						A     :						Dict   )  ->  List[Any]:
							# pipeline 1

							_start_torch_memory_measurement()

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

							lowercase_ : int							       =   torch.Generator(device='''cpu'''   ).manual_seed(0   )
							lowercase_ : Union[str, Any]							       =   pipe_a(
							    prompt_embeds=A    ,						negative_prompt_embeds=A    ,						image=A    ,						mask_image=A    ,						num_inference_steps=2    ,						generator=A    ,						output_type='''np'''    ,						)

							lowercase_ : Tuple							       =   output.images[0]

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

							lowercase_ : str							       =   torch.cuda.max_memory_allocated()
							assert mem_bytes < 10 * 10**9

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

							lowercase_ : int							       =   torch.Generator(device='''cpu'''   ).manual_seed(0   )

							lowercase_ : List[Any]							       =   floats_tensor((1, 3, 64, 64)    ,						rng=random.Random(0   )   ).to(A   )
							lowercase_ : List[Any]							       =   floats_tensor((1, 3, 2_56, 2_56)    ,						rng=random.Random(0   )   ).to(A   )
							lowercase_ : str							       =   floats_tensor((1, 3, 2_56, 2_56)    ,						rng=random.Random(1   )   ).to(A   )

							lowercase_ : List[str]							       =   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'''    ,						)

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

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

							lowercase_ : Optional[Any]							       =   torch.cuda.max_memory_allocated()
							assert mem_bytes < 4 * 10**9

							lowercase_ : int							       =   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 						lowercase    (				):
		torch.cuda.empty_cache()
		torch.cuda.reset_max_memory_allocated()
		torch.cuda.reset_peak_memory_stats()











"""simple docstring"""



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


__A		:       List[Any]  				=		{
    '''configuration_mega''': ['''MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''MegaConfig''', '''MegaOnnxConfig'''],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__A		:       List[str]  				=		[
						    '''MEGA_PRETRAINED_MODEL_ARCHIVE_LIST''',
						    '''MegaForCausalLM''',
						    '''MegaForMaskedLM''',
						    '''MegaForMultipleChoice''',
						    '''MegaForQuestionAnswering''',
						    '''MegaForSequenceClassification''',
						    '''MegaForTokenClassification''',
						    '''MegaModel''',
						    '''MegaPreTrainedModel''',
						]

if TYPE_CHECKING:
						from .configuration_mega import MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP, MegaConfig, MegaOnnxConfig

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_mega import (
												    MEGA_PRETRAINED_MODEL_ARCHIVE_LIST,
												    MegaForCausalLM,
												    MegaForMaskedLM,
												    MegaForMultipleChoice,
												    MegaForQuestionAnswering,
												    MegaForSequenceClassification,
												    MegaForTokenClassification,
												    MegaModel,
												    MegaPreTrainedModel,
												)

else:
						import sys

						__A		:       Any  				=		_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



