
import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetImgaImgPipeline,
    KandinskyVaaPriorEmbaEmbPipeline,
    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    SCREAMING_SNAKE_CASE__     (	UpperCAmelCase							,unittest.TestCase					):





	'''simple docstring'''



	_UpperCAmelCase       :		int       =					KandinskyVaaControlnetImgaImgPipeline
	_UpperCAmelCase       :		List[Any]       =					["image_embeds", "negative_image_embeds", "image", "hint"]
	_UpperCAmelCase       :		Any       =					["image_embeds", "negative_image_embeds", "image", "hint"]
	_UpperCAmelCase       :		Dict       =					[
	    "generator",
	    "height",
	    "width",
	    "strength",
	    "guidance_scale",
	    "num_inference_steps",
	    "return_dict",
	    "guidance_scale",
	    "num_images_per_prompt",
	    "output_type",
	    "return_dict",
	]
	_UpperCAmelCase       :		int       =					False



	@property
	def  A							(  self	:    int	):


				'''simple docstring'''


				return 32



	@property
	def  A							(  self	:    int	):


				'''simple docstring'''


				return 32



	@property
	def  A							(  self	:    Optional[int]	):


				'''simple docstring'''


				return self.time_input_dim



	@property
	def  A							(  self	:    Dict	):


				'''simple docstring'''


				return self.time_input_dim * 4



	@property
	def  A							(  self	:    int	):


				'''simple docstring'''


				return 100



	@property
	def  A							(  self	:    Optional[int]	):


				'''simple docstring'''


				torch.manual_seed(0	)

				_snake_case		=			{
				    '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,
				}

				_snake_case		=			UNetaDConditionModel(**lowercase	)
				return model



	@property
	def  A							(  self	:    Optional[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  A							(  self	:    Tuple	):


				'''simple docstring'''


				torch.manual_seed(0	)
				_snake_case		=			VQModel(**self.dummy_movq_kwargs	)
				return model



	def  A							(  self	:    List[Any]	):


				'''simple docstring'''


				_snake_case		=			self.dummy_unet
				_snake_case		=			self.dummy_movq

				_snake_case		=			{
				    'num_train_timesteps': 1_000,
				    'beta_schedule': 'linear',
				    'beta_start': 0.00085,
				    'beta_end': 0.012,
				    'clip_sample': False,
				    'set_alpha_to_one': False,
				    'steps_offset': 0,
				    'prediction_type': 'epsilon',
				    'thresholding': False,
				}

				_snake_case		=			DDIMScheduler(**lowercase	)

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

				return components



	def  A							(  self	:    Optional[int]			,  lowercase	:    Union[str, Any]			,  lowercase	:    List[Any]=0	):


				'''simple docstring'''


				_snake_case		=			floats_tensor((1, self.text_embedder_hidden_size)			,  rng=random.Random(lowercase	)	).to(lowercase	)
				_snake_case		=			floats_tensor((1, self.text_embedder_hidden_size)			,  rng=random.Random(seed + 1	)	).to(
				    lowercase	)
				# create init_image
				_snake_case		=			floats_tensor((1, 3, 64, 64)			,  rng=random.Random(lowercase	)	).to(lowercase	)
				_snake_case		=			image.cpu().permute(0			,  2			,  3			,  1	)[0]
				_snake_case		=			Image.fromarray(np.uinta(lowercase	)	).convert('RGB'	).resize((256, 256)	)
				# create hint
				_snake_case		=			floats_tensor((1, 3, 64, 64)			,  rng=random.Random(lowercase	)	).to(lowercase	)

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





	def  A							(  self	:    Any	):


				'''simple docstring'''


				_snake_case		=			'cpu'

				_snake_case		=			self.get_dummy_components()

				_snake_case		=			self.pipeline_class(**lowercase	)
				_snake_case		=			pipe.to(lowercase	)

				pipe.set_progress_bar_config(disable=lowercase	)

				_snake_case		=			pipe(**self.get_dummy_inputs(lowercase	)	)
				_snake_case		=			output.images

				_snake_case		=			pipe(
				    **self.get_dummy_inputs(lowercase	)			,  return_dict=lowercase			,  )[0]

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

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

				_snake_case		=			np.array(
				    [0.54985034, 0.55509365, 0.52561504, 0.5570494, 0.5593818, 0.5263979, 0.50285643, 0.5069846, 0.51196736]	)
				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    SCREAMING_SNAKE_CASE__     (	unittest.TestCase					):





	'''simple docstring'''



	def  A							(  self	:    int	):


				'''simple docstring'''


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





	def  A							(  self	:    Optional[Any]	):


				'''simple docstring'''


				_snake_case		=			load_numpy(
				    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
				    '/kandinskyv22/kandinskyv22_controlnet_img2img_robotcat_fp16.npy'	)

				_snake_case		=			load_image(
				    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/kandinsky/cat.png'	)
				_snake_case		=			init_image.resize((512, 512)	)

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

				_snake_case		=			'A robot, 4k photo'

				_snake_case		=			KandinskyVaaPriorEmbaEmbPipeline.from_pretrained(
				    'kandinsky-community/kandinsky-2-2-prior'			,  torch_dtype=torch.floataa	)
				pipe_prior.to(lowercase	)

				_snake_case		=			KandinskyVaaControlnetImgaImgPipeline.from_pretrained(
				    'kandinsky-community/kandinsky-2-2-controlnet-depth'			,  torch_dtype=torch.floataa	)
				_snake_case		=			pipeline.to(lowercase	)

				pipeline.set_progress_bar_config(disable=lowercase	)

				_snake_case		=			torch.Generator(device='cpu'	).manual_seed(0	)

				_snake_case	,					_snake_case		=			pipe_prior(
				    lowercase			,  image=lowercase			,  strength=0.85			,  generator=lowercase			,  negative_prompt=''			,  ).to_tuple()

				_snake_case		=			pipeline(
				    image=lowercase			,  image_embeds=lowercase			,  negative_image_embeds=lowercase			,  hint=lowercase			,  generator=lowercase			,  num_inference_steps=100			,  height=512			,  width=512			,  strength=0.5			,  output_type='np'			,  )

				_snake_case		=			output.images[0]

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

				assert_mean_pixel_difference(lowercase			,  lowercase	)

import unittest
from typing import Tuple

import torch

from diffusers.utils import floats_tensor, randn_tensor, torch_all_close, torch_device
from diffusers.utils.testing_utils import require_torch




@require_torch
class    SCREAMING_SNAKE_CASE__     :





	'''simple docstring'''



	@property
	def  A							(  self	:    List[str]	):


				'''simple docstring'''


				return self.get_dummy_input()



	@property
	def  A							(  self	:    Any	):


				'''simple docstring'''


				if self.block_type == "down":
							return (4, 32, 16, 16)
				elif self.block_type == "mid":
							return (4, 32, 32, 32)
				elif self.block_type == "up":
							return (4, 32, 64, 64)

				raise ValueError(f'''\'{self.block_type}\' is not a supported block_type. Set it to \'up\', \'mid\', or \'down\'.'''	)



	def  A							(  self	:    Union[str, Any]			,  lowercase	:    Any=True			,  lowercase	:    List[Any]=False			,  lowercase	:    List[str]=False			,  lowercase	:    Dict=False			,  ):


				'''simple docstring'''


				_snake_case		=			4
				_snake_case		=			32
				_snake_case		=			(32, 32)

				_snake_case		=			torch.manual_seed(0	)
				_snake_case		=			torch.device(lowercase	)
				_snake_case		=			(batch_size, num_channels) + sizes
				_snake_case		=			randn_tensor(lowercase			,  generator=lowercase			,  device=lowercase	)
				_snake_case		=			{'hidden_states': hidden_states}

				if include_temb:
							_snake_case		=			128
							_snake_case		=			randn_tensor((batch_size, temb_channels)			,  generator=lowercase			,  device=lowercase	)

				if include_res_hidden_states_tuple:
							_snake_case		=			torch.manual_seed(1	)
							_snake_case		=			(randn_tensor(lowercase			,  generator=lowercase			,  device=lowercase	),)

				if include_encoder_hidden_states:
							_snake_case		=			floats_tensor((batch_size, 32, 32)	).to(lowercase	)

				if include_skip_sample:
							_snake_case		=			randn_tensor(((batch_size, 3) + sizes)			,  generator=lowercase			,  device=lowercase	)

				return dummy_input



	def  A							(  self	:    Any	):


				'''simple docstring'''


				_snake_case		=			{
				    'in_channels': 32,
				    'out_channels': 32,
				    'temb_channels': 128,
				}
				if self.block_type == "up":
							_snake_case		=			32

				if self.block_type == "mid":
							init_dict.pop('out_channels'	)

				_snake_case		=			self.dummy_input
				return init_dict, inputs_dict



	def  A							(  self	:    Dict			,  lowercase	:    Optional[int]	):


				'''simple docstring'''


				_snake_case	,					_snake_case		=			self.prepare_init_args_and_inputs_for_common()
				_snake_case		=			self.block_class(**lowercase	)
				unet_block.to(lowercase	)
				unet_block.eval()

				with torch.no_grad():
							_snake_case		=			unet_block(**lowercase	)

				if isinstance(lowercase			,  lowercase	):
							_snake_case		=			output[0]

				self.assertEqual(output.shape			,  self.output_shape	)

				_snake_case		=			output[0, -1, -3:, -3:]
				_snake_case		=			torch.tensor(lowercase	).to(lowercase	)
				assert torch_all_close(output_slice.flatten()			,  lowercase			,  atol=5E-3	)





	@unittest.skipIf(torch_device == 'mps'			,  'Training is not supported in mps'	)
	def  A							(  self	:    Dict	):


				'''simple docstring'''


				_snake_case	,					_snake_case		=			self.prepare_init_args_and_inputs_for_common()
				_snake_case		=			self.block_class(**lowercase	)
				model.to(lowercase	)
				model.train()
				_snake_case		=			model(**lowercase	)

				if isinstance(lowercase			,  lowercase	):
							_snake_case		=			output[0]

				_snake_case		=			torch.device(lowercase	)
				_snake_case		=			randn_tensor(output.shape			,  device=lowercase	)
				_snake_case		=			torch.nn.functional.mse_loss(lowercase			,  lowercase	)
				loss.backward()