this wraps the huggingface clip processor to allow backprop through the image processing step the original processor forces conversion to pil images which is faster for image processing but breaks gradient flow we call the original processor to get the text embeddings but use our own image processing to keep images as torch tensors instantiate a vqganclip model if you want to use a custom vqgan model pass it as vqgan make an animation from the intermediate images saved during generation by default uses the images from the most recent generation the generate function if you want to use images from a different generation pass the path to the folder containing the images as inputpath add a vector transform to the base latent and returns the resulting image baselatent self latent detach requiresgrad translatent baselatent transformvector if self quantize zq self vqgan quantizetranslatent else zq translatent return self vqgan decodezq def getclipsimilarityself prompts image weightsnone clipinputs self clippreprocessortextprompts imagesimage returntensorspt paddingtrue clipoutputs self clipclipinputs similaritylogits clipoutputs logitsperimage if weights is not none similaritylogits similaritylogits weights return similaritylogits sum def getcliplossself posprompts negprompts image poslogits self getclipsimilaritypospromptsprompts image weights1 pospromptsweights if negprompts neglogits self getclipsimilaritynegpromptsprompts image weightsnegpromptsweights else neglogits torch tensor1 deviceself device loss torch logposlogits torch logneglogits return loss def optimizeclipself originalimg posprompts negprompts vector torch randnlikeself latent requiresgradtrue deviceself device optim torch optim adamvector lrself lr for i in rangeself iterations optim zerograd transformedimg self addvectorvector processedimg looppostprocesstransformedimg cliploss self getcliplossposprompts negprompts processedimg printclip loss cliploss if self log wandb logclip loss cliploss cliploss backwardretaingraphtrue optim step if self returnval image yield customtopiltransformedimg0 else yield vector def initloggingself positiveprompts negativeprompts imagepath wandb initreinittrue projectfaceeditor wandb config updatepositive prompts positiveprompts wandb config updatenegative prompts negativeprompts wandb config updatelr self lr iterations self iterations if imagepath image image openimagepath image image resize256 256 wandb logoriginal image wandb imageimage def processpromptsself prompts if not prompts return processedprompts weights if isinstanceprompts str prompts prompt strip for prompt in prompts split for prompt in prompts if isinstanceprompt tuple list processedprompt prompt0 weight floatprompt1 elif in prompt processedprompt weight prompt split weight floatweight else processedprompt prompt weight 1 0 processedprompts appendprocessedprompt weights appendweight return prompts processedprompts weights torch tensorweights deviceself device def generate self posprompts negpromptsnone imagepathnone showintermediatetrue saveintermediatefalse showfinaltrue savefinaltrue savepathnone if imagepath self latent self getlatentimagepath else self latent torch randnself latentdim deviceself device if self log self initloggingposprompts negprompts imagepath assert posprompts you must provide at least one positive prompt posprompts self processpromptsposprompts negprompts self processpromptsnegprompts if savefinal and savepath is none savepath os path join outputs joinpospromptsprompts if not os path existssavepath os makedirssavepath else savepath savepath gettimestamp os makedirssavepath self savepath savepath originalimg self vqgan decodeself latent0 if showintermediate printoriginal image showpilcustomtopiloriginalimg originalimg looppostprocessoriginalimg for iter transformedimg in enumerateself optimizecliporiginalimg posprompts negprompts if showintermediate showpiltransformedimg if saveintermediate transformedimg saveos path joinself savepath fiteriter 03d png if self log wandb logimage wandb imagetransformedimg if showfinal showpiltransformedimg if savefinal transformedimg saveos path joinself savepath fiteriter 03dfinal png this wraps the huggingface clip processor to allow backprop through the image processing step the original processor forces conversion to pil images which is faster for image processing but breaks gradient flow we call the original processor to get the text embeddings but use our own image processing to keep images as torch tensors instantiate a vqgan_clip model if you want to use a custom vqgan model pass it as vqgan make an animation from the intermediate images saved during generation by default uses the images from the most recent generation the generate function if you want to use images from a different generation pass the path to the folder containing the images as input_path add a vector transform to the base latent and returns the resulting image generate an image from the given prompts if image_path is provided the image is used as a starting point for the optimization if image_path is not provided a random latent vector is used as a starting point you must provide at least one positive prompt and optionally provide negative prompts prompts must be formatted in one of the following ways a single prompt as a string e g a smiling woman a set of prompts separated by pipes a smiling woman a woman with brown hair a set of prompts and their weights separated by colons a smiling woman 1 a woman with brown hair 3 default weight is 1 a list of prompts e g a smiling woman a woman with brown hair a list of prompts and weights e g a smiling woman 1 a woman with brown hair 3
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 ProcessorGradientFlow:
    

    def __init__(self, device: str = "cpu", clip_model: str = "openai/clip-vit-large-patch14") -> None:
        self.device = device
        self.tokenizer = CLIPTokenizerFast.from_pretrained(clip_model)
        self.image_mean = [0.48145466, 0.4578275, 0.40821073]
        self.image_std = [0.26862954, 0.26130258, 0.27577711]
        self.normalize = torchvision.transforms.Normalize(self.image_mean, self.image_std)
        self.resize = torchvision.transforms.Resize(224)
        self.center_crop = torchvision.transforms.CenterCrop(224)

    def preprocess_img(self, images):
        images = self.resize(images)
        images = self.center_crop(images)
        images = self.normalize(images)
        return images

    def __call__(self, text=None, images=None, **kwargs):
        encoding = self.tokenizer(text=text, **kwargs)
        encoding["pixel_values"] = self.preprocess_img(images)
        encoding = {key: value.to(self.device) for (key, value) in encoding.items()}
        return encoding


