Upload test_ai_integration.py

test_ai_integration.py

@@ -156,15 +156,11 @@ def test_ai_integration():
             print(f"Model size: {model_size / (1024**3):.2f} GB")
             
             # Store model in WebSocket storage with size information
-            model_data = {
-                "size_bytes": model_size,
-                "unlimited_vram": True,
-                "allow_resize": True
-            }
+            # Load model directly using AIAccelerator's load_model method
             ai_accelerator_for_loading.load_model(
-                model_name=model_id,
-                model_path=None,
-                model_data=model_data
+                model_id=model_id,
+                model=model,
+                processor=processor
             )
             
             print(f"Model '{model_id}' loaded successfully to WebSocket storage.")
@@ -185,15 +181,11 @@ def test_ai_integration():
             print("Falling back to zero-copy tensor mode...")
             # Try loading with zero-copy tensor mode
             try:
-                model_data = {
-                    "zero_copy": True,
-                    "unlimited_vram": True,
-                    "allow_resize": True
-                }
+                # Try zero-copy loading
                 ai_accelerator_for_loading.load_model(
-                    model_name=model_id,
-                    model_path=None,
-                    model_data=model_data
+                    model_id=model_id,
+                    model=None,
+                    processor=None
                 )
                 components['model_id'] = model_id
                 print("Successfully loaded model in zero-copy mode")