Fine-tune a vision model with the Loops SDK

Fine-tune a vision-language model with image and text pairs, deploy the checkpoint, and send an image prompt.

Fine-tune a vision-language model on examples that pair an image with a text prompt and an answer. The only difference from text-only supervised fine-tuning (SFT) is how the training data is built: the trainer renders the image as visual tokens and concatenates them with the prompt and answer tokens, then computes loss over the combined sequence in the same forward_backward() and optim_step() training step.

This guide fine-tunes a low-rank adaptation (LoRA) adapter on hiyouga/geometry3k, a dataset of geometry diagrams with text problems and numeric answers. The data-building and loss-masking techniques work with any image-and-text dataset. Vision training supports a subset of models; see Vision for which models accept images and how image inputs work.

​

Prerequisites

• A workspace API key.
• Python 3.12+ and uv to install dependencies and run the scripts.
• The Baseten CLI to deploy the checkpoint.

​

Install

Export the key as BASETEN_API_KEY:

macOS/Linux

export BASETEN_API_KEY="YOUR_API_KEY"

Install the Baseten CLI:

macOS or Linux

Terminal

brew tap basetenlabs/baseten
brew trust --formula basetenlabs/baseten/baseten
brew install baseten

Windows

Download and extract the Windows archive for the latest Baseten CLI release:

Terminal

$release = Invoke-RestMethod https://api.github.com/repos/basetenlabs/baseten-cli/releases/latest
$asset = $release.assets | Where-Object { $_.name -like "*_windows_amd64.zip" }
Invoke-WebRequest $asset.browser_download_url -OutFile baseten.zip
Expand-Archive -Force baseten.zip .

Move baseten.exe to a directory on your PATH.

For other platforms or a specific version, see the Baseten CLI install reference.

Create a uv project and install the dependencies:

uv init loops-vision
cd loops-vision
uv add baseten-loops datasets pillow torch torchvision transformers

​

Create the training script

Create a file named train_vision.py in the project directory:

train_vision.py

import io

from datasets import load_dataset
from PIL import Image
from transformers import AutoImageProcessor

from baseten.loops import (
    AdamParams,
    Datum,
    ImageChunk,
    ModelInput,
    ServiceClient,
    TensorData,
)

ServiceClient opens a session with Baseten. create_lora_training_client creates a training run for the base model and attaches a LoRA adapter with the rank you set. get_tokenizer returns the base model’s tokenizer, and the image processor parses image dimensions later in the script:

train_vision.py

BASE_MODEL = "Qwen/Qwen3.5-0.8B"

service_client = ServiceClient()
training_client = service_client.create_lora_training_client(
    base_model=BASE_MODEL,
    rank=16,
)
tokenizer = training_client.get_tokenizer()
image_processor = AutoImageProcessor.from_pretrained(BASE_MODEL, trust_remote_code=True)

​

Turn examples into training data

Each example in the geometry3k dataset contains an image, a problem, and an answer. The trainer uses each example as a Datum:

• model_input contains one ImageChunk with the image bytes and one text chunk with the problem and answer tokens.
• loss_fn_inputs supplies the loss targets through a TensorData named target_tokens, with one target token per input position.

ImageChunk takes the image bytes inline, as JPEG or PNG. One chunk holds one image; add more chunks for more images. Pass the image already resized: the trainer doesn’t resize images and rejects a decoded image larger than 18 MB or 50 megapixels.

Loops doesn’t fetch external image URLs. Pass the bytes inline in ImageChunk. The ImageAssetPointerChunk type exists for Tinker compatibility, but its constructor raises ValueError.

ImageChunk.expected_tokens tells the trainer how many input tokens the image becomes. The model replaces the image with exactly that many image tokens, so the number must match what the model actually produces. For Qwen VL and Muse Spark models, compute it from the image processor as prod(image_grid_thw) // merge_size**2.

Add these functions to build a Datum from a dataset example:

train_vision.py

def compute_expected_tokens(image: Image.Image) -> int:
    result = image_processor(images=[image])
    grid = result["image_grid_thw"][0]
    merge = image_processor.merge_size
    return int(grid[0] * grid[1] * grid[2]) // (merge * merge)


def to_image_chunk(image: Image.Image, expected_tokens: int) -> ImageChunk:
    buf = io.BytesIO()
    image.convert("RGB").save(buf, format="JPEG")
    return ImageChunk(
        data=buf.getvalue(),
        format="jpeg",
        expected_tokens=expected_tokens,
    )


