


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  snake_case_     (lowerCamelCase_			,		lowerCamelCase_			,		unittest.TestCase   ):
					UpperCAmelCase__      :					Tuple					   =    IFPipeline
					UpperCAmelCase__      :					Optional[int]					   =    TEXT_TO_IMAGE_PARAMS - {'''width''', '''height''', '''latents'''}
					UpperCAmelCase__      :					Optional[int]					   =    TEXT_TO_IMAGE_BATCH_PARAMS
					UpperCAmelCase__      :					int					   =    PipelineTesterMixin.required_optional_params - {'''latents'''}
					def 				lowerCamelCase__(   self     :Any   )      ->					List[Any]:
												return self._get_dummy_components()
					def 				lowerCamelCase__(   self     :List[Any]    ,__snake_case     :int    ,__snake_case     :int=0   )      ->					Optional[Any]:
												if str(__snake_case   ).startswith('mps'   ):
																			a__    =  torch.manual_seed(__snake_case   )
												else:
																			a__    =  torch.Generator(device=__snake_case   ).manual_seed(__snake_case   )

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

												return inputs
					def 				lowerCamelCase__(   self     :Dict   )      ->					str:
												self._test_save_load_optional_components()
					@unittest.skipIf(torch_device != 'cuda'    ,reason='float16 requires CUDA'   )
					def 				lowerCamelCase__(   self     :Optional[int]   )      ->					List[Any]:
												# 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 				lowerCamelCase__(   self     :Union[str, Any]   )      ->					Tuple:
												self._test_attention_slicing_forward_pass(expected_max_diff=1E-2   )
					def 				lowerCamelCase__(   self     :Any   )      ->					List[str]:
												self._test_save_load_local()
					def 				lowerCamelCase__(   self     :List[str]   )      ->					List[Any]:
												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     :Union[str, Any]   )      ->					Any:
												self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3   )







