#!/usr/bin/env python3
"""AI-шаг пайплайна атласов: генерация листа одного ракурса (SD1.5, CPU).

Модель: SD_PixelArt_SpriteSheet_Generator (Onodofthenorth, Apache-2.0).
Токены ракурсов: PixelartFSS (фронт), PixelartRSS (право), PixelartBSS (зад),
PixelartLSS (лево). 15 шагов / cfg 7.0 — рабочий режим пробы B3 (~6.6 мин/вид
на 4 потоках fp32). Промпт персонажа приходит из манифеста пайплайна.

Запуск (интерпретатор — env RPG_SD_PYTHON, дефолт venv-pxgpt):
  $RPG_SD_PYTHON apps/game/tools/aiart/sd_sheet.py <view> <seed> <outdir> <prompt>
"""
import sys
import time
from pathlib import Path

import torch
from diffusers import StableDiffusionPipeline
from huggingface_hub import snapshot_download

MODEL_ID = 'Onodofthenorth/SD_PixelArt_SpriteSheet_Generator'
STEPS = 15
CFG = 7.0
NEG = ('blurry, photo, realistic, 3d render, smooth gradients, text, watermark, '
       'multiple characters, cropped')


def main() -> None:
    view, seed, out_dir = sys.argv[1], int(sys.argv[2]), Path(sys.argv[3])
    prompt = sys.argv[4]
    out_dir.mkdir(parents=True, exist_ok=True)

    torch.set_num_threads(4)
    base = snapshot_download(MODEL_ID, local_files_only=True)
    t0 = time.time()
    pipe = StableDiffusionPipeline.from_pretrained(str(base), safety_checker=None)
    pipe.set_progress_bar_config(disable=True)
    print(f'[load] {time.time() - t0:.0f}s', flush=True)

    t0 = time.time()
    img = pipe(
        f'Pixelart{view} of {prompt}',
        negative_prompt=NEG,
        num_inference_steps=STEPS,
        guidance_scale=CFG,
        generator=torch.Generator('cpu').manual_seed(seed),
    ).images[0]
    print(f'[gen] {time.time() - t0:.0f}s', flush=True)

    out = out_dir / f'{view}_{seed}.png'
    img.convert('RGBA').save(out)
    print(f'[ok] {out}', flush=True)


if __name__ == '__main__':
    main()