diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 1ca2ac4ae..76b8bc79d 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -169,7 +169,7 @@ class ViewerConfig(PrintableConfig): start_train: bool = True """whether to immediately start training upon loading viewer if False, will just visualize dataset but you can toggle training in viewer""" - zmq_url: str = "tcp://127.0.0.1:6000" + zmq_port: int = 6000 """the zmq port to connect to for communication""" launch_bridge_server: bool = True """whether or not to launch the zmq bridge server""" diff --git a/nerfactory/viewer/server/server.py b/nerfactory/viewer/server/server.py index 229121278..4170f34ff 100644 --- a/nerfactory/viewer/server/server.py +++ b/nerfactory/viewer/server/server.py @@ -14,11 +14,10 @@ """Server bridge to faciliate interactions between python backend and javascript front end""" - -import argparse import sys from typing import Callable, List, Optional, Tuple +import dcargs import msgpack import msgpack_numpy import tornado.gen @@ -177,6 +176,7 @@ class ZMQWebSocketBridge: websocket_address: str = "0.0.0.0", websocket_port: Optional[int] = None, ): + self.zmq_port = zmq_port self.host = host self.websocket_address = websocket_address self.websocket_pool = set() @@ -185,14 +185,15 @@ class ZMQWebSocketBridge: self.pcs = set() self.video_tracks = set() - if zmq_port is None: + if self.zmq_port is None: def f(port): return self.setup_zmq(f"{DEFAULT_ZMQ_METHOD:s}://{self.host:s}:{port:d}") (self.zmq_socket, self.zmq_stream, self.zmq_url), _ = find_available_port(f, DEFAULT_ZMQ_PORT) + self.zmq_port = int(self.zmq_url.split(":")[-1]) else: - zmq_url = f"{DEFAULT_ZMQ_METHOD:s}://{self.host:s}:{zmq_port:d}" + zmq_url = f"{DEFAULT_ZMQ_METHOD:s}://{self.host:s}:{self.zmq_port:d}" self.zmq_socket, self.zmq_stream, self.zmq_url = self.setup_zmq(zmq_url) listen_kwargs = {"address": self.websocket_address} @@ -208,7 +209,7 @@ class ZMQWebSocketBridge: def __str__(self) -> str: class_name = self.__class__.__name__ - return f'{class_name} using zmq_url="{self.zmq_url}" and websocket_url="{self.websocket_url}"' + return f'{class_name} using zmq_port="{self.zmq_port}" and websocket_port="{self.websocket_port}"' def make_app(self): """Create a tornado application for the websocket server.""" @@ -292,13 +293,14 @@ class ZMQWebSocketBridge: self.ioloop.start() -def run_viewer_bridge_server(): - """Run the viewer bridge server""" - parser = argparse.ArgumentParser(description="Listen for ZeroMQ commands") - parser.add_argument("--zmq-port", "-z", type=int, nargs="?", default=6000) - parser.add_argument("--websocket-port", "-wp", type=int, nargs="?", default=7007) - args = parser.parse_args() - bridge = ZMQWebSocketBridge(zmq_port=args.zmq_port, websocket_port=args.websocket_port) +def run_viewer_bridge_server(zmq_port: int = 6000, websocket_port: int = 7007): + """Run the viewer bridge server. + + Args: + zmq_port: port to use for zmq + websocket_port: port to use for websocket + """ + bridge = ZMQWebSocketBridge(zmq_port=zmq_port, websocket_port=websocket_port) print(bridge) try: bridge.run() @@ -306,5 +308,11 @@ def run_viewer_bridge_server(): pass +def main(): + """The main entrypoint.""" + dcargs.extras.set_accent_color("bright_yellow") + dcargs.cli(run_viewer_bridge_server) + + if __name__ == "__main__": - run_viewer_bridge_server() + main() diff --git a/nerfactory/viewer/server/subprocess.py b/nerfactory/viewer/server/subprocess.py index 276db33d9..c36323e01 100644 --- a/nerfactory/viewer/server/subprocess.py +++ b/nerfactory/viewer/server/subprocess.py @@ -69,10 +69,10 @@ def run_viewer_bridge_server_as_subprocess(zmq_port: int, websocket_port: int, l time.sleep(0.5) string = f"\nThe viewer bridge server subprocess failed. Please check the log file {log_filename}.\n" string += ( - "You likely have to modify --viewer.zmq-url and/or --viewer.websocket-port in the " + "You likely have to modify --viewer.zmq-port and/or --viewer.websocket-port in the " "config to avoid conflicting ports.\n" ) - string += "Try modifying --viewer.zmq-url tcp://127.0.0.1:6000 --viewer.websocket-port 7007\n" + string += "Try modifying --viewer.zmq-port 6000 --viewer.websocket-port 7007\n" CONSOLE.print(f"[bold red]{string}") cleanup(process) # This exists the entire program. sys.exit() will only kill the thread that this runs in. diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index 62563c924..03ebe8099 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/nerfactory/viewer/server/viewer_utils.py @@ -209,10 +209,11 @@ class ViewerState: self.log_filename = log_filename if self.config.launch_bridge_server: # start the viewer bridge server - zmq_port = int(self.config.zmq_url.split(":")[-1]) websocket_port = self.config.websocket_port self.log_filename.parent.mkdir(exist_ok=True) - run_viewer_bridge_server_as_subprocess(zmq_port, websocket_port, log_filename=str(self.log_filename)) + run_viewer_bridge_server_as_subprocess( + self.config.zmq_port, websocket_port, log_filename=str(self.log_filename) + ) # TODO(ethan): move this into the writer such that it's at the bottom # of the logging stack and easy to see and click # TODO(ethan): log the output of the viewer bridge server in a file where the training logs go @@ -228,7 +229,7 @@ class ViewerState: console.print(dev_open_viewer_instructions_string) console.rule(characters="=") console.line() - self.vis = Viewer(zmq_url=self.config.zmq_url) + self.vis = Viewer(zmq_port=self.config.zmq_port) # viewer specific variables self.prev_camera_matrix = None diff --git a/nerfactory/viewer/server/visualizer.py b/nerfactory/viewer/server/visualizer.py index 291a51886..b268bb26f 100644 --- a/nerfactory/viewer/server/visualizer.py +++ b/nerfactory/viewer/server/visualizer.py @@ -33,10 +33,11 @@ class ViewerWindow: context = zmq.Context() - def __init__(self, zmq_url): - self.zmq_url = zmq_url + def __init__(self, zmq_port): + self.zmq_port = zmq_port self.client = self.context.socket(zmq.REQ) - self.client.connect(self.zmq_url) + zmq_url = f"tcp://127.0.0.1:{self.zmq_port}" + self.client.connect(zmq_url) self.assert_connected() def send(self, command): @@ -90,15 +91,15 @@ class Viewer: """Viewer class for connecting to the bridge server. Args: - zmq_url: Where to connect with ZMQ. + zmq_port: Where to connect with ZMQ. window: An already existing ViewerWindow. """ - def __init__(self, zmq_url: Optional[str] = None, window: Optional[ViewerWindow] = None): - if zmq_url is None and window is None: - raise ValueError("Must specify either zmq_url or window.") + def __init__(self, zmq_port: Optional[int] = None, window: Optional[ViewerWindow] = None): + if zmq_port is None and window is None: + raise ValueError("Must specify either zmq_port or window.") if window is None: - self.window = ViewerWindow(zmq_url=zmq_url) + self.window = ViewerWindow(zmq_port=zmq_port) else: self.window = window self.path = Path(()) diff --git a/pyproject.toml b/pyproject.toml index cdfca6213..ffa4daff8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -74,7 +74,7 @@ docs = [ ] [project.scripts] -nerfactory-viewer-bridge-server = "nerfactory.viewer.server.server:run_viewer_bridge_server" +nf-scripts-viewer-bridge-server = "nerfactory.viewer.server.server:main" [options] # equivalent to using --extra-index-url with pip, which is needed for specifying the CUDA version torch and torchvision diff --git a/scripts/run_viewer_bridge_server.py b/scripts/run_viewer_bridge_server.py deleted file mode 100755 index 7ee9a3089..000000000 --- a/scripts/run_viewer_bridge_server.py +++ /dev/null @@ -1,7 +0,0 @@ -#!/usr/bin/env python -"""View Bridge Server""" - -from nerfactory.viewer.server.server import run_viewer_bridge_server - -if __name__ == "__main__": - run_viewer_bridge_server()