




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


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

_lowerCamelCase    				=		'\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     a  (     _A    ):

     '''simple docstring'''





     lowerCAmelCase       : Union[List[List[PIL.Image.Image]], List[List[np.ndarray]]]



class     a  (     _A    ):

     '''simple docstring'''





     def __init__(			self     :    Dict      ,					__snake_case     :    PriorTransformer      ,					__snake_case     :    CLIPTextModelWithProjection      ,					__snake_case     :    CLIPTokenizer      ,					__snake_case     :    HeunDiscreteScheduler      ,					__snake_case     :    ShapERenderer      ,					):
         super().__init__()

         self.register_modules(
             prior=__snake_case      ,					text_encoder=__snake_case      ,					tokenizer=__snake_case      ,					scheduler=__snake_case      ,					renderer=__snake_case      ,					)
     def 						lowerCamelCase_     (			self     :    List[str]      ,					__snake_case     :    Dict      ,					__snake_case     :    Tuple      ,					__snake_case     :    Tuple      ,					__snake_case     :    str      ,					__snake_case     :    Optional[int]      ,					__snake_case     :    List[Any]       ):
         if latents is None:
             UpperCAmelCase_											=  randn_tensor(__snake_case      ,					generator=__snake_case      ,					device=__snake_case      ,					dtype=__snake_case       )
         else:
             if latents.shape != shape:
                 raise ValueError(F'Unexpected latents shape, got {latents.shape}, expected {shape}'       )
             UpperCAmelCase_											=  latents.to(__snake_case       )

         UpperCAmelCase_											=  latents * scheduler.init_noise_sigma
         return latents
     def 						lowerCamelCase_     (			self     :    Optional[Any]      ,					__snake_case     :    Optional[Any]=0       ):
         if is_accelerate_available():
             from accelerate import cpu_offload
         else:
             raise ImportError('''Please install accelerate via `pip install accelerate`'''       )

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

         UpperCAmelCase_											=  [self.text_encoder, self.prior]
         for cpu_offloaded_model in models:
             if cpu_offloaded_model is not None:
                 cpu_offload(__snake_case      ,					__snake_case       )
     def 						lowerCamelCase_     (			self     :    Optional[Any]      ,					__snake_case     :    Dict=0       ):
         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_											=  torch.device(F'cuda:{gpu_id}'       )

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

         UpperCAmelCase_											=  None
         for cpu_offloaded_model in [self.text_encoder, self.prior, self.renderer]:
             UpperCAmelCase_				,							UpperCAmelCase_											=  cpu_offload_with_hook(__snake_case      ,					__snake_case      ,					prev_module_hook=__snake_case       )

         if self.safety_checker is not None:
             UpperCAmelCase_				,							UpperCAmelCase_											=  cpu_offload_with_hook(self.safety_checker      ,					__snake_case      ,					prev_module_hook=__snake_case       )

         # We'll offload the last model manually.
         UpperCAmelCase_											=  hook
     @property
     def 						lowerCamelCase_     (			self     :    Optional[int]       ):
         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(__snake_case      ,					'''_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 						lowerCamelCase_     (			self     :    Optional[int]      ,					__snake_case     :    str      ,					__snake_case     :    int      ,					__snake_case     :    Union[str, Any]      ,					__snake_case     :    str      ,					):
         len(__snake_case       ) if isinstance(__snake_case      ,					__snake_case       ) else 1

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

         if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(__snake_case      ,					__snake_case       ):
             UpperCAmelCase_											=  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_											=  self.text_encoder(text_input_ids.to(__snake_case       )       )
         UpperCAmelCase_											=  text_encoder_output.text_embeds

         UpperCAmelCase_											=  prompt_embeds.repeat_interleave(__snake_case      ,					dim=0       )
         # in Shap-E it normalize the prompt_embeds and then later rescale it
         UpperCAmelCase_											=  prompt_embeds / torch.linalg.norm(__snake_case      ,					dim=-1      ,					keepdim=__snake_case       )

         if do_classifier_free_guidance:
             UpperCAmelCase_											=  torch.zeros_like(__snake_case       )

             # 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_											=  torch.cat([negative_prompt_embeds, prompt_embeds]       )

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

         return prompt_embeds




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

         if isinstance(__snake_case      ,					__snake_case       ):
             UpperCAmelCase_											=  1
         elif isinstance(__snake_case      ,					__snake_case       ):
             UpperCAmelCase_											=  len(__snake_case       )
         else:
             raise ValueError(F'`prompt` has to be of type `str` or `list` but is {type(__snake_case       )}'       )

         UpperCAmelCase_											=  self._execution_device

         UpperCAmelCase_											=  batch_size * num_images_per_prompt

         UpperCAmelCase_											=  guidance_scale > 1.0
         UpperCAmelCase_											=  self._encode_prompt(__snake_case      ,					__snake_case      ,					__snake_case      ,					__snake_case       )

         # prior

         self.scheduler.set_timesteps(__snake_case      ,					device=__snake_case       )
         UpperCAmelCase_											=  self.scheduler.timesteps

         UpperCAmelCase_											=  self.prior.config.num_embeddings
         UpperCAmelCase_											=  self.prior.config.embedding_dim

         UpperCAmelCase_											=  self.prepare_latents(
             (batch_size, num_embeddings * embedding_dim)      ,					prompt_embeds.dtype      ,					__snake_case      ,					__snake_case      ,					__snake_case      ,					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_											=  latents.reshape(latents.shape[0]      ,					__snake_case      ,					__snake_case       )

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

             UpperCAmelCase_											=  self.prior(
                 __snake_case      ,					timestep=__snake_case      ,					proj_embedding=__snake_case      ,					).predicted_image_embedding

             # remove the variance
             UpperCAmelCase_				,							UpperCAmelCase_											=  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_											=  noise_pred.chunk(2       )
                 UpperCAmelCase_											=  noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)

             UpperCAmelCase_											=  self.scheduler.step(
                 __snake_case      ,					timestep=__snake_case      ,					sample=__snake_case      ,					).prev_sample

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

         UpperCAmelCase_											=  []
         for i, latent in enumerate(__snake_case       ):
             UpperCAmelCase_											=  self.renderer.decode(
                 latent[None, :]      ,					__snake_case      ,					size=__snake_case      ,					ray_batch_size=40_96      ,					n_coarse_samples=64      ,					n_fine_samples=1_28      ,					)
             images.append(__snake_case       )

         UpperCAmelCase_											=  torch.stack(__snake_case       )

         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_											=  images.cpu().numpy()

         if output_type == "pil":
             UpperCAmelCase_											=  [self.numpy_to_pil(__snake_case       ) 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=__snake_case       )











import logging
import os
import sys
import threading
from logging import (
    CRITICAL,  # NOQA
    DEBUG,  # NOQA
    ERROR,  # NOQA
    FATAL,  # NOQA
    INFO,  # NOQA
    NOTSET,  # NOQA
    WARN,  # NOQA
    WARNING,  # NOQA
)
from typing import Optional

from tqdm import auto as tqdm_lib


_lowerCamelCase    				=		threading.Lock()
_lowerCamelCase    				=		None

_lowerCamelCase    				=		{
    'debug': logging.DEBUG,
    'info': logging.INFO,
    'warning': logging.WARNING,
    'error': logging.ERROR,
    'critical': logging.CRITICAL,
}

_lowerCamelCase    				=		logging.WARNING

_lowerCamelCase    				=		True




def       SCREAMING_SNAKE_CASE			( )						->							List[str]:
    UpperCAmelCase_											=  os.getenv('''DIFFUSERS_VERBOSITY''' ,  __UpperCamelCase   )
    if env_level_str:
        if env_level_str in log_levels:
            return log_levels[env_level_str]
        else:
            logging.getLogger().warning(
                f'Unknown option DIFFUSERS_VERBOSITY={env_level_str}, '
                f'has to be one of: { ", ".join(log_levels.keys()   ) }'   )
    return _default_log_level




def       SCREAMING_SNAKE_CASE			( )						->							str:
    return __name__.split('''.'''   )[0]




def       SCREAMING_SNAKE_CASE			( )						->							logging.Logger:
    return logging.getLogger(_get_library_name()   )




def       SCREAMING_SNAKE_CASE			( )						->							None:
    global _default_handler

    with _lock:
        if _default_handler:
            # This library has already configured the library root logger.
            return
        UpperCAmelCase_											=  logging.StreamHandler()  # Set sys.stderr as stream.
        UpperCAmelCase_											=  sys.stderr.flush

        # Apply our default configuration to the library root logger.
        UpperCAmelCase_											=  _get_library_root_logger()
        library_root_logger.addHandler(_default_handler   )
        library_root_logger.setLevel(_get_default_logging_level()   )
        UpperCAmelCase_											=  False




def       SCREAMING_SNAKE_CASE			( )						->							None:
    global _default_handler

    with _lock:
        if not _default_handler:
            return

        UpperCAmelCase_											=  _get_library_root_logger()
        library_root_logger.removeHandler(_default_handler   )
        library_root_logger.setLevel(logging.NOTSET   )
        UpperCAmelCase_											=  None




def       SCREAMING_SNAKE_CASE			( )						->							Dict:
    return log_levels




def       SCREAMING_SNAKE_CASE			( __UpperCamelCase    :		Optional[str] = None   )						->							logging.Logger:

    if name is None:
        UpperCAmelCase_											=  _get_library_name()

    _configure_library_root_logger()
    return logging.getLogger(__UpperCamelCase   )




def       SCREAMING_SNAKE_CASE			( )						->							int:
    _configure_library_root_logger()
    return _get_library_root_logger().getEffectiveLevel()




def       SCREAMING_SNAKE_CASE			( __UpperCamelCase    :		int   )						->							None:
    _configure_library_root_logger()
    _get_library_root_logger().setLevel(__UpperCamelCase   )




def       SCREAMING_SNAKE_CASE			( )						->							Optional[Any]:
    return set_verbosity(__UpperCamelCase   )




def       SCREAMING_SNAKE_CASE			( )						->							Union[str, Any]:
    return set_verbosity(__UpperCamelCase   )




def       SCREAMING_SNAKE_CASE			( )						->							Tuple:
    return set_verbosity(__UpperCamelCase   )




def       SCREAMING_SNAKE_CASE			( )						->							int:
    return set_verbosity(__UpperCamelCase   )




def       SCREAMING_SNAKE_CASE			( )						->							None:
    _configure_library_root_logger()

    assert _default_handler is not None
    _get_library_root_logger().removeHandler(_default_handler   )




def       SCREAMING_SNAKE_CASE			( )						->							None:
    _configure_library_root_logger()

    assert _default_handler is not None
    _get_library_root_logger().addHandler(_default_handler   )




def       SCREAMING_SNAKE_CASE			( __UpperCamelCase    :		logging.Handler   )						->							None:
    _configure_library_root_logger()

    assert handler is not None
    _get_library_root_logger().addHandler(__UpperCamelCase   )




def       SCREAMING_SNAKE_CASE			( __UpperCamelCase    :		logging.Handler   )						->							None:
    _configure_library_root_logger()

    assert handler is not None and handler not in _get_library_root_logger().handlers
    _get_library_root_logger().removeHandler(__UpperCamelCase   )




def       SCREAMING_SNAKE_CASE			( )						->							None:
    _configure_library_root_logger()
    UpperCAmelCase_											=  False




def       SCREAMING_SNAKE_CASE			( )						->							None:
    _configure_library_root_logger()
    UpperCAmelCase_											=  True




def       SCREAMING_SNAKE_CASE			( )						->							None:
    UpperCAmelCase_											=  _get_library_root_logger().handlers

    for handler in handlers:
        UpperCAmelCase_											=  logging.Formatter('''[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s'''   )
        handler.setFormatter(__UpperCamelCase   )




def       SCREAMING_SNAKE_CASE			( )						->							None:
    UpperCAmelCase_											=  _get_library_root_logger().handlers

    for handler in handlers:
        handler.setFormatter(__UpperCamelCase   )




def       SCREAMING_SNAKE_CASE			( self    :		Tuple ,  *__UpperCamelCase    :		Any ,  **__UpperCamelCase    :		Any   )						->							Union[str, Any]:
    UpperCAmelCase_											=  os.getenv('''DIFFUSERS_NO_ADVISORY_WARNINGS''' ,  __UpperCamelCase   )
    if no_advisory_warnings:
        return
    self.warning(*__UpperCamelCase ,  **__UpperCamelCase   )


_lowerCamelCase    				=		warning_advice



class     a  :

     '''simple docstring'''





     def __init__(			self     :    List[Any]      ,					*__snake_case     :    Tuple      ,					**__snake_case     :    str       ):  # pylint: disable=unused-argument
         UpperCAmelCase_											=  args[0] if args else None
     def __iter__(			self     :    Union[str, Any]       ):
         return iter(self._iterator       )
     def __getattr__(			self     :    Tuple      ,					__snake_case     :    str       ):

         def empty_fn(*__snake_case     :    Optional[int]      ,					**__snake_case     :    int       ):  # pylint: disable=unused-argument
             return

         return empty_fn
     def __enter__(			self     :    Optional[Any]       ):
         return self




     def __exit__(			self     :    Union[str, Any]      ,					__snake_case     :    List[Any]      ,					__snake_case     :    List[str]      ,					__snake_case     :    Dict       ):
         return



class     a  :

     '''simple docstring'''





     def __call__(			self     :    Union[str, Any]      ,					*__snake_case     :    List[str]      ,					**__snake_case     :    Dict       ):
         if _tqdm_active:
             return tqdm_lib.tqdm(*__snake_case      ,					**__snake_case       )
         else:
             return EmptyTqdm(*__snake_case      ,					**__snake_case       )
     def 						lowerCamelCase_     (			self     :    int      ,					*__snake_case     :    Dict      ,					**__snake_case     :    List[str]       ):
         UpperCAmelCase_											=  None
         if _tqdm_active:
             return tqdm_lib.tqdm.set_lock(*__snake_case      ,					**__snake_case       )




     def 						lowerCamelCase_     (			self     :    str       ):
         if _tqdm_active:
             return tqdm_lib.tqdm.get_lock()


_lowerCamelCase    				=		_tqdm_cls()




def       SCREAMING_SNAKE_CASE			( )						->							bool:
    global _tqdm_active
    return bool(_tqdm_active   )




def       SCREAMING_SNAKE_CASE			( )						->							List[str]:
    global _tqdm_active
    UpperCAmelCase_											=  True




def       SCREAMING_SNAKE_CASE			( )						->							Optional[int]:
    global _tqdm_active
    UpperCAmelCase_											=  False





