


'''simple docstring'''






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__				(							a__				,				a__				,				a__				,				unittest.TestCase  ):





							'''simple docstring'''






							lowercase__			:   Optional[Any]							=      StableDiffusionImageVariationPipeline
							lowercase__			:   int							=      IMAGE_VARIATION_PARAMS
							lowercase__			:   Any							=      IMAGE_VARIATION_BATCH_PARAMS
							lowercase__			:   Union[str, Any]							=      frozenset([]  )
							# TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
							lowercase__			:   Any							=      frozenset([]  )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							str:
														torch.manual_seed(0   )
														lowerCAmelCase__          =	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					,    )
														lowerCAmelCase__          =	PNDMScheduler(skip_prk_steps=lowerCamelCase_   )
														torch.manual_seed(0   )
														lowerCAmelCase__          =	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   )
														lowerCAmelCase__          =	CLIPVisionConfig(
														    hidden_size=32					,    projection_dim=32					,    intermediate_size=37					,    layer_norm_eps=1e-05					,    num_attention_heads=4					,    num_hidden_layers=5					,    image_size=32					,    patch_size=4					,    )
														lowerCAmelCase__          =	CLIPVisionModelWithProjection(lowerCamelCase_   )
														lowerCAmelCase__          =	CLIPImageProcessor(crop_size=32					,    size=32   )

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







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







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Optional[Any]:
														lowerCAmelCase__          =	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
														lowerCAmelCase__          =	self.get_dummy_components()
														lowerCAmelCase__          =	StableDiffusionImageVariationPipeline(**lowerCamelCase_   )
														lowerCAmelCase__          =	sd_pipe.to(lowerCamelCase_   )
														sd_pipe.set_progress_bar_config(disable=lowerCamelCase_   )

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

														assert image.shape == (1, 64, 64, 3)
														lowerCAmelCase__          =	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   __SCREAMING_SNAKE_CASE							(		self   )       ->							Optional[int]:
														lowerCAmelCase__          =	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
														lowerCAmelCase__          =	self.get_dummy_components()
														lowerCAmelCase__          =	StableDiffusionImageVariationPipeline(**lowerCamelCase_   )
														lowerCAmelCase__          =	sd_pipe.to(lowerCamelCase_   )
														sd_pipe.set_progress_bar_config(disable=lowerCamelCase_   )

														lowerCAmelCase__          =	self.get_dummy_inputs(lowerCamelCase_   )
														lowerCAmelCase__          =	2 * [inputs['''image''']]
														lowerCAmelCase__          =	sd_pipe(**lowerCamelCase_   )

														lowerCAmelCase__          =	output.images

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

														assert image.shape == (2, 64, 64, 3)
														lowerCAmelCase__          =	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   __SCREAMING_SNAKE_CASE							(		self   )       ->							List[str]:
														super().test_inference_batch_single_identical(expected_max_diff=3e-3   )




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





							'''simple docstring'''






							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Dict:
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_="cpu"					,    lowerCamelCase_=torch.floataa					,    lowerCamelCase_=0   )       ->							List[Any]:
														lowerCAmelCase__          =	torch.Generator(device=lowerCamelCase_   ).manual_seed(lowerCamelCase_   )
														lowerCAmelCase__          =	load_image(
														    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'''
														    '''/stable_diffusion_imgvar/input_image_vermeer.png'''   )
														lowerCAmelCase__          =	np.random.RandomState(lowerCamelCase_   ).standard_normal((1, 4, 64, 64)   )
														lowerCAmelCase__          =	torch.from_numpy(lowerCamelCase_   ).to(device=lowerCamelCase_					,    dtype=lowerCamelCase_   )
														lowerCAmelCase__          =	{
														    '''image''': init_image,
														    '''latents''': latents,
														    '''generator''': generator,
														    '''num_inference_steps''': 3,
														    '''guidance_scale''': 7.5,
														    '''output_type''': '''numpy''',
														}
														return inputs







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							str:
														lowerCAmelCase__          =	StableDiffusionImageVariationPipeline.from_pretrained(
														    '''lambdalabs/sd-image-variations-diffusers'''					,    safety_checker=lowerCamelCase_   )
														lowerCAmelCase__          =	sd_pipe.to(lowerCamelCase_   )
														sd_pipe.set_progress_bar_config(disable=lowerCamelCase_   )

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

														assert image.shape == (1, 5_12, 5_12, 3)
														lowerCAmelCase__          =	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   __SCREAMING_SNAKE_CASE							(		self   )       ->							str:
														lowerCAmelCase__          =	0

														def callback_fn(lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_   ) -> None:
																					lowerCAmelCase__          =	True
																					nonlocal number_of_steps
																					number_of_steps += 1
																					if step == 1:
																												lowerCAmelCase__          =	latents.detach().cpu().numpy()
																												assert latents.shape == (1, 4, 64, 64)
																												lowerCAmelCase__          =	latents[0, -3:, -3:, -1]
																												lowerCAmelCase__          =	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:
																												lowerCAmelCase__          =	latents.detach().cpu().numpy()
																												assert latents.shape == (1, 4, 64, 64)
																												lowerCAmelCase__          =	latents[0, -3:, -3:, -1]
																												lowerCAmelCase__          =	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

														lowerCAmelCase__          =	False

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

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







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Optional[int]:
														torch.cuda.empty_cache()
														torch.cuda.reset_max_memory_allocated()
														torch.cuda.reset_peak_memory_stats()

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

														lowerCAmelCase__          =	self.get_inputs(lowerCamelCase_					,    dtype=torch.floataa   )
														lowerCAmelCase__          =	pipe(**lowerCamelCase_   )

														lowerCAmelCase__          =	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  ):





							'''simple docstring'''






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







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_="cpu"					,    lowerCamelCase_=torch.floataa					,    lowerCamelCase_=0   )       ->							Dict:
														lowerCAmelCase__          =	torch.Generator(device=lowerCamelCase_   ).manual_seed(lowerCamelCase_   )
														lowerCAmelCase__          =	load_image(
														    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'''
														    '''/stable_diffusion_imgvar/input_image_vermeer.png'''   )
														lowerCAmelCase__          =	np.random.RandomState(lowerCamelCase_   ).standard_normal((1, 4, 64, 64)   )
														lowerCAmelCase__          =	torch.from_numpy(lowerCamelCase_   ).to(device=lowerCamelCase_					,    dtype=lowerCamelCase_   )
														lowerCAmelCase__          =	{
														    '''image''': init_image,
														    '''latents''': latents,
														    '''generator''': generator,
														    '''num_inference_steps''': 50,
														    '''guidance_scale''': 7.5,
														    '''output_type''': '''numpy''',
														}
														return inputs







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							int:
														lowerCAmelCase__          =	StableDiffusionImageVariationPipeline.from_pretrained('''fusing/sd-image-variations-diffusers'''   )
														sd_pipe.to(lowerCamelCase_   )
														sd_pipe.set_progress_bar_config(disable=lowerCamelCase_   )

														lowerCAmelCase__          =	self.get_inputs(lowerCamelCase_   )
														lowerCAmelCase__          =	sd_pipe(**lowerCamelCase_   ).images[0]

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







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							List[str]:
														lowerCAmelCase__          =	StableDiffusionImageVariationPipeline.from_pretrained('''fusing/sd-image-variations-diffusers'''   )
														lowerCAmelCase__          =	DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config   )
														sd_pipe.to(lowerCamelCase_   )
														sd_pipe.set_progress_bar_config(disable=lowerCamelCase_   )

														lowerCAmelCase__          =	self.get_inputs(lowerCamelCase_   )
														lowerCAmelCase__          =	25
														lowerCAmelCase__          =	sd_pipe(**lowerCamelCase_   ).images[0]

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



