Update config.py

config.py

@@ -1,7 +1,9 @@
+import argparse
 import sys
 import torch
 from multiprocessing import cpu_count
 
+
 class Config:
     def __init__(self):
         self.device = "cuda:0"
@@ -9,8 +11,42 @@ class Config:
         self.n_cpu = 0
         self.gpu_name = None
         self.gpu_mem = None
+        (
+            self.python_cmd,
+            self.listen_port,
+            self.iscolab,
+            self.noparallel,
+            self.noautoopen,
+        ) = self.arg_parse()
         self.x_pad, self.x_query, self.x_center, self.x_max = self.device_config()
 
+    @staticmethod
+    def arg_parse() -> tuple:
+        exe = sys.executable or "python"
+        parser = argparse.ArgumentParser()
+        parser.add_argument("--port", type=int, default=7865, help="Listen port")
+        parser.add_argument("--pycmd", type=str, default=exe, help="Python command")
+        parser.add_argument("--colab", action="store_true", help="Launch in colab")
+        parser.add_argument(
+            "--noparallel", action="store_true", help="Disable parallel processing"
+        )
+        parser.add_argument(
+            "--noautoopen",
+            action="store_true",
+            help="Do not open in browser automatically",
+        )
+        cmd_opts = parser.parse_args()
+
+        cmd_opts.port = cmd_opts.port if 0 <= cmd_opts.port <= 65535 else 7865
+
+        return (
+            cmd_opts.pycmd,
+            cmd_opts.port,
+            cmd_opts.colab,
+            cmd_opts.noparallel,
+            cmd_opts.noautoopen,
+        )
+
     # has_mps is only available in nightly pytorch (for now) and MasOS 12.3+.
     # check `getattr` and try it for compatibility
     @staticmethod
@@ -70,7 +106,7 @@ class Config:
             x_center = 38
             x_max = 41
 
-        if self.gpu_mem is not None and self.gpu_mem <= 4:
+        if self.gpu_mem != None and self.gpu_mem <= 4:
             x_pad = 1
             x_query = 5
             x_center = 30