#!/usr/bin/env python3
"""B4: батч AI-генерации кандидатов-плейсхолдеров (SD1.5, GPU fp16).

Модель: Retro Diffusion (LostMedia/RetroDiffusion, Public Domain Mark 1.0) —
SD1.5-файнтюн пиксель-арта, слушает промпты одиночных объектов (проба 2026-09-09:
очаг/часы — отличные, мот/свиток читаются). Листовая модель B3
(SD_PixelArt_SpriteSheet_Generator) рисует только 4 фигурки персонажей —
для объектов непригодна. RTX 3090 fp16: ~1.3 с/вид, семена бесплатны.

Список объектов — LIST (id, промпт): интерактивы-плейсхолдеры (InteractableViews),
снаряд плевка и иконки сумки 8×8 (арт-библия §4). Слоты назначает post_b4.mjs.

Запуск (интерпретатор — venv-aiaudio, там torch cu130):
  ~/.cache/rpg-ai/venv-aiaudio/bin/python gen_b4.py          # все объекты, SEEDS семян
  ... gen_b4.py signpost hearth:7                             # конкретные id[:seed]
Вывод: aiart/review-b4/raw/<id>_<seed>.png (512×512 сырьё).
"""
import sys
import time
from pathlib import Path

import torch
from diffusers import StableDiffusionPipeline

MODEL_ID = 'LostMedia/RetroDiffusion'
STEPS = 20
CFG = 7.0
SEEDS = (1, 2, 3, 4)
NEG = ('blurry, photo, realistic, smooth gradients, text, watermark, '
       'multiple objects, cropped, character, person, '
       'warm background, brown background, orange background, sunset lighting')

LIST = {
    # --- интерактивы-плейсхолдеры (InteractableViews) ---
    'signpost': 'pixel art of a weathered wooden road signpost with an arrow board, dark wood',
    'mote': 'pixel art of a small mound of gray ash with a glowing warm ember inside',
    'tone_tree': 'pixel art of a dead tree trunk wrapped in black cables with small copper bells, on a plain pale gray background',
    'resonator': 'pixel art of a copper bell resonator machine on a rusty steel frame, on a plain pale gray background',
    'bell_rope': 'pixel art of a hanging bell pull rope with a wooden handle grip',
    'hearth': 'pixel art of a stone hearth fireplace with a warm fire inside',
    'note': 'pixel art of a small crumpled paper note lying flat on the ground',
    'counter': 'pixel art of a wooden market stall counter with crates and goods on top, side view',
    'spit': 'pixel art of a small blob of gray ash, a projectile',
    # --- иконки сумки 8×8 (items.ts) ---
    'icon_clock': 'pixel art icon of a small copper pocket watch',
    'icon_cloth': 'pixel art icon of a folded waxed cloth mask',
    'icon_bellflower': 'pixel art icon of a small blue bellflower',
    'icon_salt': 'pixel art icon of a small pouch of salt',
    'icon_flask': 'pixel art icon of a small copper water flask',
    'icon_mote': 'pixel art icon of a small warm glowing ash mote',
    'icon_map': 'pixel art icon of a rolled paper map scroll',
}


def main() -> None:
    args = sys.argv[1:]
    wanted = {a.split(':')[0] for a in args} if args else set(LIST)
    seeds_of = {}
    for a in args:
        if ':' in a:
            k, s = a.split(':', 1)
            seeds_of.setdefault(k, []).extend(int(x) for x in s.split(':'))
    out_dir = Path(__file__).resolve().parent / 'review-b4' / 'raw'
    out_dir.mkdir(parents=True, exist_ok=True)

    t0 = time.time()
    pipe = StableDiffusionPipeline.from_pretrained(
        MODEL_ID, torch_dtype=torch.float16, safety_checker=None,
    ).to('cuda')
    pipe.set_progress_bar_config(disable=True)
    print(f'[load] {time.time() - t0:.1f}s', flush=True)

    for oid in [k for k in LIST if k in wanted]:
        prompt = LIST[oid]
        for seed in seeds_of.get(oid, list(SEEDS)):
            gen = torch.Generator('cuda').manual_seed(seed)
            img = pipe(
                prompt, negative_prompt=NEG, num_inference_steps=STEPS,
                guidance_scale=CFG, width=512, height=512, generator=gen,
            ).images[0]
            out = out_dir / f'{oid}_{seed}.png'
            img.convert('RGBA').save(out)
            print(f'[ok] {out.name}', flush=True)


if __name__ == '__main__':
    main()