🖼️

JoyCaptionのNF4量子化モデルを使って画像キャプショニングする(GPUメモリ16GB環境対応)

に公開

はじめに

JoyCaptionという画像キャプショニングモデルをGPUメモリ16GBの環境で動かそうとして思いの外手こずったため、備忘録として記事を残しておきます。

※2025年12月19日追記: llama-serverを使ったより手軽な方法について記事を書きました

https://zenn.dev/yuyakato/articles/e3a3935e4996aa

JoyCaptionとは?

JoyCaptionは、画像に対するキャプション付け(Image Captioning)というタスクを処理するLLaVAベースのVLM(Visual Language Model)の1つです。
一般的な使い方は、画像を入力して、テキストによる記述、画像生成モデルのためのプロンプト、タグなどを出力させることです。
公式GitHubリポジトリ(以下)には、特徴として以下の4つが挙げられています。詳しくはそちらをご参照ください。

  • Free and Open
  • Uncensored
  • Diversity
  • Minimal Filtering

https://github.com/fpgaminer/joycaption

なお、本記事の執筆時点で最新の「Beta One」を使用しています。

対象読者

本記事が対象とする読者は以下の通りです。

  • ローカル環境でJoyCaptionを動かしたい
  • ComfyUIなどのGUI環境ではなく、PythonからJoyCaptionを呼び出したい
  • GPUメモリが16GB前後のNVIDIA製GPUを使用している
  • Linuxを使用している

もしGPUメモリが20GB以上ある場合は、試せてはいませんが、公式のモデル(以下)がそのまま使用できます。
GitHubリポジトリには、GPUメモリについて最低17GB、推奨24GB以上と記載されています。

https://huggingface.co/fancyfeast/llama-joycaption-beta-one-hf-llava

上手くいった方法

色々な方法を試した結果、以下の環境、方法では上手く推論することができました。
結果を端的に言えば「NF4(4ビットNormalFloat)量子化モデルなら8GB程度のGPUメモリでJoyCaptionが動作する」です。

より詳細なパッケージ構成、バージョンについては、pyproject.tomluv pip freezeの結果を参照ください。

`pyproject.toml`の内容
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 モデル名Capabilitiesvisionが含まれているかどうかで判断できます。

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