




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 		__A      (			UpperCamelCase__					,						UpperCamelCase__					,						unittest.TestCase			):
						UpperCamelCase     						=					IFPipeline
						UpperCamelCase     						=					TEXT_TO_IMAGE_PARAMS - {"""width""", """height""", """latents"""}
						UpperCamelCase     						=					TEXT_TO_IMAGE_BATCH_PARAMS
						UpperCamelCase     						=					PipelineTesterMixin.required_optional_params - {"""latents"""}

						def     A__						(		self   :str ):
										'''simple docstring'''

										return self._get_dummy_components()

						def     A__						(		self   :Tuple    , __snake_case   :Optional[Any]    , __snake_case   :Tuple=0 ):
										'''simple docstring'''

										if str(__snake_case ).startswith("""mps""" ):
														__magic_name__	:	str								=torch.manual_seed(__snake_case )
										else:
														__magic_name__	:	Tuple								=torch.Generator(device=__snake_case ).manual_seed(__snake_case )

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

										return inputs

						def     A__						(		self   :Tuple ):
										'''simple docstring'''

										self._test_save_load_optional_components()

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

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

						def     A__						(		self   :Optional[Any] ):
										'''simple docstring'''

										self._test_attention_slicing_forward_pass(expected_max_diff=1E-2 )

						def     A__						(		self   :Tuple ):
										'''simple docstring'''

										self._test_save_load_local()

						def     A__						(		self   :Any ):
										'''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     A__						(		self   :str ):
										'''simple docstring'''

										self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3 )




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

						def     A__						(		self   :List[Any] ):
										'''simple docstring'''

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

						def     A__						(		self   :Optional[Any] ):
										'''simple docstring'''

										__magic_name__	:	Dict								=IFPipeline.from_pretrained("""DeepFloyd/IF-I-XL-v1.0"""    , variant="""fp16"""    , torch_dtype=torch.floataa )

										__magic_name__	:	Optional[Any]								=IFSuperResolutionPipeline.from_pretrained(
										    """DeepFloyd/IF-II-L-v1.0"""    , variant="""fp16"""    , torch_dtype=torch.floataa    , text_encoder=__snake_case    , tokenizer=__snake_case )

										# pre compute text embeddings and remove T5 to save memory

										pipe_a.text_encoder.to("""cuda""" )

										__magic_name__				, __magic_name__	:	Any								=pipe_a.encode_prompt("""anime turtle"""    , device="""cuda""" )

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

										__magic_name__	:	List[str]								=None
										__magic_name__	:	Dict								=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(__snake_case    , __snake_case    , __snake_case    , __snake_case )

										pipe_a.remove_all_hooks()
										pipe_a.remove_all_hooks()

										# img2img

										__magic_name__	:	Union[str, Any]								=IFImgaImgPipeline(**pipe_a.components )
										__magic_name__	:	List[Any]								=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(__snake_case    , __snake_case    , __snake_case    , __snake_case )

										pipe_a.remove_all_hooks()
										pipe_a.remove_all_hooks()

										# inpainting

										__magic_name__	:	Tuple								=IFInpaintingPipeline(**pipe_a.components )
										__magic_name__	:	List[str]								=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(__snake_case    , __snake_case    , __snake_case    , __snake_case )

						def     A__						(		self   :Optional[int]    , __snake_case   :List[str]    , __snake_case   :List[Any]    , __snake_case   :Optional[int]    , __snake_case   :Union[str, Any] ):
										'''simple docstring'''

										_start_torch_memory_measurement()

										__magic_name__	:	str								=torch.Generator(device="""cpu""" ).manual_seed(0 )
										__magic_name__	:	Any								=pipe_a(
										    prompt_embeds=__snake_case    , negative_prompt_embeds=__snake_case    , num_inference_steps=2    , generator=__snake_case    , output_type="""np"""    , )

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

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

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

										__magic_name__	:	int								=load_numpy(
										    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy""" )
										assert_mean_pixel_difference(__snake_case    , __snake_case )

										# pipeline 2

										_start_torch_memory_measurement()

										__magic_name__	:	Dict								=torch.Generator(device="""cpu""" ).manual_seed(0 )

										__magic_name__	:	Optional[Any]								=floats_tensor((1, 3, 64, 64)    , rng=random.Random(0 ) ).to(__snake_case )

										__magic_name__	:	Union[str, Any]								=pipe_a(
										    prompt_embeds=__snake_case    , negative_prompt_embeds=__snake_case    , image=__snake_case    , generator=__snake_case    , num_inference_steps=2    , output_type="""np"""    , )

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

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

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

										__magic_name__	:	Dict								=load_numpy(
										    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy""" )
										assert_mean_pixel_difference(__snake_case    , __snake_case )

						def     A__						(		self   :List[Any]    , __snake_case   :List[Any]    , __snake_case   :Optional[Any]    , __snake_case   :Any    , __snake_case   :Tuple ):
										'''simple docstring'''

										_start_torch_memory_measurement()

										__magic_name__	:	Optional[Any]								=floats_tensor((1, 3, 64, 64)    , rng=random.Random(0 ) ).to(__snake_case )

										__magic_name__	:	Dict								=torch.Generator(device="""cpu""" ).manual_seed(0 )

										__magic_name__	:	List[str]								=pipe_a(
										    prompt_embeds=__snake_case    , negative_prompt_embeds=__snake_case    , image=__snake_case    , num_inference_steps=2    , generator=__snake_case    , output_type="""np"""    , )

										__magic_name__	:	int								=output.images[0]

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

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

										__magic_name__	:	Tuple								=load_numpy(
										    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy""" )
										assert_mean_pixel_difference(__snake_case    , __snake_case )

										# pipeline 2

										_start_torch_memory_measurement()

										__magic_name__	:	Dict								=torch.Generator(device="""cpu""" ).manual_seed(0 )

										__magic_name__	:	Optional[int]								=floats_tensor((1, 3, 2_56, 2_56)    , rng=random.Random(0 ) ).to(__snake_case )
										__magic_name__	:	Union[str, Any]								=floats_tensor((1, 3, 64, 64)    , rng=random.Random(0 ) ).to(__snake_case )

										__magic_name__	:	Union[str, Any]								=pipe_a(
										    prompt_embeds=__snake_case    , negative_prompt_embeds=__snake_case    , image=__snake_case    , original_image=__snake_case    , generator=__snake_case    , num_inference_steps=2    , output_type="""np"""    , )

										__magic_name__	:	str								=output.images[0]

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

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

										__magic_name__	:	Dict								=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(__snake_case    , __snake_case )


						def     A__						(		self   :Any    , __snake_case   :Any    , __snake_case   :List[Any]    , __snake_case   :Optional[Any]    , __snake_case   :int ):
										'''simple docstring'''

										_start_torch_memory_measurement()

										__magic_name__	:	List[Any]								=floats_tensor((1, 3, 64, 64)    , rng=random.Random(0 ) ).to(__snake_case )
										__magic_name__	:	int								=floats_tensor((1, 3, 64, 64)    , rng=random.Random(1 ) ).to(__snake_case )

										__magic_name__	:	Dict								=torch.Generator(device="""cpu""" ).manual_seed(0 )
										__magic_name__	:	int								=pipe_a(
										    prompt_embeds=__snake_case    , negative_prompt_embeds=__snake_case    , image=__snake_case    , mask_image=__snake_case    , num_inference_steps=2    , generator=__snake_case    , output_type="""np"""    , )

										__magic_name__	:	Any								=output.images[0]

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

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

										__magic_name__	:	str								=load_numpy(
										    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy""" )
										assert_mean_pixel_difference(__snake_case    , __snake_case )

										# pipeline 2

										_start_torch_memory_measurement()

										__magic_name__	:	int								=torch.Generator(device="""cpu""" ).manual_seed(0 )

										__magic_name__	:	str								=floats_tensor((1, 3, 64, 64)    , rng=random.Random(0 ) ).to(__snake_case )
										__magic_name__	:	Optional[Any]								=floats_tensor((1, 3, 2_56, 2_56)    , rng=random.Random(0 ) ).to(__snake_case )
										__magic_name__	:	str								=floats_tensor((1, 3, 2_56, 2_56)    , rng=random.Random(1 ) ).to(__snake_case )

										__magic_name__	:	Optional[int]								=pipe_a(
										    prompt_embeds=__snake_case    , negative_prompt_embeds=__snake_case    , image=__snake_case    , mask_image=__snake_case    , original_image=__snake_case    , generator=__snake_case    , num_inference_steps=2    , output_type="""np"""    , )

										__magic_name__	:	int								=output.images[0]

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

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

										__magic_name__	:	Union[str, 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(__snake_case    , __snake_case )




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










import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase_		:				Union[str, Any]									=     logging.get_logger(__name__)

UpperCAmelCase_		:				Dict									=     {
    "google/pix2struct-textcaps-base": (
        "https://huggingface.co/google/pix2struct-textcaps-base/resolve/main/config.json"
    ),
}




class 		__A      (			UpperCamelCase__			):
						UpperCamelCase     						=					"""pix2struct_text_model"""
						UpperCamelCase     						=					["""past_key_values"""]
						UpperCamelCase     						=					{
						    """hidden_size""": """hidden_size""",
						    """num_attention_heads""": """num_heads""",
						    """num_hidden_layers""": """num_layers""",
						}

						def __init__(		self   :List[Any]    , __snake_case   :Dict=5_02_44    , __snake_case   :Dict=7_68    , __snake_case   :Union[str, Any]=64    , __snake_case   :Union[str, Any]=20_48    , __snake_case   :List[Any]=12    , __snake_case   :Any=12    , __snake_case   :str=32    , __snake_case   :Dict=1_28    , __snake_case   :Optional[int]=0.1    , __snake_case   :Any=1E-6    , __snake_case   :Optional[int]=1.0    , __snake_case   :List[Any]="gelu_new"    , __snake_case   :Any=0    , __snake_case   :Tuple=False    , __snake_case   :Any=0    , __snake_case   :int=1    , __snake_case   :Dict=False    , __snake_case   :Optional[Any]=True    , **__snake_case   :Dict    , ):
										'''simple docstring'''

										__magic_name__	:	List[str]								=vocab_size
										__magic_name__	:	List[str]								=hidden_size
										__magic_name__	:	Optional[int]								=d_kv
										__magic_name__	:	Dict								=d_ff
										__magic_name__	:	Tuple								=num_layers
										__magic_name__	:	str								=num_heads
										__magic_name__	:	Optional[Any]								=relative_attention_num_buckets
										__magic_name__	:	List[Any]								=relative_attention_max_distance
										__magic_name__	:	str								=dropout_rate
										__magic_name__	:	Optional[Any]								=layer_norm_epsilon
										__magic_name__	:	Union[str, Any]								=initializer_factor
										__magic_name__	:	Tuple								=use_cache

										__magic_name__	:	Any								=eos_token_id
										__magic_name__	:	int								=decoder_start_token_id

										# for backwards compatibility
										__magic_name__	:	Union[str, Any]								=dense_act_fn

										super().__init__(
										    pad_token_id=__snake_case    , eos_token_id=__snake_case    , decoder_start_token_id=__snake_case    , tie_word_embeddings=__snake_case    , is_decoder=__snake_case    , **__snake_case    , )


						@classmethod
						def     A__						(		cls   :Union[str, Any]    , __snake_case   :Union[str, os.PathLike]    , **__snake_case   :Any ):
										'''simple docstring'''

										cls._set_token_in_kwargs(__snake_case )

										__magic_name__				, __magic_name__	:	Optional[Any]								=cls.get_config_dict(__snake_case    , **__snake_case )

										# get the text config dict if we are loading from Pix2StructConfig
										if config_dict.get("""model_type""" ) == "pix2struct":
														__magic_name__	:	Tuple								=config_dict["""text_config"""]

										if "model_type" in config_dict and hasattr(cls    , """model_type""" ) and config_dict["model_type"] != cls.model_type:
														logger.warning(
														    f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
														    f"{cls.model_type}. This is not supported for all configurations of models and can yield errors." )

										return cls.from_dict(__snake_case    , **__snake_case )




class 		__A      (			UpperCamelCase__			):
						UpperCamelCase     						=					"""pix2struct_vision_model"""

						def __init__(		self   :Union[str, Any]    , __snake_case   :Tuple=7_68    , __snake_case   :Any=7_68    , __snake_case   :Tuple=20_48    , __snake_case   :Dict=64    , __snake_case   :List[str]=12    , __snake_case   :str=12    , __snake_case   :str="gelu_new"    , __snake_case   :Optional[int]=1E-6    , __snake_case   :Optional[int]=0.0    , __snake_case   :Any=0.0    , __snake_case   :Any=1E-10    , __snake_case   :List[str]=1.0    , __snake_case   :Tuple=40_96    , __snake_case   :Optional[int]=32    , __snake_case   :Union[str, Any]=1_28    , **__snake_case   :Optional[int]    , ):
										'''simple docstring'''

										super().__init__(**__snake_case )

										__magic_name__	:	Union[str, Any]								=hidden_size
										__magic_name__	:	Dict								=patch_embed_hidden_size
										__magic_name__	:	Tuple								=d_ff
										__magic_name__	:	str								=dropout_rate
										__magic_name__	:	str								=num_hidden_layers
										__magic_name__	:	Dict								=num_attention_heads
										__magic_name__	:	List[str]								=initializer_range
										__magic_name__	:	Optional[int]								=initializer_factor
										__magic_name__	:	str								=attention_dropout
										__magic_name__	:	Union[str, Any]								=layer_norm_eps
										__magic_name__	:	List[str]								=dense_act_fn
										__magic_name__	:	List[str]								=seq_len
										__magic_name__	:	str								=relative_attention_num_buckets
										__magic_name__	:	Optional[int]								=relative_attention_max_distance
										__magic_name__	:	List[str]								=d_kv


						@classmethod
						def     A__						(		cls   :Optional[Any]    , __snake_case   :Union[str, os.PathLike]    , **__snake_case   :Any ):
										'''simple docstring'''

										cls._set_token_in_kwargs(__snake_case )

										__magic_name__				, __magic_name__	:	Optional[int]								=cls.get_config_dict(__snake_case    , **__snake_case )

										# get the vision config dict if we are loading from Pix2StructConfig
										if config_dict.get("""model_type""" ) == "pix2struct":
														__magic_name__	:	int								=config_dict["""vision_config"""]

										if "model_type" in config_dict and hasattr(cls    , """model_type""" ) and config_dict["model_type"] != cls.model_type:
														logger.warning(
														    f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
														    f"{cls.model_type}. This is not supported for all configurations of models and can yield errors." )

										return cls.from_dict(__snake_case    , **__snake_case )




class 		__A      (			UpperCamelCase__			):
						UpperCamelCase     						=					"""pix2struct"""
						UpperCamelCase     						=					True

						def __init__(		self   :Union[str, Any]    , __snake_case   :Optional[int]=None    , __snake_case   :List[Any]=None    , __snake_case   :Optional[Any]=1.0    , __snake_case   :Tuple=0.02    , __snake_case   :str=False    , __snake_case   :List[str]=False    , __snake_case   :str=True    , **__snake_case   :List[str]    , ):
										'''simple docstring'''

										super().__init__(tie_word_embeddings=__snake_case    , is_encoder_decoder=__snake_case    , **__snake_case )

										if text_config is None:
														__magic_name__	:	int								={}
														logger.info("""text_config is None. Initializing the Pix2StructTextConfig with default values.""" )

										if vision_config is None:
														__magic_name__	:	str								={}
														logger.info("""vision_config is None. Initializing the Pix2StructVisionConfig with default values.""" )

										__magic_name__	:	Union[str, Any]								=PixaStructTextConfig(**__snake_case )
										__magic_name__	:	str								=PixaStructVisionConfig(**__snake_case )

										__magic_name__	:	int								=self.text_config.decoder_start_token_id
										__magic_name__	:	Optional[Any]								=self.text_config.pad_token_id
										__magic_name__	:	str								=self.text_config.eos_token_id

										__magic_name__	:	int								=initializer_factor
										__magic_name__	:	List[Any]								=initializer_range

										__magic_name__	:	Dict								=self.initializer_range
										__magic_name__	:	Union[str, Any]								=self.initializer_range

										__magic_name__	:	Tuple								=is_vqa

						@classmethod
						def     A__						(		cls   :Optional[int]    , __snake_case   :PixaStructTextConfig    , __snake_case   :PixaStructVisionConfig    , **__snake_case   :List[str] ):
										'''simple docstring'''

										return cls(text_config=text_config.to_dict()    , vision_config=vision_config.to_dict()    , **__snake_case )


						def     A__						(		self   :Dict ):
										'''simple docstring'''

										__magic_name__	:	Union[str, Any]								=copy.deepcopy(self.__dict__ )
										__magic_name__	:	Any								=self.text_config.to_dict()
										__magic_name__	:	List[str]								=self.vision_config.to_dict()
										__magic_name__	:	Optional[int]								=self.__class__.model_type
										return output




