JoyCaptionのNF4量子化モデルを使って画像キャプショニングする(GPUメモリ16GB環境対応)
はじめに
JoyCaptionという画像キャプショニングモデルをGPUメモリ16GBの環境で動かそうとして思いの外手こずったため、備忘録として記事を残しておきます。
※2025年12月19日追記: llama-serverを使ったより手軽な方法について記事を書きました
JoyCaptionとは?
JoyCaptionは、画像に対するキャプション付け(Image Captioning)というタスクを処理するLLaVAベースのVLM(Visual Language Model)の1つです。
一般的な使い方は、画像を入力して、テキストによる記述、画像生成モデルのためのプロンプト、タグなどを出力させることです。
公式GitHubリポジトリ(以下)には、特徴として以下の4つが挙げられています。詳しくはそちらをご参照ください。
- Free and Open
- Uncensored
- Diversity
- Minimal Filtering
なお、本記事の執筆時点で最新の「Beta One」を使用しています。
対象読者
本記事が対象とする読者は以下の通りです。
- ローカル環境でJoyCaptionを動かしたい
- ComfyUIなどのGUI環境ではなく、PythonからJoyCaptionを呼び出したい
- GPUメモリが16GB前後のNVIDIA製GPUを使用している
- Linuxを使用している
もしGPUメモリが20GB以上ある場合は、試せてはいませんが、公式のモデル(以下)がそのまま使用できます。
GitHubリポジトリには、GPUメモリについて最低17GB、推奨24GB以上と記載されています。
上手くいった方法
色々な方法を試した結果、以下の環境、方法では上手く推論することができました。
結果を端的に言えば「NF4(4ビットNormalFloat)量子化モデルなら8GB程度のGPUメモリでJoyCaptionが動作する」です。
- ハードウェア:
- CPU: AMD Ryzen 7 3700X(8コア)
- メモリ: 64GB
- GPU: NVIDIA GeForce RTX 4080(16GB)
- ソフトウェア:
- OS: Ubuntu 24.04.3 LTS
- Python: v3.12.9
- PyTorch: v2.9.1(CUDA 12系)
- モデル:
より詳細なパッケージ構成、バージョンについては、pyproject.toml、uv pip freezeの結果を参照ください。
`pyproject.toml`の内容
[project]
name = "joycaption"
version = "0.1.0"
requires-python = ">=3.12"
dependencies = [
"accelerate>=1.12.0",
"bitsandbytes>=0.49.0",
"pillow>=12.0.0",
"safetensors>=0.7.0",
"transformers>=4.57.3",
]
`uv pip freeze`の結果
accelerate==1.12.0
bitsandbytes==0.49.0
certifi==2025.11.12
charset-normalizer==3.4.4
filelock==3.20.0
fsspec==2025.12.0
hf-xet==1.2.0
huggingface-hub==0.36.0
idna==3.11
jinja2==3.1.6
markupsafe==3.0.3
mpmath==1.3.0
networkx==3.6.1
numpy==2.3.5
nvidia-cublas-cu12==12.8.4.1
nvidia-cuda-cupti-cu12==12.8.90
nvidia-cuda-nvrtc-cu12==12.8.93
nvidia-cuda-runtime-cu12==12.8.90
nvidia-cudnn-cu12==9.10.2.21
nvidia-cufft-cu12==11.3.3.83
nvidia-cufile-cu12==1.13.1.3
nvidia-curand-cu12==10.3.9.90
nvidia-cusolver-cu12==11.7.3.90
nvidia-cusparse-cu12==12.5.8.93
nvidia-cusparselt-cu12==0.7.1
nvidia-nccl-cu12==2.27.5
nvidia-nvjitlink-cu12==12.8.93
nvidia-nvshmem-cu12==3.3.20
nvidia-nvtx-cu12==12.8.90
packaging==25.0
pillow==12.0.0
psutil==7.1.3
pyyaml==6.0.3
regex==2025.11.3
requests==2.32.5
safetensors==0.7.0
setuptools==80.9.0
sympy==1.14.0
tokenizers==0.22.1
torch==2.9.1
tqdm==4.67.1
transformers==4.57.3
triton==3.5.1
typing-extensions==4.15.0
urllib3==2.6.2
推論に使用したコードは以下の通りです。
公式のGitHubのサンプルコードとの違いは、使用しているモデル、NF4を使っていること、既知の回避策(コード参照)です。
import sys
import torch
from PIL import Image
from transformers import (
AutoProcessor,
LlavaForConditionalGeneration,
BitsAndBytesConfig,
)
IMAGE_PATH = sys.argv[1]
PROMPT = "Write a long detailed description for this image."
MODEL_NAME = "John6666/llama-joycaption-beta-one-hf-llava-nf4"
nf4_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_quant_storage=torch.bfloat16,
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=torch.bfloat16,
)
processor = AutoProcessor.from_pretrained(MODEL_NAME)
model = LlavaForConditionalGeneration.from_pretrained(
MODEL_NAME,
torch_dtype=torch.bfloat16,
quantization_config=nf4_config,
device_map=0,
)
model.eval()
# 既知の回避策
# REF: https://github.com/fpgaminer/joycaption/issues/3#issuecomment-2619253277
attention = model.vision_tower.vision_model.head.attention
attention.out_proj = torch.nn.Linear(
attention.embed_dim, attention.embed_dim, device=model.device, dtype=torch.bfloat16
)
with torch.no_grad():
image = Image.open(IMAGE_PATH).convert("RGB")
convo = [
{"role": "system", "content": "You are a helpful image captioner."},
{"role": "user", "content": PROMPT},
]
convo_string = processor.apply_chat_template(
convo, tokenize=False, add_generation_prompt=True
)
inputs = processor(text=[convo_string], images=[image], return_tensors="pt").to(
"cuda"
)
inputs["pixel_values"] = inputs["pixel_values"].to(torch.bfloat16)
out = model.generate(
**inputs,
max_new_tokens=512,
do_sample=True,
suppress_tokens=None,
use_cache=True,
temperature=0.6,
top_k=None,
top_p=0.9,
)[0]
out = out[inputs["input_ids"].shape[1] :]
caption = processor.tokenizer.decode(
out, skip_special_tokens=True, clean_up_tokenization_spaces=False
).strip()
print(caption)
ちなみに上記環境における推論時間は約5秒、GPUメモリの使用量は7,781 MiBでした。
上手くいかなかった方法
詳細は追えていませんが、試したけれど上手く行かなかった環境、組み合わせは以下の通りです。
macOS/Linux + ollama + GGUFファイルの組み合わせ
https://huggingface.co/mradermacher/llama-joycaption-beta-one-hf-llava-GGUF などのGGUFファイルにはProjector(mmproj)が同梱されていませんでした。
そのためollamaがVLM(画像を入力するモデル)として認識せず、動作はするものの、画像を参照した推論は行えませんでした(でたらめなキャプションが生成されました)。
なお、ollamaがVLMモデルとして認識しているかどうかは、ollama show モデル名のCapabilitiesにvisionが含まれているかどうかで判断できます。
llama.cppであればProjectorなしのGGUFファイルと組み合わせできるみたいですが、未検証です。
macOS + LM Studio + GGUFファイルの組み合わせ
ollamaと同じく、LM StudioがVLMとして認識せず、動作はするものの、画像を参照した推論は行えませんでした。
Linux + PyTorch + 公式モデルの組み合わせ
OOM(Out Of Memory)となり、推論できませんでした。
vLLM + 公式モデルの組み合わせ
OOMとなり、推論できませんでした。
Linux + PyTorch + NF4量子化モデル + GPUメモリ6GBの組み合わせ
GPUメモリが6GBのNVIDIA GeForce RTX 4050では、NF4量子化モデルでもOOMとなり、推論できませんでした。
macOS + PyTorch(MPS版) + NF4量子化モデルの組み合わせ
MPSはNF4をサポートしておらず、推論できませんでした。
おわりに
JoyCaptionのNF4量子化モデルをGPUメモリ16GBの環境で動かす方法について簡単に紹介しました。
VLM界隈には詳しくはないため、頓珍漢なことをしている可能性はありますが、もしその時はコメントで優しく教えて頂けると幸いです。
なお、試したいと思っているけれど、まだできていないことは以下の通りです。また別の機会に記事にできたらと思います。
- オリジナルモデルとの精度の比較: 出力されるキャプションについて、オリジナルのモデルとどの程度異なるのか調べてみたいです。
- バッチ処理: 16GBに収まるかどうかわかりませんが、複数枚を同時に推論し、スループットが向上するかどうか調べてみたいです。
本記事が少しでも参考になったら「いいね」して頂けると励みになります。
Discussion