

import json
import os
import re
import unicodedata
from typing import Optional, Tuple

from ...tokenization_utils import PreTrainedTokenizer, _is_control, _is_punctuation, _is_whitespace
from ...utils import logging


_SCREAMING_SNAKE_CASE							 =			logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE							 =			{
    '''vocab_file''': '''vocab.json''',
    '''merges_file''': '''merges.txt''',
}

_SCREAMING_SNAKE_CASE							 =			{
    '''vocab_file''': {'''openai-gpt''': '''https://huggingface.co/openai-gpt/resolve/main/vocab.json'''},
    '''merges_file''': {'''openai-gpt''': '''https://huggingface.co/openai-gpt/resolve/main/merges.txt'''},
}

_SCREAMING_SNAKE_CASE							 =			{
    '''openai-gpt''': 5_1_2,
}





def  _snake_case		(snake_case_					:     Tuple  )							->		Union[str, Any]:
  __UpperCAmelCase	:  str		=     text.strip()
  if not text:
    return []
  __UpperCAmelCase	:  Optional[Any]		=     text.split()
  return tokens


class      lowerCamelCase_ (			_lowerCAmelCase):




 def __init__(			self  ,  _UpperCAmelCase=True  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=True  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=True  ,  ):
   if never_split is None:
     __UpperCAmelCase	:  str		=     []
   __UpperCAmelCase	:  Any		=     do_lower_case
   __UpperCAmelCase	:  Any		=     set(_UpperCAmelCase							)
   __UpperCAmelCase	:  str		=     tokenize_chinese_chars
   __UpperCAmelCase	:  List[str]		=     strip_accents
   __UpperCAmelCase	:  List[str]		=     do_split_on_punc




 def     a       (			self  ,  _UpperCAmelCase  ,  _UpperCAmelCase=None							):
   __UpperCAmelCase	:  str		=     self.never_split.union(set(_UpperCAmelCase							)							) if never_split else self.never_split
   __UpperCAmelCase	:  str		=     self._clean_text(_UpperCAmelCase							)

   # This was added on November 1st, 2018 for the multilingual and Chinese
   # models. This is also applied to the English models now, but it doesn't
   # matter since the English models were not trained on any Chinese data
   # and generally don't have any Chinese data in them (there are Chinese
   # characters in the vocabulary because Wikipedia does have some Chinese
   # words in the English Wikipedia.).
   if self.tokenize_chinese_chars:
     __UpperCAmelCase	:  Dict		=     self._tokenize_chinese_chars(_UpperCAmelCase							)
   # prevents treating the same character with different unicode codepoints as different characters
   __UpperCAmelCase	:  Dict		=     unicodedata.normalize('NFC'  ,  _UpperCAmelCase							)
   __UpperCAmelCase	:  List[Any]		=     whitespace_tokenize(_UpperCAmelCase							)
   __UpperCAmelCase	:  str		=     []
   for token in orig_tokens:
     if token not in never_split:
       if self.do_lower_case:
         __UpperCAmelCase	:  Optional[Any]		=     token.lower()
         if self.strip_accents is not False:
           __UpperCAmelCase	:  Any		=     self._run_strip_accents(_UpperCAmelCase							)
       elif self.strip_accents:
         __UpperCAmelCase	:  Tuple		=     self._run_strip_accents(_UpperCAmelCase							)
     split_tokens.extend(self._run_split_on_punc(_UpperCAmelCase  ,  _UpperCAmelCase							)							)

   __UpperCAmelCase	:  List[Any]		=     whitespace_tokenize(' '.join(_UpperCAmelCase							)							)
   return output_tokens




 def     a       (			self  ,  _UpperCAmelCase							):
   __UpperCAmelCase	:  List[str]		=     unicodedata.normalize('NFD'  ,  _UpperCAmelCase							)
   __UpperCAmelCase	:  Tuple		=     []
   for char in text:
     __UpperCAmelCase	:  List[str]		=     unicodedata.category(_UpperCAmelCase							)
     if cat == "Mn":
       continue
     output.append(_UpperCAmelCase							)
   return "".join(_UpperCAmelCase							)




 def     a       (			self  ,  _UpperCAmelCase  ,  _UpperCAmelCase=None							):
   if not self.do_split_on_punc or (never_split is not None and text in never_split):
     return [text]
   __UpperCAmelCase	:  int		=     list(_UpperCAmelCase							)
   __UpperCAmelCase	:  Tuple		=     0
   __UpperCAmelCase	:  Tuple		=     True
   __UpperCAmelCase	:  Tuple		=     []
   while i < len(_UpperCAmelCase							):
     __UpperCAmelCase	:  Any		=     chars[i]
     if _is_punctuation(_UpperCAmelCase							):
       output.append([char]							)
       __UpperCAmelCase	:  int		=     True
     else:
       if start_new_word:
         output.append([]							)
       __UpperCAmelCase	:  str		=     False
       output[-1].append(_UpperCAmelCase							)
     i += 1

   return ["".join(_UpperCAmelCase							) for x in output]




 def     a       (			self  ,  _UpperCAmelCase							):
   __UpperCAmelCase	:  int		=     []
   for char in text:
     __UpperCAmelCase	:  List[Any]		=     ord(_UpperCAmelCase							)
     if self._is_chinese_char(_UpperCAmelCase							):
       output.append(' '							)
       output.append(_UpperCAmelCase							)
       output.append(' '							)
     else:
       output.append(_UpperCAmelCase							)
   return "".join(_UpperCAmelCase							)




 def     a       (			self  ,  _UpperCAmelCase							):
   # This defines a "chinese character" as anything in the CJK Unicode block:
   #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
   #
   # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
   # despite its name. The modern Korean Hangul alphabet is a different block,
   # as is Japanese Hiragana and Katakana. Those alphabets are used to write
   # space-separated words, so they are not treated specially and handled
   # like the all of the other languages.
   if (
       (cp >= 0x4_E00 and cp <= 0x9_FFF)
       or (cp >= 0x3_400 and cp <= 0x4_DBF)  #
       or (cp >= 0x20_000 and cp <= 0x2A_6DF)  #
       or (cp >= 0x2A_700 and cp <= 0x2B_73F)  #
       or (cp >= 0x2B_740 and cp <= 0x2B_81F)  #
       or (cp >= 0x2B_820 and cp <= 0x2C_EAF)  #
       or (cp >= 0xF_900 and cp <= 0xF_AFF)
       or (cp >= 0x2F_800 and cp <= 0x2F_A1F)  #
   ):  #
     return True

   return False




 def     a       (			self  ,  _UpperCAmelCase							):
   __UpperCAmelCase	:  List[Any]		=     []
   for char in text:
     __UpperCAmelCase	:  Optional[Any]		=     ord(_UpperCAmelCase							)
     if cp == 0 or cp == 0xF_FFD or _is_control(_UpperCAmelCase							):
       continue
     if _is_whitespace(_UpperCAmelCase							):
       output.append(' '							)
     else:
       output.append(_UpperCAmelCase							)
   return "".join(_UpperCAmelCase							)





