Spaces:
Running
on
A10G
Running
on
A10G
Commit
•
c46775f
1
Parent(s):
ec4e754
Create app.py
Browse files
app.py
CHANGED
@@ -1,312 +0,0 @@
|
|
1 |
-
import sys
|
2 |
-
sys.path.append('src/blip')
|
3 |
-
sys.path.append('src/clip')
|
4 |
-
|
5 |
-
import clip
|
6 |
-
import gradio as gr
|
7 |
-
import hashlib
|
8 |
-
import math
|
9 |
-
import numpy as np
|
10 |
-
import os
|
11 |
-
import pickle
|
12 |
-
import torch
|
13 |
-
import torchvision.transforms as T
|
14 |
-
import torchvision.transforms.functional as TF
|
15 |
-
|
16 |
-
from models.blip import blip_decoder
|
17 |
-
from PIL import Image
|
18 |
-
from torch import nn
|
19 |
-
from torch.nn import functional as F
|
20 |
-
from tqdm import tqdm
|
21 |
-
|
22 |
-
from share_btn import community_icon_html, loading_icon_html, share_js
|
23 |
-
|
24 |
-
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
25 |
-
|
26 |
-
print("Loading BLIP model...")
|
27 |
-
blip_image_eval_size = 384
|
28 |
-
blip_model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_large_caption.pth'
|
29 |
-
blip_model = blip_decoder(pretrained=blip_model_url, image_size=blip_image_eval_size, vit='large', med_config='./src/blip/configs/med_config.json')
|
30 |
-
blip_model.eval()
|
31 |
-
blip_model = blip_model.to(device)
|
32 |
-
|
33 |
-
print("Loading CLIP model...")
|
34 |
-
clip_model_name = 'ViT-L/14' # https://huggingface.co/openai/clip-vit-large-patch14
|
35 |
-
clip_model, clip_preprocess = clip.load(clip_model_name, device=device)
|
36 |
-
clip_model.to(device).eval()
|
37 |
-
|
38 |
-
chunk_size = 2048
|
39 |
-
flavor_intermediate_count = 2048
|
40 |
-
|
41 |
-
|
42 |
-
class LabelTable():
|
43 |
-
def __init__(self, labels, desc):
|
44 |
-
self.labels = labels
|
45 |
-
self.embeds = []
|
46 |
-
|
47 |
-
hash = hashlib.sha256(",".join(labels).encode()).hexdigest()
|
48 |
-
|
49 |
-
os.makedirs('./cache', exist_ok=True)
|
50 |
-
cache_filepath = f"./cache/{desc}.pkl"
|
51 |
-
if desc is not None and os.path.exists(cache_filepath):
|
52 |
-
with open(cache_filepath, 'rb') as f:
|
53 |
-
data = pickle.load(f)
|
54 |
-
if data['hash'] == hash:
|
55 |
-
self.labels = data['labels']
|
56 |
-
self.embeds = data['embeds']
|
57 |
-
|
58 |
-
if len(self.labels) != len(self.embeds):
|
59 |
-
self.embeds = []
|
60 |
-
chunks = np.array_split(self.labels, max(1, len(self.labels)/chunk_size))
|
61 |
-
for chunk in tqdm(chunks, desc=f"Preprocessing {desc}" if desc else None):
|
62 |
-
text_tokens = clip.tokenize(chunk).to(device)
|
63 |
-
with torch.no_grad():
|
64 |
-
text_features = clip_model.encode_text(text_tokens).float()
|
65 |
-
text_features /= text_features.norm(dim=-1, keepdim=True)
|
66 |
-
text_features = text_features.half().cpu().numpy()
|
67 |
-
for i in range(text_features.shape[0]):
|
68 |
-
self.embeds.append(text_features[i])
|
69 |
-
|
70 |
-
with open(cache_filepath, 'wb') as f:
|
71 |
-
pickle.dump({"labels":self.labels, "embeds":self.embeds, "hash":hash}, f)
|
72 |
-
|
73 |
-
def _rank(self, image_features, text_embeds, top_count=1):
|
74 |
-
top_count = min(top_count, len(text_embeds))
|
75 |
-
similarity = torch.zeros((1, len(text_embeds))).to(device)
|
76 |
-
text_embeds = torch.stack([torch.from_numpy(t) for t in text_embeds]).float().to(device)
|
77 |
-
for i in range(image_features.shape[0]):
|
78 |
-
similarity += (image_features[i].unsqueeze(0) @ text_embeds.T).softmax(dim=-1)
|
79 |
-
_, top_labels = similarity.cpu().topk(top_count, dim=-1)
|
80 |
-
return [top_labels[0][i].numpy() for i in range(top_count)]
|
81 |
-
|
82 |
-
def rank(self, image_features, top_count=1):
|
83 |
-
if len(self.labels) <= chunk_size:
|
84 |
-
tops = self._rank(image_features, self.embeds, top_count=top_count)
|
85 |
-
return [self.labels[i] for i in tops]
|
86 |
-
|
87 |
-
num_chunks = int(math.ceil(len(self.labels)/chunk_size))
|
88 |
-
keep_per_chunk = int(chunk_size / num_chunks)
|
89 |
-
|
90 |
-
top_labels, top_embeds = [], []
|
91 |
-
for chunk_idx in tqdm(range(num_chunks)):
|
92 |
-
start = chunk_idx*chunk_size
|
93 |
-
stop = min(start+chunk_size, len(self.embeds))
|
94 |
-
tops = self._rank(image_features, self.embeds[start:stop], top_count=keep_per_chunk)
|
95 |
-
top_labels.extend([self.labels[start+i] for i in tops])
|
96 |
-
top_embeds.extend([self.embeds[start+i] for i in tops])
|
97 |
-
|
98 |
-
tops = self._rank(image_features, top_embeds, top_count=top_count)
|
99 |
-
return [top_labels[i] for i in tops]
|
100 |
-
|
101 |
-
def generate_caption(pil_image):
|
102 |
-
gpu_image = T.Compose([
|
103 |
-
T.Resize((blip_image_eval_size, blip_image_eval_size), interpolation=TF.InterpolationMode.BICUBIC),
|
104 |
-
T.ToTensor(),
|
105 |
-
T.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
|
106 |
-
])(pil_image).unsqueeze(0).to(device)
|
107 |
-
|
108 |
-
with torch.no_grad():
|
109 |
-
caption = blip_model.generate(gpu_image, sample=False, num_beams=3, max_length=20, min_length=5)
|
110 |
-
return caption[0]
|
111 |
-
|
112 |
-
def load_list(filename):
|
113 |
-
with open(filename, 'r', encoding='utf-8', errors='replace') as f:
|
114 |
-
items = [line.strip() for line in f.readlines()]
|
115 |
-
return items
|
116 |
-
|
117 |
-
def rank_top(image_features, text_array):
|
118 |
-
text_tokens = clip.tokenize([text for text in text_array]).to(device)
|
119 |
-
with torch.no_grad():
|
120 |
-
text_features = clip_model.encode_text(text_tokens).float()
|
121 |
-
text_features /= text_features.norm(dim=-1, keepdim=True)
|
122 |
-
|
123 |
-
similarity = torch.zeros((1, len(text_array)), device=device)
|
124 |
-
for i in range(image_features.shape[0]):
|
125 |
-
similarity += (image_features[i].unsqueeze(0) @ text_features.T).softmax(dim=-1)
|
126 |
-
|
127 |
-
_, top_labels = similarity.cpu().topk(1, dim=-1)
|
128 |
-
return text_array[top_labels[0][0].numpy()]
|
129 |
-
|
130 |
-
def similarity(image_features, text):
|
131 |
-
text_tokens = clip.tokenize([text]).to(device)
|
132 |
-
with torch.no_grad():
|
133 |
-
text_features = clip_model.encode_text(text_tokens).float()
|
134 |
-
text_features /= text_features.norm(dim=-1, keepdim=True)
|
135 |
-
similarity = text_features.cpu().numpy() @ image_features.cpu().numpy().T
|
136 |
-
return similarity[0][0]
|
137 |
-
|
138 |
-
def interrogate(image):
|
139 |
-
caption = generate_caption(image)
|
140 |
-
|
141 |
-
images = clip_preprocess(image).unsqueeze(0).to(device)
|
142 |
-
with torch.no_grad():
|
143 |
-
image_features = clip_model.encode_image(images).float()
|
144 |
-
image_features /= image_features.norm(dim=-1, keepdim=True)
|
145 |
-
|
146 |
-
flaves = flavors.rank(image_features, flavor_intermediate_count)
|
147 |
-
best_medium = mediums.rank(image_features, 1)[0]
|
148 |
-
best_artist = artists.rank(image_features, 1)[0]
|
149 |
-
best_trending = trendings.rank(image_features, 1)[0]
|
150 |
-
best_movement = movements.rank(image_features, 1)[0]
|
151 |
-
|
152 |
-
best_prompt = caption
|
153 |
-
best_sim = similarity(image_features, best_prompt)
|
154 |
-
|
155 |
-
def check(addition):
|
156 |
-
nonlocal best_prompt, best_sim
|
157 |
-
prompt = best_prompt + ", " + addition
|
158 |
-
sim = similarity(image_features, prompt)
|
159 |
-
if sim > best_sim:
|
160 |
-
best_sim = sim
|
161 |
-
best_prompt = prompt
|
162 |
-
return True
|
163 |
-
return False
|
164 |
-
|
165 |
-
def check_multi_batch(opts):
|
166 |
-
nonlocal best_prompt, best_sim
|
167 |
-
prompts = []
|
168 |
-
for i in range(2**len(opts)):
|
169 |
-
prompt = best_prompt
|
170 |
-
for bit in range(len(opts)):
|
171 |
-
if i & (1 << bit):
|
172 |
-
prompt += ", " + opts[bit]
|
173 |
-
prompts.append(prompt)
|
174 |
-
|
175 |
-
prompt = rank_top(image_features, prompts)
|
176 |
-
sim = similarity(image_features, prompt)
|
177 |
-
if sim > best_sim:
|
178 |
-
best_sim = sim
|
179 |
-
best_prompt = prompt
|
180 |
-
|
181 |
-
check_multi_batch([best_medium, best_artist, best_trending, best_movement])
|
182 |
-
|
183 |
-
extended_flavors = set(flaves)
|
184 |
-
for _ in tqdm(range(25), desc="Flavor chain"):
|
185 |
-
try:
|
186 |
-
best = rank_top(image_features, [f"{best_prompt}, {f}" for f in extended_flavors])
|
187 |
-
flave = best[len(best_prompt)+2:]
|
188 |
-
if not check(flave):
|
189 |
-
break
|
190 |
-
extended_flavors.remove(flave)
|
191 |
-
except:
|
192 |
-
# exceeded max prompt length
|
193 |
-
break
|
194 |
-
|
195 |
-
return best_prompt
|
196 |
-
|
197 |
-
|
198 |
-
sites = ['Artstation', 'behance', 'cg society', 'cgsociety', 'deviantart', 'dribble', 'flickr', 'instagram', 'pexels', 'pinterest', 'pixabay', 'pixiv', 'polycount', 'reddit', 'shutterstock', 'tumblr', 'unsplash', 'zbrush central']
|
199 |
-
trending_list = [site for site in sites]
|
200 |
-
trending_list.extend(["trending on "+site for site in sites])
|
201 |
-
trending_list.extend(["featured on "+site for site in sites])
|
202 |
-
trending_list.extend([site+" contest winner" for site in sites])
|
203 |
-
|
204 |
-
raw_artists = load_list('data/artists.txt')
|
205 |
-
artists = [f"by {a}" for a in raw_artists]
|
206 |
-
artists.extend([f"inspired by {a}" for a in raw_artists])
|
207 |
-
|
208 |
-
artists = LabelTable(artists, "artists")
|
209 |
-
flavors = LabelTable(load_list('data/flavors.txt'), "flavors")
|
210 |
-
mediums = LabelTable(load_list('data/mediums.txt'), "mediums")
|
211 |
-
movements = LabelTable(load_list('data/movements.txt'), "movements")
|
212 |
-
trendings = LabelTable(trending_list, "trendings")
|
213 |
-
|
214 |
-
|
215 |
-
def inference(image):
|
216 |
-
return interrogate(image), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
|
217 |
-
|
218 |
-
title = """
|
219 |
-
<div style="text-align: center; max-width: 650px; margin: 0 auto;">
|
220 |
-
<div
|
221 |
-
style="
|
222 |
-
display: inline-flex;
|
223 |
-
align-items: center;
|
224 |
-
gap: 0.8rem;
|
225 |
-
font-size: 1.75rem;
|
226 |
-
"
|
227 |
-
>
|
228 |
-
<h1 style="font-weight: 900; margin-bottom: 7px;">
|
229 |
-
CLIP Interrogator
|
230 |
-
</h1>
|
231 |
-
</div>
|
232 |
-
<p style="margin-bottom: 10px; font-size: 94%">
|
233 |
-
Want to figure out what a good prompt might be to create new images like an existing one? The CLIP Interrogator is here to get you answers!
|
234 |
-
</p>
|
235 |
-
</div>
|
236 |
-
"""
|
237 |
-
article = """
|
238 |
-
<div style="text-align: center; max-width: 650px; margin: 0 auto;">
|
239 |
-
<p>
|
240 |
-
Example art by <a href="https://pixabay.com/illustrations/watercolour-painting-art-effect-4799014/">Layers</a>
|
241 |
-
and <a href="https://pixabay.com/illustrations/animal-painting-cat-feline-pet-7154059/">Lin Tong</a>
|
242 |
-
from pixabay.com
|
243 |
-
</p>
|
244 |
-
|
245 |
-
<p>
|
246 |
-
Server busy? You can also run on <a href="https://colab.research.google.com/github/pharmapsychotic/clip-interrogator/blob/main/clip_interrogator.ipynb">Google Colab</a>
|
247 |
-
</p>
|
248 |
-
|
249 |
-
<p>
|
250 |
-
Has this been helpful to you? Follow me on twitter
|
251 |
-
<a href="https://twitter.com/pharmapsychotic">@pharmapsychotic</a>
|
252 |
-
and check out more tools at my
|
253 |
-
<a href="https://pharmapsychotic.com/tools.html">Ai generative art tools list</a>
|
254 |
-
</p>
|
255 |
-
</div>
|
256 |
-
"""
|
257 |
-
|
258 |
-
css = '''
|
259 |
-
#col-container {max-width: 700px; margin-left: auto; margin-right: auto;}
|
260 |
-
a {text-decoration-line: underline; font-weight: 600;}
|
261 |
-
.animate-spin {
|
262 |
-
animation: spin 1s linear infinite;
|
263 |
-
}
|
264 |
-
@keyframes spin {
|
265 |
-
from {
|
266 |
-
transform: rotate(0deg);
|
267 |
-
}
|
268 |
-
to {
|
269 |
-
transform: rotate(360deg);
|
270 |
-
}
|
271 |
-
}
|
272 |
-
#share-btn-container {
|
273 |
-
display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem;
|
274 |
-
}
|
275 |
-
#share-btn {
|
276 |
-
all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.25rem !important; padding-bottom: 0.25rem !important;
|
277 |
-
}
|
278 |
-
#share-btn * {
|
279 |
-
all: unset;
|
280 |
-
}
|
281 |
-
#share-btn-container div:nth-child(-n+2){
|
282 |
-
width: auto !important;
|
283 |
-
min-height: 0px !important;
|
284 |
-
}
|
285 |
-
#share-btn-container .wrap {
|
286 |
-
display: none !important;
|
287 |
-
}
|
288 |
-
'''
|
289 |
-
|
290 |
-
with gr.Blocks(css=css) as block:
|
291 |
-
with gr.Column(elem_id="col-container"):
|
292 |
-
gr.HTML(title)
|
293 |
-
|
294 |
-
input_image = gr.Image(type='pil', elem_id="input-img")
|
295 |
-
submit_btn = gr.Button("Submit")
|
296 |
-
output_text = gr.Textbox(label="Output", elem_id="output-txt")
|
297 |
-
|
298 |
-
with gr.Group(elem_id="share-btn-container"):
|
299 |
-
community_icon = gr.HTML(community_icon_html, visible=False)
|
300 |
-
loading_icon = gr.HTML(loading_icon_html, visible=False)
|
301 |
-
share_button = gr.Button("Share to community", elem_id="share-btn", visible=False)
|
302 |
-
|
303 |
-
examples=[['example01.jpg'], ['example02.jpg']]
|
304 |
-
ex = gr.Examples(examples=examples, fn=inference, inputs=input_image, outputs=[output_text, share_button, community_icon, loading_icon], cache_examples=True, run_on_click=True)
|
305 |
-
ex.dataset.headers = [""]
|
306 |
-
|
307 |
-
gr.HTML(article)
|
308 |
-
|
309 |
-
submit_btn.click(fn=inference, inputs=input_image, outputs=[output_text, share_button, community_icon, loading_icon])
|
310 |
-
share_button.click(None, [], [], _js=share_js)
|
311 |
-
|
312 |
-
block.queue(max_size=32).launch(show_api=False)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|