






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  (	__UpperCAmelCase   ,    __UpperCAmelCase   ,    unittest.TestCase					):
			_lowercase		:    Union[str, Any]			     =							IFPipeline
			_lowercase		:    Optional[int]			     =							TEXT_TO_IMAGE_PARAMS - {'''width''', '''height''', '''latents'''}
			_lowercase		:    Any			     =							TEXT_TO_IMAGE_BATCH_PARAMS
			_lowercase		:    str			     =							PipelineTesterMixin.required_optional_params - {'''latents'''}






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[Any]:
						return self._get_dummy_components()






			def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a=0)      ->       Any:

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

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

						return inputs






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       int:
						self._test_save_load_optional_components()






			@unittest.skipIf(torch_device != 'cuda'		,					reason='float16 requires CUDA')
			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Optional[Any]:
						super().test_save_load_floataa(expected_max_diff=1E-1)






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[str]:
						self._test_attention_slicing_forward_pass(expected_max_diff=1E-2)






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[Any]:
						self._test_save_load_local()






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       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    SCREAMING_SNAKE_CASE__    (       self)      ->       str:
						self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3)



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






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       int:
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Tuple:
						SCREAMING_SNAKE_CASE						= IFPipeline.from_pretrained('DeepFloyd/IF-I-XL-v1.0'		,					variant='fp16'		,					torch_dtype=torch.floataa)

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

						# pre compute text embeddings and remove T5 to save memory

						pipe_a.text_encoder.to('cuda')

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

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

						SCREAMING_SNAKE_CASE						= None
						SCREAMING_SNAKE_CASE						= 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(a		,					a		,					a		,					a)

						pipe_a.remove_all_hooks()
						pipe_a.remove_all_hooks()

						# img2img

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

						pipe_a.remove_all_hooks()
						pipe_a.remove_all_hooks()

						# inpainting

						SCREAMING_SNAKE_CASE						= IFInpaintingPipeline(**pipe_a.components)
						SCREAMING_SNAKE_CASE						= 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(a		,					a		,					a		,					a)






			def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a		,					a)      ->       List[Any]:
						_start_torch_memory_measurement()

						SCREAMING_SNAKE_CASE						= torch.Generator(device='cpu').manual_seed(0)
						SCREAMING_SNAKE_CASE						= pipe_a(
						    prompt_embeds=a		,					negative_prompt_embeds=a		,					num_inference_steps=2		,					generator=a		,					output_type='np'		,					)

						SCREAMING_SNAKE_CASE						= output.images[0]

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

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

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

						# pipeline 2

						_start_torch_memory_measurement()

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

						SCREAMING_SNAKE_CASE						= floats_tensor((1, 3, 64, 64)		,					rng=random.Random(0)).to(a)

						SCREAMING_SNAKE_CASE						= pipe_a(
						    prompt_embeds=a		,					negative_prompt_embeds=a		,					image=a		,					generator=a		,					num_inference_steps=2		,					output_type='np'		,					)

						SCREAMING_SNAKE_CASE						= output.images[0]

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

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

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






			def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a		,					a)      ->       Any:
						_start_torch_memory_measurement()

						SCREAMING_SNAKE_CASE						= floats_tensor((1, 3, 64, 64)		,					rng=random.Random(0)).to(a)

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

						SCREAMING_SNAKE_CASE						= pipe_a(
						    prompt_embeds=a		,					negative_prompt_embeds=a		,					image=a		,					num_inference_steps=2		,					generator=a		,					output_type='np'		,					)

						SCREAMING_SNAKE_CASE						= output.images[0]

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

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

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

						# pipeline 2

						_start_torch_memory_measurement()

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

						SCREAMING_SNAKE_CASE						= floats_tensor((1, 3, 256, 256)		,					rng=random.Random(0)).to(a)
						SCREAMING_SNAKE_CASE						= floats_tensor((1, 3, 64, 64)		,					rng=random.Random(0)).to(a)

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

						SCREAMING_SNAKE_CASE						= output.images[0]

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

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

						SCREAMING_SNAKE_CASE						= 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(a		,					a)







			def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a		,					a)      ->       Tuple:
						_start_torch_memory_measurement()

						SCREAMING_SNAKE_CASE						= floats_tensor((1, 3, 64, 64)		,					rng=random.Random(0)).to(a)
						SCREAMING_SNAKE_CASE						= floats_tensor((1, 3, 64, 64)		,					rng=random.Random(1)).to(a)

						SCREAMING_SNAKE_CASE						= torch.Generator(device='cpu').manual_seed(0)
						SCREAMING_SNAKE_CASE						= pipe_a(
						    prompt_embeds=a		,					negative_prompt_embeds=a		,					image=a		,					mask_image=a		,					num_inference_steps=2		,					generator=a		,					output_type='np'		,					)

						SCREAMING_SNAKE_CASE						= output.images[0]

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

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

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

						# pipeline 2

						_start_torch_memory_measurement()

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

						SCREAMING_SNAKE_CASE						= floats_tensor((1, 3, 64, 64)		,					rng=random.Random(0)).to(a)
						SCREAMING_SNAKE_CASE						= floats_tensor((1, 3, 256, 256)		,					rng=random.Random(0)).to(a)
						SCREAMING_SNAKE_CASE						= floats_tensor((1, 3, 256, 256)		,					rng=random.Random(1)).to(a)

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

						SCREAMING_SNAKE_CASE						= output.images[0]

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

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

						SCREAMING_SNAKE_CASE						= 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(a		,					a)






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












import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
				import torch

if is_vision_available():
				from PIL import Image

				from transformers import MobileViTImageProcessor