def  _snake_case		(snake_case_					:     List[str]  )							->		Union[str, Any]:
  __UpperCAmelCase	:  Any		=     set()
  __UpperCAmelCase	:  Dict		=     word[0]
  for char in word[1:]:
    pairs.add((prev_char, char)  )
    __UpperCAmelCase	:  Union[str, Any]		=     char
  return pairs





def  _snake_case		(snake_case_					:     Union[str, Any]  )							->		int:
  __UpperCAmelCase	:  int		=     text.replace('—'    ,       '-'  )
  __UpperCAmelCase	:  int		=     text.replace('–'    ,       '-'  )
  __UpperCAmelCase	:  List[Any]		=     text.replace('―'    ,       '-'  )
  __UpperCAmelCase	:  Optional[Any]		=     text.replace('…'    ,       '...'  )
  __UpperCAmelCase	:  Union[str, Any]		=     text.replace('´'    ,       '\''  )
  __UpperCAmelCase	:  Union[str, Any]		=     re.sub(R'(-+|~+|!+|"+|;+|\?+|\++|,+|\)+|\(+|\\+|\/+|\*+|\[+|\]+|}+|{+|\|+|_+)'    ,       R' \1 '    ,       snake_case_  )
  __UpperCAmelCase	:  Dict		=     re.sub(R'\s*\n\s*'    ,       ' \n '    ,       snake_case_  )
  __UpperCAmelCase	:  Optional[Any]		=     re.sub(R'[^\S\n]+'    ,       ' '    ,       snake_case_  )
  return text.strip()







