recaption.py

#!/usr/bin/env python3
"""
recaption.py — DALL-E-3-method recaptioning pipeline for training-set prep.

Walks an image directory, calls a local VLM via the Ollama API, writes
<image>.txt dense captions as sidecars, implements DALL-E-3-style caption
mixing (95% descriptive synthetic / 5% pass-through short), caps token
length per encoder target, and logs per-image quality/hallucination flags.

USAGE
-----
# Phase-0 proof run (30 images, descriptive mode):
python recaption.py --img-dir ./dataset/raw --model qwen2.5vl:7b

# Full domain run with FLUX dual-encoder output (T5-long captions):
python recaption.py --img-dir ./dataset/raw --target flux --model qwen2.5vl:7b

# Structured 4-part captions (Re-LAION method, slightly better on SD2):
python recaption.py --img-dir ./dataset/raw --style structured

# Preview mode — print first 5 captions, write nothing:
python recaption.py --img-dir ./dataset/raw --dry-run --sample 5

# With caption mixing (95% DSC / 5% short pass-through):
python recaption.py --img-dir ./dataset/raw --mix-ratio 0.05

# With trigger word injection:
python recaption.py --img-dir ./dataset/raw --trigger "ohwxperson"

DEPENDENCIES (all on the Mac / serve box)
------------------------------------------
pip install ollama pillow imagehash tqdm requests

MODEL SETUP (one-time, run on Mac where Ollama is already installed)
----------------------------------------------------------------------
ollama pull qwen2.5vl:7b   # 6.0 GB — recommended for Mac M2 16GB
# Alt: qwen2.5vl:3b (3.2 GB) if VRAM is contested during other work

CAPTION STRATEGY (DALL-E-3 paper, Betker et al. 2023)
------------------------------------------------------
Two caption types:
  DSC — Descriptive Synthetic Caption: long, every visual detail.
        Used for 95% of training images (default, --mix-ratio 0.05).
  SSC — Short Synthetic Caption: subject + action only, ~10-15 words.
        Used for the 5% pass-through to prevent overfitting on verbose
        prompt distributions.

Mixing prevents the T2I model from only responding to wall-of-text prompts.
95/5 DSC/SSC is the published DALL-E-3 optimum (ablations show monotonic
improvement with more synthetic captions up to 95%; pure synthetic hurts).

ENCODER LENGTH TARGETS
-----------------------
  sdxl:  CLIP-L/14 encodes first 77 tokens only. Front-load subject +
         key attributes in the first ~60 words. Caption is also written
         in full (no truncation) so kohya can use the OpenCLIP second
         encoder (ViT-bigG, 77 tokens) which covers more vocabulary.
         Recommended cap: 225 chars / ~50 words for safety across both
         SDXL text encoders.
  flux:  T5-XXL encodes up to 512 tokens (t5xxl_max_token_length=512
         in ai-toolkit config). No practical cap — full dense caption.
         CLIP-L in FLUX still only sees first 77 tokens, so front-loading
         still matters. Recommended cap: 450 words / ~600 chars.

HALLUCINATION FLAGS (logged, not filtered — you decide)
--------------------------------------------------------
  REPEAT  — word repetition ratio > 0.12 (same word > 12% of tokens)
  SHORT   — caption < 15 words (likely VLM refusal or truncation)
  LONG    — caption > encoder cap words
  TIMEOUT — VLM call timed out (image written to error log, no .txt)
  OK      — passes all heuristics
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import json
import logging
import os
import random
import re
import sys
import time
from collections import Counter
from pathlib import Path
from typing import Literal

# ---------------------------------------------------------------------------
# Optional imports — graceful degradation
# ---------------------------------------------------------------------------
try:
    import imagehash
    from PIL import Image as PILImage
    HAS_DEDUP = True
except ImportError:
    HAS_DEDUP = False

try:
    from tqdm import tqdm
    HAS_TQDM = True
except ImportError:
    HAS_TQDM = False

try:
    import ollama as ollama_client
    HAS_OLLAMA_SDK = True
except ImportError:
    HAS_OLLAMA_SDK = False

try:
    import requests
    HAS_REQUESTS = True
except ImportError:
    HAS_REQUESTS = False

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
SUPPORTED_EXTS = {".jpg", ".jpeg", ".png", ".webp"}
DEFAULT_MODEL = "qwen2.5vl:7b"
OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "http://localhost:11434")

# DALL-E-3 mixing ratio: 95% DSC, 5% SSC (Betker et al. 2023, §3.2)
DEFAULT_MIX_RATIO = 0.05  # fraction of images that get SSC (short caption)

# Encoder token/word caps
CAPS = {
    "sdxl": {"words": 50, "chars": 300},   # CLIP-L 77-tok = ~50 words; front-load
    "flux": {"words": 400, "chars": 2500},  # T5-XXL 512 tok = ~400 words
}

# Hallucination heuristics
REPEAT_THRESHOLD = 0.12  # fraction of tokens that are the same word
MIN_WORDS = 15
TIMEOUT_SEC = 90  # per-image VLM call timeout


# ---------------------------------------------------------------------------
# Caption prompt templates
# ---------------------------------------------------------------------------

# DSC — DALL-E-3 style: dense, descriptive, hallucination-resistant
# Structured to capture: subject → action → setting → aesthetics → camera
# Front-loads subject to survive CLIP 77-token truncation.
DSC_PROMPT_TEMPLATE = """\
You are an expert image captioner for AI training datasets. Describe the image \
in dense, specific, factual natural language. Structure your description in this \
exact order:
1. Main subject(s): what they are, appearance, any text visible
2. Action or pose if applicable
3. Setting or background details
4. Colors, lighting, and visual style
5. Camera angle and framing (e.g., close-up, wide shot, overhead)

