scripts/ltx_conditioning_dump.py

"""Write the conditioning tensor an LTX-2.5 encoder produces, so a DiT can use it.

The gate (`ltx_prompt_embedding_gate.py`) scores an encoder without rendering.
This is its sibling and shares its entire forward path: same tokenization, same
49-slot hidden stack, same `apply_aggregates`. The difference is the output -
the gate compares two encoders and prints a verdict, this writes the tensor.

Why it has to exist at all: ComfyUI cannot load our compressed encoder. Its GGUF
path gates text encoders on `TXT_ARCH_LIST`, which carries `gemma3` but not
`gemma4`, and its native quantisation kernels need compute capability 8.9. On a
V100 there is no route from a compressed LTX-2.5 encoder into a ComfyUI graph.
So the encoder runs here, out of process, and hands over the one tensor the DiT
actually consumes.

That tensor is `torch.cat((video, audio), dim=-1)`, [T, 6144] - exactly what
`LTXAVTEModel.encode_token_weights` returns, and everything downstream treats it
as one opaque block. `apply_aggregates` already produces the two halves.

Running out of process also solves the memory problem rather than working around
it. The encoder is 26 GB of BF16 and the DiT is 10.6 GB; neither fits beside the
other on a 16 GB card. Here the encoder loads, emits, and exits, and the card is
empty again before the DiT is asked for.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from ltx_prompt_embedding_gate import (
    apply_aggregates,
    assert_fully_loaded,
    check_aggregates_present,
    hidden_stack,
    hidden_state_artifact,
    register_ltx_renames,
    weights_path,
)

from pacific.ltx_av import load_event_set

#: ComfyUI pads every LTX prompt to this length on the left
#: (`Gemma3_12BTokenizer`, `comfy/text_encoders/lt.py:76`: `min_length=1024`,
#: `pad_left=True`). It is also exactly the encoder's `sliding_window`, so the
#: real tokens sit at the top of a window that is otherwise entirely padding.
COMFY_MIN_LENGTH = 1024

#: `special_tokens={"start": 2, "pad": 0}` in the same constructor. The HF
#: tokenizer for this checkpoint adds neither - verified, its first id for
#: "A humanoid robot..." is 236776, not 2 - so the start token is prepended here
#: or the encoder never sees a sentence begin.
COMFY_BOS = 2
COMFY_PAD = 0


def hidden_stack_comfy(
    model, tokenizer, prompts: list[str], device: str, min_length: int = COMFY_MIN_LENGTH
) -> list:
    """Per prompt, the raw [49, T, 3840] Transformer stack for ComfyUI tokens.

    `hidden_stack` in the gate runs the prompt unpadded, which is the right
    choice for comparing two encoders to each other: both sides get the same
    treatment and the difference between them is the only thing measured. It is
    the wrong choice for feeding a DiT, because the DiT was trained on what the
    vendor pipeline emits, and that differs in two ways:

    * a start token. ComfyUI tokenizes with `add_bos=True`; this checkpoint's HF
      tokenizer does not add one.
    * position. ComfyUI left-pads to 1024 and then takes `position_ids =
      arange(seq_len)` over the *padded* sequence (`llama.py:683`), so a
      68-token prompt occupies RoPE positions 956..1023. Unpadded it occupies
      0..67. Same tokens, different angles in every rotary layer.

    Masking makes the padding itself inert - `enable_attention_masks=True` in
    `Gemma3_12BModel` - so this is not about the pad values. It is about where
    on the rotary circle the real tokens land, and the DiT has an opinion.

    Transformers already final-normalizes only the last returned slot. Preserve
    every raw intermediate slot, then select the real-token suffix.

    The trailing slice reproduces `out[:, :, -sum(attention_mask):]` from
    `LTXAVTEModel.encode_token_weights`.
    """
    import torch

    out = []
    for prompt in prompts:
        ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
        ids = [COMFY_BOS, *ids]
        real = len(ids)
        if real > min_length:
            raise SystemExit(
                f"prompt is {real} tokens, past the {min_length} ComfyUI pads to; "
                "the vendor path would not have padded it either"
            )

        pad = min_length - real
        input_ids = torch.tensor([[COMFY_PAD] * pad + ids], device=device)
        attention_mask = torch.tensor([[0] * pad + [1] * real], device=device)
        # Passed explicitly rather than left to `transformers`, which derives
        # positions from the mask in some versions - that would put the real
        # tokens back at 0..T-1 and undo the entire point of padding here.
        position_ids = torch.arange(min_length, device=device).unsqueeze(0)

        with torch.no_grad():
            states = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                output_hidden_states=True,
            ).hidden_states
            stacked = torch.stack([state[0, -real:] for state in states])
        out.append(stacked.to(torch.float32).cpu())
    return out


def load_prompt_file(path: Path) -> list[str]:
    """One prompt per line; blank lines and `#` comments ignored.

    Prompts are long enough that a multi-line format would be nicer to read and
    far easier to get wrong - a stray blank line silently splits one prompt into
    two, and the resulting clip is of half a sentence. One line each is the
    format that cannot do that.
    """
    if path.suffix == ".jsonl":
        return [event["prompt"] for event in load_event_set(path)]
    lines = [line.strip() for line in path.read_text().splitlines()]
    return [line for line in lines if line and not line.startswith("#")]


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model-dir", required=True, help="encoder dir: config, tokenizer and model.safetensors"
    )
    parser.add_argument("--prompts", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--label", default="", help="build name recorded in the file, e.g. r45c-gptq"
    )
    parser.add_argument("--checkpoint-sha256", default="")
    parser.add_argument("--profile-hash", default="")
    parser.add_argument("--audio-aggregate-checkpoint", type=Path)
    parser.add_argument("--video-aggregate-checkpoint", type=Path)
    parser.add_argument("--hidden-stacks-output", type=Path)
    parser.add_argument(
        "--device-map",
        default="auto",
        help="26 GB of BF16 does not fit 16 GB of VRAM; 'auto' splits it",
    )
    parser.add_argument(
        "--packed-checkpoint",
        type=Path,
        help="load weights from a ltx_packed_codec artifact "
        "instead of model.safetensors; --model-dir still "
        "supplies config and tokenizer",
    )
    parser.add_argument(
        "--resident",
        action="store_true",
        help="keep the packed weights packed and dequantize in "
        "forward: 8.5 GB resident instead of 26.3 GB",
    )
    parser.add_argument(
        "--fold-pre-scale",
        action="store_true",
        help="with --resident, apply the AWQ scale to the "
        "activation instead of the weight; faster and not "
        "bit-identical",
    )
    parser.add_argument(
        "--comfy-padding",
        action="store_true",
        help="reproduce ComfyUI's BOS and 1024-token left pad; "
        "required for output a DiT will accept",
    )
    args = parser.parse_args()

    import torch
    from transformers import AutoModel, AutoTokenizer

    prompts = load_prompt_file(args.prompts)
    prompt_ids = (
        [event["id"] for event in load_event_set(args.prompts)]
        if args.prompts.suffix == ".jsonl"
        else [f"prompt-{index:02d}" for index in range(len(prompts))]
    )
    if not prompts:
        raise SystemExit(f"{args.prompts} has no prompts")
    print(f"{len(prompts)} prompts from {args.prompts}", flush=True)

    model_dir = Path(args.model_dir)
    checkpoint = args.packed_checkpoint or weights_path(model_dir)
    check_aggregates_present(checkpoint)

    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
    if args.packed_checkpoint:
        from ltx_packed_codec import load_packed_model

        model = load_packed_model(
            args.model_dir,
            args.packed_checkpoint,
            resident=args.resident,
            fold_pre_scale=args.fold_pre_scale,
        )
    else:
        register_ltx_renames()
        model, info = AutoModel.from_pretrained(
            args.model_dir,
            dtype=torch.bfloat16,
            device_map=args.device_map,
            output_loading_info=True,
        )
        assert_fully_loaded(info)
    model.eval()
    device = str(next(model.parameters()).device)
    print(f"{args.model_dir}: loaded on {device}", flush=True)

    capture = hidden_stack_comfy if args.comfy_padding else hidden_stack
    print(f"forward path: {capture.__name__}", flush=True)
    stacks = capture(model, tokenizer, prompts, device)
    slots, _, hidden = stacks[0].shape
    del model
    torch.cuda.empty_cache()
    print(f"{len(stacks)} stacks, {slots} slots x {hidden}", flush=True)
    if args.hidden_stacks_output is not None:
        args.hidden_stacks_output.parent.mkdir(parents=True, exist_ok=True)
        torch.save(
            hidden_state_artifact(
                prompt_ids=prompt_ids,
                prompts=prompts,
                comfy_padding=args.comfy_padding,
                hidden_stacks=stacks,
            ),
            args.hidden_stacks_output,
        )
        print(f"wrote {args.hidden_stacks_output}", flush=True)

    branch_checkpoints = {
        branch: path
        for branch, path in {
            "audio": args.audio_aggregate_checkpoint,
            "video": args.video_aggregate_checkpoint,
        }.items()
        if path is not None
    }
    for source in branch_checkpoints.values():
        check_aggregates_present(source)
    branches = apply_aggregates(stacks, checkpoint, branch_checkpoints)

    # The DiT sees one [T, 6144] block. Keeping the halves separate here would
    # push the concatenation - and the chance of getting its order backwards -
    # into the render script, where it would be checked by nothing.
    conditioning = []
    for prompt, branch in zip(prompts, branches, strict=True):
        merged = torch.cat((branch["video"], branch["audio"]), dim=-1)
        conditioning.append(merged.to(torch.float32).contiguous())
        print(f"  {tuple(merged.shape)}  {prompt[:60]}", flush=True)

    widths = {tuple(c.shape[1:]) for c in conditioning}
    if widths != {(6144,)}:
        raise SystemExit(f"expected a 6144-wide conditioning, got {widths}")

    args.output.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        hidden_state_artifact(
            label=args.label or model_dir.name,
            model_dir=str(model_dir),
            checkpoint=str(checkpoint),
            checkpoint_sha256=args.checkpoint_sha256,
            profile_hash=args.profile_hash,
            aggregate_checkpoints={
                branch: str(path) for branch, path in branch_checkpoints.items()
            },
            comfy_padding=args.comfy_padding,
            prompt_ids=prompt_ids,
            prompts=prompts,
            conditioning=conditioning,
        ),
        args.output,
    )
    print(f"wrote {args.output}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())