






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







     '''simple docstring'''
     SCREAMING_SNAKE_CASE   						=				IFPipeline
     SCREAMING_SNAKE_CASE   						=				TEXT_TO_IMAGE_PARAMS - {'width', 'height', 'latents'}
     SCREAMING_SNAKE_CASE   						=				TEXT_TO_IMAGE_BATCH_PARAMS
     SCREAMING_SNAKE_CASE   						=				PipelineTesterMixin.required_optional_params - {'latents'}




     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



         return self._get_dummy_components()




     def 		lowerCAmelCase_     (				self       :			Union[str, Any] ,	UpperCamelCase_       :			Tuple ,	UpperCamelCase_       :			Union[str, Any]=0				):


         """simple docstring"""



         if str(__A				).startswith("""mps"""				):
             __A         =	torch.manual_seed(__A				)
         else:
             __A         =	torch.Generator(device=__A				).manual_seed(__A				)

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

         return inputs




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


         """simple docstring"""



         self._test_save_load_optional_components()




     @unittest.skipIf(torch_device != """cuda""" ,	reason="""float16 requires CUDA"""				)
     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



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




     def 		lowerCAmelCase_     (				self       :			str				):


         """simple docstring"""



         self._test_attention_slicing_forward_pass(expected_max_diff=1e-2				)




     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         self._test_save_load_local()




     def 		lowerCAmelCase_     (				self       :			Dict				):


         """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 		lowerCAmelCase_     (				self       :			Tuple				):


         """simple docstring"""



         self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3				)


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







     '''simple docstring'''




     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



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




     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



         __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=__A ,	tokenizer=__A				)

         # 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(__A ,	__A ,	__A ,	__A				)

         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(__A ,	__A ,	__A ,	__A				)

         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(__A ,	__A ,	__A ,	__A				)




     def 		lowerCAmelCase_     (				self       :			Optional[int] ,	UpperCamelCase_       :			Dict ,	UpperCamelCase_       :			Optional[int] ,	UpperCamelCase_       :			Any ,	UpperCamelCase_       :			Union[str, Any]				):


         """simple docstring"""



         _start_torch_memory_measurement()

         __A         =	torch.Generator(device="""cpu"""				).manual_seed(0				)
         __A         =	pipe_a(
             prompt_embeds=__A ,	negative_prompt_embeds=__A ,	num_inference_steps=2 ,	generator=__A ,	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(__A ,	__A				)

         # 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(__A				)

         __A         =	pipe_a(
             prompt_embeds=__A ,	negative_prompt_embeds=__A ,	image=__A ,	generator=__A ,	num_inference_steps=2 ,	output_type="""np""" ,	)

         __A         =	output.images[0]

         assert image.shape == (256, 256, 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(__A ,	__A				)




     def 		lowerCAmelCase_     (				self       :			str ,	UpperCamelCase_       :			Optional[int] ,	UpperCamelCase_       :			List[str] ,	UpperCamelCase_       :			Union[str, Any] ,	UpperCamelCase_       :			Union[str, Any]				):


         """simple docstring"""



         _start_torch_memory_measurement()

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

         __A         =	torch.Generator(device="""cpu"""				).manual_seed(0				)

         __A         =	pipe_a(
             prompt_embeds=__A ,	negative_prompt_embeds=__A ,	image=__A ,	num_inference_steps=2 ,	generator=__A ,	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(__A ,	__A				)

         # pipeline 2

         _start_torch_memory_measurement()

         __A         =	torch.Generator(device="""cpu"""				).manual_seed(0				)

         __A         =	floats_tensor((1, 3, 256, 256) ,	rng=random.Random(0				)				).to(__A				)
         __A         =	floats_tensor((1, 3, 64, 64) ,	rng=random.Random(0				)				).to(__A				)

         __A         =	pipe_a(
             prompt_embeds=__A ,	negative_prompt_embeds=__A ,	image=__A ,	original_image=__A ,	generator=__A ,	num_inference_steps=2 ,	output_type="""np""" ,	)

         __A         =	output.images[0]

         assert image.shape == (256, 256, 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(__A ,	__A				)





     def 		lowerCAmelCase_     (				self       :			Union[str, Any] ,	UpperCamelCase_       :			Tuple ,	UpperCamelCase_       :			int ,	UpperCamelCase_       :			Any ,	UpperCamelCase_       :			Optional[int]				):


         """simple docstring"""



         _start_torch_memory_measurement()

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

         __A         =	torch.Generator(device="""cpu"""				).manual_seed(0				)
         __A         =	pipe_a(
             prompt_embeds=__A ,	negative_prompt_embeds=__A ,	image=__A ,	mask_image=__A ,	num_inference_steps=2 ,	generator=__A ,	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(__A ,	__A				)

         # 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(__A				)
         __A         =	floats_tensor((1, 3, 256, 256) ,	rng=random.Random(0				)				).to(__A				)
         __A         =	floats_tensor((1, 3, 256, 256) ,	rng=random.Random(1				)				).to(__A				)

         __A         =	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""" ,	)

         __A         =	output.images[0]

         assert image.shape == (256, 256, 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(__A ,	__A				)



def 			_SCREAMING_SNAKE_CASE						( )						->      Optional[Any]:





    """simple docstring"""

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












from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, Encoder, VectorQuantizer


@dataclass
class       __lowercase      (	lowercase_  ):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE   						=				42


class       __lowercase      (	lowercase_			,      lowercase_  ):







     '''simple docstring'''




     @register_to_config
     def __init__(				self       :			List[Any] ,	UpperCamelCase_       :			int = 3 ,	UpperCamelCase_       :			int = 3 ,	UpperCamelCase_       :			Tuple[str] = ("DownEncoderBlock2D",) ,	UpperCamelCase_       :			Tuple[str] = ("UpDecoderBlock2D",) ,	UpperCamelCase_       :			Tuple[int] = (64,) ,	UpperCamelCase_       :			int = 1 ,	UpperCamelCase_       :			str = "silu" ,	UpperCamelCase_       :			int = 3 ,	UpperCamelCase_       :			int = 32 ,	UpperCamelCase_       :			int = 256 ,	UpperCamelCase_       :			int = 32 ,	UpperCamelCase_       :			Optional[int] = None ,	UpperCamelCase_       :			float = 0.18215 ,	UpperCamelCase_       :			str = "group" ,	):


         """simple docstring"""



         super().__init__()

         # pass init params to Encoder
         __A         =	Encoder(
             in_channels=UpperCamelCase_ ,	out_channels=UpperCamelCase_ ,	down_block_types=UpperCamelCase_ ,	block_out_channels=UpperCamelCase_ ,	layers_per_block=UpperCamelCase_ ,	act_fn=UpperCamelCase_ ,	norm_num_groups=UpperCamelCase_ ,	double_z=UpperCamelCase_ ,	)

         __A         =	vq_embed_dim if vq_embed_dim is not None else latent_channels

         __A         =	nn.Convad(UpperCamelCase_ ,	UpperCamelCase_ ,	1				)
         __A         =	VectorQuantizer(UpperCamelCase_ ,	UpperCamelCase_ ,	beta=0.25 ,	remap=UpperCamelCase_ ,	sane_index_shape=UpperCamelCase_				)
         __A         =	nn.Convad(UpperCamelCase_ ,	UpperCamelCase_ ,	1				)

         # pass init params to Decoder
         __A         =	Decoder(
             in_channels=UpperCamelCase_ ,	out_channels=UpperCamelCase_ ,	up_block_types=UpperCamelCase_ ,	block_out_channels=UpperCamelCase_ ,	layers_per_block=UpperCamelCase_ ,	act_fn=UpperCamelCase_ ,	norm_num_groups=UpperCamelCase_ ,	norm_type=UpperCamelCase_ ,	)




     @apply_forward_hook
     def 		lowerCAmelCase_     (				self       :			List[str] ,	UpperCamelCase_       :			torch.FloatTensor ,	UpperCamelCase_       :			bool = True				):


         """simple docstring"""



         __A         =	self.encoder(UpperCamelCase_				)
         __A         =	self.quant_conv(UpperCamelCase_				)

         if not return_dict:
             return (h,)

         return VQEncoderOutput(latents=UpperCamelCase_				)




     @apply_forward_hook
     def 		lowerCAmelCase_     (				self       :			List[str] ,	UpperCamelCase_       :			torch.FloatTensor ,	UpperCamelCase_       :			bool = False ,	UpperCamelCase_       :			bool = True				):


         """simple docstring"""



         if not force_not_quantize:
             __A    ,	__A    ,	__A         =	self.quantize(UpperCamelCase_				)
         else:
             __A         =	h
         __A         =	self.post_quant_conv(UpperCamelCase_				)
         __A         =	self.decoder(UpperCamelCase_ ,	quant if self.config.norm_type == """spatial""" else None				)

         if not return_dict:
             return (dec,)

         return DecoderOutput(sample=UpperCamelCase_				)





     def 		lowerCAmelCase_     (				self       :			Any ,	UpperCamelCase_       :			torch.FloatTensor ,	UpperCamelCase_       :			bool = True				):


         """simple docstring"""



         __A         =	sample
         __A         =	self.encode(UpperCamelCase_				).latents
         __A         =	self.decode(UpperCamelCase_				).sample

         if not return_dict:
             return (dec,)

         return DecoderOutput(sample=UpperCamelCase_				)




