




"""simple docstring"""


import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    VideoMAEConfig,
    VideoMAEForPreTraining,
    VideoMAEForVideoClassification,
    VideoMAEImageProcessor,
)


def       UpperCAmelCase__ (snake_case__							:      Any					):
  """simple docstring"""

  _snake_case							:	Dict			=						VideoMAEConfig()

  set_architecture_configs(snake_case__	,     snake_case__					)

  if "finetuned" not in model_name:
    _snake_case							:	List[str]			=						False

  if "finetuned" in model_name:
    _snake_case							:	str			=						"""huggingface/label-files"""
    if "kinetics" in model_name:
      _snake_case							:	List[str]			=						4_00
      _snake_case							:	Any			=						"""kinetics400-id2label.json"""
    elif "ssv2" in model_name:
      _snake_case							:	int			=						1_74
      _snake_case							:	Tuple			=						"""something-something-v2-id2label.json"""
    else:
      raise ValueError("""Model name should either contain 'kinetics' or 'ssv2' in case it's fine-tuned."""					)
    _snake_case							:	List[Any]			=						json.load(open(hf_hub_download(snake_case__	,     snake_case__	,     repo_type="""dataset"""					)	,     """r"""					)					)
    _snake_case							:	int			=						{int(snake_case__					): v for k, v in idalabel.items()}
    _snake_case							:	Any			=						idalabel
    _snake_case							:	Union[str, Any]			=						{v: k for k, v in idalabel.items()}

  return config


def       UpperCAmelCase__ (snake_case__							:      Dict	,     snake_case__							:      List[Any]					):
  """simple docstring"""

  if "small" in model_name:
    _snake_case							:	List[str]			=						3_84
    _snake_case							:	Optional[Any]			=						15_36
    _snake_case							:	Dict			=						12
    _snake_case							:	Optional[int]			=						16
    _snake_case							:	Dict			=						12
    _snake_case							:	int			=						3
    _snake_case							:	Union[str, Any]			=						1_92
    _snake_case							:	Optional[int]			=						7_68
  elif "large" in model_name:
    _snake_case							:	Dict			=						10_24
    _snake_case							:	int			=						40_96
    _snake_case							:	int			=						24
    _snake_case							:	int			=						16
    _snake_case							:	Optional[int]			=						12
    _snake_case							:	Optional[int]			=						8
    _snake_case							:	Union[str, Any]			=						5_12
    _snake_case							:	Optional[Any]			=						20_48
  elif "huge" in model_name:
    _snake_case							:	Dict			=						12_80
    _snake_case							:	Any			=						51_20
    _snake_case							:	int			=						32
    _snake_case							:	Union[str, Any]			=						16
    _snake_case							:	Union[str, Any]			=						12
    _snake_case							:	int			=						8
    _snake_case							:	List[str]			=						6_40
    _snake_case							:	List[str]			=						25_60
  elif "base" not in model_name:
    raise ValueError("""Model name should include either \"small\", \"base\", \"large\", or \"huge\""""					)


