File size: 609 Bytes
d5e0d8f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | """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
}
|