#!/usr/bin/env python3
"""
Seed2 detokenizer -- turns `<seed2_N>` tokens back into an actual image.

Unlike decode_cosmos.py/decode_avclm.py (both have a dedicated neural video
decoder), Seed2Tokenizer (prototype/seed2/, vendored `ontocord/seed2`) is a
Q-Former/BLIP-2-style *understanding* tokenizer -- it has no pixel decoder of
its own. Its own README documents the only reconstruction path: look up each
token's codebook embedding, then condition a diffusion img2img pipeline
(`stabilityai/stable-diffusion-2-1-unclip`) on those embeddings to *generate*
a plausible image. This is fundamentally different from cosmos/avc_lm's
reconstruction (same lossy-but-deterministic neural codec used to encode) --
here the "decode" is itself a generative sample, not a round-trip of the
original pixels. Two different tokens can plausibly decode to visually
different images even if they'd caption similarly, and the same tokens can
decode to slightly different images across runs unless a fixed seed/latent
is used (this script fixes the latent, matching the tokenizer's own
`self.latents` buffer, so decoding is reproducible run-to-run).

Token ids are raw codebook indices, 0-8191, no vocab offset (verified:
tools/tokenizer/build_tokenizers.py's `<seed2_{i}> for i in range(8192)`
and Seed2Tokenizer.num_image_tokens == 8192 agree; unlike snac's +128266
offset or cosmos's chunking, seed2 tokens map directly).

First run downloads `stabilityai/stable-diffusion-2-1-unclip` (~5GB) from
HF -- needs internet access (compute nodes reportedly do not have it, see
REPORT.md's checkpoint-conversion note; run this on a node that does, or
pre-populate $HF_HOME first).

Usage:
    # From a raw list of ints:
    python tools/decode/decode_seed2.py --tokens 6750,2157,4657,... --output out.png

    # From a flattened JSONL / text file containing <seed2_N> tokens anywhere:
    python tools/decode/decode_seed2.py --text-file path.txt --output out.png
"""
import argparse
import os
import re
import sys

PROTOTYPE_DIR = "/e/project1/reformo/nguyen38/prototype"
# 2026-07-23: prototype/ is NOT part of the public github.com/TieuDaoChanNhan/
# finevideo-vla repo (0 files git-tracked), so external users can't reach
# _LOCAL_SEED2_DIR below. Found that Seed2Tokenizer's own vendored README
# (prototype/seed2/README.md) already documents its true public home --
# `git clone https://huggingface.co/ontocord/seed2` -- confirmed for real via
# HfApi().model_info("ontocord/seed2"): same 2 checkpoint files
# (ae.safetensors, model.safetensors) + seed2_tokenizer.py are already public
# there. No redistribution/licensing decision needed on our end (unlike a
# vendor-our-own-checkpoint approach) -- just point at the existing repo.
SEED2_HF_REPO = "ontocord/seed2"
_LOCAL_SEED2_DIR = os.path.join(PROTOTYPE_DIR, "seed2")


def _resolve_seed2_dir() -> str:
    if os.path.isdir(_LOCAL_SEED2_DIR):
        return _LOCAL_SEED2_DIR
    from huggingface_hub import snapshot_download
    print(f"Local seed2 checkpoint not found -- downloading from {SEED2_HF_REPO} "
          f"(~2.6GB, cached for future runs)...")
    return snapshot_download(repo_id=SEED2_HF_REPO)


# stabilityai/stable-diffusion-2-1-unclip returns a genuine 404 (page title
# literally "404 - Hugging Face", not a gated-access page) as of 2026-07-22 --
# confirmed via HF search API that it no longer appears under the stabilityai
# org at all (Stability AI removed it, not access-gated). Using the community
# re-upload instead: same weights/pipeline class (StableUnCLIPImg2ImgPipeline,
# safetensors, openrail++ license), created 2025-11-14 specifically as a mirror.
DIFFUSION_NAME = "sd2-community/stable-diffusion-2-1-unclip"
NUM_IMAGE_TOKENS = 8192
SEED2_QUERY_LEN = 32  # fixed Q-former query length trained into Seed2Tokenizer, 1 image's worth

_SEED2_ATOMIC_RE = re.compile(r"<seed2_(\d+)>")
_SEED2_BLOCK_RE = re.compile(r"<seed2>(.*?)</seed2>", re.DOTALL)


def extract_seed2_tokens(text: str) -> list:
    """Pull every <seed2_N> id inside every <seed2>...</seed2> block, in order,
    concatenated into one flat list. Falls back to scanning the whole text if
    no <seed2>...</seed2> wrapper is present. Only meaningful as one image's
    worth of tokens if the text has exactly one block (or zero, unwrapped) --
    for multi-block text (e.g. a multi-activity generation), use
    extract_seed2_blocks() instead so each image decodes separately."""
    blocks = _SEED2_BLOCK_RE.findall(text)
    source = " ".join(blocks) if blocks else text
    return [int(x) for x in _SEED2_ATOMIC_RE.findall(source)]