Rules:
- Write ONE continuous paragraph, no bullet points or headers
- Be specific: say "red brick wall" not "wall"
- Include every visible object that takes up >5% of the frame
- Do not invent details not visible in the image
- Do not mention image quality, resolution, or that this is a photograph/image
- Do not add subjective opinions or emotions
- Output the description only, no preamble or sign-off
- Length: 40-120 words
"""

# SSC — short caption, main subject only, ~10-15 words
# Used for the 5% mixing fraction to regularize the T2I model.
SSC_PROMPT_TEMPLATE = """\
Describe the main subject of this image in one concise sentence of 10-15 words. \
Focus only on the primary subject and their most important visual attribute. \
Output only the sentence, no preamble.
"""

# Structured 4-part template from Re-LAION paper (Hirano et al. 2025, arXiv 2507.05300)
# Shows +1.6pp VQA gain on SD2 vs random-ordered captions.
# Use with --style structured
STRUCTURED_PROMPT_TEMPLATE = """\
Describe the image using the following structure (4 sentences in total, \
written as one paragraph separated by spaces, NOT bullet points):
Sentence 1: Subjects or objects in the image, including actions if applicable.
Sentence 2: Location and setting.
Sentence 3: Image aesthetics (lighting, color palette, mood, style).
Sentence 4: Camera perspective including angle, framing, and focal point.

