from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from transformers import CLIPImageProcessor, CLIPVisionModel

from ...models import PriorTransformer
from ...pipelines import DiffusionPipeline
from ...schedulers import HeunDiscreteScheduler
from ...utils import (
    BaseOutput,
    is_accelerate_available,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from .renderer import ShapERenderer


__lowercase     :   Optional[Any]											=      logging.get_logger(__name__)  # pylint: disable=invalid-name

__lowercase     :   Union[str, Any]											=      '''
    Examples:
        ```py
        >>> from PIL import Image
        >>> import torch
        >>> from diffusers import DiffusionPipeline
        >>> from diffusers.utils import export_to_gif, load_image

        >>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        >>> repo = "openai/shap-e-img2img"
        >>> pipe = DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.float16)
        >>> pipe = pipe.to(device)

        >>> guidance_scale = 3.0
        >>> image_url = "https://hf.co/datasets/diffusers/docs-images/resolve/main/shap-e/corgi.png"
        >>> image = load_image(image_url).convert("RGB")

        >>> images = pipe(
        ...     image,
        ...     guidance_scale=guidance_scale,
        ...     num_inference_steps=64,
        ...     frame_size=256,
        ... ).images

        >>> gif_path = export_to_gif(images[0], "corgi_3d.gif")
        ```
'''


@dataclass
class 			lowerCamelCase_  (     UpperCAmelCase_				):






    '''simple docstring'''






    a__			:		Union[PIL.Image.Image, np.ndarray]


class 			lowerCamelCase_  (     UpperCAmelCase_				):






    '''simple docstring'''






    def __init__(			self				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				) ->			Tuple:
        super().__init__()

        self.register_modules(
            prior=__lowercase				,				image_encoder=__lowercase				,				image_processor=__lowercase				,				scheduler=__lowercase				,				renderer=__lowercase				,				)







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase) ->			Dict:
        if latents is None:
            __UpperCamelCase  :List[Any]    						=	randn_tensor(__lowercase				,				generator=__lowercase				,				device=__lowercase				,				dtype=__lowercase)
        else:
            if latents.shape != shape:
                raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {shape}""")
            __UpperCamelCase  :Dict    						=	latents.to(__lowercase)

        __UpperCamelCase  :List[str]    						=	latents * scheduler.init_noise_sigma
        return latents







    def     UpperCamelCase__     (			self				,				__lowercase=0) ->			Any:
        if is_accelerate_available():
            from accelerate import cpu_offload
        else:
            raise ImportError('''Please install accelerate via `pip install accelerate`''')

        __UpperCamelCase  :Tuple    						=	torch.device(f"""cuda:{gpu_id}""")

        __UpperCamelCase  :Optional[int]    						=	[self.image_encoder, self.prior]
        for cpu_offloaded_model in models:
            if cpu_offloaded_model is not None:
                cpu_offload(__lowercase				,				__lowercase)







    @property
    def     UpperCamelCase__     (			self) ->			Any:
        if self.device != torch.device('''meta''') or not hasattr(self.image_encoder				,				'''_hf_hook'''):
            return self.device
        for module in self.image_encoder.modules():
            if (
                hasattr(__lowercase				,				'''_hf_hook''')
                and hasattr(module._hf_hook				,				'''execution_device''')
                and module._hf_hook.execution_device is not None
            ):
                return torch.device(module._hf_hook.execution_device)
        return self.device







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				) ->			Optional[Any]:
        if isinstance(__lowercase				,				__lowercase) and isinstance(image[0]				,				torch.Tensor):
            __UpperCamelCase  :str    						=	torch.cat(__lowercase				,				axis=0) if image[0].ndim == 4 else torch.stack(__lowercase				,				axis=0)

        if not isinstance(__lowercase				,				torch.Tensor):
            __UpperCamelCase  :Optional[int]    						=	self.image_processor(__lowercase				,				return_tensors='''pt''').pixel_values[0].unsqueeze(0)

        __UpperCamelCase  :Any    						=	image.to(dtype=self.image_encoder.dtype				,				device=__lowercase)

        __UpperCamelCase  :List[Any]    						=	self.image_encoder(__lowercase)['''last_hidden_state''']
        __UpperCamelCase  :List[Any]    						=	image_embeds[:, 1:, :].contiguous()  # batch_size, dim, 256

        __UpperCamelCase  :str    						=	image_embeds.repeat_interleave(__lowercase				,				dim=0)

        if do_classifier_free_guidance:
            __UpperCamelCase  :Dict    						=	torch.zeros_like(__lowercase)

            # For classifier free guidance, we need to do two forward passes.
            # Here we concatenate the unconditional and text embeddings into a single batch
            # to avoid doing two forward passes
            __UpperCamelCase  :List[str]    						=	torch.cat([negative_image_embeds, image_embeds])

        return image_embeds







    @torch.no_grad()
    @replace_example_docstring(__lowercase)
    def __call__(			self				,				__lowercase				,				__lowercase = 1				,				__lowercase = 25				,				__lowercase = None				,				__lowercase = None				,				__lowercase = 4.0				,				__lowercase = 64				,				__lowercase = "pil"				,				__lowercase = True				,				) ->			List[Any]:

        if isinstance(__lowercase				,				PIL.Image.Image):
            __UpperCamelCase  :Dict    						=	1
        elif isinstance(__lowercase				,				torch.Tensor):
            __UpperCamelCase  :List[str]    						=	image.shape[0]
        elif isinstance(__lowercase				,				__lowercase) and isinstance(image[0]				,				(torch.Tensor, PIL.Image.Image)):
            __UpperCamelCase  :Any    						=	len(__lowercase)
        else:
            raise ValueError(
                f"""`image` has to be of type `PIL.Image.Image`, `torch.Tensor`, `List[PIL.Image.Image]` or `List[torch.Tensor]` but is {type(__lowercase)}""")

        __UpperCamelCase  :Optional[int]    						=	self._execution_device

        __UpperCamelCase  :Tuple    						=	batch_size * num_images_per_prompt

        __UpperCamelCase  :str    						=	guidance_scale > 1.0
        __UpperCamelCase  :Tuple    						=	self._encode_image(__lowercase				,				__lowercase				,				__lowercase				,				__lowercase)

        # prior

        self.scheduler.set_timesteps(__lowercase				,				device=__lowercase)
        __UpperCamelCase  :Any    						=	self.scheduler.timesteps

        __UpperCamelCase  :int    						=	self.prior.config.num_embeddings
        __UpperCamelCase  :str    						=	self.prior.config.embedding_dim

        __UpperCamelCase  :str    						=	self.prepare_latents(
            (batch_size, num_embeddings * embedding_dim)				,				image_embeds.dtype				,				__lowercase				,				__lowercase				,				__lowercase				,				self.scheduler				,				)

        # YiYi notes: for testing only to match ldm, we can directly create a latents with desired shape: batch_size, num_embeddings, embedding_dim
        __UpperCamelCase  :Tuple    						=	latents.reshape(latents.shape[0]				,				__lowercase				,				__lowercase)

        for i, t in enumerate(self.progress_bar(__lowercase)):
            # expand the latents if we are doing classifier free guidance
            __UpperCamelCase  :List[Any]    						=	torch.cat([latents] * 2) if do_classifier_free_guidance else latents
            __UpperCamelCase  :Dict    						=	self.scheduler.scale_model_input(__lowercase				,				__lowercase)

            __UpperCamelCase  :Union[str, Any]    						=	self.prior(
                __lowercase				,				timestep=__lowercase				,				proj_embedding=__lowercase				,				).predicted_image_embedding

            # remove the variance
            __UpperCamelCase  :Optional[int]    						=	noise_pred.split(
                scaled_model_input.shape[2]				,				dim=2)  # batch_size, num_embeddings, embedding_dim

            if do_classifier_free_guidance is not None:
                __UpperCamelCase  :Optional[int]    						=	noise_pred.chunk(2)
                __UpperCamelCase  :str    						=	noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)

            __UpperCamelCase  :str    						=	self.scheduler.step(
                __lowercase				,				timestep=__lowercase				,				sample=__lowercase				,				).prev_sample

        if output_type == "latent":
            return ShapEPipelineOutput(images=__lowercase)

        __UpperCamelCase  :Union[str, Any]    						=	[]
        for i, latent in enumerate(__lowercase):
            print()
            __UpperCamelCase  :Optional[int]    						=	self.renderer.decode(
                latent[None, :]				,				__lowercase				,				size=__lowercase				,				ray_batch_size=4_096				,				n_coarse_samples=64				,				n_fine_samples=128				,				)

            images.append(__lowercase)

        __UpperCamelCase  :Optional[Any]    						=	torch.stack(__lowercase)

        if output_type not in ["np", "pil"]:
            raise ValueError(f"""Only the output types `pil` and `np` are supported not output_type={output_type}""")

        __UpperCamelCase  :Dict    						=	images.cpu().numpy()

        if output_type == "pil":
            __UpperCamelCase  :Union[str, Any]    						=	[self.numpy_to_pil(__lowercase) for image in images]

        # Offload last model to CPU
        if hasattr(self				,				'''final_offload_hook''') and self.final_offload_hook is not None:
            self.final_offload_hook.offload()

        if not return_dict:
            return (images,)

        return ShapEPipelineOutput(images=__lowercase)






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







    def     UpperCamelCase__     (			self) ->			List[Any]:
        __UpperCamelCase  :int    						=	ids_tensor([self.batch_size, self.seq_length]				,				self.vocab_size)

        __UpperCamelCase  :int    						=	None
        if self.use_input_mask:
            __UpperCamelCase  :Optional[Any]    						=	random_attention_mask([self.batch_size, self.seq_length])

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

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

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







    def     UpperCamelCase__     (			self) ->			Optional[int]:
        __UpperCamelCase  :int    						=	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=__lowercase				,				esmfold_config={'''trunk''': {'''num_blocks''': 2}, '''fp16_esm''': False}				,				)
        return config







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase) ->			Union[str, Any]:
        __UpperCamelCase  :int    						=	EsmForProteinFolding(config=__lowercase).float()
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase  :Tuple    						=	model(__lowercase				,				attention_mask=__lowercase)
        __UpperCamelCase  :Any    						=	model(__lowercase)
        __UpperCamelCase  :Optional[Any]    						=	model(__lowercase)

        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     UpperCamelCase__     (			self) ->			Optional[Any]:
        __UpperCamelCase  :Dict    						=	self.prepare_config_and_inputs()
        (
            (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, 
        )  :List[str]    						=	config_and_inputs
        __UpperCamelCase  :Tuple    						=	{'''input_ids''': input_ids, '''attention_mask''': input_mask}
        return config, inputs_dict


@require_torch
class 			lowerCamelCase_  (     UpperCAmelCase_		,			UpperCAmelCase_		,			unittest.TestCase				):






    '''simple docstring'''






    a__			:		Optional[int]      			=   False

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







    def     UpperCamelCase__     (			self) ->			Tuple:
        __UpperCamelCase  :Dict    						=	EsmFoldModelTester(self)
        __UpperCamelCase  :Dict    						=	ConfigTester(self				,				config_class=__lowercase				,				hidden_size=37)







    def     UpperCamelCase__     (			self) ->			List[Any]:
        self.config_tester.run_common_tests()







    def     UpperCamelCase__     (			self) ->			List[str]:
        __UpperCamelCase  :Any    						=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowercase)







    @unittest.skip('''Does not support attention outputs''')
    def     UpperCamelCase__     (			self) ->			Any:
        pass







    @unittest.skip
    def     UpperCamelCase__     (			self) ->			Any:
        pass







    @unittest.skip('''Esm does not support embedding resizing''')
    def     UpperCamelCase__     (			self) ->			Union[str, Any]:
        pass







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







    @unittest.skip('''ESMFold does not support passing input embeds!''')
    def     UpperCamelCase__     (			self) ->			List[Any]:
        pass







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







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







    @unittest.skip('''ESMFold does not support head pruning.''')
    def     UpperCamelCase__     (			self) ->			Dict:
        pass







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







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







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







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







    @unittest.skip('''ESMFold only has one output format.''')
    def     UpperCamelCase__     (			self) ->			Union[str, Any]:
        pass







    @unittest.skip('''This test doesn\'t work for ESMFold and doesn\'t test core functionality''')
    def     UpperCamelCase__     (			self) ->			List[Any]:
        pass







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







    @unittest.skip('''ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.''')
    def     UpperCamelCase__     (			self) ->			int:
        pass







    @unittest.skip('''ESMFold doesn\'t support torchscript compilation.''')
    def     UpperCamelCase__     (			self) ->			str:
        pass







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







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







    @unittest.skip('''ESMFold doesn\'t support data parallel.''')
    def     UpperCamelCase__     (			self) ->			Any:
        pass







    @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.''')
    def     UpperCamelCase__     (			self) ->			Dict:
        pass


@require_torch
class 			lowerCamelCase_  (     UpperCAmelCase_				):






    '''simple docstring'''






    @slow
    def     UpperCamelCase__     (			self) ->			Tuple:
        __UpperCamelCase  :Optional[Any]    						=	EsmForProteinFolding.from_pretrained('''facebook/esmfold_v1''').float()
        model.eval()
        __UpperCamelCase  :Tuple    						=	torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]])
        __UpperCamelCase  :List[Any]    						=	model(__lowercase)['''positions''']
        __UpperCamelCase  :Optional[int]    						=	torch.tensor([2.58_28, 0.79_93, -10.93_34]				,				dtype=torch.floataa)
        self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]				,				__lowercase				,				atol=1E-4))





