



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     __snake_case						(					UpperCamelCase_						,UpperCamelCase_						,UpperCamelCase_						,unittest.TestCase						):
							_a      				=			StableDiffusionImageVariationPipeline
							_a      				=			IMAGE_VARIATION_PARAMS
							_a      				=			IMAGE_VARIATION_BATCH_PARAMS
							_a      				=			frozenset([]						)
							# TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
							_a      				=			frozenset([]						)




							def        UpperCAmelCase__      (	self					:	Optional[Any]):
									torch.manual_seed(0)
									lowerCAmelCase_      :					Optional[Any]     =			UNetaDConditionModel(
									    block_out_channels=(3_2, 6_4)    ,    layers_per_block=2    ,    sample_size=3_2    ,    in_channels=4    ,    out_channels=4    ,    down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')    ,    up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')    ,    cross_attention_dim=3_2    ,    )
									lowerCAmelCase_      :					Optional[int]     =			PNDMScheduler(skip_prk_steps=A_)
									torch.manual_seed(0)
									lowerCAmelCase_      :					Optional[int]     =			AutoencoderKL(
									    block_out_channels=[3_2, 6_4]    ,    in_channels=3    ,    out_channels=3    ,    down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']    ,    up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']    ,    latent_channels=4    ,    )
									torch.manual_seed(0)
									lowerCAmelCase_      :					int     =			CLIPVisionConfig(
									    hidden_size=3_2    ,    projection_dim=3_2    ,    intermediate_size=3_7    ,    layer_norm_eps=1e-05    ,    num_attention_heads=4    ,    num_hidden_layers=5    ,    image_size=3_2    ,    patch_size=4    ,    )
									lowerCAmelCase_      :					str     =			CLIPVisionModelWithProjection(A_)
									lowerCAmelCase_      :					Optional[Any]     =			CLIPImageProcessor(crop_size=3_2    ,    size=3_2)

									lowerCAmelCase_      :					Union[str, Any]     =			{
									    '''unet''': unet,
									    '''scheduler''': scheduler,
									    '''vae''': vae,
									    '''image_encoder''': image_encoder,
									    '''feature_extractor''': feature_extractor,
									    '''safety_checker''': None,
									}
									return components




							def        UpperCAmelCase__      (	self					:	Optional[Any]    ,    A_					:	List[str]    ,    A_					:	Tuple=0):
									lowerCAmelCase_      :					List[str]     =			floats_tensor((1, 3, 3_2, 3_2)    ,    rng=random.Random(A_))
									lowerCAmelCase_      :					List[Any]     =			image.cpu().permute(0    ,    2    ,    3    ,    1)[0]
									lowerCAmelCase_      :					List[str]     =			Image.fromarray(np.uinta(A_)).convert('''RGB''').resize((3_2, 3_2))
									if str(A_).startswith('''mps'''):
											lowerCAmelCase_      :					Tuple     =			torch.manual_seed(A_)
									else:
											lowerCAmelCase_      :					Optional[Any]     =			torch.Generator(device=A_).manual_seed(A_)
									lowerCAmelCase_      :					Optional[Any]     =			{
									    '''image''': image,
									    '''generator''': generator,
									    '''num_inference_steps''': 2,
									    '''guidance_scale''': 6.0,
									    '''output_type''': '''numpy''',
									}
									return inputs




							def        UpperCAmelCase__      (	self					:	List[Any]):
									lowerCAmelCase_      :					str     =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									lowerCAmelCase_      :					List[Any]     =			self.get_dummy_components()
									lowerCAmelCase_      :					Optional[Any]     =			StableDiffusionImageVariationPipeline(**A_)
									lowerCAmelCase_      :					List[str]     =			sd_pipe.to(A_)
									sd_pipe.set_progress_bar_config(disable=A_)

									lowerCAmelCase_      :					Dict     =			self.get_dummy_inputs(A_)
									lowerCAmelCase_      :					Dict     =			sd_pipe(**A_).images
									lowerCAmelCase_      :					List[str]     =			image[0, -3:, -3:, -1]

									assert image.shape == (1, 6_4, 6_4, 3)
									lowerCAmelCase_      :					List[str]     =			np.array([0.5239, 0.5723, 0.4796, 0.5049, 0.5550, 0.4685, 0.5329, 0.4891, 0.4921])

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




							def        UpperCAmelCase__      (	self					:	List[Any]):
									lowerCAmelCase_      :					List[str]     =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									lowerCAmelCase_      :					List[Any]     =			self.get_dummy_components()
									lowerCAmelCase_      :					Tuple     =			StableDiffusionImageVariationPipeline(**A_)
									lowerCAmelCase_      :					Dict     =			sd_pipe.to(A_)
									sd_pipe.set_progress_bar_config(disable=A_)

									lowerCAmelCase_      :					Tuple     =			self.get_dummy_inputs(A_)
									lowerCAmelCase_      :					Any     =			2 * [inputs['''image''']]
									lowerCAmelCase_      :					Tuple     =			sd_pipe(**A_)

									lowerCAmelCase_      :					List[str]     =			output.images

									lowerCAmelCase_      :					Any     =			image[-1, -3:, -3:, -1]

									assert image.shape == (2, 6_4, 6_4, 3)
									lowerCAmelCase_      :					List[Any]     =			np.array([0.6892, 0.5637, 0.5836, 0.5771, 0.6254, 0.6409, 0.5580, 0.5569, 0.5289])

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







							def        UpperCAmelCase__      (	self					:	Optional[int]):
									super().test_inference_batch_single_identical(expected_max_diff=3e-3)





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




							def        UpperCAmelCase__      (	self					:	Any):
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()




							def        UpperCAmelCase__      (	self					:	str    ,    A_					:	Union[str, Any]    ,    A_					:	Any="cpu"    ,    A_					:	Tuple=torch.floataa    ,    A_					:	str=0):
									lowerCAmelCase_      :					int     =			torch.Generator(device=A_).manual_seed(A_)
									lowerCAmelCase_      :					Tuple     =			load_image(
									    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'''
									    '''/stable_diffusion_imgvar/input_image_vermeer.png''')
									lowerCAmelCase_      :					Any     =			np.random.RandomState(A_).standard_normal((1, 4, 6_4, 6_4))
									lowerCAmelCase_      :					int     =			torch.from_numpy(A_).to(device=A_    ,    dtype=A_)
									lowerCAmelCase_      :					str     =			{
									    '''image''': init_image,
									    '''latents''': latents,
									    '''generator''': generator,
									    '''num_inference_steps''': 3,
									    '''guidance_scale''': 7.5,
									    '''output_type''': '''numpy''',
									}
									return inputs




							def        UpperCAmelCase__      (	self					:	Optional[Any]):
									lowerCAmelCase_      :					Any     =			StableDiffusionImageVariationPipeline.from_pretrained(
									    '''lambdalabs/sd-image-variations-diffusers'''    ,    safety_checker=A_)
									lowerCAmelCase_      :					Tuple     =			sd_pipe.to(A_)
									sd_pipe.set_progress_bar_config(disable=A_)

									lowerCAmelCase_      :					int     =			self.get_inputs(A_)
									lowerCAmelCase_      :					int     =			sd_pipe(**A_).images
									lowerCAmelCase_      :					List[Any]     =			image[0, -3:, -3:, -1].flatten()

									assert image.shape == (1, 5_1_2, 5_1_2, 3)
									lowerCAmelCase_      :					List[str]     =			np.array([0.8_4491, 0.9_0789, 0.7_5708, 0.7_8734, 0.8_3485, 0.7_0099, 0.6_6938, 0.6_8727, 0.6_1379])
									assert np.abs(image_slice - expected_slice).max() < 6e-3




							def        UpperCAmelCase__      (	self					:	str):
									lowerCAmelCase_      :					Union[str, Any]     =			0

									def callback_fn(A_					:	int    ,    A_					:	int    ,    A_					:	torch.FloatTensor) -> None:
											lowerCAmelCase_      :					Dict     =			True
											nonlocal number_of_steps
											number_of_steps += 1
											if step == 1:
													lowerCAmelCase_      :					List[str]     =			latents.detach().cpu().numpy()
													assert latents.shape == (1, 4, 6_4, 6_4)
													lowerCAmelCase_      :					List[Any]     =			latents[0, -3:, -3:, -1]
													lowerCAmelCase_      :					List[str]     =			np.array(
													    [-0.1621, 0.2837, -0.7979, -0.1221, -1.3057, 0.7681, -2.1191, 0.0464, 1.6309])

													assert np.abs(latents_slice.flatten() - expected_slice).max() < 5e-2
											elif step == 2:
													lowerCAmelCase_      :					Tuple     =			latents.detach().cpu().numpy()
													assert latents.shape == (1, 4, 6_4, 6_4)
													lowerCAmelCase_      :					List[Any]     =			latents[0, -3:, -3:, -1]
													lowerCAmelCase_      :					Dict     =			np.array([0.6299, 1.7500, 1.1992, -2.1582, -1.8994, 0.7334, -0.7090, 1.0137, 1.5273])

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

									lowerCAmelCase_      :					Optional[Any]     =			False

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

									lowerCAmelCase_      :					List[Any]     =			self.get_inputs(A_    ,    dtype=torch.floataa)
									pipe(**A_    ,    callback=A_    ,    callback_steps=1)
									assert callback_fn.has_been_called
									assert number_of_steps == inputs["num_inference_steps"]







							def        UpperCAmelCase__      (	self					:	Any):
									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()
									torch.cuda.reset_peak_memory_stats()

									lowerCAmelCase_      :					Any     =			'''fusing/sd-image-variations-diffusers'''
									lowerCAmelCase_      :					List[Any]     =			StableDiffusionImageVariationPipeline.from_pretrained(
									    A_    ,    safety_checker=A_    ,    torch_dtype=torch.floataa)
									lowerCAmelCase_      :					str     =			pipe.to(A_)
									pipe.set_progress_bar_config(disable=A_)
									pipe.enable_attention_slicing(1)
									pipe.enable_sequential_cpu_offload()

									lowerCAmelCase_      :					str     =			self.get_inputs(A_    ,    dtype=torch.floataa)
									lowerCAmelCase_      :					Optional[Any]     =			pipe(**A_)

									lowerCAmelCase_      :					Tuple     =			torch.cuda.max_memory_allocated()
									# make sure that less than 2.6 GB is allocated
									assert mem_bytes < 2.6 * 1_0**9





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




							def        UpperCAmelCase__      (	self					:	Union[str, Any]):
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()




							def        UpperCAmelCase__      (	self					:	int    ,    A_					:	Any    ,    A_					:	List[Any]="cpu"    ,    A_					:	Optional[int]=torch.floataa    ,    A_					:	List[str]=0):
									lowerCAmelCase_      :					Union[str, Any]     =			torch.Generator(device=A_).manual_seed(A_)
									lowerCAmelCase_      :					Optional[Any]     =			load_image(
									    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'''
									    '''/stable_diffusion_imgvar/input_image_vermeer.png''')
									lowerCAmelCase_      :					Dict     =			np.random.RandomState(A_).standard_normal((1, 4, 6_4, 6_4))
									lowerCAmelCase_      :					Optional[Any]     =			torch.from_numpy(A_).to(device=A_    ,    dtype=A_)
									lowerCAmelCase_      :					Union[str, Any]     =			{
									    '''image''': init_image,
									    '''latents''': latents,
									    '''generator''': generator,
									    '''num_inference_steps''': 5_0,
									    '''guidance_scale''': 7.5,
									    '''output_type''': '''numpy''',
									}
									return inputs




							def        UpperCAmelCase__      (	self					:	int):
									lowerCAmelCase_      :					Union[str, Any]     =			StableDiffusionImageVariationPipeline.from_pretrained('''fusing/sd-image-variations-diffusers''')
									sd_pipe.to(A_)
									sd_pipe.set_progress_bar_config(disable=A_)

									lowerCAmelCase_      :					Optional[Any]     =			self.get_inputs(A_)
									lowerCAmelCase_      :					Dict     =			sd_pipe(**A_).images[0]

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







							def        UpperCAmelCase__      (	self					:	Union[str, Any]):
									lowerCAmelCase_      :					List[str]     =			StableDiffusionImageVariationPipeline.from_pretrained('''fusing/sd-image-variations-diffusers''')
									lowerCAmelCase_      :					str     =			DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config)
									sd_pipe.to(A_)
									sd_pipe.set_progress_bar_config(disable=A_)

									lowerCAmelCase_      :					Dict     =			self.get_inputs(A_)
									lowerCAmelCase_      :					Optional[int]     =			2_5
									lowerCAmelCase_      :					str     =			sd_pipe(**A_).images[0]

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









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_video_inputs


