# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from packaging import version

from .. import __version__
from .constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD
from .doc import (
    add_code_sample_docstrings,
    add_end_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    copy_func,
    replace_return_docstrings,
)
from .generic import (
    ContextManagers,
    ExplicitEnum,
    ModelOutput,
    PaddingStrategy,
    TensorType,
    add_model_info_to_auto_map,
    cached_property,
    can_return_loss,
    expand_dims,
    find_labels,
    flatten_dict,
    infer_framework,
    is_jax_tensor,
    is_numpy_array,
    is_tensor,
    is_tf_symbolic_tensor,
    is_tf_tensor,
    is_torch_device,
    is_torch_dtype,
    is_torch_tensor,
    reshape,
    squeeze,
    strtobool,
    tensor_size,
    to_numpy,
    to_py_obj,
    transpose,
    working_or_temp_dir,
)
from .hub import (
    CLOUDFRONT_DISTRIB_PREFIX,
    DISABLE_TELEMETRY,
    HF_MODULES_CACHE,
    HUGGINGFACE_CO_PREFIX,
    HUGGINGFACE_CO_RESOLVE_ENDPOINT,
    PYTORCH_PRETRAINED_BERT_CACHE,
    PYTORCH_TRANSFORMERS_CACHE,
    S3_BUCKET_PREFIX,
    TRANSFORMERS_CACHE,
    TRANSFORMERS_DYNAMIC_MODULE_NAME,
    EntryNotFoundError,
    PushToHubMixin,
    RepositoryNotFoundError,
    RevisionNotFoundError,
    cached_file,
    default_cache_path,
    define_sagemaker_information,
    download_url,
    extract_commit_hash,
    get_cached_models,
    get_file_from_repo,
    get_full_repo_name,
    has_file,
    http_user_agent,
    is_offline_mode,
    is_remote_url,
    move_cache,
    send_example_telemetry,
    try_to_load_from_cache,
)
from .import_utils import (
    ENV_VARS_TRUE_AND_AUTO_VALUES,
    ENV_VARS_TRUE_VALUES,
    TORCH_FX_REQUIRED_VERSION,
    USE_JAX,
    USE_TF,
    USE_TORCH,
    DummyObject,
    OptionalDependencyNotAvailable,
    _LazyModule,
    ccl_version,
    direct_transformers_import,
    get_torch_version,
    is_accelerate_available,
    is_apex_available,
    is_bitsandbytes_available,
    is_bsa_available,
    is_coloredlogs_available,
    is_cython_available,
    is_datasets_available,
    is_decord_available,
    is_detectrona_available,
    is_faiss_available,
    is_flax_available,
    is_ftfy_available,
    is_in_notebook,
    is_ipex_available,
    is_jieba_available,
    is_jumanpp_available,
    is_kenlm_available,
    is_keras_nlp_available,
    is_librosa_available,
    is_natten_available,
    is_ninja_available,
    is_onnx_available,
    is_openai_available,
    is_optimum_available,
    is_pandas_available,
    is_peft_available,
    is_phonemizer_available,
    is_protobuf_available,
    is_psutil_available,
    is_pyanvml_available,
    is_pyctcdecode_available,
    is_pytesseract_available,
    is_pytest_available,
    is_pytorch_quantization_available,
    is_rjieba_available,
    is_sacremoses_available,
    is_safetensors_available,
    is_sagemaker_dp_enabled,
    is_sagemaker_mp_enabled,
    is_scipy_available,
    is_sentencepiece_available,
    is_seqio_available,
    is_sklearn_available,
    is_soundfile_availble,
    is_spacy_available,
    is_speech_available,
    is_sudachi_available,
    is_tensorflow_probability_available,
    is_tensorflow_text_available,
    is_tfaonnx_available,
    is_tf_available,
    is_timm_available,
    is_tokenizers_available,
    is_torch_available,
    is_torch_bfaa_available,
    is_torch_bfaa_cpu_available,
    is_torch_bfaa_gpu_available,
    is_torch_compile_available,
    is_torch_cuda_available,
    is_torch_fx_available,
    is_torch_fx_proxy,
    is_torch_mps_available,
    is_torch_neuroncore_available,
    is_torch_tensorrt_fx_available,
    is_torch_tfaa_available,
    is_torch_tpu_available,
    is_torchaudio_available,
    is_torchdistx_available,
    is_torchdynamo_available,
    is_torchvision_available,
    is_training_run_on_sagemaker,
    is_vision_available,
    requires_backends,
    torch_only_method,
)


