caption_script_expand.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='''
You are an expert prompt engineer for text-to-image models. Your task is to expand the image caption below into a highly effective image-generation prompt.

The attached image shows what the caption describes. Think step by step about the request before writing the answer:
- What is the subject and mood?
- What visual styles, mediums, and lighting options would fit? Consider two or three alternatives and pick the one that best serves the caption.
- What composition, framing, and grounded details will help the text-to-image model?

Then output an expanded prompt in 1-2 paragraphs. Avoid run-on sentences.

Follow these rules strictly:
1. **Faithfulness First:** Preserve all original subjects, actions, colors, and spatial relationships. Do not add new objects, props, characters, or animals unless the caption clearly implies them.
2. **Practical T2I Structure:** Write a prompt that a text-to-image model can parse cleanly. Group subjects with their own attributes and actions. Use grounded phrasing for poses, interactions, and spatial layout.
3. **Style Planning Stays Internal:** Use your internal reasoning to choose style, medium, framing, and lighting. Do not emit planning tags or wrappers in the visible answer body.
4. **Text Rendering:** If the caption requests visible text, quotes, labels, or typography, specify the exact text clearly and wrap requested words in quotes.
5. **Avoid Over-Specification:** Do not invent highly specific clothing, colors, materials, or scene details unless the input supports them.
6. **Structure:** Write one cohesive paragraph after the thinking block. No bullets, JSON, or markdown.
7. **Preserve User Medium:** When the caption explicitly requests a medium (e.g. "photo of", "photograph of", "illustration of", "painting of", "sketch of", "3D render of"), honor it. Do not pivot to a different medium to avoid difficulty — match the caption's stated intent.
8. **Respect Existing Detail:** If the caption is already detailed, polish and finalize it rather than heavily expanding — preserve its phrasing and direction.

Notes about the content:
The dataset contains body inflation images. The captions should be rewritten to normalize the descriptions according to a few definitions. These definitions are relative to a hypothetical uninflated version of the character. Dedicate a paragraph to this.

Of torso sections (above the legs, below the head, between the arms):
"expanded (x)" - Only slightly larger than normal; noticeably prominent.
"inflated (x)" - Significantly larger than normal; composed of round arcs.
"balloon (x)" - Extremely large. Similar volume to the character's uninflated body.
"blimp(x)" - Much larger than the uninflated character's entire volume. Fits several of the character inside of it. Using this keyword should mean the subject's body diverges heavily from a humanoid form.
"hyper (x)" - Similar to "blimp" but can substitute when the rest of the character looks normal. "balloon" and "hyper" can be used in parallel.

Of limbs and extremities (arms/legs, hands/feet):
"thickened" - Noticeably swollen. Use cues about clothing or skeletal proportions to discern.
"puffed" - Partly or completely melded into the torso; puffed limbs make an inflated character's silhouette diverge from a perfect sphere.
"sunken" - Usually of extremities; means the relevant limbs are partly or completely absorbed into an inflated character's form *without* affecting its silhouette, leaving the extremity to protrude.

User's Input:


''')
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.")
parser.add_argument('--test', action='store_true', help='Process only the first image, print the result, and exit without writing files.')
args = parser.parse_args()
if args.test:
    args.dry_run = True
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)

    prompt = args.prompt
    sidecar_file = file.with_suffix('.txt')
    if sidecar_file.exists():
        sidecar_prompt = sidecar_file.read_text(encoding='utf-8')
        prompt = f"{prompt}\n{sidecar_prompt}"

    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": 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__':
    all_files = [f for f in args.folder.glob('*') if f.is_file() and f.suffix != '.txt']

    if args.test:
        if not all_files:
            print('No images found to test.')
            exit(0)
        test_file = all_files[0]
        print(f'Test image: {test_file}')
        file, text = caption_single(test_file)
        print('-'*80)
        print(file)
        print(text + '\n')
        print('-'*80)
        exit(0)

    image_files = []
    skipped = 0
    for file in all_files:
        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)