class VQGAN_CLIP(nn.Module):
    def __init__(
        self,
        iterations=10,
        lr=0.01,
        vqgan=None,
        vqgan_config=None,
        vqgan_checkpoint=None,
        clip=None,
        clip_preprocessor=None,
        device=None,
        log=False,
        save_vector=True,
        return_val="image",
        quantize=True,
        save_intermediate=False,
        show_intermediate=False,
        make_grid=False,
    ) -> None:
        
        super().__init__()
        self.latent = None
        self.device = device if device else get_device()
        if vqgan:
            self.vqgan = vqgan
        else:
            self.vqgan = load_vqgan(self.device, conf_path=vqgan_config, ckpt_path=vqgan_checkpoint)
        self.vqgan.eval()
        if clip:
            self.clip = clip
        else:
            self.clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
        self.clip.to(self.device)
        self.clip_preprocessor = ProcessorGradientFlow(device=self.device)

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

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

    def _get_latent(self, path=None, img=None):
        if not (path or img):
            raise ValueError("Input either path or tensor")
        if img is not None:
            raise NotImplementedError
        x = preprocess(Image.open(path), target_image_size=256).to(self.device)
        x_processed = preprocess_vqgan(x)
        z, *_ = self.vqgan.encode(x_processed)
        return z

    def _add_vector(self, transform_vector):
        
        base_latent = self.latent.detach().requires_grad_()
        trans_latent = base_latent + transform_vector
        if self.quantize:
            z_q, *_ = self.vqgan.quantize(trans_latent)
        else:
            z_q = trans_latent
        return self.vqgan.decode(z_q)

    def _get_clip_similarity(self, prompts, image, weights=None):
        clip_inputs = self.clip_preprocessor(text=prompts, images=image, return_tensors="pt", padding=True)
        clip_outputs = self.clip(**clip_inputs)
        similarity_logits = clip_outputs.logits_per_image
        if weights is not None:
            similarity_logits = similarity_logits * weights
        return similarity_logits.sum()

    def _get_clip_loss(self, pos_prompts, neg_prompts, image):
        pos_logits = self._get_clip_similarity(pos_prompts["prompts"], image, weights=(1 / pos_prompts["weights"]))
        if neg_prompts:
            neg_logits = self._get_clip_similarity(neg_prompts["prompts"], image, weights=neg_prompts["weights"])
        else:
            neg_logits = torch.tensor([1], device=self.device)
        loss = -torch.log(pos_logits) + torch.log(neg_logits)
        return loss

    def _optimize_CLIP(self, original_img, pos_prompts, neg_prompts):
        vector = torch.randn_like(self.latent, requires_grad=True, device=self.device)
        optim = torch.optim.Adam([vector], lr=self.lr)

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

    def _init_logging(self, positive_prompts, negative_prompts, image_path):
        wandb.init(reinit=True, 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:
            image = Image.open(image_path)
            image = image.resize((256, 256))
            wandb.log("Original Image", wandb.Image(image))

    def process_prompts(self, prompts):
        if not prompts:
            return []
        processed_prompts = []
        weights = []
        if isinstance(prompts, str):
            prompts = [prompt.strip() for prompt in prompts.split("|")]
        for prompt in prompts:
            if isinstance(prompt, (tuple, list)):
                processed_prompt = prompt[0]
                weight = float(prompt[1])
            elif ":" in prompt:
                processed_prompt, weight = prompt.split(":")
                weight = float(weight)
            else:
                processed_prompt = prompt
                weight = 1.0
            processed_prompts.append(processed_prompt)
            weights.append(weight)
        return {
            "prompts": processed_prompts,
            "weights": torch.tensor(weights, device=self.device),
        }

    def generate(
        self,
        pos_prompts,
        neg_prompts=None,
        image_path=None,
        show_intermediate=True,
        save_intermediate=False,
        show_final=True,
        save_final=True,
        save_path=None,
    ):
        
        if image_path:
            self.latent = self._get_latent(image_path)
        else:
            self.latent = torch.randn(self.latent_dim, device=self.device)
        if self.log:
            self._init_logging(pos_prompts, neg_prompts, image_path)

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

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

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