if is_torch_available():
						import torch

if is_vision_available():
						from PIL import Image

						from transformers import VideoMAEImageProcessor



class     __snake_case						(					unittest.TestCase						):




							def __init__(	self					:	str    ,    A_					:	str    ,    A_					:	int=7    ,    A_					:	List[Any]=3    ,    A_					:	Dict=1_0    ,    A_					:	List[str]=1_8    ,    A_					:	List[Any]=3_0    ,    A_					:	str=4_0_0    ,    A_					:	str=True    ,    A_					:	Tuple=None    ,    A_					:	Optional[Any]=True    ,    A_					:	Tuple=[0.5, 0.5, 0.5]    ,    A_					:	Any=[0.5, 0.5, 0.5]    ,    A_					:	List[str]=None    ,    ):
									lowerCAmelCase_      :					Optional[Any]     =			size if size is not None else {'''shortest_edge''': 1_8}
									lowerCAmelCase_      :					Optional[int]     =			crop_size if crop_size is not None else {'''height''': 1_8, '''width''': 1_8}

									lowerCAmelCase_      :					Any     =			parent
									lowerCAmelCase_      :					Optional[Any]     =			batch_size
									lowerCAmelCase_      :					List[Any]     =			num_channels
									lowerCAmelCase_      :					Optional[int]     =			num_frames
									lowerCAmelCase_      :					Optional[int]     =			image_size
									lowerCAmelCase_      :					List[Any]     =			min_resolution
									lowerCAmelCase_      :					str     =			max_resolution
									lowerCAmelCase_      :					List[str]     =			do_resize
									lowerCAmelCase_      :					int     =			size
									lowerCAmelCase_      :					Union[str, Any]     =			do_normalize
									lowerCAmelCase_      :					Union[str, Any]     =			image_mean
									lowerCAmelCase_      :					List[str]     =			image_std
									lowerCAmelCase_      :					Any     =			crop_size







							def        UpperCAmelCase__      (	self					:	Any):
									return {
									    "image_mean": self.image_mean,
									    "image_std": self.image_std,
									    "do_normalize": self.do_normalize,
									    "do_resize": self.do_resize,
									    "size": self.size,
									    "crop_size": self.crop_size,
									}