@slow
@require_torch_gpu
class  snake_case_     (unittest.TestCase   ):
					def 				lowerCamelCase__(   self     :List[str]   )      ->					int:
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()
					def 				lowerCamelCase__(   self     :Dict   )      ->					List[Any]:
												# if

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

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

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

												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(__snake_case    ,__snake_case    ,__snake_case    ,__snake_case   )
					def 				lowerCamelCase__(   self     :Optional[int]    ,__snake_case     :List[Any]    ,__snake_case     :List[str]    ,__snake_case     :Optional[Any]    ,__snake_case     :int   )      ->					Any:
												# pipeline 1

												_start_torch_memory_measurement()

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

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

												a__    =  pipe_a(
												    prompt_embeds=__snake_case    ,negative_prompt_embeds=__snake_case    ,image=__snake_case    ,generator=__snake_case    ,num_inference_steps=2    ,output_type='np'    ,)

												a__    =  output.images[0]

												assert image.shape == (2_56, 2_56, 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(__snake_case    ,__snake_case   )
					def 				lowerCamelCase__(   self     :Union[str, Any]    ,__snake_case     :int    ,__snake_case     :str    ,__snake_case     :Any    ,__snake_case     :Union[str, Any]   )      ->					List[Any]:
												# pipeline 1

												_start_torch_memory_measurement()

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

												a__    =  torch.Generator(device='cpu'   ).manual_seed(0   )

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

												# pipeline 2

												_start_torch_memory_measurement()

												a__    =  torch.Generator(device='cpu'   ).manual_seed(0   )

												a__    =  floats_tensor((1, 3, 2_56, 2_56)    ,rng=random.Random(0   )   ).to(__snake_case   )
												a__    =  floats_tensor((1, 3, 64, 64)    ,rng=random.Random(0   )   ).to(__snake_case   )

												a__    =  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'    ,)

												a__    =  output.images[0]

												assert image.shape == (2_56, 2_56, 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(__snake_case    ,__snake_case   )



					def 				lowerCamelCase__(   self     :Union[str, Any]    ,__snake_case     :Dict    ,__snake_case     :int    ,__snake_case     :List[Any]    ,__snake_case     :Dict   )      ->					Union[str, Any]:
												# pipeline 1

												_start_torch_memory_measurement()

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

												a__    =  torch.Generator(device='cpu'   ).manual_seed(0   )
												a__    =  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'    ,)

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

												# 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(__snake_case   )
												a__    =  floats_tensor((1, 3, 2_56, 2_56)    ,rng=random.Random(0   )   ).to(__snake_case   )
												a__    =  floats_tensor((1, 3, 2_56, 2_56)    ,rng=random.Random(1   )   ).to(__snake_case   )

												a__    =  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'    ,)

												a__    =  output.images[0]

												assert image.shape == (2_56, 2_56, 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(__snake_case    ,__snake_case   )






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








from typing import List, Optional, Union

import numpy as np
import PIL.Image

from ...image_processing_utils import BaseImageProcessor, BatchFeature
from ...image_transforms import rescale, resize, to_channel_dimension_format
from ...image_utils import (
    ChannelDimension,
    PILImageResampling,
    get_image_size,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


snake_case							: Dict      =   logging.get_logger(__name__)







class  snake_case_     (lowerCamelCase_   ):
					UpperCAmelCase__      :					Dict					   =    ['''pixel_values''']
					def __init__(   self     :Optional[Any]    ,__snake_case     :bool = True    ,__snake_case     :int = 32    ,__snake_case     :Union[str, Any]=PILImageResampling.BILINEAR    ,__snake_case     :bool = True    ,**__snake_case     :Tuple    ,)      ->					None:
												a__    =  do_resize
												a__    =  do_rescale
												a__    =  size_divisor
												a__    =  resample
												super().__init__(**__snake_case   )
					def 				lowerCamelCase__(   self     :Union[str, Any]    ,__snake_case     :np.ndarray    ,__snake_case     :int    ,__snake_case     :Tuple    ,__snake_case     :Optional[ChannelDimension] = None    ,**__snake_case     :List[Any]   )      ->					np.ndarray:
												a__   ,      a__    =  get_image_size(__snake_case   )
												# Rounds the height and width down to the closest multiple of size_divisor
												a__    =  height // size_divisor * size_divisor
												a__    =  width // size_divisor * size_divisor
												a__    =  resize(__snake_case    ,(new_h, new_w)    ,resample=__snake_case    ,data_format=__snake_case    ,**__snake_case   )
												return image
					def 				lowerCamelCase__(   self     :List[str]    ,__snake_case     :np.ndarray    ,__snake_case     :float    ,__snake_case     :Optional[ChannelDimension] = None    ,**__snake_case     :str   )      ->					np.ndarray:
												return rescale(image=__snake_case    ,scale=__snake_case    ,data_format=__snake_case    ,**__snake_case   )



					def 				lowerCamelCase__(   self     :Tuple    ,__snake_case     :Union["PIL.Image.Image", TensorType, List["PIL.Image.Image"], List[TensorType]]    ,__snake_case     :Optional[bool] = None    ,__snake_case     :Optional[int] = None    ,__snake_case     :Union[str, Any]=None    ,__snake_case     :Optional[bool] = None    ,__snake_case     :Optional[Union[TensorType, str]] = None    ,__snake_case     :ChannelDimension = ChannelDimension.FIRST    ,**__snake_case     :List[Any]    ,)      ->					BatchFeature:
												a__    =  do_resize if do_resize is not None else self.do_resize
												a__    =  do_rescale if do_rescale is not None else self.do_rescale
												a__    =  size_divisor if size_divisor is not None else self.size_divisor
												a__    =  resample if resample is not None else self.resample

												if do_resize and size_divisor is None:
																			raise ValueError('size_divisor is required for resizing'   )

												a__    =  make_list_of_images(__snake_case   )

												if not valid_images(__snake_case   ):
																			raise ValueError('Invalid image(s)'   )

												# All transformations expect numpy arrays.
												a__    =  [to_numpy_array(__snake_case   ) for img in images]

												if do_resize:
																			a__    =  [self.resize(__snake_case    ,size_divisor=__snake_case    ,resample=__snake_case   ) for image in images]

												if do_rescale:
																			a__    =  [self.rescale(__snake_case    ,scale=1 / 2_55   ) for image in images]

												a__    =  [to_channel_dimension_format(__snake_case    ,__snake_case   ) for image in images]

												a__    =  {'pixel_values': images}
												return BatchFeature(data=__snake_case    ,tensor_type=__snake_case   )




