akhaliq HF staff commited on
Commit
2618f02
1 Parent(s): a703b1d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -3
app.py CHANGED
@@ -48,11 +48,14 @@ model_vq = model_vq.to(device)
48
 
49
 
50
 
51
- def inference(raw_image, model_n, question):
52
  if model_n == 'Image Captioning':
53
  image = transform(raw_image).unsqueeze(0).to(device)
54
  with torch.no_grad():
55
- caption = model.generate(image, sample=False, num_beams=3, max_length=20, min_length=5)
 
 
 
56
  return 'caption: '+caption[0]
57
 
58
  else:
@@ -61,7 +64,7 @@ def inference(raw_image, model_n, question):
61
  answer = model_vq(image_vq, question, train=False, inference='generate')
62
  return 'answer: '+answer[0]
63
 
64
- inputs = [gr.inputs.Image(type='pil'),gr.inputs.Radio(choices=['Image Captioning',"Visual Question Answering"], type="value", default="Image Captioning", label="Model"),"textbox"]
65
  outputs = gr.outputs.Textbox(label="Output")
66
 
67
  title = "BLIP"
 
48
 
49
 
50
 
51
+ def inference(raw_image, model_n, question, strategy):
52
  if model_n == 'Image Captioning':
53
  image = transform(raw_image).unsqueeze(0).to(device)
54
  with torch.no_grad():
55
+ if strategy == "beam search":
56
+ caption = model.generate(image, sample=False, num_beams=3, max_length=20, min_length=5)
57
+ else:
58
+ caption = model.generate(image, sample=True, top_p=0.9, max_length=20, min_length=5)
59
  return 'caption: '+caption[0]
60
 
61
  else:
 
64
  answer = model_vq(image_vq, question, train=False, inference='generate')
65
  return 'answer: '+answer[0]
66
 
67
+ inputs = [gr.inputs.Image(type='pil'),gr.inputs.Radio(choices=['Image Captioning',"Visual Question Answering"], type="value", default="Image Captioning", label="Model"),"textbox",gr.inputs.Radio(choices=['Beam search','Nucleus sampling'], type="value", default="Nucleus sampling", label="Strategy")]
68
  outputs = gr.outputs.Textbox(label="Output")
69
 
70
  title = "BLIP"