Cccccz/HY

1"""Read-only capture of exact v2 trajectories and selected AR context KV."""2 3from __future__ import annotations4 5from dataclasses import dataclass6from typing import Any, Callable, Mapping7 8import torch9 10from .schema import LATENT_HEIGHT, LATENT_WIDTH, NUM_STEPS11from .v2_schema import CONTEXT_BLOCK_IDS, block_key12 13 14def _clone_cpu(tensor: torch.Tensor) -> torch.Tensor:15    return tensor.detach().to(device="cpu").contiguous()16 17 18@dataclass19class _ActiveStep:20    chunk_id: int21    step_id: int22    tensors: dict[str, torch.Tensor]23    shared: dict[str, torch.Tensor]24 25 26class PredictorV2TeacherCapture:27    """Capture Full-DiT targets plus step-invariant selected-layer AR memory."""28 29    def __init__(30        self,31        transformer: torch.nn.Module,32        *,33        on_chunk: Callable[34            [int, dict[str, torch.Tensor], dict[int, dict[str, torch.Tensor]]], None35        ],36        num_steps: int = NUM_STEPS,37        context_block_ids: tuple[int, ...] = CONTEXT_BLOCK_IDS,38    ) -> None:39        self.transformer = transformer40        self.on_chunk = on_chunk41        self.num_steps = num_steps42        self.context_block_ids = tuple(context_block_ids)43        self.call_index = 044        self.active: _ActiveStep | None = None45        self.chunk_steps: list[_ActiveStep] = []46        self.image_condition_latent: torch.Tensor | None = None47        self.text_context: dict[int, dict[str, torch.Tensor]] | None = None48        self.chunk_context: dict[int, dict[str, torch.Tensor]] | None = None49        self._handles: list[Any] = []50 51    def __enter__(self) -> "PredictorV2TeacherCapture":52        self._handles.append(53            self.transformer.register_forward_pre_hook(self._transformer_pre, with_kwargs=True)54        )55        self._handles.append(56            self.transformer.register_forward_hook(self._transformer_post, with_kwargs=True)57        )58        self._handles.append(59            self.transformer.final_layer.register_forward_pre_hook(self._final_pre, with_kwargs=True)60        )61        return self62 63    def __exit__(self, exc_type, exc, traceback) -> bool:64        for handle in self._handles:65            handle.remove()66        self._handles.clear()67        self.active = None68        return False69 70    def _capture_context(71        self,72        kv_cache: list[Mapping[str, torch.Tensor | None]],73    ) -> dict[int, dict[str, torch.Tensor]]:74        result: dict[int, dict[str, torch.Tensor]] = {}75        for block_id in self.context_block_ids:76            cache = kv_cache[block_id]77            k_txt = cache.get("k_txt")78            v_txt = cache.get("v_txt")79            if k_txt is None or v_txt is None:80                raise RuntimeError(f"Block {block_id} text KV is unavailable")81            k_vision = cache.get("k_vision")82            v_vision = cache.get("v_vision")83            if (k_vision is None) != (v_vision is None):84                raise RuntimeError(f"Block {block_id} has incomplete vision KV")85            if k_vision is None:86                empty_shape = (2, int(k_txt.shape[1]), 0, int(k_txt.shape[3]))87                k_vision_cpu = torch.empty(empty_shape, dtype=k_txt.dtype, device="cpu")88                v_vision_cpu = torch.empty(empty_shape, dtype=v_txt.dtype, device="cpu")89            else:90                k_vision_cpu = _clone_cpu(k_vision)91                v_vision_cpu = _clone_cpu(v_vision)92            result[block_id] = {93                "k_vision": k_vision_cpu,94                "v_vision": v_vision_cpu,95            }96            if self.text_context is None:97                result[block_id]["k_txt"] = _clone_cpu(k_txt)98                result[block_id]["v_txt"] = _clone_cpu(v_txt)99        if self.text_context is None:100            self.text_context = {101                block_id: {102                    "k_txt": result[block_id].pop("k_txt"),103                    "v_txt": result[block_id].pop("v_txt"),104                }105                for block_id in self.context_block_ids106            }107        return result108 109    def _transformer_pre(self, module, args, kwargs) -> None:110        is_denoise = (111            kwargs.get("ar_vision_inference", False)112            and not kwargs.get("cache_vision", False)113        )114        if not is_denoise:115            return116        if self.active is not None:117            raise RuntimeError("Nested denoising capture is not supported")118 119        chunk_id, step_id = divmod(self.call_index, self.num_steps)120        model_input = kwargs["hidden_states"]121        if model_input.shape[1] != 65:122            raise ValueError(f"Teacher denoising input must have 65 channels, got {model_input.shape}")123        if self.image_condition_latent is None:124            self.image_condition_latent = _clone_cpu(model_input[:, 32:64, 0:1])125            mask = model_input[:, 64:65]126            if not torch.all(mask[:, :, 0] == 1) or not torch.all(mask[:, :, 1:] == 0):127                raise ValueError("Unexpected I2V condition mask in first chunk")128 129        if step_id == 0:130            if self.chunk_context is not None:131                raise RuntimeError("Previous chunk context was not flushed")132            self.chunk_context = self._capture_context(kwargs["kv_cache"])133 134        timestep = kwargs["timestep"].reshape(-1)[0:1]135        shared = {136            "action_labels": _clone_cpu(kwargs["action"].reshape(1, -1).round().long()),137            "target_viewmats": _clone_cpu(kwargs["viewmats"]),138            "target_Ks": _clone_cpu(kwargs["Ks"]),139            "rope_temporal_size": torch.tensor([int(kwargs["rope_temporal_size"])], dtype=torch.int64),140            "start_rope_start_idx": torch.tensor(141                [int(kwargs["start_rope_start_idx"])], dtype=torch.int64142            ),143        }144        self.active = _ActiveStep(145            chunk_id=chunk_id,146            step_id=step_id,147            tensors={148                "timestep": _clone_cpu(timestep.float()),149                "noisy_sample": _clone_cpu(model_input[:, :32]),150            },151            shared=shared,152        )153 154    def _final_pre(self, module, args, kwargs) -> None:155        if self.active is None:156            return157        hidden, condition = args[0], args[1]158        batch, tokens, hidden_size = hidden.shape159        spatial_tokens = LATENT_HEIGHT * LATENT_WIDTH160        if tokens % spatial_tokens:161            raise ValueError(f"Final hidden token count {tokens} is not frame-aligned")162        frames = tokens // spatial_tokens163        compact = condition.reshape(batch, frames, spatial_tokens, hidden_size)[:, :, 0]164        expanded = compact[:, :, None].expand(batch, frames, spatial_tokens, hidden_size)165        if not torch.equal(166            expanded.reshape(batch, tokens, hidden_size),167            condition.reshape(batch, tokens, hidden_size),168        ):169            raise ValueError("Final-layer condition varies inside a latent frame")170        self.active.tensors["frame_condition"] = _clone_cpu(compact)171        self.active.tensors["final_hidden"] = _clone_cpu(hidden)172 173    def _transformer_post(self, module, args, kwargs, output) -> None:174        if self.active is None:175            return176        velocity = output[0] if isinstance(output, tuple) else output177        self.active.tensors["velocity"] = _clone_cpu(velocity)178        required = {"timestep", "noisy_sample", "frame_condition", "final_hidden", "velocity"}179        missing = required.difference(self.active.tensors)180        if missing:181            raise RuntimeError(f"Incomplete v2 teacher step capture: {sorted(missing)}")182        self.chunk_steps.append(self.active)183        completed_step = self.active.step_id184        self.active = None185        self.call_index += 1186        if completed_step == self.num_steps - 1:187            self._flush_chunk()188 189    def _flush_chunk(self) -> None:190        if len(self.chunk_steps) != self.num_steps:191            raise RuntimeError(f"Expected {self.num_steps} steps, got {len(self.chunk_steps)}")192        if self.chunk_context is None:193            raise RuntimeError("Chunk context was not captured")194        chunk_id = self.chunk_steps[0].chunk_id195        if any(step.chunk_id != chunk_id for step in self.chunk_steps):196            raise RuntimeError("Captured steps cross a chunk boundary")197        tensors = dict(self.chunk_steps[0].shared)198        for step in self.chunk_steps:199            for name, tensor in step.tensors.items():200                tensors[f"step_{step.step_id}_{name}"] = tensor201        self.on_chunk(chunk_id, tensors, self.chunk_context)202        self.chunk_steps.clear()203        self.chunk_context = None204 205    def case_tensors(self) -> dict[str, torch.Tensor]:206        if self.image_condition_latent is None or self.text_context is None:207            raise RuntimeError("Case condition or text context has not been captured")208        result = {"image_condition_latent": self.image_condition_latent}209        for block_id, cache in self.text_context.items():210            result[block_key(block_id, "k_txt")] = cache["k_txt"]211            result[block_key(block_id, "v_txt")] = cache["v_txt"]212        return result213 214    @property215    def text_token_count(self) -> int:216        if self.text_context is None:217            raise RuntimeError("Text context has not been captured")218        first = self.text_context[self.context_block_ids[0]]["k_txt"]219        return int(first.shape[2])220