class 		_snake_case  (  unittest.TestCase					):






			def __init__(       self		,					a		,					a=7		,					a=3		,					a=18		,					a=30		,					a=400		,					a=True		,					a=None		,					a=True		,					a=None		,					a=True		,					)      ->       Any:
						SCREAMING_SNAKE_CASE						= size if size is not None else {'shortest_edge': 20}
						SCREAMING_SNAKE_CASE						= crop_size if crop_size is not None else {'height': 18, 'width': 18}
						SCREAMING_SNAKE_CASE						= parent
						SCREAMING_SNAKE_CASE						= batch_size
						SCREAMING_SNAKE_CASE						= num_channels
						SCREAMING_SNAKE_CASE						= image_size
						SCREAMING_SNAKE_CASE						= min_resolution
						SCREAMING_SNAKE_CASE						= max_resolution
						SCREAMING_SNAKE_CASE						= do_resize
						SCREAMING_SNAKE_CASE						= size
						SCREAMING_SNAKE_CASE						= do_center_crop
						SCREAMING_SNAKE_CASE						= crop_size
						SCREAMING_SNAKE_CASE						= do_flip_channel_order







			def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[Any]:
						return {
						    "do_resize": self.do_resize,
						    "size": self.size,
						    "do_center_crop": self.do_center_crop,
						    "crop_size": self.crop_size,
						    "do_flip_channel_order": self.do_flip_channel_order,
						}



@require_torch
@require_vision
class 		_snake_case  (  A__   ,    unittest.TestCase					):
			_lowercase		:    Optional[int]			     =							MobileViTImageProcessor if is_vision_available() else None






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[str]:
						SCREAMING_SNAKE_CASE						= MobileViTImageProcessingTester(self)






			@property
			def    SCREAMING_SNAKE_CASE__    (       self)      ->       str:
						return self.image_processor_tester.prepare_image_processor_dict()






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       str:
						SCREAMING_SNAKE_CASE						= self.image_processing_class(**self.image_processor_dict)
						self.assertTrue(hasattr(a		,					'do_resize'))
						self.assertTrue(hasattr(a		,					'size'))
						self.assertTrue(hasattr(a		,					'do_center_crop'))
						self.assertTrue(hasattr(a		,					'center_crop'))
						self.assertTrue(hasattr(a		,					'do_flip_channel_order'))






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       str:
						SCREAMING_SNAKE_CASE						= self.image_processing_class.from_dict(self.image_processor_dict)
						self.assertEqual(image_processor.size		,					{'shortest_edge': 20})
						self.assertEqual(image_processor.crop_size		,					{'height': 18, 'width': 18})

						SCREAMING_SNAKE_CASE						= self.image_processing_class.from_dict(self.image_processor_dict		,					size=42		,					crop_size=84)
						self.assertEqual(image_processor.size		,					{'shortest_edge': 42})
						self.assertEqual(image_processor.crop_size		,					{'height': 84, 'width': 84})






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Tuple:
						pass






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Tuple:
						# Initialize image_processing
						SCREAMING_SNAKE_CASE						= self.image_processing_class(**self.image_processor_dict)
						# create random PIL images
						SCREAMING_SNAKE_CASE						= prepare_image_inputs(self.image_processor_tester		,					equal_resolution=a)
						for image in image_inputs:
									self.assertIsInstance(a		,					Image.Image)

						# Test not batched input
						SCREAMING_SNAKE_CASE						= image_processing(image_inputs[0]		,					return_tensors='pt').pixel_values
						self.assertEqual(
						    encoded_images.shape		,					(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['height'],
						        self.image_processor_tester.crop_size['width'],
						    )		,					)

						# Test batched
						SCREAMING_SNAKE_CASE						= image_processing(a		,					return_tensors='pt').pixel_values
						self.assertEqual(
						    encoded_images.shape		,					(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['height'],
						        self.image_processor_tester.crop_size['width'],
						    )		,					)






			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Dict:
						# Initialize image_processing
						SCREAMING_SNAKE_CASE						= self.image_processing_class(**self.image_processor_dict)
						# create random numpy tensors
						SCREAMING_SNAKE_CASE						= prepare_image_inputs(self.image_processor_tester		,					equal_resolution=a		,					numpify=a)
						for image in image_inputs:
									self.assertIsInstance(a		,					np.ndarray)

						# Test not batched input
						SCREAMING_SNAKE_CASE						= image_processing(image_inputs[0]		,					return_tensors='pt').pixel_values
						self.assertEqual(
						    encoded_images.shape		,					(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['height'],
						        self.image_processor_tester.crop_size['width'],
						    )		,					)

						# Test batched
						SCREAMING_SNAKE_CASE						= image_processing(a		,					return_tensors='pt').pixel_values
						self.assertEqual(
						    encoded_images.shape		,					(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['height'],
						        self.image_processor_tester.crop_size['width'],
						    )		,					)







			def    SCREAMING_SNAKE_CASE__    (       self)      ->       int:
						# Initialize image_processing
						SCREAMING_SNAKE_CASE						= self.image_processing_class(**self.image_processor_dict)
						# create random PyTorch tensors
						SCREAMING_SNAKE_CASE						= prepare_image_inputs(self.image_processor_tester		,					equal_resolution=a		,					torchify=a)
						for image in image_inputs:
									self.assertIsInstance(a		,					torch.Tensor)

						# Test not batched input
						SCREAMING_SNAKE_CASE						= image_processing(image_inputs[0]		,					return_tensors='pt').pixel_values
						self.assertEqual(
						    encoded_images.shape		,					(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['height'],
						        self.image_processor_tester.crop_size['width'],
						    )		,					)

						# Test batched
						SCREAMING_SNAKE_CASE						= image_processing(a		,					return_tensors='pt').pixel_values
						self.assertEqual(
						    encoded_images.shape		,					(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['height'],
						        self.image_processor_tester.crop_size['width'],
						    )		,					)




