ynhe commited on
Commit
8974591
·
1 Parent(s): 938205c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -12,7 +12,8 @@ from util import *
12
  import gradio as gr
13
  from chatbot import *
14
  from load_internvideo import *
15
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
16
  from simplet5 import SimpleT5
17
  from models.grit_model import DenseCaptioning
18
  bot = ConversationBot()
@@ -42,7 +43,7 @@ trans_action = transform_action()
42
  topil = T.ToPILImage()
43
  print("[INFO] initialize InternVideo model success!")
44
 
45
- dense_caption_model = DenseCaptioning(device)
46
  dense_caption_model.initialize_model()
47
  print("[INFO] initialize dense caption model success!")
48
 
 
12
  import gradio as gr
13
  from chatbot import *
14
  from load_internvideo import *
15
+ device = torch.device('cuda'
16
+ if torch.cuda.is_available() else 'cpu')
17
  from simplet5 import SimpleT5
18
  from models.grit_model import DenseCaptioning
19
  bot = ConversationBot()
 
43
  topil = T.ToPILImage()
44
  print("[INFO] initialize InternVideo model success!")
45
 
46
+ dense_caption_model = DenseCaptioning(str(device))
47
  dense_caption_model.initialize_model()
48
  print("[INFO] initialize dense caption model success!")
49