Spaces:
Configuration error
Configuration error
| # 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 | |