From 5883a2d6c68cf31bd39f636de19cd1c6474a8075 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 26 May 2026 21:49:36 +0400 Subject: [PATCH] Studio: load RAG captioner via Unsloth FastVisionModel (4-bit, native path) --- studio/backend/core/rag/captioner.py | 28 +++++++++------------------- 1 file changed, 9 insertions(+), 19 deletions(-) diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py index eae20a079c..ba89706dca 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -48,31 +48,21 @@ def _load() -> tuple[Any, Any] | None: if _model is not None and _processor is not None: return _model, _processor try: - import torch - from transformers import ( - AutoModelForVision2Seq, - AutoProcessor, - BitsAndBytesConfig, - ) + from unsloth import FastVisionModel - logger.info("Loading RAG captioner: %s (4-bit)", _CAPTION_MODEL_NAME) - processor = AutoProcessor.from_pretrained( + logger.info( + "Loading RAG captioner: %s (4-bit via Unsloth)", _CAPTION_MODEL_NAME, - trust_remote_code = True, ) - quant_config = BitsAndBytesConfig( + # 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, - bnb_4bit_quant_type = "nf4", - bnb_4bit_compute_dtype = torch.bfloat16, - bnb_4bit_use_double_quant = True, - ) - model = AutoModelForVision2Seq.from_pretrained( - _CAPTION_MODEL_NAME, trust_remote_code = True, - device_map = "auto", - quantization_config = quant_config, ) - model.eval() + FastVisionModel.for_inference(model) _model = model _processor = processor return _model, _processor