Newer
Older
rpg / apps / game / tools / aiaudio / gen_audioldm2.py
"""Генерация AI-аудио (работа A, этап A1): text-to-audio через AudioLDM2 по манифесту.

Манифест — apps/game/tools/aiaudio/manifest.json: промпты, сиды, длительности,
gate — единственный источник истины пайплайна (аналог atlas-manifest.json в
aiart). Модель — AudioLDM2 (cvssp/audioldm2, CC-BY-NC-SA; решение A1 2026-09-09:
игра некоммерческая. Смена модели = новый gen-скрипт с тем же CLI).

Запуск (обычно через runner: npm run aiaudio gen [id[:seed,seed]...]):
  ~/.cache/rpg-ai/venv-aiaudio/bin/python gen_audioldm2.py [id[:seed]...]
Без аргументов — все сиды всех позиций манифеста. Выход — review/<id>_<seed>.wav
(16000 Гц, моно); луп/нормализация — post.py (runner запускает сам).

Детерминизм: seed + шаги фиксированы, промпт в манифесте; детерминизм
по-устройству (fp16 GPU ≠ fp32 CPU). GPU (RTX 3090) при наличии: fp16.
"""
import json
import sys
from pathlib import Path

import numpy as np
import torch

HERE = Path(__file__).resolve().parent
REVIEW = HERE / "review"
MANIFEST = HERE / "manifest.json"

NEG_NOISE = (
    "music, melody, rhythm, drums, speech, voices, sudden loud noises, "
    "harsh noise, distortion"
)

DEFAULT_SEED = 1


def load_items() -> dict[str, dict]:
    """Позиции манифеста по id (промпт/длительность/сиды — источник истины)."""
    data = json.loads(MANIFEST.read_text(encoding="utf-8"))
    return {it["id"]: it for it in data["items"]}


def load_pipeline():
    from diffusers import AudioLDM2Pipeline

    use_gpu = torch.cuda.is_available()
    dtype = torch.float16 if use_gpu else torch.float32
    pipe = AudioLDM2Pipeline.from_pretrained("cvssp/audioldm2", torch_dtype=dtype)
    pipe.to("cuda" if use_gpu else "cpu")
    if not use_gpu:
        torch.set_num_threads(12)
    return pipe, use_gpu


def main() -> None:
    items = load_items()

    # Аргументы — id или id:seed (один сид на аргумент, runner разворачивает сам).
    jobs: list[tuple[str, int]] = []
    for arg in sys.argv[1:]:
        ident, _, seed = arg.partition(":")
        if ident not in items:
            print(f"неизвестный id {ident!r}; доступно: {', '.join(items)}")
            sys.exit(1)
        jobs.append((ident, int(seed) if seed else DEFAULT_SEED))
    if not jobs:
        for ident, it in items.items():
            jobs.extend((ident, s) for s in it.get("seeds", [DEFAULT_SEED]))

    import diffusers  # noqa: F401  (после разбора аргументов — чтобы --help был быстрым)

    pipe, use_gpu = load_pipeline()
    dev = "cuda fp16" if use_gpu else "cpu fp32"
    REVIEW.mkdir(parents=True, exist_ok=True)

    for ident, seed in jobs:
        spec = items[ident]
        torch.manual_seed(seed)
        np.random.seed(seed)
        print(f"генерация {ident} (seed={seed}, {spec['duration']} с, {dev})...")
        audio = pipe(
            spec["positive"],
            negative_prompt=spec.get("negative", NEG_NOISE),
            num_inference_steps=100,
            guidance_scale=3.0,
            audio_length_in_s=spec["duration"],
        ).audios[0]
        audio = np.clip(audio, -1.0, 1.0)
        pcm = (audio * 32767.0).astype(np.int16)
        import wave

        out = REVIEW / f"{ident}_{seed}.wav"
        with wave.open(str(out), "wb") as w:
            w.setnchannels(1)
            w.setsampwidth(2)
            w.setframerate(16000)
            w.writeframes(pcm.tobytes())
        print(f"готово: {out.relative_to(HERE)} ({pcm.shape[0] / 16000:.1f} с)")


if __name__ == "__main__":
    main()