Cccccz/HY

1"""Read-only hooks that capture exact Full-DiT teacher trajectories."""2 3from __future__ import annotations4 5from dataclasses import dataclass6from typing import Any, Callable7 8import torch9 10from .schema import LATENT_HEIGHT, LATENT_WIDTH, NUM_STEPS11 12 13def _clone_cpu(tensor: torch.Tensor) -> torch.Tensor:14    return tensor.detach().to(device="cpu").contiguous()15 16 17@dataclass18class _ActiveStep:19    chunk_id: int20    step_id: int21    tensors: dict[str, torch.Tensor]22    shared: dict[str, torch.Tensor]23 24 25class PredictorTeacherCapture:26    """Capture denoising inputs, final hidden/condition, velocity, and dense txt features.27 28    The hook assumes a single positive AR stream (few-step guidance=1) and a fixed29    number of denoising steps per chunk. History-prefill calls are excluded through30    ``cache_vision``.31    """32 33    def __init__(34        self,35        transformer: torch.nn.Module,36        *,37        on_chunk: Callable[[int, dict[str, torch.Tensor]], None],38        num_steps: int = NUM_STEPS,39    ) -> None:40        self.transformer = transformer41        self.on_chunk = on_chunk42        self.num_steps = num_steps43        self.call_index = 044        self.active: _ActiveStep | None = None45        self.chunk_steps: list[_ActiveStep] = []46        self.current_txt: torch.Tensor | None = None47        self.cached_txt: torch.Tensor | None = None48        self.vec_txt: torch.Tensor | None = None49        self.image_condition_latent: torch.Tensor | None = None50        self._handles: list[Any] = []51        self._original_get_text_and_mask = None52 53    def __enter__(self) -> "PredictorTeacherCapture":54        self._original_get_text_and_mask = self.transformer.get_text_and_mask55 56        def wrapped_get_text_and_mask(*args, **kwargs):57            txt, text_mask, vec_txt = self._original_get_text_and_mask(*args, **kwargs)58            if self.current_txt is None:59                if txt.shape[0] != 1:60                    raise ValueError("Predictor capture currently requires text batch size 1")61                valid = text_mask[0].bool().to(txt.device)62                self.current_txt = _clone_cpu(txt[:, valid])63                self.vec_txt = _clone_cpu(vec_txt)64            return txt, text_mask, vec_txt65 66        self.transformer.get_text_and_mask = wrapped_get_text_and_mask67        self._handles.append(68            self.transformer.register_forward_pre_hook(self._transformer_pre, with_kwargs=True)69        )70        self._handles.append(71            self.transformer.register_forward_hook(self._transformer_post, with_kwargs=True)72        )73        self._handles.append(74            self.transformer.final_layer.register_forward_pre_hook(self._final_pre, with_kwargs=True)75        )76        self._handles.append(77            self.transformer.double_blocks[-1].register_forward_hook(78                self._last_block_post, with_kwargs=True79            )80        )81        return self82 83    def __exit__(self, exc_type, exc, traceback) -> bool:84        for handle in self._handles:85            handle.remove()86        self._handles.clear()87        if self._original_get_text_and_mask is not None:88            self.transformer.get_text_and_mask = self._original_get_text_and_mask89        self.active = None90        return False91 92    def _last_block_post(self, module, args, kwargs, output) -> None:93        if kwargs.get("ar_txt_inference", False):94            txt = output[0] if isinstance(output, tuple) else output95            self.cached_txt = _clone_cpu(txt)96 97    def _transformer_pre(self, module, args, kwargs) -> None:98        is_denoise = (99            kwargs.get("ar_vision_inference", False)100            and not kwargs.get("cache_vision", False)101        )102        if not is_denoise:103            return104        if self.active is not None:105            raise RuntimeError("Nested denoising capture is not supported")106 107        chunk_id, step_id = divmod(self.call_index, self.num_steps)108        model_input = kwargs["hidden_states"]109        if model_input.shape[1] != 65:110            raise ValueError(f"Teacher denoising input must have 65 channels, got {model_input.shape}")111        if self.image_condition_latent is None:112            self.image_condition_latent = _clone_cpu(model_input[:, 32:64, 0:1])113            mask = model_input[:, 64:65]114            if not torch.all(mask[:, :, 0] == 1) or not torch.all(mask[:, :, 1:] == 0):115                raise ValueError("Unexpected I2V condition mask in first chunk")116 117        timestep = kwargs["timestep"].reshape(-1)[0:1]118        shared = {119            "action_labels": _clone_cpu(kwargs["action"].reshape(1, -1).round().long()),120            "target_viewmats": _clone_cpu(kwargs["viewmats"]),121            "target_Ks": _clone_cpu(kwargs["Ks"]),122            "rope_temporal_size": torch.tensor([int(kwargs["rope_temporal_size"])], dtype=torch.int64),123            "start_rope_start_idx": torch.tensor(124                [int(kwargs["start_rope_start_idx"])], dtype=torch.int64125            ),126        }127        self.active = _ActiveStep(128            chunk_id=chunk_id,129            step_id=step_id,130            tensors={131                "timestep": _clone_cpu(timestep.float()),132                "noisy_sample": _clone_cpu(model_input[:, :32]),133            },134            shared=shared,135        )136 137    def _final_pre(self, module, args, kwargs) -> None:138        if self.active is None:139            return140        hidden, condition = args[0], args[1]141        batch, tokens, hidden_size = hidden.shape142        spatial_tokens = LATENT_HEIGHT * LATENT_WIDTH143        if tokens % spatial_tokens:144            raise ValueError(f"Final hidden token count {tokens} is not divisible by {spatial_tokens}")145        frames = tokens // spatial_tokens146        compact = condition.reshape(batch, frames, spatial_tokens, hidden_size)[:, :, 0]147        expanded = compact[:, :, None].expand(batch, frames, spatial_tokens, hidden_size)148        if not torch.equal(expanded.reshape(batch, tokens, hidden_size), condition.reshape(batch, tokens, hidden_size)):149            raise ValueError("Final-layer condition varies inside a latent frame")150        self.active.tensors["frame_condition"] = _clone_cpu(compact)151        self.active.tensors["final_hidden"] = _clone_cpu(hidden)152 153    def _transformer_post(self, module, args, kwargs, output) -> None:154        if self.active is None:155            return156        velocity = output[0] if isinstance(output, tuple) else output157        self.active.tensors["velocity"] = _clone_cpu(velocity)158        required = {"timestep", "noisy_sample", "frame_condition", "final_hidden", "velocity"}159        missing = required.difference(self.active.tensors)160        if missing:161            raise RuntimeError(f"Incomplete teacher step capture: {sorted(missing)}")162        self.chunk_steps.append(self.active)163        completed_step = self.active.step_id164        self.active = None165        self.call_index += 1166        if completed_step == self.num_steps - 1:167            self._flush_chunk()168 169    def _flush_chunk(self) -> None:170        if len(self.chunk_steps) != self.num_steps:171            raise RuntimeError(f"Expected {self.num_steps} captured steps, got {len(self.chunk_steps)}")172        chunk_id = self.chunk_steps[0].chunk_id173        if any(step.chunk_id != chunk_id for step in self.chunk_steps):174            raise RuntimeError("Captured steps cross chunk boundary")175        tensors = dict(self.chunk_steps[0].shared)176        for step in self.chunk_steps:177            for name, tensor in step.tensors.items():178                tensors[f"step_{step.step_id}_{name}"] = tensor179        self.on_chunk(chunk_id, tensors)180        self.chunk_steps.clear()181 182    def case_tensors(self) -> dict[str, torch.Tensor]:183        missing = [184            name185            for name, value in (186                ("image_condition_latent", self.image_condition_latent),187                ("current_txt", self.current_txt),188                ("cached_txt", self.cached_txt),189                ("vec_txt", self.vec_txt),190            )191            if value is None192        ]193        if missing:194            raise RuntimeError(f"Missing case captures: {missing}")195        return {196            "image_condition_latent": self.image_condition_latent,197            "current_txt": self.current_txt,198            "cached_txt": self.cached_txt,199            "vec_txt": self.vec_txt,200        }201