diff --git a/dots_ocr/model/inference.py b/dots_ocr/model/inference.py index 80f3395..1140107 100755 --- a/dots_ocr/model/inference.py +++ b/dots_ocr/model/inference.py @@ -12,6 +12,7 @@ import os def inference_with_vllm( image, prompt, + protocol="http", ip="localhost", port=8000, temperature=0.1, @@ -20,7 +21,7 @@ def inference_with_vllm( model_name='model', ): - addr = f"http://{ip}:{port}/v1" + addr = f"{protocol}://{ip}:{port}/v1" client = OpenAI(api_key="{}".format(os.environ.get("API_KEY", "0")), base_url=addr) messages = [] messages.append( diff --git a/dots_ocr/parser.py b/dots_ocr/parser.py index 54eaf49..18acdad 100755 --- a/dots_ocr/parser.py +++ b/dots_ocr/parser.py @@ -20,6 +20,7 @@ class DotsOCRParser: """ def __init__(self, + protocol='http', ip='localhost', port=8000, model_name='model', @@ -36,6 +37,7 @@ class DotsOCRParser: self.dpi = dpi # default args for vllm server + self.protocol = protocol self.ip = ip self.port = port self.model_name = model_name @@ -119,6 +121,7 @@ class DotsOCRParser: image, prompt, model_name=self.model_name, + protocol=self.protocol, ip=self.ip, port=self.port, temperature=self.temperature, @@ -347,6 +350,10 @@ def main(): metavar=('x1', 'y1', 'x2', 'y2'), help='should give this argument if you want to prompt_grounding_ocr' ) + parser.add_argument( + "--protocol", type=str, choices=['http', 'https'], default="http", + help="" + ) parser.add_argument( "--ip", type=str, default="localhost", help="" @@ -398,6 +405,7 @@ def main(): args = parser.parse_args() dots_ocr_parser = DotsOCRParser( + protocol=args.protocol, ip=args.ip, port=args.port, model_name=args.model_name,