dreamx_world_ar ¶

def forward(self, x, cam_viewmats, cam_K, seq_lens, grid_sizes, freqs,
            kv_cache=None, current_start=0, cache_start=None,
            sink_recache_after_switch=False, cache_update_policy="commit_detached"):
    """
    Args:
        x: Shape [B, L, C]
        cam_viewmats: Camera view matrices
        cam_K: Camera intrinsics
        kv_cache: Optional KV cache dict. When None, runs full attention over current chunk.
    """
    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
    if cache_start is None:
        cache_start = current_start

    q, _ = self.q_proj(x)
    q = self.norm_q(q).view(b, s, n, d)
    k, _ = self.k_proj(x)
    k = self.norm_k(k).view(b, s, n, d)
    v, _ = self.v_proj(x)
    v = v.view(b, s, n, d)

    # Apply PRoPE (Positional Rotary Position Embedding from camera parameters)
    q_t, k_t, v_t, apply_fn_o = prope_qkv(
        q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
        viewmats=cam_viewmats, Ks=cam_K)
    proped_q = q_t.transpose(1, 2)
    proped_k = k_t.transpose(1, 2)
    proped_v = v_t.transpose(1, 2)

    if kv_cache is None:
        # No cache: full attention over current chunk
        x_out = attention(proped_q, proped_k, proped_v)
    else:
        # KV cache mode with rolling cache support
        frame_seqlen = math.prod(grid_sizes[0][1:]).item()
        num_new_tokens = s
        current_end = current_start + num_new_tokens
        sink_tokens = self.sink_size * frame_seqlen
        kv_cache_size = kv_cache["k"].shape[1]
        is_recompute = (current_end <= kv_cache["global_end_index"].item()) and (current_start > 0)

        if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
                num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
            # === ROLLING MODE ===
            num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
            num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
            local_end_index = kv_cache["local_end_index"].item() + current_end - \
                kv_cache["global_end_index"].item() - num_evicted_tokens
            local_start_index = local_end_index - num_new_tokens

            if cache_update_policy != "none":
                with torch.no_grad():
                    kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
                        kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
                    kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
                        kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()

                write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
                roped_offset = max(0, write_start_index - local_start_index)
                write_len = max(0, local_end_index - write_start_index)
                if write_len > 0:
                    with torch.no_grad():
                        kv_cache["k"][:, write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset + write_len].detach()
                        kv_cache["v"][:, write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset + write_len].detach()
        else:
            # === DIRECT INSERT MODE ===
            local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
            local_start_index = local_end_index - num_new_tokens

            if cache_update_policy != "none":
                write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
                if sink_recache_after_switch:
                    write_start_index = local_start_index
                roped_offset = max(0, write_start_index - local_start_index)
                write_len = max(0, local_end_index - write_start_index)
                if write_len > 0:
                    with torch.no_grad():
                        kv_cache["k"][:, write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset + write_len].detach()
                        kv_cache["v"][:, write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset + write_len].detach()

        # Attention: sink tokens + local window
        if sink_tokens > 0:
            local_budget = self.max_attention_size - sink_tokens
            k_sink = kv_cache["k"][:, :sink_tokens].detach()
            v_sink = kv_cache["v"][:, :sink_tokens].detach()
            if local_budget > 0:
                local_start_for_window = max(sink_tokens, local_end_index - local_budget)
                k_local = kv_cache["k"][:, local_start_for_window:local_end_index].detach()
                v_local = kv_cache["v"][:, local_start_for_window:local_end_index].detach()
                k_cat = torch.cat([k_sink, k_local], dim=1)
                v_cat = torch.cat([v_sink, v_local], dim=1)
            else:
                k_cat = k_sink
                v_cat = v_sink
            x_out = attention(proped_q, k_cat, v_cat)
        else:
            window_start = max(0, local_end_index - self.max_attention_size)
            x_out = attention(
                proped_q,
                kv_cache["k"][:, window_start:local_end_index].detach(),
                kv_cache["v"][:, window_start:local_end_index].detach())

        if not is_recompute and cache_update_policy != "none":
            kv_cache["global_end_index"].fill_(current_end)
            kv_cache["local_end_index"].fill_(local_end_index)

    # Apply inverse PRoPE
    x = apply_fn_o(x_out.transpose(1, 2)).transpose(1, 2)
    x = x.flatten(2)
    x, _ = self.out_proj(x)
    return x

