code/tests/test_arm_prompt.py

"""prompt 臂（Incantation 移植，baseline/prompt.py）的自测：baseline/SPEC.md 的测试 1–4 + 模板测试。

    CUDA_VISIBLE_DEVICES=6 PYTHONPATH=/opt/dlami/nvme/zhiyangdeng/ActionRoPE DIFFSYNTH_SKIP_DOWNLOAD=True \
        /opt/dlami/nvme/zhiyangdeng/ActionRoPE/.venv/bin/python -m pytest \
        /opt/dlami/nvme/zhiyangdeng/ActionRoPE/tests/test_arm_prompt.py -s -v

这个臂零参数、零钩子，动作全在逐 cell 文本里，所以 SPEC 测试 2 的"扰动新增参数"换成"改动作从句"：
文本表（text_table_eybx_mirror.pt，mmap 打开）里真实的 UMT5 嵌入，同场景下换动作词输出必须变；
而 action_inputs 里的偏移必须对输出没有任何影响（这个臂不读它）。
实测数字追加写到 $AROPE_TEST_NUMBERS（默认 outputs/test_artifacts/arm_prompt_numbers.json）。
"""
from __future__ import annotations

import json
import os
import sys
import time

import pytest
import torch

os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")

ROOT = "/opt/dlami/nvme/zhiyangdeng/ActionRoPE"
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)
MODEL_DIR = f"{ROOT}/models/Wan2.2-TI2V-5B"
DIT_FILES = [f"{MODEL_DIR}/diffusion_pytorch_model-0000{i}-of-00003.safetensors" for i in (1, 2, 3)]
MIRROR_TABLE = f"{ROOT}/data/latent/text_table_eybx_mirror.pt"
TEST_DIR = os.environ.get("AROPE_TEST_DIR", os.path.join(ROOT, "outputs", "test_artifacts"))
NUMBERS_PATH = os.environ.get("AROPE_TEST_NUMBERS", os.path.join(TEST_DIR, "arm_prompt_numbers.json"))

C, F, LAT_H, LAT_W = 48, 21, 30, 52
L_TXT, D_TXT = 512, 4096
SCENE = "tidal_flats"


def record(**kv):
    data = {}
    try:
        with open(NUMBERS_PATH) as fp:
            data = json.load(fp)
    except (FileNotFoundError, json.JSONDecodeError):
        pass
    data.update({k: (float(v) if isinstance(v, (int, float)) and not isinstance(v, bool) else v) for k, v in kv.items()})
    os.makedirs(os.path.dirname(NUMBERS_PATH), exist_ok=True)
    with open(NUMBERS_PATH, "w") as fp:
        json.dump(data, fp, indent=2, ensure_ascii=False)
    print("\n[numbers]", json.dumps(kv, ensure_ascii=False))


def err_stats(out, ref):
    out, ref = out.float(), ref.float()
    diff = (out - ref).abs()
    return {"max_abs": diff.max().item(), "rel_l2": (diff.norm() / ref.norm().clamp_min(1e-12)).item()}


# ---------------------------------------------------------------- fixtures

@pytest.fixture(scope="module")
def dit():
    from diffsynth.pipelines.wan_video import ModelConfig, WanVideoPipeline
    from actionrope.arope import install_arope

    pipe = WanVideoPipeline.from_pretrained(
        torch_dtype=torch.bfloat16, device="cuda",
        model_configs=[ModelConfig(path=DIT_FILES)],
        tokenizer_config=None, redirect_common_files=False,
    )
    model = pipe.dit
    model.eval().requires_grad_(False)
    model._baseline_keys = set(model.state_dict().keys())
    # integration 的接法：普通 RoPE、无 mask 通道（SPEC「训练 / 推理接入」）
    install_arope(model, mask_channel=False)
    return model


@pytest.fixture(scope="module")
def arm(dit):
    from baseline.prompt import PromptArm

    a = PromptArm()
    a.install(dit)
    a.to("cuda")
    return a


