


'''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		:						List[Any]  ,					_A		:						int  ,					_A		:						str=13  ,					_A		:						Optional[Any]=7  ,					_A		:						Union[str, Any]=False  ,					_A		:						Optional[int]=True  ,					_A		:						Optional[int]=False  ,					_A		:						List[str]=False  ,					_A		:						List[str]=19  ,					_A		:						Union[str, Any]=32  ,					_A		:						List[Any]=5  ,					_A		:						Dict=4  ,					_A		:						Any=37  ,					_A		:						Union[str, Any]="gelu"  ,					_A		:						Dict=0.1  ,					_A		:						List[Any]=0.1  ,					_A		:						List[str]=512  ,					_A		:						Tuple=16  ,					_A		:						Optional[int]=2  ,					_A		:						Dict=0.02  ,					_A		:						str=3  ,					_A		:						List[Any]=4  ,					_A		:						str=None  ,					)  ->			List[Any]:
			__magic_name__	:   Tuple					     =					parent
			__magic_name__	:   Dict					     =					batch_size
			__magic_name__	:   int					     =					seq_length
			__magic_name__	:   Union[str, Any]					     =					is_training
			__magic_name__	:   Dict					     =					use_input_mask
			__magic_name__	:   List[str]					     =					use_token_type_ids
			__magic_name__	:   Optional[int]					     =					use_labels
			__magic_name__	:   Union[str, Any]					     =					vocab_size
			__magic_name__	:   Dict					     =					hidden_size
			__magic_name__	:   List[Any]					     =					num_hidden_layers
			__magic_name__	:   int					     =					num_attention_heads
			__magic_name__	:   List[Any]					     =					intermediate_size
			__magic_name__	:   Dict					     =					hidden_act
			__magic_name__	:   str					     =					hidden_dropout_prob
			__magic_name__	:   List[Any]					     =					attention_probs_dropout_prob
			__magic_name__	:   Optional[Any]					     =					max_position_embeddings
			__magic_name__	:   Dict					     =					type_vocab_size
			__magic_name__	:   int					     =					type_sequence_label_size
			__magic_name__	:   Any					     =					initializer_range
			__magic_name__	:   Union[str, Any]					     =					num_labels
			__magic_name__	:   Tuple					     =					num_choices
			__magic_name__	:   Union[str, Any]					     =					scope







		def        __lowerCAmelCase		( self		:						Dict      )  ->			List[Any]:
			__magic_name__	:   Dict					     =					ids_tensor([self.batch_size, self.seq_length]  ,					self.vocab_size      )

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

			__magic_name__	:   Optional[int]					     =					None
			__magic_name__	:   int					     =					None
			__magic_name__	:   str					     =					None
			if self.use_labels:
				__magic_name__	:   str					     =					ids_tensor([self.batch_size]  ,					self.type_sequence_label_size      )
				__magic_name__	:   Optional[int]					     =					ids_tensor([self.batch_size, self.seq_length]  ,					self.num_labels      )
				__magic_name__	:   str					     =					ids_tensor([self.batch_size]  ,					self.num_choices      )

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

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







		def        __lowerCAmelCase		( self		:						str      )  ->			List[str]:
			__magic_name__	:   List[Any]					     =					EsmConfig(
			    vocab_size=33  ,					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=_A  ,					esmfold_config={'trunk': {'num_blocks': 2}, 'fp16_esm': False}  ,					)
			return config







		def        __lowerCAmelCase		( self		:						Optional[Any]  ,					_A		:						Optional[Any]  ,					_A		:						Optional[Any]  ,					_A		:						Tuple  ,					_A		:						Tuple  ,					_A		:						Optional[Any]  ,					_A		:						Optional[int]      )  ->			Union[str, Any]:
			__magic_name__	:   Any					     =					EsmForProteinFolding(config=_A      ).float()
			model.to(_A      )
			model.eval()
			__magic_name__	:   List[str]					     =					model(_A  ,					attention_mask=_A      )
			__magic_name__	:   Any					     =					model(_A      )
			__magic_name__	:   Union[str, Any]					     =					model(_A      )

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







		def        __lowerCAmelCase		( self		:						Any      )  ->			Union[str, Any]:
			__magic_name__	:   List[Any]					     =					self.prepare_config_and_inputs()
			(
			    (
			    __magic_name__
			)     ,    (
			    __magic_name__
			)     ,    (
			    __magic_name__
			)     ,    (
			    __magic_name__
			)     ,    (
			    __magic_name__
			)     ,    (
			    __magic_name__
			)     ,    
			)	:   Dict					     =					config_and_inputs
			__magic_name__	:   Dict					     =					{'input_ids': input_ids, 'attention_mask': input_mask}
			return config, inputs_dict







@require_torch
class _lowerCamelCase	(						lowercase__		,					lowercase__		,					unittest.TestCase							):



		'''simple docstring'''







		A_     :      str				   =    False

		A_     :      Union[str, Any]				   =    (EsmForProteinFolding,) if is_torch_available() else ()
		A_     :      List[Any]				   =    ()
		A_     :      Any				   =    {} if is_torch_available() else {}
		A_     :      Optional[int]				   =    False







		def        __lowerCAmelCase		( self		:						Tuple      )  ->			int:
			__magic_name__	:   Optional[int]					     =					EsmFoldModelTester(self      )
			__magic_name__	:   List[Any]					     =					ConfigTester(self  ,					config_class=_A  ,					hidden_size=37      )







		def        __lowerCAmelCase		( self		:						Union[str, Any]      )  ->			Tuple:
			self.config_tester.run_common_tests()







		def        __lowerCAmelCase		( self		:						List[Any]      )  ->			List[Any]:
			__magic_name__	:   int					     =					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*_A      )







		@unittest.skip('Does not support attention outputs'      )
		def        __lowerCAmelCase		( self		:						List[Any]      )  ->			Tuple:
			pass







		@unittest.skip
		def        __lowerCAmelCase		( self		:						List[str]      )  ->			List[str]:
			pass







		@unittest.skip('Esm does not support embedding resizing'      )
		def        __lowerCAmelCase		( self		:						int      )  ->			str:
			pass







		@unittest.skip('Esm does not support embedding resizing'      )
		def        __lowerCAmelCase		( self		:						Any      )  ->			Any:
			pass







		@unittest.skip('ESMFold does not support passing input embeds!'      )
		def        __lowerCAmelCase		( self		:						int      )  ->			int:
			pass







		@unittest.skip('ESMFold does not support head pruning.'      )
		def        __lowerCAmelCase		( self		:						str      )  ->			int:
			pass







		@unittest.skip('ESMFold does not support head pruning.'      )
		def        __lowerCAmelCase		( self		:						str      )  ->			Union[str, Any]:
			pass







		@unittest.skip('ESMFold does not support head pruning.'      )
		def        __lowerCAmelCase		( self		:						str      )  ->			List[str]:
			pass







		@unittest.skip('ESMFold does not support head pruning.'      )
		def        __lowerCAmelCase		( self		:						List[Any]      )  ->			Optional[Any]:
			pass







		@unittest.skip('ESMFold does not support head pruning.'      )
		def        __lowerCAmelCase		( self		:						Optional[Any]      )  ->			int:
			pass







		@unittest.skip('ESMFold does not output hidden states in the normal way.'      )
		def        __lowerCAmelCase		( self		:						Any      )  ->			Tuple:
			pass







		@unittest.skip('ESMfold does not output hidden states in the normal way.'      )
		def        __lowerCAmelCase		( self		:						Dict      )  ->			int:
			pass







		@unittest.skip('ESMFold only has one output format.'      )
		def        __lowerCAmelCase		( self		:						int      )  ->			Optional[Any]:
			pass







		@unittest.skip('This test doesn\'t work for ESMFold and doesn\'t test core functionality'      )
		def        __lowerCAmelCase		( self		:						int      )  ->			Tuple:
			pass







		@unittest.skip('ESMFold does not support input chunking.'      )
		def        __lowerCAmelCase		( self		:						List[Any]      )  ->			List[Any]:
			pass







		@unittest.skip('ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.'      )
		def        __lowerCAmelCase		( self		:						Optional[Any]      )  ->			List[Any]:
			pass







		@unittest.skip('ESMFold doesn\'t support torchscript compilation.'      )
		def        __lowerCAmelCase		( self		:						List[str]      )  ->			Optional[int]:
			pass







		@unittest.skip('ESMFold doesn\'t support torchscript compilation.'      )
		def        __lowerCAmelCase		( self		:						Tuple      )  ->			List[Any]:
			pass







		@unittest.skip('ESMFold doesn\'t support torchscript compilation.'      )
		def        __lowerCAmelCase		( self		:						Union[str, Any]      )  ->			Optional[int]:
			pass







		@unittest.skip('ESMFold doesn\'t support data parallel.'      )
		def        __lowerCAmelCase		( self		:						Any      )  ->			Optional[int]:
			pass







		@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'      )
		def        __lowerCAmelCase		( self		:						Tuple      )  ->			Optional[Any]:
			pass







@require_torch
class _lowerCamelCase	(						lowercase__							):



		'''simple docstring'''







		@slow
		def        __lowerCAmelCase		( self		:						Any      )  ->			Any:
			__magic_name__	:   Any					     =					EsmForProteinFolding.from_pretrained('facebook/esmfold_v1'      ).float()
			model.eval()
			__magic_name__	:   int					     =					torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]      )
			__magic_name__	:   List[Any]					     =					model(_A      )['positions']
			__magic_name__	:   Tuple					     =					torch.tensor([2.5828, 0.7993, -10.9334]  ,					dtype=torch.floataa      )
			self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]  ,					_A  ,					atol=1E-4      )      )



'''simple docstring'''

from typing import List, Optional, Tuple, Union

import PIL
import torch
from torchvision import transforms

from diffusers.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
from diffusers.schedulers import DDIMScheduler
from diffusers.utils import randn_tensor


lowerCAmelCase				:Optional[Any]        =      transforms.Compose(
    [
        transforms.Resize((2_5_6, 2_5_6)),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5]),
    ]
)





def lowerCamelCase      (      lowerCAmelCase    :     Tuple		):





	"""simple docstring"""







	if isinstance(lowerCAmelCase		,		torch.Tensor		):
		return image
	elif isinstance(lowerCAmelCase		,		PIL.Image.Image		):
		__magic_name__	:   List[Any]					     =					[image]

	__magic_name__	:   List[Any]					     =					[trans(img.convert('RGB'		)		) for img in image]
	__magic_name__	:   Dict					     =					torch.stack(lowerCAmelCase		)
	return image





class _lowerCamelCase	(						lowercase__							):



		'''simple docstring'''







		def __init__( self		:						Optional[Any]  ,					_A		:						str  ,					_A		:						int      )  ->			Dict:
			super().__init__()

			# make sure scheduler can always be converted to DDIM
			__magic_name__	:   Optional[int]					     =					DDIMScheduler.from_config(scheduler.config      )

			self.register_modules(unet=_A  ,					scheduler=_A      )







		def        __lowerCAmelCase		( self		:						Union[str, Any]  ,					_A		:						Union[str, Any]      )  ->			Optional[int]:
			if strength < 0 or strength > 1:
				raise ValueError(F'The value of strength should in [0.0, 1.0] but is {strength}'      )







		def        __lowerCAmelCase		( self		:						Any  ,					_A		:						List[str]  ,					_A		:						Optional[Any]  ,					_A		:						int      )  ->			List[Any]:
			# get the original timestep using init_timestep
			__magic_name__	:   Tuple					     =					min(int(num_inference_steps * strength      )  ,					_A      )

			__magic_name__	:   Any					     =					max(num_inference_steps - init_timestep  ,					0      )
			__magic_name__	:   List[str]					     =					self.scheduler.timesteps[t_start:]

			return timesteps, num_inference_steps - t_start







		def        __lowerCAmelCase		( self		:						Any  ,					_A		:						str  ,					_A		:						Optional[int]  ,					_A		:						Tuple  ,					_A		:						List[str]  ,					_A		:						str  ,					_A		:						Optional[int]=None      )  ->			Dict:
			if not isinstance(_A  ,					(torch.Tensor, PIL.Image.Image, list)      ):
				raise ValueError(
				    F'`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(_A      )}'      )

			__magic_name__	:   Union[str, Any]					     =					image.to(device=_A  ,					dtype=_A      )

			if isinstance(_A  ,					_A      ) and len(_A      ) != batch_size:
				raise ValueError(
				    F'You have passed a list of generators of length {len(_A      )}, but requested an effective batch'
				    F' size of {batch_size}. Make sure the batch size matches the length of the generators.'      )

			__magic_name__	:   Tuple					     =					init_latents.shape
			__magic_name__	:   Any					     =					randn_tensor(_A  ,					generator=_A  ,					device=_A  ,					dtype=_A      )

			# get latents
			print('add noise to latents at timestep'  ,					_A      )
			__magic_name__	:   List[str]					     =					self.scheduler.add_noise(_A  ,					_A  ,					_A      )
			__magic_name__	:   List[str]					     =					init_latents

			return latents







		@torch.no_grad()
		def __call__( self		:						Tuple  ,					_A		:						Union[torch.FloatTensor, PIL.Image.Image] = None  ,					_A		:						float = 0.8  ,					_A		:						int = 1  ,					_A		:						Optional[Union[torch.Generator, List[torch.Generator]]] = None  ,					_A		:						float = 0.0  ,					_A		:						int = 50  ,					_A		:						Optional[bool] = None  ,					_A		:						Optional[str] = "pil"  ,					_A		:						bool = True  ,					)  ->			Union[ImagePipelineOutput, Tuple]:
			self.check_inputs(_A      )

			# 2. Preprocess image
			__magic_name__	:   int					     =					preprocess(_A      )

			# 3. set timesteps
			self.scheduler.set_timesteps(_A  ,					device=self.device      )
			__magic_name__     ,    __magic_name__	:   Dict					     =					self.get_timesteps(_A  ,					_A  ,					self.device      )
			__magic_name__	:   Dict					     =					timesteps[:1].repeat(_A      )

			# 4. Prepare latent variables
			__magic_name__	:   Optional[Any]					     =					self.prepare_latents(_A  ,					_A  ,					_A  ,					self.unet.dtype  ,					self.device  ,					_A      )
			__magic_name__	:   Optional[Any]					     =					latents

			# 5. Denoising loop
			for t in self.progress_bar(_A      ):
				# 1. predict noise model_output
				__magic_name__	:   Dict					     =					self.unet(_A  ,					_A      ).sample

				# 2. predict previous mean of image x_t-1 and add variance depending on eta
				# eta corresponds to η in paper and should be between [0, 1]
				# do x_t -> x_t-1
				__magic_name__	:   List[Any]					     =					self.scheduler.step(
				    _A  ,					_A  ,					_A  ,					eta=_A  ,					use_clipped_model_output=_A  ,					generator=_A  ,					).prev_sample

			__magic_name__	:   Union[str, Any]					     =					(image / 2 + 0.5).clamp(0  ,					1      )
			__magic_name__	:   List[Any]					     =					image.cpu().permute(0  ,					2  ,					3  ,					1      ).numpy()
			if output_type == "pil":
				__magic_name__	:   Dict					     =					self.numpy_to_pil(_A      )

			if not return_dict:
				return (image, latent_timestep.item())

			return ImagePipelineOutput(images=_A      )