def forward(self, x, seq_lens, grid_sizes, freqs, kv_cache,
            current_start=0, cache_start=None, sink_recache_after_switch=False):
    """
    Args:
        x: Shape [B, L, C]
        seq_lens: Shape [B]
        grid_sizes: Shape [B, 3] containing (F, H, W)
        freqs: RoPE frequencies [1024, head_dim / 2]
        kv_cache: Dict with 'k', 'v', 'global_end_index', 'local_end_index'
        current_start: Current position in the global token sequence
    """
    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
    if cache_start is None:
        cache_start = current_start

    q, _ = self.q(x)
    q = self.norm_q(q).view(b, s, n, d)
    k, _ = self.k(x)
    k = self.norm_k(k).view(b, s, n, d)
    v, _ = self.v(x)
    v = v.view(b, s, n, d)

    frame_seqlen = math.prod(grid_sizes[0][1:]).item()
    num_new_frames = grid_sizes[0][0].item()
    current_end = current_start + q.shape[1]
    sink_tokens = self.sink_size * frame_seqlen
    kv_cache_size = kv_cache["k"].shape[1]
    num_new_tokens = q.shape[1]

    cache_update_info = None
    is_recompute = current_end <= kv_cache["global_end_index"].item() and current_start > 0

    if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
            num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
        # === ROLLING MODE: cache full, evict oldest non-sink tokens ===
        num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
        num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
        local_end_index = kv_cache["local_end_index"].item() + current_end - \
            kv_cache["global_end_index"].item() - num_evicted_tokens
        local_start_index = local_end_index - num_new_tokens

        temp_k = kv_cache["k"].detach().clone()
        temp_v = kv_cache["v"].detach().clone()
        temp_k[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
            temp_k[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
        temp_v[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
            temp_v[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()

        write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
        roped_offset = max(0, write_start_index - local_start_index)
        write_len = max(0, local_end_index - write_start_index)
        if write_len > 0:
            temp_k[:, write_start_index:local_end_index] = k[:, roped_offset:roped_offset + write_len]
            temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len]

        # Block-Relativistic RoPE: query uses window-relative indices
        query_relative_indices = torch.arange(
            self.local_attn_size - num_new_frames, self.local_attn_size, device=q.device)
        roped_query = block_relativistic_rope(
            q, grid_sizes, freqs, relative_frame_indices=query_relative_indices).type_as(v)

        # Block-Relativistic RoPE: cached K uses position-in-window indices
        num_cache_frames = local_end_index // frame_seqlen
        cache_relative_indices = torch.arange(0, num_cache_frames, device=k.device)
        cache_grid_sizes = grid_sizes.clone()
        cache_grid_sizes[0, 0] = num_cache_frames
        roped_temp_k = block_relativistic_rope(
            temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen, n, d).flatten(1, 2),
            cache_grid_sizes, freqs, relative_frame_indices=cache_relative_indices).type_as(v)

        cache_update_info = {
            "action": "roll_and_insert",
            "sink_tokens": sink_tokens,
            "num_rolled_tokens": num_rolled_tokens,
            "num_evicted_tokens": num_evicted_tokens,
            "local_start_index": local_start_index,
            "local_end_index": local_end_index,
            "write_start_index": write_start_index,
            "write_end_index": local_end_index,
            "new_k": k[:, roped_offset:roped_offset + write_len],
            "new_v": v[:, roped_offset:roped_offset + write_len],
            "current_end": current_end,
            "is_recompute": is_recompute
        }
    else:
        # === DIRECT INSERT MODE: cache not yet full ===
        local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
        local_start_index = local_end_index - num_new_tokens

        temp_k = kv_cache["k"].detach().clone()
        temp_v = kv_cache["v"].detach().clone()

        write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
        if sink_recache_after_switch:
            write_start_index = local_start_index
        roped_offset = max(0, write_start_index - local_start_index)
        write_len = max(0, local_end_index - write_start_index)
        if write_len > 0:
            temp_k[:, write_start_index:local_end_index] = k[:, roped_offset:roped_offset + write_len]
            temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len]

        # RoPE with relative indices (growing sequentially before cache fills)
        current_frame_in_window = local_start_index // frame_seqlen
        query_relative_indices = torch.arange(
            current_frame_in_window, current_frame_in_window + num_new_frames, device=q.device)
        roped_query = block_relativistic_rope(
            q, grid_sizes, freqs, relative_frame_indices=query_relative_indices).type_as(v)

        num_cache_frames = local_end_index // frame_seqlen
        cache_relative_indices = torch.arange(0, num_cache_frames, device=k.device)
        cache_grid_sizes = grid_sizes.clone()
        cache_grid_sizes[0, 0] = num_cache_frames
        roped_temp_k = block_relativistic_rope(
            temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen, n, d).flatten(1, 2),
            cache_grid_sizes, freqs, relative_frame_indices=cache_relative_indices).type_as(v)

        cache_update_info = {
            "action": "direct_insert",
            "local_start_index": local_start_index,
            "local_end_index": local_end_index,
            "write_start_index": write_start_index,
            "write_end_index": local_end_index,
            "new_k": k[:, roped_offset:roped_offset + write_len],
            "new_v": v[:, roped_offset:roped_offset + write_len],
            "current_end": current_end,
            "is_recompute": is_recompute
        }

    # Attention: sink tokens + local window
    if sink_tokens > 0:
        local_budget = self.max_attention_size - sink_tokens
        k_sink = roped_temp_k[:, :sink_tokens]
        v_sink = temp_v[:, :sink_tokens]
        if local_budget > 0:
            local_start_for_window = max(sink_tokens, local_end_index - local_budget)
            k_local = roped_temp_k[:, local_start_for_window:local_end_index]
            v_local = temp_v[:, local_start_for_window:local_end_index]
            k_cat = torch.cat([k_sink, k_local], dim=1)
            v_cat = torch.cat([v_sink, v_local], dim=1)
        else:
            k_cat = k_sink
            v_cat = v_sink
        x = attention(roped_query, k_cat, v_cat)
    else:
        window_start = max(0, local_end_index - self.max_attention_size)
        x = attention(
            roped_query,
            roped_temp_k[:, window_start:local_end_index],
            temp_v[:, window_start:local_end_index])

    x = x.flatten(2)
    x, _ = self.o(x)
    return x, (current_end, local_end_index, cache_update_info)

