Spaces:
Starting
on
A100
Starting
on
A100
from fastapi import FastAPI, WebSocket, HTTPException, WebSocketDisconnect | |
from fastapi.responses import StreamingResponse, JSONResponse | |
from fastapi.middleware.cors import CORSMiddleware | |
from fastapi.staticfiles import StaticFiles | |
from fastapi import Request | |
import markdown2 | |
import logging | |
from config import config, Args | |
from connection_manager import ConnectionManager, ServerFullException | |
import uuid | |
import time | |
from types import SimpleNamespace | |
from util import pil_to_frame, bytes_to_pil, is_firefox, get_pipeline_class | |
from device import device, torch_dtype | |
import asyncio | |
import os | |
import time | |
import torch | |
THROTTLE = 1.0 / 120 | |
class App: | |
def __init__(self, config: Args, pipeline): | |
self.args = config | |
self.pipeline = pipeline | |
self.app = FastAPI() | |
self.conn_manager = ConnectionManager() | |
self.init_app() | |
def init_app(self): | |
self.app.add_middleware( | |
CORSMiddleware, | |
allow_origins=["*"], | |
allow_credentials=True, | |
allow_methods=["*"], | |
allow_headers=["*"], | |
) | |
async def websocket_endpoint(user_id: uuid.UUID, websocket: WebSocket): | |
try: | |
await self.conn_manager.connect( | |
user_id, websocket, self.args.max_queue_size | |
) | |
await handle_websocket_data(user_id) | |
except ServerFullException as e: | |
logging.error(f"Server Full: {e}") | |
finally: | |
await self.conn_manager.disconnect(user_id) | |
logging.info(f"User disconnected: {user_id}") | |
async def handle_websocket_data(user_id: uuid.UUID): | |
if not self.conn_manager.check_user(user_id): | |
return HTTPException(status_code=404, detail="User not found") | |
last_time = time.time() | |
try: | |
while True: | |
if ( | |
self.args.timeout > 0 | |
and time.time() - last_time > self.args.timeout | |
): | |
await self.conn_manager.send_json( | |
user_id, | |
{ | |
"status": "timeout", | |
"message": "Your session has ended", | |
}, | |
) | |
await self.conn_manager.disconnect(user_id) | |
return | |
data = await self.conn_manager.receive_json(user_id) | |
if data["status"] == "next_frame": | |
info = pipeline.Info() | |
params = await self.conn_manager.receive_json(user_id) | |
params = pipeline.InputParams(**params) | |
params = SimpleNamespace(**params.dict()) | |
if info.input_mode == "image": | |
image_data = await self.conn_manager.receive_bytes(user_id) | |
if len(image_data) == 0: | |
await self.conn_manager.send_json( | |
user_id, {"status": "send_frame"} | |
) | |
continue | |
params.image = bytes_to_pil(image_data) | |
await self.conn_manager.update_data(user_id, params) | |
await self.conn_manager.send_json(user_id, {"status": "wait"}) | |
except Exception as e: | |
logging.error(f"Websocket Error: {e}, {user_id} ") | |
await self.conn_manager.disconnect(user_id) | |
async def get_queue_size(): | |
queue_size = self.conn_manager.get_user_count() | |
return JSONResponse({"queue_size": queue_size}) | |
async def stream(user_id: uuid.UUID, request: Request): | |
try: | |
async def generate(): | |
last_params = SimpleNamespace() | |
while True: | |
last_time = time.time() | |
await self.conn_manager.send_json( | |
user_id, {"status": "send_frame"} | |
) | |
params = await self.conn_manager.get_latest_data(user_id) | |
if params.__dict__ == last_params.__dict__ or params is None: | |
await asyncio.sleep(THROTTLE) | |
continue | |
last_params = params | |
image = pipeline.predict(params) | |
if image is None: | |
continue | |
frame = pil_to_frame(image) | |
yield frame | |
# https://bugs.chromium.org/p/chromium/issues/detail?id=1250396 | |
if not is_firefox(request.headers["user-agent"]): | |
yield frame | |
if self.args.debug: | |
print(f"Time taken: {time.time() - last_time}") | |
return StreamingResponse( | |
generate(), | |
media_type="multipart/x-mixed-replace;boundary=frame", | |
headers={"Cache-Control": "no-cache"}, | |
) | |
except Exception as e: | |
logging.error(f"Streaming Error: {e}, {user_id} ") | |
return HTTPException(status_code=404, detail="User not found") | |
# route to setup frontend | |
async def settings(): | |
info_schema = pipeline.Info.schema() | |
info = pipeline.Info() | |
if info.page_content: | |
page_content = markdown2.markdown(info.page_content) | |
input_params = pipeline.InputParams.schema() | |
return JSONResponse( | |
{ | |
"info": info_schema, | |
"input_params": input_params, | |
"max_queue_size": self.args.max_queue_size, | |
"page_content": page_content if info.page_content else "", | |
} | |
) | |
if not os.path.exists("public"): | |
os.makedirs("public") | |
self.app.mount("/", StaticFiles(directory="public", html=True), name="public") | |
pipeline_class = get_pipeline_class(config.pipeline) | |
pipeline = pipeline_class(config, device, torch_dtype) | |
app = App(config, pipeline).app | |
if __name__ == "__main__": | |
import uvicorn | |
uvicorn.run( | |
"main:app", | |
host=config.host, | |
port=config.port, | |
reload=config.reload, | |
ssl_certfile=config.ssl_certfile, | |
ssl_keyfile=config.ssl_keyfile, | |
) | |