__UpperCamelCase	: Optional[int]    		=  'pytorch_model.bin'
__UpperCamelCase	: Union[str, Any]    		=  'pytorch_model.bin.index.json'
__UpperCamelCase	: Dict    		=  'adapter_config.json'
__UpperCamelCase	: Optional[int]    		=  'adapter_model.bin'
__UpperCamelCase	: List[str]    		=  'adapter_model.safetensors'
__UpperCamelCase	: Optional[Any]    		=  'tf_model.h5'
__UpperCamelCase	: Optional[int]    		=  'tf_model.h5.index.json'
__UpperCamelCase	: Any    		=  'model.ckpt'
__UpperCamelCase	: Tuple    		=  'flax_model.msgpack'
__UpperCamelCase	: Optional[Any]    		=  'flax_model.msgpack.index.json'
__UpperCamelCase	: Optional[int]    		=  'model.safetensors'
__UpperCamelCase	: int    		=  'model.safetensors.index.json'
__UpperCamelCase	: List[Any]    		=  'config.json'
__UpperCamelCase	: List[str]    		=  'preprocessor_config.json'
__UpperCamelCase	: int    		=  FEATURE_EXTRACTOR_NAME
__UpperCamelCase	: int    		=  'generation_config.json'
__UpperCamelCase	: Optional[Any]    		=  'modelcard.json'

__UpperCamelCase	: List[Any]    		=  '▁'
__UpperCamelCase	: Optional[Any]    		=  SENTENCEPIECE_UNDERLINE  # Kept for backward compatibility

