




"""simple docstring"""







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       lowerCAmelCase__       (			lowerCAmelCase_  ,			unittest.TestCase    ):


				"""simple docstring"""





				__UpperCAmelCase :     str       =      KandinskyVaaControlnetImgaImgPipeline
				__UpperCAmelCase :     Tuple       =      ["image_embeds", "negative_image_embeds", "image", "hint"]
				__UpperCAmelCase :     Tuple       =      ["image_embeds", "negative_image_embeds", "image", "hint"]
				__UpperCAmelCase :     List[Any]       =      [
				    "generator",
				    "height",
				    "width",
				    "strength",
				    "guidance_scale",
				    "num_inference_steps",
				    "return_dict",
				    "guidance_scale",
				    "num_images_per_prompt",
				    "output_type",
				    "return_dict",
				]
				__UpperCAmelCase :     Optional[Any]       =      False



				@property
				def       snake_case    (     self					:	Any						):
						return 3_2



				@property
				def       snake_case    (     self					:	str						):
						return 3_2



				@property
				def       snake_case    (     self					:	Optional[int]						):
						return self.time_input_dim



				@property
				def       snake_case    (     self					:	Union[str, Any]						):
						return self.time_input_dim * 4



				@property
				def       snake_case    (     self					:	Optional[int]						):
						return 1_0_0



				@property
				def       snake_case    (     self					:	str						):
						torch.manual_seed(0						)

						__lowercase				:		Tuple												=				{
						    "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,
						}

						__lowercase				:		Union[str, Any]												=				UNetaDConditionModel(**lowercase__						)
						return model



				@property
				def       snake_case    (     self					:	int						):
						return {
						    "block_out_channels": [3_2, 3_2, 6_4, 6_4],
						    "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": 1_2,
						    "out_channels": 3,
						    "up_block_types": ["AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"],
						    "vq_embed_dim": 4,
						}



				@property
				def       snake_case    (     self					:	str						):
						torch.manual_seed(0						)
						__lowercase				:		int												=				VQModel(**self.dummy_movq_kwargs						)
						return model



				def       snake_case    (     self					:	List[Any]						):
						__lowercase				:		Optional[Any]												=				self.dummy_unet
						__lowercase				:		Optional[Any]												=				self.dummy_movq

						__lowercase				:		Union[str, Any]												=				{
						    "num_train_timesteps": 1_0_0_0,
						    "beta_schedule": "linear",
						    "beta_start": 0.0_0_0_8_5,
						    "beta_end": 0.0_1_2,
						    "clip_sample": False,
						    "set_alpha_to_one": False,
						    "steps_offset": 0,
						    "prediction_type": "epsilon",
						    "thresholding": False,
						}

						__lowercase				:		List[str]												=				DDIMScheduler(**lowercase__						)

						__lowercase				:		Optional[Any]												=				{
						    "unet": unet,
						    "scheduler": scheduler,
						    "movq": movq,
						}

						return components



				def       snake_case    (     self					:	Any    ,   lowercase__					:	Optional[int]    ,   lowercase__					:	Optional[Any]=0						):
						__lowercase				:		List[str]												=				floats_tensor((1, self.text_embedder_hidden_size)    ,   rng=random.Random(lowercase__						)						).to(lowercase__						)
						__lowercase				:		Any												=				floats_tensor((1, self.text_embedder_hidden_size)    ,   rng=random.Random(seed + 1						)						).to(
						    lowercase__						)
						# create init_image
						__lowercase				:		Optional[int]												=				floats_tensor((1, 3, 6_4, 6_4)    ,   rng=random.Random(lowercase__						)						).to(lowercase__						)
						__lowercase				:		Optional[Any]												=				image.cpu().permute(0    ,   2    ,   3    ,   1						)[0]
						__lowercase				:		Union[str, Any]												=				Image.fromarray(np.uinta(lowercase__						)						).convert("RGB"						).resize((2_5_6, 2_5_6)						)
						# create hint
						__lowercase				:		str												=				floats_tensor((1, 3, 6_4, 6_4)    ,   rng=random.Random(lowercase__						)						).to(lowercase__						)

						if str(lowercase__						).startswith("mps"						):
								__lowercase				:		Tuple												=				torch.manual_seed(lowercase__						)
						else:
								__lowercase				:		Tuple												=				torch.Generator(device=lowercase__						).manual_seed(lowercase__						)
						__lowercase				:		Optional[int]												=				{
						    "image": init_image,
						    "image_embeds": image_embeds,
						    "negative_image_embeds": negative_image_embeds,
						    "hint": hint,
						    "generator": generator,
						    "height": 6_4,
						    "width": 6_4,
						    "num_inference_steps": 1_0,
						    "guidance_scale": 7.0,
						    "strength": 0.2,
						    "output_type": "np",
						}
						return inputs



				def       snake_case    (     self					:	Tuple						):
						__lowercase				:		Union[str, Any]												=				"cpu"

						__lowercase				:		str												=				self.get_dummy_components()

						__lowercase				:		Dict												=				self.pipeline_class(**lowercase__						)
						__lowercase				:		Dict												=				pipe.to(lowercase__						)

						pipe.set_progress_bar_config(disable=lowercase__						)

						__lowercase				:		Any												=				pipe(**self.get_dummy_inputs(lowercase__						)						)
						__lowercase				:		str												=				output.images

						__lowercase				:		Any												=				pipe(
						    **self.get_dummy_inputs(lowercase__						)    ,   return_dict=lowercase__    ,   )[0]

						__lowercase				:		List[str]												=				image[0, -3:, -3:, -1]
						__lowercase				:		List[str]												=				image_from_tuple[0, -3:, -3:, -1]

						assert image.shape == (1, 6_4, 6_4, 3)

						__lowercase				:		List[str]												=				np.array(
						    [0.5_4_9_8_5_0_3_4, 0.5_5_5_0_9_3_6_5, 0.5_2_5_6_1_5_0_4, 0.5_5_7_0_4_9_4, 0.5_5_9_3_8_1_8, 0.5_2_6_3_9_7_9, 0.5_0_2_8_5_6_4_3, 0.5_0_6_9_8_4_6, 0.5_1_1_9_6_7_3_6]						)
						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       snake_case    (     self					:	List[str]						):
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()



				def       snake_case    (     self					:	str						):
						__lowercase				:		Optional[int]												=				load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/kandinskyv22/kandinskyv22_controlnet_img2img_robotcat_fp16.npy"						)

						__lowercase				:		List[str]												=				load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/kandinsky/cat.png"						)
						__lowercase				:		List[Any]												=				init_image.resize((5_1_2, 5_1_2)						)

						__lowercase				:		List[Any]												=				load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/kandinskyv22/hint_image_cat.png"						)
						__lowercase				:		Any												=				torch.from_numpy(np.array(lowercase__						)						).float() / 2_5_5.0
						__lowercase				:		Tuple												=				hint.permute(2    ,   0    ,   1						).unsqueeze(0						)

						__lowercase				:		Tuple												=				"A robot, 4k photo"

						__lowercase				:		str												=				KandinskyVaaPriorEmbaEmbPipeline.from_pretrained(
						    "kandinsky-community/kandinsky-2-2-prior"    ,   torch_dtype=torch.floataa						)
						pipe_prior.to(lowercase__						)

						__lowercase				:		List[Any]												=				KandinskyVaaControlnetImgaImgPipeline.from_pretrained(
						    "kandinsky-community/kandinsky-2-2-controlnet-depth"    ,   torch_dtype=torch.floataa						)
						__lowercase				:		List[Any]												=				pipeline.to(lowercase__						)

						pipeline.set_progress_bar_config(disable=lowercase__						)

						__lowercase				:		Optional[int]												=				torch.Generator(device="cpu"						).manual_seed(0						)

						__lowercase						,__lowercase				:		Union[str, Any]												=				pipe_prior(
						    lowercase__    ,   image=lowercase__    ,   strength=0.8_5    ,   generator=lowercase__    ,   negative_prompt=""    ,   ).to_tuple()

						__lowercase				:		Dict												=				pipeline(
						    image=lowercase__    ,   image_embeds=lowercase__    ,   negative_image_embeds=lowercase__    ,   hint=lowercase__    ,   generator=lowercase__    ,   num_inference_steps=1_0_0    ,   height=5_1_2    ,   width=5_1_2    ,   strength=0.5    ,   output_type="np"    ,   )

						__lowercase				:		List[Any]												=				output.images[0]

						assert image.shape == (5_1_2, 5_1_2, 3)

						assert_mean_pixel_difference(lowercase__    ,   lowercase__						)