@pytest.fixture(scope="module")
def texts():
    """同一场景、两种动作词的真实 UMT5 嵌入（查 mirror 表）。表不在 ⇒ 用随机张量代替（形状契约不变）。"""
    from baseline.prompt import script_to_cell_prompts

    p_right = script_to_cell_prompts(SCENE, [["right", 6.0]])           # 21 × moving right
    p_mix = script_to_cell_prompts(SCENE, [["right", 2.75], ["left", 3.0]])  # 11 × right + 10 × left
    if os.path.exists(MIRROR_TABLE):
        table = torch.load(MIRROR_TABLE, map_location="cpu", weights_only=False, mmap=True)
        emb = {p: table[p] for p in set(p_right) | set(p_mix)}
        real = True
    else:
        g = torch.Generator(device="cpu").manual_seed(7)
        emb = {p: torch.randn(L_TXT, D_TXT, generator=g).to(torch.bfloat16) for p in set(p_right) | set(p_mix)}
        real = False
    uniq = sorted(emb)
    def ctx_of(prompts):
        ctx = torch.stack([emb[p] for p in prompts], 0).unsqueeze(0).to("cuda", torch.bfloat16)  # [1, 21, 512, 4096]
        ids = torch.tensor([[uniq.index(p) for p in prompts]])
        return ctx.contiguous(), ids
    return dict(right=ctx_of(p_right), mix=ctx_of(p_mix), real=real, prompts_right=p_right, prompts_mix=p_mix)


def make_latents(seed=0, device="cuda"):
    g = torch.Generator(device="cpu").manual_seed(seed)
    latents = torch.randn(1, C, F, LAT_H, LAT_W, generator=g).to(device=device, dtype=torch.bfloat16)
    timestep = torch.tensor([500.0], device=device, dtype=torch.bfloat16)
    return latents, timestep


def action_inputs(offset_px):
    """SPEC 的 action_inputs dict（这个臂不读，但按契约完整给出）。"""
    from actionrope.prompts import ACTIONS  # noqa: F401
    off_tok = offset_px / 32
    delta = torch.zeros_like(off_tok)
    delta[:, 1:] = off_tok[:, 1:] - off_tok[:, :-1]
    return {"offset_px": offset_px, "offset_tok": off_tok, "delta_tok": delta,
            "action_idx": torch.full((offset_px.shape[0], F), 7, dtype=torch.long, device=offset_px.device)}


# ---------------------------------------------------------------- 1. 零初始化等价（逐位）

@torch.no_grad()
def test_1_install_is_bitwise_identity(dit, arm, texts):
    from actionrope.arope import arope_forward

    latents, timestep = make_latents()
    ctx, ids = texts["right"]
    off = torch.zeros(1, F, 2, device="cuda")
    off[0, :, 0] = torch.linspace(0, 96, F)
    ref = arope_forward(dit, latents, timestep, ctx, context_ids=ids)
    out = arope_forward(dit, latents, timestep, ctx, context_ids=ids, arm=arm, action_inputs=action_inputs(off))
    assert out.shape == (1, C, F, LAT_H, LAT_W) and torch.isfinite(out).all()
    s = err_stats(out, ref)
    record(equiv_max_abs=s["max_abs"], equiv_rel_l2=s["rel_l2"], equiv_bitwise=bool(torch.equal(out, ref)),
           text_from_table=texts["real"])
    assert torch.equal(out, ref)

    # 所有 cell 文本相同 ⇒ 与整条 context 等价（bf16 容差；SDPA 的 batch 切分不同）——
    # 但 require_per_cell 默认拒收整条 context，先确认它会报错，再用放行版比对
    with pytest.raises(ValueError):
        arope_forward(dit, latents, timestep, ctx[:, 0], arm=arm, action_inputs=None)
    from baseline.prompt import PromptArm
    lax = PromptArm(require_per_cell=False)
    whole = arope_forward(dit, latents, timestep, ctx[:, 0], arm=lax, action_inputs=None)
    s2 = err_stats(out, whole)
    record(percell_vs_whole_rel_l2=s2["rel_l2"])
    assert s2["rel_l2"] <= 1e-2


# ---------------------------------------------------------------- 2. 动作信号有效 + 梯度

