






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      ):
      snake_case_	=				IFPipeline
      snake_case_	=				TEXT_TO_IMAGE_PARAMS - {"""width""", """height""", """latents"""}
      snake_case_	=				TEXT_TO_IMAGE_BATCH_PARAMS
      snake_case_	=				PipelineTesterMixin.required_optional_params - {"""latents"""}


      def 			__UpperCAmelCase		(						self						)->						Any:
           return self._get_dummy_components()


      def 			__UpperCAmelCase		(						self       ,    __UpperCamelCase       ,    __UpperCamelCase=0						)->						Optional[int]:
           if str(__UpperCamelCase						).startswith("mps"						):
                __lowerCAmelCase             =       torch.manual_seed(__UpperCamelCase						)
           else:
                __lowerCAmelCase             =       torch.Generator(device=__UpperCamelCase						).manual_seed(__UpperCamelCase						)

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

           return inputs


      def 			__UpperCAmelCase		(						self						)->						List[str]:
           self._test_save_load_optional_components()


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


      def 			__UpperCAmelCase		(						self						)->						int:
           self._test_save_load_local()


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



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


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


      def 			__UpperCAmelCase		(						self						)->						List[str]:
           # if

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

           __lowerCAmelCase             =       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"						)

           __lowerCAmelCase      ,			__lowerCAmelCase             =       pipe_a.encode_prompt("anime turtle"       ,    device="cuda"						)

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

           __lowerCAmelCase             =       None
           __lowerCAmelCase             =       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

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

           __lowerCAmelCase             =       IFInpaintingPipeline(**pipe_a.components						)
           __lowerCAmelCase             =       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 			__UpperCAmelCase		(						self       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase						)->						List[str]:
           # pipeline 1

           _start_torch_memory_measurement()

           __lowerCAmelCase             =       torch.Generator(device="cpu"						).manual_seed(0						)
           __lowerCAmelCase             =       pipe_a(
               prompt_embeds=__UpperCamelCase       ,    negative_prompt_embeds=__UpperCamelCase       ,    num_inference_steps=2       ,    generator=__UpperCamelCase       ,    output_type="np"       ,    )

           __lowerCAmelCase             =       output.images[0]

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

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

           __lowerCAmelCase             =       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()

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

           __lowerCAmelCase             =       floats_tensor((1, 3, 6_4, 6_4)       ,    rng=random.Random(0						)						).to(__UpperCamelCase						)

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

           __lowerCAmelCase             =       output.images[0]

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

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

           __lowerCAmelCase             =       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 			__UpperCAmelCase		(						self       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase						)->						Union[str, Any]:
           # pipeline 1

           _start_torch_memory_measurement()

           __lowerCAmelCase             =       floats_tensor((1, 3, 6_4, 6_4)       ,    rng=random.Random(0						)						).to(__UpperCamelCase						)

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

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

           __lowerCAmelCase             =       output.images[0]

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

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

           __lowerCAmelCase             =       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()

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

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

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

           __lowerCAmelCase             =       output.images[0]

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

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

           __lowerCAmelCase             =       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 			__UpperCAmelCase		(						self       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase						)->						Dict:
           # pipeline 1

           _start_torch_memory_measurement()

           __lowerCAmelCase             =       floats_tensor((1, 3, 6_4, 6_4)       ,    rng=random.Random(0						)						).to(__UpperCamelCase						)
           __lowerCAmelCase             =       floats_tensor((1, 3, 6_4, 6_4)       ,    rng=random.Random(1						)						).to(__UpperCamelCase						)

           __lowerCAmelCase             =       torch.Generator(device="cpu"						).manual_seed(0						)
           __lowerCAmelCase             =       pipe_a(
               prompt_embeds=__UpperCamelCase       ,    negative_prompt_embeds=__UpperCamelCase       ,    image=__UpperCamelCase       ,    mask_image=__UpperCamelCase       ,    num_inference_steps=2       ,    generator=__UpperCamelCase       ,    output_type="np"       ,    )

           __lowerCAmelCase             =       output.images[0]

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

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

           __lowerCAmelCase             =       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()

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

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

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

           __lowerCAmelCase             =       output.images[0]

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

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

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














import warnings
from typing import List, Optional, Union

