Spaces:

factorstudios
/

WBS1

Sleeping

App Files Files Community

Factor Studios commited on Aug 12, 2025

Commit

ba864c5

verified ·

1 Parent(s): 68afb37

Update server.py

Browse files

Files changed (1) hide show

server.py +114 -87

server.py CHANGED Viewed

@@ -7,10 +7,13 @@ import uuid
 import time
 from typing import Dict, Any, Optional
 import numpy as np
-from aiohttp import web
-import aiohttp
 from datetime import datetime
 class VirtualGPUServer:
     def __init__(self):
         self.base_path = Path(__file__).parent / "storage"
@@ -29,17 +32,12 @@ class VirtualGPUServer:
         self.memory_cache: Dict[str, Any] = {}
         # Active connections and sessions
-        self.active_connections: Dict[str, websockets.WebSocketServerProtocol] = {}
         self.active_sessions: Dict[str, Dict[str, Any]] = {}
         # Performance monitoring
         self.ops_counter = 0
         self.start_time = time.time()
-        # HTTP app
-        self.app = web.Application()
-        self.app.router.add_get('/', self.handle_index)
-        self.app.router.add_get('/files', self.handle_files)
     async def handle_vram_operation(self, operation: dict) -> dict:
         """Handle VRAM read/write operations"""
@@ -216,92 +214,121 @@ class VirtualGPUServer:
             'memory_cache_size': len(self.memory_cache)
         }
-    async def handle_index(self, request):
-        """Handle HTTP index request"""
-        stats = self.get_stats()
-        html = f"""
-        <!DOCTYPE html>
-        <html>
-        <head>
-            <title>Virtual GPU Server</title>
-            <style>
-                body {{ font-family: Arial, sans-serif; margin: 40px; }}
-                table {{ border-collapse: collapse; width: 100%; margin-top: 20px; }}
-                th, td {{ padding: 12px; text-align: left; border-bottom: 1px solid #ddd; }}
-                th {{ background-color: #f2f2f2; }}
-                .stats {{ background-color: #f9f9f9; padding: 20px; border-radius: 5px; }}
-            </style>
-        </head>
-        <body>
-            <h1>Virtual GPU Server Status</h1>
-            <div class="stats">
-                <h2>Server Statistics</h2>
-                <ul>
-                    <li>Uptime: {stats['uptime']:.2f} seconds</li>
-                    <li>Total Operations: {stats['total_operations']}</li>
-                    <li>Operations per Second: {stats['ops_per_second']:.2f}</li>
-                    <li>Active Connections: {stats['active_connections']}</li>
-                    <li>VRAM Cache Size: {stats['vram_cache_size']}</li>
-                    <li>State Cache Size: {stats['state_cache_size']}</li>
-                    <li>Memory Cache Size: {stats['memory_cache_size']}</li>
-                </ul>
-            </div>
-            <h2>Server Files</h2>
-            <iframe src="/files" style="width: 100%; height: 500px; border: none;"></iframe>
-        </body>
-        </html>
-        """
-        return web.Response(text=html, content_type='text/html')
-    async def handle_files(self, request):
-        """Handle HTTP files listing request"""
-        def format_size(size):
-            for unit in ['B', 'KB', 'MB', 'GB']:
-                if size < 1024:
-                    return f"{size:.2f} {unit}"
-                size /= 1024
-            return f"{size:.2f} TB"
-        html = ['<!DOCTYPE html><html><head>',
-                '<style>',
-                'body { font-family: Arial, sans-serif; margin: 20px; }',
-                'table { border-collapse: collapse; width: 100%; }',
-                'th, td { padding: 12px; text-align: left; border-bottom: 1px solid #ddd; }',
-                'th { background-color: #f2f2f2; }',
-                '</style></head><body>',
-                '<h2>Server Files</h2>',
-                '<table><tr><th>Path</th><th>Size</th><th>Last Modified</th></tr>']
-        for root, _, files in os.walk(self.base_path):
-            for file in files:
-                full_path = Path(root) / file
-                rel_path = full_path.relative_to(self.base_path)
-                size = format_size(os.path.getsize(full_path))
-                mtime = datetime.fromtimestamp(os.path.getmtime(full_path))
-                html.append(f'<tr><td>{rel_path}</td><td>{size}</td><td>{mtime}</td></tr>')
-        html.extend(['</table></body></html>'])
-        return web.Response(text='\n'.join(html), content_type='text/html')
-# 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__":
-    web.run_app(app, host="0.0.0.0", port=7860)

 import time
 from typing import Dict, Any, Optional
 import numpy as np
+from fastapi import FastAPI, WebSocket
+from fastapi.responses import HTMLResponse
 from datetime import datetime
