xinyu1205 commited on
Commit
6cfdc97
1 Parent(s): b0832a1

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -7
app.py CHANGED
@@ -1,4 +1,3 @@
1
- import ruamel_yaml as yaml
2
  import numpy as np
3
  import random
4
 
@@ -22,12 +21,7 @@ transform = transforms.Compose([transforms.Resize((image_size, image_size)),tran
22
  #######Swin Version
23
  pretrained = 'tag2text_swin_14m.pth'
24
 
25
- config_file = 'configs/tag2text_caption.yaml'
26
- config = yaml.load(open(config_file, 'r'), Loader=yaml.Loader)
27
-
28
- model = tag2text_caption(pretrained=pretrained, image_size=image_size, vit=config['vit'],
29
- vit_grad_ckpt=config['vit_grad_ckpt'], vit_ckpt_layer=config['vit_ckpt_layer'],
30
- prompt=config['prompt'],config=config,threshold = 0.75 )
31
 
32
  model.eval()
33
  model = model.to(device)
 
 
1
  import numpy as np
2
  import random
3
 
 
21
  #######Swin Version
22
  pretrained = 'tag2text_swin_14m.pth'
23
 
24
+ model = tag2text_caption(pretrained=pretrained, image_size=image_size, vit='swin_b' )
 
 
 
 
 
25
 
26
  model.eval()
27
  model = model.to(device)