





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


_lowerCAmelCase				:   Union[str, Any]						       =						logging.get_logger(__name__)  # pylint: disable=invalid-name

_lowerCAmelCase				:   int						       =						'''
    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        __magic_name__	(				lowerCamelCase__				):





     """simple docstring"""



     __UpperCamelCase            =  42




class        __magic_name__	(				lowerCamelCase__				):





     """simple docstring"""



     def __init__(				self							:Union[str, Any]      ,   snake_case							:PriorTransformer      ,   snake_case							:CLIPVisionModel      ,   snake_case							:CLIPImageProcessor      ,   snake_case							:HeunDiscreteScheduler      ,   snake_case							:ShapERenderer      ,   ):





        '''simple docstring'''




        super().__init__()

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



     def     SCREAMING_SNAKE_CASE		(				self							:Tuple      ,   snake_case							:Tuple      ,   snake_case							:str      ,   snake_case							:Any      ,   snake_case							:Optional[Any]      ,   snake_case							:Union[str, Any]      ,   snake_case							:Optional[Any]	):





        '''simple docstring'''




        if latents is None:
           A_       :					int        =      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}"	)
           A_       :					Tuple        =      latents.to(snake_case	)

        A_       :					Tuple        =      latents * scheduler.init_noise_sigma
        return latents



     def     SCREAMING_SNAKE_CASE		(				self							:Tuple      ,   snake_case							:Tuple=0	):





        '''simple docstring'''




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

        A_       :					Any        =      torch.device(f"cuda:{gpu_id}"	)

        A_       :					Any        =      [self.image_encoder, self.prior]
        for cpu_offloaded_model in models:
           if cpu_offloaded_model is not None:
              cpu_offload(snake_case      ,   snake_case	)



     @property
     def     SCREAMING_SNAKE_CASE		(				self							:Dict	):





        '''simple docstring'''




        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(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     SCREAMING_SNAKE_CASE		(				self							:Union[str, Any]      ,   snake_case							:Dict      ,   snake_case							:Union[str, Any]      ,   snake_case							:List[Any]      ,   snake_case							:Optional[int]      ,   ):





        '''simple docstring'''




        if isinstance(snake_case      ,   snake_case	) and isinstance(image[0]      ,   torch.Tensor	):
           A_       :					Union[str, Any]        =      torch.cat(snake_case      ,   axis=0	) if image[0].ndim == 4 else torch.stack(snake_case      ,   axis=0	)

        if not isinstance(snake_case      ,   torch.Tensor	):
           A_       :					Tuple        =      self.image_processor(snake_case      ,   return_tensors="pt"	).pixel_values[0].unsqueeze(0	)

        A_       :					Any        =      image.to(dtype=self.image_encoder.dtype      ,   device=snake_case	)

        A_       :					Any        =      self.image_encoder(snake_case	)["last_hidden_state"]
        A_       :					Tuple        =      image_embeds[:, 1:, :].contiguous()  # batch_size, dim, 256

        A_       :					Any        =      image_embeds.repeat_interleave(snake_case      ,   dim=0	)

        if do_classifier_free_guidance:
           A_       :					Optional[int]        =      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
           A_       :					int        =      torch.cat([negative_image_embeds, image_embeds]	)

        return image_embeds





     @torch.no_grad()
     @replace_example_docstring(snake_case	)
     def __call__(				self							:int      ,   snake_case							:Union[PIL.Image.Image, List[PIL.Image.Image]]      ,   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      ,   ):





        '''simple docstring'''




        if isinstance(snake_case      ,   PIL.Image.Image	):
           A_       :					str        =      1
        elif isinstance(snake_case      ,   torch.Tensor	):
           A_       :					Optional[Any]        =      image.shape[0]
        elif isinstance(snake_case      ,   snake_case	) and isinstance(image[0]      ,   (torch.Tensor, PIL.Image.Image)	):
           A_       :					List[Any]        =      len(snake_case	)
        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(snake_case	)}"	)

        A_       :					int        =      self._execution_device

        A_       :					Optional[Any]        =      batch_size * num_images_per_prompt

        A_       :					List[str]        =      guidance_scale > 1.0
        A_       :					str        =      self._encode_image(snake_case      ,   snake_case      ,   snake_case      ,   snake_case	)

        # prior

        self.scheduler.set_timesteps(snake_case      ,   device=snake_case	)
        A_       :					Tuple        =      self.scheduler.timesteps

        A_       :					int        =      self.prior.config.num_embeddings
        A_       :					Optional[Any]        =      self.prior.config.embedding_dim

        A_       :					Union[str, Any]        =      self.prepare_latents(
            (batch_size, num_embeddings * embedding_dim)      ,   image_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
        A_       :					Union[str, Any]        =      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
           A_       :					Any        =      torch.cat([latents] * 2	) if do_classifier_free_guidance else latents
           A_       :					Any        =      self.scheduler.scale_model_input(snake_case      ,   snake_case	)

           A_       :					List[Any]        =      self.prior(
               snake_case      ,   timestep=snake_case      ,   proj_embedding=snake_case      ,   ).predicted_image_embedding

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

           if do_classifier_free_guidance is not None:
              A_   ,    A_       :					int        =      noise_pred.chunk(2	)
              A_       :					str        =      noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)

           A_       :					Optional[int]        =      self.scheduler.step(
               snake_case      ,   timestep=snake_case      ,   sample=snake_case      ,   ).prev_sample

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

        A_       :					str        =      []
        for i, latent in enumerate(snake_case	):
           print()
           A_       :					List[str]        =      self.renderer.decode(
               latent[None, :]      ,   snake_case      ,   size=snake_case      ,   ray_batch_size=4_096      ,   n_coarse_samples=64      ,   n_fine_samples=128      ,   )

           images.append(snake_case	)

        A_       :					List[str]        =      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}"	)

        A_       :					Optional[int]        =      images.cpu().numpy()

        if output_type == "pil":
           A_       :					Tuple        =      [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 json
from typing import List, Optional, Tuple

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_bart import BartTokenizer


_lowerCAmelCase				:   Optional[Any]						       =						logging.get_logger(__name__)


_lowerCAmelCase				:   List[str]						       =						{'''vocab_file''': '''vocab.json''', '''merges_file''': '''merges.txt''', '''tokenizer_file''': '''tokenizer.json'''}

# See all BART models at https://huggingface.co/models?filter=bart
_lowerCAmelCase				:   int						       =						{
    '''vocab_file''': {
        '''facebook/bart-base''': '''https://huggingface.co/facebook/bart-base/resolve/main/vocab.json''',
        '''facebook/bart-large''': '''https://huggingface.co/facebook/bart-large/resolve/main/vocab.json''',
        '''facebook/bart-large-mnli''': '''https://huggingface.co/facebook/bart-large-mnli/resolve/main/vocab.json''',
        '''facebook/bart-large-cnn''': '''https://huggingface.co/facebook/bart-large-cnn/resolve/main/vocab.json''',
        '''facebook/bart-large-xsum''': '''https://huggingface.co/facebook/bart-large-xsum/resolve/main/vocab.json''',
        '''yjernite/bart_eli5''': '''https://huggingface.co/yjernite/bart_eli5/resolve/main/vocab.json''',
    },
    '''merges_file''': {
        '''facebook/bart-base''': '''https://huggingface.co/facebook/bart-base/resolve/main/merges.txt''',
        '''facebook/bart-large''': '''https://huggingface.co/facebook/bart-large/resolve/main/merges.txt''',
        '''facebook/bart-large-mnli''': '''https://huggingface.co/facebook/bart-large-mnli/resolve/main/merges.txt''',
        '''facebook/bart-large-cnn''': '''https://huggingface.co/facebook/bart-large-cnn/resolve/main/merges.txt''',
        '''facebook/bart-large-xsum''': '''https://huggingface.co/facebook/bart-large-xsum/resolve/main/merges.txt''',
        '''yjernite/bart_eli5''': '''https://huggingface.co/yjernite/bart_eli5/resolve/main/merges.txt''',
    },
    '''tokenizer_file''': {
        '''facebook/bart-base''': '''https://huggingface.co/facebook/bart-base/resolve/main/tokenizer.json''',
        '''facebook/bart-large''': '''https://huggingface.co/facebook/bart-large/resolve/main/tokenizer.json''',
        '''facebook/bart-large-mnli''': '''https://huggingface.co/facebook/bart-large-mnli/resolve/main/tokenizer.json''',
        '''facebook/bart-large-cnn''': '''https://huggingface.co/facebook/bart-large-cnn/resolve/main/tokenizer.json''',
        '''facebook/bart-large-xsum''': '''https://huggingface.co/facebook/bart-large-xsum/resolve/main/tokenizer.json''',
        '''yjernite/bart_eli5''': '''https://huggingface.co/yjernite/bart_eli5/resolve/main/tokenizer.json''',
    },
}