def test_2_action_text_changes_output_and_grads_flow(dit, arm, texts):
    from actionrope.arope import arope_forward

    latents, timestep = make_latents(seed=1)
    ctx_r, ids_r = texts["right"]
    ctx_m, ids_m = texts["mix"]
    off = torch.zeros(1, F, 2, device="cuda")
    off[0, :, 0] = torch.linspace(0, 96, F)

    with torch.no_grad():
        o_right = arope_forward(dit, latents, timestep, ctx_r, context_ids=ids_r, arm=arm, action_inputs=action_inputs(off))
        o_mix = arope_forward(dit, latents, timestep, ctx_m, context_ids=ids_m, arm=arm, action_inputs=action_inputs(off))
        # 这个臂的"动作"就是文本：cell 11–20 换成 moving left，输出必须变，且量级不变
        s = err_stats(o_mix, o_right)
        assert s["rel_l2"] > 1e-3, s
        assert abs(o_mix.float().std().item() / o_right.float().std().item() - 1) < 0.2
        # 诊断量（不断言）：换了文本的 cell 11–20 与没换的 cell 0–10 的均方差之比。基座在随机 latent 上
        # 差异经自注意力扩散到全部帧（实测 ≈0.7，未训练时没有局部性），只记录供训练后对照
        d = (o_mix.float() - o_right.float()).pow(2).mean(dim=(0, 1, 3, 4))  # [F]
        rel_changed = d[11:].mean().item() / d[:11].mean().item()
        # 非文本的动作输入必须完全不起作用（offset 改 10 倍、action_idx 改成 idle，输出逐位不变）
        ai2 = action_inputs(off * 10)
        ai2["action_idx"].zero_()
        o_off = arope_forward(dit, latents, timestep, ctx_r, context_ids=ids_r, arm=arm, action_inputs=ai2)
        assert torch.equal(o_off, o_right)
        o_none = arope_forward(dit, latents, timestep, ctx_r, context_ids=ids_r, arm=arm, action_inputs=None)
        assert torch.equal(o_none, o_right)
    record(perturbed_rel_l2=s["rel_l2"], perturbed_max_abs=s["max_abs"],
           changed_cells_vs_unchanged_msd_ratio=rel_changed, offset_ignored_bitwise=True)

    # 训练一步：梯度检查点开，梯度回到 DiT 全部参数；臂没有参数（n_new_params = 0）
    assert list(arm.parameters()) == []
    try:
        dit.train().requires_grad_(True)
        torch.cuda.synchronize()
        torch.cuda.reset_peak_memory_stats()
        t0 = time.time()
        out = arope_forward(dit, latents, timestep, ctx_m, context_ids=ids_m, arm=arm, action_inputs=action_inputs(off),
                            use_gradient_checkpointing=True)
        target = torch.randn_like(out)
        # 与 SPEC 的 loss 形式一致（帧 0 不计；权重全 1）
        loss = ((out[:, :, 1:].float() - target[:, :, 1:].float()) ** 2).mean()
        loss.backward()
        torch.cuda.synchronize()
        dt = time.time() - t0
        peak_gb = torch.cuda.max_memory_allocated() / 1024 ** 3
        assert torch.isfinite(loss)
        checks = {
            "patch_embedding.weight": dit.patch_embedding.weight,
            "blocks.0.self_attn.q.weight": dit.blocks[0].self_attn.q.weight,
            "blocks.15.cross_attn.k.weight": dit.blocks[15].cross_attn.k.weight,
            "blocks.29.ffn.2.weight": dit.blocks[29].ffn[2].weight,
            "head.head.weight": dit.head.head.weight,
            "text_embedding.0.weight": dit.text_embedding[0].weight,
        }
        grad_norms = {}
        for name, p in checks.items():
            assert p.grad is not None, name
            grad_norms[name] = p.grad.float().norm().item()
            assert grad_norms[name] > 0, name
        n_with_grad = sum(1 for p in dit.parameters() if p.grad is not None and p.grad.abs().sum() > 0)
        n_total = sum(1 for p in dit.parameters())
        grad_ok = n_with_grad == n_total
        record(loss=loss.item(), peak_mem_gb=peak_gb, fwd_bwd_sec=dt, grad_norms=grad_norms,
               dit_params_with_nonzero_grad=f"{n_with_grad}/{n_total}", grad_ok=grad_ok,
               arm_params_with_grad="0/0 (臂无参数)")
        assert grad_ok
    finally:
        for p in dit.parameters():
            p.grad = None
        dit.eval().requires_grad_(False)
        torch.cuda.empty_cache()


# ---------------------------------------------------------------- 3. 参数量

def test_3_param_count(arm):
    n = arm.n_new_params()
    n_direct = sum(p.numel() for p in arm.parameters())
    # 对照：Incantation 的 student 与 WanModel 权重同构（causal_model.py:9, :386），动作走文本 ⇒ 0；
    # linear（ReactiveGWM）≈ 0.18M、xattn（Matrix-Game 3）≈ 98.6M 只是量级参照
    record(n_new_params=n, n_new_params_direct=n_direct, upstream_new_params=0, ref_linear=180_000, ref_xattn=98_600_000)
    print(f"\n[params] prompt 臂新增参数 {n}（上游 Incantation 亦为 0；linear ≈ 0.18M，xattn ≈ 98.6M）")
    assert n == 0 and n_direct == 0
    assert arm.zero_init_check()


