Update server.py

server.py

@@ -283,24 +283,25 @@ class VirtualGPUServer:
         html.extend(['</table></body></html>'])
         return web.Response(text='\n'.join(html), content_type='text/html')
 
-async def main():
-    server = VirtualGPUServer()
-    
-    # Start WebSocket server
-    websocket_server = await websockets.serve(server.handle_connection, "0.0.0.0", 8765)
-    
-    # Start HTTP server
-    runner = web.AppRunner(server.app)
-    await runner.setup()
-    site = web.TCPSite(runner, '0.0.0.0', 8080)
-    await site.start()
-    
-    print("Virtual GPU Server running:")
-    print("- WebSocket: ws://localhost:8765")
-    print("- HTTP Interface: http://localhost:8080")
-    
-    # Run forever
-    await asyncio.Future()
+# Create the application instance
+server = VirtualGPUServer()
+app = server.app
 
+# WebSocket setup
+@app.on_startup.append
+async def start_websocket_server(app):
+    app['websocket_server'] = await websockets.serve(
+        server.handle_connection,
+        "0.0.0.0",  # Listen on all interfaces
+        8765
+    )
+    print("WebSocket server started on ws://0.0.0.0:8765")
+
+@app.on_cleanup.append
+async def cleanup_websocket_server(app):
+    app['websocket_server'].close()
+    await app['websocket_server'].wait_closed()
+
+# For running directly (development)
 if __name__ == "__main__":
-    asyncio.run(main())
+    web.run_app(app, host="0.0.0.0", port=7860)