"""ArtiFixer 캡션 임베딩 (caption.h5) — artifixer 환경에서 실행한다 (UMT5-XXL 텍스트 인코더, CPU).

  ~/micromamba/envs/artifixer/bin/python scripts/artifixer_caption.py outputs/artifixer/<준비 디렉터리> [...]

형식은 artifixer data_processing/captioning/generate_captions.py 와 같다: 데이터셋 하나 = bf16 임베딩을 uint16 으로 본 배열
(토큰 수 × 4096), attrs caption · image_indices. 텍스트 인코더는 22 GB(fp32 파일) — bf16 으로 CPU 에 올린다 (VRAM 은 쓰지 않음).
"""

import sys
from pathlib import Path

import h5py
import torch

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "third_party/artifixer"))
from data_processing.captioning.generate_captions import generate_text_embedding  # noqa: E402
from transformers import AutoTokenizer, UMT5EncoderModel  # noqa: E402

MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"


def main(dirs: list[str]) -> None:
    tok = AutoTokenizer.from_pretrained(MODEL_ID, subfolder="tokenizer")
    enc = UMT5EncoderModel.from_pretrained(MODEL_ID, subfolder="text_encoder", torch_dtype=torch.bfloat16).eval()
    for d in map(Path, dirs):
        cap = (d / "caption.txt").read_text().strip()
        with torch.inference_mode():
            emb = generate_text_embedding(cap, enc, tok, 512)
        with h5py.File(d / "caption.h5", "w") as hf:
            ds = hf.create_dataset("frames_0", data=emb)
            ds.attrs["caption"] = cap
            ds.attrs["image_indices"] = [0]
        print(f"{d / 'caption.h5'}: {emb.shape} — {cap[:60]}")


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