pharmapsychotic commited on
Commit
95d3029
1 Parent(s): ec4e754

Support both ViT-L and ViT-H!

Browse files

- use clip-interrogator as pip package
- use huggingface_hub to download preprocessed files

Files changed (3) hide show
  1. .gitignore +2 -0
  2. app.py +88 -228
  3. requirements.txt +6 -3
.gitignore ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ cache/
2
+ venv/
app.py CHANGED
@@ -1,221 +1,71 @@
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="
@@ -234,7 +84,8 @@ title = """
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>
@@ -255,19 +106,15 @@ article = """
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;
@@ -287,11 +134,13 @@ a {text-decoration-line: underline; font-weight: 600;}
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
 
@@ -300,13 +149,24 @@ with gr.Blocks(css=css) as block:
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)
 
1
+ #!/usr/bin/env python3
 
 
 
 
2
  import gradio as gr
 
 
 
3
  import os
4
+ from clip_interrogator import Config, Interrogator
5
+ from huggingface_hub import hf_hub_download
 
 
 
 
 
 
 
 
 
6
  from share_btn import community_icon_html, loading_icon_html, share_js
7
 
8
+ MODELS = ['ViT-L (best for Stable Diffusion 1.*)', 'ViT-H (best for Stable Diffusion 2.*)']
9
+
10
+ # download preprocessed files
11
+ PREPROCESS_FILES = [
12
+ 'ViT-H-14_laion2b_s32b_b79k_artists.pkl',
13
+ 'ViT-H-14_laion2b_s32b_b79k_flavors.pkl',
14
+ 'ViT-H-14_laion2b_s32b_b79k_mediums.pkl',
15
+ 'ViT-H-14_laion2b_s32b_b79k_movements.pkl',
16
+ 'ViT-H-14_laion2b_s32b_b79k_trendings.pkl',
17
+ 'ViT-L-14_openai_artists.pkl',
18
+ 'ViT-L-14_openai_flavors.pkl',
19
+ 'ViT-L-14_openai_mediums.pkl',
20
+ 'ViT-L-14_openai_movements.pkl',
21
+ 'ViT-L-14_openai_trendings.pkl',
22
+ ]
23
+ print("Download preprocessed cache files...")
24
+ for file in PREPROCESS_FILES:
25
+ path = hf_hub_download(repo_id="pharma/ci-preprocess", filename=file, cache_dir="cache")
26
+ cache_path = os.path.dirname(path)
27
+
28
+
29
+ # load BLIP and ViT-L
30
+ config = Config(cache_path=cache_path, clip_model_path="cache", clip_model_name="ViT-L-14/openai")
31
+ ci_vitl = Interrogator(config)
32
+ ci_vitl.clip_model = ci_vitl.clip_model.to("cpu")
33
+
34
+ # load ViT-H
35
+ config.blip_model = ci_vitl.blip_model
36
+ config.clip_model_name = "ViT-H-14/laion2b_s32b_b79k"
37
+ ci_vith = Interrogator(config)
38
+ ci_vith.clip_model = ci_vith.clip_model.to("cpu")
39
+
40
+
41
+ def inference(image, clip_model_name, mode):
42
+
43
+ # move selected model to GPU and other model to CPU
44
+ if clip_model_name == MODELS[0]:
45
+ ci_vith.clip_model = ci_vith.clip_model.to("cpu")
46
+ ci_vitl.clip_model = ci_vitl.clip_model.to(ci_vitl.device)
47
+ ci = ci_vitl
48
+ else:
49
+ ci_vitl.clip_model = ci_vitl.clip_model.to("cpu")
50
+ ci_vith.clip_model = ci_vith.clip_model.to(ci_vith.device)
51
+ ci = ci_vith
52
+
53
+ ci.config.blip_num_beams = 64
54
+ ci.config.chunk_size = 2048
55
+ ci.config.flavor_intermediate_count = 2048 if clip_model_name == MODELS[0] else 1024
56
+
57
+ image = image.convert('RGB')
58
+ if mode == 'best':
59
+ prompt = ci.interrogate(image)
60
+ elif mode == 'classic':
61
+ prompt = ci.interrogate_classic(image)
62
+ else:
63
+ prompt = ci.interrogate_fast(image)
64
+
65
+ return prompt, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
66
+
67
+
68
+ TITLE = """
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
69
  <div style="text-align: center; max-width: 650px; margin: 0 auto;">
70
  <div
71
  style="
 
84
  </p>
85
  </div>
86
  """
