Update ai.py

ai.py

@@ -49,9 +49,10 @@ class AIAccelerator:
         )
         self.tensor_cores_initialized = False
         
-        # Initialize model and tensor tracking
+        # Initialize model, tensor, and tokenizer tracking
         self.model_registry: Dict[str, Dict[str, Any]] = {}  # Track loaded models
         self.tensor_registry: Dict[str, Dict[str, Any]] = {}  # Track tensor metadata
+        self.tokenizer_registry: Dict[str, Any] = {}  # Track tokenizers
         self.resource_monitor = {
             'vram_used': 0,
             'active_tensors': 0,
@@ -571,25 +572,45 @@ class AIAccelerator:
     def load_model(self, model_id: str, model: Any, processor: Any):
         """Loads a model directly into WebSocket storage without CPU intermediary."""
         try:
+            if model is None and processor is None:
+                # Zero-copy mode
+                self.model_registry[model_id] = {
+                    "zero_copy": True,
+                    "websocket_mapped": True
+                }
+                self.tokenizer_registry[model_id] = None
+                self.model_loaded = True
+                return
+
             # Extract model metadata
-            model_info = {
-                "architecture": model.__class__.__name__,
-                "processor": processor.__class__.__name__,
-                "config": model.config.to_dict() if hasattr(model, "config") else {}
-            }
-            
+            try:
+                model_info = {
+                    "architecture": model.__class__.__name__ if model else "Unknown",
+                    "processor": processor.__class__.__name__ if processor else "Unknown",
+                    "config": self._serialize_model_config(model.config) if hasattr(model, "config") else {}
+                }
+            except Exception as e:
+                print(f"Warning: Error serializing model metadata: {e}")
+                model_info = {"error": str(e)}
+
+            # Verify WebSocket connection
+            if not self.storage or not self.storage.wait_for_connection():
+                raise RuntimeError("WebSocket connection not available")
+
             # Store model state in WebSocket storage
-            self.storage.store_state(f"models/{model_id}", "info", model_info)
+            if not self.storage.store_state(f"models/{model_id}/info", "info", model_info):
+                raise RuntimeError("Failed to store model info")
             
             # Map weight tensors directly to WebSocket storage
-            if hasattr(model, "state_dict"):
+            if model is not None and hasattr(model, "state_dict"):
                 model_weights = {}
                 
                 for name, param in model.state_dict().items():
                     tensor_id = f"{model_id}/weights/{name}"
                     
                     # Store tensor directly in WebSocket storage
-                    self.storage.store_tensor(tensor_id, param.detach().numpy())
+                    if not self.storage.store_tensor(tensor_id, param.detach().numpy()):
+                        raise RuntimeError(f"Failed to store tensor {name}")
                     model_weights[name] = tensor_id
                 
                 # Store only WebSocket references
@@ -601,9 +622,11 @@ class AIAccelerator:
             else:
                 # Store the entire model state in WebSocket storage
                 tensor_id = f"{model_id}/model_state"
-                self.storage.store_state(f"models/{model_id}", "state", model)
+                if not self.storage.store_state(f"models/{model_id}/state", "state", model):
+                    raise RuntimeError("Failed to store model state")
                 self.model_registry[model_id] = tensor_id
             
+            # Store tokenizer/processor
             self.tokenizer_registry[model_id] = processor
             self.model_loaded = True
             print(f"Model '{model_id}' loaded into WebSocket storage")
@@ -675,5 +698,3 @@ class AIAccelerator:
         except Exception as e:
             print(f"[ERROR] WebSocket-based inference failed for idx={idx}: {e}")
             return None
-
-