






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						:List[str] =							IFPipeline
	__snake_case						:Optional[Any] =							TEXT_TO_IMAGE_PARAMS - {"width", "height", "latents"}
	__snake_case						:Optional[int] =							TEXT_TO_IMAGE_BATCH_PARAMS
	__snake_case						:Optional[Any] =							PipelineTesterMixin.required_optional_params - {"latents"}







	def 			_a (	self   :				Dict      )		->			Optional[Any]:
		"""simple docstring"""







		return self._get_dummy_components()







	def 			_a (	self   :				Any							,					_lowerCAmelCase   :				int							,					_lowerCAmelCase   :				Optional[Any]=0      )		->			List[str]:
		"""simple docstring"""







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

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

		return inputs







	def 			_a (	self   :				str      )		->			Tuple:
		"""simple docstring"""







		self._test_save_load_optional_components()







	@unittest.skipIf(torch_device != """cuda"""							,					reason="""float16 requires CUDA"""      )
	def 			_a (	self   :				Union[str, Any]      )		->			Any:
		"""simple docstring"""







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







	def 			_a (	self   :				Union[str, Any]      )		->			List[Any]:
		"""simple docstring"""







		self._test_attention_slicing_forward_pass(expected_max_diff=1e-2      )







	def 			_a (	self   :				Union[str, Any]      )		->			Union[str, Any]:
		"""simple docstring"""







		self._test_save_load_local()







	def 			_a (	self   :				Dict      )		->			List[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   :				Dict      )		->			List[Any]:
		"""simple docstring"""







		self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3      )


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







	def 			_a (	self   :				List[Any]      )		->			Any:
		"""simple docstring"""







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







	def 			_a (	self   :				Tuple      )		->			str:
		"""simple docstring"""







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

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

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

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

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

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

		__lowercase   							=	IFInpaintingPipeline(**pipe_a.components      )
		__lowercase   							=	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 			_a (	self   :				str							,					_lowerCAmelCase   :				Tuple							,					_lowerCAmelCase   :				List[str]							,					_lowerCAmelCase   :				Optional[Any]							,					_lowerCAmelCase   :				List[Any]      )		->			Tuple:
		"""simple docstring"""







		_start_torch_memory_measurement()

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

		__lowercase   							=	output.images[0]

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

		__lowercase   							=	torch.cuda.max_memory_allocated()
		assert mem_bytes < 13 * 10**9

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

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

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

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

		__lowercase   							=	output.images[0]

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

		__lowercase   							=	torch.cuda.max_memory_allocated()
		assert mem_bytes < 4 * 10**9

		__lowercase   							=	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 			_a (	self   :				str							,					_lowerCAmelCase   :				Any							,					_lowerCAmelCase   :				Union[str, Any]							,					_lowerCAmelCase   :				List[str]							,					_lowerCAmelCase   :				Optional[Any]      )		->			Dict:
		"""simple docstring"""







		_start_torch_memory_measurement()

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

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

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

		__lowercase   							=	output.images[0]

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

		__lowercase   							=	torch.cuda.max_memory_allocated()
		assert mem_bytes < 10 * 10**9

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

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

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

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

		__lowercase   							=	output.images[0]

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

		__lowercase   							=	torch.cuda.max_memory_allocated()
		assert mem_bytes < 4 * 10**9

		__lowercase   							=	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 			_a (	self   :				str							,					_lowerCAmelCase   :				Optional[Any]							,					_lowerCAmelCase   :				Tuple							,					_lowerCAmelCase   :				Optional[int]							,					_lowerCAmelCase   :				Tuple      )		->			Tuple:
		"""simple docstring"""







		_start_torch_memory_measurement()

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

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

		__lowercase   							=	output.images[0]

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

		__lowercase   							=	torch.cuda.max_memory_allocated()
		assert mem_bytes < 10 * 10**9

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

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

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

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

		__lowercase   							=	output.images[0]

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

		__lowercase   							=	torch.cuda.max_memory_allocated()
		assert mem_bytes < 4 * 10**9

		__lowercase   							=	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 typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
   from PIL import Image

   from ..image_utils import load_image

if is_tf_available():
   from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_VISION_2_SEQ_MAPPING

if is_torch_available():
   import torch

   from ..models.auto.modeling_auto import MODEL_FOR_VISION_2_SEQ_MAPPING

UpperCAmelCase_		 =					logging.get_logger(__name__)


