




import gc
import random
import unittest

import numpy as np
import torch

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetPipeline,
    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     __lowerCAmelCase (       __magic_name__			,				unittest.TestCase ):

			"""simple docstring"""







			snake_case_   		=       KandinskyVaaControlnetPipeline
			snake_case_   		=       ['''image_embeds''', '''negative_image_embeds''', '''hint''']
			snake_case_   		=       ['''image_embeds''', '''negative_image_embeds''', '''hint''']
			snake_case_   		=       [
			    '''generator''',
			    '''height''',
			    '''width''',
			    '''latents''',
			    '''guidance_scale''',
			    '''num_inference_steps''',
			    '''return_dict''',
			    '''guidance_scale''',
			    '''num_images_per_prompt''',
			    '''output_type''',
			    '''return_dict''',
			]
			snake_case_   		=       False


			@property
			def       lowercase_    (      self							)    ->      str:



							'''simple docstring'''
							return 32


			@property
			def       lowercase_    (      self							)    ->      Tuple:



							'''simple docstring'''
							return 32


			@property
			def       lowercase_    (      self							)    ->      Tuple:



							'''simple docstring'''
							return self.time_input_dim


			@property
			def       lowercase_    (      self							)    ->      str:



							'''simple docstring'''
							return self.time_input_dim * 4


			@property
			def       lowercase_    (      self							)    ->      List[Any]:



							'''simple docstring'''
							return 100


			@property
			def       lowercase_    (      self							)    ->      Any:



							'''simple docstring'''
							torch.manual_seed(0							)

							__lowerCamelCase		       =    {
							    'in_channels': 8,
							    # Out channels is double in channels because predicts mean and variance
							    'out_channels': 8,
							    'addition_embed_type': 'image_hint',
							    '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       lowercase_    (      self							)    ->      Union[str, Any]:



							'''simple docstring'''
							return {
							    "block_out_channels": [32, 32, 64, 64],
							    "down_block_types": [
							        "DownEncoderBlock2D",
							        "DownEncoderBlock2D",
							        "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", "UpDecoderBlock2D", "UpDecoderBlock2D"],
							    "vq_embed_dim": 4,
							}


			@property
			def       lowercase_    (      self							)    ->      List[str]:



							'''simple docstring'''
							torch.manual_seed(0							)
							__lowerCamelCase		       =    VQModel(**self.dummy_movq_kwargs							)
							return model


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



							'''simple docstring'''
							__lowerCamelCase		       =    self.dummy_unet
							__lowerCamelCase		       =    self.dummy_movq

							__lowerCamelCase		       =    DDIMScheduler(
							    num_train_timesteps=1_000			,						beta_schedule='linear'			,						beta_start=0.0_00_85			,						beta_end=0.0_12			,						clip_sample=lowerCamelCase__			,						set_alpha_to_one=lowerCamelCase__			,						steps_offset=1			,						prediction_type='epsilon'			,						thresholding=lowerCamelCase__			,						)

							__lowerCamelCase		       =    {
							    'unet': unet,
							    'scheduler': scheduler,
							    'movq': movq,
							}
							return components


			def       lowercase_    (      self			,						lowerCamelCase__			,						lowerCamelCase__=0							)    ->      Union[str, Any]:



							'''simple docstring'''
							__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 hint
							__lowerCamelCase		       =    floats_tensor((1, 3, 64, 64)			,						rng=random.Random(lowerCamelCase__							)							).to(lowerCamelCase__							)

							if str(lowerCamelCase__							).startswith('mps'							):
											__lowerCamelCase		       =    torch.manual_seed(lowerCamelCase__							)
							else:
											__lowerCamelCase		       =    torch.Generator(device=lowerCamelCase__							).manual_seed(lowerCamelCase__							)
							__lowerCamelCase		       =    {
							    'image_embeds': image_embeds,
							    'negative_image_embeds': negative_image_embeds,
							    'hint': hint,
							    'generator': generator,
							    'height': 64,
							    'width': 64,
							    'guidance_scale': 4.0,
							    'num_inference_steps': 2,
							    'output_type': 'np',
							}
							return inputs






			def       lowercase_    (      self							)    ->      str:



							'''simple docstring'''
							__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_95_98_26, 0.86_82_79, 0.7_55_80_92, 0.68_76_94_67, 0.85_80_58_04, 0.65_97_74_96, 0.44_88_53_02, 0.5_95_91_11, 0.4_25_15_95]							)

							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     __lowerCAmelCase (       unittest.TestCase ):

			"""simple docstring"""







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



							'''simple docstring'''
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()






			def       lowercase_    (      self							)    ->      Any:



							'''simple docstring'''
							__lowerCamelCase		       =    load_numpy(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/kandinskyv22/kandinskyv22_controlnet_robotcat_fp16.npy'							)

							__lowerCamelCase		       =    load_image(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/kandinskyv22/hint_image_cat.png'							)
							__lowerCamelCase		       =    torch.from_numpy(np.array(lowerCamelCase__							)							).float() / 2_55.0
							__lowerCamelCase		       =    hint.permute(2			,						0			,						1							).unsqueeze(0							)

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

							__lowerCamelCase		       =    KandinskyVaaControlnetPipeline.from_pretrained(
							    'kandinsky-community/kandinsky-2-2-controlnet-depth'			,						torch_dtype=torch.floataa							)
							__lowerCamelCase		       =    pipeline.to(lowerCamelCase__							)
							pipeline.set_progress_bar_config(disable=lowerCamelCase__							)

							__lowerCamelCase		       =    'A robot, 4k photo'

							__lowerCamelCase		       =    torch.Generator(device='cuda'							).manual_seed(0							)
							__lowerCamelCase       ,						__lowerCamelCase		       =    pipe_prior(
							    lowerCamelCase__			,						generator=lowerCamelCase__			,						num_inference_steps=5			,						negative_prompt=''			,						).to_tuple()

							__lowerCamelCase		       =    torch.Generator(device='cuda'							).manual_seed(0							)
							__lowerCamelCase		       =    pipeline(
							    image_embeds=lowerCamelCase__			,						negative_image_embeds=lowerCamelCase__			,						hint=lowerCamelCase__			,						generator=lowerCamelCase__			,						num_inference_steps=100			,						output_type='np'			,						)

							__lowerCamelCase		       =    output.images[0]

							assert image.shape == (512, 512, 3)

							assert_mean_pixel_difference(lowerCamelCase__			,						lowerCamelCase__							)








def 			lowerCamelCase_		(       UpperCamelCase__				:      int      )     ->       bool:





				"""simple docstring"""


				return number & 1 == 0


if __name__ == "__main__":
					import doctest

					doctest.testmod()


