rename backend

backend.py

@@ -0,0 +1,230 @@
+# from functools import cache
+import importlib
+
+import gradio as gr
+import matplotlib.pyplot as plt
+import torch
+import torchvision
+import wandb
+from icecream import ic
+from torch import nn
+from torchvision.transforms.functional import resize
+from tqdm import tqdm
+from transformers import CLIPModel, CLIPProcessor
+import lpips
+from edit import blend_paths
+from img_processing import *
+from img_processing import custom_to_pil
+from loaders import load_default
+import glob
+
+global log
+log=False
+
+# ic.disable()
+# ic.enable()
+def get_resized_tensor(x):
+    if len(x.shape) == 2:
+        re = x.unsqueeze(0)
+    else: re = x
+    re = resize(re, (10, 10))
+    return re
+class ProcessorGradientFlow():
+    """
+    This wraps the huggingface CLIP processor to allow backprop through the image processing step.
+    The original processor forces conversion to PIL images, which breaks gradient flow. 
+    """
+    def __init__(self, device="cuda") -> None:
+        self.device = device
+        self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
+        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.center_crop(images)
+        images = self.resize(images)
+        images = self.center_crop(images)
+        images = self.normalize(images)
+        return images
+    def __call__(self, images=[], **kwargs):
+        processed_inputs = self.processor(**kwargs)
+        processed_inputs["pixel_values"] = self.preprocess_img(images)
+        processed_inputs = {key:value.to(self.device) for (key, value) in processed_inputs.items()}
+        return processed_inputs
+
+class ImagePromptOptimizer(nn.Module):
+    def __init__(self, 
+                vqgan, 
+                clip,
+                clip_preprocessor,
+                lpips_fn,
+                iterations=100,
+                lr = 0.01,
+                save_vector=True,
+                return_val="vector",
+                quantize=True,
+                make_grid=False,
+                lpips_weight = 6.2) -> None:
+                
+        super().__init__()
+        self.latent = None
+        self.device = vqgan.device
+        vqgan.eval()
+        self.vqgan = vqgan
+        self.clip = clip
+        self.iterations = iterations
+        self.lr = lr
+        self.clip_preprocessor = clip_preprocessor
+        self.make_grid = make_grid
+        self.return_val = return_val
+        self.quantize = quantize
+        self.lpips_weight = lpips_weight
+        self.perceptual_loss = lpips_fn
+    def set_latent(self, latent):
+        self.latent = latent.detach().to(self.device)
+    def set_params(self, lr, iterations, lpips_weight, reconstruction_steps, attn_mask):
+        self.attn_mask = attn_mask
+        self.iterations = iterations
+        self.lr = lr
+        self.lpips_weight = lpips_weight
+        self.reconstruction_steps = reconstruction_steps
+    def forward(self, vector):
+        base_latent = self.latent.detach().requires_grad_()
+        trans_latent = base_latent + vector
+        if self.quantize:
+            z_q, *_ = self.vqgan.quantize(trans_latent)
+        else:
+            z_q = trans_latent
+        dec = self.vqgan.decode(z_q)
+        return dec
+    def _get_clip_similarity(self, prompts, image, weights=None):
+        if isinstance(prompts, str):
+            prompts = [prompts]
+        elif not isinstance(prompts, list):
+            raise TypeError("Provide prompts as string or list of strings")
+        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:
+            similarity_logits *= weights
+        return similarity_logits.sum()
+    def get_similarity_loss(self, pos_prompts, neg_prompts, image):
+        pos_logits = self._get_clip_similarity(pos_prompts, image)
+        if neg_prompts:
+            neg_logits = self._get_clip_similarity(neg_prompts, image)
+        else:
+            neg_logits = torch.tensor([1], device=self.device)
+        loss = -torch.log(pos_logits) + torch.log(neg_logits)
+        return loss
+    def visualize(self, processed_img):
+        if self.make_grid:
+            self.index += 1
+            plt.subplot(1, 13, self.index)
+            plt.imshow(get_pil(processed_img[0]).detach().cpu())
+        else:
+            plt.imshow(get_pil(processed_img[0]).detach().cpu())
+            plt.show()
+    def attn_masking(self, grad):
+        # print("attnmask 1")
+        # print(f"input grad.shape = {grad.shape}")
+        # print(f"input grad = {get_resized_tensor(grad)}")
+        newgrad = grad
+        if self.attn_mask is not None:
+            # print("masking mult")
+            newgrad = grad * (self.attn_mask)
+        # print("output grad, ", get_resized_tensor(newgrad))
+        # print("end atn 1")
+        return newgrad
+    def attn_masking2(self, grad):
+        # print("attnmask 2")
+        # print(f"input grad.shape = {grad.shape}")
+        # print(f"input grad = {get_resized_tensor(grad)}")
+        newgrad = grad
+        if self.attn_mask is not None:
+            # print("masking mult")
+            newgrad = grad * ((self.attn_mask - 1) * -1)
+        # print("output grad, ", get_resized_tensor(newgrad))
+        # print("end atn 2")
+        return newgrad
+
+    def optimize(self, latent, pos_prompts, neg_prompts):
+        self.set_latent(latent)
+        # self.make_grid=True
+        transformed_img = self(torch.zeros_like(self.latent, requires_grad=True, device=self.device))
+        original_img = loop_post_process(transformed_img)
+        vector = torch.randn_like(self.latent, requires_grad=True, device=self.device)
+        optim = torch.optim.Adam([vector], lr=self.lr)
+        if self.make_grid:
+            plt.figure(figsize=(35, 25))
+            self.index = 1
+        for i in tqdm(range(self.iterations)):
+            optim.zero_grad()
+            transformed_img = self(vector)
+            processed_img = loop_post_process(transformed_img) #* self.attn_mask
+            processed_img.retain_grad()
+            lpips_input = processed_img.clone()
+            lpips_input.register_hook(self.attn_masking2)
+            lpips_input.retain_grad()
+            clip_clone = processed_img.clone()
+            clip_clone.register_hook(self.attn_masking)
+            clip_clone.retain_grad()
+            with torch.autocast("cuda"):
+                clip_loss = self.get_similarity_loss(pos_prompts, neg_prompts, clip_clone)
+                print("CLIP loss", clip_loss)
+                perceptual_loss = self.perceptual_loss(lpips_input, original_img.clone()) * self.lpips_weight
+                print("LPIPS loss: ", perceptual_loss)
+            if log:
+                wandb.log({"Perceptual Loss": perceptual_loss})
+                wandb.log({"CLIP Loss": clip_loss})
+            clip_loss.backward(retain_graph=True)
+            perceptual_loss.backward(retain_graph=True)
+            p2 = processed_img.grad
+            print("Sum Loss", perceptual_loss + clip_loss)
+            optim.step()
+            # if i % self.iterations // 10 == 0: 
+                # self.visualize(transformed_img)
+            yield vector
+        if self.make_grid:
+            plt.savefig(f"plot {pos_prompts[0]}.png")
+            plt.show()
+        print("lpips solo op")
+        for i in range(self.reconstruction_steps):
+            optim.zero_grad()
+            transformed_img = self(vector)
+            processed_img = loop_post_process(transformed_img) #* self.attn_mask
+            processed_img.retain_grad()
+            lpips_input = processed_img.clone()
+            lpips_input.register_hook(self.attn_masking2)
+            lpips_input.retain_grad()
+            with torch.autocast("cuda"):
+                perceptual_loss = self.perceptual_loss(lpips_input, original_img.clone()) * self.lpips_weight
+            if log:
+                wandb.log({"Perceptual Loss": perceptual_loss})
+            print("LPIPS loss: ", perceptual_loss)
+            perceptual_loss.backward(retain_graph=True)
+            optim.step()
+            yield vector
+        # torch.save(vector, "nose_vector.pt")
+        # print("")
+        # print("DISC STEPS")
+        # print("*************")
+        # for i in range(self.reconstruction_steps):
+        #     optim.zero_grad()
+        #     transformed_img = self(vector)
+        #     processed_img = loop_post_process(transformed_img) #* self.attn_mask
+        #     disc_logits = self.disc(transformed_img)
+        #     disc_loss = self.disc_loss_fn(disc_logits)
+        #     print(f"disc_loss = {disc_loss}")
+        #     if log:
+        #         wandb.log({"Disc Loss": disc_loss})
+        #     print("LPIPS loss: ", perceptual_loss)
+        #     disc_loss.backward(retain_graph=True)
+        #     optim.step()
+        #     yield vector
+        yield vector if self.return_val == "vector" else self.latent + vector