def to_datum(example, tokenizer):
    image = example["images"][0]
    problem = example["problem"]
    answer = example["answer"]

    expected_tokens = compute_expected_tokens(image)
    img_chunk = to_image_chunk(image, expected_tokens)

    # geometry3k wraps the prompt in an <image> placeholder; strip it.
    prompt = problem.replace("<image>", "")
    p = tokenizer.encode(prompt, add_special_tokens=False)
    a = tokenizer.encode(answer, add_special_tokens=False)

    text_chunk = ModelInput.from_ints(p + a[:-1]).chunks[0]
    targets = [-100] * (expected_tokens + len(p) - 1) + list(a)

    return Datum(
        model_input=ModelInput(chunks=[img_chunk, text_chunk]),
        loss_fn_inputs={
            "target_tokens": TensorData(
                data=targets, dtype="int64", shape=[len(targets)]
            )
        },
    )

The trainer expands each ImageChunk into expected_tokens copies of the model’s image-placeholder token ID and appends the text chunk. Because to_datum omits the final answer token from the input, the input sequence is [image: expected_tokens] [prompt: len(p)] [answer: len(a) - 1].

​

Mask the image and prompt tokens

The input and target_tokens have the same length: expected_tokens + len(p) + len(a) - 1. Mask the image and prompt positions with -100, which cross-entropy ignores, except for the last prompt position. That position targets the first answer token. The remaining positions target the subsequent answer tokens, including the final answer token:

targets = [-100] * (expected_tokens + len(p) - 1) + list(a)

​

Add the training loop

At each training step, forward_backward() runs the model on a batch and accumulates gradients. Then optim_step() applies them with AdamParams. Add the loop to train_vision.py:

train_vision.py

BATCH_SIZE = 4

dataset = load_dataset("hiyouga/geometry3k", split="train[:4]")
data = [to_datum(ex, tokenizer) for ex in dataset]
print(f"prepared {len(data)} examples")

for step, start in enumerate(range(0, len(data), BATCH_SIZE), 1):
    batch = data[start : start + BATCH_SIZE]
    fb = training_client.forward_backward(data=batch).result(timeout=600.0)
    training_client.optim_step(AdamParams(learning_rate=4e-5)).result(timeout=600.0)
    print(f"step {step} loss={fb.loss:.4f}")

For a complete fine-tune, use the full split. Image examples use more memory than text-only examples because the model processes image tensors and text together. Start with a small batch size, and increase it if GPU memory allows.

​

Save a checkpoint

Save the tuned adapter as a named checkpoint with save_weights_for_sampler(). The checkpoint stays available after the training session ends. Add the save call to train_vision.py:

train_vision.py

save_resp = training_client.save_weights_for_sampler(name="vision-step-1").result(
    timeout=600.0
)
print(f"saved checkpoint at {save_resp.path}")
training_client.close()

​

Run the training script

Qwen3.5-0.8B trains in the dedicated mode, so this run bills for GPU time until you deactivate it. Run the script:

uv run python train_vision.py

In the dedicated mode, a new run waits for its trainer to start before the first training step. To reuse a trainer deployment that’s already running, see Reuse a warm run.

Output

prepared 4 examples
step 1 loss=5.4734
saved checkpoint at bt://loops:e3m1863/sampler_weights/vision-step-1

​

Shut down the run

This example uses dedicated capacity. Closing the SDK client doesn’t deactivate the run. To shut down its trainer and paired sampler and stop GPU billing, replace <run_id> with the run ID after bt://loops: in your checkpoint path:

1. Deactivate the run, and confirm when prompted:

   baseten loops run deactivate --run-id "<run_id>"
   

2. Check that the run’s status.name is INACTIVE:

   baseten loops run describe --run-id "<run_id>" --output json
   

The saved checkpoint remains available for deployment. For the REST API option, see Deactivate a run.

​

Deploy the checkpoint

Deploy the checkpoint to Dedicated Inference. The deployment loads your LoRA adapter on top of the base weights and serves an OpenAI-compatible /v1/chat/completions endpoint for image inputs. Before you deploy, add an hf_access_token to workspace secrets. The deployment uses this token to download the base weights from Hugging Face, and it bills for its GPU while it’s active.

List the checkpoints for your run. The checkpoint path includes the run ID after bt://loops:; replace <run_id> with yours, not the example ID e3m1863:

baseten loops checkpoint list --run-id "<run_id>"

Output

                          Checkpoints for run: e3m1863
╭─────────┬──────────┬─────────┬─────────┬──────┬─────────┬─────────╮
│ ID      │ Name     │ Run     │ Target  │ Type │ Size    │ Created │
├─────────┼──────────┼─────────┼─────────┼──────┼─────────┼─────────┤
│ MPXxYJP │ vision…  │ e3m1863 │ sampler │ lora │ 20.6 MB  │ 2026-… │
╰─────────┴──────────┴─────────┴─────────┴──────┴─────────┴─────────╯

