





"""simple docstring"""
import os
from glob import glob

import imageio
import torch
import torchvision
import wandb
from img_processing import custom_to_pil, loop_post_process, preprocess, preprocess_vqgan
from loaders import load_vqgan
from PIL import Image
from torch import nn

from transformers import CLIPModel, CLIPTokenizerFast
from utils import get_device, get_timestamp, show_pil






class  _SCREAMING_SNAKE_CASE:




       def __init__(						self		,SCREAMING_SNAKE_CASE__ = "cpu"		,SCREAMING_SNAKE_CASE__ = "openai/clip-vit-large-patch14"							)      ->				None:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  device
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  CLIPTokenizerFast.from_pretrained(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  [0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3]
             __SCREAMING_SNAKE_CASE       :int		    =  [0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1]
             __SCREAMING_SNAKE_CASE       :List[str]		    =  torchvision.transforms.Normalize(self.image_mean		,self.image_std							)
             __SCREAMING_SNAKE_CASE       :Tuple		    =  torchvision.transforms.Resize(2_24							)
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  torchvision.transforms.CenterCrop(2_24							)




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__							)      ->				Union[str, Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :str		    =  self.resize(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :Dict		    =  self.center_crop(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :Tuple		    =  self.normalize(SCREAMING_SNAKE_CASE__							)
             return images




       def __call__(						self		,SCREAMING_SNAKE_CASE__=None		,SCREAMING_SNAKE_CASE__=None		,**SCREAMING_SNAKE_CASE__							)      ->				Union[str, Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  self.tokenizer(text=SCREAMING_SNAKE_CASE__		,**SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :Any		    =  self.preprocess_img(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  {key: value.to(self.device							) for (key, value) in encoding.items()}
             return encoding






class  _SCREAMING_SNAKE_CASE(       nn.Module  ):




       def __init__(						self		,SCREAMING_SNAKE_CASE__=10		,SCREAMING_SNAKE_CASE__=0.0_1		,SCREAMING_SNAKE_CASE__=None		,SCREAMING_SNAKE_CASE__=None		,SCREAMING_SNAKE_CASE__=None		,SCREAMING_SNAKE_CASE__=None		,SCREAMING_SNAKE_CASE__=None		,SCREAMING_SNAKE_CASE__=None		,SCREAMING_SNAKE_CASE__=False		,SCREAMING_SNAKE_CASE__=True		,SCREAMING_SNAKE_CASE__="image"		,SCREAMING_SNAKE_CASE__=True		,SCREAMING_SNAKE_CASE__=False		,SCREAMING_SNAKE_CASE__=False		,SCREAMING_SNAKE_CASE__=False		,)      ->				None:





             """simple docstring"""






             super().__init__()
             __SCREAMING_SNAKE_CASE       :List[str]		    =  None
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  device if device else get_device()
             if vqgan:
                   __SCREAMING_SNAKE_CASE       :Any		    =  vqgan
             else:
                   __SCREAMING_SNAKE_CASE       :Tuple		    =  load_vqgan(self.device		,conf_path=SCREAMING_SNAKE_CASE__		,ckpt_path=SCREAMING_SNAKE_CASE__							)
             self.vqgan.eval()
             if clip:
                   __SCREAMING_SNAKE_CASE       :Optional[int]		    =  clip
             else:
                   __SCREAMING_SNAKE_CASE       :Tuple		    =  CLIPModel.from_pretrained('''openai/clip-vit-base-patch32'''							)
             self.clip.to(self.device							)
             __SCREAMING_SNAKE_CASE       :Tuple		    =  ProcessorGradientFlow(device=self.device							)

             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  iterations
             __SCREAMING_SNAKE_CASE       :List[str]		    =  lr
             __SCREAMING_SNAKE_CASE       :List[str]		    =  log
             __SCREAMING_SNAKE_CASE       :int		    =  make_grid
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  return_val
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  quantize
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  self.vqgan.decoder.z_shape




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__=None		,SCREAMING_SNAKE_CASE__=None		,SCREAMING_SNAKE_CASE__=5		,SCREAMING_SNAKE_CASE__=True							)      ->				int:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :List[str]		    =  []
             if output_path is None:
                   __SCREAMING_SNAKE_CASE       :int		    =  '''./animation.gif'''
             if input_path is None:
                   __SCREAMING_SNAKE_CASE       :Any		    =  self.save_path
             __SCREAMING_SNAKE_CASE       :List[str]		    =  sorted(glob(input_path + '''/*'''							)							)
             if not len(SCREAMING_SNAKE_CASE__							):
                   raise ValueError(
                       '''No images found in save path, aborting (did you pass save_intermediate=True to the generate'''
                       ''' function?)'''							)
             if len(SCREAMING_SNAKE_CASE__							) == 1:
                   print('''Only one image found in save path, (did you pass save_intermediate=True to the generate function?)'''							)
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  total_duration / len(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :str		    =  [frame_duration] * len(SCREAMING_SNAKE_CASE__							)
             if extend_frames:
                   __SCREAMING_SNAKE_CASE       :Dict		    =  1.5
                   __SCREAMING_SNAKE_CASE       :Dict		    =  3
             for file_name in paths:
                   if file_name.endswith('''.png'''							):
                         images.append(imageio.imread(SCREAMING_SNAKE_CASE__							)							)
             imageio.mimsave(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,duration=SCREAMING_SNAKE_CASE__							)
             print(f'''gif saved to {output_path}'''							)




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__=None		,SCREAMING_SNAKE_CASE__=None							)      ->				Optional[int]:





             """simple docstring"""






             if not (path or img):
                   raise ValueError('''Input either path or tensor'''							)
             if img is not None:
                   raise NotImplementedError
             __SCREAMING_SNAKE_CASE       :Tuple		    =  preprocess(Image.open(SCREAMING_SNAKE_CASE__							)		,target_image_size=2_56							).to(self.device							)
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  preprocess_vqgan(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE    ,*__SCREAMING_SNAKE_CASE       :int		    =  self.vqgan.encode(SCREAMING_SNAKE_CASE__							)
             return z




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__							)      ->				Optional[int]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :List[Any]		    =  self.latent.detach().requires_grad_()
             __SCREAMING_SNAKE_CASE       :Dict		    =  base_latent + transform_vector
             if self.quantize:
                   __SCREAMING_SNAKE_CASE    ,*__SCREAMING_SNAKE_CASE       :str		    =  self.vqgan.quantize(SCREAMING_SNAKE_CASE__							)
             else:
                   __SCREAMING_SNAKE_CASE       :Optional[int]		    =  trans_latent
             return self.vqgan.decode(SCREAMING_SNAKE_CASE__							)




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=None							)      ->				str:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  self.clip_preprocessor(text=SCREAMING_SNAKE_CASE__		,images=SCREAMING_SNAKE_CASE__		,return_tensors='''pt'''		,padding=SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  self.clip(**SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  clip_outputs.logits_per_image
             if weights is not None:
                   __SCREAMING_SNAKE_CASE       :int		    =  similarity_logits * weights
             return similarity_logits.sum()




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)      ->				Optional[int]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :int		    =  self._get_clip_similarity(pos_prompts['''prompts''']		,SCREAMING_SNAKE_CASE__		,weights=(1 / pos_prompts['''weights'''])							)
             if neg_prompts:
                   __SCREAMING_SNAKE_CASE       :str		    =  self._get_clip_similarity(neg_prompts['''prompts''']		,SCREAMING_SNAKE_CASE__		,weights=neg_prompts['''weights''']							)
             else:
                   __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  torch.tensor([1]		,device=self.device							)
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  -torch.log(SCREAMING_SNAKE_CASE__							) + torch.log(SCREAMING_SNAKE_CASE__							)
             return loss




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)      ->				List[Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :List[str]		    =  torch.randn_like(self.latent		,requires_grad=SCREAMING_SNAKE_CASE__		,device=self.device							)
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  torch.optim.Adam([vector]		,lr=self.lr							)

             for i in range(self.iterations							):
                   optim.zero_grad()
                   __SCREAMING_SNAKE_CASE       :Any		    =  self._add_vector(SCREAMING_SNAKE_CASE__							)
                   __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  loop_post_process(SCREAMING_SNAKE_CASE__							)
                   __SCREAMING_SNAKE_CASE       :Tuple		    =  self._get_CLIP_loss(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)
                   print('''CLIP loss'''		,SCREAMING_SNAKE_CASE__							)
                   if self.log:
                         wandb.log({'''CLIP Loss''': clip_loss}							)
                   clip_loss.backward(retain_graph=SCREAMING_SNAKE_CASE__							)
                   optim.step()
                   if self.return_val == "image":
                         yield custom_to_pil(transformed_img[0]							)
                   else:
                         yield vector




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)      ->				Optional[int]:





             """simple docstring"""






             wandb.init(reinit=SCREAMING_SNAKE_CASE__		,project='''face-editor'''							)
             wandb.config.update({'''Positive Prompts''': positive_prompts}							)
             wandb.config.update({'''Negative Prompts''': negative_prompts}							)
             wandb.config.update({'''lr''': self.lr, '''iterations''': self.iterations}							)
             if image_path:
                   __SCREAMING_SNAKE_CASE       :List[Any]		    =  Image.open(SCREAMING_SNAKE_CASE__							)
                   __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  image.resize((2_56, 2_56)							)
                   wandb.log('''Original Image'''		,wandb.Image(SCREAMING_SNAKE_CASE__							)							)




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__							)      ->				Any:





             """simple docstring"""






             if not prompts:
                   return []
             __SCREAMING_SNAKE_CASE       :int		    =  []
             __SCREAMING_SNAKE_CASE       :List[str]		    =  []
             if isinstance(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							):
                   __SCREAMING_SNAKE_CASE       :Optional[int]		    =  [prompt.strip() for prompt in prompts.split('''|'''							)]
             for prompt in prompts:
                   if isinstance(SCREAMING_SNAKE_CASE__		,(tuple, list)							):
                         __SCREAMING_SNAKE_CASE       :Tuple		    =  prompt[0]
                         __SCREAMING_SNAKE_CASE       :Optional[int]		    =  float(prompt[1]							)
                   elif ":" in prompt:
                         __SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE       :List[str]		    =  prompt.split(''':'''							)
                         __SCREAMING_SNAKE_CASE       :List[str]		    =  float(SCREAMING_SNAKE_CASE__							)
                   else:
                         __SCREAMING_SNAKE_CASE       :str		    =  prompt
                         __SCREAMING_SNAKE_CASE       :Optional[int]		    =  1.0
                   processed_prompts.append(SCREAMING_SNAKE_CASE__							)
                   weights.append(SCREAMING_SNAKE_CASE__							)
             return {
                 "prompts": processed_prompts,
                 "weights": torch.tensor(SCREAMING_SNAKE_CASE__		,device=self.device							),
             }




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=None		,SCREAMING_SNAKE_CASE__=None		,SCREAMING_SNAKE_CASE__=True		,SCREAMING_SNAKE_CASE__=False		,SCREAMING_SNAKE_CASE__=True		,SCREAMING_SNAKE_CASE__=True		,SCREAMING_SNAKE_CASE__=None		,)      ->				Union[str, Any]:





             """simple docstring"""






             if image_path:
                   __SCREAMING_SNAKE_CASE       :str		    =  self._get_latent(SCREAMING_SNAKE_CASE__							)
             else:
                   __SCREAMING_SNAKE_CASE       :str		    =  torch.randn(self.latent_dim		,device=self.device							)
             if self.log:
                   self._init_logging(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)

             assert pos_prompts, "You must provide at least one positive prompt."
             __SCREAMING_SNAKE_CASE       :Any		    =  self.process_prompts(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  self.process_prompts(SCREAMING_SNAKE_CASE__							)
             if save_final and save_path is None:
                   __SCREAMING_SNAKE_CASE       :Optional[int]		    =  os.path.join('''./outputs/'''		,'''_'''.join(pos_prompts['''prompts''']							)							)
             if not os.path.exists(SCREAMING_SNAKE_CASE__							):
                   os.makedirs(SCREAMING_SNAKE_CASE__							)
             else:
                   __SCREAMING_SNAKE_CASE       :str		    =  save_path + '''_''' + get_timestamp()
                   os.makedirs(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  save_path

             __SCREAMING_SNAKE_CASE       :Any		    =  self.vqgan.decode(self.latent							)[0]
             if show_intermediate:
                   print('''Original Image'''							)
                   show_pil(custom_to_pil(SCREAMING_SNAKE_CASE__							)							)

             __SCREAMING_SNAKE_CASE       :str		    =  loop_post_process(SCREAMING_SNAKE_CASE__							)
             for iter, transformed_img in enumerate(self._optimize_CLIP(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)							):
                   if show_intermediate:
                         show_pil(SCREAMING_SNAKE_CASE__							)
                   if save_intermediate:
                         transformed_img.save(os.path.join(self.save_path		,f'''iter_{iter:03d}.png'''							)							)
                   if self.log:
                         wandb.log({'''Image''': wandb.Image(SCREAMING_SNAKE_CASE__							)}							)
             if show_final:
                   show_pil(SCREAMING_SNAKE_CASE__							)
             if save_final:
                   transformed_img.save(os.path.join(self.save_path		,f'''iter_{iter:03d}_final.png'''							)							)






"""simple docstring"""
import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


lowerCamelCase_       =		logging.get_logger(__name__)

lowerCamelCase_       =		{"vocab_file": "spiece.model"}

lowerCamelCase_       =		{
    "vocab_file": {
        "bert_for_seq_generation": (
            "https://huggingface.co/google/bert_for_seq_generation_L-24_bbc_encoder/resolve/main/spiece.model"
        ),
    }
}

lowerCamelCase_       =		{"bert_for_seq_generation": 5_1_2}






class  _SCREAMING_SNAKE_CASE(       A  ):
       SCREAMING_SNAKE_CASE_       :				str     						=  VOCAB_FILES_NAMES
       SCREAMING_SNAKE_CASE_       :				Optional[Any]     						=  PRETRAINED_VOCAB_FILES_MAP
       SCREAMING_SNAKE_CASE_       :				Optional[int]     						=  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       SCREAMING_SNAKE_CASE_       :				List[int]     						=  []
       SCREAMING_SNAKE_CASE_       :				List[Any]     						=  ['''input_ids''', '''attention_mask''']




       def __init__(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__="<s>"		,SCREAMING_SNAKE_CASE__="</s>"		,SCREAMING_SNAKE_CASE__="<unk>"		,SCREAMING_SNAKE_CASE__="<pad>"		,SCREAMING_SNAKE_CASE__="<::::>"		,SCREAMING_SNAKE_CASE__ = None		,**SCREAMING_SNAKE_CASE__		,)      ->				None:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  {} if sp_model_kwargs is None else sp_model_kwargs

             # Add extra_ids to the special token list
             super().__init__(
                 bos_token=SCREAMING_SNAKE_CASE__		,eos_token=SCREAMING_SNAKE_CASE__		,unk_token=SCREAMING_SNAKE_CASE__		,pad_token=SCREAMING_SNAKE_CASE__		,sep_token=SCREAMING_SNAKE_CASE__		,sp_model_kwargs=self.sp_model_kwargs		,**SCREAMING_SNAKE_CASE__		,)

             __SCREAMING_SNAKE_CASE       :Dict		    =  vocab_file

             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  spm.SentencePieceProcessor(**self.sp_model_kwargs							)
             self.sp_model.Load(SCREAMING_SNAKE_CASE__							)




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





             """simple docstring"""






             return self.sp_model.get_piece_size()




       def    _UpperCamelCase					(						self							)      ->				int:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :List[Any]		    =  {self.convert_ids_to_tokens(SCREAMING_SNAKE_CASE__							): i for i in range(self.vocab_size							)}
             vocab.update(self.added_tokens_encoder							)
             return vocab




       def __getstate__(						self							)      ->				Dict:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :int		    =  self.__dict__.copy()
             __SCREAMING_SNAKE_CASE       :Dict		    =  None
             return state




       def __setstate__(						self		,SCREAMING_SNAKE_CASE__							)      ->				Optional[Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :List[str]		    =  d

             # for backward compatibility
             if not hasattr(self		,'''sp_model_kwargs'''							):
                   __SCREAMING_SNAKE_CASE       :str		    =  {}

             __SCREAMING_SNAKE_CASE       :str		    =  spm.SentencePieceProcessor(**self.sp_model_kwargs							)
             self.sp_model.Load(self.vocab_file							)




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__							)      ->				List[str]:





             """simple docstring"""






             return self.sp_model.encode(SCREAMING_SNAKE_CASE__		,out_type=SCREAMING_SNAKE_CASE__							)




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__							)      ->				int:





             """simple docstring"""






             return self.sp_model.piece_to_id(SCREAMING_SNAKE_CASE__							)




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__							)      ->				str:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  self.sp_model.IdToPiece(SCREAMING_SNAKE_CASE__							)
             return token




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__							)      ->				List[Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :int		    =  []
             __SCREAMING_SNAKE_CASE       :List[str]		    =  ''''''
             for token in tokens:
                   # make sure that special tokens are not decoded using sentencepiece model
                   if token in self.all_special_tokens:
                         out_string += self.sp_model.decode(SCREAMING_SNAKE_CASE__							) + token
                         __SCREAMING_SNAKE_CASE       :Dict		    =  []
                   else:
                         current_sub_tokens.append(SCREAMING_SNAKE_CASE__							)
             out_string += self.sp_model.decode(SCREAMING_SNAKE_CASE__							)
             return out_string.strip()




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__ = None							)      ->				Tuple[str]:





             """simple docstring"""






             if not os.path.isdir(SCREAMING_SNAKE_CASE__							):
                   logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''							)
                   return
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  os.path.join(
                 SCREAMING_SNAKE_CASE__		,(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']							)

             if os.path.abspath(self.vocab_file							) != os.path.abspath(SCREAMING_SNAKE_CASE__							) and os.path.isfile(self.vocab_file							):
                   copyfile(self.vocab_file		,SCREAMING_SNAKE_CASE__							)
             elif not os.path.isfile(self.vocab_file							):
                   with open(SCREAMING_SNAKE_CASE__		,'''wb'''							) as fi:
                         __SCREAMING_SNAKE_CASE       :Tuple		    =  self.sp_model.serialized_model_proto()
                         fi.write(SCREAMING_SNAKE_CASE__							)

             return (out_vocab_file,)