From fd1c5ac526dc646bd0e568142f8f384f39c31c14 Mon Sep 17 00:00:00 2001
From: Elite Encoder <john@eliteencoder.net>
Date: Tue, 15 Apr 2025 22:45:17 -0400
Subject: [PATCH] exclude app.py from comfystream package restore torch init to
 app.py

---
 nodes/server_manager.py                   | 2 +-
 {src/comfystream/server => server}/app.py | 5 +++++
 src/comfystream/__init__.py               | 3 ---
 3 files changed, 6 insertions(+), 4 deletions(-)
 rename {src/comfystream/server => server}/app.py (99%)

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'