diff --git a/nodes/server_manager.py b/nodes/server_manager.py index 454b1dc0..a7436d33 100644 --- a/nodes/server_manager.py +++ b/nodes/server_manager.py @@ -148,7 +148,7 @@ async def start(self, port=None, host=None): self.host = host # Get the path to the ComfyStream server directory and script - server_dir = Path(__file__).parent.parent / "src" / "comfystream" / "server" + server_dir = Path(__file__).parent.parent / "server" server_script = server_dir / "app.py" logging.info(f"Server script: {server_script}") diff --git a/src/comfystream/server/app.py b/server/app.py similarity index 99% rename from src/comfystream/server/app.py rename to server/app.py index 676b3b74..a52f8061 100644 --- a/src/comfystream/server/app.py +++ b/server/app.py @@ -4,6 +4,11 @@ import logging import os import sys +import torch + +# Initialize CUDA before any other imports to prevent core dump. +if torch.cuda.is_available(): + torch.cuda.init() from aiohttp import web from aiortc import ( diff --git a/src/comfystream/__init__.py b/src/comfystream/__init__.py index 5e1b9fac..b58bf2e4 100644 --- a/src/comfystream/__init__.py +++ b/src/comfystream/__init__.py @@ -1,7 +1,6 @@ from .client import ComfyStreamClient from .pipeline import Pipeline from .server.utils import temporary_log_level -from .server.app import VideoStreamTrack, AudioStreamTrack from .server.utils import FPSMeter from .server.metrics import MetricsManager, StreamStatsManager @@ -9,8 +8,6 @@ 'ComfyStreamClient', 'Pipeline', 'temporary_log_level', - 'VideoStreamTrack', - 'AudioStreamTrack', 'FPSMeter', 'MetricsManager', 'StreamStatsManager'