#!/usr/bin/env python3
"""Инференс MOSS-Image без текстового энкодера 9B: 0.8B + адаптер вместо 15.6 ГБ.

Условие для DiT считается здесь и уходит в пайплайн официальными аргументами
(`prompt_embeds` / `negative_prompt_embeds`), поэтому 15.6 ГБ учителя не нужны вовсе.
CFG включён (guidance 3.0), негатив по умолчанию пустой — его тоже считает адаптер.

  python3 src/example.py --prompt "a red fox in a snowy forest at dusk, 85mm" --out fox.png
  python3 src/example.py --prompt '{"layout":[{"text":"SALE","box":[40,40,420,120]}]}' --out sale.png
"""
import argparse
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from adapter_lib import load_adapter                      # noqa: E402
from moss_lib import (MOSS_DIR, STUDENT, Student, build_pipeline, condition,   # noqa: E402
                      empty_cache, generate, load_tokenizer)

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEVICE = "cuda"


def main():
    ap = argparse.ArgumentParser(description="MOSS-Image с адаптером 0.8B вместо энкодера 9B")
    ap.add_argument("--prompt", required=True)
    ap.add_argument("--negative", default="", help="негативный промпт (по умолчанию пустой, как в пайплайне)")
    ap.add_argument("--out", default="out.png")
    ap.add_argument("--adapter", default=next(
        (p for p in (os.path.join(ROOT, "adapter.safetensors"),
                     os.path.join(ROOT, "out", "adapter_v1.safetensors")) if os.path.exists(p)),
        os.path.join(ROOT, "adapter.safetensors")))
    ap.add_argument("--moss-dir", default=MOSS_DIR)
    ap.add_argument("--student", default=STUDENT)
    ap.add_argument("--size", type=int, default=1024)
    ap.add_argument("--steps", type=int, default=30)
    ap.add_argument("--guidance", type=float, default=3.0)
    ap.add_argument("--seed", type=int, default=1234)
    args = ap.parse_args()

    tok = load_tokenizer(args.moss_dir)
    adapter, meta = load_adapter(args.adapter, device=DEVICE)
    print(f"адаптер: {sum(p.numel() for p in adapter.parameters()) / 1e6:.1f}M, шаг "
          f"{meta.get('step', '?')}, срезы {meta.get('student_layers')}", flush=True)
    student = Student(args.student, [int(x) for x in str(meta["student_layers"]).split(",")])

    pos = condition(adapter, student, tok, [args.prompt])[0]
    neg = condition(adapter, student, tok, [args.negative])[0]
    n_tok = pos[1].sum().item()
    print(f"условие: {n_tok} токенов (негатив {neg[1].sum().item()})", flush=True)
    if n_tok > int(meta.get("attn_max_len", 10 ** 6)):
        raise SystemExit(f"в промпте {n_tok} токенов, а таблица позиций ветки внимания "
                         f"обучена до {meta['attn_max_len']} — сократите промпт")
    del student, adapter                              # 0.8B и адаптер больше не нужны
    empty_cache()

    pipe = build_pipeline(args.moss_dir)
    pipe.to(DEVICE)
    print(f"DiT на карте, VRAM {torch.cuda.memory_allocated() / 2**30:.1f} ГиБ", flush=True)
    image = generate(pipe, pos, neg, size=args.size, steps=args.steps,
                     seed=args.seed, guidance=args.guidance)
    image.save(args.out)
    print(f"{args.size}px, {args.steps} шагов, сид {args.seed} -> {args.out} {image.size}, "
          f"пик VRAM {torch.cuda.max_memory_allocated() / 2**30:.1f} ГиБ", flush=True)


if __name__ == "__main__":
    main()