File size: 4,809 Bytes
510ab6b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | """Opt-in forward hooks used by the offline Full-DiT teacher builder."""
from __future__ import annotations
from typing import Any
import torch
from .schema import CANDIDATE_BLOCK_IDS
def clone_bf16_cpu(tensor: torch.Tensor) -> torch.Tensor:
return tensor.detach().to(device="cpu", dtype=torch.bfloat16).contiguous()
class PredictorV4TeacherCapture:
"""Capture head input and clean pre-KV features without modifying Wan.
The builder explicitly marks denoising and clean calls. Hooks therefore do
not infer call order and cannot accidentally record the final clean pass as
a fifth denoising step.
"""
def __init__(
self,
model: torch.nn.Module,
block_ids: tuple[int, ...] = CANDIDATE_BLOCK_IDS,
) -> None:
self.model = model
self.block_ids = tuple(int(value) for value in block_ids)
self._handles: list[Any] = []
self._mode: str | None = None
self._active_step: int | None = None
self._final_hidden: torch.Tensor | None = None
self._clean_features: dict[int, torch.Tensor] = {}
def __enter__(self) -> "PredictorV4TeacherCapture":
if self._handles:
raise RuntimeError("capture hooks are already installed")
self._handles.append(
self.model.head.register_forward_pre_hook(self._head_pre_hook, with_kwargs=True)
)
for block_id in self.block_ids:
try:
projection = self.model.blocks[block_id].self_attn.k
except IndexError as exc:
raise ValueError(f"Wan model does not contain block {block_id}") from exc
self._handles.append(
projection.register_forward_pre_hook(
self._make_clean_k_pre_hook(block_id),
with_kwargs=True,
)
)
return self
def __exit__(self, exc_type, exc, traceback) -> bool:
for handle in self._handles:
handle.remove()
self._handles.clear()
self._mode = None
self._active_step = None
self._final_hidden = None
self._clean_features.clear()
return False
def _head_pre_hook(self, module, args, kwargs) -> None:
if self._mode != "denoise":
return
if self._final_hidden is not None:
raise RuntimeError(f"head input captured twice for step {self._active_step}")
if not args or not torch.is_tensor(args[0]):
raise RuntimeError("Wan head did not receive a tensor hidden state")
self._final_hidden = clone_bf16_cpu(args[0])
def _make_clean_k_pre_hook(self, block_id: int):
def hook(module, args, kwargs) -> None:
if self._mode != "clean":
return
if block_id in self._clean_features:
raise RuntimeError(f"clean block {block_id} was captured twice")
if not args or not torch.is_tensor(args[0]):
raise RuntimeError(f"block {block_id} self_attn.k lacks a tensor input")
self._clean_features[block_id] = clone_bf16_cpu(args[0])
return hook
def begin_denoise(self, step_id: int) -> None:
if self._mode is not None:
raise RuntimeError(f"cannot begin denoise while capture mode is {self._mode}")
self._mode = "denoise"
self._active_step = int(step_id)
self._final_hidden = None
def finish_denoise(self) -> torch.Tensor:
if self._mode != "denoise":
raise RuntimeError("finish_denoise called without begin_denoise")
value = self._final_hidden
step_id = self._active_step
self._mode = None
self._active_step = None
self._final_hidden = None
if value is None:
raise RuntimeError(f"missing Wan head input for denoising step {step_id}")
return value
def begin_clean(self) -> None:
if self._mode is not None:
raise RuntimeError(f"cannot begin clean capture while mode is {self._mode}")
self._mode = "clean"
self._clean_features = {}
def finish_clean(self) -> dict[int, torch.Tensor]:
if self._mode != "clean":
raise RuntimeError("finish_clean called without begin_clean")
result = self._clean_features
self._mode = None
self._clean_features = {}
missing = set(self.block_ids).difference(result)
if missing:
raise RuntimeError(f"missing clean prefeatures for blocks {sorted(missing)}")
return result
def abort_active_call(self) -> None:
"""Reset hook state after a model exception before re-raising it."""
self._mode = None
self._active_step = None
self._final_hidden = None
self._clean_features = {}
|