Quick fix to zmq ip address (#795)

quick fix to zmq host / ip address

nerfstudio/configs/base_config.py

@@ -170,7 +170,7 @@ class ViewerConfig(PrintableConfig):
     """whether or not to launch the bridge server"""
     websocket_port: Optional[int] = 7007
     """the default websocket port to connect to"""
-    ip_address: str = "0.0.0.0"
+    ip_address: str = "127.0.0.1"
     """the ip address where the bridge server is running"""
     num_rays_per_chunk: int = 32768
     """number of rays per chunk to render with viewer"""

nerfstudio/viewer/server/server.py

@@ -100,14 +100,14 @@ class ZMQWebSocketBridge:
 
     context = zmq.Context()
 
-    def __init__(self, zmq_port: int, websocket_port: int):
+    def __init__(self, zmq_port: int, websocket_port: int, ip_address: str):
         self.zmq_port = zmq_port
         self.websocket_pool = set()
         self.app = self.make_app()
         self.ioloop = tornado.ioloop.IOLoop.current()
 
         # zmq
-        zmq_url = f"tcp://0.0.0.0:{self.zmq_port:d}"
+        zmq_url = f"tcp://{ip_address}:{self.zmq_port:d}"
         self.zmq_socket, self.zmq_stream, self.zmq_url = self.setup_zmq(zmq_url)
 
         # websocket
@@ -198,13 +198,16 @@ class ZMQWebSocketBridge:
         self.ioloop.start()
 
 
-def run_viewer_bridge_server(zmq_port: int = 6000, websocket_port: int = 7007, use_ngrok: bool = False):
+def run_viewer_bridge_server(
+    zmq_port: int = 6000, websocket_port: int = 7007, ip_address: str = "127.0.0.1", use_ngrok: bool = False
+):
     """Run the viewer bridge server.
 
     Args:
         zmq_port: port to use for zmq
         websocket_port: port to use for websocket
-        host: host to connect to
+        ip_address: host to connect to
+        use_ngrok: whether to use ngrok to expose the zmq port
     """
 
     # whether to launch pyngrok or not
@@ -214,7 +217,7 @@ def run_viewer_bridge_server(zmq_port: int = 6000, websocket_port: int = 7007, u
         http_tunnel = ngrok.connect(addr=str(zmq_port), proto="tcp")
         print(http_tunnel)
 
-    bridge = ZMQWebSocketBridge(zmq_port=zmq_port, websocket_port=websocket_port)
+    bridge = ZMQWebSocketBridge(zmq_port=zmq_port, websocket_port=websocket_port, ip_address=ip_address)
     print(bridge)
     try:
         bridge.run()

nerfstudio/viewer/server/subprocess.py

@@ -34,13 +34,17 @@ CONSOLE = Console()
 
 
 def run_viewer_bridge_server_as_subprocess(
-    websocket_port: int, zmq_port: Optional[int] = None, log_filename: Union[str, None] = None
+    websocket_port: int,
+    zmq_port: Optional[int] = None,
+    ip_address: str = "127.0.0.1",
+    log_filename: Union[str, None] = None,
 ):
     """Runs the viewer bridge server as a subprocess.
 
     Args:
         zmq_port: Port to use for the ZMQ server.
         websocket_port: Port to use for the websocket server.
+        ip_address: host to connect to
         log_filename: Filename to use for the log file. If None, no log file is created.
 
     Returns:
@@ -60,6 +64,8 @@ def run_viewer_bridge_server_as_subprocess(
     args.append(str(zmq_port))
     args.append("--websocket-port")
     args.append(str(websocket_port))
+    args.append("--ip-address")
+    args.append(str(ip_address))
     # supress output if no log filename is specified
     logfile = open(  # pylint: disable=consider-using-with
         log_filename if log_filename else os.devnull, "w", encoding="utf8"

nerfstudio/viewer/server/viewer_utils.py

@@ -223,7 +223,10 @@ class ViewerState:
             assert self.config.websocket_port is not None
             self.log_filename.parent.mkdir(exist_ok=True)
             zmq_port = run_viewer_bridge_server_as_subprocess(
-                self.config.websocket_port, zmq_port=self.config.zmq_port, log_filename=str(self.log_filename)
+                self.config.websocket_port,
+                zmq_port=self.config.zmq_port,
+                ip_address=self.config.ip_address,
+                log_filename=str(self.log_filename),
             )
             # TODO(ethan): log the output of the viewer bridge server in a file where the training logs go
             CONSOLE.line()
@@ -235,8 +238,8 @@ class ViewerState:
             CONSOLE.print(f"[Public] Open the viewer at {self.viewer_url}")
             CONSOLE.rule(characters="=")
             CONSOLE.line()
-            self.vis = Viewer(zmq_port=zmq_port)
-            self.vis_webrtc_thread = Viewer(zmq_port=zmq_port)
+            self.vis = Viewer(zmq_port=zmq_port, ip_address=self.config.ip_address)
+            self.vis_webrtc_thread = Viewer(zmq_port=zmq_port, ip_address=self.config.ip_address)
         else:
             assert self.config.zmq_port is not None
             self.vis = Viewer(zmq_port=self.config.zmq_port, ip_address=self.config.ip_address)

nerfstudio/viewer/server/visualizer.py

@@ -38,7 +38,7 @@ class ViewerWindow:
 
     context = zmq.Context()
 
-    def __init__(self, zmq_port, ip_address="0.0.0.0"):
+    def __init__(self, zmq_port, ip_address="127.0.0.1"):
         self.zmq_port = zmq_port
         self.client = self.context.socket(zmq.REQ)
         zmq_url = f"tcp://{ip_address}:{self.zmq_port}"
@@ -70,7 +70,7 @@ class ViewerWindow:
         )
         return umsgpack.unpackb(self.client.recv())
 
-    def timeout_ping(self, timeout_in_sec: int = 5):
+    def timeout_ping(self, timeout_in_sec: int = 15):
         """Timeout if ping fails to complete in timeout_in_secs seconds"""
 
         res = [Exception(f"Couldn't connect to the viewer Bridge Server in {timeout_in_sec} seconds. Exiting.")]
@@ -91,7 +91,7 @@ class ViewerWindow:
             raise ret
         return ret
 
-    def assert_connected(self, timeout_in_sec: int = 5):
+    def assert_connected(self, timeout_in_sec: int = 15):
         """Check if the connection was established properly within some time.
 
         Args:
@@ -117,7 +117,7 @@ class Viewer:
     """
 
     def __init__(
-        self, zmq_port: Optional[int] = None, window: Optional[ViewerWindow] = None, ip_address: str = "0.0.0.0"
+        self, zmq_port: Optional[int] = None, window: Optional[ViewerWindow] = None, ip_address: str = "127.0.0.1"
     ):
         if zmq_port is None and window is None:
             raise ValueError("Must specify either zmq_port or window.")