Output the 4 sentences only, no labels, no preamble.
"""


def get_prompt(style: str, caption_type: Literal["dsc", "ssc"]) -> str:
    if caption_type == "ssc":
        return SSC_PROMPT_TEMPLATE
    if style == "structured":
        return STRUCTURED_PROMPT_TEMPLATE
    return DSC_PROMPT_TEMPLATE


# ---------------------------------------------------------------------------
# VLM call
# ---------------------------------------------------------------------------

def call_vlm(
    image_path: Path,
    prompt: str,
    model: str,
    timeout: int = TIMEOUT_SEC,
) -> str | None:
    """
    Call the Ollama vision endpoint.
    Prefers the ollama Python SDK if installed; falls back to raw HTTP via
    the ``requests`` library. If neither dependency is available the function
    logs a clear error and returns None (caller records TIMEOUT flag).

    Install at least one backend:
        pip install ollama          # SDK (recommended)
        pip install requests        # raw-HTTP fallback
    """
    if not HAS_OLLAMA_SDK and not HAS_REQUESTS:
        logging.error(
            "No VLM backend available: install 'ollama' (SDK) or 'requests' "
            "(HTTP fallback) — pip install ollama"
        )
        return None

    img_bytes = image_path.read_bytes()
    img_b64 = base64.b64encode(img_bytes).decode("utf-8")

    if HAS_OLLAMA_SDK:
        try:
            resp = ollama_client.chat(
                model=model,
                messages=[
                    {
                        "role": "user",
                        "content": prompt,
                        "images": [img_bytes],  # SDK accepts raw bytes
                    }
                ],
                options={"temperature": 0.2, "num_predict": 512},
            )
            return resp.message.content.strip()
        except Exception as exc:
            logging.warning("ollama SDK call failed (%s), retrying via HTTP", exc)

    # Fallback: raw HTTP to /api/generate (requires 'requests')
    if not HAS_REQUESTS:
        logging.error(
            "ollama SDK unavailable and 'requests' not installed — "
            "pip install requests"
        )
        return None

    payload = {
        "model": model,
        "prompt": prompt,
        "images": [img_b64],
        "stream": False,
        "options": {"temperature": 0.2, "num_predict": 512},
    }
    try:
        r = requests.post(
            f"{OLLAMA_HOST}/api/generate",
            json=payload,
            timeout=timeout,
        )
        r.raise_for_status()
        return r.json().get("response", "").strip()
    except Exception as exc:
        logging.error("VLM call failed for %s: %s", image_path.name, exc)
        return None


# ---------------------------------------------------------------------------
# Quality / hallucination heuristics
# ---------------------------------------------------------------------------

def quality_flag(caption: str, target: str) -> str:
    """
    Returns one of: OK | REPEAT | SHORT | LONG
    Caller combines with TIMEOUT for the full flag set.
    """
    words = caption.lower().split()
    if len(words) < MIN_WORDS:
        return "SHORT"

    # Word repetition check
    freq = Counter(words)
    most_common_count = freq.most_common(1)[0][1]
    if most_common_count / len(words) > REPEAT_THRESHOLD:
        return "REPEAT"

    cap_words = CAPS[target]["words"]
    if len(words) > cap_words:
        return "LONG"

    return "OK"


def truncate_caption(caption: str, target: str) -> str:
    """
    Truncate to word cap while keeping whole sentences where possible.
    For SDXL: hard 50-word cap; for FLUX: 400-word cap (rarely hit).
    """
    cap = CAPS[target]["words"]
    words = caption.split()
    if len(words) <= cap:
        return caption
    truncated = " ".join(words[:cap])
    # Try to end on a sentence boundary
    last_period = truncated.rfind(".")
    if last_period > len(truncated) * 0.6:
        return truncated[: last_period + 1]
    return truncated


def inject_trigger(caption: str, trigger: str | None) -> str:
    if not trigger:
        return caption
    # ai-toolkit convention: [trigger] placeholder in caption text
    # Here we prepend "trigger, " for natural language captions
    return f"{trigger}, {caption}"


# ---------------------------------------------------------------------------
# Near-duplicate detection
# ---------------------------------------------------------------------------

def compute_phash(path: Path) -> str | None:
    if not HAS_DEDUP:
        return None
    try:
        img = PILImage.open(path).convert("RGB")
        return str(imagehash.phash(img))
    except Exception:
        return None


def find_duplicates(
    image_paths: list[Path],
    hamming_threshold: int = 8,
) -> set[Path]:
    """
    Returns set of near-duplicate paths to SKIP (keeping first occurrence).
    Uses perceptual hash with hamming distance <= threshold.
    """
    if not HAS_DEDUP:
        logging.warning("imagehash/Pillow not installed — dedup skipped")
        return set()

    seen: list[tuple[str, Path]] = []
    duplicates: set[Path] = set()

    for p in image_paths:
        h = compute_phash(p)
        if h is None:
            continue
        ih = imagehash.hex_to_hash(h)
        is_dup = False
        for existing_h_str, _ in seen:
            existing_h = imagehash.hex_to_hash(existing_h_str)
            if ih - existing_h <= hamming_threshold:
                is_dup = True
                break
        if is_dup:
            duplicates.add(p)
        else:
            seen.append((h, p))

    return duplicates


# ---------------------------------------------------------------------------
# Aspect-ratio bucket assignment (for downstream kohya/ai-toolkit toml)
# ---------------------------------------------------------------------------

# Standard SDXL and FLUX training buckets
SDXL_BUCKETS = [
    (1024, 1024), (1152, 896), (896, 1152),
    (1216, 832), (832, 1216), (1344, 768),
    (768, 1344), (1536, 640), (640, 1536),
]

FLUX_BUCKETS = [
    (1024, 1024), (1280, 768), (768, 1280),
    (1536, 640), (640, 1536), (1152, 896),
    (896, 1152), (1024, 768), (768, 1024),
]


def assign_bucket(
    width: int, height: int, target: str
) -> tuple[int, int]:
    buckets = FLUX_BUCKETS if target == "flux" else SDXL_BUCKETS
    img_ar = width / height
    best = min(buckets, key=lambda b: abs((b[0] / b[1]) - img_ar))
    return best


# ---------------------------------------------------------------------------
# Main pipeline
# ---------------------------------------------------------------------------

def recaption(
    img_dir: Path,
    out_dir: Path | None,
    model: str,
    target: str,
    style: str,
    mix_ratio: float,
    trigger: str | None,
    overwrite: bool,
    dry_run: bool,
    sample: int | None,
    dedup_threshold: int,
    seed: int,
) -> None:
    random.seed(seed)
    out_dir = out_dir or img_dir
    # Dry-run must not create directories or write any files.
    if not dry_run:
        out_dir.mkdir(parents=True, exist_ok=True)

    log_path = out_dir / "recaption_log.jsonl"

    # Gather images
    all_images = sorted(
        p for p in img_dir.iterdir()
        if p.suffix.lower() in SUPPORTED_EXTS
    )
    if not all_images:
        logging.error("No images found in %s", img_dir)
        sys.exit(1)

    if sample:
        all_images = all_images[:sample]

    logging.info("Found %d images in %s", len(all_images), img_dir)

    # Near-duplicate removal
    duplicates = find_duplicates(all_images, dedup_threshold)
    if duplicates:
        logging.info(
            "Dedup: skipping %d near-duplicates (hamming <= %d)",
            len(duplicates), dedup_threshold,
        )

    to_process = [p for p in all_images if p not in duplicates]

    # Determine which images get SSC (short) vs DSC (descriptive)
    n_ssc = max(1, int(len(to_process) * mix_ratio)) if mix_ratio > 0 else 0
    ssc_indices = set(random.sample(range(len(to_process)), n_ssc)) if n_ssc else set()
    logging.info(
        "Caption plan: %d DSC + %d SSC (mix_ratio=%.2f, DALL-E-3 method)",
        len(to_process) - n_ssc, n_ssc, mix_ratio,
    )

    iterator = enumerate(to_process)
    if HAS_TQDM and not dry_run:
        iterator = enumerate(tqdm(to_process, desc="captioning"))

    stats = {"ok": 0, "short": 0, "repeat": 0, "long": 0, "timeout": 0, "skip": 0}
    t0 = time.time()

    for idx, img_path in iterator:
        txt_path = out_dir / (img_path.stem + ".txt")

        if not overwrite and txt_path.exists():
            stats["skip"] += 1
            continue

        caption_type: Literal["dsc", "ssc"] = "ssc" if idx in ssc_indices else "dsc"
        prompt = get_prompt(style, caption_type)

        # Dry-run: print first N and exit
        if dry_run:
            print(f"\n--- {img_path.name} [{caption_type.upper()}] ---")
            print(f"  Prompt snippet: {prompt[:120].strip()}...")
            continue

        caption = call_vlm(img_path, prompt, model)

        if caption is None:
            flag = "TIMEOUT"
            stats["timeout"] += 1
        else:
            flag = quality_flag(caption, target)
            stats[flag.lower()] += 1
            caption = truncate_caption(caption, target)
            caption = inject_trigger(caption, trigger)

        # Bucket info for sidecar metadata
        bucket: tuple[int, int] | None = None
        if HAS_DEDUP:
            try:
                img = PILImage.open(img_path)
                bucket = assign_bucket(img.width, img.height, target)
            except Exception:
                pass

        # Write sidecar .txt
        if caption and flag != "TIMEOUT":
            txt_path.write_text(caption, encoding="utf-8")

        # Write log entry
        log_entry = {
            "image": img_path.name,
            "caption_type": caption_type,
            "flag": flag,
            "word_count": len(caption.split()) if caption else 0,
            "bucket": list(bucket) if bucket else None,
            "model": model,
            "target": target,
            "ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        }
        with log_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(log_entry) + "\n")

    elapsed = time.time() - t0
    total = len(to_process)
    logging.info(
        "Done: %d images in %.1fs (%.1f s/img) | "
        "ok=%d short=%d repeat=%d long=%d timeout=%d skip=%d",
        total, elapsed, elapsed / max(total, 1),
        stats["ok"], stats["short"], stats["repeat"],
        stats["long"], stats["timeout"], stats["skip"],
    )
    if dry_run:
        logging.info("Dry-run complete — no files written, no VLM calls made")
    else:
        logging.info("Captions written to %s", out_dir)
        logging.info("Quality log at %s", log_path)


# ---------------------------------------------------------------------------
# Dataset config helpers
# ---------------------------------------------------------------------------

def print_dataset_configs(img_dir: Path, target: str) -> None:
    """
    Print ready-to-use kohya and ai-toolkit dataset config snippets.
    """
    abs_dir = img_dir.resolve()

    print("\n" + "="*60)
    print("KOHYA SD-SCRIPTS dataset.toml (paste into your config)")
    print("="*60)
    if target == "sdxl":
        print(f"""
