Peijie commited on
Commit
9b455a7
·
1 Parent(s): 11e3b9b

specify torch/huggingface version.

Browse files
Files changed (2) hide show
  1. app.py +11 -4
  2. requirements.txt +3 -3
app.py CHANGED
@@ -1,9 +1,12 @@
1
  import os
2
  import io
3
  #DEBUG
4
- os.environ["TORCH_USE_CUDA_DSA"] = "1"
5
 
6
  import torch
 
 
 
7
  import json
8
  import base64
9
  import gradio as gr
@@ -15,13 +18,17 @@ from plots import get_pre_define_colors
15
  from utils.load_model import load_xclip
16
  from utils.predict import xclip_pred
17
 
18
-
19
  #! Huggingface does not allow load model to main process, so we need to load the model when needed, it may not help in improve the speed of the app.
20
  DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
21
- print(f"Not at Huggingface demo, load model to main process.")
 
 
 
 
 
22
  XCLIP, OWLVIT_PRECESSOR = load_xclip(DEVICE)
23
 
24
- print(f"Device: {DEVICE}")
25
 
26
  XCLIP_DESC_PATH = "data/jsons/bs_cub_desc.json"
27
  XCLIP_DESC = json.load(open(XCLIP_DESC_PATH, "r"))
 
1
  import os
2
  import io
3
  #DEBUG
4
+ os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
5
 
6
  import torch
7
+ import torchvision
8
+ import transformers
9
+ import logging
10
  import json
11
  import base64
12
  import gradio as gr
 
18
  from utils.load_model import load_xclip
19
  from utils.predict import xclip_pred
20
 
21
+ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
22
  #! Huggingface does not allow load model to main process, so we need to load the model when needed, it may not help in improve the speed of the app.
23
  DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
24
+ logging.info(f"Using device: {DEVICE}")
25
+ # get the torch, torchvision, and transformers version
26
+ logging.info(f"torch version: {torch.__version__}")
27
+ logging.info(f"torchvision version: {torchvision.__version__}")
28
+ logging.info(f"transformers version: {transformers.__version__}")
29
+
30
  XCLIP, OWLVIT_PRECESSOR = load_xclip(DEVICE)
31
 
 
32
 
33
  XCLIP_DESC_PATH = "data/jsons/bs_cub_desc.json"
34
  XCLIP_DESC = json.load(open(XCLIP_DESC_PATH, "r"))
requirements.txt CHANGED
@@ -1,9 +1,9 @@
1
- # torch
2
- # torchvision
3
  # gradio
4
  numpy
5
  Pillow
6
- transformers
7
  ftfy
8
  regex
9
  pandas
 
1
+ torch>=2.6.0
2
+ torchvision>=0.21.0
3
  # gradio
4
  numpy
5
  Pillow
6
+ transformers==4.47.1
7
  ftfy
8
  regex
9
  pandas