Spaces:
Running
Running
File size: 3,470 Bytes
2a4f1d3 9559a52 0410591 2a4f1d3 9559a52 2a4f1d3 9086f73 9559a52 0410591 2a4f1d3 0410591 2a4f1d3 0410591 2a4f1d3 0410591 2a4f1d3 0410591 9559a52 0410591 9559a52 35f3879 9559a52 2a4f1d3 35f3879 2a4f1d3 9559a52 2a4f1d3 9559a52 2a4f1d3 9559a52 2a4f1d3 9559a52 2a4f1d3 0f6bee9 9559a52 0f6bee9 9559a52 0f6bee9 9559a52 0f6bee9 9559a52 0f6bee9 9559a52 0f6bee9 9559a52 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
from fastapi import FastAPI, Query
from transformers import (
AutoProcessor,
AutoModelForCausalLM,
AutoTokenizer,
)
from transformers import Qwen2_5_VLForConditionalGeneration
from qwen_vl_utils import process_vision_info
import torch
import logging
logging.basicConfig(level=logging.INFO)
app = FastAPI()
checkpoint = "Qwen/Qwen2.5-VL-3B-Instruct"
min_pixels = 256*28*28
max_pixels = 1280*28*28
processor = AutoProcessor.from_pretrained(
checkpoint,
min_pixels=min_pixels,
max_pixels=max_pixels
)
qwen_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
checkpoint,
torch_dtype=torch.bfloat16,
device_map="auto",
# attn_implementation="flash_attention_2",
)
# LLaMA Model Setup
llama_model_name = "Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2"
llama_tokenizer = AutoTokenizer.from_pretrained(llama_model_name)
llama_model = AutoModelForCausalLM.from_pretrained(
llama_model_name, torch_dtype=torch.float16, device_map="auto"
)
@app.get("/")
def read_root():
return {"message": "API is live. Use the /predict, /chat, or /llama_chat endpoints."}
@app.get("/predict")
def predict(image_url: str = Query(...), prompt: str = Query(...)):
messages = [
{"role": "system", "content": "You are a helpful assistant with vision abilities."},
{
"role": "user",
"content": [
{"type": "image", "image": image_url},
{"type": "text", "text": prompt},
],
},
]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
image_inputs, video_inputs = process_vision_info(messages)
inputs = processor(
text=[text],
images=image_inputs,
videos=video_inputs,
padding=True,
return_tensors="pt",
).to(qwen_model.device)
with torch.no_grad():
generated_ids = qwen_model.generate(**inputs, max_new_tokens=128)
generated_ids_trimmed = [
out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs["input_ids"], generated_ids)
]
output_texts = processor.batch_decode(
generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
return {"response": output_texts[0]}
@app.get("/chat")
def chat(prompt: str = Query(...)):
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": [{"type": "text", "text": prompt}]},
]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = processor(
text=[text],
padding=True,
return_tensors="pt",
).to(qwen_model.device)
with torch.no_grad():
generated_ids = qwen_model.generate(**inputs, max_new_tokens=128)
generated_ids_trimmed = [
out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs["input_ids"], generated_ids)
]
output_texts = processor.batch_decode(
generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
return {"response": output_texts[0]}
@app.get("/llama_chat")
def llama_chat(prompt: str = Query(...)):
inputs = llama_tokenizer(prompt, return_tensors="pt").to(llama_model.device)
with torch.no_grad():
outputs = llama_model.generate(**inputs, max_new_tokens=128)
response = llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
return {"response": response}
|