def       UpperCAmelCase__ (snake_case__							:      Tuple					):
  """simple docstring"""

  if "encoder." in name:
    _snake_case							:	int			=						name.replace("""encoder."""	,     """"""					)
  if "cls_token" in name:
    _snake_case							:	Dict			=						name.replace("""cls_token"""	,     """videomae.embeddings.cls_token"""					)
  if "decoder_pos_embed" in name:
    _snake_case							:	Dict			=						name.replace("""decoder_pos_embed"""	,     """decoder.decoder_pos_embed"""					)
  if "pos_embed" in name and "decoder" not in name:
    _snake_case							:	Tuple			=						name.replace("""pos_embed"""	,     """videomae.embeddings.position_embeddings"""					)
  if "patch_embed.proj" in name:
    _snake_case							:	Any			=						name.replace("""patch_embed.proj"""	,     """videomae.embeddings.patch_embeddings.projection"""					)
  if "patch_embed.norm" in name:
    _snake_case							:	Tuple			=						name.replace("""patch_embed.norm"""	,     """videomae.embeddings.norm"""					)
  if "decoder.blocks" in name:
    _snake_case							:	str			=						name.replace("""decoder.blocks"""	,     """decoder.decoder_layers"""					)
  if "blocks" in name:
    _snake_case							:	List[Any]			=						name.replace("""blocks"""	,     """videomae.encoder.layer"""					)
  if "attn.proj" in name:
    _snake_case							:	str			=						name.replace("""attn.proj"""	,     """attention.output.dense"""					)
  if "attn" in name and "bias" not in name:
    _snake_case							:	Dict			=						name.replace("""attn"""	,     """attention.self"""					)
  if "attn" in name:
    _snake_case							:	int			=						name.replace("""attn"""	,     """attention.attention"""					)
  if "norm1" in name:
    _snake_case							:	str			=						name.replace("""norm1"""	,     """layernorm_before"""					)
  if "norm2" in name:
    _snake_case							:	List[Any]			=						name.replace("""norm2"""	,     """layernorm_after"""					)
  if "mlp.fc1" in name:
    _snake_case							:	str			=						name.replace("""mlp.fc1"""	,     """intermediate.dense"""					)
  if "mlp.fc2" in name:
    _snake_case							:	Any			=						name.replace("""mlp.fc2"""	,     """output.dense"""					)
  if "decoder_embed" in name:
    _snake_case							:	Union[str, Any]			=						name.replace("""decoder_embed"""	,     """decoder.decoder_embed"""					)
  if "decoder_norm" in name:
    _snake_case							:	str			=						name.replace("""decoder_norm"""	,     """decoder.decoder_norm"""					)
  if "decoder_pred" in name:
    _snake_case							:	List[Any]			=						name.replace("""decoder_pred"""	,     """decoder.decoder_pred"""					)
  if "norm.weight" in name and "decoder" not in name and "fc" not in name:
    _snake_case							:	Tuple			=						name.replace("""norm.weight"""	,     """videomae.layernorm.weight"""					)
  if "norm.bias" in name and "decoder" not in name and "fc" not in name:
    _snake_case							:	Any			=						name.replace("""norm.bias"""	,     """videomae.layernorm.bias"""					)
  if "head" in name and "decoder" not in name:
    _snake_case							:	Tuple			=						name.replace("""head"""	,     """classifier"""					)

  return name


def       UpperCAmelCase__ (snake_case__							:      Optional[int]	,     snake_case__							:      List[str]					):
  """simple docstring"""

  for key in orig_state_dict.copy().keys():
    _snake_case							:	int			=						orig_state_dict.pop(snake_case__					)

    if key.startswith("""encoder."""					):
      _snake_case							:	Tuple			=						key.replace("""encoder."""	,     """"""					)

    if "qkv" in key:
      _snake_case							:	str			=						key.split("""."""					)
      if key.startswith("""decoder.blocks"""					):
        _snake_case							:	int			=						config.decoder_hidden_size
        _snake_case							:	Tuple			=						int(key_split[2]					)
        _snake_case							:	str			=						"""decoder.decoder_layers."""
        if "weight" in key:
          _snake_case							:	Optional[int]			=						val[:dim, :]
          _snake_case							:	Optional[int]			=						val[dim : dim * 2, :]
          _snake_case							:	Optional[Any]			=						val[-dim:, :]
      else:
        _snake_case							:	Union[str, Any]			=						config.hidden_size
        _snake_case							:	Union[str, Any]			=						int(key_split[1]					)
        _snake_case							:	str			=						"""videomae.encoder.layer."""
        if "weight" in key:
          _snake_case							:	List[Any]			=						val[:dim, :]
          _snake_case							:	Tuple			=						val[dim : dim * 2, :]
          _snake_case							:	Dict			=						val[-dim:, :]
    else:
      _snake_case							:	List[Any]			=						val

  return orig_state_dict


def       UpperCAmelCase__ ():
  """simple docstring"""

  _snake_case							:	Dict			=						hf_hub_download(
      repo_id="""hf-internal-testing/spaghetti-video"""	,     filename="""eating_spaghetti.npy"""	,     repo_type="""dataset"""					)
  _snake_case							:	Optional[int]			=						np.load(snake_case__					)
  return list(snake_case__					)


