1234567891011121314151617181920212223242526272829303132333435363738394041424344454647 |
- import base64
- import os
- import random
- from pathlib import Path
- import typer
- import uvicorn
- import main
- app = typer.Typer()
- KEY_FILE = Path.cwd() / ".webui_secret_key"
- @app.command()
- def serve(
- host: str = "0.0.0.0",
- port: int = 8080,
- ):
- if os.getenv("WEBUI_SECRET_KEY") is None:
- typer.echo(
- "Loading WEBUI_SECRET_KEY from file, not provided as an environment variable."
- )
- if not KEY_FILE.exists():
- typer.echo(f"Generating a new secret key and saving it to {KEY_FILE}")
- KEY_FILE.write_bytes(base64.b64encode(random.randbytes(12)))
- typer.echo(f"Loading WEBUI_SECRET_KEY from {KEY_FILE}")
- os.environ["WEBUI_SECRET_KEY"] = KEY_FILE.read_text()
- if os.getenv("USE_CUDA_DOCKER", "false") == "true":
- typer.echo(
- "CUDA is enabled, appending LD_LIBRARY_PATH to include torch/cudnn & cublas libraries."
- )
- LD_LIBRARY_PATH = os.getenv("LD_LIBRARY_PATH", "").split(":")
- os.environ["LD_LIBRARY_PATH"] = ":".join(
- LD_LIBRARY_PATH
- + [
- "/usr/local/lib/python3.11/site-packages/torch/lib",
- "/usr/local/lib/python3.11/site-packages/nvidia/cudnn/lib",
- ]
- )
- uvicorn.run(main.app, host=host, port=port, forwarded_allow_ips="*")
- if __name__ == "__main__":
- app()
|