predictor_training/cache.py

"""Small cache adapters for offline Predictor-v4 training."""

from __future__ import annotations

from typing import Mapping

import torch


def build_cross_attention_cache(
    text_kv: Mapping[int, Mapping[str, torch.Tensor]],
    source_block_ids: tuple[int, ...],
) -> dict[int, dict[str, torch.Tensor | bool]]:
    """Wrap saved Teacher text K/V in Wan's initialized-cache structure."""
    return {
        int(block_id): {
            "k": text_kv[int(block_id)]["k"],
            "v": text_kv[int(block_id)]["v"],
            "is_init": True,
        }
        for block_id in source_block_ids
    }