Update test_ai_integration.py

test_ai_integration.py

@@ -156,15 +156,15 @@ 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
+            }
             ai_accelerator_for_loading.load_model(
-                model_id=model_id,
-                model=model,
-                processor=processor,
-                model_config={
-                    "size_bytes": model_size,
-                    "unlimited_vram": True,
-                    "allow_resize": True
-                }
+                model_name=model_id,
+                model_path=None,
+                model_data=model_data
             )
             
             print(f"Model '{model_id}' loaded successfully to WebSocket storage.")
@@ -185,15 +185,15 @@ 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
+                }
                 ai_accelerator_for_loading.load_model(
-                    model_id=model_id,
-                    model=None,  # Use zero-copy mode
-                    processor=None,
-                    model_config={
-                        "zero_copy": True,
-                        "unlimited_vram": True,
-                        "allow_resize": True
-                    }
+                    model_name=model_id,
+                    model_path=None,
+                    model_data=model_data
                 )
                 components['model_id'] = model_id
                 print("Successfully loaded model in zero-copy mode")