class      lowerCamelCase_ (			_lowerCAmelCase):
 _lowerCAmelCase									=	VOCAB_FILES_NAMES
 _lowerCAmelCase									=	PRETRAINED_VOCAB_FILES_MAP
 _lowerCAmelCase									=	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 _lowerCAmelCase									=	['''input_ids''', '''attention_mask''']




 def __init__(			self  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase="<unk>"  ,  **_UpperCAmelCase							):
   super().__init__(unk_token=_UpperCAmelCase  ,  **_UpperCAmelCase							)

   try:
     import ftfy
     from spacy.lang.en import English

     __UpperCAmelCase	:  Dict		=     English()
     __UpperCAmelCase	:  List[Any]		=     _nlp.tokenizer
     __UpperCAmelCase	:  int		=     ftfy.fix_text
   except ImportError:
     logger.warning('ftfy or spacy is not installed using BERT BasicTokenizer instead of SpaCy & ftfy.'							)
     __UpperCAmelCase	:  List[str]		=     BasicTokenizer(do_lower_case=_UpperCAmelCase							)
     __UpperCAmelCase	:  Dict		=     None

   with open(_UpperCAmelCase  ,  encoding='utf-8'							) as vocab_handle:
     __UpperCAmelCase	:  List[Any]		=     json.load(_UpperCAmelCase							)
   __UpperCAmelCase	:  Dict		=     {v: k for k, v in self.encoder.items()}
   with open(_UpperCAmelCase  ,  encoding='utf-8'							) as merges_handle:
     __UpperCAmelCase	:  List[Any]		=     merges_handle.read().split('\n'							)[1:-1]
   __UpperCAmelCase	:  Union[str, Any]		=     [tuple(merge.split()							) for merge in merges]
   __UpperCAmelCase	:  int		=     dict(zip(_UpperCAmelCase  ,  range(len(_UpperCAmelCase							)							)							)							)
   __UpperCAmelCase	:  Dict		=     {}




 @property
 def     a       (			self							):
   return True




 @property
 def     a       (			self							):
   return len(self.encoder							)




 def     a       (			self							):
   return dict(self.encoder  ,  **self.added_tokens_encoder							)




 def     a       (			self  ,  _UpperCAmelCase							):
   __UpperCAmelCase	:  List[Any]		=     tuple(token[:-1]							) + (token[-1] + '</w>',)
   if token in self.cache:
     return self.cache[token]
   __UpperCAmelCase	:  Optional[int]		=     get_pairs(_UpperCAmelCase							)

   if not pairs:
     return token + "</w>"

   while True:
     __UpperCAmelCase	:  str		=     min(_UpperCAmelCase  ,  key=lambda _UpperCAmelCase							: self.bpe_ranks.get(_UpperCAmelCase  ,  float('inf'							)							)							)
     if bigram not in self.bpe_ranks:
       break
     __UpperCAmelCase,			__UpperCAmelCase	:  List[str]		=     bigram
     __UpperCAmelCase	:  List[Any]		=     []
     __UpperCAmelCase	:  Union[str, Any]		=     0
     while i < len(_UpperCAmelCase							):
       try:
         __UpperCAmelCase	:  List[Any]		=     word.index(_UpperCAmelCase  ,  _UpperCAmelCase							)
       except ValueError:
         new_word.extend(word[i:]							)
         break
       else:
         new_word.extend(word[i:j]							)
         __UpperCAmelCase	:  Dict		=     j

       if word[i] == first and i < len(_UpperCAmelCase							) - 1 and word[i + 1] == second:
         new_word.append(first + second							)
         i += 2
       else:
         new_word.append(word[i]							)
         i += 1
     __UpperCAmelCase	:  Tuple		=     tuple(_UpperCAmelCase							)
     __UpperCAmelCase	:  Optional[int]		=     new_word
     if len(_UpperCAmelCase							) == 1:
       break
     else:
       __UpperCAmelCase	:  int		=     get_pairs(_UpperCAmelCase							)
   __UpperCAmelCase	:  List[Any]		=     ' '.join(_UpperCAmelCase							)
   if word == "\n  </w>":
     __UpperCAmelCase	:  List[str]		=     '\n</w>'
   __UpperCAmelCase	:  Tuple		=     word
   return word




 def     a       (			self  ,  _UpperCAmelCase							):
   __UpperCAmelCase	:  int		=     []
   if self.fix_text is None:
     # Using BERT's BasicTokenizer
     __UpperCAmelCase	:  str		=     self.nlp.tokenize(_UpperCAmelCase							)
     for token in text:
       split_tokens.extend(list(self.bpe(_UpperCAmelCase							).split(' '							)							)							)
   else:
     # Using SpaCy & ftfy (original tokenization process of OpenAI GPT)
     __UpperCAmelCase	:  str		=     self.nlp(text_standardize(self.fix_text(_UpperCAmelCase							)							)							)
     for token in text:
       split_tokens.extend(list(self.bpe(token.text.lower()							).split(' '							)							)							)
   return split_tokens




 def     a       (			self  ,  _UpperCAmelCase							):
   return self.encoder.get(_UpperCAmelCase  ,  self.encoder.get(self.unk_token							)							)




 def     a       (			self  ,  _UpperCAmelCase							):
   return self.decoder.get(_UpperCAmelCase  ,  self.unk_token							)




 def     a       (			self  ,  _UpperCAmelCase							):
   __UpperCAmelCase	:  str		=     ''.join(_UpperCAmelCase							).replace('</w>'  ,  ' '							).strip()
   return out_string




 def     a       (			self  ,  _UpperCAmelCase  ,  _UpperCAmelCase = None							):
   if not os.path.isdir(_UpperCAmelCase							):
     logger.error(F"Vocabulary path ({save_directory}) should be a directory"							)
     return
   __UpperCAmelCase	:  List[str]		=     os.path.join(
       _UpperCAmelCase  ,  (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']							)
   __UpperCAmelCase	:  Dict		=     os.path.join(
       _UpperCAmelCase  ,  (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['merges_file']							)

   with open(_UpperCAmelCase  ,  'w'  ,  encoding='utf-8'							) as f:
     f.write(json.dumps(self.encoder  ,  indent=2  ,  sort_keys=_UpperCAmelCase  ,  ensure_ascii=_UpperCAmelCase							) + '\n'							)

   __UpperCAmelCase	:  Tuple		=     0
   with open(_UpperCAmelCase  ,  'w'  ,  encoding='utf-8'							) as writer:
     writer.write('#version: 0.2\n'							)
     for bpe_tokens, token_index in sorted(self.bpe_ranks.items()  ,  key=lambda _UpperCAmelCase							: kv[1]							):
       if index != token_index:
         logger.warning(
             F"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
             ' Please check that the tokenizer is not corrupted!'							)
         __UpperCAmelCase	:  Optional[int]		=     token_index
       writer.write(' '.join(_UpperCAmelCase							) + '\n'							)
       index += 1

   return vocab_file, merge_file








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

import torch
from transformers import CLIPTextModelWithProjection, CLIPTokenizer
from transformers.models.clip.modeling_clip import CLIPTextModelOutput

from diffusers.models import PriorTransformer
from diffusers.pipelines import DiffusionPipeline, StableDiffusionImageVariationPipeline
from diffusers.schedulers import UnCLIPScheduler
from diffusers.utils import logging, randn_tensor


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





def  _snake_case		(self					:     Optional[Any]    ,       snake_case_					:     Optional[Any]    ,       snake_case_					:     Dict    ,       snake_case_					:     Optional[int]    ,       snake_case_					:     Optional[Any]  )							->		Union[str, Any]:
  __UpperCAmelCase	:  Any		=     image.to(device=snake_case_  )
  __UpperCAmelCase	:  Tuple		=     image  # take image as image_embeddings
  __UpperCAmelCase	:  str		=     image_embeddings.unsqueeze(1  )

  # duplicate image embeddings for each generation per prompt, using mps friendly method
  __UpperCAmelCase,			__UpperCAmelCase,			__UpperCAmelCase	:  int		=     image_embeddings.shape
  __UpperCAmelCase	:  int		=     image_embeddings.repeat(1    ,       snake_case_    ,       1  )
  __UpperCAmelCase	:  Optional[Any]		=     image_embeddings.view(bs_embed * num_images_per_prompt    ,       snake_case_    ,       -1  )

  if do_classifier_free_guidance:
    __UpperCAmelCase	:  List[Any]		=     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	:  Union[str, Any]		=     torch.cat([uncond_embeddings, image_embeddings]  )

  return image_embeddings







class      lowerCamelCase_ (			_lowerCAmelCase):




 def __init__(			self  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase = None  ,  ):
   super().__init__()

   __UpperCAmelCase	:  List[Any]		=     {'image_encoder': None} if decoder_pipe_kwargs is None else decoder_pipe_kwargs

   __UpperCAmelCase	:  Optional[Any]		=     decoder_pipe_kwargs.get('torch_dtype'  ,  _UpperCAmelCase							) or prior.dtype

   __UpperCAmelCase	:  str		=     StableDiffusionImageVariationPipeline.from_pretrained(
       'lambdalabs/sd-image-variations-diffusers'  ,  **_UpperCAmelCase							)

   # replace `_encode_image` method
   __UpperCAmelCase	:  Dict		=     types.MethodType(_encode_image  ,  self.decoder_pipe							)

   self.register_modules(
       prior=_UpperCAmelCase  ,  tokenizer=_UpperCAmelCase  ,  text_encoder=_UpperCAmelCase  ,  prior_scheduler=_UpperCAmelCase  ,  )




 def     a       (			self  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase = None  ,  _UpperCAmelCase = None  ,  ):
   if text_model_output is None:
     __UpperCAmelCase	:  Any		=     len(_UpperCAmelCase							) if isinstance(_UpperCAmelCase  ,  _UpperCAmelCase							) else 1
     # get prompt text embeddings
     __UpperCAmelCase	:  str		=     self.tokenizer(
         _UpperCAmelCase  ,  padding='max_length'  ,  max_length=self.tokenizer.model_max_length  ,  return_tensors='pt'  ,  )
     __UpperCAmelCase	:  int		=     text_inputs.input_ids
     __UpperCAmelCase	:  str		=     text_inputs.attention_mask.bool().to(_UpperCAmelCase							)

     if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
       __UpperCAmelCase	:  Any		=     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}"							)
       __UpperCAmelCase	:  int		=     text_input_ids[:, : self.tokenizer.model_max_length]

     __UpperCAmelCase	:  Any		=     self.text_encoder(text_input_ids.to(_UpperCAmelCase							)							)

     __UpperCAmelCase	:  str		=     text_encoder_output.text_embeds
     __UpperCAmelCase	:  str		=     text_encoder_output.last_hidden_state

   else:
     __UpperCAmelCase	:  str		=     text_model_output[0].shape[0]
     __UpperCAmelCase,			__UpperCAmelCase	:  Any		=     text_model_output[0], text_model_output[1]
     __UpperCAmelCase	:  Optional[Any]		=     text_attention_mask

   __UpperCAmelCase	:  Optional[Any]		=     text_embeddings.repeat_interleave(_UpperCAmelCase  ,  dim=0							)
   __UpperCAmelCase	:  Tuple		=     text_encoder_hidden_states.repeat_interleave(_UpperCAmelCase  ,  dim=0							)
   __UpperCAmelCase	:  Optional[Any]		=     text_mask.repeat_interleave(_UpperCAmelCase  ,  dim=0							)

   if do_classifier_free_guidance:
     __UpperCAmelCase	:  int		=     [''] * batch_size

     __UpperCAmelCase	:  Optional[int]		=     self.tokenizer(
         _UpperCAmelCase  ,  padding='max_length'  ,  max_length=self.tokenizer.model_max_length  ,  truncation=_UpperCAmelCase  ,  return_tensors='pt'  ,  )
     __UpperCAmelCase	:  List[str]		=     uncond_input.attention_mask.bool().to(_UpperCAmelCase							)
     __UpperCAmelCase	:  int		=     self.text_encoder(uncond_input.input_ids.to(_UpperCAmelCase							)							)

     __UpperCAmelCase	:  Optional[Any]		=     uncond_embeddings_text_encoder_output.text_embeds
     __UpperCAmelCase	:  List[str]		=     uncond_embeddings_text_encoder_output.last_hidden_state

     # duplicate unconditional embeddings for each generation per prompt, using mps friendly method

     __UpperCAmelCase	:  List[str]		=     uncond_embeddings.shape[1]
     __UpperCAmelCase	:  Tuple		=     uncond_embeddings.repeat(1  ,  _UpperCAmelCase							)
     __UpperCAmelCase	:  List[str]		=     uncond_embeddings.view(batch_size * num_images_per_prompt  ,  _UpperCAmelCase							)

     __UpperCAmelCase	:  Tuple		=     uncond_text_encoder_hidden_states.shape[1]
     __UpperCAmelCase	:  Union[str, Any]		=     uncond_text_encoder_hidden_states.repeat(1  ,  _UpperCAmelCase  ,  1							)
     __UpperCAmelCase	:  Any		=     uncond_text_encoder_hidden_states.view(
         batch_size * num_images_per_prompt  ,  _UpperCAmelCase  ,  -1							)
     __UpperCAmelCase	:  int		=     uncond_text_mask.repeat_interleave(_UpperCAmelCase  ,  dim=0							)

     # done duplicates

     # 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([uncond_embeddings, text_embeddings]							)
     __UpperCAmelCase	:  str		=     torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]							)

     __UpperCAmelCase	:  Any		=     torch.cat([uncond_text_mask, text_mask]							)

   return text_embeddings, text_encoder_hidden_states, text_mask




 @property
 def     a       (			self							):
   if self.device != torch.device('meta'							) or not hasattr(self.prior  ,  '_hf_hook'							):
     return self.device
   for module in self.prior.modules():
     if (
         hasattr(_UpperCAmelCase  ,  '_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  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase							):
   if latents is None:
     __UpperCAmelCase	:  Any		=     randn_tensor(_UpperCAmelCase  ,  generator=_UpperCAmelCase  ,  device=_UpperCAmelCase  ,  dtype=_UpperCAmelCase							)
   else:
     if latents.shape != shape:
       raise ValueError(F"Unexpected latents shape, got {latents.shape}, expected {shape}"							)
     __UpperCAmelCase	:  Optional[int]		=     latents.to(_UpperCAmelCase							)

   __UpperCAmelCase	:  str		=     latents * scheduler.init_noise_sigma
   return latents




 def     a       (			self  ,  _UpperCAmelCase = None							):
   self.decoder_pipe.to(_UpperCAmelCase							)
   super().to(_UpperCAmelCase							)




 @torch.no_grad()
 def __call__(			self  ,  _UpperCAmelCase = None  ,  _UpperCAmelCase = None  ,  _UpperCAmelCase = None  ,  _UpperCAmelCase = 1  ,  _UpperCAmelCase = 25  ,  _UpperCAmelCase = None  ,  _UpperCAmelCase = None  ,  _UpperCAmelCase = None  ,  _UpperCAmelCase = None  ,  _UpperCAmelCase = 4.0  ,  _UpperCAmelCase = 8.0  ,  _UpperCAmelCase = 50  ,  _UpperCAmelCase = 1  ,  _UpperCAmelCase = 0.0  ,  _UpperCAmelCase = "pil"  ,  _UpperCAmelCase = True  ,  ):
   if prompt is not None:
     if isinstance(_UpperCAmelCase  ,  _UpperCAmelCase							):
       __UpperCAmelCase	:  List[Any]		=     1
     elif isinstance(_UpperCAmelCase  ,  _UpperCAmelCase							):
       __UpperCAmelCase	:  Tuple		=     len(_UpperCAmelCase							)
     else:
       raise ValueError(F"`prompt` has to be of type `str` or `list` but is {type(_UpperCAmelCase							)}"							)
   else:
     __UpperCAmelCase	:  Tuple		=     text_model_output[0].shape[0]

   __UpperCAmelCase	:  Any		=     self._execution_device

   __UpperCAmelCase	:  Any		=     batch_size * num_images_per_prompt

   __UpperCAmelCase	:  List[str]		=     prior_guidance_scale > 1.0 or decoder_guidance_scale > 1.0

   __UpperCAmelCase,			__UpperCAmelCase,			__UpperCAmelCase	:  Optional[Any]		=     self._encode_prompt(
       _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase							)

   # prior

   self.prior_scheduler.set_timesteps(_UpperCAmelCase  ,  device=_UpperCAmelCase							)
   __UpperCAmelCase	:  Dict		=     self.prior_scheduler.timesteps

   __UpperCAmelCase	:  Tuple		=     self.prior.config.embedding_dim

   __UpperCAmelCase	:  Dict		=     self.prepare_latents(
       (batch_size, embedding_dim)  ,  text_embeddings.dtype  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  self.prior_scheduler  ,  )

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

     __UpperCAmelCase	:  Union[str, Any]		=     self.prior(
         _UpperCAmelCase  ,  timestep=_UpperCAmelCase  ,  proj_embedding=_UpperCAmelCase  ,  encoder_hidden_states=_UpperCAmelCase  ,  attention_mask=_UpperCAmelCase  ,  ).predicted_image_embedding

     if do_classifier_free_guidance:
       __UpperCAmelCase,			__UpperCAmelCase	:  List[str]		=     predicted_image_embedding.chunk(2							)
       __UpperCAmelCase	:  Optional[Any]		=     predicted_image_embedding_uncond + prior_guidance_scale * (
           predicted_image_embedding_text - predicted_image_embedding_uncond
       )

     if i + 1 == prior_timesteps_tensor.shape[0]:
       __UpperCAmelCase	:  int		=     None
     else:
       __UpperCAmelCase	:  Any		=     prior_timesteps_tensor[i + 1]

     __UpperCAmelCase	:  Union[str, Any]		=     self.prior_scheduler.step(
         _UpperCAmelCase  ,  timestep=_UpperCAmelCase  ,  sample=_UpperCAmelCase  ,  generator=_UpperCAmelCase  ,  prev_timestep=_UpperCAmelCase  ,  ).prev_sample

   __UpperCAmelCase	:  Union[str, Any]		=     self.prior.post_process_latents(_UpperCAmelCase							)

   __UpperCAmelCase	:  Any		=     prior_latents

   __UpperCAmelCase	:  List[Any]		=     self.decoder_pipe(
       image=_UpperCAmelCase  ,  height=_UpperCAmelCase  ,  width=_UpperCAmelCase  ,  num_inference_steps=_UpperCAmelCase  ,  guidance_scale=_UpperCAmelCase  ,  generator=_UpperCAmelCase  ,  output_type=_UpperCAmelCase  ,  return_dict=_UpperCAmelCase  ,  num_images_per_prompt=_UpperCAmelCase  ,  eta=_UpperCAmelCase  ,  )
   return output





