| 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: |
| 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: |
| 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 |
|
|