'''simple docstring'''






from .hash_table import HashTable
from .number_theory.prime_numbers import is_prime, next_prime




class 		a__				(							a__  ):





							'''simple docstring'''






							def __init__(		self					,    *lowerCamelCase_					,    **lowerCamelCase_   )       ->							Optional[int]:
														super().__init__(*lowerCamelCase_					,    **lowerCamelCase_   )







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_   )       ->							Optional[int]:
														lowerCAmelCase__          =	(
														    next_prime(value % self.size_table   )
														    if not is_prime(value % self.size_table   )
														    else value % self.size_table
														)  # gt = bigger than
														return next_prime_gt - (data % next_prime_gt)







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_   )       ->							Optional[int]:
														return (increment * self.__hash_function_a(lowerCamelCase_					,    lowerCamelCase_   )) % self.size_table







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_=None   )       ->							Any:
														lowerCAmelCase__          =	1
														lowerCAmelCase__          =	self.hash_function(lowerCamelCase_   )

														while self.values[new_key] is not None and self.values[new_key] != key:
																					lowerCAmelCase__          =	(
																					    self.__hash_double_function(lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_   )
																					    if self.balanced_factor() >= self.lim_charge
																					    else None
																					)
																					if new_key is None:
																												break
																					else:
																												i += 1

														return new_key