viper/physics_encoder.py

"""VIPER visual physics encoder.

Sec. 3.2 of the paper: a frozen MLLM (Qwen3-VL-4B-Instruct) reads the reference
video together with a set of *learnable query tokens* `q_learnable`.  The query
tokens are concatenated after the system-prompt tokens and the reference visual
tokens; after the MLLM forward pass their final hidden states are taken as the
reference physics representation ``v_ref``.  A 3-layer connector ``g_theta``
projects ``v_ref`` into physics condition tokens ``c_p`` that live in the DiT
context space (5120-d for Wan2.2-A14B).

Only ``q_learnable`` and ``g_theta`` are trainable here -- the MLLM stays frozen
in all three stages.
"""

from __future__ import annotations

import torch
import torch.nn as nn

# The physics-probing system prompt.  The paper describes the MLLM as extracting
# "material response, contact interaction, deformation and motion trajectory";
# we phrase the prompt around exactly those axes.
PHYSICS_SYSTEM_PROMPT = (
    "You are a physics analyst. Watch the video and reason about its physical "
    "process: the material properties of the objects, how they deform under "
    "force, the contact and collision interactions between them, and the motion "
    "trajectory over time. Summarize the underlying physical dynamics, ignoring "
    "the specific appearance, identity and background of the objects."
)