+# Create FastAPI instance
+app = FastAPI()
 class VirtualGPUServer:
     def __init__(self):
         self.base_path = Path(__file__).parent / "storage"
         self.memory_cache: Dict[str, Any] = {}
         # Active connections and sessions
+        self.active_connections: Dict[str, WebSocket] = {}
         self.active_sessions: Dict[str, Dict[str, Any]] = {}
         # Performance monitoring
         self.ops_counter = 0
         self.start_time = time.time()
     async def handle_vram_operation(self, operation: dict) -> dict:
         """Handle VRAM read/write operations"""
             'memory_cache_size': len(self.memory_cache)
         }
+server = VirtualGPUServer()
+@app.get("/", response_class=HTMLResponse)
+async def handle_index():
+    """Handle HTTP index request"""
+    stats = server.get_stats()
+    html = f"""
+    <!DOCTYPE html>
+    <html>
+    <head>
+        <title>Virtual GPU Server</title>
+        <style>
+            body {{ font-family: Arial, sans-serif; margin: 40px; }}
+            table {{ border-collapse: collapse; width: 100%; margin-top: 20px; }}
+            th, td {{ padding: 12px; text-align: left; border-bottom: 1px solid #ddd; }}
+            th {{ background-color: #f2f2f2; }}
+            .stats {{ background-color: #f9f9f9; padding: 20px; border-radius: 5px; }}
+        </style>
+    </head>
+    <body>
+        <h1>Virtual GPU Server Status</h1>
+        <div class="stats">
+            <h2>Server Statistics</h2>
+            <ul>
+                <li>Uptime: {stats['uptime']:.2f} seconds</li>
+                <li>Total Operations: {stats['total_operations']}</li>
+                <li>Operations per Second: {stats['ops_per_second']:.2f}</li>
+                <li>Active Connections: {stats['active_connections']}</li>
+                <li>VRAM Cache Size: {stats['vram_cache_size']}</li>
+                <li>State Cache Size: {stats['state_cache_size']}</li>
+                <li>Memory Cache Size: {stats['memory_cache_size']}</li>
+            </ul>
+        </div>
+        <h2>Server Files</h2>
+        <iframe src="/files" style="width: 100%; height: 500px; border: none;"></iframe>
+    </body>
+    </html>
+    """
+    return HTMLResponse(content=html)
+@app.get("/files", response_class=HTMLResponse)
+async def handle_files():
+    """Handle HTTP files listing request"""
+    def format_size(size):
+        for unit in ['B', 'KB', 'MB', 'GB']:
+            if size < 1024:
+                return f"{size:.2f} {unit}"
+            size /= 1024
+        return f"{size:.2f} TB"
+    html = ['<!DOCTYPE html><html><head>',
+            '<style>',
+            'body { font-family: Arial, sans-serif; margin: 20px; }',
+            'table { border-collapse: collapse; width: 100%; }',
+            'th, td { padding: 12px; text-align: left; border-bottom: 1px solid #ddd; }',
+            'th { background-color: #f2f2f2; }',
+            '</style></head><body>',
+            '<h2>Server Files</h2>',
+            '<table><tr><th>Path</th><th>Size</th><th>Last Modified</th></tr>']
+    for root, _, files in os.walk(server.base_path):
+        for file in files:
+            full_path = Path(root) / file
+            rel_path = full_path.relative_to(server.base_path)
+            size = format_size(os.path.getsize(full_path))
+            mtime = datetime.fromtimestamp(os.path.getmtime(full_path))
+            html.append(f'<tr><td>{rel_path}</td><td>{size}</td><td>{mtime}</td></tr>')
+    html.extend(['</table></body></html>'])
+    return HTMLResponse(content='\n'.join(html))
+# WebSocket endpoint
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    session_id = str(uuid.uuid4())
+    server.active_connections[session_id] = websocket
+    server.active_sessions[session_id] = {
+        'start_time': time.time(),
+        'ops_count': 0
+    }
+    try:
+        while True:
+            message = await websocket.receive_json()
+            # Route operation to appropriate handler
+            operation_type = message.get('operation')
+            if operation_type == 'vram':
+                response = await server.handle_vram_operation(message)
+            elif operation_type == 'state':
+                response = await server.handle_state_operation(message)
+            elif operation_type == 'cache':
+                response = await server.handle_cache_operation(message)
+            else:
+                response = {
+                    'status': 'error',
+                    'message': 'Unknown operation type'
+                }
+            # Update statistics
+            server.ops_counter += 1
+            server.active_sessions[session_id]['ops_count'] += 1
+            # Send response
+            await websocket.send_json(response)
+    except Exception as e:
+        print(f"WebSocket error: {e}")
+    finally:
+        # Cleanup on disconnect
+        del server.active_connections[session_id]
+        del server.active_sessions[session_id]
 # For running directly (development)
 if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("server:app", host="0.0.0.0", port=7860, reload=True)