Spaces:
Sleeping
Sleeping
File size: 3,875 Bytes
2a4f1d3 9559a52 0410591 2a4f1d3 9559a52 2a4f1d3 9559a52 0410591 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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
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()
# 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",
# )
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}
|