
import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModelWithProjection

from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    PNDMScheduler,
    StableDiffusionImageVariationPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
class      A__       (						lowerCAmelCase__   ,						lowerCAmelCase__   ,						lowerCAmelCase__   ,						unittest.TestCase    ):
							lowerCAmelCase__		:						Optional[int]         =							StableDiffusionImageVariationPipeline
							lowerCAmelCase__		:						List[Any]         =							IMAGE_VARIATION_PARAMS
							lowerCAmelCase__		:						Union[str, Any]         =							IMAGE_VARIATION_BATCH_PARAMS
							lowerCAmelCase__		:						Any         =							frozenset([]    )
							# TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
							lowerCAmelCase__		:						Any         =							frozenset([]    )





							def       a__			(			self  :       List[str]    )						->		Optional[int]:



											"""simple docstring"""


											torch.manual_seed(0    )
											__lowercase								= UNetaDConditionModel(
											    block_out_channels=(32, 64)					,			layers_per_block=2					,			sample_size=32					,			in_channels=4					,			out_channels=4					,			down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')					,			up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')					,			cross_attention_dim=32					,			)
											__lowercase								= PNDMScheduler(skip_prk_steps=_UpperCAmelCase    )
											torch.manual_seed(0    )
											__lowercase								= AutoencoderKL(
											    block_out_channels=[32, 64]					,			in_channels=3					,			out_channels=3					,			down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']					,			up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']					,			latent_channels=4					,			)
											torch.manual_seed(0    )
											__lowercase								= CLIPVisionConfig(
											    hidden_size=32					,			projection_dim=32					,			intermediate_size=37					,			layer_norm_eps=1e-0_5					,			num_attention_heads=4					,			num_hidden_layers=5					,			image_size=32					,			patch_size=4					,			)
											__lowercase								= CLIPVisionModelWithProjection(_UpperCAmelCase    )
											__lowercase								= CLIPImageProcessor(crop_size=32					,			size=32    )

											__lowercase								= {
											    'unet': unet,
											    'scheduler': scheduler,
											    'vae': vae,
											    'image_encoder': image_encoder,
											    'feature_extractor': feature_extractor,
											    'safety_checker': None,
											}
											return components





							def       a__			(			self  :       Dict					,			_UpperCAmelCase  :       Tuple					,			_UpperCAmelCase  :       Optional[int]=0    )						->		Any:



											"""simple docstring"""


											__lowercase								= floats_tensor((1, 3, 32, 32)					,			rng=random.Random(_UpperCAmelCase    )    )
											__lowercase								= image.cpu().permute(0					,			2					,			3					,			1    )[0]
											__lowercase								= Image.fromarray(np.uinta(_UpperCAmelCase    )    ).convert('RGB'    ).resize((32, 32)    )
											if str(_UpperCAmelCase    ).startswith('mps'    ):
															__lowercase								= torch.manual_seed(_UpperCAmelCase    )
											else:
															__lowercase								= torch.Generator(device=_UpperCAmelCase    ).manual_seed(_UpperCAmelCase    )
											__lowercase								= {
											    'image': image,
											    'generator': generator,
											    'num_inference_steps': 2,
											    'guidance_scale': 6.0,
											    'output_type': 'numpy',
											}
											return inputs





							def       a__			(			self  :       Tuple    )						->		List[Any]:



											"""simple docstring"""


											__lowercase								= 'cpu'  # ensure determinism for the device-dependent torch.Generator
											__lowercase								= self.get_dummy_components()
											__lowercase								= StableDiffusionImageVariationPipeline(**_UpperCAmelCase    )
											__lowercase								= sd_pipe.to(_UpperCAmelCase    )
											sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase    )

											__lowercase								= self.get_dummy_inputs(_UpperCAmelCase    )
											__lowercase								= sd_pipe(**_UpperCAmelCase    ).images
											__lowercase								= image[0, -3:, -3:, -1]

											assert image.shape == (1, 64, 64, 3)
											__lowercase								= np.array([0.5_239, 0.5_723, 0.4_796, 0.5_049, 0.5_550, 0.4_685, 0.5_329, 0.4_891, 0.4_921]    )

											assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-3





							def       a__			(			self  :       Any    )						->		Any:



											"""simple docstring"""


											__lowercase								= 'cpu'  # ensure determinism for the device-dependent torch.Generator
											__lowercase								= self.get_dummy_components()
											__lowercase								= StableDiffusionImageVariationPipeline(**_UpperCAmelCase    )
											__lowercase								= sd_pipe.to(_UpperCAmelCase    )
											sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase    )

											__lowercase								= self.get_dummy_inputs(_UpperCAmelCase    )
											__lowercase								= 2 * [inputs['image']]
											__lowercase								= sd_pipe(**_UpperCAmelCase    )

											__lowercase								= output.images

											__lowercase								= image[-1, -3:, -3:, -1]

											assert image.shape == (2, 64, 64, 3)
											__lowercase								= np.array([0.6_892, 0.5_637, 0.5_836, 0.5_771, 0.6_254, 0.6_409, 0.5_580, 0.5_569, 0.5_289]    )

											assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-3





							def       a__			(			self  :       Dict    )						->		str:



											"""simple docstring"""


											super().test_inference_batch_single_identical(expected_max_diff=3e-3    )






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





							def       a__			(			self  :       Optional[int]    )						->		str:



											"""simple docstring"""


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





							def       a__			(			self  :       Tuple					,			_UpperCAmelCase  :       Union[str, Any]					,			_UpperCAmelCase  :       Union[str, Any]="cpu"					,			_UpperCAmelCase  :       Optional[Any]=torch.floataa					,			_UpperCAmelCase  :       Optional[int]=0    )						->		List[Any]:



											"""simple docstring"""


											__lowercase								= torch.Generator(device=_UpperCAmelCase    ).manual_seed(_UpperCAmelCase    )
											__lowercase								= load_image(
											    'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
											    '/stable_diffusion_imgvar/input_image_vermeer.png'    )
											__lowercase								= np.random.RandomState(_UpperCAmelCase    ).standard_normal((1, 4, 64, 64)    )
											__lowercase								= torch.from_numpy(_UpperCAmelCase    ).to(device=_UpperCAmelCase					,			dtype=_UpperCAmelCase    )
											__lowercase								= {
											    'image': init_image,
											    'latents': latents,
											    'generator': generator,
											    'num_inference_steps': 3,
											    'guidance_scale': 7.5,
											    'output_type': 'numpy',
											}
											return inputs





							def       a__			(			self  :       Dict    )						->		str:



											"""simple docstring"""


											__lowercase								= StableDiffusionImageVariationPipeline.from_pretrained(
											    'lambdalabs/sd-image-variations-diffusers'					,			safety_checker=_UpperCAmelCase    )
											__lowercase								= sd_pipe.to(_UpperCAmelCase    )
											sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase    )

											__lowercase								= self.get_inputs(_UpperCAmelCase    )
											__lowercase								= sd_pipe(**_UpperCAmelCase    ).images
											__lowercase								= image[0, -3:, -3:, -1].flatten()

											assert image.shape == (1, 5_12, 5_12, 3)
											__lowercase								= np.array([0.84_491, 0.90_789, 0.75_708, 0.78_734, 0.83_485, 0.70_099, 0.66_938, 0.68_727, 0.61_379]    )
											assert np.abs(image_slice - expected_slice    ).max() < 6e-3





							def       a__			(			self  :       Optional[int]    )						->		List[str]:



											"""simple docstring"""


											__lowercase								= 0

											def callback_fn(_UpperCAmelCase  :       int					,			_UpperCAmelCase  :       int					,			_UpperCAmelCase  :       torch.FloatTensor    ) -> None:
															__lowercase								= True
															nonlocal number_of_steps
															number_of_steps += 1
															if step == 1:
																			__lowercase								= latents.detach().cpu().numpy()
																			assert latents.shape == (1, 4, 64, 64)
																			__lowercase								= latents[0, -3:, -3:, -1]
																			__lowercase								= np.array(
																			    [-0.1_621, 0.2_837, -0.7_979, -0.1_221, -1.3_057, 0.7_681, -2.1_191, 0.0_464, 1.6_309]    )

																			assert np.abs(latents_slice.flatten() - expected_slice    ).max() < 5e-2
															elif step == 2:
																			__lowercase								= latents.detach().cpu().numpy()
																			assert latents.shape == (1, 4, 64, 64)
																			__lowercase								= latents[0, -3:, -3:, -1]
																			__lowercase								= np.array([0.6_299, 1.7_500, 1.1_992, -2.1_582, -1.8_994, 0.7_334, -0.7_090, 1.0_137, 1.5_273]    )

																			assert np.abs(latents_slice.flatten() - expected_slice    ).max() < 5e-2

											__lowercase								= False

											__lowercase								= StableDiffusionImageVariationPipeline.from_pretrained(
											    'fusing/sd-image-variations-diffusers'					,			safety_checker=_UpperCAmelCase					,			torch_dtype=torch.floataa					,			)
											pipe.to(_UpperCAmelCase    )
											pipe.set_progress_bar_config(disable=_UpperCAmelCase    )
											pipe.enable_attention_slicing()

											__lowercase								= self.get_inputs(_UpperCAmelCase					,			dtype=torch.floataa    )
											pipe(**_UpperCAmelCase					,			callback=_UpperCAmelCase					,			callback_steps=1    )
											assert callback_fn.has_been_called
											assert number_of_steps == inputs["num_inference_steps"]





							def       a__			(			self  :       Dict    )						->		str:



											"""simple docstring"""


											torch.cuda.empty_cache()
											torch.cuda.reset_max_memory_allocated()
											torch.cuda.reset_peak_memory_stats()

											__lowercase								= 'fusing/sd-image-variations-diffusers'
											__lowercase								= StableDiffusionImageVariationPipeline.from_pretrained(
											    _UpperCAmelCase					,			safety_checker=_UpperCAmelCase					,			torch_dtype=torch.floataa    )
											__lowercase								= pipe.to(_UpperCAmelCase    )
											pipe.set_progress_bar_config(disable=_UpperCAmelCase    )
											pipe.enable_attention_slicing(1    )
											pipe.enable_sequential_cpu_offload()

											__lowercase								= self.get_inputs(_UpperCAmelCase					,			dtype=torch.floataa    )
											__lowercase								= pipe(**_UpperCAmelCase    )

											__lowercase								= torch.cuda.max_memory_allocated()
											# make sure that less than 2.6 GB is allocated
											assert mem_bytes < 2.6 * 10**9






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





							def       a__			(			self  :       Optional[int]    )						->		List[str]:



											"""simple docstring"""


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





							def       a__			(			self  :       List[Any]					,			_UpperCAmelCase  :       Any					,			_UpperCAmelCase  :       Union[str, Any]="cpu"					,			_UpperCAmelCase  :       Dict=torch.floataa					,			_UpperCAmelCase  :       List[Any]=0    )						->		Union[str, Any]:



											"""simple docstring"""


											__lowercase								= torch.Generator(device=_UpperCAmelCase    ).manual_seed(_UpperCAmelCase    )
											__lowercase								= load_image(
											    'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
											    '/stable_diffusion_imgvar/input_image_vermeer.png'    )
											__lowercase								= np.random.RandomState(_UpperCAmelCase    ).standard_normal((1, 4, 64, 64)    )
											__lowercase								= torch.from_numpy(_UpperCAmelCase    ).to(device=_UpperCAmelCase					,			dtype=_UpperCAmelCase    )
											__lowercase								= {
											    'image': init_image,
											    'latents': latents,
											    'generator': generator,
											    'num_inference_steps': 50,
											    'guidance_scale': 7.5,
											    'output_type': 'numpy',
											}
											return inputs





							def       a__			(			self  :       Tuple    )						->		Any:



											"""simple docstring"""


											__lowercase								= StableDiffusionImageVariationPipeline.from_pretrained('fusing/sd-image-variations-diffusers'    )
											sd_pipe.to(_UpperCAmelCase    )
											sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase    )

											__lowercase								= self.get_inputs(_UpperCAmelCase    )
											__lowercase								= sd_pipe(**_UpperCAmelCase    ).images[0]

											__lowercase								= load_numpy(
											    'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
											    '/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy'    )
											__lowercase								= np.abs(expected_image - image    ).max()
											assert max_diff < 1e-3





							def       a__			(			self  :       Dict    )						->		Optional[int]:



											"""simple docstring"""


											__lowercase								= StableDiffusionImageVariationPipeline.from_pretrained('fusing/sd-image-variations-diffusers'    )
											__lowercase								= DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config    )
											sd_pipe.to(_UpperCAmelCase    )
											sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase    )

											__lowercase								= self.get_inputs(_UpperCAmelCase    )
											__lowercase								= 25
											__lowercase								= sd_pipe(**_UpperCAmelCase    ).images[0]

											__lowercase								= load_numpy(
											    'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
											    '/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy'    )
											__lowercase								= np.abs(expected_image - image    ).max()
											assert max_diff < 1e-3







