cleanup and refactoring

ImageState.py

@@ -63,24 +63,15 @@ class ImageState:
     def _decode_latent_to_pil(self, latent):
         current_im = self.vqgan.decode(latent.to(self.device))[0]
         return custom_to_pil(current_im)
-    # def _get_current_vector_transforms(self):
-    #     current_vector_transforms = (self.blue_eyes, self.lip_size, self.hair_gp, self.asian_transform, sum(self.current_prompt_transforms))
-    #     return (self.blend_latent, current_vector_transforms)
-    # @cache
     def get_mask(self, img, mask=None):
         if img and "mask" in img and img["mask"] is not None:
             attn_mask = torchvision.transforms.ToTensor()(img["mask"])
             attn_mask = torch.ceil(attn_mask[0].to(self.device))
-            plt.imshow(attn_mask.detach().cpu(), cmap="Blues")
-            plt.show()
-            torch.save(attn_mask, "test_mask.pt")
             print("mask set successfully")
-            # attn_mask = self.rescale_mask(attn_mask)
             print(type(attn_mask))
             print(attn_mask.shape)
         else:
             attn_mask = mask
-            print("mask in apply ", get_resized_tensor(attn_mask), get_resized_tensor(attn_mask).shape)
         return attn_mask
     def set_mask(self, img):
         attn_mask = self.get_mask(img)

animation.py

@@ -8,7 +8,6 @@ def clear_img_dir():
         os.mkdir(img_dir)
     for filename in glob.glob(img_dir+"/*"):
         os.remove(filename)
-    
 
 def create_gif(total_duration, extend_frames, folder="./img_history", gif_name="face_edit.gif"):
     images = []
@@ -23,12 +22,8 @@ def create_gif(total_duration, extend_frames, folder="./img_history", gif_name="
         if file_name.endswith('.png'):
             file_path = os.path.join(folder, file_name)
             images.append(imageio.imread(file_path))
-    # images[0] = images[0].set_meta_data({'duration': 1})
-    # images[-1] = images[-1].set_meta_data({'duration': 1})
     imageio.mimsave(gif_name, images, duration=durations)
     return gif_name
 
 if __name__ == "__main__":
-    # clear_img_dir()
-    create_gif()
-# make_animation()
+    create_gif()

app.py

@@ -4,7 +4,7 @@ import sys
 
 import wandb
 
-from configs import set_major_global, set_major_local, set_small_local
+from presets import set_major_global, set_major_local, set_small_local
 
 sys.path.append("taming-transformers")
 import functools

app_backend.py

@@ -81,7 +81,7 @@ class ImagePromptOptimizer(nn.Module):
         self.make_grid = make_grid
         self.return_val = return_val
         self.quantize = quantize
-        self.disc = load_disc(self.device)
+        # self.disc = load_disc(self.device)
         self.lpips_weight = lpips_weight
         self.perceptual_loss = lpips.LPIPS(net='vgg').to(self.device)
     def disc_loss_fn(self, logits):
@@ -89,7 +89,7 @@ class ImagePromptOptimizer(nn.Module):
     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._attn_mask = attn_mask
         self.iterations = iterations
         self.lr = lr
         self.lpips_weight = lpips_weight
@@ -131,32 +131,29 @@ class ImagePromptOptimizer(nn.Module):
         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)}")
+    def _attn_mask(self, 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")
+        if self._attn_mask is not None:
+            newgrad = grad * (self._attn_mask)
         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)}")
+    def _attn_mask_inverse(self, 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")
+        if self._attn_mask is not None:
+            newgrad = grad * ((self._attn_mask - 1) * -1)
         return newgrad
+    def _get_next_inputs(self, transformed_img):
+        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_mask_inverse)
+        lpips_input.retain_grad()
+        clip_input = processed_img.clone()
+        clip_input.register_hook(self._attn_mask)
+        clip_input.retain_grad()
+        return processed_img, lpips_input, clip_input
 
     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)
@@ -167,27 +164,14 @@ class ImagePromptOptimizer(nn.Module):
         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()
