awacke1 commited on
Commit
973f818
1 Parent(s): 0e4281e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -13
app.py CHANGED
@@ -2,22 +2,15 @@ import os
2
  import gradio as gr
3
  import torch
4
  import PIL
5
- import transformers
6
- transformers.utils.move_cache()
7
 
8
  from flamingo_mini import FlamingoConfig, FlamingoModel, FlamingoProcessor
9
 
10
-
11
-
12
  EXAMPLES_DIR = 'examples'
13
  DEFAULT_PROMPT = "<image>"
14
-
15
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
16
-
17
  model = FlamingoModel.from_pretrained('dhansmair/flamingo-mini')
18
  model.to(device)
19
  model.eval()
20
-
21
  processor = FlamingoProcessor(model.config, load_vision_processor=True)
22
 
23
  # setup some example images
@@ -27,21 +20,16 @@ if os.path.isdir(EXAMPLES_DIR):
27
  path = EXAMPLES_DIR + "/" + file
28
  examples.append([path, DEFAULT_PROMPT])
29
 
30
-
31
  def predict_caption(image, prompt):
32
  assert isinstance(prompt, str)
33
-
34
  features = processor.extract_features(image).to(device)
35
  caption = model.generate_captions(processor,
36
  visual_features=features,
37
  prompt=prompt)
38
-
39
  if isinstance(caption, list):
40
  caption = caption[0]
41
-
42
  return caption
43
-
44
-
45
  iface = gr.Interface(fn=predict_caption,
46
  inputs=[gr.Image(type="pil"), gr.Textbox(value=DEFAULT_PROMPT, label="Prompt")],
47
  examples=examples,
 
2
  import gradio as gr
3
  import torch
4
  import PIL
 
 
5
 
6
  from flamingo_mini import FlamingoConfig, FlamingoModel, FlamingoProcessor
7
 
 
 
8
  EXAMPLES_DIR = 'examples'
9
  DEFAULT_PROMPT = "<image>"
 
10
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
11
  model = FlamingoModel.from_pretrained('dhansmair/flamingo-mini')
12
  model.to(device)
13
  model.eval()
 
14
  processor = FlamingoProcessor(model.config, load_vision_processor=True)
15
 
16
  # setup some example images
 
20
  path = EXAMPLES_DIR + "/" + file
21
  examples.append([path, DEFAULT_PROMPT])
22
 
 
23
  def predict_caption(image, prompt):
24
  assert isinstance(prompt, str)
 
25
  features = processor.extract_features(image).to(device)
26
  caption = model.generate_captions(processor,
27
  visual_features=features,
28
  prompt=prompt)
 
29
  if isinstance(caption, list):
30
  caption = caption[0]
 
31
  return caption
32
+
 
33
  iface = gr.Interface(fn=predict_caption,
34
  inputs=[gr.Image(type="pil"), gr.Textbox(value=DEFAULT_PROMPT, label="Prompt")],
35
  examples=examples,