File size: 1,331 Bytes
b4942cf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import argparse
import os.path
import gradio as gr
from gradio.components import Textbox, Image
from ovis.serve.runner import RunnerArguments, OvisRunner
class Server:
def __init__(self, runner: OvisRunner):
self.runner = runner
def __call__(self, image, text):
response = self.runner.run([image, text])
output = response["output"]
return output
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Ovis Server')
parser.add_argument('--model_path', type=str, required=True)
parser.add_argument('--flagging_dir', type=str, default=os.path.expanduser('~/ovis-flagged'))
parser.add_argument('--max_partition', type=int, default=9)
parser.add_argument('--port', type=int, required=True)
args = parser.parse_args()
os.makedirs(args.flagging_dir, exist_ok=True)
runner_args = RunnerArguments(
model_path=args.model_path,
max_partition=args.max_partition
)
demo = gr.Interface(
fn=Server(OvisRunner(runner_args)),
inputs=[Image(type='pil', label='image'),
Textbox(placeholder='Enter your text here...', label='prompt')],
outputs=gr.Markdown(),
title=args.model_path.split('/')[-1],
flagging_dir=args.flagging_dir
)
demo.launch(server_port=args.port)
|