



"""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     UpperCAmelCase    :



   def __init__(       self    :	Dict							,       __lowerCamelCase    :	str = "cpu"							,       __lowerCamelCase    :	str = "openai/clip-vit-large-patch14"				):



     """simple docstring"""
     _snake_case   	=      device
     _snake_case   	=      CLIPTokenizerFast.from_pretrained(__lowerCamelCase				)
     _snake_case   	=      [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]
     _snake_case   	=      [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]
     _snake_case   	=      torchvision.transforms.Normalize(self.image_mean							,       self.image_std				)
     _snake_case   	=      torchvision.transforms.Resize(2_2_4				)
     _snake_case   	=      torchvision.transforms.CenterCrop(2_2_4				)



   def    __UpperCAmelCase    (       self    :	int							,       __lowerCamelCase    :	Optional[Any]				):



     """simple docstring"""
     _snake_case   	=      self.resize(__lowerCamelCase				)
     _snake_case   	=      self.center_crop(__lowerCamelCase				)
     _snake_case   	=      self.normalize(__lowerCamelCase				)
     return images





   def __call__(       self    :	Optional[int]							,       __lowerCamelCase    :	List[Any]=None							,       __lowerCamelCase    :	int=None							,       **__lowerCamelCase    :	List[Any]				):



     """simple docstring"""
     _snake_case   	=      self.tokenizer(text=__lowerCamelCase							,       **__lowerCamelCase				)
     _snake_case   	=      self.preprocess_img(__lowerCamelCase				)
     _snake_case   	=      {key: value.to(self.device				) for (key, value) in encoding.items()}
     return encoding



class     UpperCAmelCase    (	nn.Module  ):



   def __init__(       self    :	Union[str, Any]							,       __lowerCamelCase    :	Union[str, Any]=1_0							,       __lowerCamelCase    :	Optional[Any]=0.0_1							,       __lowerCamelCase    :	Dict=None							,       __lowerCamelCase    :	int=None							,       __lowerCamelCase    :	Optional[Any]=None							,       __lowerCamelCase    :	Optional[Any]=None							,       __lowerCamelCase    :	int=None							,       __lowerCamelCase    :	Tuple=None							,       __lowerCamelCase    :	List[Any]=False							,       __lowerCamelCase    :	int=True							,       __lowerCamelCase    :	int="image"							,       __lowerCamelCase    :	Union[str, Any]=True							,       __lowerCamelCase    :	Optional[Any]=False							,       __lowerCamelCase    :	Tuple=False							,       __lowerCamelCase    :	Optional[int]=False							,       ):



     """simple docstring"""
     super().__init__()
     _snake_case   	=      None
     _snake_case   	=      device if device else get_device()
     if vqgan:
       _snake_case   	=      vqgan
     else:
       _snake_case   	=      load_vqgan(self.device							,       conf_path=__lowerCamelCase							,       ckpt_path=__lowerCamelCase				)
     self.vqgan.eval()
     if clip:
       _snake_case   	=      clip
     else:
       _snake_case   	=      CLIPModel.from_pretrained('''openai/clip-vit-base-patch32'''				)
     self.clip.to(self.device				)
     _snake_case   	=      ProcessorGradientFlow(device=self.device				)

     _snake_case   	=      iterations
     _snake_case   	=      lr
     _snake_case   	=      log
     _snake_case   	=      make_grid
     _snake_case   	=      return_val
     _snake_case   	=      quantize
     _snake_case   	=      self.vqgan.decoder.z_shape



   def    __UpperCAmelCase    (       self    :	int							,       __lowerCamelCase    :	Optional[int]=None							,       __lowerCamelCase    :	Any=None							,       __lowerCamelCase    :	Optional[Any]=5							,       __lowerCamelCase    :	Tuple=True				):



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



   def    __UpperCAmelCase    (       self    :	Optional[int]							,       __lowerCamelCase    :	List[str]=None							,       __lowerCamelCase    :	Any=None				):



     """simple docstring"""
     if not (path or img):
       raise ValueError('''Input either path or tensor'''				)
     if img is not None:
       raise NotImplementedError
     _snake_case   	=      preprocess(Image.open(__lowerCamelCase				)							,       target_image_size=2_5_6				).to(self.device				)
     _snake_case   	=      preprocess_vqgan(__lowerCamelCase				)
     _snake_case	,     *_snake_case   	=      self.vqgan.encode(__lowerCamelCase				)
     return z



   def    __UpperCAmelCase    (       self    :	List[Any]							,       __lowerCamelCase    :	Dict				):



     """simple docstring"""
     _snake_case   	=      self.latent.detach().requires_grad_()
     _snake_case   	=      base_latent + transform_vector
     if self.quantize:
       _snake_case	,     *_snake_case   	=      self.vqgan.quantize(__lowerCamelCase				)
     else:
       _snake_case   	=      trans_latent
     return self.vqgan.decode(__lowerCamelCase				)



   def    __UpperCAmelCase    (       self    :	Optional[Any]							,       __lowerCamelCase    :	Tuple							,       __lowerCamelCase    :	List[Any]							,       __lowerCamelCase    :	Any=None				):



     """simple docstring"""
     _snake_case   	=      self.clip_preprocessor(text=__lowerCamelCase							,       images=__lowerCamelCase							,       return_tensors='''pt'''							,       padding=__lowerCamelCase				)
     _snake_case   	=      self.clip(**__lowerCamelCase				)
     _snake_case   	=      clip_outputs.logits_per_image
     if weights is not None:
       _snake_case   	=      similarity_logits * weights
     return similarity_logits.sum()



   def    __UpperCAmelCase    (       self    :	Optional[int]							,       __lowerCamelCase    :	str							,       __lowerCamelCase    :	List[Any]							,       __lowerCamelCase    :	Tuple				):



     """simple docstring"""
     _snake_case   	=      self._get_clip_similarity(pos_prompts['''prompts''']							,       __lowerCamelCase							,       weights=(1 / pos_prompts['''weights'''])				)
     if neg_prompts:
       _snake_case   	=      self._get_clip_similarity(neg_prompts['''prompts''']							,       __lowerCamelCase							,       weights=neg_prompts['''weights''']				)
     else:
       _snake_case   	=      torch.tensor([1]							,       device=self.device				)
     _snake_case   	=      -torch.log(__lowerCamelCase				) + torch.log(__lowerCamelCase				)
     return loss



   def    __UpperCAmelCase    (       self    :	Optional[int]							,       __lowerCamelCase    :	Optional[int]							,       __lowerCamelCase    :	Optional[int]							,       __lowerCamelCase    :	Dict				):



     """simple docstring"""
     _snake_case   	=      torch.randn_like(self.latent							,       requires_grad=__lowerCamelCase							,       device=self.device				)
     _snake_case   	=      torch.optim.Adam([vector]							,       lr=self.lr				)

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



   def    __UpperCAmelCase    (       self    :	Optional[Any]							,       __lowerCamelCase    :	Optional[int]							,       __lowerCamelCase    :	Tuple							,       __lowerCamelCase    :	List[str]				):



     """simple docstring"""
     wandb.init(reinit=__lowerCamelCase							,       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:
       _snake_case   	=      Image.open(__lowerCamelCase				)
       _snake_case   	=      image.resize((2_5_6, 2_5_6)				)
       wandb.log('''Original Image'''							,       wandb.Image(__lowerCamelCase				)				)



   def    __UpperCAmelCase    (       self    :	Dict							,       __lowerCamelCase    :	Any				):



     """simple docstring"""
     if not prompts:
       return []
     _snake_case   	=      []
     _snake_case   	=      []
     if isinstance(__lowerCamelCase							,       __lowerCamelCase				):
       _snake_case   	=      [prompt.strip() for prompt in prompts.split('''|'''				)]
     for prompt in prompts:
       if isinstance(__lowerCamelCase							,       (tuple, list)				):
         _snake_case   	=      prompt[0]
         _snake_case   	=      float(prompt[1]				)
       elif ":" in prompt:
         _snake_case	,     _snake_case   	=      prompt.split(''':'''				)
         _snake_case   	=      float(__lowerCamelCase				)
       else:
         _snake_case   	=      prompt
         _snake_case   	=      1.0
       processed_prompts.append(__lowerCamelCase				)
       weights.append(__lowerCamelCase				)
     return {
         "prompts": processed_prompts,
         "weights": torch.tensor(__lowerCamelCase							,       device=self.device				),
     }





   def    __UpperCAmelCase    (       self    :	int							,       __lowerCamelCase    :	Any							,       __lowerCamelCase    :	Tuple=None							,       __lowerCamelCase    :	List[Any]=None							,       __lowerCamelCase    :	Dict=True							,       __lowerCamelCase    :	int=False							,       __lowerCamelCase    :	Optional[int]=True							,       __lowerCamelCase    :	Union[str, Any]=True							,       __lowerCamelCase    :	Optional[int]=None							,       ):



     """simple docstring"""
     if image_path:
       _snake_case   	=      self._get_latent(__lowerCamelCase				)
     else:
       _snake_case   	=      torch.randn(self.latent_dim							,       device=self.device				)
     if self.log:
       self._init_logging(__lowerCamelCase							,       __lowerCamelCase							,       __lowerCamelCase				)

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

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

     _snake_case   	=      loop_post_process(__lowerCamelCase				)
     for iter, transformed_img in enumerate(self._optimize_CLIP(__lowerCamelCase							,       __lowerCamelCase							,       __lowerCamelCase				)				):
       if show_intermediate:
         show_pil(__lowerCamelCase				)
       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(__lowerCamelCase				)}				)
     if show_final:
       show_pil(__lowerCamelCase				)
     if save_final:
       transformed_img.save(os.path.join(self.save_path							,       f"""iter_{iter:03d}_final.png"""				)				)













'''simple docstring'''
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase__			=      {
    'configuration_funnel': ['FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP', 'FunnelConfig'],
    'convert_funnel_original_tf_checkpoint_to_pytorch': [],
    'tokenization_funnel': ['FunnelTokenizer'],
}

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCamelCase__			=      ['FunnelTokenizerFast']

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCamelCase__			=      [
     'FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST',
     'FunnelBaseModel',
     'FunnelForMaskedLM',
     'FunnelForMultipleChoice',
     'FunnelForPreTraining',
     'FunnelForQuestionAnswering',
     'FunnelForSequenceClassification',
     'FunnelForTokenClassification',
     'FunnelModel',
     'FunnelPreTrainedModel',
     'load_tf_weights_in_funnel',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCamelCase__			=      [
     'TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST',
     'TFFunnelBaseModel',
     'TFFunnelForMaskedLM',
     'TFFunnelForMultipleChoice',
     'TFFunnelForPreTraining',
     'TFFunnelForQuestionAnswering',
     'TFFunnelForSequenceClassification',
     'TFFunnelForTokenClassification',
     'TFFunnelModel',
     'TFFunnelPreTrainedModel',
 ]


if TYPE_CHECKING:
 from .configuration_funnel import FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP, FunnelConfig
 from .tokenization_funnel import FunnelTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_funnel_fast import FunnelTokenizerFast

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_funnel import (
      FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
      FunnelBaseModel,
      FunnelForMaskedLM,
      FunnelForMultipleChoice,
      FunnelForPreTraining,
      FunnelForQuestionAnswering,
      FunnelForSequenceClassification,
      FunnelForTokenClassification,
      FunnelModel,
      FunnelPreTrainedModel,
      load_tf_weights_in_funnel,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_funnel import (
      TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFFunnelBaseModel,
      TFFunnelForMaskedLM,
      TFFunnelForMultipleChoice,
      TFFunnelForPreTraining,
      TFFunnelForQuestionAnswering,
      TFFunnelForSequenceClassification,
      TFFunnelForTokenClassification,
      TFFunnelModel,
      TFFunnelPreTrainedModel,
  )

else:
 import sys

 UpperCamelCase__			=      _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




