iridescentX/openui

1import base642import os3import random4from pathlib import Path5 6import typer7import uvicorn8 9app = typer.Typer()10 11KEY_FILE = Path.cwd() / ".webui_secret_key"12if (frontend_build_dir := Path(__file__).parent / "frontend").exists():13    os.environ["FRONTEND_BUILD_DIR"] = str(frontend_build_dir)14 15 16@app.command()17def serve(18    host: str = "0.0.0.0",19    port: int = 8080,20):21    if os.getenv("WEBUI_SECRET_KEY") is None:22        typer.echo(23            "Loading WEBUI_SECRET_KEY from file, not provided as an environment variable."24        )25        if not KEY_FILE.exists():26            typer.echo(f"Generating a new secret key and saving it to {KEY_FILE}")27            KEY_FILE.write_bytes(base64.b64encode(random.randbytes(12)))28        typer.echo(f"Loading WEBUI_SECRET_KEY from {KEY_FILE}")29        os.environ["WEBUI_SECRET_KEY"] = KEY_FILE.read_text()30 31    if os.getenv("USE_CUDA_DOCKER", "false") == "true":32        typer.echo(33            "CUDA is enabled, appending LD_LIBRARY_PATH to include torch/cudnn & cublas libraries."34        )35        LD_LIBRARY_PATH = os.getenv("LD_LIBRARY_PATH", "").split(":")36        os.environ["LD_LIBRARY_PATH"] = ":".join(37            LD_LIBRARY_PATH38            + [39                "/usr/local/lib/python3.11/site-packages/torch/lib",40                "/usr/local/lib/python3.11/site-packages/nvidia/cudnn/lib",41            ]42        )43    import main  # we need set environment variables before importing main44 45    uvicorn.run(main.app, host=host, port=port, forwarded_allow_ips="*")46 47 48@app.command()49def dev(50    host: str = "0.0.0.0",51    port: int = 8080,52    reload: bool = True,53):54    uvicorn.run(55        "main:app", host=host, port=port, reload=reload, forwarded_allow_ips="*"56    )57 58 59if __name__ == "__main__":60    app()61