def       UpperCAmelCase__ (snake_case__							:      List[Any]	,     snake_case__							:      List[str]	,     snake_case__							:      Union[str, Any]	,     snake_case__							:      Tuple					):
  """simple docstring"""

  _snake_case							:	Union[str, Any]			=						get_videomae_config(snake_case__					)

  if "finetuned" in model_name:
    _snake_case							:	Tuple			=						VideoMAEForVideoClassification(snake_case__					)
  else:
    _snake_case							:	int			=						VideoMAEForPreTraining(snake_case__					)

  # download original checkpoint, hosted on Google Drive
  _snake_case							:	int			=						"""pytorch_model.bin"""
  gdown.cached_download(snake_case__	,     snake_case__	,     quiet=snake_case__					)
  _snake_case							:	List[Any]			=						torch.load(snake_case__	,     map_location="""cpu"""					)
  if "model" in files:
    _snake_case							:	int			=						files["""model"""]
  else:
    _snake_case							:	Tuple			=						files["""module"""]
  _snake_case							:	List[str]			=						convert_state_dict(snake_case__	,     snake_case__					)

  model.load_state_dict(snake_case__					)
  model.eval()

  # verify model on basic input
  _snake_case							:	List[Any]			=						VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]	,     image_std=[0.5, 0.5, 0.5]					)
  _snake_case							:	int			=						prepare_video()
  _snake_case							:	str			=						image_processor(snake_case__	,     return_tensors="""pt"""					)

  if "finetuned" not in model_name:
    _snake_case							:	int			=						hf_hub_download(repo_id="""hf-internal-testing/bool-masked-pos"""	,     filename="""bool_masked_pos.pt"""					)
    _snake_case							:	Tuple			=						torch.load(snake_case__					)

  _snake_case							:	int			=						model(**snake_case__					)
  _snake_case							:	Optional[Any]			=						outputs.logits

  _snake_case							:	Optional[Any]			=						[
      """videomae-small-finetuned-kinetics""",
      """videomae-small-finetuned-ssv2""",
      # Kinetics-400 checkpoints (short = pretrained only for 800 epochs instead of 1600)
      """videomae-base-short""",
      """videomae-base-short-finetuned-kinetics""",
      """videomae-base""",
      """videomae-base-finetuned-kinetics""",
      """videomae-large""",
      """videomae-large-finetuned-kinetics""",
      """videomae-huge-finetuned-kinetics""",
      # Something-Something-v2 checkpoints (short = pretrained only for 800 epochs instead of 2400)
      """videomae-base-short-ssv2""",
      """videomae-base-short-finetuned-ssv2""",
      """videomae-base-ssv2""",
      """videomae-base-finetuned-ssv2""",
  ]

  # NOTE: logits were tested with image_mean and image_std equal to [0.5, 0.5, 0.5] and [0.5, 0.5, 0.5]
  if model_name == "videomae-small-finetuned-kinetics":
    _snake_case							:	Tuple			=						torch.Size([1, 4_00]					)
    _snake_case							:	int			=						torch.tensor([-0.92_91, -0.40_61, -0.93_07]					)
  elif model_name == "videomae-small-finetuned-ssv2":
    _snake_case							:	Dict			=						torch.Size([1, 1_74]					)
    _snake_case							:	Any			=						torch.tensor([0.26_71, -0.46_89, -0.82_35]					)
  elif model_name == "videomae-base":
    _snake_case							:	str			=						torch.Size([1, 14_08, 15_36]					)
    _snake_case							:	List[Any]			=						torch.tensor([[0.77_39, 0.79_68, 0.70_89], [0.67_01, 0.74_87, 0.62_09], [0.42_87, 0.51_58, 0.47_73]]					)
  elif model_name == "videomae-base-short":
    _snake_case							:	int			=						torch.Size([1, 14_08, 15_36]					)
    _snake_case							:	Optional[int]			=						torch.tensor([[0.79_94, 0.96_12, 0.85_08], [0.74_01, 0.89_58, 0.83_02], [0.58_62, 0.74_68, 0.73_25]]					)
    # we verified the loss both for normalized and unnormalized targets for this one
    _snake_case							:	Any			=						torch.tensor([0.51_42]					) if config.norm_pix_loss else torch.tensor([0.64_69]					)
  elif model_name == "videomae-large":
    _snake_case							:	Optional[int]			=						torch.Size([1, 14_08, 15_36]					)
    _snake_case							:	int			=						torch.tensor([[0.71_49, 0.79_97, 0.69_66], [0.67_68, 0.78_69, 0.69_48], [0.51_39, 0.62_21, 0.56_05]]					)
  elif model_name == "videomae-large-finetuned-kinetics":
    _snake_case							:	Tuple			=						torch.Size([1, 4_00]					)
    _snake_case							:	Optional[Any]			=						torch.tensor([0.07_71, 0.00_11, -0.36_25]					)
  elif model_name == "videomae-huge-finetuned-kinetics":
    _snake_case							:	Tuple			=						torch.Size([1, 4_00]					)
    _snake_case							:	List[str]			=						torch.tensor([0.24_33, 0.16_32, -0.48_94]					)
  elif model_name == "videomae-base-short-finetuned-kinetics":
    _snake_case							:	int			=						torch.Size([1, 4_00]					)
    _snake_case							:	Dict			=						torch.tensor([0.65_88, 0.09_90, -0.24_93]					)
  elif model_name == "videomae-base-finetuned-kinetics":
    _snake_case							:	Dict			=						torch.Size([1, 4_00]					)
    _snake_case							:	Tuple			=						torch.tensor([0.36_69, -0.06_88, -0.24_21]					)
  elif model_name == "videomae-base-short-ssv2":
    _snake_case							:	List[str]			=						torch.Size([1, 14_08, 15_36]					)
    _snake_case							:	int			=						torch.tensor([[0.47_12, 0.52_96, 0.57_86], [0.22_78, 0.27_29, 0.40_26], [0.03_52, 0.07_30, 0.25_06]]					)
  elif model_name == "videomae-base-short-finetuned-ssv2":
    _snake_case							:	Optional[int]			=						torch.Size([1, 1_74]					)
    _snake_case							:	Dict			=						torch.tensor([-0.05_37, -0.15_39, -0.32_66]					)
  elif model_name == "videomae-base-ssv2":
    _snake_case							:	Tuple			=						torch.Size([1, 14_08, 15_36]					)
    _snake_case							:	Dict			=						torch.tensor([[0.81_31, 0.87_27, 0.85_46], [0.73_66, 0.93_77, 0.88_70], [0.59_35, 0.88_74, 0.85_64]]					)
  elif model_name == "videomae-base-finetuned-ssv2":
    _snake_case							:	str			=						torch.Size([1, 1_74]					)
    _snake_case							:	List[Any]			=						torch.tensor([0.19_61, -0.83_37, -0.63_89]					)
  else:
    raise ValueError(F"Model name not supported. Should be one of {model_names}"					)

  # verify logits
  assert logits.shape == expected_shape
  if "finetuned" in model_name:
    assert torch.allclose(logits[0, :3]	,     snake_case__	,     atol=1e-4					)
  else:
    print("""Logits:"""	,     logits[0, :3, :3]					)
    assert torch.allclose(logits[0, :3, :3]	,     snake_case__	,     atol=1e-4					)
  print("""Logits ok!"""					)

  # verify loss, if applicable
  if model_name == "videomae-base-short":
    _snake_case							:	Union[str, Any]			=						outputs.loss
    assert torch.allclose(snake_case__	,     snake_case__	,     atol=1e-4					)
    print("""Loss ok!"""					)

  if pytorch_dump_folder_path is not None:
    print(F"Saving model and image processor to {pytorch_dump_folder_path}"					)
    image_processor.save_pretrained(snake_case__					)
    model.save_pretrained(snake_case__					)

  if push_to_hub:
    print("""Pushing to the hub..."""					)
    model.push_to_hub(snake_case__	,     organization="""nielsr"""					)


