Samet Yilmaz commited on
Commit
6cae924
1 Parent(s): f9fa47c

Add pillow

Browse files
Files changed (2) hide show
  1. app.py +22 -1
  2. requirements.txt +1 -0
app.py CHANGED
@@ -1,20 +1,41 @@
1
  from vllm import LLM, SamplingParams
2
  import gradio as gr
 
 
 
 
 
 
3
  repo_id = "mistral-community/pixtral-12b-240910" #Replace to the model you would like to use
4
  sampling_params = SamplingParams(max_tokens=8192, temperature=0.7)
5
  max_tokens_per_img = 4096
6
  max_img_per_msg = 5
7
 
 
 
 
 
 
 
 
 
8
  # @spaces.GPU #[uncomment to use ZeroGPU]
9
  def infer(image_url, prompt, progress=gr.Progress(track_tqdm=True)):
10
  # tokenize image urls and text
11
  llm = LLM(model="mistralai/Pixtral-12B-2409",
12
  tokenizer_mode="mistral",
 
 
13
  limit_mm_per_prompt={"image": max_img_per_msg}) # Name or path of your model
 
 
 
 
 
14
  messages = [
15
  {
16
  "role": "user",
17
- "content": [{"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": image_url}}]
18
  },
19
  ]
20
 
 
1
  from vllm import LLM, SamplingParams
2
  import gradio as gr
3
+
4
+ from PIL import Image
5
+ from io import BytesIO
6
+ import base64
7
+ import requests
8
+
9
  repo_id = "mistral-community/pixtral-12b-240910" #Replace to the model you would like to use
10
  sampling_params = SamplingParams(max_tokens=8192, temperature=0.7)
11
  max_tokens_per_img = 4096
12
  max_img_per_msg = 5
13
 
14
+ def encode_image(image: Image.Image, image_format="PNG") -> str:
15
+ im_file = BytesIO()
16
+ image.save(im_file, format=image_format)
17
+ im_bytes = im_file.getvalue()
18
+ im_64 = base64.b64encode(im_bytes).decode("utf-8")
19
+ return im_64
20
+
21
+
22
  # @spaces.GPU #[uncomment to use ZeroGPU]
23
  def infer(image_url, prompt, progress=gr.Progress(track_tqdm=True)):
24
  # tokenize image urls and text
25
  llm = LLM(model="mistralai/Pixtral-12B-2409",
26
  tokenizer_mode="mistral",
27
+ max_model_len=65536,
28
+ max_num_batched_tokens=max_img_per_msg * max_tokens_per_img,
29
  limit_mm_per_prompt={"image": max_img_per_msg}) # Name or path of your model
30
+
31
+ image = Image.open(BytesIO(requests.get(image_url).content))
32
+ image = image.resize((3844, 2408))
33
+ new_image_url = f"data:image/png;base64,{encode_image(image, image_format='PNG')}"
34
+
35
  messages = [
36
  {
37
  "role": "user",
38
+ "content": [{"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": new_image_url}}]
39
  },
40
  ]
41
 
requirements.txt CHANGED
@@ -1 +1,2 @@
1
  vllm==0.6.1
 
 
1
  vllm==0.6.1
2
+ Pillow