pharmapsychotic
2 years ago
4 changed files with 8 additions and 341 deletions
@ -1,341 +1,18 @@ |
|||||||
import sys |
import sys |
||||||
|
|
||||||
sys.path.append("src/clip") |
|
||||||
sys.path.append("src/blip") |
|
||||||
|
|
||||||
import os |
|
||||||
import hashlib |
|
||||||
import math |
|
||||||
import numpy as np |
|
||||||
import pickle |
|
||||||
from tqdm import tqdm |
|
||||||
from PIL import Image |
from PIL import Image |
||||||
import torch |
|
||||||
from torchvision import transforms |
|
||||||
from torchvision.transforms.functional import InterpolationMode |
|
||||||
import clip |
|
||||||
from models.blip import blip_decoder |
|
||||||
from cog import BasePredictor, Input, Path |
from cog import BasePredictor, Input, Path |
||||||
|
|
||||||
|
sys.path.extend(["src/clip", "src/blip"]) |
||||||
|
|
||||||
DATA_PATH = "data" |
from clip_interrogator import Interrogator, Config |
||||||
chunk_size = 2048 |
|
||||||
flavor_intermediate_count = 2048 |
|
||||||
blip_image_eval_size = 384 |
|
||||||
|
|
||||||
|
|
||||||
class Predictor(BasePredictor): |
class Predictor(BasePredictor): |
||||||
def setup(self): |
def setup(self): |
||||||
"""Load the model into memory to make running multiple predictions efficient""" |
config = Config(device="cuda:0", clip_model_name='ViT-L/14') |
||||||
|
self.ci = Interrogator(config) |
||||||
self.device = "cuda:0" |
|
||||||
|
|
||||||
print("Loading BLIP model...") |
|
||||||
self.blip_model = blip_decoder( |
|
||||||
pretrained="weights/model_large_caption.pth", # downloaded with wget https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_large_caption.pth |
|
||||||
image_size=blip_image_eval_size, |
|
||||||
vit="large", |
|
||||||
med_config="src/blip/configs/med_config.json", |
|
||||||
) |
|
||||||
self.blip_model.eval() |
|
||||||
self.blip_model = self.blip_model.to(self.device) |
|
||||||
|
|
||||||
print("Loading CLIP model...") |
|
||||||
self.clip_models, self.clip_preprocess = {}, {} |
|
||||||
for clip_model_name in [ |
|
||||||
"ViT-B/32", |
|
||||||
"ViT-B/16", |
|
||||||
"ViT-L/14", |
|
||||||
"ViT-L/14@336px", |
|
||||||
"RN101", |
|
||||||
"RN50", |
|
||||||
"RN50x4", |
|
||||||
"RN50x16", |
|
||||||
"RN50x64", |
|
||||||
]: |
|
||||||
( |
|
||||||
self.clip_models[clip_model_name], |
|
||||||
self.clip_preprocess[clip_model_name], |
|
||||||
) = clip.load(clip_model_name, device=self.device) |
|
||||||
self.clip_models[clip_model_name].cuda().eval() |
|
||||||
|
|
||||||
sites = [ |
|
||||||
"Artstation", |
|
||||||
"behance", |
|
||||||
"cg society", |
|
||||||
"cgsociety", |
|
||||||
"deviantart", |
|
||||||
"dribble", |
|
||||||
"flickr", |
|
||||||
"instagram", |
|
||||||
"pexels", |
|
||||||
"pinterest", |
|
||||||
"pixabay", |
|
||||||
"pixiv", |
|
||||||
"polycount", |
|
||||||
"reddit", |
|
||||||
"shutterstock", |
|
||||||
"tumblr", |
|
||||||
"unsplash", |
|
||||||
"zbrush central", |
|
||||||
] |
|
||||||
self.trending_list = [site for site in sites] |
|
||||||
self.trending_list.extend(["trending on " + site for site in sites]) |
|
||||||
self.trending_list.extend(["featured on " + site for site in sites]) |
|
||||||
self.trending_list.extend([site + " contest winner" for site in sites]) |
|
||||||
raw_artists = load_list(f"{DATA_PATH}/artists.txt") |
|
||||||
self.artists = [f"by {a}" for a in raw_artists] |
|
||||||
self.artists.extend([f"inspired by {a}" for a in raw_artists]) |
|
||||||
|
|
||||||
def predict( |
def predict(self, image: Path = Input(description="Input image")) -> str: |
||||||
self, |
|
||||||
image: Path = Input(description="Input image"), |
|
||||||
clip_model_name: str = Input( |
|
||||||
default="ViT-L/14", |
|
||||||
choices=[ |
|
||||||
"ViT-B/32", |
|
||||||
"ViT-B/16", |
|
||||||
"ViT-L/14", |
|
||||||
"ViT-L/14@336px", |
|
||||||
"RN101", |
|
||||||
"RN50", |
|
||||||
"RN50x4", |
|
||||||
"RN50x16", |
|
||||||
"RN50x64", |
|
||||||
], |
|
||||||
description="Choose a clip model.", |
|
||||||
), |
|
||||||
) -> str: |
|
||||||
"""Run a single prediction on the model""" |
"""Run a single prediction on the model""" |
||||||
clip_model = self.clip_models[clip_model_name] |
|
||||||
clip_preprocess = self.clip_preprocess[clip_model_name] |
|
||||||
|
|
||||||
artists = LabelTable(self.artists, "artists", clip_model_name, clip_model) |
|
||||||
flavors = LabelTable( |
|
||||||
load_list(f"{DATA_PATH}/flavors.txt"), |
|
||||||
"flavors", |
|
||||||
clip_model_name, |
|
||||||
clip_model, |
|
||||||
) |
|
||||||
mediums = LabelTable( |
|
||||||
load_list(f"{DATA_PATH}/mediums.txt"), |
|
||||||
"mediums", |
|
||||||
clip_model_name, |
|
||||||
clip_model, |
|
||||||
) |
|
||||||
movements = LabelTable( |
|
||||||
load_list(f"{DATA_PATH}/movements.txt"), |
|
||||||
"movements", |
|
||||||
clip_model_name, |
|
||||||
clip_model, |
|
||||||
) |
|
||||||
trendings = LabelTable( |
|
||||||
self.trending_list, "trendings", clip_model_name, clip_model |
|
||||||
) |
|
||||||
|
|
||||||
image = Image.open(str(image)).convert("RGB") |
image = Image.open(str(image)).convert("RGB") |
||||||
|
return self.ci.interrogate(image) |
||||||
labels = [flavors, mediums, artists, trendings, movements] |
|
||||||
|
|
||||||
prompt = interrogate( |
|
||||||
image, |
|
||||||
clip_model_name, |
|
||||||
clip_preprocess, |
|
||||||
clip_model, |
|
||||||
self.blip_model, |
|
||||||
*labels, |
|
||||||
) |
|
||||||
|
|
||||||
return prompt |
|
||||||
|
|
||||||
|
|
||||||
class LabelTable: |
|
||||||
def __init__(self, labels, desc, clip_model_name, clip_model): |
|
||||||
self.labels = labels |
|
||||||
self.embeds = [] |
|
||||||
|
|
||||||
hash = hashlib.sha256(",".join(labels).encode()).hexdigest() |
|
||||||
|
|
||||||
os.makedirs("./cache", exist_ok=True) |
|
||||||
cache_filepath = f"./cache/{desc}.pkl" |
|
||||||
if desc is not None and os.path.exists(cache_filepath): |
|
||||||
with open(cache_filepath, "rb") as f: |
|
||||||
data = pickle.load(f) |
|
||||||
if data.get("hash") == hash and data.get("model") == clip_model_name: |
|
||||||
self.labels = data["labels"] |
|
||||||
self.embeds = data["embeds"] |
|
||||||
|
|
||||||
if len(self.labels) != len(self.embeds): |
|
||||||
self.embeds = [] |
|
||||||
chunks = np.array_split(self.labels, max(1, len(self.labels) / chunk_size)) |
|
||||||
for chunk in tqdm(chunks, desc=f"Preprocessing {desc}" if desc else None): |
|
||||||
text_tokens = clip.tokenize(chunk).cuda() |
|
||||||
with torch.no_grad(): |
|
||||||
text_features = clip_model.encode_text(text_tokens).float() |
|
||||||
text_features /= text_features.norm(dim=-1, keepdim=True) |
|
||||||
text_features = text_features.half().cpu().numpy() |
|
||||||
for i in range(text_features.shape[0]): |
|
||||||
self.embeds.append(text_features[i]) |
|
||||||
|
|
||||||
with open(cache_filepath, "wb") as f: |
|
||||||
pickle.dump( |
|
||||||
{ |
|
||||||
"labels": self.labels, |
|
||||||
"embeds": self.embeds, |
|
||||||
"hash": hash, |
|
||||||
"model": clip_model_name, |
|
||||||
}, |
|
||||||
f, |
|
||||||
) |
|
||||||
|
|
||||||
def _rank(self, image_features, text_embeds, device="cuda", top_count=1): |
|
||||||
top_count = min(top_count, len(text_embeds)) |
|
||||||
similarity = torch.zeros((1, len(text_embeds))).to(device) |
|
||||||
text_embeds = ( |
|
||||||
torch.stack([torch.from_numpy(t) for t in text_embeds]).float().to(device) |
|
||||||
) |
|
||||||
for i in range(image_features.shape[0]): |
|
||||||
similarity += (image_features[i].unsqueeze(0) @ text_embeds.T).softmax( |
|
||||||
dim=-1 |
|
||||||
) |
|
||||||
_, top_labels = similarity.cpu().topk(top_count, dim=-1) |
|
||||||
return [top_labels[0][i].numpy() for i in range(top_count)] |
|
||||||
|
|
||||||
def rank(self, image_features, top_count=1): |
|
||||||
if len(self.labels) <= chunk_size: |
|
||||||
tops = self._rank(image_features, self.embeds, top_count=top_count) |
|
||||||
return [self.labels[i] for i in tops] |
|
||||||
|
|
||||||
num_chunks = int(math.ceil(len(self.labels) / chunk_size)) |
|
||||||
keep_per_chunk = int(chunk_size / num_chunks) |
|
||||||
|
|
||||||
top_labels, top_embeds = [], [] |
|
||||||
for chunk_idx in tqdm(range(num_chunks)): |
|
||||||
start = chunk_idx * chunk_size |
|
||||||
stop = min(start + chunk_size, len(self.embeds)) |
|
||||||
tops = self._rank( |
|
||||||
image_features, self.embeds[start:stop], top_count=keep_per_chunk |
|
||||||
) |
|
||||||
top_labels.extend([self.labels[start + i] for i in tops]) |
|
||||||
top_embeds.extend([self.embeds[start + i] for i in tops]) |
|
||||||
|
|
||||||
tops = self._rank(image_features, top_embeds, top_count=top_count) |
|
||||||
return [top_labels[i] for i in tops] |
|
||||||
|
|
||||||
|
|
||||||
def generate_caption(pil_image, blip_model, device="cuda"): |
|
||||||
gpu_image = ( |
|
||||||
transforms.Compose( |
|
||||||
[ |
|
||||||
transforms.Resize( |
|
||||||
(blip_image_eval_size, blip_image_eval_size), |
|
||||||
interpolation=InterpolationMode.BICUBIC, |
|
||||||
), |
|
||||||
transforms.ToTensor(), |
|
||||||
transforms.Normalize( |
|
||||||
(0.48145466, 0.4578275, 0.40821073), |
|
||||||
(0.26862954, 0.26130258, 0.27577711), |
|
||||||
), |
|
||||||
] |
|
||||||
)(pil_image) |
|
||||||
.unsqueeze(0) |
|
||||||
.to(device) |
|
||||||
) |
|
||||||
|
|
||||||
with torch.no_grad(): |
|
||||||
caption = blip_model.generate( |
|
||||||
gpu_image, sample=False, num_beams=3, max_length=20, min_length=5 |
|
||||||
) |
|
||||||
return caption[0] |
|
||||||
|
|
||||||
|
|
||||||
def rank_top(image_features, text_array, clip_model, device="cuda"): |
|
||||||
text_tokens = clip.tokenize([text for text in text_array]).cuda() |
|
||||||
with torch.no_grad(): |
|
||||||
text_features = clip_model.encode_text(text_tokens).float() |
|
||||||
text_features /= text_features.norm(dim=-1, keepdim=True) |
|
||||||
|
|
||||||
similarity = torch.zeros((1, len(text_array)), device=device) |
|
||||||
for i in range(image_features.shape[0]): |
|
||||||
similarity += (image_features[i].unsqueeze(0) @ text_features.T).softmax(dim=-1) |
|
||||||
|
|
||||||
_, top_labels = similarity.cpu().topk(1, dim=-1) |
|
||||||
return text_array[top_labels[0][0].numpy()] |
|
||||||
|
|
||||||
|
|
||||||
def similarity(image_features, text, clip_model): |
|
||||||
text_tokens = clip.tokenize([text]).cuda() |
|
||||||
with torch.no_grad(): |
|
||||||
text_features = clip_model.encode_text(text_tokens).float() |
|
||||||
text_features /= text_features.norm(dim=-1, keepdim=True) |
|
||||||
similarity = text_features.cpu().numpy() @ image_features.cpu().numpy().T |
|
||||||
return similarity[0][0] |
|
||||||
|
|
||||||
|
|
||||||
def load_list(filename): |
|
||||||
with open(filename, "r", encoding="utf-8", errors="replace") as f: |
|
||||||
items = [line.strip() for line in f.readlines()] |
|
||||||
return items |
|
||||||
|
|
||||||
|
|
||||||
def interrogate(image, clip_model_name, clip_preprocess, clip_model, blip_model, *args): |
|
||||||
flavors, mediums, artists, trendings, movements = args |
|
||||||
caption = generate_caption(image, blip_model) |
|
||||||
|
|
||||||
images = clip_preprocess(image).unsqueeze(0).cuda() |
|
||||||
with torch.no_grad(): |
|
||||||
image_features = clip_model.encode_image(images).float() |
|
||||||
image_features /= image_features.norm(dim=-1, keepdim=True) |
|
||||||
|
|
||||||
flaves = flavors.rank(image_features, flavor_intermediate_count) |
|
||||||
best_medium = mediums.rank(image_features, 1)[0] |
|
||||||
best_artist = artists.rank(image_features, 1)[0] |
|
||||||
best_trending = trendings.rank(image_features, 1)[0] |
|
||||||
best_movement = movements.rank(image_features, 1)[0] |
|
||||||
|
|
||||||
best_prompt = caption |
|
||||||
best_sim = similarity(image_features, best_prompt, clip_model) |
|
||||||
|
|
||||||
def check(addition): |
|
||||||
nonlocal best_prompt, best_sim |
|
||||||
prompt = best_prompt + ", " + addition |
|
||||||
sim = similarity(image_features, prompt, clip_model) |
|
||||||
if sim > best_sim: |
|
||||||
best_sim = sim |
|
||||||
best_prompt = prompt |
|
||||||
return True |
|
||||||
return False |
|
||||||
|
|
||||||
def check_multi_batch(opts): |
|
||||||
nonlocal best_prompt, best_sim |
|
||||||
prompts = [] |
|
||||||
for i in range(2 ** len(opts)): |
|
||||||
prompt = best_prompt |
|
||||||
for bit in range(len(opts)): |
|
||||||
if i & (1 << bit): |
|
||||||
prompt += ", " + opts[bit] |
|
||||||
prompts.append(prompt) |
|
||||||
|
|
||||||
t = LabelTable(prompts, None, clip_model_name, clip_model) |
|
||||||
best_prompt = t.rank(image_features, 1)[0] |
|
||||||
best_sim = similarity(image_features, best_prompt, clip_model) |
|
||||||
|
|
||||||
check_multi_batch([best_medium, best_artist, best_trending, best_movement]) |
|
||||||
|
|
||||||
extended_flavors = set(flaves) |
|
||||||
for _ in tqdm(range(25), desc="Flavor chain"): |
|
||||||
try: |
|
||||||
best = rank_top( |
|
||||||
image_features, |
|
||||||
[f"{best_prompt}, {f}" for f in extended_flavors], |
|
||||||
clip_model, |
|
||||||
) |
|
||||||
flave = best[len(best_prompt) + 2 :] |
|
||||||
if not check(flave): |
|
||||||
break |
|
||||||
extended_flavors.remove(flave) |
|
||||||
except: |
|
||||||
# exceeded max prompt length |
|
||||||
break |
|
||||||
|
|
||||||
return best_prompt |
|
||||||
|
|||||||
Loading…
Reference in new issue