Spaces:
Build error
Build error
File size: 4,328 Bytes
7c66cbc dc7242a 7c66cbc dc7242a 7c66cbc dc7242a 7c66cbc e1e7a87 7c66cbc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
from typing import Annotated, Any, Generator
from pathlib import Path
from gymnasium.wrappers.record_video import RecordVideo
from litrl.env.make import make
from litrl.common.agent import RandomAgent
from litrl.env.typing import SingleAgentId
from fastapi import Depends, FastAPI, Request, status
from fastapi.exceptions import RequestValidationError
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from litrl.env.connect_four import Board
from loguru import logger
from fastapi.responses import StreamingResponse, RedirectResponse
from src.app_state import AppState
from src.typing import CpuConfig
from src.huggingface.huggingface_client import HuggingFaceClient
def stream_mp4(mp4_path: Path) -> StreamingResponse:
def iter_file()-> Generator[bytes, Any, None]:
with mp4_path.open(mode="rb") as env_file:
yield from env_file
return StreamingResponse(content=iter_file(), media_type="video/mp4")
def create_app() -> FastAPI:
app = FastAPI()
@app.get('/')
async def to_docs():
return RedirectResponse("/docs")
@app.post("/", response_model=int)
def bot_action(
board: Board,
cpuConfig: CpuConfig,
app_state: Annotated[AppState, Depends(dependency=AppState)],
) -> int:
app_state.set_config(cpu_config=cpuConfig)
app_state.set_board(board=board)
return app_state.get_action()
@app.post(path=f"/game", response_model=str)
def bot_action(
env_id: SingleAgentId,
) -> str:
env = RecordVideo(
env=make(id=env_id, render_mode="rgb_array"),
video_folder="tmp",
)
env.reset(seed=123)
agent = RandomAgent[Any, Any]()
terminated, truncated = False, False
while not (terminated or truncated):
action = agent.get_action(env=env)
_, _, terminated, truncated, _ = env.step(action=action)
env.render()
env.video_recorder.close()
return stream_mp4(mp4_path=Path(env.video_recorder.path))
@app.get(path=f"/hfmp4")
def fh_stream(
env_id: SingleAgentId,
hf_client: Annotated[HuggingFaceClient, Depends(dependency=HuggingFaceClient)],
) -> StreamingResponse:
hf_client.mp4_paths[env_id]
return stream_mp4(mp4_path=hf_client.mp4_paths[env_id])
@app.get(path=f"/mp4")
def bot_action(
env_id: SingleAgentId,
) -> StreamingResponse:
env = make(id=env_id, render_mode="rgb_array")
env = RecordVideo(
env=env,
video_folder="tmp",
)
env.reset(seed=123)
agent = RandomAgent[Any, Any]()
terminated, truncated = False, False
while not (terminated or truncated):
action = agent.get_action(env=env)
_, _, terminated, truncated, _ = env.step(action=action)
env.render()
env.video_recorder.close()
return stream_mp4(mp4_path=Path(env.video_recorder.path))
@app.exception_handler(exc_class_or_status_code=RequestValidationError)
async def validation_exception_handler(
request: Request, exc: RequestValidationError
) -> JSONResponse:
logger.debug(f"url: {request.url}")
if hasattr(request, "_body"):
logger.debug(f"body: {request._body}")
logger.debug(f"header: {request.headers}")
logger.error(f"{request}: {exc}")
exc_str = f"{exc}".replace("\n", " ").replace(" ", " ")
content = {"status_code": 10422, "message": exc_str, "data": None}
return JSONResponse(
content=content, status_code=status.HTTP_422_UNPROCESSABLE_ENTITY
)
app.add_middleware(
middleware_class=CORSMiddleware,
allow_origins="*",
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
return app
if __name__ == "__main__":
import uvicorn
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--host", type=str, default="0.0.0.0")
parser.add_argument("--port", type=int, default=7860)
args = parser.parse_args()
config = uvicorn.Config(app=create_app(), host=args.host, port=args.port, log_level="info")
server = uvicorn.Server(config=config)
server.run() |