def extract_seed2_blocks(text: str) -> list:
    """Like extract_seed2_tokens(), but keeps each <seed2>...</seed2> block's
    ids separate (one list per image) instead of concatenating them -- needed
    because Seed2Tokenizer only ever decodes exactly SEED2_QUERY_LEN=32 tokens
    as one image (see decode_seed2_tokens()). Falls back to treating the whole
    text as one block if no wrapper is present."""
    blocks = _SEED2_BLOCK_RE.findall(text)
    if not blocks:
        return [[int(x) for x in _SEED2_ATOMIC_RE.findall(text)]] if _SEED2_ATOMIC_RE.search(text) else []
    return [[int(x) for x in _SEED2_ATOMIC_RE.findall(b)] for b in blocks]


def _load_seed2_tokenizer():
    """Same import/shim sequence as data_prep/synth_llava/tokenize_seed2.py --
    reused verbatim rather than re-derived (see that file's docstring for why
    each shim exists: a transformers-version helper-function move, and a
    Qformer.cls=None crash).

    NOTE: tokenize_seed2.py only ever calls .encode_image() (only needs
    prototype/pipeline.py's thin Seed2Tokenizer wrapper), but decode needs
    .decode()/.from_pretrained(), which only exist on the real HF PreTrainedModel
    class in seed2_tokenizer.py itself -- pipeline.py's wrapper doesn't have
    them. Copying the import verbatim silently returned the wrong class here
    (caught 2026-07-22: AttributeError, no from_pretrained) -- must return
    _seed2_tokenizer.Seed2Tokenizer directly, not prototype/pipeline.py's.

    2026-07-23: no longer hardcodes PROTOTYPE_DIR -- uses whatever
    _resolve_seed2_dir() finds (local cluster copy, or a fresh download from
    the public ontocord/seed2 HF repo). Returns that dir alongside the class
    so callers know where to point Seed2Tokenizer.from_pretrained().

    chdir target differs by branch: init_tokenizer() inside seed2_tokenizer.py
    does BertTokenizer.from_pretrained("./seed2/bert-base-uncased") -- a
    relative lookup that expects cwd to be seed2_dir's PARENT (matching how
    the local cluster copy is laid out: PROTOTYPE_DIR/seed2/bert-base-uncased).
    Chdir'ing into seed2_dir itself (one level too deep) breaks that lookup --
    caught 2026-07-23 testing the new encode_seed2.py against the local
    branch specifically (OSError: can't load './seed2/bert-base-uncased').
    The downloaded ontocord/seed2 snapshot has no bert-base-uncased subfolder
    at all (verified via its real file listing) yet works anyway -- empirically
    that branch's relative lookup resolves some other way (not fully
    root-caused), so only the local branch needs the parent-dir chdir fix."""
    seed2_dir = _resolve_seed2_dir()
    os.chdir(PROTOTYPE_DIR if seed2_dir == _LOCAL_SEED2_DIR else seed2_dir)

    import transformers.modeling_utils as _modeling_utils
    import transformers.pytorch_utils as _pytorch_utils
    for _name in ("apply_chunking_to_forward", "find_pruneable_heads_and_indices", "prune_linear_layer"):
        if not hasattr(_modeling_utils, _name):
            setattr(_modeling_utils, _name, getattr(_pytorch_utils, _name))

    sys.path.insert(0, seed2_dir)
    import seed2_tokenizer as _seed2_tokenizer

    def _safe_get_output_embeddings(self):
        return None if self.cls is None else self.cls.predictions.decoder

    def _safe_set_output_embeddings(self, new_embeddings):
        if self.cls is not None:
            self.cls.predictions.decoder = new_embeddings

    for _cls in (_seed2_tokenizer.BertLMHeadModel, _seed2_tokenizer.BertForMaskedLM):
        _cls.get_output_embeddings = _safe_get_output_embeddings
        _cls.set_output_embeddings = _safe_set_output_embeddings

    return _seed2_tokenizer.Seed2Tokenizer, seed2_dir


