[Lingbot World] Bound image condition storage and extend temporal RoPE (#6838)

Signed-off-by: wtz2333 <2955110911@qq.com>

vllm_omni/diffusion/models/lingbot_world/camera.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Camera trajectory loading and geometry for LingBot World v2 conditioning."""
 
 from __future__ import annotations
@@ -19,7 +19,6 @@ import torch
 
 _REFERENCE_HEIGHT = 480
 _REFERENCE_WIDTH = 832
-_MAX_ACTION_FRAMES = 117
 _MAX_SOURCE_ACTION_FRAMES = 4096
 _MAX_NPY_HEADER_BYTES = 10_000
 _DIRECTORY_OPEN_FLAGS = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
@@ -207,9 +206,7 @@ def _load_bounded_npy(
             array = np.load(snapshot_file, allow_pickle=False)
         except (EOFError, TypeError, ValueError):
             raise ValueError(f"{filename} does not contain a valid numeric NPY array.") from None
-    # Official LingBot trajectories may be longer than one request. Only
-    # materialize the bounded prefix that any supported request can consume.
-    result: np.ndarray = np.asarray(array[:_MAX_ACTION_FRAMES], dtype=np.float32)
+    result: np.ndarray = np.asarray(array, dtype=np.float32)
     return result
 
 

vllm_omni/diffusion/models/lingbot_world/pipeline.py

@@ -7,7 +7,7 @@ from __future__ import annotations
 import math
 import os
 from collections.abc import Callable, Iterable, Iterator, Sequence
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass, replace
 from typing import TYPE_CHECKING, Any, ClassVar, cast
 
@@ -15,6 +15,7 @@ import numpy as np
 import PIL.Image
 import torch
 import torch.nn.functional as F
+from diffusers.models.autoencoders.autoencoder_kl_wan import CACHE_T, WanAttentionBlock, WanResample, WanResidualBlock
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import AutoTokenizer, UMT5EncoderModel
@@ -97,7 +98,6 @@ LINGBOT_DMD_TIMESTEPS = (1000, 750, 500, 250)
 _CAMERA_SPATIAL_FOLD = 8
 _MAX_PIXEL_AREA = 480 * 832
 _MAX_SOURCE_IMAGE_PIXELS = 4096 * 4096
-_MAX_RAW_FRAMES = 117
 _MAX_SEQUENCE_LENGTH = 512
 _ACTION_ROOT_ENV = "VLLM_OMNI_LINGBOT_ACTION_ROOT"
 _PREPROCESSED_CAMERA_KEY = "_lingbot_camera_trajectory"
@@ -150,12 +150,13 @@ class _LingBotRequestInputs:
 
 @dataclass
 class _LingBotARSessionState:
-    """Small model-owned state; attention tensors remain runner-owned."""
+    """Model-owned state; attention tensors remain runner-owned."""
 
     next_chunk_index: int = 0
     prompt: str | None = None
     generator_state: torch.Tensor | None = None
-    image_condition: torch.Tensor | None = None
+    encoder_cache: list[torch.Tensor | None] | None = None
+    pending_encoder_cache: list[torch.Tensor | None] | None = None
     camera_tail: CameraTrajectory | None = None
     camera_pitch: float = 0.0
 
@@ -577,6 +578,7 @@ class LingBotWorldCausalDMDPipeline(
         self.setup_diffusion_pipeline_profiler(
             profiler_targets=[
                 "vae.encode",
+                "_prepare_condition_chunk",
                 "vae.decode",
                 # Streaming decode never reaches vae.decode, so it needs its
                 # own stage or a streamed rollout reports no decode time at
@@ -609,11 +611,11 @@ class LingBotWorldCausalDMDPipeline(
         recent_window_frames = total_window_frames - sink_frames
         if recent_window_frames <= 0:
             raise ValueError("LingBot AR-Diffusion cache needs a positive recent window after reserving sink frames.")
-        horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1
+        condition_latent_frames = int(self.transformer.config.num_frames_per_block)
         condition_channels = self.vae_scale_factor_temporal + int(self.transformer.config.out_channels)
         condition_bytes_per_session = (
             condition_channels
-            * horizon_latent_frames
+            * condition_latent_frames
             * latent_height
             * latent_width
             * torch.empty((), dtype=self.transformer.dtype).element_size()
@@ -635,10 +637,34 @@ class LingBotWorldCausalDMDPipeline(
                 ),
             ),
             model_owned_state_bytes_per_session=(
-                condition_bytes_per_session + self._streaming_decode_bytes_per_session()
+                condition_bytes_per_session
+                + 2 * self._condition_encoder_cache_bytes()  # Committed and in-flight histories.
+                + self._streaming_decode_bytes_per_session()
             ),
         )
 
+    def _condition_encoder_cache_bytes(self) -> int:
+        """Bound one Wan encoder history from its per-convolution input grids."""
+        if getattr(self.vae.config, "patch_size", None) is not None:
+            raise ValueError("Stateful LingBot conditioning requires an unpatched Wan VAE.")
+        encoder = self.vae.encoder
+        height, width = self._ar_height, self._ar_width
+        elements = CACHE_T * encoder.conv_in.in_channels * height * width
+        for layer in encoder.down_blocks:
+            if isinstance(layer, WanResidualBlock):
+                elements += CACHE_T * (layer.conv1.in_channels + layer.conv2.in_channels) * height * width
+            elif isinstance(layer, WanResample) and layer.mode in ("downsample2d", "downsample3d"):
+                height, width = height // 2, width // 2
+                if layer.mode == "downsample3d":
+                    # Temporal downsampling retains one frame after spatial downsampling.
+                    elements += layer.time_conv.in_channels * height * width
+            elif not isinstance(layer, WanAttentionBlock):
+                raise ValueError("Stateful LingBot conditioning requires the Wan encoder cache layout.")
+        for layer in encoder.mid_block.resnets:
+            elements += CACHE_T * (layer.conv1.in_channels + layer.conv2.in_channels) * height * width
+        elements += CACHE_T * encoder.conv_out.in_channels * height * width
+        return elements * torch.empty((), dtype=self.vae.dtype).element_size()
+
     def _streaming_decode_bytes_per_session(self) -> int:
         """Resident streaming-decoder bytes one session can hold, for admission.
 
@@ -678,12 +704,18 @@ class LingBotWorldCausalDMDPipeline(
         finally:
             self._ar_diffusion_kv_state = None
 
+    def _release_condition_encoder_state(self, session_id: str) -> None:
+        state = self._ar_sessions.pop(session_id, None)
+        if state is not None:
+            state.encoder_cache = None
+            state.pending_encoder_cache = None
+
     def reset_ar_diffusion_session(self, session_id: str) -> None:
-        self._ar_sessions.pop(session_id, None)
+        self._release_condition_encoder_state(session_id)
         self._release_streaming_decode_state(session_id)
 
     def close_ar_diffusion_session(self, session_id: str) -> None:
-        self._ar_sessions.pop(session_id, None)
+        self._release_condition_encoder_state(session_id)
         self._release_streaming_decode_state(session_id)
 
     def _parse_request(self, req: DiffusionRequestBatch) -> _LingBotRequestInputs:
@@ -810,8 +842,6 @@ class LingBotWorldCausalDMDPipeline(
         num_frames = getattr(sampling, "num_frames", None)
         if isinstance(num_frames, bool) or not isinstance(num_frames, int) or num_frames <= 0:
             raise ValueError(f"num_frames must be a positive integer, got {num_frames!r}.")
-        if num_frames > _MAX_RAW_FRAMES:
-            raise ValueError(f"num_frames must not exceed {_MAX_RAW_FRAMES}.")
         temporal_factor = self.vae_scale_factor_temporal
         if (num_frames - 1) % temporal_factor:
             raise ValueError(
@@ -914,16 +944,16 @@ class LingBotWorldCausalDMDPipeline(
             raise RuntimeError(
                 f"vae.encode returned an incompatible image latent shape: got {tuple(latent_condition.shape)}."
             )
+        return self._condition_from_latents(latent_condition, first_frame=True, dtype=dtype)
+
+    def _condition_from_latents(
+        self, latent_condition: torch.Tensor, *, first_frame: bool, dtype: torch.dtype
+    ) -> torch.Tensor:
         latent_mean, latent_std = self._vae_latent_stats(latent_condition)
         latent_condition = (latent_condition - latent_mean) / latent_std
-        temporal_mask = latent_condition.new_zeros(
-            1,
-            self.vae_scale_factor_temporal,
-            inputs.num_latent_frames,
-            latent_condition.shape[-2],
-            latent_condition.shape[-1],
-        )
-        temporal_mask[:, :, 0] = 1
+        temporal_mask = latent_condition.new_zeros(1, self.vae_scale_factor_temporal, *latent_condition.shape[2:])
+        if first_frame:
+            temporal_mask[:, :, 0] = 1
         condition = torch.cat((temporal_mask, latent_condition), dim=1).to(dtype=dtype)
         if condition.shape[1] != 20:
             raise RuntimeError(
@@ -932,6 +962,78 @@ class LingBotWorldCausalDMDPipeline(
             )
         return condition
 
+    @torch.no_grad()
+    def _prepare_condition_chunk(
+        self,
+        inputs: _LingBotRequestInputs,
+        *,
+        start_frame: int,
+        encoder_cache: list[torch.Tensor | None] | None,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, list[torch.Tensor | None]]:
+        """Continue the causal encoder without modifying its committed history."""
+        if getattr(self.vae.config, "patch_size", None) is not None or self.vae_scale_factor_temporal != 4:
+            raise ValueError("Stateful LingBot conditioning requires an unpatched Wan VAE with temporal factor 4.")
+        if getattr(self.vae, "use_tiling", False) and (
+            inputs.height > self.vae.tile_sample_min_height or inputs.width > self.vae.tile_sample_min_width
+        ):
+            raise ValueError("Stateful LingBot conditioning does not support tiled VAE encoding.")
+        if (encoder_cache is None) != (start_frame == 0):
+            raise ValueError("LingBot condition encoder history must match the current temporal position.")
+
+        count = int(self.vae._cached_conv_counts["encoder"])
+        # Wan replaces cache entries with new tensors; a separate list keeps the
+        # committed history intact until this AR block succeeds.
+        pending_cache = [None] * count if encoder_cache is None else list(encoder_cache)
+        if len(pending_cache) != count:
+            raise ValueError("LingBot condition encoder cache has an incompatible number of slots.")
+        block_frames = int(self.transformer.config.num_frames_per_block)
+        zeros = torch.zeros(
+            1,
+            3,
+            self.vae_scale_factor_temporal,
+            inputs.height,
+            inputs.width,
+            device=self.device,
+            dtype=self.vae.dtype,
+        )
+        image = (
+            self._prepare_image_tensor(inputs.image, height=inputs.height, width=inputs.width)
+            .unsqueeze(2)
+            .to(dtype=self.vae.dtype)
+            if start_frame == 0
+            else None
+        )
+        context = getattr(self.vae, "_execution_context", None)
+        encoded_frames = []
+        with context() if callable(context) else nullcontext():
+            for index in range(block_frames):
+                # Match Wan's encode loop: one opening pixel frame, then four
+                # fresh zero pixel frames per subsequent latent frame.
+                pixel_frames = image if start_frame == 0 and index == 0 else zeros
+                encoded_frames.append(self.vae.encoder(pixel_frames, feat_cache=pending_cache, feat_idx=[0]))
+            # The posterior projection is pointwise in time. GPU bitwise parity
+            # with the whole-clip projection still requires the real-weight gate.
+            moments = self.vae.quant_conv(torch.cat(encoded_frames, dim=2))
+        latent_condition = moments.chunk(2, dim=1)[0]
+        expected_shape = (
+            1,
+            self.transformer.config.out_channels,
+            block_frames,
+            inputs.height // self.vae_scale_factor_spatial,
+            inputs.width // self.vae_scale_factor_spatial,
+        )
+        if latent_condition.shape != expected_shape:
+            raise RuntimeError(
+                f"Wan condition encoder returned {tuple(latent_condition.shape)}, expected {expected_shape}."
+            )
+        condition = self._condition_from_latents(
+            latent_condition,
+            first_frame=start_frame == 0,
+            dtype=dtype,
+        )
+        return condition, pending_cache
+
     def _prepare_camera(
         self,
         inputs: _LingBotRequestInputs,
@@ -1150,6 +1252,11 @@ class LingBotWorldCausalDMDPipeline(
             raise RuntimeError("LingBot typed ticks require ARDiffusionEngine session binding.")
         if tick is None and self._ar_diffusion_kv_state is not None:
             raise ValueError("LingBot ARDiffusionEngine requests must carry ar_diffusion_tick.")
+        if tick is not None and inputs.output_type != "latent":
+            raise ValueError(
+                "LingBot realtime ticks currently require output_type='latent'; "
+                "stateful streaming VAE decode is a separate integration step."
+            )
         session_state: _LingBotARSessionState | None = None
         if tick is not None:
             if tick.prompt is not None and " ".join(tick.prompt.split()) != " ".join(inputs.prompt.split()):
@@ -1166,16 +1273,6 @@ class LingBotWorldCausalDMDPipeline(
                     "fixed cache geometry "
                     f"{self._ar_height}x{self._ar_width}."
                 )
-            horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1
-            max_realtime_ticks = horizon_latent_frames // block_frames
-            if tick.chunk_index >= max_realtime_ticks:
-                raise ValueError(
-                    "LingBot realtime generation currently supports at most "
-                    f"{max_realtime_ticks} ticks per generation epoch "
-                    f"(chunk_index 0 through {max_realtime_ticks - 1}) because "
-                    f"the image-condition horizon is {_MAX_RAW_FRAMES} pixel "
-                    "frames; reset or create a session to start a new world."
-                )
             session_state = self._ar_sessions.setdefault(
                 tick.session_id,
                 _LingBotARSessionState(),
@@ -1200,26 +1297,12 @@ class LingBotWorldCausalDMDPipeline(
             condition = self._prepare_condition(inputs, dtype=dtype)
         else:
             assert session_state is not None
-            if session_state.image_condition is None:
-                horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1
-                session_state.image_condition = self._prepare_condition(
-                    replace(
-                        inputs,
-                        num_frames=_MAX_RAW_FRAMES,
-                        num_latent_frames=horizon_latent_frames,
-                    ),
-                    dtype=dtype,
-                )
-            block_frames = int(self.transformer.config.num_frames_per_block)
-            condition_start = tick.chunk_index * block_frames
-            condition_stop = condition_start + block_frames
-            if condition_stop > session_state.image_condition.shape[2]:
-                raise ValueError("LingBot chunk_index exceeds the configured causal image condition horizon.")
-            condition = session_state.image_condition[
-                :,
-                :,
-                condition_start:condition_stop,
-            ]
+            condition, pending_encoder_cache = self._prepare_condition_chunk(
+                inputs,
+                start_frame=tick.chunk_index * block_frames,
+                encoder_cache=session_state.encoder_cache,
+                dtype=dtype,
+            )
         camera_pitch: float | None = None
         if inputs.camera_actions is not None:
             assert session_state is not None
@@ -1299,6 +1382,7 @@ class LingBotWorldCausalDMDPipeline(
             assert session_state is not None
             session_state.prompt = inputs.prompt
             session_state.generator_state = inputs.generator.get_state()
+            session_state.encoder_cache = pending_encoder_cache
             session_state.camera_tail = camera_tail
             if camera_pitch is not None:
                 session_state.camera_pitch = camera_pitch
@@ -1307,11 +1391,6 @@ class LingBotWorldCausalDMDPipeline(
         # Phase 4: either expose model-space latents or invert the checkpoint's
         # latent normalization and decode to pixel-space video.
         if tick is not None:
-            if inputs.output_type != "latent":
-                raise ValueError(
-                    "LingBot realtime ticks currently require output_type='latent'; "
-                    "stateful streaming VAE decode is a separate integration step."
-                )
             output = {
                 "payload": {"latents": generated_latents},
                 "metadata": {"ar_diffusion": ARDiffusionChunkMetadata.from_tick(tick).to_dict()},
@@ -1460,12 +1539,6 @@ class LingBotWorldCausalDMDPipeline(
             self._cached_video_processor = processor
         return processor
 
-    def _horizon_latent_frames(self) -> int:
-        return (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1
-
-    def _max_realtime_chunks(self) -> int:
-        return self._horizon_latent_frames() // int(self.transformer.config.num_frames_per_block)
-
     def _require_bound_ar_state(self) -> None:
         if self._ar_diffusion_kv_state is None:
             raise RuntimeError("LingBot step execution requires AR-Diffusion session binding.")
@@ -1492,13 +1565,6 @@ class LingBotWorldCausalDMDPipeline(
             )
         block_frames = int(self.transformer.config.num_frames_per_block)
         total_chunks = inputs.num_latent_frames // block_frames
-        max_chunks = self._max_realtime_chunks()
-        if total_chunks > max_chunks:
-            raise ValueError(
-                "LingBot step execution currently supports at most "
-                f"{max_chunks} chunks because the image-condition horizon is "
-                f"{_MAX_RAW_FRAMES} pixel frames."
-            )
         if inputs.camera_action_script is not None and len(inputs.camera_action_script) != total_chunks:
             raise ValueError(
                 "camera_action_script must contain one action list per generated chunk; "
@@ -1510,14 +1576,6 @@ class LingBotWorldCausalDMDPipeline(
             max_sequence_length=inputs.max_sequence_length,
             dtype=dtype,
         )
-        image_condition = self._prepare_condition(
-            replace(
-                inputs,
-                num_frames=_MAX_RAW_FRAMES,
-                num_latent_frames=self._horizon_latent_frames(),
-            ),
-            dtype=dtype,
-        )
         camera_trajectory_cache = None
         camera_embedding_cache = None
         if inputs.camera_trajectory is not None:
@@ -1546,7 +1604,6 @@ class LingBotWorldCausalDMDPipeline(
         state.chunk_num_steps = len(LINGBOT_DMD_TIMESTEPS)
         state.extra = {
             "inputs": inputs,
-            "image_condition": image_condition,
             "schedule": _build_shifted_flow_schedule(flow_shift=inputs.flow_shift),
             "dtype": dtype,
             "block_frames": block_frames,
@@ -1557,7 +1614,12 @@ class LingBotWorldCausalDMDPipeline(
             "camera_embedding_cache": camera_embedding_cache,
         }
         self._ar_text_caches(prompt_embeds, invalidate=False)
-        self._prepare_next_chunk(state)
+        self._ar_sessions[state.request_id] = _LingBotARSessionState()
+        try:
+            self._prepare_next_chunk(state)
+        except Exception:
+            self._release_condition_encoder_state(state.request_id)
+            raise
         return state
 
     def _prepare_next_chunk(self, state: StepRequestState) -> None:
@@ -1566,10 +1628,15 @@ class LingBotWorldCausalDMDPipeline(
         block_frames = int(extra["block_frames"])
         start_frame = state.chunk_index * block_frames
         stop_frame = start_frame + block_frames
-        image_condition = extra["image_condition"]
-        if stop_frame > image_condition.shape[2]:
-            raise ValueError("LingBot chunk_index exceeds the configured causal image condition horizon.")
-        condition = image_condition[:, :, start_frame:stop_frame]
+        session_state = self._ar_sessions[state.request_id]
+        if session_state.next_chunk_index != state.chunk_index:
+            raise ValueError("LingBot condition chunks must be contiguous within the session.")
+        condition, pending_encoder_cache = self._prepare_condition_chunk(
+            inputs,
+            start_frame=start_frame,
+            encoder_cache=session_state.encoder_cache,
+            dtype=extra["dtype"],
+        )
         previous = extra.get("camera_tail")
         if extra.get("camera_action_script") is not None:
             chunk_actions = extra["camera_action_script"][state.chunk_index]
@@ -1639,6 +1706,7 @@ class LingBotWorldCausalDMDPipeline(
         )
         state.step_in_chunk = 0
         state.step_index = 0
+        session_state.pending_encoder_cache = pending_encoder_cache
 
     def denoise_step(
         self,
@@ -1702,6 +1770,9 @@ class LingBotWorldCausalDMDPipeline(
         latents = state.latents
         if latents is None:
             raise RuntimeError("LingBot step execution requires latents before post_decode.")
+        session_state = self._ar_sessions[state.request_id]
+        if session_state.pending_encoder_cache is None:
+            raise RuntimeError("LingBot condition encoder has no prepared chunk to commit.")
         self._dmd_blocks.commit_block_kv(
             latents=latents,
             condition=extra["condition"],
@@ -1737,10 +1808,15 @@ class LingBotWorldCausalDMDPipeline(
                 ).to_dict()
             },
         }
+        session_state.encoder_cache = session_state.pending_encoder_cache
+        session_state.pending_encoder_cache = None
+        session_state.next_chunk_index += 1
         state.chunk_index += 1
         finished = state.request_denoise_completed
         if not finished:
             self._prepare_next_chunk(state)
+        else:
+            self._release_condition_encoder_state(state.request_id)
         return DiffusionOutput(
             output=output,
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,

vllm_omni/diffusion/models/lingbot_world/transformer.py

@@ -749,7 +749,7 @@ def _sinusoidal_embedding(dim: int, timestep: torch.Tensor) -> torch.Tensor:
     return torch.cat((phase.cos(), phase.sin()), dim=1)
 
 
-def _rope_axis(max_seq_len: int, dim: int) -> tuple[torch.Tensor, torch.Tensor]:
+def _rope_axis(max_seq_len: int, dim: int, *, start: int = 0) -> tuple[torch.Tensor, torch.Tensor]:
     if dim == 0:
         empty = torch.empty(max_seq_len, 0, dtype=torch.float32)
         return empty, empty.clone()
@@ -759,7 +759,7 @@ def _rope_axis(max_seq_len: int, dim: int) -> tuple[torch.Tensor, torch.Tensor]:
         10000,
         torch.arange(0, dim, 2, dtype=torch.float64) / dim,
     )
-    phase = torch.outer(torch.arange(max_seq_len, dtype=torch.float64), frequencies)
+    phase = torch.outer(torch.arange(start, start + max_seq_len, dtype=torch.float64), frequencies)
     return phase.cos().float(), phase.sin().float()
 
 
@@ -1028,23 +1028,29 @@ class CausalLingBotWorldTransformer3DModel(nn.Module):
         dtype: torch.dtype,
         device: torch.device,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        if start_frame + frames > self.config.rope_max_seq_len:
-            raise ValueError("Temporal RoPE positions exceed rope_max_seq_len.")
         if height > self.config.rope_max_seq_len or width > self.config.rope_max_seq_len:
             raise ValueError("Spatial RoPE positions exceed rope_max_seq_len.")
+        if start_frame + frames <= self._rope_temporal_cosine.shape[0]:
+            temporal_cosine = self._rope_temporal_cosine[start_frame : start_frame + frames]
+            temporal_sine = self._rope_temporal_sine[start_frame : start_frame + frames]
+        else:
+            # Keep absolute positions and checkpoint frequencies, computing only
+            # this block instead of growing a table for the whole session.
+            temporal_dim = self._rope_temporal_cosine.shape[1] * 2
+            temporal_cosine, temporal_sine = _rope_axis(frames, temporal_dim, start=start_frame)
+            temporal_cosine = temporal_cosine.to(self._rope_temporal_cosine)
+            temporal_sine = temporal_sine.to(self._rope_temporal_sine)
 
         def expand_axis(table: torch.Tensor, axis: str) -> torch.Tensor:
             if axis == "temporal":
-                return (
-                    table[start_frame : start_frame + frames].view(frames, 1, 1, -1).expand(frames, height, width, -1)
-                )
+                return table.view(frames, 1, 1, -1).expand(frames, height, width, -1)
             if axis == "height":
                 return table[:height].view(1, height, 1, -1).expand(frames, height, width, -1)
             return table[:width].view(1, 1, width, -1).expand(frames, height, width, -1)
 
         cosine = torch.cat(
             (
-                expand_axis(self._rope_temporal_cosine, "temporal"),
+                expand_axis(temporal_cosine, "temporal"),
                 expand_axis(self._rope_height_cosine, "height"),
                 expand_axis(self._rope_width_cosine, "width"),
             ),
@@ -1052,7 +1058,7 @@ class CausalLingBotWorldTransformer3DModel(nn.Module):
         )
         sine = torch.cat(
             (
-                expand_axis(self._rope_temporal_sine, "temporal"),
+                expand_axis(temporal_sine, "temporal"),
                 expand_axis(self._rope_height_sine, "height"),
                 expand_axis(self._rope_width_sine, "width"),
             ),