@add_end_docstrings(_A)
class 	lowerCamelCase__	(       _A):



 """simple docstring"""

 def __init__(    self					:					Optional[int]			,						*__lowerCAmelCase					:					List[Any]			,						**__lowerCAmelCase					:					List[str]      )   ->	List[str]:
    super().__init__(*__lowerCAmelCase			,						**__lowerCAmelCase      )
    requires_backends(self			,						'''vision'''      )
    self.check_model_type(
        TF_MODEL_FOR_VISION_2_SEQ_MAPPING if self.framework == '''tf''' else MODEL_FOR_VISION_2_SEQ_MAPPING      )

 def     snake_case_      (    self					:					Any			,						__lowerCAmelCase					:					Dict=None			,						__lowerCAmelCase					:					List[str]=None			,						__lowerCAmelCase					:					Optional[Any]=None      )   ->	int:
    _A      =       {}
    _A      =       {}

    if prompt is not None:
       _A      =       prompt

    if generate_kwargs is not None:
       _A      =       generate_kwargs
    if max_new_tokens is not None:
       if "generate_kwargs" not in forward_kwargs:
          _A      =       {}
       if "max_new_tokens" in forward_kwargs["generate_kwargs"]:
          raise ValueError(
              '''\'max_new_tokens\' is defined twice, once in \'generate_kwargs\' and once as a direct parameter,'''
              ''' please use only one'''      )
       _A      =       max_new_tokens
    return preprocess_params, forward_kwargs, {}

 def __call__(    self					:					List[str]			,						__lowerCAmelCase					:					Union[str, List[str], "Image.Image", List["Image.Image"]]			,						**__lowerCAmelCase					:					Union[str, Any]      )   ->	Union[str, Any]:
    return super().__call__(__lowerCAmelCase			,						**__lowerCAmelCase      )

 def     snake_case_      (    self					:					List[Any]			,						__lowerCAmelCase					:					str			,						__lowerCAmelCase					:					Union[str, Any]=None      )   ->	int:
    _A      =       load_image(__lowerCAmelCase      )

    if prompt is not None:
       if not isinstance(__lowerCAmelCase			,						__lowerCAmelCase      ):
          raise ValueError(
              f'''Received an invalid text input, got - {type(__lowerCAmelCase      )} - but expected a single string. '''
              '''Note also that one single text can be provided for conditional image to text generation.'''      )

       _A      =       self.model.config.model_type

       if model_type == "git":
          _A      =       self.image_processor(images=__lowerCAmelCase			,						return_tensors=self.framework      )
          _A      =       self.tokenizer(text=__lowerCAmelCase			,						add_special_tokens=__lowerCAmelCase      ).input_ids
          _A      =       [self.tokenizer.cls_token_id] + input_ids
          _A      =       torch.tensor(__lowerCAmelCase      ).unsqueeze(0      )
          model_inputs.update({'''input_ids''': input_ids}      )

       elif model_type == "pix2struct":
          _A      =       self.image_processor(images=__lowerCAmelCase			,						header_text=__lowerCAmelCase			,						return_tensors=self.framework      )

       elif model_type != "vision-encoder-decoder":
          # vision-encoder-decoder does not support conditional generation
          _A      =       self.image_processor(images=__lowerCAmelCase			,						return_tensors=self.framework      )
          _A      =       self.tokenizer(__lowerCAmelCase			,						return_tensors=self.framework      )
          model_inputs.update(__lowerCAmelCase      )

       else:
          raise ValueError(f'''Model type {model_type} does not support conditional text generation'''      )

    else:
       _A      =       self.image_processor(images=__lowerCAmelCase			,						return_tensors=self.framework      )

    if self.model.config.model_type == "git" and prompt is None:
       _A      =       None

    return model_inputs

 def     snake_case_      (    self					:					Optional[int]			,						__lowerCAmelCase					:					Tuple			,						__lowerCAmelCase					:					Dict=None      )   ->	str:
    # Git model sets `model_inputs["input_ids"] = None` in `preprocess` (when `prompt=None`). In batch model, the
    # pipeline will group them into a list of `None`, which fail `_forward`. Avoid this by checking it first.
    if (
        "input_ids" in model_inputs
        and isinstance(model_inputs['''input_ids''']			,						__lowerCAmelCase      )
        and all(x is None for x in model_inputs['''input_ids''']      )
    ):
       _A      =       None

    if generate_kwargs is None:
       _A      =       {}
    # FIXME: We need to pop here due to a difference in how `generation.py` and `generation.tf_utils.py`
    #  parse inputs. In the Tensorflow version, `generate` raises an error if we don't use `input_ids` whereas
    #  the PyTorch version matches it with `self.model.main_input_name` or `self.model.encoder.main_input_name`
    #  in the `_prepare_model_inputs` method.
    _A      =       model_inputs.pop(self.model.main_input_name      )
    _A      =       self.model.generate(__lowerCAmelCase			,						**__lowerCAmelCase			,						**__lowerCAmelCase      )
    return model_outputs

 def     snake_case_      (    self					:					Dict			,						__lowerCAmelCase					:					Any      )   ->	Union[str, Any]:
    _A      =       []
    for output_ids in model_outputs:
       _A      =       {
           '''generated_text''': self.tokenizer.decode(
               __lowerCAmelCase			,						skip_special_tokens=__lowerCAmelCase			,						)
       }
       records.append(__lowerCAmelCase      )
    return records



