"""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()}