__UpperCamelCase	: str    		=  [
    [[0, 1, 0, 1], [1, 0, 0, 1]]
] * 2  # Needs to have 0s and 1s only since XLM uses it for langs too.
__UpperCamelCase	: Union[str, Any]    		=  [[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]]
__UpperCamelCase	: Dict    		=  [[1, 1, 1, 1, 1], [1, 1, 1, 0, 0], [0, 0, 0, 1, 1]]
def    A			(					_lowercase		):
						if version.parse(_lowercase		) < version.parse(_lowercase		):
												if "dev" in min_version:
																		SCREAMING_SNAKE_CASE							:					Dict           =							(
																		    '''This example requires a source install from HuggingFace Transformers (see '''
																		    '''`https://huggingface.co/docs/transformers/installation#install-from-source`),'''
																		)
												else:
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							f"""This example requires a minimum version of {min_version},"""
												error_message += f""" but the version found is {__version__}.\n"""
												raise ImportError(
												    error_message
												    + '''Check out https://github.com/huggingface/transformers/tree/main/examples#important-note for the examples corresponding to other '''
												    '''versions of HuggingFace Transformers.'''		)






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__		(     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			:       Tuple		):





												'''simple docstring'''





												return self._get_dummy_components()


						def  __A ( self			:       int					,	UpperCamelCase__			:       Dict					,	UpperCamelCase__			:       int=0		):





												'''simple docstring'''





												if str(UpperCamelCase__		).startswith('''mps'''		):
																		SCREAMING_SNAKE_CASE							:					Tuple           =							torch.manual_seed(UpperCamelCase__		)
												else:
																		SCREAMING_SNAKE_CASE							:					Dict           =							torch.Generator(device=UpperCamelCase__		).manual_seed(UpperCamelCase__		)

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

												return inputs


						def  __A ( self			:       List[str]		):





												'''simple docstring'''





												self._test_save_load_optional_components()


						@unittest.skipIf(torch_device != '''cuda'''					,	reason='''float16 requires CUDA'''		)
						def  __A ( self			:       Any		):





												'''simple docstring'''





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


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





												'''simple docstring'''





												self._test_attention_slicing_forward_pass(expected_max_diff=1E-2		)


						def  __A ( self			:       List[Any]		):





												'''simple docstring'''





												self._test_save_load_local()


						def  __A ( self			:       List[str]		):





												'''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			:       Tuple		):





												'''simple docstring'''





												self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3		)






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


						def  __A ( self			:       Optional[Any]		):





												'''simple docstring'''





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


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





												'''simple docstring'''





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

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

												# pre compute text embeddings and remove T5 to save memory

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

												SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE							:					Dict           =							pipe_a.encode_prompt('''anime turtle'''					,	device='''cuda'''		)

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

												SCREAMING_SNAKE_CASE							:					Tuple           =							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(UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__		)

												pipe_a.remove_all_hooks()
												pipe_a.remove_all_hooks()

												# img2img

												SCREAMING_SNAKE_CASE							:					Optional[int]           =							IFImgaImgPipeline(**pipe_a.components		)
												SCREAMING_SNAKE_CASE							:					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(UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__		)

												pipe_a.remove_all_hooks()
												pipe_a.remove_all_hooks()

												# inpainting

												SCREAMING_SNAKE_CASE							:					Tuple           =							IFInpaintingPipeline(**pipe_a.components		)
												SCREAMING_SNAKE_CASE							:					Optional[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(UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__		)


						def  __A ( self			:       Tuple					,	UpperCamelCase__			:       int					,	UpperCamelCase__			:       Union[str, Any]					,	UpperCamelCase__			:       Union[str, Any]					,	UpperCamelCase__			:       Dict		):





												'''simple docstring'''





												_start_torch_memory_measurement()

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

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

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

												SCREAMING_SNAKE_CASE							:					Tuple           =							torch.cuda.max_memory_allocated()
												assert mem_bytes < 13 * 10**9

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							load_numpy(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy'''		)
												assert_mean_pixel_difference(UpperCamelCase__					,	UpperCamelCase__		)

												# pipeline 2

												_start_torch_memory_measurement()

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

												SCREAMING_SNAKE_CASE							:					int           =							floats_tensor((1, 3, 64, 64)					,	rng=random.Random(0		)		).to(UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Any           =							pipe_a(
												    prompt_embeds=UpperCamelCase__					,	negative_prompt_embeds=UpperCamelCase__					,	image=UpperCamelCase__					,	generator=UpperCamelCase__					,	num_inference_steps=2					,	output_type='''np'''					,	)

												SCREAMING_SNAKE_CASE							:					Tuple           =							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							:					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(UpperCamelCase__					,	UpperCamelCase__		)


						def  __A ( self			:       Optional[int]					,	UpperCamelCase__			:       str					,	UpperCamelCase__			:       int					,	UpperCamelCase__			:       Optional[Any]					,	UpperCamelCase__			:       str		):





												'''simple docstring'''





												_start_torch_memory_measurement()

												SCREAMING_SNAKE_CASE							:					Tuple           =							floats_tensor((1, 3, 64, 64)					,	rng=random.Random(0		)		).to(UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Any           =							torch.Generator(device='''cpu'''		).manual_seed(0		)

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

												SCREAMING_SNAKE_CASE							:					int           =							output.images[0]

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

												SCREAMING_SNAKE_CASE							:					Optional[int]           =							torch.cuda.max_memory_allocated()
												assert mem_bytes < 10 * 10**9

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

												# pipeline 2

												_start_torch_memory_measurement()

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

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							floats_tensor((1, 3, 256, 256)					,	rng=random.Random(0		)		).to(UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					List[Any]           =							floats_tensor((1, 3, 64, 64)					,	rng=random.Random(0		)		).to(UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					str           =							pipe_a(
												    prompt_embeds=UpperCamelCase__					,	negative_prompt_embeds=UpperCamelCase__					,	image=UpperCamelCase__					,	original_image=UpperCamelCase__					,	generator=UpperCamelCase__					,	num_inference_steps=2					,	output_type='''np'''					,	)

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

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

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

												SCREAMING_SNAKE_CASE							:					Optional[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(UpperCamelCase__					,	UpperCamelCase__		)



						def  __A ( self			:       List[str]					,	UpperCamelCase__			:       Optional[int]					,	UpperCamelCase__			:       Optional[int]					,	UpperCamelCase__			:       Union[str, Any]					,	UpperCamelCase__			:       Optional[Any]		):





												'''simple docstring'''





												_start_torch_memory_measurement()

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

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

												SCREAMING_SNAKE_CASE							:					Any           =							output.images[0]

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

												SCREAMING_SNAKE_CASE							:					Any           =							torch.cuda.max_memory_allocated()
												assert mem_bytes < 10 * 10**9

												SCREAMING_SNAKE_CASE							:					List[Any]           =							load_numpy(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy'''		)
												assert_mean_pixel_difference(UpperCamelCase__					,	UpperCamelCase__		)

												# pipeline 2

												_start_torch_memory_measurement()

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

												SCREAMING_SNAKE_CASE							:					List[str]           =							floats_tensor((1, 3, 64, 64)					,	rng=random.Random(0		)		).to(UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					int           =							floats_tensor((1, 3, 256, 256)					,	rng=random.Random(0		)		).to(UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							floats_tensor((1, 3, 256, 256)					,	rng=random.Random(1		)		).to(UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Dict           =							pipe_a(
												    prompt_embeds=UpperCamelCase__					,	negative_prompt_embeds=UpperCamelCase__					,	image=UpperCamelCase__					,	mask_image=UpperCamelCase__					,	original_image=UpperCamelCase__					,	generator=UpperCamelCase__					,	num_inference_steps=2					,	output_type='''np'''					,	)

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

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

												SCREAMING_SNAKE_CASE							:					Any           =							torch.cuda.max_memory_allocated()
												assert mem_bytes < 4 * 10**9

												SCREAMING_SNAKE_CASE							:					Tuple           =							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(UpperCamelCase__					,	UpperCamelCase__		)






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