import unittest

from transformers import load_tool
from transformers.tools.agent_types import AGENT_TYPE_MAPPING

from .test_tools_common import ToolTesterMixin, output_types
class      A__       (						unittest.TestCase   ,						lowerCAmelCase__    ):





							def       a__			(			self  :       Any    )						->		Dict:



											"""simple docstring"""


											__lowercase								= load_tool('translation'    )
											self.tool.setup()
											__lowercase								= load_tool('translation'					,			remote=_UpperCAmelCase    )





							def       a__			(			self  :       int    )						->		Optional[Any]:



											"""simple docstring"""


											__lowercase								= self.tool('Hey, what\'s up?'					,			src_lang='English'					,			tgt_lang='French'    )
											self.assertEqual(_UpperCAmelCase					,			'- Hé, comment ça va?'    )





							def       a__			(			self  :       Tuple    )						->		Optional[int]:



											"""simple docstring"""


											__lowercase								= self.remote_tool('Hey, what\'s up?'					,			src_lang='English'					,			tgt_lang='French'    )
											self.assertEqual(_UpperCAmelCase					,			'- Hé, comment ça va?'    )





							def       a__			(			self  :       Union[str, Any]    )						->		Tuple:



											"""simple docstring"""


											__lowercase								= self.tool(text='Hey, what\'s up?'					,			src_lang='English'					,			tgt_lang='French'    )
											self.assertEqual(_UpperCAmelCase					,			'- Hé, comment ça va?'    )





							def       a__			(			self  :       Tuple    )						->		Optional[Any]:



											"""simple docstring"""


											__lowercase								= self.remote_tool(text='Hey, what\'s up?'					,			src_lang='English'					,			tgt_lang='French'    )
											self.assertEqual(_UpperCAmelCase					,			'- Hé, comment ça va?'    )





							def       a__			(			self  :       List[Any]    )						->		int:



											"""simple docstring"""


											__lowercase								= ['Hey, what\'s up?', 'English', 'Spanish']
											__lowercase								= self.tool(*_UpperCAmelCase    )

											# There is a single output
											if len(self.tool.outputs    ) == 1:
															__lowercase								= [outputs]

											self.assertListEqual(output_types(_UpperCAmelCase    )					,			self.tool.outputs    )





							def       a__			(			self  :       List[str]    )						->		Any:



											"""simple docstring"""


											__lowercase								= ['Hey, what\'s up?', 'English', 'Spanish']
											__lowercase								= self.tool(*_UpperCAmelCase    )

											if not isinstance(_UpperCAmelCase					,			_UpperCAmelCase    ):
															__lowercase								= [outputs]

											self.assertEqual(len(_UpperCAmelCase    )					,			len(self.tool.outputs    )    )

											for output, output_type in zip(_UpperCAmelCase					,			self.tool.outputs    ):
															__lowercase								= AGENT_TYPE_MAPPING[output_type]
															self.assertTrue(isinstance(_UpperCAmelCase					,			_UpperCAmelCase    )    )





							def       a__			(			self  :       Tuple    )						->		List[Any]:



											"""simple docstring"""


											__lowercase								= ['Hey, what\'s up?', 'English', 'Spanish']

											__lowercase								= []

											for _input, input_type in zip(_UpperCAmelCase					,			self.tool.inputs    ):
															if isinstance(_UpperCAmelCase					,			_UpperCAmelCase    ):
																			_inputs.append([AGENT_TYPE_MAPPING[_input_type](_input    ) for _input_type in input_type]    )
															else:
																			_inputs.append(AGENT_TYPE_MAPPING[input_type](_input    )    )

        # Should not raise an error
											__lowercase								= self.tool(*_UpperCAmelCase    )

											if not isinstance(_UpperCAmelCase					,			_UpperCAmelCase    ):
															__lowercase								= [outputs]

											self.assertEqual(len(_UpperCAmelCase    )					,			len(self.tool.outputs    )    )