Replace <checkpoint_id> with the ID of your vision-step-1 checkpoint from the list output. The deploy command asks for a model name, GPU type, GPU count, and Hugging Face secret name. The default secret name is hf_access_token.

baseten loops checkpoint deploy --checkpoint-id "<checkpoint_id>"

Output

Successfully created deployment: deployment-1
Model ID: q86o6283
Deployment ID: qvxp6gg
Deployment succeeded.

Use the Model ID and Deployment ID from your deploy output for <model_id> and <deployment_id> in all remaining commands. Check the deployment status:

baseten model deployment describe \
  --model-id "<model_id>" \
  --deployment-id "<deployment_id>" \
  --jq '.status'

Output

"SCALED_TO_ZERO"

The deployment is ready when its status is ACTIVE or SCALED_TO_ZERO. If the status is SCALED_TO_ZERO, the first request starts a replica and takes longer than later requests.

​

Send an image prompt

Send the image as base64-encoded data in the OpenAI image_url format.

The example sends this geometry diagram with the text prompt “Find x.”:

Create build_request.py to encode the image and build the request:

build_request.py

import base64
import io
import json

from datasets import load_dataset
from PIL import Image

ds = load_dataset("hiyouga/geometry3k", split="train[:1]")
img = ds[0]["images"][0]
buf = io.BytesIO()
img.convert("RGB").save(buf, format="JPEG")
b64 = base64.b64encode(buf.getvalue()).decode()
problem = ds[0]["problem"].replace("<image>", "")

payload = {
    "model": "vision-step-1",
    "messages": [
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{b64}"},
                },
                {"type": "text", "text": problem},
            ],
        }
    ],
    "max_tokens": 64,
    "chat_template_kwargs": {"enable_thinking": False},
}

with open("request.json", "w") as f:
    json.dump(payload, f)

Run the script:

uv run python build_request.py

In the payload, model is the checkpoint name. Send the request:

cURL

curl -X POST \
  "https://model-<model_id>.api.baseten.co/deployment/<deployment_id>/sync/v1/chat/completions" \
  -H "Authorization: Bearer $BASETEN_API_KEY" \
  -H "Content-Type: application/json" \
  -d @request.json

Baseten CLI

baseten model predict --model-id "<model_id>" --deployment-id "<deployment_id>" --file request.json

Output

{
  "choices": [
    {
      "message": {
        "role": "assistant",
        "content": "We are given a circle with two chords intersecting inside it, labeled as circles with radii from the center to the endpoints.\n\nThe diagram shows:\n\n- A central vertical chord.\n- Two intersecting chords on the left side.\n- The labels are:\n  - 4 and 6 are"
      },
      "finish_reason": "length"
    }
  ],
  "usage": {"prompt_tokens": 81, "completion_tokens": 64, "total_tokens": 145}
}

One training step on four examples has little effect on the base model. Use more examples and training steps to improve accuracy on the geometry task.

​

Send a text-only prompt

Omit the image content block:

cURL

curl -X POST \
  "https://model-<model_id>.api.baseten.co/deployment/<deployment_id>/sync/v1/chat/completions" \
  -H "Authorization: Bearer $BASETEN_API_KEY" \
  -H "Content-Type: application/json" \
  -d '{
    "model": "vision-step-1",
    "messages": [
      {"role": "user", "content": "What is the capital of France?"}
    ],
    "max_tokens": 32
  }'

Baseten CLI

baseten model predict --model-id "<model_id>" --deployment-id "<deployment_id>" \
  --data '{
    "model": "vision-step-1",
    "messages": [
      {"role": "user", "content": "What is the capital of France?"}
    ],
    "max_tokens": 32
  }'

​

Deactivate the deployment

Deactivate the deployment when you finish testing:

baseten model deployment deactivate --model-id "<model_id>" --deployment-id "<deployment_id>"

For scripts, pass --yes to skip the confirmation prompt. Deactivation doesn’t delete the checkpoint, so you can deploy it again.

​

Next steps

• Deploy a checkpoint: The full deploy flow, including the dashboard option and the --config flag for non-interactive deploys.
• Train on a chat dataset: The text-only SFT guide. Vision training follows the same training step.
• Loops concepts: Sessions, trainers, samplers, checkpoints, and how weight sync works.
• Migrate from Tinker: If you’re porting a Tinker vision recipe, the ImageChunk type is call-compatible. Loops doesn’t support ImageAssetPointerChunk.