Nodes/AITocha Vision/Load Caption Model (AITocha)
ComfyUI Node

Load Caption Model (AITocha)

Load JoyCaption the PyTorch way — no llama.cpp, no compile step

By AITocha·Created 2 months ago·Updated 2 months ago· 3
Load Caption Model (AITocha)
    • vlm
    model_idfancyfeast/llama-joycaption-beta-one-hf-llava
    quantization4-bit

    Load Caption Model (AITocha) is the boring half of a two-node setup, and boring is exactly what you want here. All it does is load a vision-language model into VRAM and hand you a handle you can feed to the pack's other node for caption after caption. The reason it exists is more interesting than the node itself.

    That model is JoyCaption Beta One - fpgaminer's 8B LLaVA-style captioner, the community's default for uncensored natural-language captions (build a LoRA dataset on Flux or another LLM-encoded base and this is what most people reach for). The classic way to run it in ComfyUI is through the GGUF nodes, which lean on llama-cpp-python compiled with CUDA. On a lot of RunPod, Vast, and Docker images that build comes out CPU-only, and a CPU-only JoyCaption silently returns nothing - which is how your workflow dies with an AttributeError: 'NoneType' object has no attribute 'replace' downstream in CLIPTextEncode. This pack is the author's fix: load the same JoyCaption the standard PyTorch way, via HuggingFace transformers + bitsandbytes, straight onto the GPU. No compilation, no GGUF, no silent CPU fallback.

    How it works

    You pick a model and a quantization level, hit Queue, and the node builds a BitsAndBytesConfig - 4-bit is nf4 with double quantization, 8-bit is plain load_in_8bit, bf16 just loads at torch_dtype=bfloat16 - then pulls the LlavaForConditionalGeneration weights with device_map="auto" and wraps them in an AITOCHA_VLM handle.

    Two implementation details worth knowing. The model is cached at module level keyed by (model_id, quantization), so re-running a queue doesn't reload the weights - your second run starts fast. But the cache holds one model at a time: switch quantization and the previous one is evicted to keep VRAM bounded, so expect a ~20–30 second reload on that first run after a change. And the imports are lazy, so ComfyUI startup doesn't slow down just because the pack is installed.

    Inputs that matter

    There are only two, and you'll set both:

    • model_id - defaults to fancyfeast/llama-joycaption-beta-one-hf-llava. The older alpha-two-hf-llava is the second choice if you want a lighter, less current model.
    • quantization - 4-bit (≈6 GB VRAM), 8-bit (≈10 GB), or bf16 (≈16 GB). 4-bit is the sensible default; bf16 buys a little speed and fidelity only if you've got the headroom.

    The output is a single vlm handle of the custom AITOCHA_VLM type. Don't expect it to plug into anything outside this pack - it only connects to the Generate Caption (AITocha) node.

    Install

    ComfyUI Manager, search "AITocha Vision", install, restart. Or manually:

    cd ComfyUI/custom_nodes
    git clone https://github.com/AITocha/ComfyUI-AITocha-Vision.git
    cd ComfyUI-AITocha-Vision
    pip install -r requirements.txt
    

    Restart ComfyUI. The dependencies are transformers, accelerate, bitsandbytes, and Pillow - nothing exotic, but it does assume a CUDA GPU. The first run downloads the model (around 14 GB) into your HF_HOME cache (default ~/.cache/huggingface); on a rented pod set HF_HOME=/workspace/huggingface so it survives restarts.

    Gotchas

    Keep bitsandbytes current - the 8-bit path in particular has a habit of demanding a newer version than you have (pip install -U bitsandbytes is the fix people actually hit). And if you're only captioning one image per run, you don't need this node at all: use the pack's Auto Caption (AITocha - one node) and skip the loader entirely. The split setup pays off when you caption many images in one graph - load once, caption N times.

    CategoryAITocha/Vision

    Inputs (2)

    NameTypeDefaultDescription
    model_idCOMBOfancyfeast/llama-joycaption-beta-one-hf-llava2 options: fancyfeast/llama-joycaption-beta-one-hf-llava, fancyfeast/llama-joycaption-alpha-two-hf-llava
    quantizationCOMBO4-bit3 options: 4-bit, 8-bit, bf16

    Outputs (1)

    NameTypeDescription
    vlmAITOCHA_VLM