initial commit

.gitattributes

@@ -0,0 +1 @@
+vqgan_only.pt filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,9 @@
+__pycache__
+__pycache__/*
+img_history
+aligned
+wandb
+women
+men
+gradio_cache*
+.DS_Store

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "taming"]
+	path = taming
+	url = https://github.com/CompVis/taming-transformers.git

ImageState.py

@@ -0,0 +1,192 @@
+# from align import align_from_path
+from animation import clear_img_dir
+from app_backend import ImagePromptOptimizer, log
+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 app_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
+
+num = 0
+class PromptTransformHistory():
+    def __init__(self, iterations) -> None:
+        self.iterations = iterations
+        self.transforms = []
+
+class ImageState:
+    def __init__(self, vqgan, prompt_optimizer: ImagePromptOptimizer) -> None:
+        self.vqgan = vqgan
+        self.device = vqgan.device
+        self.blend_latent = None
+        self.quant = True
+        self.path1 = None
+        self.path2 = None
+        self.transform_history = []
+        self.attn_mask = None
+        self.prompt_optim = prompt_optimizer
+        self._load_vectors()
+        self.init_transforms()
+    def _load_vectors(self):
+        self.lip_vector = torch.load("./latent_vectors/lipvector.pt", map_location=self.device)
+        self.red_blue_vector = torch.load("./latent_vectors/2blue_eyes.pt", map_location=self.device)
+        self.green_purple_vector = torch.load("./latent_vectors/nose_vector.pt", map_location=self.device)
+        self.asian_vector = torch.load("./latent_vectors/asian10.pt", map_location=self.device)
+    def init_transforms(self):
+        self.blue_eyes = torch.zeros_like(self.lip_vector)
+        self.lip_size = torch.zeros_like(self.lip_vector)
+        self.asian_transform = torch.zeros_like(self.lip_vector)
+        self.current_prompt_transforms = [torch.zeros_like(self.lip_vector)]
+        self.hair_gp = torch.zeros_like(self.lip_vector)
+    def clear_transforms(self):
+        global num
+        self.init_transforms()
+        clear_img_dir()
+        num = 0
+        return self._render_all_transformations()
+    def _apply_vector(self, src, vector):
+        new_latent = torch.lerp(src, src + vector, 1)
+        return new_latent
+    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)
+        self.attn_mask = attn_mask
+            # attn_mask = torch.ones_like(img, device=self.device)
+        x = attn_mask.clone()
+        x = x.detach().cpu()
+        x = torch.clamp(x, -1., 1.)
+        x = (x + 1.)/2.
+        x = x.numpy()
+        x = (255*x).astype(np.uint8)
+        x = Image.fromarray(x, "L")
+        return x
+    @torch.no_grad()
+    def _render_all_transformations(self, return_twice=True):
+        global num
+        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:
+            new_latent, _, _ = self.vqgan.quantize(new_latent.to(self.device))
+        image = self._decode_latent_to_pil(new_latent)
+        img_dir = "./img_history"
+        if not os.path.exists(img_dir):
+            os.mkdir(img_dir)
+        image.save(f"./img_history/img_{num:06}.png")
+        num += 1
+        return (image, image) if return_twice else image
+    def apply_gp_vector(self, weight):
+        self.hair_gp = weight * self.green_purple_vector
+        return self._render_all_transformations()
+    def apply_rb_vector(self, weight):
+        self.blue_eyes = weight * self.red_blue_vector
+        return self._render_all_transformations()
+    def apply_lip_vector(self, weight):
+        self.lip_size = weight * self.lip_vector
+        return self._render_all_transformations()
+    def update_requant(self, val):
+        print(f"val = {val}")
+        self.quant = val
+        return self._render_all_transformations()
+    def apply_gender_vector(self, weight):
+        self.asian_transform = weight * self.asian_vector
+        return self._render_all_transformations()
+    def update_images(self, path1, path2, blend_weight):
+        if path1 is None and path2 is None:
+            return None
+        if path1 is None: path1 = path2
+        if path2 is None: path2 = path1
+        self.path1, self.path2 = path1, path2
+        # self.aligned_path1 = align_from_path(path1)
+        # self.aligned_path2 = align_from_path(path2)
+        return self.blend(blend_weight)
+    @torch.no_grad()
+    def blend(self, weight):
+        _, latent = blend_paths(self.vqgan, self.path1, self.path2, weight=weight, show=False, device=self.device)
+        self.blend_latent = latent
+        return self._render_all_transformations()
+    @torch.no_grad()
+    def rewind(self, index):
+        if not self.transform_history:
+            print("no history")
+            return self._render_all_transformations()
+        prompt_transform = self.transform_history[-1]
+        latent_index = int(index / 100 * (prompt_transform.iterations - 1))
+        print(latent_index)
+        self.current_prompt_transforms[-1] = prompt_transform.transforms[latent_index]
+        # print(self.current_prompt_transform)
+        # print(self.current_prompt_transforms.mean())
+        return self._render_all_transformations()
+    def rescale_mask(self, mask):
+        rep = mask.clone()
+        rep[mask < 0.03] = -1000000
+        rep[mask >= 0.03] = 1
+        return rep
+    def apply_prompts(self, positive_prompts, negative_prompts, lr, iterations, lpips_weight, reconstruction_steps):
+        transform_log = PromptTransformHistory(iterations + reconstruction_steps)
+        transform_log.transforms.append(torch.zeros_like(self.blend_latent, requires_grad=False))
+        self.current_prompt_transforms.append(torch.zeros_like(self.blend_latent, requires_grad=False))
+        if log:
+            wandb.init(reinit=True, project="face-editor")
+            wandb.config.update({"Positive Prompts": positive_prompts})
+            wandb.config.update({"Negative Prompts": negative_prompts})
+            wandb.config.update(dict(
+                lr=lr,
+                iterations=iterations,
+                lpips_weight=lpips_weight
+            ))
+        positive_prompts = [prompt.strip() for prompt in positive_prompts.split("|")]
+        negative_prompts = [prompt.strip() for prompt in negative_prompts.split("|")]
+        self.prompt_optim.set_params(lr, iterations, lpips_weight, attn_mask=self.attn_mask, reconstruction_steps=reconstruction_steps)
+        for i, transform in enumerate(self.prompt_optim.optimize(self.blend_latent,
+                                                                positive_prompts,
+                                                                negative_prompts)):
+            transform_log.transforms.append(transform.clone().detach())
+            self.current_prompt_transforms[-1] = transform
+            with torch.no_grad():
+                image = self._render_all_transformations(return_twice=False)
+            if log:
+                wandb.log({"image": wandb.Image(image)})
+            yield (image, image)
+        if log:
+            wandb.finish()
+        self.attn_mask = None
+        self.transform_history.append(transform_log)
+        # transform = self.prompt_optim.optimize(self.blend_latent,
+                                                # positive_prompts,
+                                                # negative_prompts)
+        # self.prompt_transforms = transform
+        # return self._render_all_transformations()

animation.py

@@ -0,0 +1,34 @@
+import imageio
+import glob
+import os
+
+def clear_img_dir():
+    img_dir = "./img_history"
+    if not os.path.exists(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 = []
+    paths = glob.glob(folder + "/*")
+    frame_duration = total_duration / len(paths)
+    print(len(paths), "frame dur", frame_duration)
+    durations = [frame_duration] * len(paths)
+    if extend_frames:
+        durations [0] = 1.5
+        durations [-1] = 3
+    for file_name in os.listdir(folder):
+        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()

app.py

@@ -0,0 +1,195 @@
+import glob
+import os
+import sys
+
+import wandb
+
+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
+
+import edit
+# import importlib
+# importlib.reload(edit)
+from app_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 = "cuda"
+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)
+state = ImageState(vqgan, promptoptim)
+def set_img_from_example(img):
+    return state.update_images(img, img, 0)
+def get_cleared_mask():
+    return gr.Image.update(value=None)
+    # mask.clear()
+with gr.Blocks(css="styles.css") as demo:
+    with gr.Row():
+        with gr.Column(scale=1):
+            blue_eyes = gr.Slider(
+                label="Blue Eyes",
+                minimum=-.8,
+                maximum=3.,
+                value=0,
+                step=0.1,
+            )
+            # hair_green_purple = gr.Slider(
+            #     label="hair green<->purple ",
+            #     minimum=-.8,
+            #     maximum=.8,
+            #     value=0,
+            #     step=0.1,
+            # )
+            lip_size = gr.Slider(
+                label="Lip Size",
+                minimum=-1.9,
+                value=0,
+                maximum=1.9,
+                step=0.1,
+            )
+            blend_weight = gr.Slider(
+                label="Blend faces: 0 is base image, 1 is the second img",
+                minimum=-0.,
+                value=0,
+                maximum=1.,
+                step=0.1,
+            )
+            # requantize = gr.Checkbox(
+            #     label="Requantize Latents (necessary using text prompts)",
+            #     value=True,
+            # )
+            asian_weight = gr.Slider(
+                minimum=-2.,
+                value=0,
+                label="Asian",
+                maximum=2.,
+                step=0.07,
+            )
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown(value="""## Image Upload
+                                    For best results, crop the photos like in the example pictures""", show_label=False)
+                    with gr.Row():
+                        base_img = gr.Image(label="Base Image", type="filepath")
+                        blend_img = gr.Image(label="Image for face blending (optional)", type="filepath")
+                    # gr.Markdown("## Image Examples")
+                with gr.Accordion(label="Add Mask", open=False):
+                    mask = gr.Image(tool="sketch", interactive=True)
+                    gr.Markdown(value="Note: You must clear the mask using the rewind button every time you want to change the mask (this is a gradio bug)")
+                    set_mask = gr.Button(value="Set mask")
+                    gr.Text(value="this image shows the mask passed to the model when you press set mask (debugging purposes)")
+                    testim = gr.Image()
+                    clear_mask = gr.Button(value="Clear mask")
+                    clear_mask.click(get_cleared_mask, outputs=mask)
+            with gr.Row():
+                gr.Examples(
+                    examples=glob.glob("test_pics/*"),
+                    inputs=base_img,
+                    outputs=blend_img,
+                    fn=set_img_from_example,
+                    # cache_examples=True,
+                    ) 
+        with gr.Column(scale=1):
+            out = gr.Image()
+            rewind = gr.Slider(value=100,
+                                label="Rewind back through a prompt transform: Use this to scroll through the iterations of your prompt transformation.",
+                                minimum=0,
+                                maximum=100)
+
+            apply_prompts = gr.Button(value="Apply Prompts", elem_id="apply")
+            clear = gr.Button(value="Clear all transformations (irreversible)", elem_id="warning")
+            with gr.Accordion(label="Save Animation", open=False):
+                gr.Text(value="Creates an animation of all the steps in the editing process", show_label=False)
+                duration = gr.Number(value=10, label="Duration of the animation in seconds")
+                extend_frames = gr.Checkbox(value=True, label="Make first and last frame longer")
+                gif = gr.File(interactive=False)
+                create_animation = gr.Button(value="Create Animation")
+                create_animation.click(create_gif, inputs=[duration, extend_frames], outputs=gif)
+
+        with gr.Column(scale=1):
+            gr.Markdown(value="""## Text Prompting
+            See readme for a prompting guide. Use the '|' symbol to separate prompts. Use the "Add mask" section to make local edits. Negative prompts are highly recommended""", show_label=False)
+            positive_prompts = gr.Textbox(label="Positive prompts",
+                                            value="a picture of a woman with a very big nose | a picture of a woman with a large wide nose | a woman with an extremely prominent nose")
+            negative_prompts = gr.Textbox(label="Negative prompts",
+                                            value="a picture of a person with a tiny nose | a picture of a person with a very thin nose")
+            gen_prompts = gr.Button(value="🎲 Random prompts")
+            gen_prompts.click(get_random_prompts, outputs=[positive_prompts, negative_prompts])
+            with gr.Row():
+                with gr.Column():
+                    gr.Text(value="Prompt Editing Configuration", show_label=False)
+                    with gr.Row():
+                        gr.Markdown(value="## Preset Configs", show_label=False)
+                    with gr.Row():
+                        with gr.Column():
+                            small_local = gr.Button(value="Small Masked Changes (e.g. add lipstick)", elem_id="small_local").style(full_width=False)
+                        with gr.Column():
+                            major_local = gr.Button(value="Major Masked Changes (e.g. change hair color or nose size)").style(full_width=False)
+                        with gr.Column():
+                            major_global = gr.Button(value="Major Global Changes (e.g. change race / gender").style(full_width=False)
+                    iterations = gr.Slider(minimum=10,
+                                            maximum=300,
+                                            step=1,
+                                            value=20,
+                                            label="Iterations: How many steps the model will take to modify the image. Try starting small and seeing how the results turn out, you can always resume with afterwards",)
+                    learning_rate = gr.Slider(minimum=1e-3,
+                                            maximum=6e-1,
+                                            value=1e-2,
+                                            label="Learning Rate: How strong the change in each step will be (you should raise this for bigger changes (for example, changing hair color), and lower it for more minor changes. Raise if changes aren't strong enough")
+                with gr.Accordion(label="Advanced Prompt Editing Options", open=False):
+                    lpips_weight = gr.Slider(minimum=0,
+                                            maximum=50,
+                                            value=1,
+                                            label="Perceptual similarity weight (Keeps areas outside of the mask looking similar to the original. Increase if the rest of the image is changing too much while you're trying to change make a localized edit")
+                    reconstruction_steps = gr.Slider(minimum=0,
+                                            maximum=50,
+                                            value=15,
+                                            step=1,
+                                            label="Steps to run at the end of the optimization, optimizing only the masked perceptual loss. If the edit is changing the identity too much, this setting will run steps at the end that will 'pull' the image back towards the original identity")
+                    # discriminator_steps = gr.Slider(minimum=0,
+                    #                         maximum=50,
+                    #                         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, outputs=[out, mask])
+    asian_weight.change(state.apply_gender_vector, inputs=[asian_weight], outputs=[out, mask])
+    lip_size.change(state.apply_lip_vector, inputs=[lip_size], outputs=[out, mask])
+    # hair_green_purple.change(state.apply_gp_vector, inputs=[hair_green_purple], outputs=[out, mask])
+    blue_eyes.change(state.apply_rb_vector, inputs=[blue_eyes], outputs=[out, mask])
+
+    blend_weight.change(state.blend, inputs=[blend_weight], outputs=[out, mask])
+    # requantize.change(state.update_requant, inputs=[requantize], outputs=[out, mask])
+
+
+    base_img.change(state.update_images, inputs=[base_img, blend_img, blend_weight], outputs=[out, mask])
+    blend_img.change(state.update_images, inputs=[base_img, blend_img, blend_weight], outputs=[out, mask])
+
+    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])
+    small_local.click(set_major_global, outputs=[iterations, learning_rate, lpips_weight, reconstruction_steps])
+    apply_prompts.click(state.apply_prompts, inputs=[positive_prompts, negative_prompts, learning_rate, iterations, lpips_weight, reconstruction_steps], outputs=[out, mask])
+    rewind.change(state.rewind, inputs=[rewind], outputs=[out, mask])
+    set_mask.click(state.set_mask, inputs=mask, outputs=testim)
+demo.queue()
+demo.launch(debug=True, inbrowser=True)
+# if __name__ == "__main__":
+    # import argparse
+    # parser = argparse.ArgumentParser()
+    # parser.add_argument('--debug', action='store_true', default=False, help='Enable debugging output')
+    # args = parser.parse_args()
+    # # if args.debug:
+    # #     state=None
+    # #     promptoptim=None
+    # # else: 
+    # main()

app_backend.py

@@ -0,0 +1,243 @@
+# 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.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):
+        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)
+                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)
+            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
+                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)
+            # 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

@@ -0,0 +1,7 @@
+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

@@ -0,0 +1,69 @@
+import os
+import sys
+
+from img_processing import custom_to_pil, preprocess, preprocess_vqgan
+
+sys.path.append("taming-transformers")
+import glob
+
+import gradio as gr
+import matplotlib.pyplot as plt
+import PIL
+import taming
+import torch
+
+from loaders import load_config
+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"
+    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
+
+    
+def blend_paths(model, path1, path2, quantize=False, weight=0.5, show=True, device="cuda"):
+    x = preprocess(PIL.Image.open(path1), target_image_size=256).to(device)
+    y = preprocess(PIL.Image.open(path2), target_image_size=256).to(device)
+    x_latent, y_latent = get_embedding(model, path=path1, device=device), get_embedding(model, path=path2, device=device)
+    z = torch.lerp(x_latent, y_latent, weight)
+    if quantize:
+        z = model.quantize(z)[0]
+    decoded = model.decode(z)[0]
+    if show:
+        plt.figure(figsize=(10, 20))
+        plt.subplot(1, 3, 1)
+        plt.imshow(x.cpu().permute(0, 2, 3, 1)[0])
+        plt.subplot(1, 3, 2)
+        plt.imshow(custom_to_pil(decoded))
+        plt.subplot(1, 3, 3)
+        plt.imshow(y.cpu().permute(0, 2, 3, 1)[0])
+        plt.show()
+    return custom_to_pil(decoded), z
+
+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!",
+    )
+

img_processing.py

@@ -0,0 +1,72 @@
+import io
+import os
+import sys
+
+import numpy as np
+import PIL
+import requests
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as T
+import torchvision.transforms.functional as TF
+from PIL import Image, ImageDraw, ImageFont
+
+
+def download_image(url):
+    resp = requests.get(url)
+    resp.raise_for_status()
+    return PIL.Image.open(io.BytesIO(resp.content))
+
+
+def preprocess(img, target_image_size=256, map_dalle=False):
+    s = min(img.size)
+    
+    if s < target_image_size:
+        raise ValueError(f'min dim for image {s} < {target_image_size}')
+        
+    r = target_image_size / s
+    s = (round(r * img.size[1]), round(r * img.size[0]))
+    img = TF.resize(img, s, interpolation=PIL.Image.LANCZOS)
+    img = TF.center_crop(img, output_size=2 * [target_image_size])
+    img = torch.unsqueeze(T.ToTensor()(img), 0)
+    return img
+
+def preprocess_vqgan(x):
+  x = 2.*x - 1.
+  return x
+
+def custom_to_pil(x, process=True, mode="RGB"):
+  x = x.detach().cpu()
+  if process:
+    x = torch.clamp(x, -1., 1.)
+    x = (x + 1.)/2.
+  x = x.permute(1,2,0).numpy()
+  if process:
+    x = (255*x).astype(np.uint8)
+  x = Image.fromarray(x)
+  if not x.mode == mode:
+    x = x.convert(mode)
+  return x
+
+def get_pil(x):
+  x = torch.clamp(x, -1., 1.)
+  x = (x + 1.)/2.
+  x = x.permute(1,2,0)
+  return x
+
+def loop_post_process(x):
+  x = get_pil(x.squeeze())
+  return x.permute(2, 0, 1).unsqueeze(0)
+
+def stack_reconstructions(input, x0, x1, x2, x3, titles=[]):
+  assert input.size == x1.size == x2.size == x3.size
+  w, h = input.size[0], input.size[1]
+  img = Image.new("RGB", (5*w, h))
+  img.paste(input, (0,0))
+  img.paste(x0, (1*w,0))
+  img.paste(x1, (2*w,0))
+  img.paste(x2, (3*w,0))
+  img.paste(x3, (4*w,0))
+  for i, title in enumerate(titles):
+    ImageDraw.Draw(img).text((i*w, 0), f'{title}', (255, 255, 255), font=font) # coordinates, text, color, font
+  return img

loaders.py

@@ -0,0 +1,97 @@
+import importlib
+
+import numpy as np
+import taming
+import torch
+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)
+  if display:
+    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)
+    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()
+
+def load_ffhq():
+    conf = "2020-11-09T13-33-36_faceshq_vqgan/configs/2020-11-09T13-33-36-project.yaml"
+    ckpt = "2020-11-09T13-33-36_faceshq_vqgan/checkpoints/last.ckpt"
+    vqgan = load_model(load_config(conf), ckpt, True, True)[0]
+
+def reconstruct_with_vqgan(x, model):
+  # could also use model(x) for reconstruction but use explicit encoding and decoding here
+  z, _, [_, _, indices] = model.encode(x)
+  print(f"VQGAN --- {model.__class__.__name__}: latent shape: {z.shape[2:]}")
+  xrec = model.decode(z)
+  return xrec
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+def instantiate_from_config(config):
+
+    if not "target" in config:
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+def load_model_from_config(config, sd, gpu=True, eval_mode=True):
+    model = instantiate_from_config(config)
+    if sd is not None:
+        model.load_state_dict(sd)
+    if gpu:
+        model.cuda()
+    if eval_mode:
+        model.eval()
+    return {"model": model}
+
+
+def load_model(config, ckpt, gpu, eval_mode):
+    # load the specified checkpoint
+    if ckpt:
+        pl_sd = torch.load(ckpt, map_location="cpu")
+        global_step = pl_sd["global_step"]
+        print(f"loaded model from global step {global_step}.")
+    else:
+        pl_sd = {"state_dict": None}
+        global_step = None
+    model = load_model_from_config(config.model, pl_sd["state_dict"], gpu=gpu, eval_mode=eval_mode)["model"]
+    return model, global_step

masking.py

@@ -0,0 +1,32 @@
+import os
+import sys
+
+import matplotlib.pyplot as plt
+import torch
+
+sys.path.append("taming-transformers")
+import functools
+
+import gradio as gr
+from transformers import CLIPModel, CLIPProcessor
+
+import edit
+# import importlib
+# importlib.reload(edit)
+from app_backend import ImagePromptOptimizer, ImageState, ProcessorGradientFlow
+from loaders import load_default
+
+device = "cuda"
+vqgan = load_default(device)
+vqgan.eval()
+processor = ProcessorGradientFlow(device=device)
+clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+clip.to(device)
+promptoptim = ImagePromptOptimizer(vqgan, clip, processor, quantize=True)
+state = ImageState(vqgan, promptoptim)
+mask = torch.load("eyebrow_mask.pt")
+x = state.blend("./test_data/face.jpeg", "./test_data/face2.jpeg", 0.5)
+plt.imshow(x)
+plt.show()
+state.apply_prompts("a picture of a woman with big eyebrows", "", 0.009, 40, None, mask=mask)
+print('done')

prompts.py

@@ -0,0 +1,17 @@
+import random
+class PromptSet:
+    def __init__(self, pos, neg, config=None):
+        self.positive = pos
+        self.negative = neg
+        self.config = config
+example_prompts = (
+    PromptSet("a picture of a woman with light blonde hair", "a picture of a person with dark hair | a picture of a person with brown hair"),
+    PromptSet("A picture of a woman with very thick eyebrows", "a picture of a person with very thin eyebrows | a picture of a person with no eyebrows"),
+    PromptSet("A picture of a woman wearing bright red lipstick", "a picture of a person wearing no lipstick | a picture of a person wearing dark lipstick"),
+    PromptSet("A picture of a beautiful chinese woman | a picture of a Japanese woman | a picture of an Asian woman", "a picture of a white woman | a picture of an Indian woman | a picture of a black woman"),
+    PromptSet("A picture of a handsome man | a picture of a masculine man", "a picture of a woman | a picture of a feminine person"),
+    PromptSet("A picture of a woman with a very big nose", "a picture of a person with a small nose | a picture of a person with a normal nose"),
+)
+def get_random_prompts():
+    prompt = random.choice(example_prompts)
+    return prompt.positive, prompt.negative

requirements.txt

@@ -0,0 +1,27 @@
+taming-transformers
+einops
+gradio
+icecream
+imageio
+matplotlib
+more_itertools
+numpy
+omegaconf
+opencv_python_headless
+Pillow
+prompts
+pudb
+pytorch_lightning
+PyYAML
+requests
+scikit_image
+scipy
+setuptools
+streamlit
+torch
+torchvision
+tqdm
+transformers
+typing_extensions
+wandb
+lpips

unwrapped.yaml

@@ -0,0 +1,37 @@
+model:
+  target: taming.models.vqgan.VQModel
+  params:
+    embed_dim: 256
+    n_embed: 1024
+    ddconfig:
+      double_z: false
+      z_channels: 256
+      resolution: 256
+      in_channels: 3
+      out_ch: 3
+      ch: 128
+      ch_mult:
+      - 1
+      - 1
+      - 2
+      - 2
+      - 4
+      num_res_blocks: 2
+      attn_resolutions:
+      - 16
+      dropout: 0.0
+    lossconfig:
+      target: taming.modules.losses.vqperceptual.DummyLoss
+data:
+  target: cutlit.DataModuleFromConfig
+  params:
+    batch_size: 24
+    num_workers: 24
+    train:
+      target: taming.data.faceshq.CelebAHQTrain
+      params:
+        size: 256
+    validation:
+      target: taming.data.faceshq.CelebAHQValidation
+      params:
+        size: 256

utils.py

@@ -0,0 +1,18 @@
+import time
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+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():
+        device = "mps"
+    return (device)

vqgan_latent_ops.py

@@ -0,0 +1,14 @@
+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)