methods/cache_strategy/WorldCache/runtime.py

from __future__ import annotations

import types
from typing import List, Optional, Tuple, Union

import torch
import torch.amp as amp
import torch.nn.functional as F
from einops import rearrange

from methods.cache_strategy.common import (
    WorldCacheConfig,
    _maybe_concat_condition_mask,
    _maybe_embed_action,
    initialize_worldcache_state,
    prepare_cache_runtime_model,
)

try:
    try:
        from cosmos_predict2.conditioner import DataType
    except ImportError:
        from cosmos_predict2._src.predict2.conditioner import DataType
except Exception:  # pragma: no cover - test fallback for minimal environments
    from enum import Enum

    class DataType(Enum):
        VIDEO = "video"
        IMAGE = "image"

try:
    try:
        from imaginaire.utils import log
    except ImportError:
        from cosmos_predict2._src.imaginaire.utils import log
except Exception:  # pragma: no cover - test fallback for minimal environments
    class _FallbackLog:
        @staticmethod
        def info(*args, **kwargs):
            pass

    log = _FallbackLog()


def estimate_optical_flow(prev_img_tensor, curr_img_tensor, scale_factor=0.5):
    """GPU-native Lucas-Kanade optical flow. (B,C,H,W) -> (B,H,W,2)."""
    original_h, original_w = prev_img_tensor.shape[2], prev_img_tensor.shape[3]
    if scale_factor != 1.0:
        prev_img_tensor = F.interpolate(
            prev_img_tensor,
            scale_factor=scale_factor,
            mode="bilinear",
            align_corners=False,
        )
        curr_img_tensor = F.interpolate(
            curr_img_tensor,
            scale_factor=scale_factor,
            mode="bilinear",
            align_corners=False,
        )
    i1 = prev_img_tensor.mean(dim=1, keepdim=True)
    i2 = curr_img_tensor.mean(dim=1, keepdim=True)
    k_x = torch.tensor(
        [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]],
        dtype=i1.dtype,
        device=i1.device,
    ).view(1, 1, 3, 3)
    k_y = torch.tensor(
        [[-1, -2, -1], [0, 0, 0], [1, 2, 1]],
        dtype=i1.dtype,
        device=i1.device,
    ).view(1, 1, 3, 3)
    i_x = F.conv2d(i1, k_x, padding=1)
    i_y = F.conv2d(i1, k_y, padding=1)
    i_t = i2 - i1
    win = 21
    avg = torch.nn.AvgPool2d(kernel_size=win, stride=1, padding=win // 2)
    s_ix2 = avg(i_x * i_x)
    s_iy2 = avg(i_y * i_y)
    s_ixiy = avg(i_x * i_y)
    s_ixit = avg(i_x * i_t)
    s_iyit = avg(i_y * i_t)
    det = s_ix2 * s_iy2 - s_ixiy * s_ixiy + 1e-6
    u = -(s_iy2 * s_ixit - s_ixiy * s_iyit) / det
    v = -(s_ix2 * s_iyit - s_ixiy * s_ixit) / det
    flow = torch.cat((u, v), dim=1)
    if scale_factor != 1.0:
        flow = F.interpolate(
            flow,
            size=(original_h, original_w),
            mode="bilinear",
            align_corners=False,
        ) * (1.0 / scale_factor)
    return flow.permute(0, 2, 3, 1)


def warp_feature(feature_tensor, flow_tensor):
    if flow_tensor is None:
        return feature_tensor
    bsz, _channels, height, width = feature_tensor.shape
    if flow_tensor.ndim == 3:
        flow_tensor = flow_tensor.unsqueeze(0).expand(bsz, -1, -1, -1)
    elif flow_tensor.ndim == 4 and flow_tensor.shape[0] == 1 and bsz > 1:
        flow_tensor = flow_tensor.expand(bsz, -1, -1, -1)
    flow_tensor = flow_tensor.to(device=feature_tensor.device, dtype=feature_tensor.dtype)
    yy, xx = torch.meshgrid(
        torch.arange(height, device=feature_tensor.device, dtype=feature_tensor.dtype),
        torch.arange(width, device=feature_tensor.device, dtype=feature_tensor.dtype),
        indexing="ij",
    )
    xx = xx.unsqueeze(0).expand(bsz, -1, -1)
    yy = yy.unsqueeze(0).expand(bsz, -1, -1)
    gx = 2.0 * (xx + flow_tensor[..., 0]) / max(width - 1, 1) - 1.0
    gy = 2.0 * (yy + flow_tensor[..., 1]) / max(height - 1, 1) - 1.0
    grid = torch.stack((gx, gy), dim=3)
    return F.grid_sample(
        feature_tensor,
        grid,
        mode="bilinear",
        padding_mode="reflection",
        align_corners=True,
    )


def compute_hf_drift(prev_img, curr_img):
    channels = prev_img.shape[1]
    kernel = torch.tensor(
        [[0, -1, 0], [-1, 4, -1], [0, -1, 0]],
        dtype=prev_img.dtype,
        device=prev_img.device,
    ).view(1, 1, 3, 3).repeat(channels, 1, 1, 1)
    prev_hf = F.conv2d(prev_img, kernel, padding=1, groups=channels)
    curr_hf = F.conv2d(curr_img, kernel, padding=1, groups=channels)
    return (prev_hf - curr_hf).abs().mean()


def compute_saliency_map(features):
    if features.ndim == 5:
        bsz, time, height, width, channels = features.shape
        features = rearrange(features, "b t h w d -> b (t d) h w")
    saliency = torch.std(features, dim=1, keepdim=True)
    bsz = saliency.shape[0]
    saliency_flat = saliency.view(bsz, -1)
    mn = saliency_flat.min(dim=1, keepdim=True)[0].view(bsz, 1, 1, 1)
    mx = saliency_flat.max(dim=1, keepdim=True)[0].view(bsz, 1, 1, 1)
    return (saliency - mn) / (mx - mn + 1e-6)


def compute_optimal_gamma(delta_curr, delta_prev):
    bsz = delta_curr.shape[0]
    delta_curr_flat = delta_curr.view(bsz, -1)
    delta_prev_flat = delta_prev.view(bsz, -1)
    numer = (delta_curr_flat * delta_prev_flat).sum(dim=1, keepdim=True)
    denom = (delta_prev_flat * delta_prev_flat).sum(dim=1, keepdim=True)
    gamma = torch.clamp(numer / (denom + 1e-6), 0.2, 1.8)
    return gamma.view(bsz, 1, 1, 1, 1)


def worldcache_mini_train_dit_forward(
    self,
    x_B_C_T_H_W: torch.Tensor,
    timesteps_B_T: torch.Tensor,
    crossattn_emb: torch.Tensor,
    fps: Optional[torch.Tensor] = None,
    padding_mask: Optional[torch.Tensor] = None,
    data_type: Optional[DataType] = DataType.VIDEO,
    intermediate_feature_ids: Optional[List[int]] = None,
    img_context_emb: Optional[torch.Tensor] = None,
    condition_video_input_mask_B_C_T_H_W: Optional[torch.Tensor] = None,
    **kwargs,
) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]]]:
    del intermediate_feature_ids
    assert isinstance(data_type, DataType), f"Expected DataType, got {type(data_type)}."

    if kwargs.get("timestep_scale") is None and hasattr(self, "timestep_scale"):
        timesteps_B_T = timesteps_B_T * self.timestep_scale
    x_B_C_T_H_W = _maybe_concat_condition_mask(
        self,
        x_B_C_T_H_W,
        is_video=data_type == DataType.VIDEO,
        condition_video_input_mask_B_C_T_H_W=condition_video_input_mask_B_C_T_H_W,
    )

    x_B_T_H_W_D, rope_emb_L_1_1_D, extra_pos_emb = self.prepare_embedded_sequence(
        x_B_C_T_H_W,
        fps=fps,
        padding_mask=padding_mask,
    )

    if self.crossattn_proj is not None:
        crossattn_emb = self.crossattn_proj(crossattn_emb)

    if img_context_emb is not None:
        assert self.extra_image_context_dim is not None
        img_context_emb = self.img_context_proj(img_context_emb)
        context_input = (crossattn_emb, img_context_emb)
    else:
        context_input = crossattn_emb

    with amp.autocast("cuda", enabled=getattr(self, "use_wan_fp32_strategy", False), dtype=torch.float32):
        if timesteps_B_T.ndim == 1:
            timesteps_B_T = timesteps_B_T.unsqueeze(1)
        t_embedding_B_T_D, adaln_lora_B_T_3D = self.t_embedder(timesteps_B_T)
        t_embedding_B_T_D, adaln_lora_B_T_3D = _maybe_embed_action(
            self,
            t_embedding_B_T_D,
            adaln_lora_B_T_3D,
            kwargs,
        )
        t_embedding_B_T_D = self.t_embedding_norm(t_embedding_B_T_D)

    self.affline_scale_log_info = {"t_embedding_B_T_D": t_embedding_B_T_D.detach()}
    self.affline_emb = t_embedding_B_T_D
    self.crossattn_emb = crossattn_emb

    if extra_pos_emb is not None:
        assert x_B_T_H_W_D.shape == extra_pos_emb.shape

    batch, time, _height, _width, _dim = x_B_T_H_W_D.shape
    block_kwargs = {
        "emb_B_T_D": t_embedding_B_T_D,
        "crossattn_emb": context_input,
        "rope_emb_L_1_1_D": rope_emb_L_1_1_D,
        "adaln_lora_B_T_3D": adaln_lora_B_T_3D,
        "extra_per_block_pos_emb": extra_pos_emb,
    }

    skip_forward = False
    ori_x = x_B_T_H_W_D
    residual_x = None

    is_parallel_cfg = getattr(self, "worldcache_parallel_cfg", False)
    current_idx = 0 if is_parallel_cfg else self.cnt % 2
    test_x = x_B_T_H_W_D.clone()

    if self.cnt >= int(self.worldcache_num_steps * self.worldcache_ret_ratio):
        probe_depth = self.worldcache_probe_depth
        for blk in self.blocks[:probe_depth]:
            test_x = blk(test_x, **block_kwargs)

        if self.previous_input[current_idx] is not None and self.previous_internal_states[current_idx] is not None:
            delta_x = (x_B_T_H_W_D - self.previous_input[current_idx]).abs().mean() / (
                self.previous_input[current_idx].abs().mean() + 1e-8
            )
            delta_y = (test_x - self.previous_internal_states[current_idx]).abs().mean() / (
                self.previous_internal_states[current_idx].abs().mean() + 1e-8
            )
            self.accumulated_rel_l1_distance[current_idx] += delta_y

            if getattr(self, "worldcache_saliency_enabled", False):
                diff_map = rearrange(
                    (test_x - self.previous_internal_states[current_idx]).abs(),
                    "b t h w d -> b (t d) h w",
                ).mean(dim=1, keepdim=True)
                saliency_map = compute_saliency_map(self.previous_internal_states[current_idx])
                beta = getattr(self, "worldcache_saliency_weight", 5.0)
                weighted_drift = (diff_map * (1.0 + beta * saliency_map)).mean()
                denom = self.previous_internal_states[current_idx].abs().mean() + 1e-6
                weighted_rel = weighted_drift / denom
                self.accumulated_rel_l1_distance[current_idx] += weighted_rel - delta_y
                delta_y = weighted_rel

            if getattr(self, "worldcache_aduc_enabled", False) and not is_parallel_cfg:
                actual_step = self.cnt // 2
                step_ratio = actual_step / max(getattr(self, "worldcache_num_steps", 35), 1)
                if current_idx == 1 and step_ratio > getattr(self, "worldcache_aduc_start", 0.5):
                    if len(self.previous_output) > 1 and self.previous_output[1] is not None:
                        self.cnt += 1
                        return self.previous_output[1]

            input_velocity = delta_x
            alpha = getattr(self, "worldcache_motion_sensitivity", 5.0)
            dynamic_thresh = self.worldcache_rel_l1_thresh / (1.0 + alpha * input_velocity)

            if getattr(self, "worldcache_dynamic_decay", False):
                num_steps = max(getattr(self, "worldcache_num_steps", 35), 1)
                u_ratio = num_steps / 35.0
                base_mult = (u_ratio**2) / 6.0 + u_ratio / 2.0 + 10.0 / 3.0
                dynamic_thresh *= 1.0 + base_mult * (self.cnt / num_steps)

            hf_ok = True
            if getattr(self, "worldcache_hf_enabled", False):
                current_input = rearrange(x_B_T_H_W_D, "b t h w d -> b (t d) h w")
                previous_input = rearrange(self.previous_input[current_idx], "b t h w d -> b (t d) h w")
                hf_drift = compute_hf_drift(previous_input, current_input)
                if hf_drift > getattr(self, "worldcache_hf_thresh", 0.01):
                    hf_ok = False

            if self.accumulated_rel_l1_distance[current_idx] < dynamic_thresh and hf_ok:
                if is_parallel_cfg and batch > 1:
                    diff_full = (test_x - self.previous_internal_states[current_idx]).abs()
                    half_batch = batch // 2
                    drift_cond = diff_full[:half_batch].mean() / (
                        self.previous_internal_states[current_idx][:half_batch].abs().mean() + 1e-6
                    )
                    drift_uncond = diff_full[half_batch:].mean() / (
                        self.previous_internal_states[current_idx][half_batch:].abs().mean() + 1e-6
                    )
                    if drift_cond < dynamic_thresh and drift_uncond < dynamic_thresh:
                        skip_forward = True
                        self.worldcache_step_skipped_count += 1
                        self.resume_flag[current_idx] = False
                        residual_x = self.residual_cache[current_idx]
                    else:
                        self.resume_flag[current_idx] = True
                        self.accumulated_rel_l1_distance[current_idx] = 0
                        self.previous_internal_states[current_idx] = test_x.clone()
                else:
                    skip_forward = True
                    self.worldcache_step_skipped_count += 1
                    self.resume_flag[current_idx] = False
                    residual_x = self.residual_cache[current_idx]
            else:
                self.resume_flag[current_idx] = True
                self.accumulated_rel_l1_distance[current_idx] = 0
                self.previous_internal_states[current_idx] = test_x.clone()

    if skip_forward:
        if len(self.residual_window[current_idx]) >= 2:
            current_residual = test_x - x_B_T_H_W_D
            if getattr(self, "worldcache_osi_enabled", False):
                dt = current_residual - self.probe_residual_window[current_idx][-2]
                ds = (
                    self.probe_residual_window[current_idx][-1]
                    - self.probe_residual_window[current_idx][-2]
                )
                gamma = compute_optimal_gamma(dt, ds)
            else:
                numer = (current_residual - self.probe_residual_window[current_idx][-2]).abs().mean()
                denom = (
                    self.probe_residual_window[current_idx][-1]
                    - self.probe_residual_window[current_idx][-2]
                ).abs().mean()
                gamma = (numer / denom).clip(1, 2) if denom > 1e-6 else 1.0
            x_B_T_H_W_D = x_B_T_H_W_D + self.residual_window[current_idx][-2] + gamma * (
                self.residual_window[current_idx][-1] - self.residual_window[current_idx][-2]
            )
        else:
            x_B_T_H_W_D = x_B_T_H_W_D + residual_x

        self.previous_internal_states[current_idx] = test_x
        self.previous_input[current_idx] = ori_x

        if (
            getattr(self, "worldcache_flow_enabled", False)
            and self.previous_input[current_idx] is not None
            and residual_x is not None
        ):
            current_input = rearrange(x_B_T_H_W_D, "b t h w d -> b (t d) h w")
            previous_input = rearrange(self.previous_input[current_idx], "b t h w d -> b (t d) h w")
            flow = estimate_optical_flow(
                previous_input,
                current_input,
                scale_factor=getattr(self, "worldcache_flow_scale", 0.5),
            )
            if flow is not None:
                residual_input = rearrange(residual_x, "b t h w d -> b (t d) h w")
                warped_residual = warp_feature(residual_input, flow)
                warped = rearrange(warped_residual, "b (t d) h w -> b t h w d", t=time)
                x_B_T_H_W_D = x_B_T_H_W_D - residual_x + warped
    else:
        if self.resume_flag[current_idx]:
            x_B_T_H_W_D = test_x
            remaining = self.blocks[self.worldcache_probe_depth :]
        else:
            remaining = self.blocks

        for i, blk in enumerate(remaining):
            x_B_T_H_W_D = blk(x_B_T_H_W_D, **block_kwargs)
            real_idx = i if not self.resume_flag[current_idx] else i + self.worldcache_probe_depth
            if real_idx == self.worldcache_probe_depth - 1:
                self.previous_internal_states[current_idx] = x_B_T_H_W_D.clone()

        residual_x = x_B_T_H_W_D - ori_x
        self.residual_cache[current_idx] = residual_x
        probe_residual = (
            self.previous_internal_states[current_idx] - ori_x
            if self.previous_internal_states[current_idx] is not None
            else residual_x
        )
        self.probe_residual_cache[current_idx] = probe_residual
        self.previous_input[current_idx] = ori_x
        self.previous_output[current_idx] = x_B_T_H_W_D

        if len(self.residual_window[current_idx]) <= 2:
            self.residual_window[current_idx].append(residual_x)
            self.probe_residual_window[current_idx].append(probe_residual)
        else:
            self.residual_window[current_idx][-2] = self.residual_window[current_idx][-1]
            self.residual_window[current_idx][-1] = residual_x
            self.probe_residual_window[current_idx][-2] = self.probe_residual_window[current_idx][-1]
            self.probe_residual_window[current_idx][-1] = probe_residual

    x_out = self.final_layer(
        x_B_T_H_W_D,
        t_embedding_B_T_D,
        adaln_lora_B_T_3D=adaln_lora_B_T_3D,
    )
    x_B_C_Tt_Hp_Wp = self.unpatchify(x_out)

    self.cnt += 1
    if self.cnt >= self.worldcache_num_steps:
        rate = self.worldcache_step_skipped_count / max(self.worldcache_num_steps, 1)
        log.info(
            f"[WorldCache] Skipped {self.worldcache_step_skipped_count}/{self.worldcache_num_steps} "
            f"({rate:.1%}) | alpha={getattr(self, 'worldcache_motion_sensitivity', 'N/A')}"
        )
        initialize_worldcache_state(self, num_steps=self.worldcache_num_steps)

    if current_idx < len(self.previous_output):
        self.previous_output[current_idx] = x_B_C_Tt_Hp_Wp.clone()

    return x_B_C_Tt_Hp_Wp


