# ------------------------------------------------------------------- # This source file is available under the terms of the # Pimcore Open Core License (POCL) # Full copyright and license information is available in # LICENSE.md which is distributed with this source code. # # @copyright Copyright (c) Pimcore GmbH (https://www.pimcore.com) # @license Pimcore Open Core License (POCL) # ------------------------------------------------------------------- import logging import os from typing import Annotated import torch from fastapi import Body, FastAPI, Path, Request from .classification import ClassificationTaskService from .embeddings import ( EmbeddingsRequest, EmbeddingsResponse, ImageEmbeddingTaskService, TextEmbeddingTaskService, VectorSizeResponse, ) from .middleware import MaxBodySizeMiddleware from .openapi_examples import ( EMBEDDINGS_RESPONSES, IMAGE_EMBEDDING_EXAMPLES, TEXT_EMBEDDING_EXAMPLES, ) from .text_to_image import TextToImageTaskService from .translation_task import TranslationTaskService app = FastAPI( title="Pimcore Local Inference Service", description=( "HF inference-provider-compatible inference for models not available at HF inference " "providers (translation, classification, image-to-text, and text/image embeddings)." ), version="1.0.0", ) # Cap request bodies so an oversized upload can't exhaust memory before a handler runs. A batch of # base64 images is the largest legitimate payload; 100 MB covers a full image batch with headroom. # Override via MAX_BODY_SIZE_BYTES. MAX_BODY_SIZE_BYTES = int(os.environ.get("MAX_BODY_SIZE_BYTES", 100 * 1024 * 1024)) app.add_middleware(MaxBodySizeMiddleware, max_body_size=MAX_BODY_SIZE_BYTES) # The DEBUG records say which extraction branch each embedding model took (pooler_output vs # get_image_features vs masked mean-pool) — the first thing you need when a model's embeddings look # wrong — but they are far too chatty to leave on in production, hence a knob instead of a constant. # getLevelName() returns a "Level " string for anything it doesn't recognise, so an # unparseable LOG_LEVEL degrades to INFO rather than taking the service down at import. _requested_log_level = logging.getLevelName(os.environ.get("LOG_LEVEL", "INFO").strip().upper()) LOG_LEVEL: int = _requested_log_level if isinstance(_requested_log_level, int) else logging.INFO logging.basicConfig(format="%(asctime)s %(levelname)-8s %(message)s", level=LOG_LEVEL) logger = logging.getLogger(__name__) # Singleton service instances: models and processors are cached on the instance, so reusing one # instance across requests keeps the (expensive) weights loaded instead of reloading per call. image_embedding_service = ImageEmbeddingTaskService(logger) text_embedding_service = TextEmbeddingTaskService(logger) translation_service = TranslationTaskService(logger) image_to_text_service = TextToImageTaskService(logger) zero_shot_image_classification_service = ClassificationTaskService(logger, "zero-shot-image-classification") image_classification_service = ClassificationTaskService(logger, "image-classification") zero_shot_text_classification_service = ClassificationTaskService(logger, "zero-shot-classification") text_classification_service = ClassificationTaskService(logger, "text-classification") @app.get("/gpu_check") async def gpu_check() -> dict: """Check if a GPU is available.""" available = torch.cuda.is_available() logger.info("GPU is available" if available else "GPU is not available") return {"success": True, "gpu": "GPU is available" if available else "GPU not available"} # ========================= # Translation Task # ========================= @app.post( "/translation/{model_name:path}", openapi_extra={ "requestBody": { "content": { "application/json": { "example": { "inputs": "Hello, world! foo bar", "parameters": {"repetition_penalty": 1.6}, } } } } }, ) async def translate( request: Request, model_name: str = Path( ..., description="The name of the translation model (e.g. Helsinki-NLP/opus-mt-en-de)", examples=["Helsinki-NLP/opus-mt-en-de"], ), ): """Execute translation tasks.""" return await translation_service.translate(request, model_name.rstrip("/")) # ========================= # Zero-Shot Image Classification Task # ========================= @app.post( "/zero-shot-image-classification/{model_name:path}", openapi_extra={ "requestBody": { "content": { "application/json": { "example": { "inputs": "base64_encoded_image_string", "parameters": {"candidate_labels": "green, yellow, blue, white, silver"}, } } } } }, ) async def zero_shot_image_classification( request: Request, model_name: str = Path( ..., description="The name of the zero-shot classification model (e.g., openai/clip-vit-large-patch14-336)", examples=["openai/clip-vit-large-patch14-336"], ), ): """Execute zero-shot image classification tasks.""" return await zero_shot_image_classification_service.classify(request, model_name.rstrip("/")) # ========================= # Image Classification Task # ========================= @app.post( "/image-classification/{model_name:path}", openapi_extra={ "requestBody": {"content": {"application/json": {"example": {"inputs": "base64_encoded_image_string"}}}} }, ) async def image_classification( request: Request, model_name: str = Path( ..., description="The name of the image classification model (e.g., pimcore/car-countries-classification)", examples=["pimcore/car-countries-classification"], ), ): """Execute image classification tasks.""" return await image_classification_service.classify(request, model_name.rstrip("/")) # ========================= # Zero-Shot Text Classification Task # ========================= @app.post( "/zero-shot-text-classification/{model_name:path}", openapi_extra={ "requestBody": { "content": { "application/json": { "example": { "inputs": "text to classify", "parameters": {"candidate_labels": "green, yellow, blue, white, silver"}, } } } } }, ) async def zero_shot_text_classification( request: Request, model_name: str = Path( ..., description="The name of the zero-shot text classification model (e.g., facebook/bart-large-mnli)", examples=["facebook/bart-large-mnli"], ), ): """Execute zero-shot text classification tasks.""" return await zero_shot_text_classification_service.classify(request, model_name.rstrip("/")) # ========================= # Text Classification Task # ========================= @app.post( "/text-classification/{model_name:path}", openapi_extra={"requestBody": {"content": {"application/json": {"example": {"inputs": "text to classify"}}}}}, ) async def text_classification( request: Request, model_name: str = Path( ..., description="The name of the text classification model (e.g., pimcore/car-class-classification)", examples=["pimcore/car-class-classification"], ), ): """Execute text classification tasks.""" return await text_classification_service.classify(request, model_name.rstrip("/")) # ========================= # Image to Text Task # ========================= @app.post( "/image-to-text/{model_name:path}", # This handler takes a raw `Request` (it accepts either of two content types), so FastAPI has # no typed model to generate a body schema from — without this the endpoint documents no body # at all and Swagger UI offers no way to send an image. Both branches of # `TextToImageTaskService.get_encoded_image` are documented here; keep them in sync with it. openapi_extra={ "requestBody": { "content": { "multipart/form-data": { "schema": { "type": "object", "properties": { "image": { "type": "string", "format": "binary", "description": "Image file to upload", } }, "required": ["image"], } }, "image/*": { "schema": { "type": "string", "format": "binary", "description": "Raw image bytes posted directly, with an image/* content type", } }, } } }, ) async def image_to_text( request: Request, model_name: str = Path( ..., description="The name of the image-to-text model (e.g., Salesforce/blip-image-captioning-base)", examples=["Salesforce/blip-image-captioning-base"], ), ): """Execute image-to-text tasks.""" return await image_to_text_service.extract(request, model_name.rstrip("/")) # ========================= # Text embeddings # ========================= @app.post("/text-embedding", response_model=EmbeddingsResponse, responses=EMBEDDINGS_RESPONSES) async def text_embedding( body: Annotated[EmbeddingsRequest, Body(openapi_examples=TEXT_EMBEDDING_EXAMPLES)], ) -> EmbeddingsResponse: """Text embeddings: `{"model": ..., "input": "..."|[...]}` -> `{"data": [{index, embedding}]}`. The model is taken from the body; `input` may be one string or a list, and a list is embedded in a single forward pass. Every vector comes back L2-normalized. """ return await text_embedding_service.generate_embeddings(body) # ========================= # Image embeddings # ========================= @app.post("/image-embedding", response_model=EmbeddingsResponse, responses=EMBEDDINGS_RESPONSES) async def image_embedding( body: Annotated[EmbeddingsRequest, Body(openapi_examples=IMAGE_EMBEDDING_EXAMPLES)], ) -> EmbeddingsResponse: """Image embeddings, same shape as /text-embedding: `input` is one base64-encoded image or a list of them. An image that cannot be decoded is skipped rather than failing the batch — it is omitted from `data`, so its `index` is absent instead of shifting the others. """ return await image_embedding_service.generate_embeddings(body) # ========================= # Embedding Vector Size # ========================= @app.get("/embedding-vector-size/{model_name:path}", response_model=VectorSizeResponse) async def embedding_vector_size( model_name: str = Path( ..., description="The name of the embedding model.", examples=["nomic-ai/nomic-embed-text-v1.5"], ), ) -> VectorSizeResponse: """Return the dimensionality of the embeddings a model produces.""" # Either embedding service works — they share the base class model cache. return await image_embedding_service.get_embedding_vector_size(model_name.rstrip("/"))