Cccccz/HY

1"""Read-only Full-DiT capture of targets and joint-window Context prefeatures."""2 3from __future__ import annotations4 5from typing import Any, Callable6 7import torch8 9from .prefeature_schema import CONTEXT_BLOCK_IDS10from .v2_hooks import PredictorV2TeacherCapture, _clone_cpu11 12 13class PredictorPrefeatureTeacherCapture(PredictorV2TeacherCapture):14    """Capture ``img_modulated`` before selected Teacher K/V projections."""15 16    def __init__(17        self,18        transformer: torch.nn.Module,19        *,20        on_chunk: Callable[[int, dict[str, torch.Tensor], dict[int, dict[str, torch.Tensor]]], None],21        num_steps: int = 4,22        context_block_ids: tuple[int, ...] = CONTEXT_BLOCK_IDS,23        capture_direct_kv: bool = False,24    ) -> None:25        super().__init__(26            transformer,27            on_chunk=on_chunk,28            num_steps=num_steps,29            context_block_ids=context_block_ids,30        )31        self.capture_direct_kv = bool(capture_direct_kv)32        self.pending_context: dict[int, dict[str, torch.Tensor]] | None = None33        self.pending_direct_kv: dict[int, dict[str, torch.Tensor]] | None = None34        self._prefill_active = False35        self._prefill_features: dict[int, torch.Tensor] = {}36        self._prefill_metadata: dict[str, torch.Tensor] = {}37 38    def __enter__(self) -> "PredictorPrefeatureTeacherCapture":39        super().__enter__()40        for block_id in self.context_block_ids:41            block = self.transformer.double_blocks[block_id]42            self._handles.append(43                block.img_attn_k.register_forward_pre_hook(44                    self._make_k_pre_hook(block_id), with_kwargs=True45                )46            )47        return self48 49    def _make_k_pre_hook(self, block_id: int):50        def hook(module, args, kwargs) -> None:51            if not self._prefill_active:52                return53            if block_id in self._prefill_features:54                raise RuntimeError(f"Context block {block_id} was captured twice")55            if not args or not torch.is_tensor(args[0]):56                raise RuntimeError(f"Missing img_modulated input for block {block_id}")57            self._prefill_features[block_id] = _clone_cpu(args[0])58 59        return hook60 61    def _transformer_pre(self, module, args, kwargs) -> None:62        if kwargs.get("ar_vision_inference", False) and kwargs.get("cache_vision", False):63            if self.pending_context is not None or self._prefill_active:64                raise RuntimeError("A Context prefill was not consumed before the next prefill")65            frame_indices = kwargs.get("context_frame_indices")66            if frame_indices is None:67                raise RuntimeError("Context prefill lacks context_frame_indices metadata")68            if not torch.is_tensor(frame_indices):69                frame_indices = torch.tensor(frame_indices, dtype=torch.int64)70            frame_indices = frame_indices.detach().to(device="cpu", dtype=torch.int64).reshape(-1)71            frames = int(kwargs["hidden_states"].shape[2])72            if frame_indices.numel() != frames:73                raise ValueError("Context frame metadata does not match prefill tensor")74            self._prefill_metadata = {75                "selected_frame_indices": frame_indices.contiguous(),76                "context_viewmats": _clone_cpu(kwargs["viewmats"]),77                "context_Ks": _clone_cpu(kwargs["Ks"]),78                "rope_temporal_size": torch.tensor(79                    [int(kwargs["rope_temporal_size"])], dtype=torch.int6480                ),81                "start_rope_start_idx": torch.tensor(82                    [int(kwargs["start_rope_start_idx"])], dtype=torch.int6483                ),84            }85            self._prefill_features = {}86            self._prefill_active = True87            return88        super()._transformer_pre(module, args, kwargs)89        if self.active is not None and self.active.step_id == 0:90            self.chunk_context = self._consume_context()91 92    def _transformer_post(self, module, args, kwargs, output) -> None:93        if self._prefill_active:94            self._prefill_active = False95            missing = set(self.context_block_ids).difference(self._prefill_features)96            if missing:97                raise RuntimeError(f"Missing prefeatures for blocks {sorted(missing)}")98            token_counts = {int(value.shape[1]) for value in self._prefill_features.values()}99            if len(token_counts) != 1:100                raise RuntimeError("Selected Context blocks have different window lengths")101            tokens = next(iter(token_counts))102            mask = torch.ones((1, tokens), dtype=torch.bool)103            self.pending_context = {104                block_id: {105                    "img_modulated": feature,106                    "context_valid_mask": mask.clone(),107                    **{name: value.clone() for name, value in self._prefill_metadata.items()},108                }109                for block_id, feature in self._prefill_features.items()110            }111            if self.capture_direct_kv:112                self.pending_direct_kv = {113                    block_id: {114                        "k_vision": _clone_cpu(output[block_id]["k_vision"]),115                        "v_vision": _clone_cpu(output[block_id]["v_vision"]),116                    }117                    for block_id in self.context_block_ids118                }119            self._prefill_features = {}120            self._prefill_metadata = {}121            return122        super()._transformer_post(module, args, kwargs, output)123 124    def _consume_context(self) -> dict[int, dict[str, torch.Tensor]]:125        if self.pending_context is not None:126            context = self.pending_context127            self.pending_context = None128            return context129        # The first chunk has no history prefill.130        empty_metadata = {131            "context_valid_mask": torch.ones((1, 0), dtype=torch.bool),132            "selected_frame_indices": torch.empty((0,), dtype=torch.int64),133            "context_viewmats": torch.empty((1, 0, 4, 4), dtype=torch.bfloat16),134            "context_Ks": torch.empty((1, 0, 3, 3), dtype=torch.bfloat16),135            "rope_temporal_size": torch.tensor([0], dtype=torch.int64),136            "start_rope_start_idx": torch.tensor([0], dtype=torch.int64),137        }138        return {139            block_id: {140                "img_modulated": torch.empty((1, 0, 2048), dtype=torch.bfloat16),141                **{name: value.clone() for name, value in empty_metadata.items()},142            }143            for block_id in self.context_block_ids144        }145 146    def _capture_context(self, kv_cache):147        """Capture only case-level text KV; Context features come from block hooks."""148        if self.text_context is None:149            text_context: dict[int, dict[str, torch.Tensor]] = {}150            for block_id in self.context_block_ids:151                cache = kv_cache[block_id]152                k_txt, v_txt = cache.get("k_txt"), cache.get("v_txt")153                if k_txt is None or v_txt is None:154                    raise RuntimeError(f"Block {block_id} text KV is unavailable")155                text_context[block_id] = {156                    "k_txt": _clone_cpu(k_txt),157                    "v_txt": _clone_cpu(v_txt),158                }159            self.text_context = text_context160        return {block_id: {} for block_id in self.context_block_ids}161 162    def take_direct_kv(self) -> dict[int, dict[str, torch.Tensor]] | None:163        value = self.pending_direct_kv164        self.pending_direct_kv = None165        return value166