vlm_server.py 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  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, attempting to use LMDeploy as the inference engine for VLM server.")
  29. try:
  30. import lmdeploy
  31. inference_engine = 'lmdeploy'
  32. # Success message moved after successful import
  33. logger.info("Using LMDeploy as the inference engine for VLM server.")
  34. except ImportError:
  35. logger.error("Neither vLLM nor LMDeploy is installed. Please install at least one of them.")
  36. sys.exit(1)
  37. if inference_engine == 'vllm':
  38. try:
  39. import vllm
  40. except ImportError:
  41. logger.error("vLLM is not installed. Please install vLLM or choose LMDeploy as the inference engine.")
  42. sys.exit(1)
  43. vllm_server()
  44. elif inference_engine == 'lmdeploy':
  45. try:
  46. import lmdeploy
  47. except ImportError:
  48. logger.error("LMDeploy is not installed. Please install LMDeploy or choose vLLM as the inference engine.")
  49. sys.exit(1)
  50. lmdeploy_server()
  51. if __name__ == "__main__":
  52. openai_server()