`tools/gen_images.py`

51 lines
#!/usr/bin/env python3
"""Generate one news photo per story with local SDXL (free, runs on Apple MPS).

  ~/.venvs/sdxl/bin/python tools/gen_images.py prompts.json

prompts.json is a list of {"id": ..., "imagePrompt": ...}. Writes images/<id>.jpg
(800px wide JPEG). Skips ids that already have an image, so it is safe to re-run.
Weights: SDXL_CKPT env var, else the archived sd_xl_base_1.0 on /Volumes/Henry.
"""
import hashlib, json, os, sys, time

import torch
from diffusers import StableDiffusionXLPipeline

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "..", "images")
CKPT = os.environ.get("SDXL_CKPT") or \
    "/Volumes/Henry/mac2-archive/2026-06-26-reclaim/ComfyUI-checkpoints/sd_xl_base_1.0.safetensors"
STYLE = ("award-winning news photograph, photojournalism, natural light, 35mm, "
         "shallow depth of field, realistic, high detail")
NEGATIVE = ("text, words, letters, caption, watermark, logo, signage, typography, "
            "cartoon, illustration, painting, deformed, blurry, lowres, close-up face")


def main(path):
    items = json.load(open(path))
    os.makedirs(OUT, exist_ok=True)
    todo = [s for s in items if not os.path.exists(os.path.join(OUT, s["id"] + ".jpg"))]
    print(f"{len(items)} stories, {len(todo)} to generate", flush=True)
    if not todo:
        return
    device = "mps" if torch.backends.mps.is_available() else "cpu"
    pipe = StableDiffusionXLPipeline.from_single_file(
        CKPT, torch_dtype=torch.float16 if device == "mps" else torch.float32)
    pipe.to(device)
    pipe.set_progress_bar_config(disable=True)
    for i, s in enumerate(todo, 1):
        t0 = time.time()
        # Stable per-story seed, so a re-run of one id reproduces the same photo.
        seed = int(hashlib.sha1(s["id"].encode()).hexdigest()[:8], 16)
        img = pipe(prompt=f'{s["imagePrompt"]}, {STYLE}', negative_prompt=NEGATIVE,
                   width=1344, height=768, num_inference_steps=28, guidance_scale=6.5,
                   generator=torch.Generator("cpu").manual_seed(seed)).images[0]
        img = img.resize((800, 457))
        img.save(os.path.join(OUT, s["id"] + ".jpg"), quality=80, optimize=True)
        print(f"[{i}/{len(todo)}] {s['id']} {time.time() - t0:.0f}s", flush=True)


if __name__ == "__main__":
    main(sys.argv[1])