refactor(server, pipeline): fix workspace path mapping, add --config flag option

.vscode/launch.json

@@ -41,7 +41,7 @@
             "name": "Run ComfyStream BYOC",
             "type": "debugpy",
             "request": "launch",
-            "cwd": "/workspace/ComfyUI",
+            "cwd": "/workspace/comfystream",
             "program": "/workspace/comfystream/server/byoc.py",
             "console": "integratedTerminal",
             "args": [

README.md

@@ -160,6 +160,8 @@ If you only have a subset of those UDP ports available, you can use the `--media
 python server/app.py --workspace <COMFY_WORKSPACE> --media-ports 1024,1025,...
 ```
 
+> Tip: Use `--workspace` (preferred). `--cwd` remains a compatible alias and honors `COMFYUI_CWD`.
+
 If you are running the server in a restrictive network environment where this is not possible, you will need to use a TURN server.
 
 At the moment, the server supports using Twilio's TURN servers (although it is easy to make the update to support arbitrary TURN servers):

docker/entrypoint.sh

@@ -29,6 +29,7 @@ WORKSPACE_STORAGE="/app/storage"
 COMFYUI_DIR="/workspace/ComfyUI"
 MODELS_DIR="$COMFYUI_DIR/models"
 OUTPUT_DIR="$COMFYUI_DIR/output"
+export COMFYUI_CWD="$COMFYUI_DIR"
 
 # Initialize variables to track which services to start
 START_COMFYUI=false

install.py

@@ -74,8 +74,10 @@ def download_and_extract_ui_files(version: str):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Install custom node requirements")
     parser.add_argument(
+        "--cwd",
         "--workspace",
-        default=os.environ.get("COMFY_UI_WORKSPACE", None),
+        dest="workspace",
+        default=os.environ.get("COMFYUI_CWD"),
         required=False,
         help="Set Comfy workspace",
     )
@@ -100,6 +102,9 @@ def download_and_extract_ui_files(version: str):
                 break
             current = os.path.dirname(current)
 
+    if workspace is not None:
+        os.environ.setdefault("COMFYUI_CWD", workspace)
+
     logger.info("Installing comfystream package...")
     subprocess.check_call([sys.executable, "-m", "pip", "install", "-e", "."])
 

server/app.py

@@ -654,15 +654,26 @@ async def on_startup(app: web.Application):
     if app["media_ports"]:
         patch_loop_datagram(app["media_ports"])
 
+    comfy_kwargs = {}
+    if app.get("config"):
+        # Pass config directly to ComfyUI; do not override with other ComfyUI flags
+        comfy_kwargs["config"] = app["config"]
+    else:
+        comfy_kwargs.update(
+            {
+                "cwd": app["workspace"],
+                "disable_cuda_malloc": True,
+                "gpu_only": True,
+                "preview_method": "none",
+                "logging_level": app.get("logging_level", None),
+                "blacklist_custom_nodes": ["ComfyUI-Manager"],
+            }
+        )
+
     app["pipeline"] = Pipeline(
         width=512,
         height=512,
-        cwd=app["workspace"],
-        disable_cuda_malloc=True,
-        gpu_only=True,
-        preview_method="none",
-        comfyui_inference_log_level=app.get("comfyui_inference_log_level", None),
-        blacklist_custom_nodes=["ComfyUI-Manager"],
+        **comfy_kwargs,
     )
     await app["pipeline"].initialize()
     app["pcs"] = set()
@@ -681,7 +692,21 @@ async def on_shutdown(app: web.Application):
     parser.add_argument("--port", default=8889, help="Set the signaling port")
     parser.add_argument("--media-ports", default=None, help="Set the UDP ports for WebRTC media")
     parser.add_argument("--host", default="127.0.0.1", help="Set the host")
-    parser.add_argument("--workspace", default=None, required=True, help="Set Comfy workspace")
+    parser.add_argument(
+        "-c",
+        "--config",
+        dest="config",
+        default=None,
+        help="Path to ComfyUI config file (yaml/json/ini/conf). When provided, it is passed directly to ComfyUI.",
+    )
+    parser.add_argument(
+        "--workspace",
+        "--cwd",
+        dest="workspace",
+        default=os.environ.get("COMFYUI_CWD"),
+        required=True,
+        help="Set ComfyUI workspace directory (preferred; alias: --cwd)",
+    )
     parser.add_argument(
         "--log-level",
         default="INFO",
@@ -707,10 +732,10 @@ async def on_shutdown(app: web.Application):
         help="Set the global logging level for ComfyUI",
     )
     parser.add_argument(
-        "--comfyui-inference-log-level",
+        "--logging-level",
         default=None,
         choices=logging._nameToLevel.keys(),
-        help="Set the logging level for ComfyUI inference",
+        help="Set the logging level for ComfyUI (passed through to ComfyUI Configuration)",
     )
     args = parser.parse_args()
 
@@ -720,9 +745,14 @@ async def on_shutdown(app: web.Application):
         datefmt="%H:%M:%S",
     )
 