def apply_worldcache(model, config: WorldCacheConfig):
    model.worldcache_enabled = True
    model.worldcache_num_steps = config.num_steps
    model.worldcache_rel_l1_thresh = config.rel_l1_thresh
    model.worldcache_ret_ratio = config.ret_ratio
    model.worldcache_probe_depth = config.probe_depth
    model.worldcache_motion_sensitivity = config.motion_sensitivity
    model.worldcache_flow_enabled = config.flow_enabled
    model.worldcache_flow_scale = config.flow_scale
    model.worldcache_hf_enabled = config.hf_enabled
    model.worldcache_hf_thresh = config.hf_thresh
    model.worldcache_saliency_enabled = config.saliency_enabled
    model.worldcache_saliency_weight = config.saliency_weight
    model.worldcache_osi_enabled = config.osi_enabled
    model.worldcache_dynamic_decay = config.dynamic_decay
    model.worldcache_aduc_enabled = config.aduc_enabled
    model.worldcache_aduc_start = config.aduc_start
    model.worldcache_parallel_cfg = config.parallel_cfg

    initialize_worldcache_state(model, num_steps=config.num_steps)
    prepare_cache_runtime_model(model)
    model.forward = types.MethodType(worldcache_mini_train_dit_forward, model)

    log.info(
        f"[WorldCache] Applied: steps={config.num_steps} thresh={config.rel_l1_thresh} "
        f"ret_ratio={config.ret_ratio} probe_depth={config.probe_depth} "
        f"alpha={config.motion_sensitivity} flow={config.flow_enabled}({config.flow_scale}) "
        f"hf={config.hf_enabled} saliency={config.saliency_enabled} osi={config.osi_enabled} "
        f"decay={config.dynamic_decay} aduc={config.aduc_enabled}({config.aduc_start}) "
        f"parallel_cfg={config.parallel_cfg}"
    )
    return model