class Connector(nn.Module):
    """3-layer MLP connector g_theta: MLLM hidden dim -> DiT context dim."""

    def __init__(self, in_dim: int, out_dim: int, hidden_dim: int | None = None):
        super().__init__()
        hidden_dim = hidden_dim or out_dim
        # NOTE on ordering: the LayerNorm sits *before* the final projection, not
        # after.  Normalising after a zero-initialised output layer would be a
        # degenerate fixed point -- LayerNorm(0) is 0 with a 1/sqrt(eps) Jacobian,
        # and it also destroys the scale that makes zero-init a no-op in the first
        # place.  With the norm first, the output layer starts as an exact no-op
        # yet still receives healthy gradients (dL/dW = dL/dy . x^T with x != 0).
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(hidden_dim, hidden_dim),
            nn.GELU(approximate="tanh"),
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, out_dim),
        )
        # Zero-init the output projection so that at step 0 the physics tokens are
        # exactly zero and the frozen DiT reproduces the base model's behaviour --
        # stage 1 then starts from "does not disturb the visual prior".
        nn.init.zeros_(self.net[-1].weight)
        nn.init.zeros_(self.net[-1].bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class VisualPhysicsEncoder(nn.Module):
    """Frozen Qwen3-VL + learnable query tokens + connector.

    Parameters
    ----------
    mllm:
        A ``Qwen3VLForConditionalGeneration`` (or the inner model).  Kept frozen.
    num_queries:
        Number of learnable query tokens ``|q_learnable|``.
    out_dim:
        DiT context dimension (5120 for Wan2.2-A14B).
    """

    def __init__(self, mllm, num_queries: int = 64, out_dim: int = 5120):
        super().__init__()
        self.mllm = mllm
        self.num_queries = num_queries

        cfg = mllm.config
        text_cfg = getattr(cfg, "text_config", cfg)
        self.hidden_size = text_cfg.hidden_size

        # q_learnable -- initialised small so the frozen MLLM sees an in-distribution
        # perturbation rather than noise.
        self.query_tokens = nn.Parameter(
            torch.randn(num_queries, self.hidden_size) * 0.02
        )
        self.connector = Connector(self.hidden_size, out_dim)

        # Freeze the MLLM.
        self.mllm.requires_grad_(False)
        self.mllm.eval()

        self._hook_handle = None
        self._install_query_hook()

    # ------------------------------------------------------------------
    # Query injection
    # ------------------------------------------------------------------
    def _embed_tokens_module(self):
        """Locate the text embedding module across transformers layouts."""
        m = self.mllm
        for path in (
            "model.language_model.embed_tokens",   # Qwen3VLModel -> Qwen3VLTextModel
            "model.model.embed_tokens",
            "language_model.model.embed_tokens",
            "model.embed_tokens",
        ):
            obj = m
            try:
                for part in path.split("."):
                    obj = getattr(obj, part)
                if isinstance(obj, nn.Embedding):
                    return obj
            except AttributeError:
                continue
        raise RuntimeError("Could not locate embed_tokens on the MLLM")

    def _install_query_hook(self):
        """Overwrite the last ``num_queries`` token embeddings with q_learnable.

        Injecting at the embedding output (rather than passing ``inputs_embeds``)
        keeps all of Qwen3-VL's native vision-token merging and 3D-mRoPE index
        computation intact, while still letting gradients reach ``query_tokens``.
        The queries are appended at the very end of the prompt, so they occupy
        the final ``num_queries`` positions and are never vision positions.
        """
        emb = self._embed_tokens_module()

        def hook(module, args, output):
            if not self._inject:
                return output
            n = self.num_queries
            q = self.query_tokens.to(output.dtype)
            out = output.clone()
            out[:, -n:, :] = q.unsqueeze(0).expand(output.shape[0], -1, -1)
            return out

        self._inject = False
        self._hook_handle = emb.register_forward_hook(hook)

    # ------------------------------------------------------------------
    def forward(self, mllm_inputs: dict) -> torch.Tensor:
        """Encode a reference video into physics condition tokens.

        ``mllm_inputs`` is the processor output for
        (system prompt + reference video), already extended with ``num_queries``
        placeholder tokens at the end (see ``build_mllm_inputs``).

        Returns ``c_p`` of shape ``[B, num_queries, out_dim]``.
        """
        self._inject = True
        try:
            out = self.mllm.model(
                **mllm_inputs,
                output_hidden_states=True,
                return_dict=True,
                use_cache=False,
            )
        finally:
            self._inject = False

        hidden = out.last_hidden_state          # [B, L, H]
        v_ref = hidden[:, -self.num_queries:, :]  # [B, Nq, H]
        return self.connector(v_ref.to(self.connector.net[0].weight.dtype))

    # ------------------------------------------------------------------
    def trainable_parameters(self):
        yield self.query_tokens
        yield from self.connector.parameters()


def build_mllm_inputs(processor, video_frames, num_queries: int, device,
                      system_prompt: str = PHYSICS_SYSTEM_PROMPT,
                      max_pixels: int = 200704):
    """Build processor inputs for (system prompt + video) + query placeholders.

    ``video_frames``: uint8 numpy array [T, H, W, 3] (a single reference video).

    The learnable queries need real sequence slots so that attention and mRoPE
    treat them as ordinary trailing text positions.  We append ``num_queries``
    copies of a harmless token; the embedding hook then overwrites their vectors.

    ``max_pixels`` caps the per-frame visual token count.  This matters because,
    unlike ordinary MLLM inference, we backpropagate through the MLLM to reach
    ``q_learnable`` -- so the video token count directly drives training memory.
    """
    messages = [
        {"role": "system", "content": [{"type": "text", "text": system_prompt}]},
        {
            "role": "user",
            "content": [
                {"type": "video"},
                {"type": "text", "text": "Describe the physical process in this video."},
            ],
        },
    ]
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    # Cap visual tokens on the video processor itself; passing max_pixels through
    # processor(**kwargs) is silently ignored by this transformers version.
    vp = getattr(processor, "video_processor", None)
    if vp is not None and getattr(vp, "max_pixels", None) != max_pixels:
        vp.max_pixels = max_pixels
        if hasattr(vp, "size") and isinstance(vp.size, dict):
            vp.size = {**vp.size, "longest_edge": max_pixels}

    inputs = processor(text=[text], videos=[video_frames], return_tensors="pt")

    # Append placeholder slots for the learnable queries.
    tok = processor.tokenizer
    pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
    b = inputs["input_ids"].shape[0]
    pad_block = torch.full((b, num_queries), pad_id, dtype=inputs["input_ids"].dtype)
    inputs["input_ids"] = torch.cat([inputs["input_ids"], pad_block], dim=1)
    if "attention_mask" in inputs:
        inputs["attention_mask"] = torch.cat(
            [inputs["attention_mask"], torch.ones_like(pad_block)], dim=1
        )
    return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in inputs.items()}