def decode_seed2_tokens(token_ids: list, output_path: str, guidance_scale: float = 10.0,
                         num_inference_steps: int = 20) -> None:
    if not token_ids:
        raise ValueError("No seed2 tokens to decode")
    bad = [t for t in token_ids if not (0 <= t < NUM_IMAGE_TOKENS)]
    if bad:
        raise ValueError(f"Token ids out of range [0, {NUM_IMAGE_TOKENS}): {bad[:5]}...")
    # Seed2Tokenizer's Q-former was trained with a fixed 32 query tokens/image
    # (see pos_embed_image.repeat(query_output_up.shape[0], ...) in
    # seed2_tokenizer.py -- shape[0] must be 1 image's worth of queries, i.e.
    # len(token_ids)==32). Any other count desyncs from the diffusion
    # pipeline's own batch-size assumption deep in the UNet (hit for real
    # 2026-07-22: 3-token and 96-token spans both crashed with a "tensor a
    # must match tensor b" RuntimeError -- see
    # samples/qwen3_1.7b_vla_v2_eval/2026-07-22_full_eval/SUMMARY.md, tests
    # 02_agent_continuation and 07_full_chain_from_scratch). Fail clearly
    # instead of letting that opaque error surface from inside the UNet.
    if len(token_ids) != SEED2_QUERY_LEN:
        raise ValueError(
            f"Got {len(token_ids)} seed2 tokens, but Seed2Tokenizer only decodes exactly "
            f"{SEED2_QUERY_LEN} tokens at a time (1 image's fixed Q-former query length). "
            f"If this span was extracted from text containing multiple <seed2>...</seed2> blocks "
            f"(e.g. a multi-activity generation), decode each block separately."
        )

    import torch
    from diffusers import StableUnCLIPImg2ImgPipeline

    Seed2Tokenizer, seed2_dir = _load_seed2_tokenizer()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32

    print(f"Loading {DIFFUSION_NAME} (first run downloads ~5GB)...")
    pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(DIFFUSION_NAME, torch_dtype=dtype).to(device)

    print("Loading Seed2Tokenizer...")
    tokenizer = Seed2Tokenizer.from_pretrained(seed2_dir, torch_dtype=dtype).to(device)

    # get_codebook_entry() does self.embedding(indices) with no batch handling of
    # its own -- dim 0 of `indices` becomes the batch dim downstream (repeat(),
    # etc.). A flat (N,) tensor is silently read as N separate 1-token images
    # instead of 1 image made of N query tokens, which desyncs from the diffusion
    # pipeline's own (batch=1, x2 for CFG) conditioning shape and crashes deep in
    # the UNet (caught 2026-07-22: "tensor a (2) must match tensor b (64)" for a
    # 32-token input, i.e. 32 x 2 leaking through). Needs an explicit batch dim.
    indices = torch.tensor(token_ids, dtype=torch.long, device=device).unsqueeze(0)
    image = tokenizer.decode(pipe, indices, guidance_scale=guidance_scale,
                              num_inference_steps=num_inference_steps)[0]
    image.save(output_path)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--tokens", help="Comma-separated list of raw <seed2_N> ids (0-8191)")
    ap.add_argument("--text-file", help="Plain text / JSONL file containing <seed2_N> tokens anywhere in it")
    ap.add_argument("--guidance-scale", type=float, default=10.0)
    ap.add_argument("--num-inference-steps", type=int, default=20)
    ap.add_argument("--output", required=True)
    args = ap.parse_args()

    # _load_seed2_tokenizer() os.chdir()s into PROTOTYPE_DIR (needs to be cwd for
    # its own relative "./seed2" lookups) -- resolve a relative --output against
    # the original cwd *before* that happens, or it silently lands in prototype/
    # instead (caught 2026-07-22: full generation succeeded, only image.save() failed).
    output_path = os.path.abspath(args.output)

    if args.tokens:
        blocks = [[int(x) for x in args.tokens.split(",")]]
    elif args.text_file:
        blocks = extract_seed2_blocks(open(args.text_file, encoding="utf-8").read())
    else:
        ap.error("Provide --tokens or --text-file")

    if not blocks:
        ap.error("No <seed2_N> tokens found")

    if len(blocks) == 1:
        print(f"Decoding {len(blocks[0])} seed2 tokens...")
        decode_seed2_tokens(blocks[0], output_path, args.guidance_scale, args.num_inference_steps)
        print(f"Saved: {output_path}")
        return

    # Multiple <seed2>...</seed2> blocks (e.g. a multi-activity generation) --
    # decode each as its own image rather than concatenating and erroring
    # (added 2026-07-22 after a 96-token/3-block span crashed decode_seed2_tokens).
    print(f"Found {len(blocks)} separate <seed2>...</seed2> blocks -- decoding each as its own image.")
    stem, ext = os.path.splitext(output_path)
    for i, block in enumerate(blocks):
        block_output = f"{stem}_{i}{ext}"
        print(f"\nBlock {i}: {len(block)} tokens -> {block_output}")
        try:
            decode_seed2_tokens(block, block_output, args.guidance_scale, args.num_inference_steps)
            print(f"Saved: {block_output}")
        except ValueError as e:
            print(f"Skipped block {i}: {e}")


if __name__ == "__main__":
    main()