predictor_data/capture.py

"""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 = {}