__init__.py 1.7 KB

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