Spaces:

pimcore
/

local-inference-service

Running

App Files Files Community

local-inference-service / src /main.py

fashxp

lukmzig

BPT adjustments: batched embedding API, normalized vectors, resource limits, lint/type/test gate (#3)

ed962ac 3 months ago

Raw History Blame Contribute Delete

11.5 kB

# -------------------------------------------------------------------
# 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 <name>" 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("/"))

	# -------------------------------------------------------------------
	# 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 <name>" 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("/"))