|
from asyncio import gather |
|
from base64 import b64decode |
|
from binascii import Error as BinasciiError |
|
from contextlib import asynccontextmanager |
|
from io import BytesIO |
|
from logging import Formatter, INFO, StreamHandler, getLogger |
|
from pathlib import Path |
|
from random import randint |
|
from typing import AsyncGenerator |
|
from uuid import UUID |
|
|
|
from PIL.Image import open as image_open |
|
from fastapi import FastAPI, Request |
|
from fastapi.responses import HTMLResponse, JSONResponse |
|
from httpx import AsyncClient |
|
from starlette.responses import Response |
|
|
|
logger = getLogger('NVIDIA_VLM_API') |
|
logger.setLevel(INFO) |
|
handler = StreamHandler() |
|
handler.setLevel(INFO) |
|
formatter = Formatter('%(asctime)s | %(levelname)s : %(message)s', datefmt='%d.%m.%Y %H:%M:%S') |
|
handler.setFormatter(formatter) |
|
logger.addHandler(handler) |
|
|
|
logger.info('инициализация приложения...') |
|
|
|
INVOKE_URLS = [ |
|
'https://ai.api.nvidia.com/v1/vlm/microsoft/phi-3-vision-128k-instruct', |
|
'https://ai.api.nvidia.com/v1/vlm/nvidia/neva-22b', |
|
'https://ai.api.nvidia.com/v1/vlm/nvidia/vila' |
|
] |
|
|
|
ASSETS_URL = 'https://api.nvcf.nvidia.com/v2/nvcf/assets' |
|
|
|
|
|
def get_extension(filename: str) -> str: |
|
return Path(filename).suffix[1:].lower() |
|
|
|
|
|
async def upload_asset(client: AsyncClient, media_file_bytes: bytes, description: str, api_key: str) -> UUID: |
|
headers = { |
|
'Authorization': f'Bearer {api_key}', |
|
'Content-Type': 'application/json', |
|
'accept': 'application/json', |
|
} |
|
|
|
authorize_response = await client.post( |
|
ASSETS_URL, |
|
headers=headers, |
|
json={'contentType': 'image/jpeg', 'description': description}, |
|
timeout=30, |
|
) |
|
authorize_response.raise_for_status() |
|
authorize_res = authorize_response.json() |
|
response = await client.put( |
|
authorize_res.get('uploadUrl'), |
|
content=media_file_bytes, |
|
headers={'x-amz-meta-nvcf-asset-description': description, 'content-type': 'image/jpeg'}, |
|
timeout=300, |
|
) |
|
|
|
response.raise_for_status() |
|
return UUID(authorize_res.get('assetId')) |
|
|
|
|
|
async def delete_asset(client: AsyncClient, asset_id: UUID, api_key: str) -> None: |
|
headers = {'Authorization': f'Bearer {api_key}'} |
|
response = await client.delete(f'{ASSETS_URL}/{asset_id}', headers=headers, timeout=30) |
|
response.raise_for_status() |
|
|
|
|
|
async def chat_with_media_nvcf(infer_url: str, media_file_bytes: bytes, query: str, api_key: str) -> str | None: |
|
try: |
|
async with AsyncClient(follow_redirects=True, timeout=45) as client: |
|
asset_list = [] |
|
asset_id = await upload_asset(client, media_file_bytes, 'Reference media file', api_key) |
|
asset_list.append(str(asset_id)) |
|
media_content = f'<img src="data:image/jpeg;asset_id,{asset_id}" />' |
|
asset_seq = ','.join(asset_list) |
|
headers = { |
|
'Authorization': f'Bearer {api_key}', |
|
'Content-Type': 'application/json', |
|
'NVCF-INPUT-ASSET-REFERENCES': asset_seq, |
|
'NVCF-FUNCTION-ASSET-IDS': asset_seq, |
|
'Accept': 'application/json', |
|
} |
|
|
|
payload = { |
|
'max_tokens': 1024, |
|
'temperature': 0.65, |
|
'top_p': 0.95, |
|
'seed': randint(0, 999999999), |
|
'messages': [{'role': 'user', 'content': f'{query} {media_content}'}], |
|
'stream': False, |
|
"model": infer_url.split('/v1/vlm/')[-1] |
|
} |
|
|
|
response = await client.post(infer_url, headers=headers, json=payload) |
|
response_json = response.json() |
|
answer = response_json.get('choices', [{}])[0].get('message', {}).get('content', None) |
|
for asset_id in asset_list: |
|
await delete_asset(client, UUID(asset_id), api_key) |
|
return answer |
|
except Exception as exc: |
|
print(exc) |
|
return None |
|
|
|
|
|
def base64_to_jpeg_bytes(base64_str: str) -> bytes: |
|
try: |
|
if ',' not in base64_str: |
|
raise ValueError('недопустимый формат строки base64') |
|
base64_data = base64_str.split(',', 1)[1] |
|
binary_data = b64decode(base64_data) |
|
with image_open(BytesIO(binary_data)) as img: |
|
with BytesIO() as jpeg_bytes: |
|
img.convert('RGB').save(jpeg_bytes, format='JPEG', quality=90, optimize=True) |
|
return jpeg_bytes.getvalue() |
|
except (BinasciiError, OSError) as e: |
|
raise ValueError('данные не являются корректным изображением') from e |
|
|
|
|
|
async def get_captions(image_base64_str: str, query: str, api_key: str) -> dict[str, str]: |
|
media_file_bytes = base64_to_jpeg_bytes(image_base64_str) |
|
tasks = [chat_with_media_nvcf(url, media_file_bytes, query, api_key) for url in INVOKE_URLS] |
|
results = await gather(*tasks) |
|
return dict(zip((url.split('/v1/vlm/')[-1] for url in INVOKE_URLS), results)) |
|
|
|
|
|
@asynccontextmanager |
|
async def app_lifespan(_) -> AsyncGenerator: |
|
logger.info('запуск приложения') |
|
try: |
|
logger.info('старт API') |
|
yield |
|
finally: |
|
logger.info('приложение завершено') |
|
|
|
|
|
app = FastAPI(lifespan=app_lifespan, title='NVIDIA_VLM_API') |
|
|
|
banned_endpoints = [ |
|
'/openapi.json', |
|
'/docs', |
|
'/docs/oauth2-redirect', |
|
'swagger_ui_redirect', |
|
'/redoc', |
|
] |
|
|
|
|
|
@app.middleware('http') |
|
async def block_banned_endpoints(request: Request, call_next): |
|
logger.debug(f'получен запрос: {request.url.path}') |
|
if request.url.path in banned_endpoints: |
|
logger.warning(f'запрещенный endpoint: {request.url.path}') |
|
return Response(status_code=403) |
|
response = await call_next(request) |
|
return response |
|
|
|
|
|
@app.post('/v1/describe') |
|
async def describe_v1(request: Request): |
|
logger.info('запрос `describe_v1`') |
|
body = await request.json() |
|
headers = request.headers |
|
authorization: str = headers.get('Authorization') or headers.get('authorization') |
|
nvapi_key = authorization.removeprefix('Bearer ').strip() |
|
if not authorization or not nvapi_key: |
|
return JSONResponse({'caption': 'в запросе нужно передать заголовок `Authorization: Bearer <NVAPI_KEY>`'}, status_code=401) |
|
|
|
content_text = '' |
|
image_data = '' |
|
|
|
messages = body.get('messages', []) |
|
for message in messages: |
|
role = message.get('role') |
|
content = message.get('content') |
|
|
|
if role in ['system', 'user']: |
|
if isinstance(content, str): |
|
content_text += content + ' ' |
|
elif isinstance(content, list): |
|
for item in content: |
|
if item.get('type') == 'text': |
|
content_text += item.get('text', '') + ' ' |
|
elif item.get('type') == 'image_url': |
|
image_url = item.get('image_url', {}) |
|
url = image_url.get('url') |
|
if url and url.startswith('data:image/'): |
|
image_data = url |
|
image_data, content_text = image_data.strip(), content_text.strip() |
|
|
|
if not content_text or not image_data: |
|
return JSONResponse({'caption': 'изображение должно быть передано как строка base64 `data:image/jpeg;base64,{base64_img}` а также текст'}, status_code=400) |
|
try: |
|
return JSONResponse(await get_captions(image_data, content_text, nvapi_key), status_code=200) |
|
except Exception as e: |
|
return JSONResponse({'caption': str(e)}, status_code=500) |
|
|
|
|
|
@app.get('/') |
|
async def root(): |
|
return HTMLResponse('ну пролапс, ну и что', status_code=200) |
|
|
|
|
|
if __name__ == '__main__': |
|
from uvicorn import run as uvicorn_run |
|
|
|
logger.info('запуск сервера uvicorn') |
|
uvicorn_run(app, host='0.0.0.0', port=7860) |
|
|
