Update ai.py

ai.py

@@ -24,14 +24,16 @@ class AIAccelerator:
     processing capabilities of the vGPU for AI workloads.
     """
     
-    def __init__(self, vram=None, num_sms: int = 800, cores_per_sm: int = 222):
-        """Initialize AI Accelerator with electron-speed awareness and WebSocket storage."""
+    def __init__(self, vram=None, num_sms: int = 800, cores_per_sm: int = 222, storage=None):
+        """Initialize AI Accelerator with electron-speed awareness and shared WebSocket storage."""
         from electron_speed import TARGET_SWITCHES_PER_SEC, TRANSISTORS_ON_CHIP, drift_velocity
-        from websocket_storage import WebSocketGPUStorage
         
-        self.storage = WebSocketGPUStorage()
-        if not self.storage.wait_for_connection():
-            raise RuntimeError("Could not connect to GPU storage server")
+        self.storage = storage  # Use the shared storage instance
+        if self.storage is None:
+            from websocket_storage import WebSocketGPUStorage
+            self.storage = WebSocketGPUStorage()  # Only create new if not provided
+            if not self.storage.wait_for_connection():
+                raise RuntimeError("Could not connect to GPU storage server")
             
         self.vram = vram
         self.num_sms = num_sms