"""Read-only hooks that capture exact Full-DiT teacher trajectories.""" from __future__ import annotations from dataclasses import dataclass from typing import Any, Callable import torch from .schema import LATENT_HEIGHT, LATENT_WIDTH, NUM_STEPS def _clone_cpu(tensor: torch.Tensor) -> torch.Tensor: return tensor.detach().to(device="cpu").contiguous() @dataclass class _ActiveStep: chunk_id: int step_id: int tensors: dict[str, torch.Tensor] shared: dict[str, torch.Tensor] class PredictorTeacherCapture: """Capture denoising inputs, final hidden/condition, velocity, and dense txt features. The hook assumes a single positive AR stream (few-step guidance=1) and a fixed number of denoising steps per chunk. History-prefill calls are excluded through ``cache_vision``. """ def __init__( self, transformer: torch.nn.Module, *, on_chunk: Callable[[int, dict[str, torch.Tensor]], None], num_steps: int = NUM_STEPS, ) -> None: self.transformer = transformer self.on_chunk = on_chunk self.num_steps = num_steps self.call_index = 0 self.active: _ActiveStep | None = None self.chunk_steps: list[_ActiveStep] = [] self.current_txt: torch.Tensor | None = None self.cached_txt: torch.Tensor | None = None self.vec_txt: torch.Tensor | None = None self.image_condition_latent: torch.Tensor | None = None self._handles: list[Any] = [] self._original_get_text_and_mask = None def __enter__(self) -> "PredictorTeacherCapture": self._original_get_text_and_mask = self.transformer.get_text_and_mask def wrapped_get_text_and_mask(*args, **kwargs): txt, text_mask, vec_txt = self._original_get_text_and_mask(*args, **kwargs) if self.current_txt is None: if txt.shape[0] != 1: raise ValueError("Predictor capture currently requires text batch size 1") valid = text_mask[0].bool().to(txt.device) self.current_txt = _clone_cpu(txt[:, valid]) self.vec_txt = _clone_cpu(vec_txt) return txt, text_mask, vec_txt self.transformer.get_text_and_mask = wrapped_get_text_and_mask self._handles.append( self.transformer.register_forward_pre_hook(self._transformer_pre, with_kwargs=True) ) self._handles.append( self.transformer.register_forward_hook(self._transformer_post, with_kwargs=True) ) self._handles.append( self.transformer.final_layer.register_forward_pre_hook(self._final_pre, with_kwargs=True) ) self._handles.append( self.transformer.double_blocks[-1].register_forward_hook( self._last_block_post, with_kwargs=True ) ) return self def __exit__(self, exc_type, exc, traceback) -> bool: for handle in self._handles: handle.remove() self._handles.clear() if self._original_get_text_and_mask is not None: self.transformer.get_text_and_mask = self._original_get_text_and_mask self.active = None return False def _last_block_post(self, module, args, kwargs, output) -> None: if kwargs.get("ar_txt_inference", False): txt = output[0] if isinstance(output, tuple) else output self.cached_txt = _clone_cpu(txt) def _transformer_pre(self, module, args, kwargs) -> None: is_denoise = ( kwargs.get("ar_vision_inference", False) and not kwargs.get("cache_vision", False) ) if not is_denoise: return if self.active is not None: raise RuntimeError("Nested denoising capture is not supported") chunk_id, step_id = divmod(self.call_index, self.num_steps) model_input = kwargs["hidden_states"] if model_input.shape[1] != 65: raise ValueError(f"Teacher denoising input must have 65 channels, got {model_input.shape}") if self.image_condition_latent is None: self.image_condition_latent = _clone_cpu(model_input[:, 32:64, 0:1]) mask = model_input[:, 64:65] if not torch.all(mask[:, :, 0] == 1) or not torch.all(mask[:, :, 1:] == 0): raise ValueError("Unexpected I2V condition mask in first chunk") timestep = kwargs["timestep"].reshape(-1)[0:1] shared = { "action_labels": _clone_cpu(kwargs["action"].reshape(1, -1).round().long()), "target_viewmats": _clone_cpu(kwargs["viewmats"]), "target_Ks": _clone_cpu(kwargs["Ks"]), "rope_temporal_size": torch.tensor([int(kwargs["rope_temporal_size"])], dtype=torch.int64), "start_rope_start_idx": torch.tensor( [int(kwargs["start_rope_start_idx"])], dtype=torch.int64 ), } self.active = _ActiveStep( chunk_id=chunk_id, step_id=step_id, tensors={ "timestep": _clone_cpu(timestep.float()), "noisy_sample": _clone_cpu(model_input[:, :32]), }, shared=shared, ) def _final_pre(self, module, args, kwargs) -> None: if self.active is None: return hidden, condition = args[0], args[1] batch, tokens, hidden_size = hidden.shape spatial_tokens = LATENT_HEIGHT * LATENT_WIDTH if tokens % spatial_tokens: raise ValueError(f"Final hidden token count {tokens} is not divisible by {spatial_tokens}") frames = tokens // spatial_tokens compact = condition.reshape(batch, frames, spatial_tokens, hidden_size)[:, :, 0] expanded = compact[:, :, None].expand(batch, frames, spatial_tokens, hidden_size) if not torch.equal(expanded.reshape(batch, tokens, hidden_size), condition.reshape(batch, tokens, hidden_size)): raise ValueError("Final-layer condition varies inside a latent frame") self.active.tensors["frame_condition"] = _clone_cpu(compact) self.active.tensors["final_hidden"] = _clone_cpu(hidden) def _transformer_post(self, module, args, kwargs, output) -> None: if self.active is None: return velocity = output[0] if isinstance(output, tuple) else output self.active.tensors["velocity"] = _clone_cpu(velocity) required = {"timestep", "noisy_sample", "frame_condition", "final_hidden", "velocity"} missing = required.difference(self.active.tensors) if missing: raise RuntimeError(f"Incomplete teacher step capture: {sorted(missing)}") self.chunk_steps.append(self.active) completed_step = self.active.step_id self.active = None self.call_index += 1 if completed_step == self.num_steps - 1: self._flush_chunk() def _flush_chunk(self) -> None: if len(self.chunk_steps) != self.num_steps: raise RuntimeError(f"Expected {self.num_steps} captured steps, got {len(self.chunk_steps)}") chunk_id = self.chunk_steps[0].chunk_id if any(step.chunk_id != chunk_id for step in self.chunk_steps): raise RuntimeError("Captured steps cross chunk boundary") tensors = dict(self.chunk_steps[0].shared) for step in self.chunk_steps: for name, tensor in step.tensors.items(): tensors[f"step_{step.step_id}_{name}"] = tensor self.on_chunk(chunk_id, tensors) self.chunk_steps.clear() def case_tensors(self) -> dict[str, torch.Tensor]: missing = [ name for name, value in ( ("image_condition_latent", self.image_condition_latent), ("current_txt", self.current_txt), ("cached_txt", self.cached_txt), ("vec_txt", self.vec_txt), ) if value is None ] if missing: raise RuntimeError(f"Missing case captures: {missing}") return { "image_condition_latent": self.image_condition_latent, "current_txt": self.current_txt, "cached_txt": self.cached_txt, "vec_txt": self.vec_txt, }