Offline Inference Pixtral#

Source vllm-project/vllm.

  1# ruff: noqa
  2import argparse
  3
  4from vllm import LLM
  5from vllm.sampling_params import SamplingParams
  6
  7# This script is an offline demo for running Pixtral.
  8#
  9# If you want to run a server/client setup, please follow this code:
 10#
 11# - Server:
 12#
 13# ```bash
 14# vllm serve mistralai/Pixtral-12B-2409 --tokenizer_mode mistral --limit_mm_per_prompt 'image=4' --max_num_batched_tokens 16384
 15# ```
 16#
 17# - Client:
 18#
 19# ```bash
 20# curl --location 'http://<your-node-url>:8000/v1/chat/completions' \
 21# --header 'Content-Type: application/json' \
 22# --header 'Authorization: Bearer token' \
 23# --data '{
 24#     "model": "mistralai/Pixtral-12B-2409",
 25#     "messages": [
 26#       {
 27#         "role": "user",
 28#         "content": [
 29#             {"type" : "text", "text": "Describe this image in detail please."},
 30#             {"type": "image_url", "image_url": {"url": "https://s3.amazonaws.com/cms.ipressroom.com/338/files/201808/5b894ee1a138352221103195_A680%7Ejogging-edit/A680%7Ejogging-edit_hero.jpg"}},
 31#             {"type" : "text", "text": "and this one as well. Answer in French."},
 32#             {"type": "image_url", "image_url": {"url": "https://www.wolframcloud.com/obj/resourcesystem/images/a0e/a0ee3983-46c6-4c92-b85d-059044639928/6af8cfb971db031b.png"}}
 33#         ]
 34#       }
 35#     ]
 36#   }'
 37# ```
 38#
 39# Usage:
 40#     python demo.py simple
 41#     python demo.py advanced
 42
 43
 44def run_simple_demo():
 45    model_name = "mistralai/Pixtral-12B-2409"
 46    sampling_params = SamplingParams(max_tokens=8192)
 47
 48    llm = LLM(model=model_name, tokenizer_mode="mistral")
 49
 50    prompt = "Describe this image in one sentence."
 51    image_url = "https://picsum.photos/id/237/200/300"
 52
 53    messages = [
 54        {
 55            "role":
 56            "user",
 57            "content": [
 58                {
 59                    "type": "text",
 60                    "text": prompt
 61                },
 62                {
 63                    "type": "image_url",
 64                    "image_url": {
 65                        "url": image_url
 66                    }
 67                },
 68            ],
 69        },
 70    ]
 71    outputs = llm.chat(messages, sampling_params=sampling_params)
 72
 73    print(outputs[0].outputs[0].text)
 74
 75
 76def run_advanced_demo():
 77    model_name = "mistralai/Pixtral-12B-2409"
 78    max_img_per_msg = 5
 79    max_tokens_per_img = 4096
 80
 81    sampling_params = SamplingParams(max_tokens=8192, temperature=0.7)
 82    llm = LLM(
 83        model=model_name,
 84        tokenizer_mode="mistral",
 85        limit_mm_per_prompt={"image": max_img_per_msg},
 86        max_num_batched_tokens=max_img_per_msg * max_tokens_per_img,
 87    )
 88
 89    prompt = "Describe the following image."
 90
 91    url_1 = "https://huggingface.co/datasets/patrickvonplaten/random_img/resolve/main/yosemite.png"
 92    url_2 = "https://picsum.photos/seed/picsum/200/300"
 93    url_3 = "https://picsum.photos/id/32/512/512"
 94
 95    messages = [
 96        {
 97            "role":
 98            "user",
 99            "content": [
100                {
101                    "type": "text",
102                    "text": prompt
103                },
104                {
105                    "type": "image_url",
106                    "image_url": {
107                        "url": url_1
108                    }
109                },
110                {
111                    "type": "image_url",
112                    "image_url": {
113                        "url": url_2
114                    }
115                },
116            ],
117        },
118        {
119            "role": "assistant",
120            "content": "The images show nature.",
121        },
122        {
123            "role": "user",
124            "content": "More details please and answer only in French!.",
125        },
126        {
127            "role": "user",
128            "content": [
129                {
130                    "type": "image_url",
131                    "image_url": {
132                        "url": url_3
133                    }
134                },
135            ],
136        },
137    ]
138
139    outputs = llm.chat(messages=messages, sampling_params=sampling_params)
140    print(outputs[0].outputs[0].text)
141
142
143def main():
144    parser = argparse.ArgumentParser(
145        description="Run a demo in simple or advanced mode.")
146
147    parser.add_argument(
148        "mode",
149        choices=["simple", "advanced"],
150        help="Specify the demo mode: 'simple' or 'advanced'",
151    )
152
153    args = parser.parse_args()
154
155    if args.mode == "simple":
156        print("Running simple demo...")
157        run_simple_demo()
158    elif args.mode == "advanced":
159        print("Running advanced demo...")
160        run_advanced_demo()
161
162
163if __name__ == "__main__":
164    main()