


'''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    __magic_name__				:






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




          '''simple docstring'''
          _UpperCAmelCase  :  Dict				       =device
          _UpperCAmelCase  :  int				       =CLIPTokenizerFast.from_pretrained(snake_case)
          _UpperCAmelCase  :  List[Any]				       =[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]
          _UpperCAmelCase  :  List[str]				       =[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]
          _UpperCAmelCase  :  Dict				       =torchvision.transforms.Normalize(self.image_mean    ,				self.image_std)
          _UpperCAmelCase  :  List[str]				       =torchvision.transforms.Resize(2_2_4)
          _UpperCAmelCase  :  Dict				       =torchvision.transforms.CenterCrop(2_2_4)






       def      lowerCAmelCase						(     self    ,				snake_case)		->    str:




          '''simple docstring'''
          _UpperCAmelCase  :  Optional[int]				       =self.resize(snake_case)
          _UpperCAmelCase  :  Union[str, Any]				       =self.center_crop(snake_case)
          _UpperCAmelCase  :  Optional[Any]				       =self.normalize(snake_case)
          return images






       def __call__(     self    ,				snake_case=None    ,				snake_case=None    ,				**snake_case)		->    Optional[Any]:




          '''simple docstring'''
          _UpperCAmelCase  :  Any				       =self.tokenizer(text=snake_case    ,				**snake_case)
          _UpperCAmelCase  :  Any				       =self.preprocess_img(snake_case)
          _UpperCAmelCase  :  Tuple				       ={key: value.to(self.device) for (key, value) in encoding.items()}
          return encoding






class    __magic_name__				(						nn.Module							):






       def __init__(     self    ,				snake_case=1_0    ,				snake_case=0.01    ,				snake_case=None    ,				snake_case=None    ,				snake_case=None    ,				snake_case=None    ,				snake_case=None    ,				snake_case=None    ,				snake_case=False    ,				snake_case=True    ,				snake_case="image"    ,				snake_case=True    ,				snake_case=False    ,				snake_case=False    ,				snake_case=False    ,				)		->    None:




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

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






       def      lowerCAmelCase						(     self    ,				snake_case=None    ,				snake_case=None    ,				snake_case=5    ,				snake_case=True)		->    Optional[int]:




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






       def      lowerCAmelCase						(     self    ,				snake_case=None    ,				snake_case=None)		->    Any:




          '''simple docstring'''
          if not (path or img):
             raise ValueError('Input either path or tensor')
          if img is not None:
             raise NotImplementedError
          _UpperCAmelCase  :  Dict				       =preprocess(Image.open(snake_case)    ,				target_image_size=2_5_6).to(self.device)
          _UpperCAmelCase  :  Optional[Any]				       =preprocess_vqgan(snake_case)
          _UpperCAmelCase       ,			*_UpperCAmelCase  :  List[str]				       =self.vqgan.encode(snake_case)
          return z






       def      lowerCAmelCase						(     self    ,				snake_case)		->    Optional[Any]:




          '''simple docstring'''
          _UpperCAmelCase  :  Any				       =self.latent.detach().requires_grad_()
          _UpperCAmelCase  :  Optional[int]				       =base_latent + transform_vector
          if self.quantize:
             _UpperCAmelCase       ,			*_UpperCAmelCase  :  Any				       =self.vqgan.quantize(snake_case)
          else:
             _UpperCAmelCase  :  List[Any]				       =trans_latent
          return self.vqgan.decode(snake_case)






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case    ,				snake_case=None)		->    str:




          '''simple docstring'''
          _UpperCAmelCase  :  Tuple				       =self.clip_preprocessor(text=snake_case    ,				images=snake_case    ,				return_tensors='pt'    ,				padding=snake_case)
          _UpperCAmelCase  :  int				       =self.clip(**snake_case)
          _UpperCAmelCase  :  Dict				       =clip_outputs.logits_per_image
          if weights is not None:
             _UpperCAmelCase  :  int				       =similarity_logits * weights
          return similarity_logits.sum()






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case    ,				snake_case)		->    Optional[Any]:




          '''simple docstring'''
          _UpperCAmelCase  :  Tuple				       =self._get_clip_similarity(pos_prompts['prompts']    ,				snake_case    ,				weights=(1 / pos_prompts['weights']))
          if neg_prompts:
             _UpperCAmelCase  :  Union[str, Any]				       =self._get_clip_similarity(neg_prompts['prompts']    ,				snake_case    ,				weights=neg_prompts['weights'])
          else:
             _UpperCAmelCase  :  List[Any]				       =torch.tensor([1]    ,				device=self.device)
          _UpperCAmelCase  :  str				       =-torch.log(snake_case) + torch.log(snake_case)
          return loss






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case    ,				snake_case)		->    List[Any]:




          '''simple docstring'''
          _UpperCAmelCase  :  Dict				       =torch.randn_like(self.latent    ,				requires_grad=snake_case    ,				device=self.device)
          _UpperCAmelCase  :  Dict				       =torch.optim.Adam([vector]    ,				lr=self.lr)

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






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case    ,				snake_case)		->    Union[str, Any]:




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






       def      lowerCAmelCase						(     self    ,				snake_case)		->    Optional[int]:




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






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case=None    ,				snake_case=None    ,				snake_case=True    ,				snake_case=False    ,				snake_case=True    ,				snake_case=True    ,				snake_case=None    ,				)		->    List[Any]:




          '''simple docstring'''
          if image_path:
             _UpperCAmelCase  :  List[str]				       =self._get_latent(snake_case)
          else:
             _UpperCAmelCase  :  List[Any]				       =torch.randn(self.latent_dim    ,				device=self.device)
          if self.log:
             self._init_logging(snake_case    ,				snake_case    ,				snake_case)

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

          _UpperCAmelCase  :  Optional[int]				       =self.vqgan.decode(self.latent)[0]
          if show_intermediate:
             print('Original Image')
             show_pil(custom_to_pil(snake_case))

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





'''simple docstring'''






import math
import time
from typing import Dict, List, Optional

from torch.utils.data import Dataset

from transformers import SeqaSeqTrainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
 import torch_xla.core.xla_model as xm
 import torch_xla.debug.metrics as met






class    __magic_name__				(						lowerCAmelCase							):






       def __init__(     self    ,				*snake_case    ,				snake_case=None    ,				snake_case=None    ,				**snake_case)		->    Tuple:




          '''simple docstring'''
          super().__init__(*snake_case    ,				**snake_case)
          _UpperCAmelCase  :  Any				       =eval_examples
          _UpperCAmelCase  :  Optional[int]				       =post_process_function






       def      lowerCAmelCase						(     self    ,				snake_case = None    ,				snake_case=None    ,				snake_case = None    ,				snake_case = "eval"    ,				**snake_case    ,				)		->    Dict[str, float]:




          '''simple docstring'''
          _UpperCAmelCase  :  Tuple				       =gen_kwargs.copy()
          _UpperCAmelCase  :  int				       =(
              gen_kwargs['max_length'] if gen_kwargs.get('max_length') is not None else self.args.generation_max_length
          )
          _UpperCAmelCase  :  Dict				       =(
              gen_kwargs['num_beams'] if gen_kwargs.get('num_beams') is not None else self.args.generation_num_beams
          )
          _UpperCAmelCase  :  Dict				       =gen_kwargs

          _UpperCAmelCase  :  Tuple				       =self.eval_dataset if eval_dataset is None else eval_dataset
          _UpperCAmelCase  :  Union[str, Any]				       =self.get_eval_dataloader(snake_case)
          _UpperCAmelCase  :  Tuple				       =self.eval_examples if eval_examples is None else eval_examples

          # Temporarily disable metric computation, we will do it in the loop here.
          _UpperCAmelCase  :  Dict				       =self.compute_metrics
          _UpperCAmelCase  :  Optional[Any]				       =None
          _UpperCAmelCase  :  str				       =time.time()
          _UpperCAmelCase  :  int				       =self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
          try:
             _UpperCAmelCase  :  str				       =eval_loop(
                 snake_case    ,				description='Evaluation'    ,				prediction_loss_only=True if compute_metrics is None else None    ,				ignore_keys=snake_case    ,				metric_key_prefix=snake_case    ,				)
          finally:
             _UpperCAmelCase  :  str				       =compute_metrics
          _UpperCAmelCase  :  Optional[int]				       =self.args.eval_batch_size * self.args.world_size
          if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
             start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
          output.metrics.update(
              speed_metrics(
                  snake_case    ,				snake_case    ,				num_samples=output.num_samples    ,				num_steps=math.ceil(output.num_samples / total_batch_size)    ,				))

          if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
             # Only the main node write the results by default
             _UpperCAmelCase  :  Tuple				       =self.post_process_function(snake_case    ,				snake_case    ,				snake_case)
             _UpperCAmelCase  :  Union[str, Any]				       =self.compute_metrics(snake_case)

             # Prefix all keys with metric_key_prefix + '_'
             for key in list(metrics.keys()):
                if not key.startswith(f"{metric_key_prefix}_"):
                   _UpperCAmelCase  :  Any				       =metrics.pop(snake_case)

             metrics.update(output.metrics)
          else:
             _UpperCAmelCase  :  str				       =output.metrics

          if self.args.should_log:
             # Only the main node log the results by default
             self.log(snake_case)

          if self.args.tpu_metrics_debug or self.args.debug:
             # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
             xm.master_print(met.metrics_report())

          _UpperCAmelCase  :  Any				       =self.callback_handler.on_evaluate(self.args    ,				self.state    ,				self.control    ,				snake_case)
          return metrics






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case    ,				snake_case=None    ,				snake_case = "test"    ,				**snake_case)		->    Dict:




          '''simple docstring'''
          _UpperCAmelCase  :  Optional[int]				       =gen_kwargs.copy()

          _UpperCAmelCase  :  int				       =self.get_test_dataloader(snake_case)

          # Temporarily disable metric computation, we will do it in the loop here.
          _UpperCAmelCase  :  str				       =self.compute_metrics
          _UpperCAmelCase  :  Union[str, Any]				       =None
          _UpperCAmelCase  :  Union[str, Any]				       =time.time()
          _UpperCAmelCase  :  Tuple				       =self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
          try:
             _UpperCAmelCase  :  Any				       =eval_loop(
                 snake_case    ,				description='Prediction'    ,				prediction_loss_only=True if compute_metrics is None else None    ,				ignore_keys=snake_case    ,				metric_key_prefix=snake_case    ,				)
          finally:
             _UpperCAmelCase  :  Any				       =compute_metrics

          _UpperCAmelCase  :  Dict				       =self.args.eval_batch_size * self.args.world_size
          if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
             start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
          output.metrics.update(
              speed_metrics(
                  snake_case    ,				snake_case    ,				num_samples=output.num_samples    ,				num_steps=math.ceil(output.num_samples / total_batch_size)    ,				))
          if self.post_process_function is None or self.compute_metrics is None:
             return output

          _UpperCAmelCase  :  Tuple				       =self.post_process_function(snake_case    ,				snake_case    ,				snake_case    ,				'predict')
          _UpperCAmelCase  :  Union[str, Any]				       =self.compute_metrics(snake_case)

          # Prefix all keys with metric_key_prefix + '_'
          for key in list(metrics.keys()):
             if not key.startswith(f"{metric_key_prefix}_"):
                _UpperCAmelCase  :  Any				       =metrics.pop(snake_case)
          metrics.update(output.metrics)
          return PredictionOutput(predictions=predictions.predictions    ,				label_ids=predictions.label_ids    ,				metrics=snake_case)