+    logger.info(f"Using ComfyUI workspace: {args.workspace}")
+    os.environ.setdefault("COMFYUI_CWD", args.workspace)
+
     app = web.Application()
     app["media_ports"] = args.media_ports.split(",") if args.media_ports else None
     app["workspace"] = args.workspace
+    app["logging_level"] = args.logging_level
+    app["config"] = args.config
 
     app.on_startup.append(on_startup)
     app.on_shutdown.append(on_shutdown)
@@ -768,7 +798,5 @@ def force_print(*args, **kwargs):
     timeout_filter = ComfyStreamTimeoutFilter()
     logging.getLogger("comfy.cmd.execution").addFilter(timeout_filter)
     logging.getLogger("comfystream").addFilter(timeout_filter)
-    if args.comfyui_inference_log_level:
-        app["comfyui_inference_log_level"] = args.comfyui_inference_log_level
 
     web.run_app(app, host=args.host, port=int(args.port), print=force_print)

server/byoc.py

@@ -30,10 +30,19 @@ def main():
     )
     parser.add_argument("--port", default=8000, help="Set the server port")
     parser.add_argument("--host", default="0.0.0.0", help="Set the host")
+    parser.add_argument(
+        "-c",
+        "--config",
+        dest="config",
+        default=None,
+        help="Path to ComfyUI config file (yaml/json/ini/conf). When provided, it is passed directly to ComfyUI.",
+    )
     parser.add_argument(
         "--workspace",
-        default=os.getcwd() + "/../ComfyUI",
-        help="Set Comfy workspace (Default: ../ComfyUI)",
+        "--cwd",
+        dest="workspace",
+        default=os.environ.get("COMFYUI_CWD", os.getcwd() + "/../ComfyUI"),
+        help="Set ComfyUI workspace directory (preferred; alias: --cwd)",
     )
     parser.add_argument(
         "--log-level",
@@ -48,10 +57,10 @@ def main():
         help="Set the global logging level for ComfyUI",
     )
     parser.add_argument(
-        "--comfyui-inference-log-level",
+        "--logging-level",
         default=None,
         choices=logging._nameToLevel.keys(),
-        help="Set the logging level for ComfyUI inference",
+        help="Set the logging level for ComfyUI (passed through to ComfyUI Configuration)",
     )
     parser.add_argument(
         "--disable-frame-skip",
@@ -80,6 +89,9 @@ def main():
     )
     logging.getLogger("comfy.model_detection").setLevel(logging.WARNING)
 
+    logger.info(f"Using ComfyUI workspace: {args.workspace}")
+    os.environ.setdefault("COMFYUI_CWD", args.workspace)
+
     # Allow overriding of ComfyUI log levels.
     if args.comfyui_log_level:
         log_level = logging._nameToLevel.get(args.comfyui_log_level.upper())
@@ -104,13 +116,13 @@ def force_print(*args, **kwargs):
     frame_processor = ComfyStreamFrameProcessor(
         width=args.width,
         height=args.height,
+        config=args.config,
         workspace=args.workspace,
         disable_cuda_malloc=True,
         gpu_only=True,
         preview_method="none",
         blacklist_custom_nodes=["ComfyUI-Manager"],
-        logging_level=args.comfyui_log_level,
-        comfyui_inference_log_level=args.comfyui_inference_log_level,
+        logging_level=args.logging_level or args.comfyui_log_level,
     )
 
     # Create frame skip configuration only if enabled

server/frame_processor.py

@@ -259,16 +259,23 @@ async def load_model(self, **kwargs):
         params = {**self._load_params, **kwargs}
 
         if self.pipeline is None:
+            comfy_kwargs: Dict[str, Any] = {
+                "cwd": params.get("workspace", os.getcwd()),
+                "disable_cuda_malloc": params.get("disable_cuda_malloc", True),
+                "gpu_only": params.get("gpu_only", True),
+                "preview_method": params.get("preview_method", "none"),
+                "logging_level": params.get("logging_level", "INFO"),
+                "blacklist_custom_nodes": ["ComfyUI-Manager"],
+            }
+
+            # If a ComfyUI config file is provided, prefer it and pass it through.
+            if params.get("config"):
+                comfy_kwargs = {"config": params["config"]}
+
             self.pipeline = Pipeline(
                 width=int(params.get("width", 512)),
                 height=int(params.get("height", 512)),
-                cwd=params.get("workspace", os.getcwd()),
-                disable_cuda_malloc=params.get("disable_cuda_malloc", True),
-                gpu_only=params.get("gpu_only", True),
-                preview_method=params.get("preview_method", "none"),
-                comfyui_inference_log_level=params.get("comfyui_inference_log_level", "INFO"),
-                logging_level=params.get("comfyui_inference_log_level", "INFO"),
-                blacklist_custom_nodes=["ComfyUI-Manager"],
+                **comfy_kwargs,
             )
             await self.pipeline.initialize()
 

src/comfystream/__init__.py

