| """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, |
| } |
|
|
