


'''simple docstring'''






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()




class 						a__							(				a__							,							unittest.TestCase    ):





							'''simple docstring'''






							lowercase__     :  List[str]      		= KandinskyVaaImgaImgPipeline
							lowercase__     :  Any      		= ["image_embeds", "negative_image_embeds", "image"]
							lowercase__     :  Union[str, Any]      		= [
							    "image_embeds",
							    "negative_image_embeds",
							    "image",
							]
							lowercase__     :  List[str]      		= [
							    "generator",
							    "height",
							    "width",
							    "strength",
							    "guidance_scale",
							    "num_inference_steps",
							    "return_dict",
							    "guidance_scale",
							    "num_images_per_prompt",
							    "output_type",
							    "return_dict",
							]
							lowercase__     :  Union[str, Any]      		= False







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														return 32







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														return 32







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														return self.time_input_dim







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														return self.time_input_dim * 4







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														return 1_00







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Dict:
														torch.manual_seed(0  )

														lowerCAmelCase__			  =      {
														    '''in_channels''': 4,
														    # Out channels is double in channels because predicts mean and variance
														    '''out_channels''': 8,
														    '''addition_embed_type''': '''image''',
														    '''down_block_types''': ('''ResnetDownsampleBlock2D''', '''SimpleCrossAttnDownBlock2D'''),
														    '''up_block_types''': ('''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''),
														    '''mid_block_type''': '''UNetMidBlock2DSimpleCrossAttn''',
														    '''block_out_channels''': (self.block_out_channels_a, self.block_out_channels_a * 2),
														    '''layers_per_block''': 1,
														    '''encoder_hid_dim''': self.text_embedder_hidden_size,
														    '''encoder_hid_dim_type''': '''image_proj''',
														    '''cross_attention_dim''': self.cross_attention_dim,
														    '''attention_head_dim''': 4,
														    '''resnet_time_scale_shift''': '''scale_shift''',
														    '''class_embed_type''': None,
														}

														lowerCAmelCase__			  =      UNetaDConditionModel(**lowerCamelCase_  )
														return model







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														return {
														    "block_out_channels": [32, 64],
														    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
														    "in_channels": 3,
														    "latent_channels": 4,
														    "layers_per_block": 1,
														    "norm_num_groups": 8,
														    "norm_type": "spatial",
														    "num_vq_embeddings": 12,
														    "out_channels": 3,
														    "up_block_types": [
														        "AttnUpDecoderBlock2D",
														        "UpDecoderBlock2D",
														    ],
														    "vq_embed_dim": 4,
														}







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      VQModel(**self.dummy_movq_kwargs  )
														return model







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														lowerCAmelCase__			  =      self.dummy_unet
														lowerCAmelCase__			  =      self.dummy_movq

														lowerCAmelCase__			  =      {
														    '''num_train_timesteps''': 10_00,
														    '''beta_schedule''': '''linear''',
														    '''beta_start''': 0.00_085,
														    '''beta_end''': 0.012,
														    '''clip_sample''': False,
														    '''set_alpha_to_one''': False,
														    '''steps_offset''': 0,
														    '''prediction_type''': '''epsilon''',
														    '''thresholding''': False,
														}

														lowerCAmelCase__			  =      DDIMScheduler(**lowerCamelCase_  )

														lowerCAmelCase__			  =      {
														    '''unet''': unet,
														    '''scheduler''': scheduler,
														    '''movq''': movq,
														}

														return components







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_=0  )			->				Dict:
														lowerCAmelCase__			  =      floats_tensor((1, self.text_embedder_hidden_size)   ,       rng=random.Random(lowerCamelCase_  )  ).to(lowerCamelCase_  )
														lowerCAmelCase__			  =      floats_tensor((1, self.text_embedder_hidden_size)   ,       rng=random.Random(seed + 1  )  ).to(
														    lowerCamelCase_  )
														# create init_image
														lowerCAmelCase__			  =      floats_tensor((1, 3, 64, 64)   ,       rng=random.Random(lowerCamelCase_  )  ).to(lowerCamelCase_  )
														lowerCAmelCase__			  =      image.cpu().permute(0   ,       2   ,       3   ,       1  )[0]
														lowerCAmelCase__			  =      Image.fromarray(np.uinta(lowerCamelCase_  )  ).convert('''RGB'''  ).resize((2_56, 2_56)  )

														if str(lowerCamelCase_  ).startswith('''mps'''  ):
																					lowerCAmelCase__			  =      torch.manual_seed(lowerCamelCase_  )
														else:
																					lowerCAmelCase__			  =      torch.Generator(device=lowerCamelCase_  ).manual_seed(lowerCamelCase_  )
														lowerCAmelCase__			  =      {
														    '''image''': init_image,
														    '''image_embeds''': image_embeds,
														    '''negative_image_embeds''': negative_image_embeds,
														    '''generator''': generator,
														    '''height''': 64,
														    '''width''': 64,
														    '''num_inference_steps''': 10,
														    '''guidance_scale''': 7.0,
														    '''strength''': 0.2,
														    '''output_type''': '''np''',
														}
														return inputs







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														lowerCAmelCase__			  =      '''cpu'''

														lowerCAmelCase__			  =      self.get_dummy_components()

														lowerCAmelCase__			  =      self.pipeline_class(**lowerCamelCase_  )
														lowerCAmelCase__			  =      pipe.to(lowerCamelCase_  )

														pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      pipe(**self.get_dummy_inputs(lowerCamelCase_  )  )
														lowerCAmelCase__			  =      output.images

														lowerCAmelCase__			  =      pipe(
														    **self.get_dummy_inputs(lowerCamelCase_  )   ,       return_dict=lowerCamelCase_   ,       )[0]

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

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

														lowerCAmelCase__			  =      np.array(
														    [0.6_199_778, 0.63_984_406, 0.46_145_785, 0.62_944_984, 0.5_622_215, 0.47_306_132, 0.47_441_456, 0.4_607_606, 0.48_719_263]  )
														assert (
														    np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2
														), F""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
														assert (
														    np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 1e-2
														), F""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""




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





							'''simple docstring'''






							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														# clean up the VRAM after each test
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														lowerCAmelCase__			  =      load_numpy(
														    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
														    '''/kandinskyv22/kandinskyv22_img2img_frog.npy'''  )

														lowerCAmelCase__			  =      load_image(
														    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''  )
														lowerCAmelCase__			  =      '''A red cartoon frog, 4k'''

														lowerCAmelCase__			  =      KandinskyVaaPriorPipeline.from_pretrained(
														    '''kandinsky-community/kandinsky-2-2-prior'''   ,       torch_dtype=torch.floataa  )
														pipe_prior.to(lowerCamelCase_  )

														lowerCAmelCase__			  =      KandinskyVaaImgaImgPipeline.from_pretrained(
														    '''kandinsky-community/kandinsky-2-2-decoder'''   ,       torch_dtype=torch.floataa  )
														lowerCAmelCase__			  =      pipeline.to(lowerCamelCase_  )

														pipeline.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      torch.Generator(device='''cpu'''  ).manual_seed(0  )
														lowerCAmelCase__  ,   lowerCAmelCase__			  =      pipe_prior(
														    lowerCamelCase_   ,       generator=lowerCamelCase_   ,       num_inference_steps=5   ,       negative_prompt=''''''   ,       ).to_tuple()

														lowerCAmelCase__			  =      pipeline(
														    image=lowerCamelCase_   ,       image_embeds=lowerCamelCase_   ,       negative_image_embeds=lowerCamelCase_   ,       generator=lowerCamelCase_   ,       num_inference_steps=1_00   ,       height=7_68   ,       width=7_68   ,       strength=0.2   ,       output_type='''np'''   ,       )

														lowerCAmelCase__			  =      output.images[0]

														assert image.shape == (7_68, 7_68, 3)

														assert_mean_pixel_difference(lowerCamelCase_   ,       lowerCamelCase_  )

# Imports
import numpy as np





class   UpperCamelCase_      :







						'''simple docstring'''



						def __init__(			self							:Optional[int]		,    lowerCAmelCase__							:int=None		,    lowerCAmelCase__							:List[str]=None		,    lowerCAmelCase__							:int=None		,    lowerCAmelCase__							:Optional[int]=None		,    lowerCAmelCase__							:List[str]=None     )  ->int:
										self.set_matricies(red=lowerCAmelCase__		,    green=lowerCAmelCase__		,    blue=lowerCAmelCase__		,    red_edge=lowerCAmelCase__		,    nir=lowerCAmelCase__     )





						def        SCREAMING_SNAKE_CASE(			self							:Optional[Any]		,    lowerCAmelCase__							:List[Any]=None		,    lowerCAmelCase__							:Any=None		,    lowerCAmelCase__							:Tuple=None		,    lowerCAmelCase__							:Optional[Any]=None		,    lowerCAmelCase__							:Optional[Any]=None     )  ->List[Any]:
										if red is not None:
														lowercase         =      red
										if green is not None:
														lowercase         =      green
										if blue is not None:
														lowercase         =      blue
										if red_edge is not None:
														lowercase         =      red_edge
										if nir is not None:
														lowercase         =      nir
										return True





						def        SCREAMING_SNAKE_CASE(			self							:Optional[Any]		,    lowerCAmelCase__							:Dict=""		,    lowerCAmelCase__							:Optional[Any]=None		,    lowerCAmelCase__							:Union[str, Any]=None		,    lowerCAmelCase__							:int=None		,    lowerCAmelCase__							:int=None		,    lowerCAmelCase__							:Union[str, Any]=None     )  ->Optional[Any]:
										self.set_matricies(red=lowerCAmelCase__		,    green=lowerCAmelCase__		,    blue=lowerCAmelCase__		,    red_edge=lowerCAmelCase__		,    nir=lowerCAmelCase__     )
										lowercase         =      {
										    "ARVI2": self.arvaa,
										    "CCCI": self.ccci,
										    "CVI": self.cvi,
										    "GLI": self.gli,
										    "NDVI": self.ndvi,
										    "BNDVI": self.bndvi,
										    "redEdgeNDVI": self.red_edge_ndvi,
										    "GNDVI": self.gndvi,
										    "GBNDVI": self.gbndvi,
										    "GRNDVI": self.grndvi,
										    "RBNDVI": self.rbndvi,
										    "PNDVI": self.pndvi,
										    "ATSAVI": self.atsavi,
										    "BWDRVI": self.bwdrvi,
										    "CIgreen": self.ci_green,
										    "CIrededge": self.ci_rededge,
										    "CI": self.ci,
										    "CTVI": self.ctvi,
										    "GDVI": self.gdvi,
										    "EVI": self.evi,
										    "GEMI": self.gemi,
										    "GOSAVI": self.gosavi,
										    "GSAVI": self.gsavi,
										    "Hue": self.hue,
										    "IVI": self.ivi,
										    "IPVI": self.ipvi,
										    "I": self.i,
										    "RVI": self.rvi,
										    "MRVI": self.mrvi,
										    "MSAVI": self.m_savi,
										    "NormG": self.norm_g,
										    "NormNIR": self.norm_nir,
										    "NormR": self.norm_r,
										    "NGRDI": self.ngrdi,
										    "RI": self.ri,
										    "S": self.s,
										    "IF": self._if,
										    "DVI": self.dvi,
										    "TVI": self.tvi,
										    "NDRE": self.ndre,
										}

										try:
														return funcs[index]()
										except KeyError:
														print("Index not in the list!"     )
														return False





						def        SCREAMING_SNAKE_CASE(			self							:Any     )  ->int:
										return -0.18 + (1.17 * ((self.nir - self.red) / (self.nir + self.red)))





						def        SCREAMING_SNAKE_CASE(			self							:List[str]     )  ->int:
										return ((self.nir - self.redEdge) / (self.nir + self.redEdge)) / (
										    (self.nir - self.red) / (self.nir + self.red)
										)





						def        SCREAMING_SNAKE_CASE(			self							:int     )  ->Tuple:
										return self.nir * (self.red / (self.green**2))





						def        SCREAMING_SNAKE_CASE(			self							:List[str]     )  ->Dict:
										return (2 * self.green - self.red - self.blue) / (
										    2 * self.green + self.red + self.blue
										)





						def        SCREAMING_SNAKE_CASE(			self							:str     )  ->Optional[int]:
										return (self.nir - self.red) / (self.nir + self.red)





						def        SCREAMING_SNAKE_CASE(			self							:str     )  ->int:
										return (self.nir - self.blue) / (self.nir + self.blue)





						def        SCREAMING_SNAKE_CASE(			self							:List[str]     )  ->List[str]:
										return (self.redEdge - self.red) / (self.redEdge + self.red)





						def        SCREAMING_SNAKE_CASE(			self							:Optional[int]     )  ->str:
										return (self.nir - self.green) / (self.nir + self.green)





						def        SCREAMING_SNAKE_CASE(			self							:int     )  ->Any:
										return (self.nir - (self.green + self.blue)) / (
										    self.nir + (self.green + self.blue)
										)





						def        SCREAMING_SNAKE_CASE(			self							:Union[str, Any]     )  ->List[Any]:
										return (self.nir - (self.green + self.red)) / (
										    self.nir + (self.green + self.red)
										)





						def        SCREAMING_SNAKE_CASE(			self							:Optional[Any]     )  ->Union[str, Any]:
										return (self.nir - (self.blue + self.red)) / (self.nir + (self.blue + self.red))





						def        SCREAMING_SNAKE_CASE(			self							:Any     )  ->Union[str, Any]:
										return (self.nir - (self.green + self.red + self.blue)) / (
										    self.nir + (self.green + self.red + self.blue)
										)





						def        SCREAMING_SNAKE_CASE(			self							:Tuple		,    lowerCAmelCase__							:Union[str, Any]=0.08		,    lowerCAmelCase__							:Dict=1.22		,    lowerCAmelCase__							:Tuple=0.03     )  ->Tuple:
										return a * (
										    (self.nir - a * self.red - b)
										    / (a * self.nir + self.red - a * b + x * (1 + a**2))
										)





						def        SCREAMING_SNAKE_CASE(			self							:Dict     )  ->Any:
										return (0.1 * self.nir - self.blue) / (0.1 * self.nir + self.blue)





						def        SCREAMING_SNAKE_CASE(			self							:Optional[Any]     )  ->int:
										return (self.nir / self.green) - 1





						def        SCREAMING_SNAKE_CASE(			self							:int     )  ->List[str]:
										return (self.nir / self.redEdge) - 1





						def        SCREAMING_SNAKE_CASE(			self							:Union[str, Any]     )  ->List[Any]:
										return (self.red - self.blue) / self.red





						def        SCREAMING_SNAKE_CASE(			self							:Any     )  ->int:
										lowercase         =      self.ndvi()
										return ((ndvi + 0.5) / (abs(ndvi + 0.5     ))) * (abs(ndvi + 0.5     ) ** (1 / 2))





						def        SCREAMING_SNAKE_CASE(			self							:List[str]     )  ->Optional[Any]:
										return self.nir - self.green





						def        SCREAMING_SNAKE_CASE(			self							:str     )  ->int:
										return 2.5 * (
										    (self.nir - self.red) / (self.nir + 6 * self.red - 7.5 * self.blue + 1)
										)





						def        SCREAMING_SNAKE_CASE(			self							:List[str]     )  ->Optional[Any]:
										lowercase         =      (2 * (self.nir**2 - self.red**2) + 1.5 * self.nir + 0.5 * self.red) / (
										    self.nir + self.red + 0.5
										)
										return n * (1 - 0.25 * n) - (self.red - 0.1_25) / (1 - self.red)





						def        SCREAMING_SNAKE_CASE(			self							:Union[str, Any]		,    lowerCAmelCase__							:Tuple=0.16     )  ->Any:
										return (self.nir - self.green) / (self.nir + self.green + y)





						def        SCREAMING_SNAKE_CASE(			self							:Dict		,    lowerCAmelCase__							:Any=0.5     )  ->str:
										return ((self.nir - self.green) / (self.nir + self.green + n)) * (1 + n)





						def        SCREAMING_SNAKE_CASE(			self							:int     )  ->List[str]:
										return np.arctan(
										    ((2 * self.red - self.green - self.blue) / 30.5) * (self.green - self.blue)     )





						def        SCREAMING_SNAKE_CASE(			self							:Any		,    lowerCAmelCase__							:str=None		,    lowerCAmelCase__							:Tuple=None     )  ->Dict:
										return (self.nir - b) / (a * self.red)





						def        SCREAMING_SNAKE_CASE(			self							:Optional[Any]     )  ->Dict:
										return (self.nir / ((self.nir + self.red) / 2)) * (self.ndvi() + 1)





						def        SCREAMING_SNAKE_CASE(			self							:str     )  ->List[Any]:
										return (self.red + self.green + self.blue) / 30.5





						def        SCREAMING_SNAKE_CASE(			self							:Tuple     )  ->Optional[Any]:
										return self.nir / self.red





						def        SCREAMING_SNAKE_CASE(			self							:Optional[int]     )  ->Dict:
										return (self.rvi() - 1) / (self.rvi() + 1)





						def        SCREAMING_SNAKE_CASE(			self							:Optional[Any]     )  ->Dict:
										return (
										    (2 * self.nir + 1)
										    - ((2 * self.nir + 1) ** 2 - 8 * (self.nir - self.red)) ** (1 / 2)
										) / 2





						def        SCREAMING_SNAKE_CASE(			self							:str     )  ->int:
										return self.green / (self.nir + self.red + self.green)





						def        SCREAMING_SNAKE_CASE(			self							:Any     )  ->Any:
										return self.nir / (self.nir + self.red + self.green)





						def        SCREAMING_SNAKE_CASE(			self							:int     )  ->Tuple:
										return self.red / (self.nir + self.red + self.green)





						def        SCREAMING_SNAKE_CASE(			self							:List[Any]     )  ->List[Any]:
										return (self.green - self.red) / (self.green + self.red)





						def        SCREAMING_SNAKE_CASE(			self							:List[str]     )  ->str:
										return (self.red - self.green) / (self.red + self.green)





						def        SCREAMING_SNAKE_CASE(			self							:Union[str, Any]     )  ->List[str]:
										lowercase         =      np.max([np.max(self.red     ), np.max(self.green     ), np.max(self.blue     )]     )
										lowercase         =      np.min([np.min(self.red     ), np.min(self.green     ), np.min(self.blue     )]     )
										return (max_value - min_value) / max_value





						def        SCREAMING_SNAKE_CASE(			self							:List[str]     )  ->Any:
										return (2 * self.red - self.green - self.blue) / (self.green - self.blue)





						def        SCREAMING_SNAKE_CASE(			self							:List[str]     )  ->List[Any]:
										return self.nir / self.red





						def        SCREAMING_SNAKE_CASE(			self							:Any     )  ->List[str]:
										return (self.ndvi() + 0.5) ** (1 / 2)





						def        SCREAMING_SNAKE_CASE(			self							:Optional[int]     )  ->Union[str, Any]:
										return (self.nir - self.redEdge) / (self.nir + self.redEdge)
