import random
import unittest

import torch

from diffusers import IFSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin





@skip_mps
class 					__snake_case    (     _lowerCamelCase       ,_lowerCamelCase       ,unittest.TestCase      ):
					__lowerCamelCase			   =  IFSuperResolutionPipeline
					__lowerCamelCase			   =  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""width""", """height"""}
					__lowerCamelCase			   =  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
					__lowerCamelCase			   =  PipelineTesterMixin.required_optional_params - {"""latents"""}




					def      __a     (  self						)	->    List[Any]:



									'''simple docstring'''







									return self._get_superresolution_dummy_components()




					def      __a     (  self       ,			__UpperCamelCase       ,			__UpperCamelCase=0						)	->    Optional[Any]:



									'''simple docstring'''







									if str(__UpperCamelCase						).startswith('mps'						):
													snake_case__   :      Optional[int]   				= torch.manual_seed(__UpperCamelCase						)
									else:
													snake_case__   :      Dict   				= torch.Generator(device=__UpperCamelCase						).manual_seed(__UpperCamelCase						)

									snake_case__   :      Tuple   				= floats_tensor((1, 3, 32, 32)       ,			rng=random.Random(__UpperCamelCase						)						).to(__UpperCamelCase						)

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

									return inputs




					@unittest.skipIf(
					    torch_device != 'cuda' or not is_xformers_available()       ,			reason='XFormers attention is only available with CUDA and `xformers` installed'       ,			)
					def      __a     (  self						)	->    Union[str, Any]:



									'''simple docstring'''







									self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3						)




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



									'''simple docstring'''







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




					def      __a     (  self						)	->    List[str]:



									'''simple docstring'''







									self._test_attention_slicing_forward_pass(expected_max_diff=1E-2						)




					def      __a     (  self						)	->    Dict:



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



import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionDualGuidedPipeline
from diffusers.utils.testing_utils import load_image, nightly, require_torch_gpu, torch_device


lowerCAmelCase__	:  Union[str, Any]						=  False





@nightly
@require_torch_gpu
class 					__snake_case    (     unittest.TestCase      ):




					def      __a     (  self						)	->    str:



									'''simple docstring'''







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




					def      __a     (  self						)	->    Optional[int]:



									'''simple docstring'''







									snake_case__   :      str   				= VersatileDiffusionDualGuidedPipeline.from_pretrained('shi-labs/versatile-diffusion'						)
									# remove text_unet
									pipe.remove_unused_weights()
									pipe.to(__UpperCamelCase						)
									pipe.set_progress_bar_config(disable=__UpperCamelCase						)

									snake_case__   :      Union[str, Any]   				= load_image(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'						)

									snake_case__   :      Dict   				= torch.manual_seed(0						)
									snake_case__   :      Tuple   				= pipe(
									    prompt='first prompt'       ,			image=__UpperCamelCase       ,			text_to_image_strength=0.7_5       ,			generator=__UpperCamelCase       ,			guidance_scale=7.5       ,			num_inference_steps=2       ,			output_type='numpy'       ,			).images

									with tempfile.TemporaryDirectory() as tmpdirname:
													pipe.save_pretrained(__UpperCamelCase						)
													snake_case__   :      List[Any]   				= VersatileDiffusionDualGuidedPipeline.from_pretrained(__UpperCamelCase						)

									pipe.to(__UpperCamelCase						)
									pipe.set_progress_bar_config(disable=__UpperCamelCase						)

									snake_case__   :      str   				= generator.manual_seed(0						)
									snake_case__   :      List[str]   				= pipe(
									    prompt='first prompt'       ,			image=__UpperCamelCase       ,			text_to_image_strength=0.7_5       ,			generator=__UpperCamelCase       ,			guidance_scale=7.5       ,			num_inference_steps=2       ,			output_type='numpy'       ,			).images

									assert np.abs(image - new_image						).sum() < 1E-5, "Models don't have the same forward pass"




					def      __a     (  self						)	->    Optional[Any]:



									'''simple docstring'''







									snake_case__   :      Union[str, Any]   				= VersatileDiffusionDualGuidedPipeline.from_pretrained('shi-labs/versatile-diffusion'						)
									pipe.remove_unused_weights()
									pipe.to(__UpperCamelCase						)
									pipe.set_progress_bar_config(disable=__UpperCamelCase						)

									snake_case__   :      str   				= 'cyberpunk 2077'
									snake_case__   :      Dict   				= load_image(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'						)
									snake_case__   :      Union[str, Any]   				= torch.manual_seed(0						)
									snake_case__   :      Union[str, Any]   				= pipe(
									    prompt=__UpperCamelCase       ,			image=__UpperCamelCase       ,			text_to_image_strength=0.7_5       ,			generator=__UpperCamelCase       ,			guidance_scale=7.5       ,			num_inference_steps=50       ,			output_type='numpy'       ,			).images

									snake_case__   :      int   				= image[0, 253:256, 253:256, -1]

									assert image.shape == (1, 512, 512, 3)
									snake_case__   :      List[str]   				= np.array([0.0_7_8_7, 0.0_8_4_9, 0.0_8_2_6, 0.0_8_1_2, 0.0_8_0_7, 0.0_7_9_5, 0.0_8_1_8, 0.0_7_9_8, 0.0_7_7_9]						)

									assert np.abs(image_slice.flatten() - expected_slice						).max() < 1E-2