from ...image_utils import ImageInput
from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType



class 				_UpperCamelCase      (a_      ):
      snake_case_	=				["""image_processor""", """tokenizer"""]
      snake_case_	=				"""FlavaImageProcessor"""
      snake_case_	=				("""BertTokenizer""", """BertTokenizerFast""")


      def __init__(						self       ,    __UpperCamelCase=None       ,    __UpperCamelCase=None       ,    **__UpperCamelCase						)->						Tuple:
           __lowerCAmelCase             =       None
           if "feature_extractor" in kwargs:
                warnings.warn(
                    "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
                    " instead."       ,    __UpperCamelCase       ,    )
                __lowerCAmelCase             =       kwargs.pop("feature_extractor"						)

           __lowerCAmelCase             =       image_processor if image_processor is not None else feature_extractor
           if image_processor is None:
                raise ValueError("You need to specify an `image_processor`."						)
           if tokenizer is None:
                raise ValueError("You need to specify a `tokenizer`."						)

           super().__init__(__UpperCamelCase       ,    __UpperCamelCase						)
           __lowerCAmelCase             =       self.image_processor


      def __call__(						self       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    __UpperCamelCase = True       ,    __UpperCamelCase = False       ,    __UpperCamelCase = False       ,    __UpperCamelCase = None       ,    __UpperCamelCase = 0       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    __UpperCamelCase = False       ,    __UpperCamelCase = False       ,    __UpperCamelCase = False       ,    __UpperCamelCase = False       ,    __UpperCamelCase = True       ,    __UpperCamelCase = None       ,    **__UpperCamelCase       ,    )->						int:

           if text is None and images is None:
                raise ValueError("You have to specify either text or images. Both cannot be none."						)

           if text is not None:
                __lowerCAmelCase             =       self.tokenizer(
                    text=__UpperCamelCase       ,    add_special_tokens=__UpperCamelCase       ,    padding=__UpperCamelCase       ,    truncation=__UpperCamelCase       ,    max_length=__UpperCamelCase       ,    stride=__UpperCamelCase       ,    pad_to_multiple_of=__UpperCamelCase       ,    return_token_type_ids=__UpperCamelCase       ,    return_attention_mask=__UpperCamelCase       ,    return_overflowing_tokens=__UpperCamelCase       ,    return_special_tokens_mask=__UpperCamelCase       ,    return_offsets_mapping=__UpperCamelCase       ,    return_length=__UpperCamelCase       ,    verbose=__UpperCamelCase       ,    return_tensors=__UpperCamelCase       ,    **__UpperCamelCase       ,    )
           if images is not None:
                __lowerCAmelCase             =       self.image_processor(
                    __UpperCamelCase       ,    return_image_mask=__UpperCamelCase       ,    return_codebook_pixels=__UpperCamelCase       ,    return_tensors=__UpperCamelCase       ,    **__UpperCamelCase       ,    )

           if text is not None and images is not None:
                encoding.update(__UpperCamelCase						)
                return encoding
           elif text is not None:
                return encoding
           else:
                return BatchEncoding(data=dict(**__UpperCamelCase						)       ,    tensor_type=__UpperCamelCase						)


      def 			__UpperCAmelCase		(						self       ,    *__UpperCamelCase       ,    **__UpperCamelCase						)->						int:
           return self.tokenizer.batch_decode(*__UpperCamelCase       ,    **__UpperCamelCase						)


      def 			__UpperCAmelCase		(						self       ,    *__UpperCamelCase       ,    **__UpperCamelCase						)->						Optional[int]:
           return self.tokenizer.decode(*__UpperCamelCase       ,    **__UpperCamelCase						)


      @property
      def 			__UpperCAmelCase		(						self						)->						Any:
           __lowerCAmelCase             =       self.tokenizer.model_input_names
           __lowerCAmelCase             =       self.image_processor.model_input_names
           return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names						)						)


      @property
      def 			__UpperCAmelCase		(						self						)->						Dict:
           warnings.warn(
               "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."       ,    __UpperCamelCase       ,    )
           return self.image_processor_class


      @property
      def 			__UpperCAmelCase		(						self						)->						Union[str, Any]:
           warnings.warn(
               "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."       ,    __UpperCamelCase       ,    )
           return self.image_processor