if __name__ == "__main__":
     A_      			=       argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '''--checkpoint_url''',
         default='''https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4''',
         type=str,
         help=(
             '''URL of the original PyTorch checkpoint (on Google Drive) you\'d like to convert. Should be a direct'''
             ''' download link.'''
         ),
     )
     parser.add_argument(
         '''--pytorch_dump_folder_path''',
         default='''/Users/nielsrogge/Documents/VideoMAE/Test''',
         type=str,
         help='''Path to the output PyTorch model directory.''',
     )
     parser.add_argument('''--model_name''', default='''videomae-base''', type=str, help='''Name of the model.''')
     parser.add_argument(
         '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
     )

     A_      			=       parser.parse_args()
     convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)








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







 '''simple docstring'''





 def __init__(       self:	List[Any],   a_:	str = "cpu",   a_:	str = "openai/clip-vit-large-patch14"	):




   '''simple docstring'''



   _snake_case							:	Optional[int]			=						device
   _snake_case							:	str			=						CLIPTokenizerFast.from_pretrained(a_	)
   _snake_case							:	Union[str, Any]			=						[0.48_145_466, 0.4_578_275, 0.40_821_073]
   _snake_case							:	Optional[int]			=						[0.26_862_954, 0.26_130_258, 0.27_577_711]
   _snake_case							:	str			=						torchvision.transforms.Normalize(self.image_mean,   self.image_std	)
   _snake_case							:	Optional[int]			=						torchvision.transforms.Resize(224	)
   _snake_case							:	str			=						torchvision.transforms.CenterCrop(224	)


 def 				UpperCamelCase_     (       self:	List[str],   a_:	str	):




   '''simple docstring'''



   _snake_case							:	Optional[int]			=						self.resize(a_	)
   _snake_case							:	List[Any]			=						self.center_crop(a_	)
   _snake_case							:	Optional[Any]			=						self.normalize(a_	)
   return images


 def __call__(       self:	Any,   a_:	Optional[int]=None,   a_:	str=None,   **a_:	str	):




   '''simple docstring'''



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







