From 298a95432da41e98a3975db9708338d656e3ddf4 Mon Sep 17 00:00:00 2001
From: Will Miao <13051207myq@gmail.com>
Date: Wed, 23 Jul 2025 18:02:38 +0800
Subject: [PATCH] feat: Integrate WebSocket routes for download progress
 tracking in standalone manager

---
 py/routes/base_model_routes.py | 3 ---
 standalone.py                  | 6 ++++++
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/py/routes/base_model_routes.py b/py/routes/base_model_routes.py
index d438d69e..8cee15f3 100644
--- a/py/routes/base_model_routes.py
+++ b/py/routes/base_model_routes.py
@@ -8,7 +8,6 @@ from typing import Dict
 import jinja2
 
 from ..utils.routes_common import ModelRouteUtils
-from ..services.service_registry import ServiceRegistry
 from ..services.websocket_manager import ws_manager
 from ..services.settings_manager import settings
 from ..config import config
@@ -516,8 +515,6 @@ class BaseModelRoutes(ABC):
                     'error': 'Download ID is required'
                 }, status=400)
             
-            # Get progress information from websocket manager
-            from ..services.websocket_manager import ws_manager
             progress_data = ws_manager.get_download_progress(download_id)
             
             if progress_data is None:
diff --git a/standalone.py b/standalone.py
index a809aaad..2ad50ade 100644
--- a/standalone.py
+++ b/standalone.py
@@ -320,6 +320,7 @@ class StandaloneLoraManager(LoraManager):
         from py.routes.misc_routes import MiscRoutes
         from py.routes.example_images_routes import ExampleImagesRoutes
         from py.routes.stats_routes import StatsRoutes
+        from py.services.websocket_manager import ws_manager
         
 
         register_default_model_types()
@@ -335,6 +336,11 @@ class StandaloneLoraManager(LoraManager):
         UpdateRoutes.setup_routes(app)
         MiscRoutes.setup_routes(app)
         ExampleImagesRoutes.setup_routes(app)
+
+        # Setup WebSocket routes that are shared across all model types
+        app.router.add_get('/ws/fetch-progress', ws_manager.handle_connection)
+        app.router.add_get('/ws/download-progress', ws_manager.handle_download_connection)
+        app.router.add_get('/ws/init-progress', ws_manager.handle_init_connection)
         
         # Schedule service initialization
         app.on_startup.append(lambda app: cls._initialize_services())