OpenAI Chat Client for Multimodal#
Refer to the trtllm-serve documentation for starting a server.
Source NVIDIA/TensorRT-LLM.
1
2import os
3from pathlib import Path
4
5from openai import OpenAI
6from PIL import Image
7
8from tensorrt_llm.inputs import (encode_base64_content_from_url,
9 encode_base64_image)
10
11client = OpenAI(
12 base_url="http://localhost:8000/v1",
13 api_key="tensorrt_llm",
14)
15
16llm_models_root = Path(os.environ.get("LLM_MODELS_ROOT"))
17
18if llm_models_root is not None:
19 # These are paths on the server's own filesystem. trtllm-serve refuses
20 # local paths unless it is started with
21 # TRTLLM_MEDIA_ALLOWED_LOCAL_PATH=$LLM_MODELS_ROOT/multimodals/test_data
22 # (or another directory that contains the files).
23 multimodal_test_data_path = llm_models_root / "multimodals" / "test_data"
24 image_url1 = str(multimodal_test_data_path / "seashore.png")
25 image_url2 = str(multimodal_test_data_path / "inpaint.png")
26 video_url = str(multimodal_test_data_path / "OAI-sora-tokyo-walk.mp4")
27 image64 = encode_base64_image(
28 Image.open(multimodal_test_data_path / "seashore.png"))
29else:
30 image_url1 = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore.png"
31 image_url2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png"
32 video_url = "https://huggingface.co/datasets/Efficient-Large-Model/VILA-inference-demos/resolve/main/OAI-sora-tokyo-walk.mp4"
33 image64 = encode_base64_content_from_url(
34 "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore.png"
35 )
36
37# SINGLE IMAGE INFERENCE
38response = client.chat.completions.create(
39 model="Qwen2.5-VL-3B-Instruct",
40 messages=[{
41 "role": "system",
42 "content": "you are a helpful assistant"
43 }, {
44 "role":
45 "user",
46 "content": [{
47 "type": "text",
48 "text": "Describe the natural environment in the image."
49 }, {
50 "type": "image_url",
51 "image_url": {
52 "url": image_url1
53 }
54 }]
55 }],
56 max_tokens=64,
57)
58print(response)
59
60# MULTI IMAGE INFERENCE
61response = client.chat.completions.create(
62 model="Qwen2.5-VL-3B-Instruct",
63 messages=[{
64 "role": "system",
65 "content": "you are a helpful assistant"
66 }, {
67 "role":
68 "user",
69 "content": [{
70 "type": "text",
71 "text": "Tell me the difference between two images"
72 }, {
73 "type": "image_url",
74 "image_url": {
75 "url": image_url2
76 }
77 }, {
78 "type": "image_url",
79 "image_url": {
80 "url": image_url1
81 }
82 }]
83 }],
84 max_tokens=64,
85)
86print(response)
87
88# SINGLE VIDEO INFERENCE
89response = client.chat.completions.create(
90 model="Qwen2.5-VL-3B-Instruct",
91 messages=[{
92 "role": "system",
93 "content": "you are a helpful assistant"
94 }, {
95 "role":
96 "user",
97 "content": [{
98 "type": "text",
99 "text": "Tell me what you see in the video briefly."
100 }, {
101 "type": "video_url",
102 "video_url": {
103 "url": video_url
104 }
105 }]
106 }],
107 max_tokens=64,
108)
109print(response)
110
111# IMAGE EMBED INFERENCE
112response = client.chat.completions.create(
113 model="Qwen2.5-VL-3B-Instruct",
114 messages=[{
115 "role": "system",
116 "content": "you are a helpful assistant"
117 }, {
118 "role":
119 "user",
120 "content": [{
121 "type": "text",
122 "text": "Describe the natural environment in the image."
123 }, {
124 "type": "image_url",
125 "image_url": {
126 "url": "data:image/png;base64," + image64
127 }
128 }]
129 }],
130 max_tokens=64,
131)
132print(response)