#gradio_openai_vlm_webserver.py
import argparse
import base64
import gradio as gr
from openai import OpenAI

parser = argparse.ArgumentParser(description='GLM-4V Chatbot Interface')
parser.add_argument('--model-url', type=str, default='http://localhost:8000/v1', help='Model URL')
parser.add_argument('-m', '--model', type=str, required=True, help='Model path')
parser.add_argument("--host", type=str, default=None)
parser.add_argument("--port", type=int, default=8001)

args = parser.parse_args()

openai_api_key = "EMPTY"
openai_api_base = args.model_url

client = OpenAI(api_key=openai_api_key, base_url=openai_api_base)

def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def predict(message, history, image):
    history_openai_format = []
    
    if image:
        base64_image = encode_image(image)
        content = [
            {"type": "text", "text": message},
            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
        ]
        history_openai_format.append({"role": "user", "content": content})
    else:
        history_openai_format.append({"role": "user", "content": message})
    
    try:
        chat_response = client.chat.completions.create(
            model=args.model,
            messages=history_openai_format
        )
        return chat_response.choices[0].message.content
    except Exception as e:
        return f"An error occurred: {str(e)}"
gr.ChatInterface(
    predict,
    additional_inputs=[
        gr.Image(label="Upload Image", type="filepath",scale=1, width="50%")
    ],
    title="GLM-4V Chatbot",
    description="Chat with the GLM-4V model and upload images for visual understanding."
).queue().launch(server_name=args.host, server_port=args.port, share=True)