+            processed_img, lpips_input, clip_input = self._get_next_inputs(transformed_img)
             with torch.autocast("cuda"):
-                clip_loss = self.get_similarity_loss(pos_prompts, neg_prompts, clip_clone)
+                clip_loss = self.get_similarity_loss(pos_prompts, neg_prompts, clip_input)
                 print("CLIP loss", clip_loss)
                 perceptual_loss = self.perceptual_loss(lpips_input, original_img.clone()) * self.lpips_weight
                 print("LPIPS loss: ", perceptual_loss)
-                with torch.no_grad():
-                    disc_logits = self.disc(transformed_img)
-                    disc_loss = self.disc_loss_fn(disc_logits)
-                    print(f"disc_loss = {disc_loss}")
-                    disc_loss2 = self.disc(processed_img)
             if log:
                 wandb.log({"Perceptual Loss": perceptual_loss})
-                wandb.log({"Discriminator Loss": disc_loss})
                 wandb.log({"CLIP Loss": clip_loss})
             clip_loss.backward(retain_graph=True)
             perceptual_loss.backward(retain_graph=True)
@@ -207,7 +191,7 @@ class ImagePromptOptimizer(nn.Module):
             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.register_hook(self._attn_mask_inverse)
             lpips_input.retain_grad()
             with torch.autocast("cuda"):
                 perceptual_loss = self.perceptual_loss(lpips_input, original_img.clone()) * self.lpips_weight
@@ -216,28 +200,10 @@ class ImagePromptOptimizer(nn.Module):
                     disc_loss = self.disc_loss_fn(disc_logits)
                     print(f"disc_loss = {disc_loss}")
                     disc_loss2 = self.disc(processed_img)
-            # print(f"disc_loss2 = {disc_loss2}")
             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

configs.py

@@ -1,7 +0,0 @@
-import gradio as gr
-def set_small_local():
-    return (gr.Slider.update(value=25), gr.Slider.update(value=0.15), gr.Slider.update(value=1), gr.Slider.update(value=4))
-def set_major_local():
-    return (gr.Slider.update(value=25), gr.Slider.update(value=0.25), gr.Slider.update(value=35), gr.Slider.update(value=10))
-def set_major_global():
-    return (gr.Slider.update(value=30), gr.Slider.update(value=0.1), gr.Slider.update(value=2), gr.Slider.update(value=0.2))

edit.py

@@ -17,13 +17,13 @@ from utils import get_device
 
 
 def get_embedding(model, path=None, img=None, device="cpu"):
-    assert path is None or img is None, "Input either path or tensor"
+    assert path or img, "Input either path or tensor"
     if img is not None:
         raise NotImplementedError
     x = preprocess(PIL.Image.open(path), target_image_size=256).to(device)
     x_processed = preprocess_vqgan(x)
-    x_latent, _, [_, _, indices] = model.encode(x_processed)
-    return x_latent
+    z, _, [_, _, indices] = model.encode(x_processed)
+    return z
 
     
 def blend_paths(model, path1, path2, quantize=False, weight=0.5, show=True, device="cuda"):
