


from typing import Callable, List, Optional, Tuple, Union

import torch
from transformers import CLIPTextModel, CLIPTokenizer

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin, TransformeraDModel, VQModel
from ...schedulers import VQDiffusionScheduler
from ...utils import logging
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


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



class 							_UpperCamelCase (						__snake_case,__snake_case						):

   '''simple docstring'''







   @register_to_config
   def __init__(  self  : Any       ,  _lowerCamelCase  : Optional[int]       ,  _lowerCamelCase  : Tuple = None       ,  _lowerCamelCase  : Tuple = None							):






         '''simple docstring'''



         super().__init__()

         __lowerCamelCase     :      List[Any]       						=						learnable

         if self.learnable:
               assert hidden_size is not None, "learnable=True requires `hidden_size` to be set"
               assert length is not None, "learnable=True requires `length` to be set"

               __lowerCamelCase     :      Any       						=						torch.zeros(__UpperCamelCase       ,  __UpperCamelCase							)
         else:
               __lowerCamelCase     :      int       						=						None

         __lowerCamelCase     :      Optional[int]       						=						torch.nn.Parameter(__UpperCamelCase							)



class 							_UpperCamelCase (						__snake_case						):

   '''simple docstring'''







   a_				:							int    = 42
   a_				:							List[Any]    = 42
   a_				:							str    = 42
   a_				:							str    = 42
   a_				:							Tuple    = 42
   a_				:							List[Any]    = 42

   def __init__(  self  : Dict       ,  _lowerCamelCase  : Any       ,  _lowerCamelCase  : List[Any]       ,  _lowerCamelCase  : Optional[Any]       ,  _lowerCamelCase  : Tuple       ,  _lowerCamelCase  : Union[str, Any]       ,  _lowerCamelCase  : int       ,  ):






         '''simple docstring'''



         super().__init__()

         self.register_modules(
             vqvae=__UpperCamelCase       ,  transformer=__UpperCamelCase       ,  text_encoder=__UpperCamelCase       ,  tokenizer=__UpperCamelCase       ,  scheduler=__UpperCamelCase       ,  learned_classifier_free_sampling_embeddings=__UpperCamelCase       ,  )

   def   _snake_case    (  self  : int       ,  _lowerCamelCase  : List[str]       ,  _lowerCamelCase  : Any       ,  _lowerCamelCase  : str							):






         '''simple docstring'''



         __lowerCamelCase     :      Any       						=						len(__UpperCamelCase							) if isinstance(__UpperCamelCase       ,  __UpperCamelCase							) else 1

         # get prompt text embeddings
         __lowerCamelCase     :      Dict       						=						self.tokenizer(
             __UpperCamelCase       ,  padding="""max_length"""       ,  max_length=self.tokenizer.model_max_length       ,  return_tensors="""pt"""       ,  )
         __lowerCamelCase     :      Optional[Any]       						=						text_inputs.input_ids

         if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
               __lowerCamelCase     :      str       						=						self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]							)
               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}"""							)
               __lowerCamelCase     :      Tuple       						=						text_input_ids[:, : self.tokenizer.model_max_length]
         __lowerCamelCase     :      Optional[Any]       						=						self.text_encoder(text_input_ids.to(self.device							)							)[0]

         # NOTE: This additional step of normalizing the text embeddings is from VQ-Diffusion.
         # While CLIP does normalize the pooled output of the text transformer when combining
         # the image and text embeddings, CLIP does not directly normalize the last hidden state.
         #
         # CLIP normalizing the pooled output.
         # https://github.com/huggingface/transformers/blob/d92e22d1f28324f513f3080e5c47c071a3916721/src/transformers/models/clip/modeling_clip.py#L1052-L1053
         __lowerCamelCase     :      Tuple       						=						prompt_embeds / prompt_embeds.norm(dim=-1       ,  keepdim=__UpperCamelCase							)

         # duplicate text embeddings for each generation per prompt
         __lowerCamelCase     :      List[str]       						=						prompt_embeds.repeat_interleave(__UpperCamelCase       ,  dim=0							)

         if do_classifier_free_guidance:
               if self.learned_classifier_free_sampling_embeddings.learnable:
                     __lowerCamelCase     :      List[str]       						=						self.learned_classifier_free_sampling_embeddings.embeddings
                     __lowerCamelCase     :      Union[str, Any]       						=						negative_prompt_embeds.unsqueeze(0							).repeat(__UpperCamelCase       ,  1       ,  1							)
               else:
                     __lowerCamelCase     :      int       						=						[""""""] * batch_size

                     __lowerCamelCase     :      int       						=						text_input_ids.shape[-1]
                     __lowerCamelCase     :      Optional[int]       						=						self.tokenizer(
                         __UpperCamelCase       ,  padding="""max_length"""       ,  max_length=__UpperCamelCase       ,  truncation=__UpperCamelCase       ,  return_tensors="""pt"""       ,  )
                     __lowerCamelCase     :      Dict       						=						self.text_encoder(uncond_input.input_ids.to(self.device							)							)[0]
                     # See comment for normalizing text embeddings
                     __lowerCamelCase     :      Optional[Any]       						=						negative_prompt_embeds / negative_prompt_embeds.norm(dim=-1       ,  keepdim=__UpperCamelCase							)

               # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
               __lowerCamelCase     :      Dict       						=						negative_prompt_embeds.shape[1]
               __lowerCamelCase     :      Dict       						=						negative_prompt_embeds.repeat(1       ,  __UpperCamelCase       ,  1							)
               __lowerCamelCase     :      Any       						=						negative_prompt_embeds.view(batch_size * num_images_per_prompt       ,  __UpperCamelCase       ,  -1							)

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

         return prompt_embeds

   @torch.no_grad()
   def __call__(  self  : Optional[Any]       ,  _lowerCamelCase  : Any       ,  _lowerCamelCase  : Tuple = 1_0_0       ,  _lowerCamelCase  : List[str] = 5.0       ,  _lowerCamelCase  : List[Any] = 1.0       ,  _lowerCamelCase  : Union[str, Any] = 1       ,  _lowerCamelCase  : int = None       ,  _lowerCamelCase  : Dict = None       ,  _lowerCamelCase  : Optional[Any] = "pil"       ,  _lowerCamelCase  : List[Any] = True       ,  _lowerCamelCase  : Tuple = None       ,  _lowerCamelCase  : Dict = 1       ,  ):






         '''simple docstring'''



         if isinstance(__UpperCamelCase       ,  __UpperCamelCase							):
               __lowerCamelCase     :      List[Any]       						=						1
         elif isinstance(__UpperCamelCase       ,  __UpperCamelCase							):
               __lowerCamelCase     :      Tuple       						=						len(__UpperCamelCase							)
         else:
               raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(__UpperCamelCase							)}"""							)

         __lowerCamelCase     :      int       						=						batch_size * num_images_per_prompt

         __lowerCamelCase     :      Optional[int]       						=						guidance_scale > 1.0

         __lowerCamelCase     :      Any       						=						self._encode_prompt(__UpperCamelCase       ,  __UpperCamelCase       ,  __UpperCamelCase							)

         if (callback_steps is None) or (
             callback_steps is not None and (not isinstance(__UpperCamelCase       ,  __UpperCamelCase							) or callback_steps <= 0)
         ):
               raise ValueError(
                   F"""`callback_steps` has to be a positive integer but is {callback_steps} of type"""
                   F""" {type(__UpperCamelCase							)}."""							)

         # get the initial completely masked latents unless the user supplied it

         __lowerCamelCase     :      Any       						=						(batch_size, self.transformer.num_latent_pixels)
         if latents is None:
               __lowerCamelCase     :      Tuple       						=						self.transformer.num_vector_embeds - 1
               __lowerCamelCase     :      List[Any]       						=						torch.full(__UpperCamelCase       ,  __UpperCamelCase							).to(self.device							)
         else:
               if latents.shape != latents_shape:
                     raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {latents_shape}"""							)
               if (latents < 0).any() or (latents >= self.transformer.num_vector_embeds).any():
                     raise ValueError(
                         """Unexpected latents value(s). All latents be valid embedding indices i.e. in the range 0,"""
                         F""" {self.transformer.num_vector_embeds - 1} (inclusive)."""							)
               __lowerCamelCase     :      Tuple       						=						latents.to(self.device							)

         # set timesteps
         self.scheduler.set_timesteps(__UpperCamelCase       ,  device=self.device							)

         __lowerCamelCase     :      str       						=						self.scheduler.timesteps.to(self.device							)

         __lowerCamelCase     :      Union[str, Any]       						=						latents

         for i, t in enumerate(self.progress_bar(__UpperCamelCase							)							):
               # expand the sample if we are doing classifier free guidance
               __lowerCamelCase     :      Union[str, Any]       						=						torch.cat([sample] * 2							) if do_classifier_free_guidance else sample

               # predict the un-noised image
               # model_output == `log_p_x_0`
               __lowerCamelCase     :      int       						=						self.transformer(__UpperCamelCase       ,  encoder_hidden_states=__UpperCamelCase       ,  timestep=__UpperCamelCase							).sample

               if do_classifier_free_guidance:
                     __lowerCamelCase						,  __lowerCamelCase     :      Optional[int]       						=						model_output.chunk(2							)
                     __lowerCamelCase     :      List[Any]       						=						model_output_uncond + guidance_scale * (model_output_text - model_output_uncond)
                     model_output -= torch.logsumexp(__UpperCamelCase       ,  dim=1       ,  keepdim=__UpperCamelCase							)

               __lowerCamelCase     :      Any       						=						self.truncate(__UpperCamelCase       ,  __UpperCamelCase							)

               # remove `log(0)`'s (`-inf`s)
               __lowerCamelCase     :      List[Any]       						=						model_output.clamp(-7_0							)

               # compute the previous noisy sample x_t -> x_t-1
               __lowerCamelCase     :      List[Any]       						=						self.scheduler.step(__UpperCamelCase       ,  timestep=__UpperCamelCase       ,  sample=__UpperCamelCase       ,  generator=__UpperCamelCase							).prev_sample

               # call the callback, if provided
               if callback is not None and i % callback_steps == 0:
                     callback(__UpperCamelCase       ,  __UpperCamelCase       ,  __UpperCamelCase							)

         __lowerCamelCase     :      Union[str, Any]       						=						self.vqvae.config.vq_embed_dim
         __lowerCamelCase     :      Optional[Any]       						=						(batch_size, self.transformer.height, self.transformer.width, embedding_channels)
         __lowerCamelCase     :      Dict       						=						self.vqvae.quantize.get_codebook_entry(__UpperCamelCase       ,  shape=__UpperCamelCase							)
         __lowerCamelCase     :      int       						=						self.vqvae.decode(__UpperCamelCase       ,  force_not_quantize=__UpperCamelCase							).sample

         __lowerCamelCase     :      List[Any]       						=						(image / 2 + 0.5).clamp(0       ,  1							)
         __lowerCamelCase     :      Dict       						=						image.cpu().permute(0       ,  2       ,  3       ,  1							).numpy()

         if output_type == "pil":
               __lowerCamelCase     :      List[str]       						=						self.numpy_to_pil(__UpperCamelCase							)

         if not return_dict:
               return (image,)

         return ImagePipelineOutput(images=__UpperCamelCase							)


   def   _snake_case    (  self  : List[Any]       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : Dict							):






         '''simple docstring'''



         __lowerCamelCase						,  __lowerCamelCase     :      int       						=						torch.sort(__UpperCamelCase       ,  1       ,  descending=__UpperCamelCase							)
         __lowerCamelCase     :      Dict       						=						torch.exp(__UpperCamelCase							)
         __lowerCamelCase     :      Any       						=						sorted_p_x_0.cumsum(dim=1							) < truncation_rate

         # Ensure that at least the largest probability is not zeroed out
         __lowerCamelCase     :      str       						=						torch.full_like(keep_mask[:, 0:1, :]       ,  __UpperCamelCase							)
         __lowerCamelCase     :      List[Any]       						=						torch.cat((all_true, keep_mask)       ,  dim=1							)
         __lowerCamelCase     :      Optional[Any]       						=						keep_mask[:, :-1, :]

         __lowerCamelCase     :      Optional[int]       						=						keep_mask.gather(1       ,  indices.argsort(1							)							)

         __lowerCamelCase     :      Optional[int]       						=						log_p_x_0.clone()

         __lowerCamelCase     :      Dict       						=						-torch.inf  # -inf = log(0)

         return rv







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__UpperCamelCase    :						Any     =							{
    'configuration_timesformer': ['TIMESFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP', 'TimesformerConfig'],
}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __UpperCamelCase    :						int     =							[
        'TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
        'TimesformerModel',
        'TimesformerForVideoClassification',
        'TimesformerPreTrainedModel',
    ]

if TYPE_CHECKING:
    from .configuration_timesformer import TIMESFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, TimesformerConfig

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_timesformer import (
            TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
            TimesformerForVideoClassification,
            TimesformerModel,
            TimesformerPreTrainedModel,
        )

else:
    import sys

    __UpperCamelCase    :						List[Any]     =							_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