"""simple docstring"""







import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers.models.esm.modeling_esmfold import EsmForProteinFolding




class       lowerCAmelCase__       :


				"""simple docstring"""





				def __init__(     self					:	Tuple    ,   lowercase__					:	Tuple    ,   lowercase__					:	Tuple=1_3    ,   lowercase__					:	Optional[int]=7    ,   lowercase__					:	List[str]=False    ,   lowercase__					:	Tuple=True    ,   lowercase__					:	int=False    ,   lowercase__					:	List[str]=False    ,   lowercase__					:	Optional[Any]=1_9    ,   lowercase__					:	int=3_2    ,   lowercase__					:	List[Any]=5    ,   lowercase__					:	Optional[int]=4    ,   lowercase__					:	Any=3_7    ,   lowercase__					:	Tuple="gelu"    ,   lowercase__					:	int=0.1    ,   lowercase__					:	Tuple=0.1    ,   lowercase__					:	List[Any]=5_1_2    ,   lowercase__					:	List[Any]=1_6    ,   lowercase__					:	Union[str, Any]=2    ,   lowercase__					:	List[str]=0.0_2    ,   lowercase__					:	List[Any]=3    ,   lowercase__					:	Any=4    ,   lowercase__					:	Optional[Any]=None    ,   ):
						__lowercase				:		int												=				parent
						__lowercase				:		Tuple												=				batch_size
						__lowercase				:		Optional[int]												=				seq_length
						__lowercase				:		str												=				is_training
						__lowercase				:		List[Any]												=				use_input_mask
						__lowercase				:		Any												=				use_token_type_ids
						__lowercase				:		str												=				use_labels
						__lowercase				:		Dict												=				vocab_size
						__lowercase				:		Optional[int]												=				hidden_size
						__lowercase				:		List[Any]												=				num_hidden_layers
						__lowercase				:		int												=				num_attention_heads
						__lowercase				:		Union[str, Any]												=				intermediate_size
						__lowercase				:		Dict												=				hidden_act
						__lowercase				:		Union[str, Any]												=				hidden_dropout_prob
						__lowercase				:		Optional[int]												=				attention_probs_dropout_prob
						__lowercase				:		Optional[Any]												=				max_position_embeddings
						__lowercase				:		List[str]												=				type_vocab_size
						__lowercase				:		str												=				type_sequence_label_size
						__lowercase				:		List[str]												=				initializer_range
						__lowercase				:		Optional[Any]												=				num_labels
						__lowercase				:		Tuple												=				num_choices
						__lowercase				:		Optional[Any]												=				scope



				def       snake_case    (     self					:	Optional[int]						):
						__lowercase				:		Tuple												=				ids_tensor([self.batch_size, self.seq_length]    ,   self.vocab_size						)

						__lowercase				:		Union[str, Any]												=				None
						if self.use_input_mask:
								__lowercase				:		List[str]												=				random_attention_mask([self.batch_size, self.seq_length]						)

						__lowercase				:		Dict												=				None
						__lowercase				:		Any												=				None
						__lowercase				:		Optional[Any]												=				None
						if self.use_labels:
								__lowercase				:		Tuple												=				ids_tensor([self.batch_size]    ,   self.type_sequence_label_size						)
								__lowercase				:		List[str]												=				ids_tensor([self.batch_size, self.seq_length]    ,   self.num_labels						)
								__lowercase				:		Optional[Any]												=				ids_tensor([self.batch_size]    ,   self.num_choices						)

						__lowercase				:		Optional[Any]												=				self.get_config()

						return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels



				def       snake_case    (     self					:	Any						):
						__lowercase				:		Dict												=				EsmConfig(
						    vocab_size=3_3    ,   hidden_size=self.hidden_size    ,   pad_token_id=1    ,   num_hidden_layers=self.num_hidden_layers    ,   num_attention_heads=self.num_attention_heads    ,   intermediate_size=self.intermediate_size    ,   hidden_act=self.hidden_act    ,   hidden_dropout_prob=self.hidden_dropout_prob    ,   attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,   max_position_embeddings=self.max_position_embeddings    ,   type_vocab_size=self.type_vocab_size    ,   initializer_range=self.initializer_range    ,   is_folding_model=lowercase__    ,   esmfold_config={"trunk": {"num_blocks": 2}, "fp16_esm": False}    ,   )
						return config



				def       snake_case    (     self					:	str    ,   lowercase__					:	List[Any]    ,   lowercase__					:	Optional[int]    ,   lowercase__					:	List[Any]    ,   lowercase__					:	int    ,   lowercase__					:	List[Any]    ,   lowercase__					:	str						):
						__lowercase				:		List[Any]												=				EsmForProteinFolding(config=lowercase__						).float()
						model.to(lowercase__						)
						model.eval()
						__lowercase				:		List[str]												=				model(lowercase__    ,   attention_mask=lowercase__						)
						__lowercase				:		Any												=				model(lowercase__						)
						__lowercase				:		Any												=				model(lowercase__						)

						self.parent.assertEqual(result.positions.shape    ,   (8, self.batch_size, self.seq_length, 1_4, 3)						)
						self.parent.assertEqual(result.angles.shape    ,   (8, self.batch_size, self.seq_length, 7, 2)						)



				def       snake_case    (     self					:	Tuple						):
						__lowercase				:		List[str]												=				self.prepare_config_and_inputs()
						(
						    (
						    __lowercase
						)						,(
						    __lowercase
						)						,(
						    __lowercase
						)						,(
						    __lowercase
						)						,(
						    __lowercase
						)						,(
						    __lowercase
						)						,
						)				:		Any												=				config_and_inputs
						__lowercase				:		int												=				{"input_ids": input_ids, "attention_mask": input_mask}
						return config, inputs_dict