87
+
88
+ ARTICLE = """
89
  <div style="text-align: center; max-width: 650px; margin: 0 auto;">
90
  <p>
91
  Example art by <a href="https://pixabay.com/illustrations/watercolour-painting-art-effect-4799014/">Layers</a>
 
106
  </div>
107
  """
108
 
109
+ CSS = '''
110
  #col-container {max-width: 700px; margin-left: auto; margin-right: auto;}
111
  a {text-decoration-line: underline; font-weight: 600;}
112
  .animate-spin {
113
  animation: spin 1s linear infinite;
114
  }
115
  @keyframes spin {
116
+ from { transform: rotate(0deg); }
117
+ to { transform: rotate(360deg); }
 
 
 
 
118
  }
119
  #share-btn-container {
120
  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;
 
134
  }
135
  '''
136
 
137
+ with gr.Blocks(css=CSS) as block:
138
  with gr.Column(elem_id="col-container"):
139
+ gr.HTML(TITLE)
140
 
141
  input_image = gr.Image(type='pil', elem_id="input-img")
142
+ input_model = gr.Dropdown(MODELS, value=MODELS[0], label='CLIP Model')
143
+ input_mode = gr.Radio(['best', 'fast'], value='best', label='Mode')
144
  submit_btn = gr.Button("Submit")
145
  output_text = gr.Textbox(label="Output", elem_id="output-txt")
146
 
 
149
  loading_icon = gr.HTML(loading_icon_html, visible=False)
150
  share_button = gr.Button("Share to community", elem_id="share-btn", visible=False)
151
 
152
+ examples=[['example01.jpg', MODELS[0], 'best'], ['example02.jpg', MODELS[0], 'best']]
153
+ ex = gr.Examples(
154
+ examples=examples,
155
+ fn=inference,
156
+ inputs=[input_image, input_model, input_mode],
157
+ outputs=[output_text, share_button, community_icon, loading_icon],
158
+ cache_examples=True,
159
+ run_on_click=True
160
+ )
161
  ex.dataset.headers = [""]
162
 
163
+ gr.HTML(ARTICLE)
164
 
165
+ submit_btn.click(
166
+ fn=inference,
167
+ inputs=[input_image, input_model, input_mode],
168
+ outputs=[output_text, share_button, community_icon, loading_icon]
169
+ )
170
  share_button.click(None, [], [], _js=share_js)
171
 
172
  block.queue(max_size=32).launch(show_api=False)
requirements.txt CHANGED
@@ -1,11 +1,14 @@
1
- --extra-index-url https://download.pytorch.org/whl/cu113
2
  torch
3
  torchvision
4
 
5
  fairscale
6
  ftfy
 
 
7
  Pillow
8
  timm
9
  transformers==4.15.0
10
- -e git+https://github.com/openai/CLIP.git@main#egg=clip
11
- -e git+https://github.com/pharmapsychotic/BLIP.git@main#egg=blip
 
 
1
+ --extra-index-url https://download.pytorch.org/whl/cu117
2
  torch
3
  torchvision
4
 
5
  fairscale
6
  ftfy
7
+ gradio
8
+ huggingface-hub
9
  Pillow
10
  timm
11
  transformers==4.15.0
12
+ open_clip_torch
13
+ clip-interrogator==0.3.1
14
+ -e git+https://github.com/pharmapsychotic/BLIP.git@lib#egg=blip