@require_torch
@require_vision
class     __snake_case						(					UpperCamelCase_						,unittest.TestCase						):
							_a      				=			VideoMAEImageProcessor if is_vision_available() else None




							def        UpperCAmelCase__      (	self					:	Tuple):
									lowerCAmelCase_      :					int     =			VideoMAEImageProcessingTester(self)




							@property
							def        UpperCAmelCase__      (	self					:	Tuple):
									return self.image_processor_tester.prepare_image_processor_dict()




							def        UpperCAmelCase__      (	self					:	Optional[int]):
									lowerCAmelCase_      :					int     =			self.image_processing_class(**self.image_processor_dict)
									self.assertTrue(hasattr(A_    ,    '''image_mean'''))
									self.assertTrue(hasattr(A_    ,    '''image_std'''))
									self.assertTrue(hasattr(A_    ,    '''do_normalize'''))
									self.assertTrue(hasattr(A_    ,    '''do_resize'''))
									self.assertTrue(hasattr(A_    ,    '''do_center_crop'''))
									self.assertTrue(hasattr(A_    ,    '''size'''))




							def        UpperCAmelCase__      (	self					:	Dict):
									lowerCAmelCase_      :					int     =			self.image_processing_class.from_dict(self.image_processor_dict)
									self.assertEqual(image_processor.size    ,    {'''shortest_edge''': 1_8})
									self.assertEqual(image_processor.crop_size    ,    {'''height''': 1_8, '''width''': 1_8})

									lowerCAmelCase_      :					str     =			self.image_processing_class.from_dict(self.image_processor_dict    ,    size=4_2    ,    crop_size=8_4)
									self.assertEqual(image_processor.size    ,    {'''shortest_edge''': 4_2})
									self.assertEqual(image_processor.crop_size    ,    {'''height''': 8_4, '''width''': 8_4})




							def        UpperCAmelCase__      (	self					:	Dict):
									pass




							def        UpperCAmelCase__      (	self					:	str):
									# Initialize image_processing
									lowerCAmelCase_      :					int     =			self.image_processing_class(**self.image_processor_dict)
									# create random PIL videos
									lowerCAmelCase_      :					Optional[int]     =			prepare_video_inputs(self.image_processor_tester    ,    equal_resolution=A_)
									for video in video_inputs:
											self.assertIsInstance(A_    ,    A_)
											self.assertIsInstance(video[0]    ,    Image.Image)

									# Test not batched input
									lowerCAmelCase_      :					int     =			image_processing(video_inputs[0]    ,    return_tensors='''pt''').pixel_values
									self.assertEqual(
									    encoded_videos.shape    ,    (
									        1,
									        self.image_processor_tester.num_frames,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size['''height'''],
									        self.image_processor_tester.crop_size['''width'''],
									    )    ,    )

									# Test batched
									lowerCAmelCase_      :					Tuple     =			image_processing(A_    ,    return_tensors='''pt''').pixel_values
									self.assertEqual(
									    encoded_videos.shape    ,    (
									        self.image_processor_tester.batch_size,
									        self.image_processor_tester.num_frames,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size['''height'''],
									        self.image_processor_tester.crop_size['''width'''],
									    )    ,    )




							def        UpperCAmelCase__      (	self					:	int):
									# Initialize image_processing
									lowerCAmelCase_      :					Dict     =			self.image_processing_class(**self.image_processor_dict)
									# create random numpy tensors
									lowerCAmelCase_      :					str     =			prepare_video_inputs(self.image_processor_tester    ,    equal_resolution=A_    ,    numpify=A_)
									for video in video_inputs:
											self.assertIsInstance(A_    ,    A_)
											self.assertIsInstance(video[0]    ,    np.ndarray)

									# Test not batched input
									lowerCAmelCase_      :					Tuple     =			image_processing(video_inputs[0]    ,    return_tensors='''pt''').pixel_values
									self.assertEqual(
									    encoded_videos.shape    ,    (
									        1,
									        self.image_processor_tester.num_frames,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size['''height'''],
									        self.image_processor_tester.crop_size['''width'''],
									    )    ,    )

									# Test batched
									lowerCAmelCase_      :					Union[str, Any]     =			image_processing(A_    ,    return_tensors='''pt''').pixel_values
									self.assertEqual(
									    encoded_videos.shape    ,    (
									        self.image_processor_tester.batch_size,
									        self.image_processor_tester.num_frames,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size['''height'''],
									        self.image_processor_tester.crop_size['''width'''],
									    )    ,    )







							def        UpperCAmelCase__      (	self					:	List[str]):
									# Initialize image_processing
									lowerCAmelCase_      :					Tuple     =			self.image_processing_class(**self.image_processor_dict)
									# create random PyTorch tensors
									lowerCAmelCase_      :					Tuple     =			prepare_video_inputs(self.image_processor_tester    ,    equal_resolution=A_    ,    torchify=A_)
									for video in video_inputs:
											self.assertIsInstance(A_    ,    A_)
											self.assertIsInstance(video[0]    ,    torch.Tensor)

									# Test not batched input
									lowerCAmelCase_      :					List[str]     =			image_processing(video_inputs[0]    ,    return_tensors='''pt''').pixel_values
									self.assertEqual(
									    encoded_videos.shape    ,    (
									        1,
									        self.image_processor_tester.num_frames,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size['''height'''],
									        self.image_processor_tester.crop_size['''width'''],
									    )    ,    )

									# Test batched
									lowerCAmelCase_      :					List[Any]     =			image_processing(A_    ,    return_tensors='''pt''').pixel_values
									self.assertEqual(
									    encoded_videos.shape    ,    (
									        self.image_processor_tester.batch_size,
									        self.image_processor_tester.num_frames,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size['''height'''],
									        self.image_processor_tester.crop_size['''width'''],
									    )    ,    )