# ---------------------------------------------------------------- 4. state_dict 键集

def test_4_state_dict_keys(dit, arm):
    sd = arm.state_dict()
    assert len(sd) == 0, sorted(sd)
    prefixed = {f"arm.{k}" for k in sd}
    assert not (prefixed & dit._baseline_keys)
    # strict 加载空 state_dict 成功；DiT 键集与原模型一致（mask_channel=False）
    arm.load_state_dict({}, strict=True)
    assert set(dit.state_dict().keys()) == dit._baseline_keys
    from baseline.prompt import PromptArm
    fresh = PromptArm()
    fresh.load_state_dict(sd, strict=True)
    record(state_dict_keys=sorted(sd), dit_keys_unchanged=True)


# ---------------------------------------------------------------- 5. 推理模板（CPU）

def test_5_prompt_template_matches_dataset_and_table():
    from baseline.prompt import (CELLS, DEFAULT_DURATION_S, TIME_STEP_S, script_to_action_spec,
                                 script_to_cell_prompts, script_to_timeline)
    from actionrope.prompts import mirror_action, split_prompt

    assert TIME_STEP_S == 0.25 and DEFAULT_DURATION_S == 1.0  # Incantation inference.py:85, :140
    # 展开规则：round(秒/0.25)，缺省 1 s = 4 cell；不足补最后一个；超出截断；空脚本 ⇒ idle
    assert script_to_timeline([["right", 2.5], "left"]) == [7] * 10 + [3] * 11
    assert script_to_timeline([["right", 0.1]]) == [7] * CELLS            # max(1, round(0.4)) = 1，再补齐
    assert script_to_timeline([["up", 10.0]]) == [1] * CELLS
    assert script_to_timeline([]) == [0] * CELLS
    assert script_to_timeline([["idle", 1.0], "ur", [6, 0.5]]) == [0] * 4 + [8] * 4 + [6] * 2 + [6] * 11

    script = [["right", 2.5], "left", ["idle", 1.75]]
    ps = script_to_cell_prompts(SCENE, script)
    assert len(ps) == CELLS
    # right 2.5 s = 10 cell (0–9)，left 缺省 1 s = 4 cell (10–13)，idle 1.75 s = 7 cell (14–20)
    assert ps[0].endswith("the player is moving right.") and ps[9].endswith("moving right.")
    assert ps[10].endswith("moving left.") and ps[13].endswith("moving left.") and ps[14].endswith("standing idle.")
    assert ps[20].endswith("standing idle.")
    # 串符合数据集模板（split_prompt 能拆），且是"真实方向"口径：镜像一次就是数据集原标签
    for p in ps:
        scene, action = split_prompt(p)
        assert scene == "in a pale tidal flat strewn with wreckage"
        assert split_prompt(mirror_action(p))[0] == scene
    assert script_to_action_spec(script) == "right:1.0:10,left:1.0:4,idle:0.0:7"

    # 转场：cell [8, 12) 用 burn 从句，之后纯 B
    pt = script_to_cell_prompts(SCENE, ["right"], trigger="burn:cave:8:4")
    assert "burning away into" in pt[8] and "burning away into" in pt[11] and "burning away into" not in pt[12]
    assert split_prompt(pt[12])[0] != split_prompt(pt[0])[0]

    # 与 infer.py 的 compose_cell_prompts / velocity_to_label(mirror=False) 一致（infer 依赖缺时跳过这一段）
    try:
        from actionrope import infer as I
        import numpy as np
        from actionrope import geometry as G
        labels = script_to_timeline(script)
        vel = np.stack([G.ACTION_DIRS[l] * np.array([69.24, 41.08]) for l in labels])
        assert [I.velocity_to_label(v, mirror=False) for v in vel] == labels
        assert I.compose_cell_prompts(SCENE, None, labels) == ps
        assert I.compose_cell_prompts(SCENE, "burn:cave:8:4", [7] * CELLS) == pt
        infer_consistent = True
    except ImportError as e:  # pragma: no cover
        infer_consistent = f"skipped: {e}"

    # 生成的串都要能在 mirror 文本表里查到（推理不挂 T5）
    if os.path.exists(MIRROR_TABLE):
        keys = set(torch.load(MIRROR_TABLE, map_location="cpu", weights_only=False, mmap=True).keys())
        missing = [p for p in ps + pt if p not in keys]
        assert not missing, missing[:3]
        in_table = True
    else:
        in_table = "table missing"
    record(template_in_mirror_table=in_table, template_consistent_with_infer=infer_consistent)