def forward(self, x=None, t=None, context=None, seq_len=None, y=None, y_camera=None,
            kv_cache=None, crossattn_cache=None, current_start=0,
            cache_start=0, cache_update_policy="commit_detached",
            hidden_states=None, encoder_hidden_states=None, timestep=None, **kwargs):
    """
    Causal inference with KV caching.
    See Algorithm 2 of CausVid (https://arxiv.org/abs/2412.07772).

    Args:
        x: List of input video tensors [C_in, F, H, W]
        t: Timestep tensor [B, L]
        context: List of text embeddings [L, C]
        seq_len: Maximum sequence length for positional encoding
        y: Optional conditional video inputs (I2V mode)
        y_camera: Camera parameters dict {'viewmats': ..., 'K': ...}
        kv_cache: List of KV cache dicts per transformer block
        crossattn_cache: List of cross-attention cache dicts
        current_start: Current position in global token sequence
        cache_start: Cache start position
        cache_update_policy: Cache update strategy ('commit_detached' or 'none')

    Returns:
        Stacked output tensors [B, C_out, F, H/8, W/8]
    """
    if x is None and hidden_states is not None:
        x = [sample for sample in hidden_states]
    if t is None and timestep is not None:
        t = timestep
    if context is None and encoder_hidden_states is not None:
        if isinstance(encoder_hidden_states, torch.Tensor):
            context = [sample for sample in encoder_hidden_states]
        else:
            context = encoder_hidden_states
    if seq_len is None:
        if torch.is_tensor(t):
            seq_len = int(t.shape[1]) if t.dim() > 1 else int(t.numel())
        elif x is not None:
            sample = x[0]
            seq_len = (sample.shape[1] // self.patch_size[0]) * (sample.shape[2] // self.patch_size[1]) * (sample.shape[3] // self.patch_size[2])
    if x is None or t is None or context is None or seq_len is None:
        raise ValueError("DreamXWorldARTransformer3DModel requires x/t/context/seq_len or FastVideo aliases")

    device = self.patch_embedding.weight.device
    if self.freqs.is_meta or self.freqs.device != device:
        d = self.dim // self.num_heads
        self.freqs = torch.cat([
            rope_params(1024, d - 4 * (d // 6)),
            rope_params(1024, 2 * (d // 6)),
            rope_params(1024, 2 * (d // 6)),
        ], dim=1).to(device)

    if y is not None:
        x = [torch.cat([u, v], dim=0) for u, v in zip(x, y, strict=True)]

    # patch embedding
    x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
    grid_sizes = torch.stack(
        [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
    x = [u.flatten(2).transpose(1, 2) for u in x]
    seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
    assert seq_lens.max() <= seq_len
    x = torch.cat(x)

    # time embedding
    e = self.time_embedding(
        sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x))
    e0 = self.time_projection(e).unflatten(
        1, (6, self.dim)).unflatten(dim=0, sizes=t.shape)

    # text embedding
    context_lens = None
    context = self.text_embedding(
        torch.stack([
            torch.cat(
                [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
            for u in context
        ]))

    # camera parameters
    if y_camera is not None and isinstance(y_camera, dict):
        cam_viewmats = y_camera['viewmats']
        cam_K = y_camera['K']
    else:
        cam_viewmats = None
        cam_K = None

    block_kwargs = dict(
        e=e0, seq_lens=seq_lens, grid_sizes=grid_sizes, freqs=self.freqs,
        context=context, context_lens=context_lens,
        cam_viewmats=cam_viewmats, cam_K=cam_K,
        cache_update_policy=cache_update_policy,
    )

    cache_update_infos = []
    for block_index, block in enumerate(self.blocks):
        block_kwargs.update({
            "kv_cache": kv_cache[block_index] if kv_cache is not None else None,
            "crossattn_cache": crossattn_cache[block_index] if crossattn_cache is not None else None,
            "current_start": current_start,
            "cache_start": cache_start,
        })
        x, block_cache_update_info = block(x, **block_kwargs)
        if kv_cache is not None:
            cache_update_infos.append((block_index, block_cache_update_info))

    # Apply deferred cache updates
    if kv_cache is not None and cache_update_infos and cache_update_policy != "none":
        self._apply_cache_updates(kv_cache, cache_update_infos)

    # head & unpatchify
    x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2))
    x = self.unpatchify(x, grid_sizes)
    return torch.stack(x)

def init_weights(self):
    """Initialize model parameters using Xavier initialization."""
    for m in self.modules():
        if isinstance(m, (nn.Linear, ReplicatedLinear)):
            nn.init.xavier_uniform_(m.weight)
            if m.bias is not None:
                nn.init.zeros_(m.bias)

    nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
    for m in self.text_embedding.modules():
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=.02)
    for m in self.time_embedding.modules():
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=.02)

    nn.init.zeros_(self.head.head.weight)

def unpatchify(self, x, grid_sizes):
    """Reconstruct video tensors from patch embeddings."""
    c = self.out_dim
    out = []
    for u, v in zip(x, grid_sizes.tolist(), strict=True):
        u = u[:math.prod(v)].view(*v, *self.patch_size, c)
        u = torch.einsum('fhwpqrc->cfphqwr', u)
        u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size, strict=True)])
        out.append(u)
    return out