Update test_ai_integration.py

test_ai_integration.py

@@ -78,7 +78,21 @@ def test_ai_integration():
         'ai_accelerators': [],
         'model_id': None,
         'vram': None,
-        'storage': None
+        'storage': None,
+        'model_config': None,
+        'tensor_registry': {},
+        'initialized': False
+    }
+    
+    # Initialize global tensor registry
+    global_tensor_registry = {
+        'model_tensors': {},
+        'runtime_tensors': {},
+        'placeholder_tensors': {},
+        'stats': {
+            'total_vram_used': 0,
+            'active_tensors': 0
+        }
     }
     
     # Increase file descriptor limit
@@ -95,16 +109,29 @@ def test_ai_integration():
     try:
         # Use WebSocket connection manager for proper resource handling
         with websocket_manager() as storage:
+            components['storage'] = storage  # Save storage reference
+            
             # Initialize virtual GPU stack with unlimited WebSocket storage
             chip_for_loading = Chip(chip_id=0, vram_size_gb=None)  # Unlimited storage
+            components['chips'].append(chip_for_loading)
             
             # Initialize VRAM with WebSocket storage
             vram = VirtualVRAM()
             vram.storage = storage  # Share WebSocket connection
+            components['vram'] = vram
             
-            # Set up AI accelerator
+            # Set up AI accelerator with proper initialization
             ai_accelerator_for_loading = chip_for_loading.ai_accelerator
             ai_accelerator_for_loading.vram = vram  # Use WebSocket-backed VRAM
+            ai_accelerator_for_loading.initialize_tensor_cores()  # Ensure tensor cores are ready
+            components['ai_accelerators'].append(ai_accelerator_for_loading)
+
+            # Initialize model registry in WebSocket storage
+            storage.store_state("model_registry", "state", {
+                "initialized": True,
+                "max_vram": None,  # Unlimited
+                "active_models": {}
+            })
 
         # Load BLIP-2 Large model directly to WebSocket storage
         from transformers import AutoModelForCausalLM, AutoProcessor