@require_torch
class       lowerCAmelCase__       (			lowerCAmelCase_  ,			lowerCAmelCase_  ,			unittest.TestCase    ):


				"""simple docstring"""





				__UpperCAmelCase :     Union[str, Any]       =      False

				__UpperCAmelCase :     List[Any]       =      (EsmForProteinFolding,) if is_torch_available() else ()
				__UpperCAmelCase :     int       =      ()
				__UpperCAmelCase :     Tuple       =      {} if is_torch_available() else {}
				__UpperCAmelCase :     Optional[Any]       =      False



				def       snake_case    (     self					:	Tuple						):
						__lowercase				:		Tuple												=				EsmFoldModelTester(self						)
						__lowercase				:		Dict												=				ConfigTester(self    ,   config_class=lowercase__    ,   hidden_size=3_7						)



				def       snake_case    (     self					:	Dict						):
						self.config_tester.run_common_tests()



				def       snake_case    (     self					:	Optional[int]						):
						__lowercase				:		Tuple												=				self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*lowercase__						)



				@unittest.skip("Does not support attention outputs"						)
				def       snake_case    (     self					:	Union[str, Any]						):
						pass



				@unittest.skip
				def       snake_case    (     self					:	Tuple						):
						pass



				@unittest.skip("Esm does not support embedding resizing"						)
				def       snake_case    (     self					:	Optional[int]						):
						pass



				@unittest.skip("Esm does not support embedding resizing"						)
				def       snake_case    (     self					:	List[str]						):
						pass



				@unittest.skip("ESMFold does not support passing input embeds!"						)
				def       snake_case    (     self					:	int						):
						pass



				@unittest.skip("ESMFold does not support head pruning."						)
				def       snake_case    (     self					:	List[Any]						):
						pass



				@unittest.skip("ESMFold does not support head pruning."						)
				def       snake_case    (     self					:	Any						):
						pass



				@unittest.skip("ESMFold does not support head pruning."						)
				def       snake_case    (     self					:	Optional[Any]						):
						pass



				@unittest.skip("ESMFold does not support head pruning."						)
				def       snake_case    (     self					:	List[str]						):
						pass



				@unittest.skip("ESMFold does not support head pruning."						)
				def       snake_case    (     self					:	List[str]						):
						pass



				@unittest.skip("ESMFold does not output hidden states in the normal way."						)
				def       snake_case    (     self					:	Optional[Any]						):
						pass



				@unittest.skip("ESMfold does not output hidden states in the normal way."						)
				def       snake_case    (     self					:	Optional[Any]						):
						pass



				@unittest.skip("ESMFold only has one output format."						)
				def       snake_case    (     self					:	Tuple						):
						pass



				@unittest.skip("This test doesn't work for ESMFold and doesn't test core functionality"						)
				def       snake_case    (     self					:	Any						):
						pass



				@unittest.skip("ESMFold does not support input chunking."						)
				def       snake_case    (     self					:	str						):
						pass



				@unittest.skip("ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments."						)
				def       snake_case    (     self					:	Any						):
						pass



				@unittest.skip("ESMFold doesn't support torchscript compilation."						)
				def       snake_case    (     self					:	Any						):
						pass



				@unittest.skip("ESMFold doesn't support torchscript compilation."						)
				def       snake_case    (     self					:	Tuple						):
						pass



				@unittest.skip("ESMFold doesn't support torchscript compilation."						)
				def       snake_case    (     self					:	List[str]						):
						pass



				@unittest.skip("ESMFold doesn't support data parallel."						)
				def       snake_case    (     self					:	Optional[Any]						):
						pass



				@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."						)
				def       snake_case    (     self					:	Any						):
						pass





@require_torch
class       lowerCAmelCase__       (			lowerCAmelCase_    ):


				"""simple docstring"""





				@slow
				def       snake_case    (     self					:	Union[str, Any]						):
						__lowercase				:		Tuple												=				EsmForProteinFolding.from_pretrained("facebook/esmfold_v1"						).float()
						model.eval()
						__lowercase				:		Optional[int]												=				torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]						)
						__lowercase				:		str												=				model(lowercase__						)["positions"]
						__lowercase				:		Union[str, Any]												=				torch.tensor([2.5_8_2_8, 0.7_9_9_3, -1_0.9_3_3_4]    ,   dtype=torch.floataa						)
						self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]    ,   lowercase__    ,   atol=1e-4						)						)
