anyantudre commited on
Commit
2ff6071
·
verified ·
1 Parent(s): fd6a0d8

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +50 -0
app.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """expose an endpoint that accepts image_url and prompt, runs inference with Qwen2.5-VL, and returns a text response"""
2
+
3
+ from fastapi import FastAPI, Query
4
+ from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
5
+ from qwen_vl_utils import process_vision_info
6
+ import torch
7
+
8
+ app = FastAPI()
9
+
10
+ checkpoint = "Qwen/Qwen2.5-VL-3B-Instruct"
11
+ min_pixels = 256*28*28
12
+ max_pixels = 1280*28*28
13
+ processor = AutoProcessor.from_pretrained(
14
+ checkpoint,
15
+ min_pixels=min_pixels,
16
+ max_pixels=max_pixels
17
+ )
18
+ model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
19
+ checkpoint,
20
+ torch_dtype=torch.bfloat16,
21
+ device_map="auto",
22
+ # attn_implementation="flash_attention_2",
23
+ )
24
+
25
+ @app.get("/")
26
+ def read_root():
27
+ return {"message": "API is live. Use the /predict endpoint."}
28
+
29
+ @app.get("/predict")
30
+ def predict(image_url: str = Query(...), prompt: str = Query(...)):
31
+ messages = [
32
+ {"role": "system", "content": "You are a helpful assistant with vision abilities."},
33
+ {"role": "user", "content": [{"type": "image", "image": image_url}, {"type": "text", "text": prompt}]},
34
+ ]
35
+ text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
36
+ image_inputs, video_inputs = process_vision_info(messages)
37
+ inputs = processor(
38
+ text=[text],
39
+ images=image_inputs,
40
+ videos=video_inputs,
41
+ padding=True,
42
+ return_tensors="pt",
43
+ ).to(model.device)
44
+ with torch.no_grad():
45
+ generated_ids = model.generate(**inputs, max_new_tokens=128)
46
+ generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
47
+ output_texts = processor.batch_decode(
48
+ generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
49
+ )
50
+ return {"response": output_texts[0]}