Steven Anderson commited on
Commit
eaa9560
1 Parent(s): e94e516
Files changed (1) hide show
  1. app.py +16 -17
app.py CHANGED
@@ -4,7 +4,7 @@ import clip
4
  from PIL import Image
5
 
6
 
7
- def process(device, image, prompt):
8
  print("Inferring...")
9
  image = preprocess(image).unsqueeze(0).to(device)
10
  print(image)
@@ -21,19 +21,18 @@ def process(device, image, prompt):
21
  return dict(zip(prompts, probs[0]))
22
 
23
 
24
- if __name__ == "__main__":
25
- print("Getting device...")
26
- device = "cuda" if torch.cuda.is_available() else "cpu"
27
- print("Loading model...")
28
- model, preprocess = clip.load("ViT-B/32", device=device)
29
- print("Loaded model.")
30
-
31
- iface = gr.Interface(
32
- fn=lambda i, p: process(device, i, p),
33
- inputs=[
34
- gr.Image(),
35
- gr.Textbox(lines=5, label="Prompts (newline-separated)"),
36
- ],
37
- outputs="label",
38
- )
39
- iface.launch()
 
4
  from PIL import Image
5
 
6
 
7
+ def process(model, preprocess, device, image, prompt):
8
  print("Inferring...")
9
  image = preprocess(image).unsqueeze(0).to(device)
10
  print(image)
 
21
  return dict(zip(prompts, probs[0]))
22
 
23
 
24
+ print("Getting device...")
25
+ device = "cuda" if torch.cuda.is_available() else "cpu"
26
+ print("Loading model...")
27
+ model, preprocess = clip.load("ViT-B/32", device=device)
28
+ print("Loaded model.")
29
+
30
+ iface = gr.Interface(
31
+ fn=lambda i, p: process(model, preprocess, device, i, p),
32
+ inputs=[
33
+ gr.Image(),
34
+ gr.Textbox(lines=5, label="Prompts (newline-separated)"),
35
+ ],
36
+ outputs="label",
37
+ )
38
+ iface.launch()