# Authors: Hui Ren (rhfeiyang.github.io)
import os
import numpy as np
from torchvision import transforms
import torch
import torch.nn.functional as F
import torch.nn as nn
from torch.autograd import Function
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
from collections import OrderedDict
from transformers import BatchFeature
import clip
import copy
import lpips
from transformers import ViTImageProcessor, ViTModel
## CSD_CLIP
def convert_weights_float(model: nn.Module):
"""Convert applicable model parameters to fp32"""
def _convert_weights_to_fp32(l):
if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
l.weight.data = l.weight.data.float()
if l.bias is not None:
l.bias.data = l.bias.data.float()
if isinstance(l, nn.MultiheadAttention):
for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
tensor = getattr(l, attr)
if tensor is not None:
tensor.data = tensor.data.float()
for name in ["text_projection", "proj"]:
if hasattr(l, name):
attr = getattr(l, name)
if attr is not None:
attr.data = attr.data.float()
model.apply(_convert_weights_to_fp32)
class ReverseLayerF(Function):
@staticmethod
def forward(ctx, x, alpha):
ctx.alpha = alpha
return x.view_as(x)
@staticmethod
def backward(ctx, grad_output):
output = grad_output.neg() * ctx.alpha
return output, None
## taken from https://github.com/moein-shariatnia/OpenAI-CLIP/blob/master/modules.py
class ProjectionHead(nn.Module):
def __init__(
self,
embedding_dim,
projection_dim,
dropout=0
):
super().__init__()
self.projection = nn.Linear(embedding_dim, projection_dim)
self.gelu = nn.GELU()
self.fc = nn.Linear(projection_dim, projection_dim)
self.dropout = nn.Dropout(dropout)
self.layer_norm = nn.LayerNorm(projection_dim)
def forward(self, x):
projected = self.projection(x)
x = self.gelu(projected)
x = self.fc(x)
x = self.dropout(x)
x = x + projected
x = self.layer_norm(x)
return x
def convert_state_dict(state_dict):
new_state_dict = OrderedDict()
for k, v in state_dict.items():
if k.startswith("module."):
k = k.replace("module.", "")
new_state_dict[k] = v
return new_state_dict
def init_weights(m):
if isinstance(m, nn.Linear):
torch.nn.init.xavier_uniform_(m.weight)
if m.bias is not None:
nn.init.normal_(m.bias, std=1e-6)
class Metric(nn.Module):
def __init__(self):
super().__init__()
self.image_preprocess = None
def load_image(self, image_path):
with open(image_path, 'rb') as f:
image = Image.open(f).convert("RGB")
return image
def load_image_path(self, image_path):
if isinstance(image_path, str):
# should be a image folder path
images_file = os.listdir(image_path)
images = [os.path.join(image_path, image) for image in images_file if
image.endswith(".jpg") or image.endswith(".png")]
if isinstance(image_path[0], str):
images = [self.load_image(image) for image in image_path]
elif isinstance(image_path[0], np.ndarray):
images = [Image.fromarray(image) for image in image_path]
elif isinstance(image_path[0], Image.Image):
images = image_path
else:
raise Exception("Invalid input")
return images
def preprocess_image(self, image, **kwargs):
if (isinstance(image, str) and os.path.isdir(image)) or (isinstance(image, list) and (isinstance(image[0], Image.Image) or isinstance(image[0], np.ndarray) or os.path.isfile(image[0]))):
input_data = self.load_image_path(image)
input_data = [self.image_preprocess(image, **kwargs) for image in input_data]
input_data = torch.stack(input_data)
elif os.path.isfile(image):
input_data = self.load_image(image)
input_data = self.image_preprocess(input_data, **kwargs)
input_data = input_data.unsqueeze(0)
elif isinstance(image, torch.Tensor):
raise Exception("Unsupported input")
return input_data
class Clip_Basic_Metric(Metric):
def __init__(self):
super().__init__()
self.tensor_preprocess = transforms.Compose([
transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC),
# transforms.rescale
transforms.Normalize(mean=[-1.0, -1.0, -1.0], std=[2.0, 2.0, 2.0]),
transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
])
self.image_preprocess = transforms.Compose([
transforms.Resize(size=224, interpolation=transforms.InterpolationMode.BICUBIC),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
])
class Clip_metric(Clip_Basic_Metric):
@torch.no_grad()
def __init__(self, target_style_prompt: str=None, clip_model_name="openai/clip-vit-large-patch14", device="cuda",
bath_size=8, alpha=0.5):
super().__init__()
self.device = device
self.alpha = alpha
self.model = (CLIPModel.from_pretrained(clip_model_name)).to(device)
self.processor = CLIPProcessor.from_pretrained(clip_model_name)
self.tokenizer = self.processor.tokenizer
self.image_processor = self.processor.image_processor
# self.style_class_features = self.get_text_features(self.styles).cpu()
self.style_class_features=[]
# self.noise_prompt_features = self.get_text_features("Noise")
self.model.eval()
self.batch_size = bath_size
if target_style_prompt is not None:
self.ref_style_features = self.get_text_features(target_style_prompt)
else:
self.ref_style_features = None
self.ref_image_style_prototype = None
def get_text_features(self, text):
prompt_encoding = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(self.device)
prompt_features = self.model.get_text_features(**prompt_encoding).to(self.device)
prompt_features = F.normalize(prompt_features, p=2, dim=-1)
return prompt_features
def get_image_features(self, images):
# if isinstance(image, torch.Tensor):
# self.tensor_transform(image)
# else:
# image_features = self.image_processor(image, return_tensors="pt", padding=True).to(self.device, non_blocking=True)
images = self.load_image_path(images)
if isinstance(images, torch.Tensor):
images = self.tensor_preprocess(images)
data = {"pixel_values": images}
image_features = BatchFeature(data=data, tensor_type="pt")
else:
image_features = self.image_processor(images, return_tensors="pt", padding=True).to(self.device,
non_blocking=True)
image_features = self.model.get_image_features(**image_features).to(self.device)
image_features = F.normalize(image_features, p=2, dim=-1)
return image_features
def img_text_similarity(self, image_features, text=None):
if text is not None:
prompt_feature = self.get_text_features(text)
if isinstance(text, str):
prompt_feature = prompt_feature.repeat(len(image_features), 1)
else:
prompt_feature = self.ref_style_features
similarity_each = torch.einsum("nc, nc -> n", image_features, prompt_feature)
return similarity_each
def forward(self, output_imgs, prompt=None):
image_features = self.get_image_features(output_imgs)
# print(image_features)
style_score = self.img_text_similarity(image_features.mean(dim=0, keepdim=True))
if prompt is not None:
content_score = self.img_text_similarity(image_features, prompt)
score = self.alpha * style_score + (1 - self.alpha) * content_score
return {"score": score, "style_score": style_score, "content_score": content_score}
else:
return {"style_score": style_score}
def content_score(self, output_imgs, prompt):
self.to(self.device)
image_features = self.get_image_features(output_imgs)
content_score_details = self.img_text_similarity(image_features, prompt)
self.to("cpu")
return {"CLIP_content_score": content_score_details.mean().cpu(), "CLIP_content_score_details": content_score_details.cpu()}
class CSD_CLIP(Clip_Basic_Metric):
"""backbone + projection head"""
def __init__(self, name='vit_large',content_proj_head='default', ckpt_path = "data/weights/CSD-checkpoint.pth", device="cuda",
alpha=0.5, **kwargs):
super(CSD_CLIP, self).__init__()
self.alpha = alpha
self.content_proj_head = content_proj_head
self.device = device
if name == 'vit_large':
clipmodel, _ = clip.load("ViT-L/14")
self.backbone = clipmodel.visual
self.embedding_dim = 1024
elif name == 'vit_base':
clipmodel, _ = clip.load("ViT-B/16")
self.backbone = clipmodel.visual
self.embedding_dim = 768
self.feat_dim = 512
else:
raise Exception('This model is not implemented')
convert_weights_float(self.backbone)
self.last_layer_style = copy.deepcopy(self.backbone.proj)
if content_proj_head == 'custom':
self.last_layer_content = ProjectionHead(self.embedding_dim,self.feat_dim)
self.last_layer_content.apply(init_weights)
else:
self.last_layer_content = copy.deepcopy(self.backbone.proj)
self.backbone.proj = None
self.backbone.requires_grad_(False)
self.last_layer_style.requires_grad_(False)
self.last_layer_content.requires_grad_(False)
self.backbone.eval()
if ckpt_path is not None:
self.load_ckpt(ckpt_path)
self.to("cpu")
def load_ckpt(self, ckpt_path):
checkpoint = torch.load(ckpt_path, map_location="cpu")
state_dict = convert_state_dict(checkpoint['model_state_dict'])
msg = self.load_state_dict(state_dict, strict=False)
print(f"=> loaded CSD_CLIP checkpoint with msg {msg}")
@property
def dtype(self):
return self.backbone.conv1.weight.dtype
def get_image_features(self, input_data, get_style=True,get_content=False,feature_alpha=None):
if isinstance(input_data, torch.Tensor):
input_data = self.tensor_preprocess(input_data)
elif (isinstance(input_data, str) and os.path.isdir(input_data)) or (isinstance(input_data, list) and (isinstance(input_data[0], Image.Image) or isinstance(input_data[0], np.ndarray) or os.path.isfile(input_data[0]))):
input_data = self.load_image_path(input_data)
input_data = [self.image_preprocess(image) for image in input_data]
input_data = torch.stack(input_data)
elif os.path.isfile(input_data):
input_data = self.load_image(input_data)
input_data = self.image_preprocess(input_data)
input_data = input_data.unsqueeze(0)
input_data = input_data.to(self.device)
style_output = None
feature = self.backbone(input_data)
if get_style:
style_output = feature @ self.last_layer_style
# style_output = style_output.mean(dim=0)
style_output = nn.functional.normalize(style_output, dim=-1, p=2)
content_output=None
if get_content:
if feature_alpha is not None:
reverse_feature = ReverseLayerF.apply(feature, feature_alpha)
else:
reverse_feature = feature
# if alpha is not None:
if self.content_proj_head == 'custom':
content_output = self.last_layer_content(reverse_feature)
else:
content_output = reverse_feature @ self.last_layer_content
content_output = nn.functional.normalize(content_output, dim=-1, p=2)
return feature, content_output, style_output
@torch.no_grad()
def define_ref_image_style_prototype(self, ref_image_path: str):
self.to(self.device)
_, _, self.ref_style_feature = self.get_image_features(ref_image_path)
self.to("cpu")
# self.ref_style_feature = self.ref_style_feature.mean(dim=0)
@torch.no_grad()
def forward(self, styled_data):
self.to(self.device)
# get_content_feature = original_data is not None
_, content_output, style_output = self.get_image_features(styled_data, get_content=False)
style_similarities = style_output @ self.ref_style_feature.T
mean_style_similarities = style_similarities.mean(dim=-1)
mean_style_similarity = mean_style_similarities.mean()
max_style_similarities_v, max_style_similarities_id = style_similarities.max(dim=-1)
max_style_similarity = max_style_similarities_v.mean()
self.to("cpu")
return {"CSD_similarity_mean": mean_style_similarity, "CSD_similarity_max": max_style_similarity, "CSD_similarity_mean_details": mean_style_similarities,
"CSD_similarity_max_v_details": max_style_similarities_v, "CSD_similarity_max_id_details": max_style_similarities_id}
def get_style_loss(self, styled_data):
_, _, style_output = self.get_image_features(styled_data, get_style=True, get_content=False)
style_similarity = (style_output @ self.ref_style_feature).mean()
loss = 1 - style_similarity
return loss.mean()
class LPIPS_metric(Metric):
def __init__(self, type="vgg", device="cuda"):
super(LPIPS_metric, self).__init__()
self.lpips = lpips.LPIPS(net=type)
self.device = device
self.image_preprocess = transforms.Compose([
transforms.Resize(256, interpolation=transforms.InterpolationMode.BICUBIC),
transforms.CenterCrop(256),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
self.to("cpu")
@torch.no_grad()
def forward(self, img1, img2):
self.to(self.device)
differences = []
for i in range(0, len(img1), 50):
img1_batch = img1[i:i+50]
img2_batch = img2[i:i+50]
img1_batch = self.preprocess_image(img1_batch).to(self.device)
img2_batch = self.preprocess_image(img2_batch).to(self.device)
differences.append(self.lpips(img1_batch, img2_batch).squeeze())
differences = torch.cat(differences)
difference = differences.mean()
# similarity = 1 - difference
self.to("cpu")
return {"LPIPS_content_difference": difference, "LPIPS_content_difference_details": differences}
class Vit_metric(Metric):
def __init__(self, device="cuda"):
super(Vit_metric, self).__init__()
self.device = device
self.model = ViTModel.from_pretrained('facebook/dino-vitb8').eval()
self.image_processor = ViTImageProcessor.from_pretrained('facebook/dino-vitb8')
self.to("cpu")
def get_image_features(self, images):
# if isinstance(image, torch.Tensor):
# self.tensor_transform(image)
# else:
# image_features = self.image_processor(image, return_tensors="pt", padding=True).to(self.device, non_blocking=True)
images = self.load_image_path(images)
batch_size = 20
all_image_features = []
for i in range(0, len(images), batch_size):
image_batch = images[i:i+batch_size]
if isinstance(image_batch, torch.Tensor):
image_batch = self.tensor_preprocess(image_batch)
data = {"pixel_values": image_batch}
image_processed = BatchFeature(data=data, tensor_type="pt")
else:
image_processed = self.image_processor(image_batch, return_tensors="pt").to(self.device)
image_features = self.model(**image_processed).last_hidden_state.flatten(start_dim=1)
image_features = F.normalize(image_features, p=2, dim=-1)
all_image_features.append(image_features)
all_image_features = torch.cat(all_image_features)
return all_image_features
@torch.no_grad()
def content_metric(self, img1, img2):
self.to(self.device)
if not(isinstance(img1, torch.Tensor) and len(img1.shape) == 2):
img1 = self.get_image_features(img1)
if not(isinstance(img2, torch.Tensor) and len(img2.shape) == 2):
img2 = self.get_image_features(img2)
similarities = torch.einsum("nc, nc -> n", img1, img2)
similarity = similarities.mean()
# self.to("cpu")
return {"Vit_content_similarity": similarity, "Vit_content_similarity_details": similarities}
# style
@torch.no_grad()
def define_ref_image_style_prototype(self, ref_image_path: str):
self.to(self.device)
self.ref_style_feature = self.get_image_features(ref_image_path)
self.to("cpu")
@torch.no_grad()
def style_metric(self, styled_data):
self.to(self.device)
if isinstance(styled_data, torch.Tensor) and len(styled_data.shape) == 2:
style_output = styled_data
else:
style_output = self.get_image_features(styled_data)
style_similarities = style_output @ self.ref_style_feature.T
mean_style_similarities = style_similarities.mean(dim=-1)
mean_style_similarity = mean_style_similarities.mean()
max_style_similarities_v, max_style_similarities_id = style_similarities.max(dim=-1)
max_style_similarity = max_style_similarities_v.mean()
# self.to("cpu")
return {"Vit_style_similarity_mean": mean_style_similarity, "Vit_style_similarity_max": max_style_similarity, "Vit_style_similarity_mean_details": mean_style_similarities,
"Vit_style_similarity_max_v_details": max_style_similarities_v, "Vit_style_similarity_max_id_details": max_style_similarities_id}
@torch.no_grad()
def forward(self, styled_data, original_data=None):
self.to(self.device)
styled_features = self.get_image_features(styled_data)
ret ={}
if original_data is not None:
content_metric = self.content_metric(styled_features, original_data)
ret["Vit_content"] = content_metric
style_metric = self.style_metric(styled_features)
ret["Vit_style"] = style_metric
self.to("cpu")
return ret
class StyleContentMetric(nn.Module):
def __init__(self, style_ref_image_folder, device="cuda"):
super(StyleContentMetric, self).__init__()
self.device = device
self.clip_style_metric = CSD_CLIP(device=device)
self.ref_image_file = os.listdir(style_ref_image_folder)
self.ref_image_file = [i for i in self.ref_image_file if i.endswith(".jpg") or i.endswith(".png")]
self.ref_image_file.sort()
self.ref_image_file = np.array(self.ref_image_file)
ref_image_file_path = [os.path.join(style_ref_image_folder, i) for i in self.ref_image_file]
self.clip_style_metric.define_ref_image_style_prototype(ref_image_file_path)
self.vit_metric = Vit_metric(device=device)
self.vit_metric.define_ref_image_style_prototype(ref_image_file_path)
self.lpips_metric = LPIPS_metric(device=device)
self.clip_content_metric = Clip_metric(alpha=0, target_style_prompt=None)
self.to("cpu")
def forward(self, styled_data, original_data=None, content_caption=None):
ret ={}
csd_score = self.clip_style_metric(styled_data)
csd_score["max_query"] = self.ref_image_file[csd_score["CSD_similarity_max_id_details"].cpu()].tolist()
torch.cuda.empty_cache()
ret["Style_CSD"] = csd_score
vit_score = self.vit_metric(styled_data, original_data)
torch.cuda.empty_cache()
vit_style = vit_score["Vit_style"]
vit_style["max_query"] = self.ref_image_file[vit_style["Vit_style_similarity_max_id_details"].cpu()].tolist()
ret["Style_VIT"] = vit_style
if original_data is not None:
vit_content = vit_score["Vit_content"]
ret["Content_VIT"] = vit_content
lpips_score = self.lpips_metric(styled_data, original_data)
torch.cuda.empty_cache()
ret["Content_LPIPS"] = lpips_score
if content_caption is not None:
clip_content = self.clip_content_metric.content_score(styled_data, content_caption)
ret["Content_CLIP"] = clip_content
torch.cuda.empty_cache()
for type_key, type_value in ret.items():
for key, value in type_value.items():
if isinstance(value, torch.Tensor):
if value.numel() == 1:
ret[type_key][key] = round(value.item(), 4)
else:
ret[type_key][key] = value.tolist()
ret[type_key][key] = [round(v, 4) for v in ret[type_key][key]]
self.to("cpu")
ret["ref_image_file"] = self.ref_image_file.tolist()
return ret
if __name__ == "__main__":
with torch.no_grad():
metric = StyleContentMetric("/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/clip_dissection/Art_styles/camille-pissarro/impressionism/split_5/paintings")
score = metric("/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/gustav-klimt_Art_Nouveau/subset3/converted_photo/500",
"/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/gustav-klimt_Art_Nouveau/subset3/paintings")
print(score)
lpips = LPIPS_metric()
score = lpips("/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/gustav-klimt_Art_Nouveau/subset3/paintings",
"/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/gustav-klimt_Art_Nouveau/subset3/converted_photo/500")
print("lpips", score)
clip_metric = CSD_CLIP()
clip_metric.define_ref_image_style_prototype(
"/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/gustav-klimt_Art_Nouveau/subset1/paintings")
score = clip_metric(
"/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/gustav-klimt_Art_Nouveau/subset3/converted_photo/500")
print("subset3-subset3_sd14_converted", score)
score = clip_metric(
"/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/imgFolder/clip_filtered_remain_500")
print("subset3-photo", score)
score = clip_metric(
"/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/gustav-klimt_Art_Nouveau/subset1/paintings")
print("subset3-subset1", score)
score = clip_metric(
"/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/andy-warhol/pop_art/subset1/paintings")
print("subset3-andy", score)
# score = clip_metric("/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/gustav-klimt_Art_Nouveau/subset3/paintings", "A painting")
# print("subset3",score)
# score_subset2 = clip_metric("/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/gustav-klimt_Art_Nouveau/subset2/paintings")
# print("subset2",score_subset2)
# score_subset3 = clip_metric("/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/gustav-klimt_Art_Nouveau/subset3/paintings")
# print("subset3",score_subset3)
#
# score_subset3_converted = clip_metric("/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/gustav-klimt_Art_Nouveau/subset3/converted_photo/500")
# print("subset3-subset3_sd14_converted" , score_subset3_converted)
#
# score_subset3_coco_converted = clip_metric("/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/gustav-klimt_Art_Nouveau/subset3/coco_converted_photo/500")
# print("subset3-subset3_coco_converted" , score_subset3_coco_converted)
#
# clip_metric = Clip_metric("/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/imgFolder/sketch_500")
# score = clip_metric("/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/imgFolder/clip_filtered_remain_500")
# print("photo500_1-sketch" ,score)
#
# clip_metric = Clip_metric("/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/imgFolder/clip_filtered_remain_500")
# score = clip_metric("/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/imgFolder/clip_filtered_remain_500_new")
# print("photo500_1-photo500_2" ,score)
# from custom_datasets.imagepair import ImageSet
# import matplotlib.pyplot as plt
# dataset = ImageSet(folder = "/data/vision/torralba/scratch/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/gustav-klimt_Art_Nouveau/subset3/paintings",
# caption_path="/data/vision/torralba/scratch/huiren/code/diffusion/stable_diffusion/custom_datasets/wikiart/data/gustav-klimt_Art_Nouveau/subset3/captions",
# keep_in_mem=False)
# for sample in dataset:
# score = clip_metric.content_score(sample["image"], sample["caption"][0])
# plt.imshow(sample["image"])
# plt.title(f"score: {round(score.item(),2)}\n prompt: {sample['caption'][0]}")
# plt.show()