cuda support

app.py

@@ -4,7 +4,7 @@ import argparse
 import glob
 import multiprocessing as mp
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = ""
+#os.environ["CUDA_VISIBLE_DEVICES"] = ""
 os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
 
 # fmt: off
@@ -28,6 +28,7 @@ from detectron2.utils.logger import setup_logger
 from cat_seg import add_cat_seg_config
 from demo.predictor import VisualizationDemo
 import gradio as gr
+import torch
 from matplotlib.backends.backend_agg import FigureCanvasAgg as fc
 
 # constants
@@ -41,6 +42,8 @@ def setup_cfg(args):
     add_cat_seg_config(cfg)
     cfg.merge_from_file(args.config_file)
     cfg.merge_from_list(args.opts)
+    if torch.cuda.is_available():
+        cfg.MODEL.DEVICE = "cuda"
     cfg.freeze()
     return cfg
 
@@ -62,14 +65,14 @@ def get_parser():
     parser.add_argument(
         "--opts",
         help="Modify config options using the command-line 'KEY VALUE' pairs",
-        default=["MODEL.WEIGHTS", "model_final.pth", 
+        default=(
+        ["MODEL.WEIGHTS", "model_final.pth", 
         "MODEL.SEM_SEG_HEAD.TRAIN_CLASS_JSON", "datasets/voc20.json",
         "MODEL.SEM_SEG_HEAD.TEST_CLASS_JSON", "datasets/voc20.json",
         "TEST.SLIDING_WINDOW", "True",
         "MODEL.SEM_SEG_HEAD.POOLING_SIZES", "[1,1]",
-        "MODEL.DEVICE", "cpu",
-        "MODEL.PROMPT_ENSEMBLE_TYPE", "single"
-        ],
+        "MODEL.PROMPT_ENSEMBLE_TYPE", "single",
+        "MODEL.DEVICE", "cpu"]),
         nargs=argparse.REMAINDER,
     )
     return parser
@@ -103,7 +106,7 @@ if __name__ == "__main__":
         description="""## CAT-Seg Demo
 Welcome to the CAT-Seg Demo! Here, we present the CAT-Seg with ViT-L model for open-vocabulary semantic segmentation.
 
-Please note that this is an optimized version of the full model, and as such, its performance may be limited compared to the full model.
+Please note that this is an optimized version of the full model, and as such, its performance may be limited compared to the full model. Also, the demo runs on a CPU, so it may take a little time to process your image.
 
 To get started, simply upload an image and a comma-separated list of categories, and let the model work its magic!""")
     iface.launch()

requirements.txt

@@ -7,7 +7,7 @@ imageio==2.4.1
 timm==0.8.3.dev0
 regex
 einops
-torch==1.13.0+cpu
-torchvision==0.14.0+cpu
-torchaudio==0.13.0
---extra-index-url https://download.pytorch.org/whl/cpu
+torch==1.13.1+cu116
+torchvision==0.14.1+cu116
+torchaudio==0.13.1
+--extra-index-url https://download.pytorch.org/whl/cu116