Ashish Soni commited on
Commit
68d8fd5
·
1 Parent(s): d0f1e1c
Files changed (1) hide show
  1. generate.py +25 -0
generate.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch, time
2
+ import clip
3
+ from PIL import Image
4
+ import lightning as L
5
+
6
+ device = "cuda" if torch.cuda.is_available() else "cpu"
7
+ model, preprocess = clip.load("ViT-B/32", device=device)
8
+
9
+ image = preprocess(Image.open("CLIP.png")).unsqueeze(0).to(device)
10
+ text = clip.tokenize(["a diagram", "a dog", "a cat"]).to(device)
11
+
12
+ start_time = time.time()
13
+
14
+ with torch.no_grad():
15
+ image_features = model.encode_image(image)
16
+ text_features = model.encode_text(text)
17
+
18
+ logits_per_image, logits_per_text = model(image, text)
19
+ probs = logits_per_image.softmax(dim=-1).cpu().numpy()
20
+
21
+ end_time = time.time()
22
+
23
+ print("Label probs:", probs) # prints: [[0.9927937 0.00421068 0.00299572]]
24
+
25
+ print(f"Prediction time: {end_time - start_time} seconds")