class       lowercase(  nn.Module				):







 '''simple docstring'''





 def __init__(       self:	List[Any],   a_:	List[Any]=10,   a_:	Optional[Any]=0.01,   a_:	List[str]=None,   a_:	str=None,   a_:	Any=None,   a_:	Tuple=None,   a_:	List[str]=None,   a_:	List[str]=None,   a_:	str=False,   a_:	List[str]=True,   a_:	Any="image",   a_:	Optional[Any]=True,   a_:	Dict=False,   a_:	List[str]=False,   a_:	Optional[int]=False,   ):




   '''simple docstring'''



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

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


 def 				UpperCamelCase_     (       self:	Tuple,   a_:	str=None,   a_:	Dict=None,   a_:	Dict=5,   a_:	Dict=True	):




   '''simple docstring'''



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


 def 				UpperCamelCase_     (       self:	str,   a_:	Tuple=None,   a_:	Optional[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							:	int			=						preprocess(Image.open(a_	),   target_image_size=256	).to(self.device	)
   _snake_case							:	int			=						preprocess_vqgan(a_	)
   _snake_case	,      *_snake_case							:	List[Any]			=						self.vqgan.encode(a_	)
   return z


 def 				UpperCamelCase_     (       self:	Union[str, Any],   a_:	Optional[Any]	):




   '''simple docstring'''



   _snake_case							:	Optional[int]			=						self.latent.detach().requires_grad_()
   _snake_case							:	Tuple			=						base_latent + transform_vector
   if self.quantize:
     _snake_case	,      *_snake_case							:	Any			=						self.vqgan.quantize(a_	)
   else:
     _snake_case							:	List[Any]			=						trans_latent
   return self.vqgan.decode(a_	)


 def 				UpperCamelCase_     (       self:	List[Any],   a_:	Any,   a_:	Union[str, Any],   a_:	Dict=None	):




   '''simple docstring'''



   _snake_case							:	Tuple			=						self.clip_preprocessor(text=a_,   images=a_,   return_tensors="""pt""",   padding=a_	)
   _snake_case							:	Any			=						self.clip(**a_	)
   _snake_case							:	str			=						clip_outputs.logits_per_image
   if weights is not None:
     _snake_case							:	Any			=						similarity_logits * weights
   return similarity_logits.sum()


 def 				UpperCamelCase_     (       self:	Any,   a_:	Any,   a_:	List[str],   a_:	Dict	):




   '''simple docstring'''



   _snake_case							:	List[Any]			=						self._get_clip_similarity(pos_prompts["""prompts"""],   a_,   weights=(1 / pos_prompts["""weights"""])	)
   if neg_prompts:
     _snake_case							:	List[str]			=						self._get_clip_similarity(neg_prompts["""prompts"""],   a_,   weights=neg_prompts["""weights"""]	)
   else:
     _snake_case							:	Tuple			=						torch.tensor([1],   device=self.device	)
   _snake_case							:	int			=						-torch.log(a_	) + torch.log(a_	)
   return loss


 def 				UpperCamelCase_     (       self:	Optional[Any],   a_:	Tuple,   a_:	Union[str, Any],   a_:	List[str]	):




   '''simple docstring'''



   _snake_case							:	Tuple			=						torch.randn_like(self.latent,   requires_grad=a_,   device=self.device	)
   _snake_case							:	Dict			=						torch.optim.Adam([vector],   lr=self.lr	)

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


 def 				UpperCamelCase_     (       self:	int,   a_:	Any,   a_:	Union[str, Any],   a_:	Optional[int]	):




   '''simple docstring'''



   wandb.init(reinit=a_,   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							:	Any			=						Image.open(a_	)
     _snake_case							:	str			=						image.resize((256, 256)	)
     wandb.log("""Original Image""",   wandb.Image(a_	)	)


 def 				UpperCamelCase_     (       self:	str,   a_:	List[Any]	):




   '''simple docstring'''



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


 def 				UpperCamelCase_     (       self:	Dict,   a_:	List[Any],   a_:	List[Any]=None,   a_:	Optional[Any]=None,   a_:	Optional[Any]=True,   a_:	Dict=False,   a_:	Optional[Any]=True,   a_:	Optional[Any]=True,   a_:	Any=None,   ):




   '''simple docstring'''



   if image_path:
     _snake_case							:	Union[str, Any]			=						self._get_latent(a_	)
   else:
     _snake_case							:	Any			=						torch.randn(self.latent_dim,   device=self.device	)
   if self.log:
     self._init_logging(a_,   a_,   a_	)

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

   _snake_case							:	List[Any]			=						self.vqgan.decode(self.latent	)[0]
   if show_intermediate:
     print("""Original Image"""	)
     show_pil(custom_to_pil(a_	)	)

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


