lbcache/patch.py

"""Route ``WanModelFast.forward`` through a cache method.

While the controller is marked active (denoising forwards) the preamble (patch
embedding, time / text / camera embeddings) and the tail (``head``, ``unpatchify``)
are reproduced verbatim from ``wan/modules/model_fast.py`` and only the 30-block
loop is handed to the controller; an output-level method (velocity ``reuse``) is
instead handed the whole stock forward as a callable.  The per-chunk context pass
runs the untouched original forward.
"""

import types

import torch
import torch.nn.functional as torch_F
from einops import rearrange

from .methods import StepCtx


def parse_schedule(schedule, num_steps):
    # 'R' spells out "reuse" in the naive-cache baselines (FRFF / FRRF / FRRR);
    # it is the same thing as 'x': a step the method may serve from cache.
    s = schedule.strip().upper().replace("?", "X").replace("R", "X")
    if len(s) != num_steps or set(s) - {"F", "X"}:
        raise ValueError(f"schedule {schedule!r} must be {num_steps} chars of F/x")
    if s[0] != "F":
        raise ValueError(f"schedule {schedule!r}: step 0 must be F")
    return tuple(i for i, c in enumerate(s) if c == "F")


def schedule_string(forced, num_steps):
    return "".join("F" if i in forced else "x" for i in range(num_steps))


class CacheController:
    def __init__(self, method, num_steps=4, forced_steps=(0, -1), first_chunk_forced_steps=None):
        self.method = method
        self.num_steps = num_steps
        self.forced = {s % num_steps for s in forced_steps}
        self.schedule = schedule_string(self.forced, num_steps)
        cacheable = [s for s in range(num_steps) if s not in self.forced]
        self.last_cacheable_step = max(cacheable) if cacheable else -1
        self.first_chunk_forced = (None if first_chunk_forced_steps is None
                                   else set(first_chunk_forced_steps))
        self.first_chunk_schedule = None
        if self.first_chunk_forced is not None:
            n0 = max(num_steps, max(self.first_chunk_forced, default=-1) + 1)
            self.first_chunk_schedule = schedule_string(self.first_chunk_forced, n0)
            c0 = [s for s in range(n0) if s not in self.first_chunk_forced]
            self.first_chunk_last_cacheable = max(c0) if c0 else -1
        self.active = False
        self.block_idx = -1
        self.step_idx = -1
        self.records = []

    def reset_video(self):
        self.method.reset_video()
        self.records = []

    def begin_chunk(self, block_idx):
        self.block_idx = block_idx
        self.method.begin_chunk(block_idx)

    def denoise_step(self, step_idx):
        self.step_idx = step_idx
        self.active = True

    def end_step(self):
        self.active = False

    def forced_now(self):
        if self.first_chunk_forced is not None and self.block_idx == 0:
            return self.first_chunk_forced
        return self.forced

    def run(self, ctx_kwargs):
        first = self.first_chunk_forced is not None and self.block_idx == 0
        ctx = StepCtx(block_idx=self.block_idx, step_idx=self.step_idx,
                      forced_full=self.step_idx in self.forced_now(),
                      last_cacheable_step=(self.first_chunk_last_cacheable if first
                                           else self.last_cacheable_step), **ctx_kwargs)
        out, frac = self.method.forward(ctx)
        self.records.append({"block": self.block_idx, "step": self.step_idx,
                             "compute_fraction": float(frac)})
        return out

    def summary(self):
        d = self.records
        if not d:
            return {}
        compute = sum(r["compute_fraction"] for r in d)
        middle = [r for r in d if r["step"] not in (self.first_chunk_forced if (
            self.first_chunk_forced is not None and r["block"] == 0) else self.forced)]
        active = [r["compute_fraction"] for r in middle
                  if 1e-9 < r["compute_fraction"] < 1 - 1e-9]
        n_mid = len(middle) or 1
        return {"denoise_forwards": len(d), "compute_equivalent_forwards": compute,
                "middle_steps": len(middle),
                "middle_compute_equivalent": sum(r["compute_fraction"] for r in middle),
                "active_step_ratio": len(active) / n_mid,
                "empty_step_ratio": sum(r["compute_fraction"] <= 1e-9 for r in middle) / n_mid,
                "full_step_ratio": sum(r["compute_fraction"] >= 1 - 1e-9 for r in middle) / n_mid,
                "mean_selected_fraction_active": (sum(active) / len(active)) if active else 0.0,
                "flops_speedup_estimate": len(d) / compute if compute else float("inf")}


