vlm_server.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960
  1. import click
  2. import sys
  3. from loguru import logger
  4. def vllm_server():
  5. from mineru.model.vlm.vllm_server import main
  6. main()
  7. def lmdeploy_server():
  8. from mineru.model.vlm.lmdeploy_server import main
  9. main()
  10. @click.command(context_settings=dict(ignore_unknown_options=True, allow_extra_args=True))
  11. @click.option(
  12. '-e',
  13. '--engine',
  14. 'inference_engine',
  15. type=click.Choice(['auto', 'vllm', 'lmdeploy']),
  16. default='auto',
  17. help='Select the inference engine used to accelerate VLM inference, default is "auto".',
  18. )
  19. @click.pass_context
  20. def openai_server(ctx, inference_engine):
  21. sys.argv = [sys.argv[0]] + ctx.args
  22. if inference_engine == 'auto':
  23. try:
  24. import vllm
  25. inference_engine = 'vllm'
  26. logger.info("Using vLLM as the inference engine for VLM server.")
  27. except ImportError:
  28. logger.info("vLLM not found, falling back to LMDeploy as the inference engine for VLM server.")
  29. try:
  30. import lmdeploy
  31. inference_engine = 'lmdeploy'
  32. logger.info("Using LMDeploy as the inference engine for VLM server.")
  33. except ImportError:
  34. logger.error("Neither vLLM nor LMDeploy is installed. Please install at least one of them.")
  35. sys.exit(1)
  36. if inference_engine == 'vllm':
  37. try:
  38. import vllm
  39. except ImportError:
  40. logger.error("vLLM is not installed. Please install vLLM or choose LMDeploy as the inference engine.")
  41. sys.exit(1)
  42. vllm_server()
  43. elif inference_engine == 'lmdeploy':
  44. try:
  45. import lmdeploy
  46. except ImportError:
  47. logger.error("LMDeploy is not installed. Please install LMDeploy or choose vLLM as the inference engine.")
  48. sys.exit(1)
  49. lmdeploy_server()
  50. if __name__ == "__main__":
  51. openai_server()