@@ -47,23 +47,12 @@ def blend_paths(model, path1, path2, quantize=False, weight=0.5, show=True, devi
 
 if __name__ == "__main__":
     device = get_device()
-    # conf_path = "logs/2021-04-23T18-11-19_celebahq_transformer/configs/2021-04-23T18-11-19-project.yaml"
     ckpt_path = "logs/2021-04-23T18-11-19_celebahq_transformer/checkpoints/last.ckpt"
-    # ckpt_path = "./faceshq/faceshq.pt"
     conf_path = "./unwrapped.yaml"
-    # conf_path = "./faceshq/faceshq.yaml"
     config = load_config(conf_path, display=False)
     model = taming.models.vqgan.VQModel(**config.model.params)
     sd = torch.load("./vqgan_only.pt", map_location="mps")
     model.load_state_dict(sd, strict=True)
     model.to(device)
     blend_paths(model, "./test_data/face.jpeg", "./test_data/face2.jpeg", quantize=False, weight=.5)
-    plt.show()
-
-    demo = gr.Interface(
-        get_image,
-        inputs=gr.inputs.Image(label="UploadZz a black and white face", type="filepath"),
-        outputs="image",
-        title="Upload a black and white face and get a colorized image!",
-    )
-
+    plt.show()

img_processing.py

@@ -32,7 +32,7 @@ def preprocess(img, target_image_size=256, map_dalle=False):
     return img
 
 def preprocess_vqgan(x):
-  x = 2.*x - 1.
+  x = 2. * x - 1.
   return x
 
 def custom_to_pil(x, process=True, mode="RGB"):

loaders.py

@@ -1,5 +1,4 @@
 import importlib
-
 import numpy as np
 import taming
 import torch
@@ -7,9 +6,8 @@ import yaml
 from omegaconf import OmegaConf
 from PIL import Image
 from taming.models.vqgan import VQModel
-
 from utils import get_device
-# import discriminator
+
 
 def load_config(config_path, display=False):
   config = OmegaConf.load(config_path)
@@ -17,37 +15,23 @@ def load_config(config_path, display=False):
     print(yaml.dump(OmegaConf.to_container(config)))
   return config
 
-# def load_disc(device):
-#     dconf = load_config("disc_config.yaml")
-#     sd = torch.load("disc.pt", map_location=device)
-#     # print(sd.keys())
-#     model = discriminator.NLayerDiscriminator()
-#     model.load_state_dict(sd, strict=True)
-#     model.to(device)
-#     return model
-    # print(dconf.keys())
-
 def load_default(device):
-    # device = get_device()
     ckpt_path = "logs/2021-04-23T18-11-19_celebahq_transformer/checkpoints/last.ckpt"
     conf_path = "./unwrapped.yaml"
     config = load_config(conf_path, display=False)
     model = taming.models.vqgan.VQModel(**config.model.params)
-    sd = torch.load("./vqgan_only.pt", map_location=device)
+    sd = torch.load("./model_checkpoints/vqgan_only.pt", map_location=device)
     model.load_state_dict(sd, strict=True)
     model.to(device)
     return model
 
 
 def load_vqgan(config, ckpt_path=None, is_gumbel=False):
-  if is_gumbel:
-    model = GumbelVQ(**config.model.params)
-  else:
     model = VQModel(**config.model.params)
-  if ckpt_path is not None:
-    sd = torch.load(ckpt_path, map_location="cpu")["state_dict"]
-    missing, unexpected = model.load_state_dict(sd, strict=False)
-  return model.eval()
+    if ckpt_path is not None:
+        sd = torch.load(ckpt_path, map_location="cpu")["state_dict"]
+        missing, unexpected = model.load_state_dict(sd, strict=False)
+    return model.eval()
 
 def load_ffhq():
     conf = "2020-11-09T13-33-36_faceshq_vqgan/configs/2020-11-09T13-33-36-project.yaml"

utils.py

@@ -7,10 +7,10 @@ import torch.nn.functional as F
 from skimage.color import lab2rgb, rgb2lab
 from torch import nn
 
-
 def freeze_module(module):
     for param in module.parameters():
       param.requires_grad = False
+
 def get_device():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     if torch.backends.mps.is_available() and torch.backends.mps.is_built():

vqgan_latent_ops.py

@@ -1,14 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from gradient_flow_ops import ReplaceGrad
-
-replace_grad = ReplaceGrad.apply
-
-def vector_quantize(x, codebook):
-
-    d = x.pow(2).sum(dim=-1, keepdim=True) + codebook.pow(2).sum(dim=1) - 2 * x @ codebook.T
-    indices = d.argmin(-1)
-    x_q = F.one_hot(indices, codebook.shape[0]).to(d.dtype) @ codebook
-    return replace_grad(x_q, x)

vqgan_only.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8e39472bae4489764c0ffc70ba84ec7815f245781020ce55cc2e7adc60e580e4
-size 288690579