def _cached_forward(self, x, t, context, seq_len, y=None, dit_cond_dict=None,
                    kv_cache=None, crossattn_cache=None, current_start=0,
                    max_attention_size=1_000_000, frame_seqlen=None,
                    cross_attn_first_call=None):
    def run_full():
        return self._orig_forward(
            x, t, context, seq_len, y=y, dit_cond_dict=dit_cond_dict, kv_cache=kv_cache,
            crossattn_cache=crossattn_cache, current_start=current_start,
            max_attention_size=max_attention_size, frame_seqlen=frame_seqlen,
            cross_attn_first_call=cross_attn_first_call)

    ctrl = getattr(self, "_cache_ctrl", None)
    if ctrl is None or not ctrl.active:
        return run_full()
    if getattr(ctrl.method, "level", "blocks") == "output":
        return ctrl.run(dict(model=self, run_full=run_full, x=None, e0=None, kwargs=None,
                             kv_cache=kv_cache, crossattn_cache=crossattn_cache,
                             current_start=current_start, grid_sizes=None))

    from wan.modules.model import sinusoidal_embedding_1d

    # -- preamble, verbatim from WanModelFast.forward ---------------------------------
    if self.model_type == 'i2v':
        assert y is not None
    device = self.patch_embedding.weight.device
    if self.freqs.device != device:
        self.freqs = self.freqs.to(device)
    if y is not None:
        x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
    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)
    if t.dim() == 1:
        t = t.expand(t.size(0), seq_lens)
    with torch.amp.autocast('cuda', dtype=torch.float32):
        bt = t.size(0)
        t = t.flatten()
        e = self.time_embedding(
            sinusoidal_embedding_1d(self.freq_dim,
                                    t).unflatten(0, (bt, seq_lens)).float())
        e0 = self.time_projection(e).unflatten(2, (6, self.dim))
        assert e.dtype == torch.float32 and e0.dtype == torch.float32
    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
        ]))
    if dit_cond_dict is not None and "c2ws_plucker_emb" in dit_cond_dict:
        c2ws_plucker_emb = dit_cond_dict["c2ws_plucker_emb"]
        c2ws_plucker_emb = [
            rearrange(
                i,
                '1 c (f c1) (h c2) (w c3) -> 1 (f h w) (c c1 c2 c3)',
                c1=self.patch_size[0],
                c2=self.patch_size[1],
                c3=self.patch_size[2],
            ) for i in c2ws_plucker_emb
        ]
        c2ws_plucker_emb = torch.cat(c2ws_plucker_emb, dim=1)
        c2ws_plucker_emb = self.patch_embedding_wancamctrl(c2ws_plucker_emb)
        c2ws_hidden_states = self.c2ws_hidden_states_layer2(
            torch_F.silu(self.c2ws_hidden_states_layer1(c2ws_plucker_emb)))
        dit_cond_dict = dict(dit_cond_dict)
        dit_cond_dict["c2ws_plucker_emb"] = (
            c2ws_plucker_emb + c2ws_hidden_states)
    kwargs = dict(
        e=e0,
        seq_lens=seq_lens,
        grid_sizes=grid_sizes,
        freqs=self.freqs,
        context=context,
        context_lens=context_lens,
        dit_cond_dict=dit_cond_dict,
        max_attention_size=max_attention_size,
        frame_seqlen=frame_seqlen,
        cross_attn_first_call=cross_attn_first_call)

    x = ctrl.run(dict(model=self, run_full=run_full, x=x, e0=e0, kwargs=kwargs,
                      kv_cache=kv_cache, crossattn_cache=crossattn_cache,
                      current_start=current_start, grid_sizes=grid_sizes))

    x = self.head(x, e)
    x = self.unpatchify(x, grid_sizes)
    return [u.float() for u in x]


def install(model, controller):
    if not hasattr(model, "_orig_forward"):
        model._orig_forward = model.forward
        model.forward = types.MethodType(_cached_forward, model)
    model._cache_ctrl = controller
    return model