@@ -1,9 +1,12 @@
-from .client import ComfyStreamClient
-from .exceptions import ComfyStreamAudioBufferError, ComfyStreamInputTimeoutError
-from .pipeline import Pipeline
-from .pipeline_state import PipelineState, PipelineStateManager
-from .server.metrics import MetricsManager, StreamStatsManager
-from .server.utils import FPSMeter, temporary_log_level
+from comfy_compatibility.imports import MAIN_PY, SITE_PACKAGES, ImportContext
+
+with ImportContext("comfy", "comfy_extras", "comfy.vendor", order=[SITE_PACKAGES, MAIN_PY]):
+    from .client import ComfyStreamClient
+    from .exceptions import ComfyStreamAudioBufferError, ComfyStreamInputTimeoutError
+    from .pipeline import Pipeline
+    from .pipeline_state import PipelineState, PipelineStateManager
+    from .server.metrics import MetricsManager, StreamStatsManager
+    from .server.utils import FPSMeter, temporary_log_level
 
 __all__ = [
     "ComfyStreamClient",

src/comfystream/pipeline.py

@@ -36,7 +36,7 @@ def __init__(
         self,
         width: int = 512,
         height: int = 512,
-        comfyui_inference_log_level: Optional[int] = None,
+        logging_level: Optional[int] = None,
         auto_warmup: bool = False,
         bootstrap_default_prompt: bool = True,
         **kwargs,
@@ -46,12 +46,11 @@ def __init__(
         Args:
             width: Width of the video frames (default: 512)
             height: Height of the video frames (default: 512)
-            comfyui_inference_log_level: The logging level for ComfyUI inference.
-                Defaults to None, using the global ComfyUI log level.
+            logging_level: The logging level for ComfyUI (passed to Configuration).
             auto_warmup: Whether to run warmup automatically after prompts are set.
             bootstrap_default_prompt: Whether to run the default workflow once during
                 initialization to start ComfyUI before prompts are applied.
-            **kwargs: Additional arguments to pass to the ComfyStreamClient
+            **kwargs: Additional arguments to pass to the ComfyStreamClient to configure ComfyUI
         """
         self.client = ComfyStreamClient(**kwargs)
         self.width = width
@@ -64,7 +63,7 @@ def __init__(
 
         self.processed_audio_buffer = np.array([], dtype=np.int16)
 
-        self._comfyui_inference_log_level = comfyui_inference_log_level
+        self._comfy_logging_level = logging_level
         self._cached_modalities: Optional[Set[str]] = None
         self._cached_io_capabilities: Optional[WorkflowModality] = None
         self.state_manager = PipelineStateManager(self.client)
@@ -624,7 +623,7 @@ async def get_processed_video_frame(self) -> av.VideoFrame:
             return frame
 
         # Get processed output from client
-        async with temporary_log_level("comfy", self._comfyui_inference_log_level):
+        async with temporary_log_level("comfy", self._comfy_logging_level):
             out_tensor = await self.client.get_video_output()
 
         processed_frame = self.video_postprocess(out_tensor)
@@ -657,7 +656,7 @@ async def get_processed_audio_frame(self) -> av.AudioFrame:
 
         # Process audio if needed
         if frame.samples > len(self.processed_audio_buffer):
-            async with temporary_log_level("comfy", self._comfyui_inference_log_level):
+            async with temporary_log_level("comfy", self._comfy_logging_level):
                 out_tensor = await self.client.get_audio_output()
             self.processed_audio_buffer = np.concatenate([self.processed_audio_buffer, out_tensor])
 
@@ -681,7 +680,7 @@ async def get_text_output(self) -> str | None:
         if not self.produces_text_output():
             return None
 
-        async with temporary_log_level("comfy", self._comfyui_inference_log_level):
+        async with temporary_log_level("comfy", self._comfy_logging_level):
             out_text = await self.client.get_text_output()
 
         return out_text

src/comfystream/scripts/README.md

@@ -14,15 +14,18 @@
 
 From the repository root:
 
-> The `--workspace` flag is optional and will default to `$COMFY_UI_WORKSPACE` or `~/comfyui`.
+> The `--workspace` flag is optional and will default to `$COMFYUI_CWD` or `~/comfyui`.
 
 ### Install custom nodes
+
 ```bash
 python src/comfystream/scripts/setup_nodes.py --workspace /path/to/comfyui
 ```
+
 > The optional flag `--pull-branches` can be used to ensure the latest git changes are pulled for any custom nodes defined with a `branch` in nodes.yaml
 
 ### Download models and compile tensorrt engines
+
 ```bash
 python src/comfystream/scripts/setup_models.py --workspace /path/to/comfyui
 ```
@@ -42,7 +45,7 @@ nodes:
       - "tensorrt"
 ```
 
-> The `branch` property can be substituted with a SHA-256 commit hash for pinning custom node versions 
+> The `branch` property can be substituted with a SHA-256 commit hash for pinning custom node versions
 
 ### Models (models.yaml)
 
@@ -70,6 +73,6 @@ workspace/
 
 ## Environment Variables
 
-- `COMFY_UI_WORKSPACE`: Base directory for installation
+- `COMFYUI_CWD`: Base directory for installation
 - `PYTHONPATH`: Defaults to workspace directory
 - `CUSTOM_NODES_PATH`: Custom nodes directory

src/comfystream/scripts/setup_models.py

@@ -11,9 +11,11 @@
 def parse_args():
     parser = argparse.ArgumentParser(description="Setup ComfyUI models")
     parser.add_argument(
+        "--cwd",
         "--workspace",
-        default=os.environ.get("COMFY_UI_WORKSPACE", os.path.expanduser("~/comfyui")),
-        help="ComfyUI workspace directory (default: ~/comfyui or $COMFY_UI_WORKSPACE)",
+        dest="workspace",
+        default=os.environ.get("COMFYUI_CWD", os.path.expanduser("~/comfyui")),
+        help="ComfyUI workspace directory (default: ~/comfyui or $COMFYUI_CWD)",
     )
     return parser.parse_args()
 

src/comfystream/scripts/setup_nodes.py

@@ -11,9 +11,11 @@
 def parse_args():
     parser = argparse.ArgumentParser(description="Setup ComfyUI nodes and models")
     parser.add_argument(
+        "--cwd",
         "--workspace",
-        default=os.environ.get("COMFY_UI_WORKSPACE", Path("~/comfyui").expanduser()),
-        help="ComfyUI workspace directory (default: ~/comfyui or $COMFY_UI_WORKSPACE)",
+        dest="workspace",
+        default=os.environ.get("COMFYUI_CWD", Path("~/comfyui").expanduser()),
+        help="ComfyUI workspace directory (default: ~/comfyui or $COMFYUI_CWD)",
     )
     parser.add_argument(
         "--pull-branches",
@@ -25,7 +27,7 @@ def parse_args():
 
 
 def setup_environment(workspace_dir):
-    os.environ["COMFY_UI_WORKSPACE"] = str(workspace_dir)
+    os.environ["COMFYUI_CWD"] = str(workspace_dir)
     os.environ["PYTHONPATH"] = str(workspace_dir)
     os.environ["CUSTOM_NODES_PATH"] = str(workspace_dir / "custom_nodes")