diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py index a0fcb5cc9d..2e79ae5d22 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -1,146 +1,109 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Figure captioning via a small generative VLM (PaddleOCR-VL, 4-bit). +"""Figure captioning via the user's currently-loaded chat VLM. -Defensive by design: any load or per-image failure returns an empty -string so the ingestion pipeline can fall back to its prior page-text -caption behaviour rather than crashing. +No separate vision model is loaded. If the chat model is vision-capable +(detected at ingestion-enqueue time and passed in as ``vlm_url`` / +``vlm_model``), we call its OpenAI-compatible ``/v1/chat/completions`` +with the figure as a base64 ``image_url``. If no vision-capable chat +model is loaded, captioning is skipped entirely — the ingestion falls +back to the parser's page-text ``nearest_caption``. -Lifecycle: lives inside the ingestion subprocess (`_subprocess_worker` -in `ingestion.py`). The model loads on first `caption_images` call and -is released when the subprocess exits — no impact on the parent chat -model. +Defensive: any per-image request failure returns an empty string so +ingestion stays resilient. """ from __future__ import annotations +import base64 import logging -import threading -from typing import Any +from io import BytesIO +from typing import Optional logger = logging.getLogger(__name__) -# Pre-quantized Unsloth bnb-4bit repo (see unsloth/models/mapper.py). -# Using this name directly skips the FLOAT_TO_INT_MAPPER redirect so the -# initial download fetches the 4-bit weights (~1.5 GB) instead of bf16 -# (~5 GB followed by on-the-fly quantization). -_CAPTION_MODEL_NAME = "unsloth/Qwen3-VL-2B-Instruct-unsloth-bnb-4bit" _PROMPT = ( "Describe this figure in <=60 words. Focus on factual content " "(axes, labels, captions, visible text, main objects). " "Do not speculate beyond what is visible." ) _MAX_NEW_TOKENS = 120 - -_lock = threading.Lock() -_model: Any | None = None -_processor: Any | None = None -_load_failed: bool = False +# Downscale large images so the base64 payload stays manageable; the chat +# model's prefill cost scales with image-tile count, not pixel count, but +# very large inputs still bloat the JSON body. 1600 px on the long side +# matches PR #5351's chat-composer extractor. +_MAX_IMAGE_SIZE = 1600 +_REQUEST_TIMEOUT_SECONDS = 120.0 -def _load() -> tuple[Any, Any] | None: - """Lazy-load PaddleOCR-VL. Returns (model, processor) or None on failure. +def _image_to_data_url(blob: bytes) -> str: + from PIL import Image - Sentinel-cached: once a load failure occurs in this subprocess we - don't keep retrying for every batch of images. - """ - global _model, _processor, _load_failed - if _load_failed: - return None - with _lock: - if _model is not None and _processor is not None: - return _model, _processor - try: - from unsloth import FastVisionModel - - logger.info( - "Loading RAG captioner: %s (4-bit via Unsloth)", - _CAPTION_MODEL_NAME, - ) - # FastVisionModel handles 4-bit quantization, dtype selection, - # and inference-mode wiring. Returns (model, processor) — for - # vision models the "tokenizer" slot carries the processor. - model, processor = FastVisionModel.from_pretrained( - model_name = _CAPTION_MODEL_NAME, - load_in_4bit = True, - trust_remote_code = True, - ) - FastVisionModel.for_inference(model) - _model = model - _processor = processor - return _model, _processor - except Exception as exc: - logger.warning( - "RAG captioner %s failed to load (%s). Falling back to " - "page-text captions for this ingestion.", - _CAPTION_MODEL_NAME, - exc, - ) - _load_failed = True - return None + img = Image.open(BytesIO(blob)).convert("RGB") + if max(img.size) > _MAX_IMAGE_SIZE: + img.thumbnail((_MAX_IMAGE_SIZE, _MAX_IMAGE_SIZE)) + buf = BytesIO() + img.save(buf, format = "JPEG", quality = 88) + encoded = base64.b64encode(buf.getvalue()).decode("ascii") + return f"data:image/jpeg;base64,{encoded}" -def caption_images(image_bytes_list: list[bytes]) -> list[str]: - """Generate one short caption per image. Same-length output. +def caption_images( + image_bytes_list: list[bytes], + *, + vlm_url: Optional[str] = None, + vlm_model: Optional[str] = None, +) -> list[str]: + """Caption each image via the loaded chat VLM. - Returns ``""`` for any image whose captioning failed (or all images - if the model couldn't load). Never raises — ingestion must stay - resilient to VLM unavailability. + ``vlm_url`` / ``vlm_model`` come from the parent's chat-backend probe. + When either is missing (no model loaded, or loaded model is text-only), + returns an empty string per image so the caller falls back to its + parser-provided caption. Never raises. """ if not image_bytes_list: return [] - loaded = _load() - if loaded is None: + if not vlm_url or not vlm_model: return ["" for _ in image_bytes_list] - model, processor = loaded - from io import BytesIO - - import torch - from PIL import Image + import httpx + endpoint = f"{vlm_url.rstrip('/')}/v1/chat/completions" out: list[str] = [] - for blob in image_bytes_list: - try: - image = Image.open(BytesIO(blob)).convert("RGB") - messages = [ - { - "role": "user", - "content": [ - {"type": "image", "image": image}, - {"type": "text", "text": _PROMPT}, + with httpx.Client(timeout = _REQUEST_TIMEOUT_SECONDS) as client: + for blob in image_bytes_list: + try: + data_url = _image_to_data_url(blob) + payload = { + "model": vlm_model, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": _PROMPT}, + { + "type": "image_url", + "image_url": {"url": data_url}, + }, + ], + } ], + "max_tokens": _MAX_NEW_TOKENS, + "temperature": 0.0, } - ] - # Qwen3-VL's processor pipes image + text through the chat - # template in one call when tokenize=True / return_dict=True. - inputs = processor.apply_chat_template( - messages, - tokenize = True, - add_generation_prompt = True, - return_dict = True, - return_tensors = "pt", - ) - inputs = { - k: (v.to(model.device) if hasattr(v, "to") else v) - for k, v in inputs.items() - } - input_ids_len = int(inputs["input_ids"].shape[1]) - with torch.no_grad(): - output_ids = model.generate( - **inputs, - max_new_tokens = _MAX_NEW_TOKENS, - do_sample = False, + response = client.post(endpoint, json = payload) + response.raise_for_status() + data = response.json() + content = ( + data.get("choices", [{}])[0].get("message", {}).get("content", "") ) - # Decode only the newly-generated tokens, not the prompt echo. - new_tokens = output_ids[:, input_ids_len:] - caption = processor.batch_decode( - new_tokens, - skip_special_tokens = True, - )[0] - out.append(caption.strip()) - except Exception as exc: # noqa: BLE001 - logger.warning("caption_images: skipping one image: %s", exc) - out.append("") + out.append(content.strip() if isinstance(content, str) else "") + except Exception as exc: # noqa: BLE001 + logger.warning( + "caption_images: per-image request to %s failed: %s", + endpoint, + exc, + ) + out.append("") return out diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index faf55fccb8..df33c488ae 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -60,6 +60,8 @@ def _subprocess_worker( chunking_strategy: str = "standard", mode: str = "text", document_id: str = "", + vlm_url: str | None = None, + vlm_model: str | None = None, ) -> None: try: from core.rag.chunking import chunk_pages @@ -117,6 +119,8 @@ def _subprocess_worker( model_name = model_name, out_queue = out_queue, first_index = text_count, + vlm_url = vlm_url, + vlm_model = vlm_model, ) out_queue.put({"type": "complete", "num_chunks": text_count + image_count}) except Exception as exc: # noqa: BLE001 @@ -190,6 +194,8 @@ def _stream_image_chunks( model_name: str, out_queue, first_index: int, + vlm_url: str | None = None, + vlm_model: str | None = None, ) -> int: """Persist images, emit image+caption chunks; pairs share pair_group.""" from core.rag.captioner import caption_images @@ -223,10 +229,15 @@ def _stream_image_chunks( if not paths: return 0 - # VLM-generated captions; falls back to the parser's nearest_caption - # (page-text blob) when the VLM is unavailable or fails for an image. + # VLM-generated captions via the loaded chat VLM (when one is loaded + # and vision-capable). Falls back to the parser's nearest_caption + # (page-text blob) when no VLM is available or a request fails. out_queue.put({"type": "progress", "stage": "caption_images", "progress": 0.87}) - vlm_captions = caption_images(bytes_for_encoding) + vlm_captions = caption_images( + bytes_for_encoding, + vlm_url = vlm_url, + vlm_model = vlm_model, + ) captions: list[str] = [ ( vlm_captions[i].strip() @@ -638,6 +649,34 @@ def _pump( state.push_event({"type": "error", "error": final_error}) +def _probe_loaded_vlm() -> tuple[str | None, str | None]: + """Best-effort: return (base_url, model_name) when a vision-capable + chat model is currently loaded via llama-server; (None, None) otherwise. + + Used to caption figures with the user's chat VLM instead of loading + a dedicated captioning model — no extra VRAM, no extra download. + Only the llama-server backend is supported today; transformers / + unsloth in-process VLMs would need a different bridge. + """ + try: + from core.inference.llama_cpp import get_llama_cpp_backend + except Exception: + return None, None + try: + backend = get_llama_cpp_backend() + except Exception: + return None, None + if not getattr(backend, "is_loaded", False): + return None, None + if not getattr(backend, "is_vision", False): + return None, None + base_url = getattr(backend, "base_url", None) + model_id = getattr(backend, "model_identifier", None) + if not base_url or not model_id: + return None, None + return base_url, model_id + + def enqueue_ingestion( document_id: str, stored_path: Path, @@ -657,6 +696,22 @@ def enqueue_ingestion( or resolve_embedder(mode, chunking_strategy) or RAG_EMBEDDING_MODEL ) + # Probe before forking the subprocess so the loaded-model info is + # captured in the parent's process state, then passed to the child. + vlm_url, vlm_model = (None, None) + if mode == "multimodal": + vlm_url, vlm_model = _probe_loaded_vlm() + if vlm_url: + logger.info( + "RAG ingest: will caption figures via loaded chat VLM %s at %s", + vlm_model, + vlm_url, + ) + else: + logger.info( + "RAG ingest: no vision-capable chat model loaded; " + "skipping figure captioning (fallback to page-text)." + ) job_id = str(uuid4()) with get_connection() as conn: conn.execute( @@ -686,6 +741,8 @@ def enqueue_ingestion( chunking_strategy, mode, document_id, + vlm_url, + vlm_model, ), daemon = True, )