caption_script.py

# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "pillow>=12.2.0",
#     "openai>=1.0.0",
# ]
# ///
import argparse
from pathlib import Path
import math
import base64
import io
import concurrent.futures

from openai import OpenAI
from PIL import Image


MAX_SEQ_LENGTH = 16384
MAX_IMAGE_PIXELS = 1024**2
BATCH_SIZE = 4

parser = argparse.ArgumentParser()
parser.add_argument('--model', type=str, required=True,
                    help="Model identifier for the OpenAI-compatible API.")
parser.add_argument('--folder', type=Path, required=True)
parser.add_argument('--api_base', type=str, default=None,
                    help="Base URL for the OpenAI-compatible API.")
parser.add_argument('--api_key', type=str, required=True,
                    help="API key for the OpenAI-compatible API.")
parser.add_argument('--system_prompt', type=str)
parser.add_argument('--prompt', type=str, default='''
Write a descriptive caption for this body inflation image that would allow a modern, natural-language-based image generation model to create something like it.
The subject is inflated into a compound shape; we're looking to capture it in detail to train a LoRA on multi-part body inflation as a concept. The prompt will be in four parts.

The first part is an overview.
For example: "<name> from <series> is <doing thing>. <name> is <a brief description of the character>. <name>'s body is inflated to a huge round shape."

The second part is the inflation section. Start from the largest/most prominent region and then move outwards.
For example:
"<name>'s belly is inflated into a wide sphere that pushes out in all directions. <name>'s arms and legs are massive and round. <name>'s [forearms, shoulders, upper arms] are pumped full, <their> [thighs, hips, calves] plump and bloated. <name>'s breasts are huge and bouncy, squeezing out from between <their> thickened arms."

The third part describes the subject's pose, taking into consideration what you said before. Include the locations of the subject's extremities to help orient the MMDiT.
For example:
"<name> stands upright in a wide stance, <their> arms resting between <their> breasts and belly. <name> holds a small helium canister in <their> puffy hands. <name>'s head is on top of <their> body, <briefly mention expression>."

For the fourth part, briefly describe the art style, background and <name>'s expression.

Note:
- Every part is mandatory. Don't use headers or any formatting, just separate with newlines.
- Keep to 2 grammatical clauses per sentence as described. The text encoder in use is Qwen3 0.6B, so it's very small and easy to confuse.
- If you can't be *certain* of the character's identity, simply describing them is enough; just use a consistent synechdoche in place of <name> (e.g. "the woman".) For anthropomorphic characters, that synechdoche should include their species (e.g. "the dog man".)
- Don't caption dialogue, text, or watermarks. Pretend they don't exist. The encoder WILL hyperfocus on them if you do.
- You needn't follow the examples exactly; describe what you *see,* and write the caption in the general shape of the examples.
''')
parser.add_argument('--prompt_file', type=Path, default=None,
                    help="Path to a text file containing the prompt. Overrides --prompt if given.")
parser.add_argument('--max_tokens', type=int, default=8192)
parser.add_argument('--temperature', type=float, default=0.7)
parser.add_argument('--overwrite', action='store_true', help='Re-caption images even if a .txt file already exists.')
parser.add_argument('--dry_run', action='store_true', help="Don't write any files.")
args = parser.parse_args()
assert args.folder.is_dir()

if args.prompt_file is not None:
    assert args.prompt_file.is_file(), f"Prompt file not found: {args.prompt_file}"
    args.prompt = args.prompt_file.read_text(encoding='utf-8')

client = OpenAI(base_url=args.api_base, api_key=args.api_key)


def rgba_to_rgb(
    image: Image.Image, background_color=(255, 255, 255)) -> Image.Image:
    """Convert an RGBA image to RGB with filled background color."""
    assert image.mode == "RGBA"
    converted = Image.new("RGB", image.size, background_color)
    converted.paste(image, mask=image.split()[3])  # 3 is the alpha channel
    return converted


def convert_image_mode(image: Image.Image, to_mode: str):
    if image.mode == to_mode:
        return image
    elif image.mode == "RGBA" and to_mode == "RGB":
        return rgba_to_rgb(image)
    else:
        return image.convert(to_mode)


def limit_size(img):
    w, h = img.size
    num_pixels = w*h
    if num_pixels > MAX_IMAGE_PIXELS:
        scale = math.sqrt(num_pixels / MAX_IMAGE_PIXELS)
        img.thumbnail((int(w/scale), int(h/scale)))
    return img


def image_to_base64(image: Image.Image) -> str:
    buffer = io.BytesIO()
    image.save(buffer, format="JPEG")
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


def caption_single(file: Path) -> tuple[Path, str]:
    img = Image.open(file)
    img = convert_image_mode(limit_size(img), 'RGB')
    b64 = image_to_base64(img)

    messages = []
    if args.system_prompt:
        messages.append({"role": "system", "content": args.system_prompt})
    messages.append({
        "role": "user",
        "content": [
            {
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{b64}"},
            },
            {
                "type": "text",
                "text": args.prompt,
            },
        ],
    })

    response = client.chat.completions.create(
        model=args.model,
        messages=messages,
        max_tokens=args.max_tokens,
        temperature=args.temperature,
    )

    text = response.choices[0].message.content or ""
    txt_file = file.with_suffix('.txt')
    if not args.dry_run:
        with open(txt_file, 'w', encoding='utf-8') as f:
            f.write(text)
    return file, text


if __name__ == '__main__':
    image_files = []
    skipped = 0
    for file in args.folder.glob('*'):
        if file.is_dir() or file.suffix == '.txt':
            continue
        if not args.overwrite and file.with_suffix('.txt').exists():
            skipped += 1
            continue
        image_files.append(file)

    if skipped:
        print(f'Skipped {skipped} image(s) that already have captions (use --overwrite to re-caption).')

    num_images = len(image_files)
    for i in range(0, num_images, BATCH_SIZE):
        batch = image_files[i:i+BATCH_SIZE]
        print(f'Progress: {i} / {num_images}')
        with concurrent.futures.ThreadPoolExecutor(max_workers=BATCH_SIZE) as executor:
            futures = [executor.submit(caption_single, f) for f in batch]
            for future in concurrent.futures.as_completed(futures):
                file, text = future.result()
                if file == batch[0]:
                    print('-'*80)
                    print(file)
                    print(text + '\n')
                    print('-'*80)