[general]
resolution = 1024
caption_extension = ".txt"
shuffle_caption = false      # natural language — do NOT shuffle
enable_bucket = true
bucket_no_upscale = true
min_bucket_reso = 512
max_bucket_reso = 2048

[[datasets]]
[[datasets.subsets]]
image_dir = "{abs_dir}"
num_repeats = 1
""")
    else:
        print(f"""
# ai-toolkit config.yaml dataset block (FLUX)
datasets:
  - folder_path: "{abs_dir}"
    caption_ext: "txt"          # sidecars written by recaption.py
    default_caption: "[trigger]"
    resolution: [512, 512]
    enable_bucket: true
    shuffle_tokens: false        # natural language — order matters
""")
    print("="*60 + "\n")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument(
        "--img-dir", type=Path, required=True,
        help="Directory of training images (.jpg/.jpeg/.png/.webp)",
    )
    p.add_argument(
        "--out-dir", type=Path, default=None,
        help="Where to write .txt sidecars (default: same as img-dir)",
    )
    p.add_argument(
        "--model", default=DEFAULT_MODEL,
        help=f"Ollama vision model (default: {DEFAULT_MODEL})",
    )
    p.add_argument(
        "--target", choices=["sdxl", "flux"], default="sdxl",
        help="Encoder target — sets caption length cap (default: sdxl)",
    )
    p.add_argument(
        "--style", choices=["dense", "structured"], default="dense",
        help=(
            "dense = DALL-E-3 DSC free-form (default); "
            "structured = Re-LAION 4-part template (+1.6pp on SD2)"
        ),
    )
    p.add_argument(
        "--mix-ratio", type=float, default=DEFAULT_MIX_RATIO,
        help=(
            f"Fraction of images that get SSC (short) captions "
            f"(default: {DEFAULT_MIX_RATIO}, i.e. 5%% SSC / 95%% DSC per DALL-E-3)"
        ),
    )
    p.add_argument(
        "--trigger", default=None,
        help='Trigger word to prepend (e.g. "ohwxperson"). '
             'Use [trigger] placeholder in custom captions for ai-toolkit compat.',
    )
    p.add_argument(
        "--overwrite", action="store_true",
        help="Re-caption images that already have a .txt sidecar",
    )
    p.add_argument(
        "--dry-run", action="store_true",
        help=(
            "Preview mode: print prompt snippets for up to --sample images. "
            "Makes NO VLM calls, writes NO files, and creates NO directories. "
            "Safe to run on any --img-dir."
        ),
    )
    p.add_argument(
        "--sample", type=int, default=None,
        help="Process only first N images, N >= 0 (for Phase-0 proof run)",
    )
    p.add_argument(
        "--dedup-threshold", type=int, default=8,
        help=(
            "Perceptual hash hamming distance threshold for near-duplicate "
            "removal (0=exact, 8=default, 20=loose). Requires imagehash+Pillow."
        ),
    )
    p.add_argument(
        "--seed", type=int, default=42,
        help="Random seed for SSC selection",
    )
    p.add_argument(
        "--print-configs", action="store_true",
        help="Print kohya/ai-toolkit dataset config snippets and exit",
    )
    p.add_argument(
        "--verbose", action="store_true",
        help="Verbose logging",
    )
    return p.parse_args()


def main() -> None:
    args = parse_args()
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
        datefmt="%H:%M:%S",
    )

    # --- Input validation ---
    if not args.img_dir.exists():
        logging.error("--img-dir %s does not exist or is not accessible", args.img_dir)
        sys.exit(1)
    if not args.img_dir.is_dir():
        logging.error("--img-dir %s is not a directory", args.img_dir)
        sys.exit(1)

    if args.sample is not None and args.sample < 0:
        logging.error("--sample must be >= 0, got %d", args.sample)
        sys.exit(1)

    if not (0.0 <= args.mix_ratio <= 1.0):
        logging.error(
            "--mix-ratio must be in [0.0, 1.0], got %.4f", args.mix_ratio
        )
        sys.exit(1)

    if args.print_configs:
        print_dataset_configs(args.img_dir, args.target)
        sys.exit(0)

    recaption(
        img_dir=args.img_dir,
        out_dir=args.out_dir,
        model=args.model,
        target=args.target,
        style=args.style,
        mix_ratio=args.mix_ratio,
        trigger=args.trigger,
        overwrite=args.overwrite,
        dry_run=args.dry_run,
        sample=args.sample,
        dedup_threshold=args.dedup_threshold,
        seed=args.seed,
    )


if __name__ == "__main__":
    main()