Create app.py

app.py

@@ -0,0 +1,211 @@
+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)