import contextlib import asyncio import hashlib import io import os import tempfile import threading from typing import Annotated, Any import cv2 import httpx import numpy as np from fastapi import FastAPI, File, Form, Header, HTTPException, UploadFile from fastapi.responses import JSONResponse from pydantic import BaseModel, Field APP = None AI_SORT_MODEL = None AI_SORT_PROCESSOR = None AI_SORT_LOCK = threading.Lock() MODEL_NAME = os.getenv("INSIGHTFACE_MODEL", "buffalo_l") AI_SORT_MODEL_NAME = os.getenv("AI_SORT_MODEL", "google/siglip2-base-patch16-224") DET_SIZE = int(os.getenv("INSIGHTFACE_DET_SIZE", "640")) MODEL_ROOT = os.path.abspath(os.getenv("INSIGHTFACE_HOME", os.path.join(os.getcwd(), ".insightface"))) MAX_IMAGE_BYTES = int(os.getenv("MAX_IMAGE_BYTES", str(24 * 1024 * 1024))) WORKER_TOKEN = os.getenv("WORKER_TOKEN", "").strip() HTTP_TIMEOUT_SECONDS = float(os.getenv("HTTP_TIMEOUT_SECONDS", "300")) app = FastAPI(title="Lensmora Image Worker", version="1.0.0") def providers(): configured = os.getenv("INSIGHTFACE_PROVIDERS", "CPUExecutionProvider") return [item.strip() for item in configured.split(",") if item.strip()] def warm_model(): global APP if APP is not None: return APP import insightface with open(os.devnull, "w") as devnull: with contextlib.redirect_stdout(devnull): face_app = insightface.app.FaceAnalysis( name=MODEL_NAME, root=MODEL_ROOT, providers=providers(), allowed_modules=["detection", "recognition"], ) face_app.prepare(ctx_id=-1, det_size=(DET_SIZE, DET_SIZE)) APP = face_app return APP def ai_sort_device(): import torch configured = os.getenv("AI_SORT_DEVICE", "auto").strip().lower() if configured and configured != "auto": return configured return "cuda" if torch.cuda.is_available() else "cpu" def warm_ai_sort_model(): global AI_SORT_MODEL, AI_SORT_PROCESSOR if AI_SORT_MODEL is not None and AI_SORT_PROCESSOR is not None: return AI_SORT_MODEL, AI_SORT_PROCESSOR with AI_SORT_LOCK: if AI_SORT_MODEL is not None and AI_SORT_PROCESSOR is not None: return AI_SORT_MODEL, AI_SORT_PROCESSOR import torch from transformers import AutoModel, AutoProcessor device = ai_sort_device() processor = AutoProcessor.from_pretrained(AI_SORT_MODEL_NAME) dtype = torch.float16 if device.startswith("cuda") else torch.float32 model = AutoModel.from_pretrained(AI_SORT_MODEL_NAME, torch_dtype=dtype).to(device) model.eval() AI_SORT_PROCESSOR = processor AI_SORT_MODEL = model return AI_SORT_MODEL, AI_SORT_PROCESSOR def require_auth(authorization: str | None, worker_token: str | None): if not WORKER_TOKEN: return auth = (authorization or "").strip() bearer = f"Bearer {WORKER_TOKEN}" if auth == bearer or (worker_token or "").strip() == WORKER_TOKEN: return raise HTTPException(status_code=401, detail="Invalid image worker token.") def normalise(vector): arr = np.asarray(vector, dtype="float32") norm = float(np.linalg.norm(arr)) if norm > 0: arr = arr / norm return [round(float(value), 8) for value in arr.tolist()] def extract_faces_from_path(image_path: str): image = cv2.imread(image_path) if image is None: raise HTTPException(status_code=422, detail="Could not read uploaded image.") return extract_faces_from_image(image) def extract_faces_from_image(image): faces = warm_model().get(image) serialised = [] for face in faces: embedding = getattr(face, "normed_embedding", None) if embedding is None: embedding = getattr(face, "embedding", None) if embedding is None: continue x1, y1, x2, y2 = [float(value) for value in face.bbox] serialised.append({ "embedding": normalise(embedding), "box": { "x": round(x1, 2), "y": round(y1, 2), "width": round(max(0.0, x2 - x1), 2), "height": round(max(0.0, y2 - y1), 2), }, "score": round(float(getattr(face, "det_score", 0.0) or 0.0), 6), }) serialised.sort(key=lambda item: item["box"]["width"] * item["box"]["height"], reverse=True) return serialised class UploadTarget(BaseModel): key: str url: str public_url: str = "" content_type: str = "image/jpeg" class ImageCallbacks(BaseModel): progress_url: str = "" crop_targets_url: str = "" token: str = "" class ThumbnailProfile(BaseModel): max_long_edge: int = Field(default=720, ge=256, le=4096) quality: int = Field(default=80, ge=50, le=96) class ProcessImageRequest(BaseModel): job_id: str photo_id: str event_id: str source_url: str upload_targets: dict[str, UploadTarget] callbacks: ImageCallbacks | None = None model: str | None = None thumbnail_profile: ThumbnailProfile = Field(default_factory=ThumbnailProfile) preview_profile: ThumbnailProfile = Field(default_factory=lambda: ThumbnailProfile(max_long_edge=1800, quality=84)) class AiSortImage(BaseModel): id: str url: str class AiSortText(BaseModel): id: str text: str class AiSortEmbedRequest(BaseModel): model: str | None = None images: list[AiSortImage] = Field(default_factory=list, max_length=64) texts: list[AiSortText] = Field(default_factory=list, max_length=64) async def post_callback(url: str, token: str, payload: dict[str, Any]): if not url or not token: return headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"} try: async with httpx.AsyncClient(timeout=30) as client: await client.post(url, headers=headers, json=payload) except Exception as exc: print(f"[image-worker] progress callback failed: {exc}", flush=True) async def report_progress(callbacks: ImageCallbacks | None, stage: str, progress: int, message: str = ""): if not callbacks: return await post_callback(callbacks.progress_url, callbacks.token, { "stage": stage, "progress": progress, "message": message, }) async def download_source_to_file(source_url: str, suffix: str = ".jpg") -> tuple[str, int, str]: total = 0 sha = hashlib.sha256() handle = tempfile.NamedTemporaryFile(delete=False, suffix=suffix) try: async with httpx.AsyncClient(timeout=HTTP_TIMEOUT_SECONDS, follow_redirects=True) as client: async with client.stream("GET", source_url) as response: if response.status_code >= 400: raise HTTPException(status_code=422, detail=f"Could not download source image: {response.status_code}") async for chunk in response.aiter_bytes(1024 * 1024): if not chunk: continue total += len(chunk) if total > MAX_IMAGE_BYTES: raise HTTPException(status_code=413, detail="Image is too large for the worker.") sha.update(chunk) handle.write(chunk) return handle.name, total, sha.hexdigest() except Exception: with contextlib.suppress(FileNotFoundError): os.unlink(handle.name) raise finally: handle.close() async def download_ai_sort_image(client: httpx.AsyncClient, item: AiSortImage): response = await client.get(item.url, follow_redirects=True) if response.status_code >= 400: raise HTTPException(status_code=422, detail=f"Could not download AI sort image {item.id}: {response.status_code}") if len(response.content) > MAX_IMAGE_BYTES: raise HTTPException(status_code=413, detail=f"AI sort image {item.id} is too large.") from PIL import Image try: image = Image.open(io.BytesIO(response.content)).convert("RGB") image.load() except Exception as exc: raise HTTPException(status_code=422, detail=f"Could not decode AI sort image {item.id}: {exc}") from exc return item.id, image def embed_ai_sort_content(images, texts): import torch def pooled_feature_tensor(output): """Handle tensor and ModelOutput return types across Transformers releases.""" if isinstance(output, torch.Tensor): return output for attribute in ("pooler_output", "image_embeds", "text_embeds"): value = getattr(output, attribute, None) if isinstance(value, torch.Tensor): return value if isinstance(output, (tuple, list)) and output and isinstance(output[0], torch.Tensor): return output[0] raise TypeError(f"Unsupported embedding output type: {type(output).__name__}") model, processor = warm_ai_sort_model() device = ai_sort_device() image_output = [] text_output = [] with AI_SORT_LOCK, torch.inference_mode(): if images: image_ids = [item[0] for item in images] inputs = processor(images=[item[1] for item in images], return_tensors="pt") inputs = {key: value.to(device) for key, value in inputs.items()} features = pooled_feature_tensor(model.get_image_features(**inputs)) features = torch.nn.functional.normalize(features.float(), dim=-1).cpu().numpy() image_output = [{"id": item_id, "embedding": normalise(vector)} for item_id, vector in zip(image_ids, features)] if texts: text_ids = [item.id for item in texts] inputs = processor( text=[item.text for item in texts], padding="max_length", truncation=True, max_length=64, return_tensors="pt", ) inputs = {key: value.to(device) for key, value in inputs.items()} features = pooled_feature_tensor(model.get_text_features(**inputs)) features = torch.nn.functional.normalize(features.float(), dim=-1).cpu().numpy() text_output = [{"id": item_id, "embedding": normalise(vector)} for item_id, vector in zip(text_ids, features)] return image_output, text_output def encode_thumbnail(image, profile: ThumbnailProfile) -> bytes: height, width = image.shape[:2] long_edge = max(width, height) if long_edge > profile.max_long_edge: scale = profile.max_long_edge / float(long_edge) image = cv2.resize(image, (max(1, int(width * scale)), max(1, int(height * scale))), interpolation=cv2.INTER_AREA) ok, encoded = cv2.imencode(".jpg", image, [int(cv2.IMWRITE_JPEG_QUALITY), int(profile.quality)]) if not ok: raise HTTPException(status_code=500, detail="Could not encode thumbnail.") return encoded.tobytes() async def put_bytes(target: UploadTarget, data: bytes): headers = {"Content-Type": target.content_type or "image/jpeg"} async with httpx.AsyncClient(timeout=HTTP_TIMEOUT_SECONDS) as client: response = await client.put(target.url, headers=headers, content=data) if response.status_code >= 400: raise HTTPException(status_code=502, detail=f"Could not upload artifact {target.key}: {response.status_code} {response.text[:200]}") async def request_crop_targets(callbacks: ImageCallbacks | None, faces: list[dict[str, Any]]) -> list[dict[str, Any]]: if not callbacks or not callbacks.crop_targets_url or not callbacks.token or not faces: return [] headers = {"Authorization": f"Bearer {callbacks.token}", "Content-Type": "application/json"} payload = {"faces": [{"index": index, "box": face.get("box", {})} for index, face in enumerate(faces)]} async with httpx.AsyncClient(timeout=60) as client: response = await client.post(callbacks.crop_targets_url, headers=headers, json=payload) if response.status_code >= 400: raise HTTPException(status_code=502, detail=f"Could not create crop upload targets: {response.status_code} {response.text[:200]}") return response.json().get("targets", []) def crop_face(image, box: dict[str, Any]) -> bytes | None: img_h, img_w = image.shape[:2] x = float(box.get("x", 0)) y = float(box.get("y", 0)) width = float(box.get("width", 0)) height = float(box.get("height", 0)) if width <= 0 or height <= 0: return None pad = 0.4 size = max(width, height) * (1 + pad * 2) cx = x + width / 2 cy = y + height / 2 crop_x = max(0, int(round(cx - size / 2))) crop_y = max(0, int(round(cy - size / 2))) crop_size = int(round(size)) crop_size = min(crop_size, img_w - crop_x, img_h - crop_y) if crop_size < 20: return None crop = image[crop_y:crop_y + crop_size, crop_x:crop_x + crop_size] crop = cv2.resize(crop, (200, 200), interpolation=cv2.INTER_AREA) ok, encoded = cv2.imencode(".jpg", crop, [int(cv2.IMWRITE_JPEG_QUALITY), 82]) if not ok: return None return encoded.tobytes() async def store_upload(upload: UploadFile) -> str: total = 0 suffix = os.path.splitext(upload.filename or "")[1].lower() or ".jpg" handle = tempfile.NamedTemporaryFile(delete=False, suffix=suffix) try: while True: chunk = await upload.read(1024 * 1024) if not chunk: break total += len(chunk) if total > MAX_IMAGE_BYTES: raise HTTPException(status_code=413, detail="Image is too large for the worker.") handle.write(chunk) return handle.name except Exception: with contextlib.suppress(FileNotFoundError): os.unlink(handle.name) raise finally: handle.close() @app.get("/health") def health(): return { "ok": True, "service": "lensmora-image-worker", "model": f"insightface-{MODEL_NAME}", "ready": APP is not None, "detSize": DET_SIZE, "aiSort": { "model": AI_SORT_MODEL_NAME, "ready": AI_SORT_MODEL is not None, "device": ai_sort_device(), }, } @app.post("/v1/warmup") def warmup( authorization: Annotated[str | None, Header()] = None, x_worker_token: Annotated[str | None, Header()] = None, ): require_auth(authorization, x_worker_token) warm_model() return {"ok": True, "model": f"insightface-{MODEL_NAME}", "detSize": DET_SIZE} @app.post("/v1/ai-sort/warmup") async def warmup_ai_sort( authorization: Annotated[str | None, Header()] = None, x_worker_token: Annotated[str | None, Header()] = None, ): require_auth(authorization, x_worker_token) await asyncio.to_thread(warm_ai_sort_model) return {"ok": True, "model": AI_SORT_MODEL_NAME, "device": ai_sort_device()} @app.post("/v1/ai-sort/embed") async def embed_for_ai_sort( request: AiSortEmbedRequest, authorization: Annotated[str | None, Header()] = None, x_worker_token: Annotated[str | None, Header()] = None, ): require_auth(authorization, x_worker_token) if request.model and request.model != AI_SORT_MODEL_NAME: raise HTTPException(status_code=409, detail=f"AI sorting worker model is {AI_SORT_MODEL_NAME}, request asked for {request.model}.") if not request.images and not request.texts: raise HTTPException(status_code=400, detail="At least one image or text item is required.") downloaded = [] errors = [] if request.images: limits = httpx.Limits(max_connections=min(16, len(request.images)), max_keepalive_connections=8) async with httpx.AsyncClient(timeout=HTTP_TIMEOUT_SECONDS, limits=limits) as client: results = await asyncio.gather( *[download_ai_sort_image(client, item) for item in request.images], return_exceptions=True, ) for item, result in zip(request.images, results): if isinstance(result, Exception): detail = getattr(result, "detail", None) or str(result) errors.append({"id": item.id, "error": str(detail)[:500]}) else: downloaded.append(result) image_embeddings, text_embeddings = await asyncio.to_thread(embed_ai_sort_content, downloaded, request.texts) dimensions = len((image_embeddings or text_embeddings)[0]["embedding"]) if (image_embeddings or text_embeddings) else 768 return { "ok": True, "model": AI_SORT_MODEL_NAME, "dimensions": dimensions, "device": ai_sort_device(), "imageEmbeddings": image_embeddings, "textEmbeddings": text_embeddings, "errors": errors, } @app.post("/v1/faces/extract") async def extract_faces( image: Annotated[UploadFile, File()], model: Annotated[str | None, Form()] = None, authorization: Annotated[str | None, Header()] = None, x_worker_token: Annotated[str | None, Header()] = None, ): require_auth(authorization, x_worker_token) if model and model != MODEL_NAME: raise HTTPException(status_code=409, detail=f"Worker model is {MODEL_NAME}, request asked for {model}.") image_path = await store_upload(image) try: faces = extract_faces_from_path(image_path) return {"faces": faces, "model": f"insightface-{MODEL_NAME}", "faceCount": len(faces)} finally: with contextlib.suppress(FileNotFoundError): os.unlink(image_path) @app.post("/v2/images/process") async def process_image_direct( request: ProcessImageRequest, authorization: Annotated[str | None, Header()] = None, x_worker_token: Annotated[str | None, Header()] = None, ): require_auth(authorization, x_worker_token) if request.model and request.model != MODEL_NAME: raise HTTPException(status_code=409, detail=f"Worker model is {MODEL_NAME}, request asked for {request.model}.") thumbnail_target = request.upload_targets.get("thumbnail") if not thumbnail_target: raise HTTPException(status_code=400, detail="thumbnail upload target is required.") preview_target = request.upload_targets.get("preview") image_path = "" try: await report_progress(request.callbacks, "downloading", 10) suffix = os.path.splitext(request.source_url.split("?", 1)[0])[1].lower() or ".jpg" image_path, source_size, sha256 = await download_source_to_file(request.source_url, suffix=suffix) await report_progress(request.callbacks, "decoding", 22) image = cv2.imread(image_path) if image is None: raise HTTPException(status_code=422, detail="Could not read downloaded image.") height, width = image.shape[:2] await report_progress(request.callbacks, "thumbnailing", 34) thumbnail_bytes = encode_thumbnail(image, request.thumbnail_profile) await put_bytes(thumbnail_target, thumbnail_bytes) preview_bytes = b"" if preview_target: await report_progress(request.callbacks, "previewing", 42) preview_bytes = encode_thumbnail(image, request.preview_profile) await put_bytes(preview_target, preview_bytes) await report_progress(request.callbacks, "detecting_faces", 52) faces = extract_faces_from_image(image) await report_progress(request.callbacks, "allocating_crops", 68) crop_targets = await request_crop_targets(request.callbacks, faces) crop_targets_by_index = {int(target.get("index", index)): target for index, target in enumerate(crop_targets)} output_faces = [] total_faces = max(1, len(faces)) for index, face in enumerate(faces): target = crop_targets_by_index.get(index) crop_public_url = "" crop_key = "" face_id = "" if target: crop_bytes = crop_face(image, face.get("box", {})) if crop_bytes: upload_target = UploadTarget( key=target.get("key", ""), url=target.get("url", ""), public_url=target.get("public_url", ""), content_type="image/jpeg", ) await put_bytes(upload_target, crop_bytes) crop_public_url = upload_target.public_url crop_key = upload_target.key face_id = target.get("face_id") or target.get("faceId") or "" output_faces.append({ **face, "face_id": face_id, "crop_key": crop_key, "crop_url": crop_public_url, }) if len(faces) > 0: await report_progress(request.callbacks, "uploading_crops", 70 + int(((index + 1) / total_faces) * 15)) await report_progress(request.callbacks, "finalizing", 90) return { "ok": True, "processor": "lensmora-image-worker-v2", "model": f"insightface-{MODEL_NAME}", "metadata": { "width": int(width), "height": int(height), "source_size_bytes": int(source_size), "sha256": sha256, }, "thumbnail": { "key": thumbnail_target.key, "public_url": thumbnail_target.public_url, "url": thumbnail_target.public_url, "size_bytes": len(thumbnail_bytes), }, "preview": { "key": preview_target.key, "public_url": preview_target.public_url, "url": preview_target.public_url, "size_bytes": len(preview_bytes), } if preview_target else None, "faces": output_faces, "faceCount": len(output_faces), } finally: if image_path: with contextlib.suppress(FileNotFoundError): os.unlink(image_path) @app.exception_handler(Exception) async def unhandled_exception(_request, exc): if isinstance(exc, HTTPException): raise exc return JSONResponse(status_code=500, content={"detail": {"error": str(exc)}})