Update websocket_storage.py

websocket_storage.py

@@ -6,8 +6,25 @@ import threading
 from queue import Queue
 import time
 import asyncio
+import hashlib
+
 class WebSocketGPUStorage:
-    def __init__(self, url: str = "wss://factorst-wbs1.hf.space/ws"):  # Default to local WebSocket server
+    # Singleton instance
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls, url: str = "wss://factorst-wbs1.hf.space/ws"):
+        with cls._lock:
+            if cls._instance is None:
+                cls._instance = super().__new__(cls)
+                cls._instance._init_singleton(url)
+            return cls._instance
+
+    def _init_singleton(self, url: str):
+        """Initialize the singleton instance"""
+        if hasattr(self, 'initialized'):
+            return
+        
         self.url = url
         self.websocket = None
         self.connected = False
@@ -20,7 +37,6 @@ class WebSocketGPUStorage:
         self.last_error_time = 0
         self.max_retries = 5
         self.tensor_registry: Dict[str, Dict[str, Any]] = {}  # Track tensor metadata
-        self.resource_monitor = {'vram_used': 0, 'active_tensors': 0}
         self.model_registry: Dict[str, Dict[str, Any]] = {}  # Track loaded models
         self.resource_monitor = {
             'vram_used': 0,
@@ -31,6 +47,11 @@ class WebSocketGPUStorage:
         # Start WebSocket connection in a separate thread
         self.ws_thread = threading.Thread(target=self._run_websocket_loop, daemon=True)
         self.ws_thread.start()
+        self.initialized = True
+        
+    def __init__(self, url: str = "wss://factorst-wbs1.hf.space/ws"):
+        """This will actually just return the singleton instance"""
+        pass
 
     def _run_websocket_loop(self):
         self._loop = asyncio.new_event_loop()