reduce memory usage

ImageState.py

@@ -1,6 +1,6 @@
 # from align import align_from_path
 from animation import clear_img_dir
-from app_backend import ImagePromptOptimizer, log
+from backend import ImagePromptOptimizer, log
 import importlib
 import gradio as gr
 import matplotlib.pyplot as plt
@@ -13,12 +13,13 @@ from torchvision.transforms.functional import resize
 from tqdm import tqdm
 from transformers import CLIPModel, CLIPProcessor
 import lpips
-from app_backend import get_resized_tensor
+from backend import get_resized_tensor
 from edit import blend_paths
 from img_processing import *
 from img_processing import custom_to_pil
 from loaders import load_default
-
+# from app import vqgan
+global vqgan
 num = 0
 class PromptTransformHistory():
     def __init__(self, iterations) -> None:
@@ -27,6 +28,7 @@ class PromptTransformHistory():
 
 class ImageState:
     def __init__(self, vqgan, prompt_optimizer: ImagePromptOptimizer) -> None:
+        # global vqgan
         self.vqgan = vqgan
         self.device = vqgan.device
         self.blend_latent = None
@@ -59,6 +61,7 @@ class ImageState:
         new_latent = torch.lerp(src, src + vector, 1)
         return new_latent
     def _decode_latent_to_pil(self, latent):
+        # global vqgan
         current_im = self.vqgan.decode(latent.to(self.device))[0]
         return custom_to_pil(current_im)
     # def _get_current_vector_transforms(self):
@@ -95,6 +98,7 @@ class ImageState:
     @torch.no_grad()
     def _render_all_transformations(self, return_twice=True):
         global num
+        # global vqgan
         current_vector_transforms = (self.blue_eyes, self.lip_size, self.hair_gp, self.asian_transform, sum(self.current_prompt_transforms))
         new_latent = self.blend_latent + sum(current_vector_transforms)
         if self.quant:

app.py

@@ -3,29 +3,33 @@ import os
 import sys
 
 import wandb
+import torch
 
 from configs import set_major_global, set_major_local, set_small_local
 
 sys.path.append("taming-transformers")
-import functools
 
 import gradio as gr
 from transformers import CLIPModel, CLIPProcessor
+from lpips import LPIPS
 
 import edit
-from app_backend import ImagePromptOptimizer, ProcessorGradientFlow
+from backend import ImagePromptOptimizer, ProcessorGradientFlow
 from ImageState import ImageState
 from loaders import load_default
 from animation import create_gif
 from prompts import get_random_prompts
 
-device = "cpu"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+global vqgan
 vqgan = load_default(device)
 vqgan.eval()
 processor = ProcessorGradientFlow(device=device)
-clip = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
-clip.to(device)
-promptoptim = ImagePromptOptimizer(vqgan, clip, processor, quantize=True)
+# clip = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
+lpips_fn = LPIPS(net='vgg').to(device)
+clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
+promptoptim = ImagePromptOptimizer(vqgan, clip, processor, lpips_fn=lpips_fn, quantize=True)
 def set_img_from_example(state, img):
     return state.update_images(img, img, 0)
 def get_cleared_mask():
@@ -40,7 +44,8 @@ class StateWrapper:
     def apply_lip_vector(state, *args, **kwargs):
         return state, *state[0].apply_lip_vector(*args, **kwargs)
     def apply_prompts(state, *args, **kwargs):
-        return state, *state[0].apply_prompts(*args, **kwargs)
+        for image in state[0].apply_prompts(*args, **kwargs):
+            yield state, *image
     def apply_rb_vector(state, *args, **kwargs):
         return state, *state[0].apply_rb_vector(*args, **kwargs)
     def blend(state, *args, **kwargs):
@@ -56,7 +61,7 @@ class StateWrapper:
     def rewind(state, *args, **kwargs):
         return state, *state[0].rewind(*args, **kwargs)
     def set_mask(state, *args, **kwargs):
-        return state, *state[0].set_mask(*args, **kwargs)
+        return state, state[0].set_mask(*args, **kwargs)
     def update_images(state, *args, **kwargs):
         return state, *state[0].update_images(*args, **kwargs)
     def update_requant(state, *args, **kwargs):
@@ -191,7 +196,7 @@ with gr.Blocks(css="styles.css") as demo:
                     #                         step=1,
                     #                         value=0,
                     #                         label="Steps to run at the end, optimizing only the discriminator loss. This helps to reduce artefacts, but because the model is trained on CelebA, this will make your generations look more like generic white celebrities")
-    # clear.click(state.clear_transforms, inputs=[state], outputs=[state, out, mask])
+    clear.click(state.clear_transforms, inputs=[state], outputs=[state, out, mask])
     asian_weight.change(StateWrapper.apply_asian_vector, inputs=[state, asian_weight], outputs=[state, out, mask])
     lip_size.change(StateWrapper.apply_lip_vector, inputs=[state, lip_size], outputs=[state, out, mask])
     # hair_green_purple.change(StateWrapper.apply_gp_vector, inputs=[state, hair_green_purple], outputs=[state, out, mask])
@@ -200,11 +205,11 @@ with gr.Blocks(css="styles.css") as demo:
     # requantize.change(StateWrapper.update_requant, inputs=[state, requantize], outputs=[state, out, mask])
     base_img.change(StateWrapper.update_images, inputs=[state, base_img, blend_img, blend_weight], outputs=[state, out, mask])
     blend_img.change(StateWrapper.update_images, inputs=[state, base_img, blend_img, blend_weight], outputs=[state, out, mask])
-    small_local.click(set_small_local, outputs=[state, iterations, learning_rate, lpips_weight, reconstruction_steps])
-    major_local.click(set_major_local, outputs=[state, iterations, learning_rate, lpips_weight, reconstruction_steps])
-    major_global.click(set_major_global, outputs=[state, iterations, learning_rate, lpips_weight, reconstruction_steps])
+    small_local.click(set_small_local, outputs=[iterations, learning_rate, lpips_weight, reconstruction_steps])
+    major_local.click(set_major_local, outputs=[iterations, learning_rate, lpips_weight, reconstruction_steps])
+    major_global.click(set_major_global, outputs=[iterations, learning_rate, lpips_weight, reconstruction_steps])
     apply_prompts.click(StateWrapper.apply_prompts, inputs=[state, positive_prompts, negative_prompts, learning_rate, iterations, lpips_weight, reconstruction_steps], outputs=[state, out, mask])
     rewind.change(StateWrapper.rewind, inputs=[state, rewind], outputs=[state, out, mask])
-    set_mask.click(StateWrapper.set_mask, inputs=mask, outputs=testim)
+    set_mask.click(StateWrapper.set_mask, inputs=[state, mask], outputs=[state, testim])
 demo.queue()
 demo.launch(debug=True, enable_queue=True)

app_backend.py

@@ -1,230 +0,0 @@
-# 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,
-                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.LPIPS(net='vgg').to(self.device)
-    def disc_loss_fn(self, logits):
-        return -torch.mean(logits)
-    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

masking.py

@@ -13,7 +13,7 @@ from transformers import CLIPModel, CLIPProcessor
 import edit
 # import importlib
 # importlib.reload(edit)
-from app_backend import ImagePromptOptimizer, ImageState, ProcessorGradientFlow
+from backend import ImagePromptOptimizer, ImageState, ProcessorGradientFlow
 from loaders import load_default
 
 device = "cuda"