




"""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 				snake_case_			(					lowerCamelCase_	, lowerCamelCase_	, unittest.TestCase			):







   """simple docstring"""




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







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







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_       ,					lowerCamelCase_=0)				->   Optional[int]:
       if str(lowerCamelCase_).startswith('''mps'''):
           UpperCamelCase              =	torch.manual_seed(lowerCamelCase_)
       else:
           UpperCamelCase              =	torch.Generator(device=lowerCamelCase_).manual_seed(lowerCamelCase_)

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

       return inputs







   def   UpperCAmelCase__			(		self)				->   Optional[Any]:
       self._test_save_load_optional_components()







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







   def   UpperCAmelCase__			(		self)				->   List[Any]:
       self._test_save_load_local()







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

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







   """simple docstring"""




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







   def   UpperCAmelCase__			(		self)				->   int:
       # if

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

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

       # pre compute text embeddings and remove T5 to save memory

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

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

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

       UpperCamelCase              =	None
       UpperCamelCase              =	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(lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_)

       pipe_a.remove_all_hooks()
       pipe_a.remove_all_hooks()

       # img2img

       UpperCamelCase              =	IFImgaImgPipeline(**pipe_a.components)
       UpperCamelCase              =	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(lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_)

       pipe_a.remove_all_hooks()
       pipe_a.remove_all_hooks()

       # inpainting

       UpperCamelCase              =	IFInpaintingPipeline(**pipe_a.components)
       UpperCamelCase              =	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(lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_)







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_)				->   Tuple:
       # pipeline 1

       _start_torch_memory_measurement()

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

       UpperCamelCase              =	output.images[0]

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

       UpperCamelCase              =	torch.cuda.max_memory_allocated()
       assert mem_bytes < 1_3 * 1_0**9

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

       # pipeline 2

       _start_torch_memory_measurement()

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

       UpperCamelCase              =	floats_tensor((1, 3, 6_4, 6_4)       ,					rng=random.Random(0)).to(lowerCamelCase_)

       UpperCamelCase              =	pipe_a(
           prompt_embeds=lowerCamelCase_       ,					negative_prompt_embeds=lowerCamelCase_       ,					image=lowerCamelCase_       ,					generator=lowerCamelCase_       ,					num_inference_steps=2       ,					output_type='''np'''       ,					)

       UpperCamelCase              =	output.images[0]

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

       UpperCamelCase              =	torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 1_0**9

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







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_)				->   Tuple:
       # pipeline 1

       _start_torch_memory_measurement()

       UpperCamelCase              =	floats_tensor((1, 3, 6_4, 6_4)       ,					rng=random.Random(0)).to(lowerCamelCase_)

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

       UpperCamelCase              =	pipe_a(
           prompt_embeds=lowerCamelCase_       ,					negative_prompt_embeds=lowerCamelCase_       ,					image=lowerCamelCase_       ,					num_inference_steps=2       ,					generator=lowerCamelCase_       ,					output_type='''np'''       ,					)

       UpperCamelCase              =	output.images[0]

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

       UpperCamelCase              =	torch.cuda.max_memory_allocated()
       assert mem_bytes < 1_0 * 1_0**9

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

       # pipeline 2

       _start_torch_memory_measurement()

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

       UpperCamelCase              =	floats_tensor((1, 3, 2_5_6, 2_5_6)       ,					rng=random.Random(0)).to(lowerCamelCase_)
       UpperCamelCase              =	floats_tensor((1, 3, 6_4, 6_4)       ,					rng=random.Random(0)).to(lowerCamelCase_)

       UpperCamelCase              =	pipe_a(
           prompt_embeds=lowerCamelCase_       ,					negative_prompt_embeds=lowerCamelCase_       ,					image=lowerCamelCase_       ,					original_image=lowerCamelCase_       ,					generator=lowerCamelCase_       ,					num_inference_steps=2       ,					output_type='''np'''       ,					)

       UpperCamelCase              =	output.images[0]

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

       UpperCamelCase              =	torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 1_0**9

       UpperCamelCase              =	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(lowerCamelCase_       ,					lowerCamelCase_)







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_)				->   int:
       # pipeline 1

       _start_torch_memory_measurement()

       UpperCamelCase              =	floats_tensor((1, 3, 6_4, 6_4)       ,					rng=random.Random(0)).to(lowerCamelCase_)
       UpperCamelCase              =	floats_tensor((1, 3, 6_4, 6_4)       ,					rng=random.Random(1)).to(lowerCamelCase_)

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

       UpperCamelCase              =	output.images[0]

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

       UpperCamelCase              =	torch.cuda.max_memory_allocated()
       assert mem_bytes < 1_0 * 1_0**9

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

       # pipeline 2

       _start_torch_memory_measurement()

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

       UpperCamelCase              =	floats_tensor((1, 3, 6_4, 6_4)       ,					rng=random.Random(0)).to(lowerCamelCase_)
       UpperCamelCase              =	floats_tensor((1, 3, 2_5_6, 2_5_6)       ,					rng=random.Random(0)).to(lowerCamelCase_)
       UpperCamelCase              =	floats_tensor((1, 3, 2_5_6, 2_5_6)       ,					rng=random.Random(1)).to(lowerCamelCase_)

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

       UpperCamelCase              =	output.images[0]

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

       UpperCamelCase              =	torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 1_0**9

       UpperCamelCase              =	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(lowerCamelCase_       ,					lowerCamelCase_)





def      __snake_case (				):


    """simple docstring"""







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






from ... import PretrainedConfig


__A							:			Optional[Any]         =			{
    "sijunhe/nezha-cn-base": "https://huggingface.co/sijunhe/nezha-cn-base/resolve/main/config.json",
}







class      lowerCamelCase(						__snake_case	):
     '''simple docstring'''


     __magic_name__   =			NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP
     __magic_name__   =			'nezha'




     def __init__(    self				,     snake_case_=2_1128				,     snake_case_=768				,     snake_case_=12				,     snake_case_=12				,     snake_case_=3072				,     snake_case_="gelu"				,     snake_case_=0.1				,     snake_case_=0.1				,     snake_case_=512				,     snake_case_=64				,     snake_case_=2				,     snake_case_=0.02				,     snake_case_=1E-12				,     snake_case_=0.1				,     snake_case_=0				,     snake_case_=2				,     snake_case_=3				,     snake_case_=True				,     **snake_case_				,     ):
         super().__init__(pad_token_id=snake_case_				,     bos_token_id=snake_case_				,     eos_token_id=snake_case_				,     **snake_case_							)

         _A     =							vocab_size
         _A     =							hidden_size
         _A     =							num_hidden_layers
         _A     =							num_attention_heads
         _A     =							hidden_act
         _A     =							intermediate_size
         _A     =							hidden_dropout_prob
         _A     =							attention_probs_dropout_prob
         _A     =							max_position_embeddings
         _A     =							max_relative_position
         _A     =							type_vocab_size
         _A     =							initializer_range
         _A     =							layer_norm_eps
         _A     =							classifier_dropout
         _A     =							use_cache





