



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

import numpy as np
import PIL
import torch
from transformers import CLIPTextModelWithProjection, CLIPTokenizer

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


UpperCAmelCase__			:   str      			=    logging.get_logger(__name__)  # pylint: disable=invalid-name

UpperCAmelCase__			:   Optional[Any]      			=    '\n    Examples:\n        ```py\n        >>> import torch\n        >>> from diffusers import DiffusionPipeline\n        >>> from diffusers.utils import export_to_gif\n\n        >>> device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n        >>> repo = \"openai/shap-e\"\n        >>> pipe = DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.float16)\n        >>> pipe = pipe.to(device)\n\n        >>> guidance_scale = 15.0\n        >>> prompt = \"a shark\"\n\n        >>> images = pipe(\n        ...     prompt,\n        ...     guidance_scale=guidance_scale,\n        ...     num_inference_steps=64,\n        ...     frame_size=256,\n        ... ).images\n\n        >>> gif_path = export_to_gif(images[0], \"shark_3d.gif\")\n        ```\n'





@dataclass
class 							lowerCAmelCase_		(		_snake_case     ):
		SCREAMING_SNAKE_CASE_       :			Optional[Any]					       =	4_2





class 							lowerCAmelCase_		(		_snake_case     ):







		def __init__(    self  :   int       ,		UpperCAmelCase_  :   PriorTransformer       ,		UpperCAmelCase_  :   CLIPTextModelWithProjection       ,		UpperCAmelCase_  :   CLIPTokenizer       ,		UpperCAmelCase_  :   HeunDiscreteScheduler       ,		UpperCAmelCase_  :   ShapERenderer       ,		)       ->		Union[str, Any]:


						'''simple docstring'''







						super().__init__()

						self.register_modules(
						    prior=lowerCAmelCase__       ,		text_encoder=lowerCAmelCase__       ,		tokenizer=lowerCAmelCase__       ,		scheduler=lowerCAmelCase__       ,		renderer=lowerCAmelCase__       ,		)







		def 					a_  (    self  :   Union[str, Any]       ,		UpperCAmelCase_  :   Any       ,		UpperCAmelCase_  :   str       ,		UpperCAmelCase_  :   Dict       ,		UpperCAmelCase_  :   Tuple       ,		UpperCAmelCase_  :   Union[str, Any]       ,		UpperCAmelCase_  :   Optional[int]					)       ->		List[Any]:


						'''simple docstring'''







						if latents is None:
										_UpperCAmelCase						:  Dict     					=	randn_tensor(lowerCAmelCase__       ,		generator=lowerCAmelCase__       ,		device=lowerCAmelCase__       ,		dtype=lowerCAmelCase__					)
						else:
										if latents.shape != shape:
														raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {shape}'''					)
										_UpperCAmelCase						:  Optional[Any]     					=	latents.to(lowerCAmelCase__					)

						_UpperCAmelCase						:  Optional[Any]     					=	latents * scheduler.init_noise_sigma
						return latents







		def 					a_  (    self  :   Union[str, Any]       ,		UpperCAmelCase_  :   Optional[int]=0					)       ->		Dict:


						'''simple docstring'''







						if is_accelerate_available():
										from accelerate import cpu_offload
						else:
										raise ImportError('''Please install accelerate via `pip install accelerate`'''					)

						_UpperCAmelCase						:  List[Any]     					=	torch.device(F'''cuda:{gpu_id}'''					)

						_UpperCAmelCase						:  List[str]     					=	[self.text_encoder, self.prior]
						for cpu_offloaded_model in models:
										if cpu_offloaded_model is not None:
														cpu_offload(lowerCAmelCase__       ,		lowerCAmelCase__					)







		def 					a_  (    self  :   List[str]       ,		UpperCAmelCase_  :   List[str]=0					)       ->		List[str]:


						'''simple docstring'''







						if is_accelerate_available() and is_accelerate_version('''>='''       ,		'''0.17.0.dev0'''					):
										from accelerate import cpu_offload_with_hook
						else:
										raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'''					)

						_UpperCAmelCase						:  Dict     					=	torch.device(F'''cuda:{gpu_id}'''					)

						if self.device.type != "cpu":
										self.to('''cpu'''       ,		silence_dtype_warnings=lowerCAmelCase__					)
										torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

						_UpperCAmelCase						:  int     					=	None
						for cpu_offloaded_model in [self.text_encoder, self.prior, self.renderer]:
										_UpperCAmelCase					, _UpperCAmelCase						:  List[Any]     					=	cpu_offload_with_hook(lowerCAmelCase__       ,		lowerCAmelCase__       ,		prev_module_hook=lowerCAmelCase__					)

						if self.safety_checker is not None:
										_UpperCAmelCase					, _UpperCAmelCase						:  List[str]     					=	cpu_offload_with_hook(self.safety_checker       ,		lowerCAmelCase__       ,		prev_module_hook=lowerCAmelCase__					)

						# We'll offload the last model manually.
						_UpperCAmelCase						:  List[Any]     					=	hook







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


						'''simple docstring'''







						if self.device != torch.device('''meta'''					) or not hasattr(self.text_encoder       ,		'''_hf_hook'''					):
										return self.device
						for module in self.text_encoder.modules():
										if (
										    hasattr(lowerCAmelCase__       ,		'''_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 					a_  (    self  :   Any       ,		UpperCAmelCase_  :   int       ,		UpperCAmelCase_  :   Any       ,		UpperCAmelCase_  :   Dict       ,		UpperCAmelCase_  :   Any       ,		)       ->		Tuple:


						'''simple docstring'''







						len(lowerCAmelCase__					) if isinstance(lowerCAmelCase__       ,		lowerCAmelCase__					) else 1

						# YiYi Notes: set pad_token_id to be 0, not sure why I can't set in the config file
						_UpperCAmelCase						:  List[Any]     					=	0
						# get prompt text embeddings
						_UpperCAmelCase						:  Dict     					=	self.tokenizer(
						    lowerCAmelCase__       ,		padding='''max_length'''       ,		max_length=self.tokenizer.model_max_length       ,		truncation=lowerCAmelCase__       ,		return_tensors='''pt'''       ,		)
						_UpperCAmelCase						:  List[Any]     					=	text_inputs.input_ids
						_UpperCAmelCase						:  str     					=	self.tokenizer(lowerCAmelCase__       ,		padding='''longest'''       ,		return_tensors='''pt'''					).input_ids

						if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(lowerCAmelCase__       ,		lowerCAmelCase__					):
										_UpperCAmelCase						:  Any     					=	self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]					)
										logger.warning(
										    '''The following part of your input was truncated because CLIP can only handle sequences up to'''
										    F''' {self.tokenizer.model_max_length} tokens: {removed_text}'''					)

						_UpperCAmelCase						:  str     					=	self.text_encoder(text_input_ids.to(lowerCAmelCase__					)					)
						_UpperCAmelCase						:  Any     					=	text_encoder_output.text_embeds

						_UpperCAmelCase						:  Dict     					=	prompt_embeds.repeat_interleave(lowerCAmelCase__       ,		dim=0					)
						# in Shap-E it normalize the prompt_embeds and then later rescale it
						_UpperCAmelCase						:  Any     					=	prompt_embeds / torch.linalg.norm(lowerCAmelCase__       ,		dim=-1       ,		keepdim=lowerCAmelCase__					)

						if do_classifier_free_guidance:
										_UpperCAmelCase						:  Optional[int]     					=	torch.zeros_like(lowerCAmelCase__					)

										# 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_prompt_embeds, prompt_embeds]					)

						# Rescale the features to have unit variance
						_UpperCAmelCase						:  int     					=	math.sqrt(prompt_embeds.shape[1]					) * prompt_embeds

						return prompt_embeds







		@torch.no_grad()
		@replace_example_docstring(lowerCAmelCase__					)
		def __call__(    self  :   Optional[int]       ,		UpperCAmelCase_  :   str       ,		UpperCAmelCase_  :   int = 1       ,		UpperCAmelCase_  :   int = 25       ,		UpperCAmelCase_  :   Optional[Union[torch.Generator, List[torch.Generator]]] = None       ,		UpperCAmelCase_  :   Optional[torch.FloatTensor] = None       ,		UpperCAmelCase_  :   float = 4.0       ,		UpperCAmelCase_  :   int = 64       ,		UpperCAmelCase_  :   Optional[str] = "pil"       ,		UpperCAmelCase_  :   bool = True       ,		)       ->		Union[str, Any]:


						'''simple docstring'''







						if isinstance(lowerCAmelCase__       ,		lowerCAmelCase__					):
										_UpperCAmelCase						:  Optional[Any]     					=	1
						elif isinstance(lowerCAmelCase__       ,		lowerCAmelCase__					):
										_UpperCAmelCase						:  Any     					=	len(lowerCAmelCase__					)
						else:
										raise ValueError(F'''`prompt` has to be of type `str` or `list` but is {type(lowerCAmelCase__					)}'''					)

						_UpperCAmelCase						:  List[Any]     					=	self._execution_device

						_UpperCAmelCase						:  Optional[Any]     					=	batch_size * num_images_per_prompt

						_UpperCAmelCase						:  List[str]     					=	guidance_scale > 1.0
						_UpperCAmelCase						:  Dict     					=	self._encode_prompt(lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__					)

						# prior

						self.scheduler.set_timesteps(lowerCAmelCase__       ,		device=lowerCAmelCase__					)
						_UpperCAmelCase						:  Union[str, Any]     					=	self.scheduler.timesteps

						_UpperCAmelCase						:  Any     					=	self.prior.config.num_embeddings
						_UpperCAmelCase						:  Dict     					=	self.prior.config.embedding_dim

						_UpperCAmelCase						:  List[str]     					=	self.prepare_latents(
						    (batch_size, num_embeddings * embedding_dim)       ,		prompt_embeds.dtype       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		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						:  int     					=	latents.reshape(latents.shape[0]       ,		lowerCAmelCase__       ,		lowerCAmelCase__					)

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

										_UpperCAmelCase						:  Tuple     					=	self.prior(
										    lowerCAmelCase__       ,		timestep=lowerCAmelCase__       ,		proj_embedding=lowerCAmelCase__       ,		).predicted_image_embedding

										# remove the variance
										_UpperCAmelCase					, _UpperCAmelCase						:  List[str]     					=	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					, _UpperCAmelCase						:  Tuple     					=	noise_pred.chunk(2					)
														_UpperCAmelCase						:  Optional[int]     					=	noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)

										_UpperCAmelCase						:  Tuple     					=	self.scheduler.step(
										    lowerCAmelCase__       ,		timestep=lowerCAmelCase__       ,		sample=lowerCAmelCase__       ,		).prev_sample

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

						_UpperCAmelCase						:  Tuple     					=	[]
						for i, latent in enumerate(lowerCAmelCase__					):
										_UpperCAmelCase						:  List[Any]     					=	self.renderer.decode(
										    latent[None, :]       ,		lowerCAmelCase__       ,		size=lowerCAmelCase__       ,		ray_batch_size=4096       ,		n_coarse_samples=64       ,		n_fine_samples=128       ,		)
										images.append(lowerCAmelCase__					)

						_UpperCAmelCase						:  List[Any]     					=	torch.stack(lowerCAmelCase__					)

						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						:  List[Any]     					=	images.cpu().numpy()

						if output_type == "pil":
										_UpperCAmelCase						:  Union[str, Any]     					=	[self.numpy_to_pil(lowerCAmelCase__					) 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=lowerCAmelCase__					)



import json
import os
from collections.abc import Mapping
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from .imports import is_safetensors_available
def   _lowerCAmelCase(						__A ,							__A ,							__A ,							__A=None      ):
  UpperCAmelCase          =	None
  # Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
  if str(weight.dtype      ) == "torch.bfloat16":
    # Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
    UpperCAmelCase          =	weight.view(torch.intaa      )
    UpperCAmelCase          =	"bfloat16"
  UpperCAmelCase          =	weight.cpu().numpy()
  UpperCAmelCase          =	os.path.join(__A ,							F"{weight_name}.dat"      )
  if index is not None:
    if dtype is None:
      UpperCAmelCase          =	str(array.dtype      )
    UpperCAmelCase          =	{"dtype": dtype, "shape": list(array.shape      )}
  if array.ndim == 0:
    UpperCAmelCase          =	array[None]
  UpperCAmelCase          =	np.memmap(__A ,							dtype=array.dtype ,							mode="w+" ,							shape=array.shape      )
  UpperCAmelCase          =	array[:]
  file_array.flush()
  return index
def   _lowerCAmelCase(						__A ,							__A      ):
  UpperCAmelCase          =	tuple(weight_info["shape"]      )
  if shape == ():
    # NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
    UpperCAmelCase          =	(1,)

  UpperCAmelCase          =	weight_info["dtype"]
  if dtype == "bfloat16":
    # NumPy does not support bfloat16 so this was saved as a int16
    UpperCAmelCase          =	"int16"

  UpperCAmelCase          =	np.memmap(__A ,							dtype=__A ,							shape=__A ,							mode="r"      )

  if len(weight_info["shape"]      ) == 0:
    UpperCAmelCase          =	weight[0]
  UpperCAmelCase          =	torch.tensor(__A      )
  if weight_info["dtype"] == "bfloat16":
    UpperCAmelCase          =	weight.view(torch.bfloataa      )

  return weight
def   _lowerCAmelCase(						__A ,							__A      ):
  if index is None or len(__A      ) == 0:
    # Nothing to save
    return

  UpperCAmelCase          =	os.path.join(__A ,							"index.json"      )
  if os.path.isfile(__A      ):
    with open(__A ,							"r" ,							encoding="utf-8"      ) as f:
      UpperCAmelCase          =	json.load(__A      )
  else:
    UpperCAmelCase          =	{}
  current_index.update(__A      )

  with open(__A ,							"w" ,							encoding="utf-8"      ) as f:
    json.dump(__A ,							__A ,							indent=2      )
def   _lowerCAmelCase(						__A ,							__A      ):
  os.makedirs(__A ,							exist_ok=__A      )
  UpperCAmelCase          =	{}
  for name, parameter in state_dict.items():
    UpperCAmelCase          =	offload_weight(__A ,							__A ,							__A ,							index=__A      )

  # Update index
  save_offload_index(__A ,							__A      )






class     __magic_name__     (							_snake_case     ):


 def __init__(    self		:			Tuple      ,    lowerCAmelCase__		:			Mapping      ,    lowerCAmelCase__		:			str     )     ->   Tuple:
   UpperCAmelCase          =	dataset
   UpperCAmelCase          =	prefix


 def __getitem__(    self		:			Dict      ,    lowerCAmelCase__		:			Any     )     ->   Optional[int]:
   return self.dataset[f"{self.prefix}{key}"]


 def __iter__(    self		:			Optional[int]     )     ->   Any:
   return iter([key for key in self.dataset if key.startswith(self.prefix     )]     )


 def __len__(    self		:			Union[str, Any]     )     ->   int:
   return len(self.dataset     )






class     __magic_name__     (							_snake_case     ):


 def __init__(    self		:			Optional[int]      ,    lowerCAmelCase__		:			Dict[str, torch.Tensor] = None      ,    lowerCAmelCase__		:			Optional[Union[str, os.PathLike]] = None      ,    lowerCAmelCase__		:			Mapping = None      ,    lowerCAmelCase__		:			Dict=None      ,    )     ->   Dict:
   if state_dict is None and save_folder is None:
     raise ValueError("Need either a `state_dict` or a `save_folder` containing offloaded weights."     )

   UpperCAmelCase          =	{} if state_dict is None else state_dict
   UpperCAmelCase          =	save_folder
   if index is None and save_folder is not None:
     with open(os.path.join(lowerCAmelCase__      ,    "index.json"     )     ) as f:
       UpperCAmelCase          =	json.load(lowerCAmelCase__     )
   UpperCAmelCase          =	{} if index is None else index
   UpperCAmelCase          =	list(self.state_dict.keys()     )
   self.all_keys.extend([key for key in self.index if key not in self.all_keys]     )
   UpperCAmelCase          =	device


 def __getitem__(    self		:			str      ,    lowerCAmelCase__		:			str     )     ->   Dict:
   # State dict gets priority
   if key in self.state_dict:
     return self.state_dict[key]
   UpperCAmelCase          =	self.index[key]
   if weight_info.get("safetensors_file"     ) is not None:
     if not is_safetensors_available():
       raise ImportError("These offloaded weights require the use of safetensors: `pip install safetensors`."     )

     from safetensors import safe_open

     UpperCAmelCase          =	"cpu" if self.device is None else self.device
     with safe_open(weight_info["safetensors_file"]      ,    framework="pt"      ,    device=lowerCAmelCase__     ) as f:
       UpperCAmelCase          =	f.get_tensor(weight_info.get("weight_name"      ,    lowerCAmelCase__     )     )

     if "dtype" in weight_info:
       return tensor.to(getattr(lowerCAmelCase__      ,    weight_info["dtype"]     )     )
     else:
       return tensor

   UpperCAmelCase          =	os.path.join(self.save_folder      ,    f"{key}.dat"     )
   return load_offloaded_weight(lowerCAmelCase__      ,    lowerCAmelCase__     )


 def __iter__(    self		:			str     )     ->   Dict:
   return iter(self.all_keys     )


 def __len__(    self		:			Optional[Any]     )     ->   List[Any]:
   return len(self.all_keys     )

def   _lowerCAmelCase(						__A ,							__A      ):
  UpperCAmelCase          =	{}
  for module_name in submodule_names:
    # We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
    # submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
    result.update(
        {
            key: param
            for key, param in state_dict.items()
            if key == module_name or key.startswith(module_name + "."      )
        }      )
  return result