_lowerCAmelCase				:   Optional[int]						       =						{
    '''facebook/bart-base''': 1_024,
    '''facebook/bart-large''': 1_024,
    '''facebook/bart-large-mnli''': 1_024,
    '''facebook/bart-large-cnn''': 1_024,
    '''facebook/bart-large-xsum''': 1_024,
    '''yjernite/bart_eli5''': 1_024,
}


class        __magic_name__	(				lowerCamelCase__				):





     """simple docstring"""



     __UpperCamelCase            =  VOCAB_FILES_NAMES
     __UpperCamelCase            =  PRETRAINED_VOCAB_FILES_MAP
     __UpperCamelCase            =  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     __UpperCamelCase            =  ['''input_ids''', '''attention_mask''']
     __UpperCamelCase            =  BartTokenizer



     def __init__(				self							:List[str]      ,   snake_case							:Tuple=None      ,   snake_case							:Union[str, Any]=None      ,   snake_case							:Union[str, Any]=None      ,   snake_case							:List[Any]="replace"      ,   snake_case							:List[str]="<s>"      ,   snake_case							:Optional[int]="</s>"      ,   snake_case							:Union[str, Any]="</s>"      ,   snake_case							:Optional[Any]="<s>"      ,   snake_case							:List[Any]="<unk>"      ,   snake_case							:Optional[Any]="<pad>"      ,   snake_case							:Dict="<mask>"      ,   snake_case							:int=False      ,   snake_case							:List[Any]=True      ,   **snake_case							:Union[str, Any]      ,   ):





        '''simple docstring'''




        super().__init__(
            snake_case      ,   snake_case      ,   tokenizer_file=snake_case      ,   errors=snake_case      ,   bos_token=snake_case      ,   eos_token=snake_case      ,   sep_token=snake_case      ,   cls_token=snake_case      ,   unk_token=snake_case      ,   pad_token=snake_case      ,   mask_token=snake_case      ,   add_prefix_space=snake_case      ,   trim_offsets=snake_case      ,   **snake_case      ,   )

        A_       :					Optional[Any]        =      json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()	)
        if pre_tok_state.get("add_prefix_space"      ,   snake_case	) != add_prefix_space:
           A_       :					Dict        =      getattr(snake_case      ,   pre_tok_state.pop("type"	)	)
           A_       :					List[str]        =      add_prefix_space
           A_       :					int        =      pre_tok_class(**snake_case	)

        A_       :					Any        =      add_prefix_space

        # the pre_tokenizer is already updated in the GPT2TokenizerFast `__init__`
        A_       :					Tuple        =      "post_processor"
        A_       :					Union[str, Any]        =      getattr(self.backend_tokenizer      ,   snake_case      ,   snake_case	)
        if tokenizer_component_instance:
           A_       :					Any        =      json.loads(tokenizer_component_instance.__getstate__()	)

           # The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
           if "sep" in state:
              A_       :					List[Any]        =      tuple(state["sep"]	)
           if "cls" in state:
              A_       :					str        =      tuple(state["cls"]	)

           A_       :					int        =      False

           if state.get("add_prefix_space"      ,   snake_case	) != add_prefix_space:
              A_       :					List[Any]        =      add_prefix_space
              A_       :					Union[str, Any]        =      True

           if state.get("trim_offsets"      ,   snake_case	) != trim_offsets:
              A_       :					int        =      trim_offsets
              A_       :					str        =      True

           if changes_to_apply:
              A_       :					Tuple        =      getattr(snake_case      ,   state.pop("type"	)	)
              A_       :					Union[str, Any]        =      component_class(**snake_case	)
              setattr(self.backend_tokenizer      ,   snake_case      ,   snake_case	)



     @property
     def     SCREAMING_SNAKE_CASE		(				self							:List[Any]	):





        '''simple docstring'''




        if self._mask_token is None:
           if self.verbose:
              logger.error("Using mask_token, but it is not set yet."	)
           return None
        return str(self._mask_token	)



     @mask_token.setter
     def     SCREAMING_SNAKE_CASE		(				self							:Tuple      ,   snake_case							:Dict	):





        '''simple docstring'''




        A_       :					Optional[int]        =      AddedToken(snake_case      ,   lstrip=snake_case      ,   rstrip=snake_case	) if isinstance(snake_case      ,   snake_case	) else value
        A_       :					List[Any]        =      value



     def     SCREAMING_SNAKE_CASE		(				self							:Tuple      ,   *snake_case							:str      ,   **snake_case							:str	):





        '''simple docstring'''




        A_       :					Tuple        =      kwargs.get("is_split_into_words"      ,   snake_case	)

        if is_split_into_words and not self.add_prefix_space:
           raise ValueError(
               f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
               "to use it with pretokenized inputs."	)

        return super()._batch_encode_plus(*snake_case      ,   **snake_case	)



     def     SCREAMING_SNAKE_CASE		(				self							:Union[str, Any]      ,   *snake_case							:Any      ,   **snake_case							:Optional[Any]	):





        '''simple docstring'''




        A_       :					List[Any]        =      kwargs.get("is_split_into_words"      ,   snake_case	)

        if is_split_into_words and not self.add_prefix_space:
           raise ValueError(
               f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
               "to use it with pretokenized inputs."	)

        return super()._encode_plus(*snake_case      ,   **snake_case	)



     def     SCREAMING_SNAKE_CASE		(				self							:int      ,   snake_case							:str      ,   snake_case							:Optional[str] = None	):





        '''simple docstring'''




        A_       :					Any        =      self._tokenizer.model.save(snake_case      ,   name=snake_case	)
        return tuple(snake_case	)



     def     SCREAMING_SNAKE_CASE		(				self							:Dict      ,   snake_case							:List[str]      ,   snake_case							:Optional[int]=None	):





        '''simple docstring'''




        A_       :					str        =      [self.bos_token_id] + token_ids_a + [self.eos_token_id]
        if token_ids_a is None:
           return output

        return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]





     def     SCREAMING_SNAKE_CASE		(				self							:int      ,   snake_case							:List[int]      ,   snake_case							:Optional[List[int]] = None	):





        '''simple docstring'''




        A_       :					Dict        =      [self.sep_token_id]
        A_       :					List[str]        =      [self.cls_token_id]

        if token_ids_a is None:
           return len(cls + token_ids_a + sep	) * [0]
        return len(cls + token_ids_a + sep + sep + token_ids_a + sep	) * [0]
