Spaces:
Sleeping
Sleeping
File size: 3,394 Bytes
2a4f1d3 9559a52 2a4f1d3 9559a52 2a4f1d3 9559a52 35f3879 9559a52 2a4f1d3 9559a52 2a4f1d3 9559a52 2a4f1d3 9559a52 2a4f1d3 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 |
from fastapi import FastAPI, Query
from transformers import (
AutoProcessor,
AutoModelForCausalLM,
AutoTokenizer,
)
from qwen_vl_utils import process_vision_info
import torch
import logging
logging.basicConfig(level=logging.INFO)
app = FastAPI()
# Qwen2.5-VL Model Setup
qwen_checkpoint = "Qwen/Qwen2.5-VL-7B-Instruct"
min_pixels = 256 * 28 * 28
max_pixels = 1280 * 28 * 28
processor = AutoProcessor.from_pretrained(
qwen_checkpoint,
min_pixels=min_pixels,
max_pixels=max_pixels,
)
qwen_model = AutoModelForCausalLM.from_pretrained(
qwen_checkpoint,
torch_dtype=torch.bfloat16,
device_map="auto",
)
# 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}
|