wanstreamer/stream.py

"""Few-step block-causal streaming: the deployment loop, and the loop the
self-forcing trainer rolls out through.

One class serves both, because that identity *is* the method. Self-forcing only
means anything if the context the student is trained on is the context it will
actually receive -- its own previous output, encoded into the same K/V cache, at
the same noise levels, with the same number of steps.

Per streaming unit (a block of `block_frames` latent frames):

    z = noise
    for t, t_next in schedule:            # e.g. 1.0 -> .75 -> .5 -> .25 -> 0
        x0 = z - t * G(z, t | clean K/V of world + committed events)
        z  = (1 - t_next) * x0 + t_next * eps            (fresh eps; last step
                                                          leaves x0 as-is)
    G(x0, t=0)  -> write this block's clean K/V, commit

The re-run at t=0 is what keeps the cache clean: every committed key/value in
the prefix was produced from a *clean* latent, which is the condition Stage 1
trained under. It costs one extra forward per unit and is worth it.

The world's K/V is pinned; eviction under a bounded window drops the oldest
events only, so memory and per-unit cost are flat in stream length while scene
and subject identity persist (v0.3's "world + event stream").
"""
import time

import torch

from . import blockcausal as bc
from .core import make_rope_table, make_cache, latent_geometry
from .graphrunner import GraphedForward, SteadyStateGraphs


def few_step_schedule(num_steps, t_max=1.0, shift=1.0):
    """Flow fractions, highest first.

    shift == 1 gives the uniform spacing DMD-distilled few-step video models use
    (CausVid / Self-Forcing: 1.0, .75, .5, .25). shift > 1 reproduces Wan's own
    inference schedule, which is the right one for evaluating a many-step,
    not-yet-distilled model.
    """
    ts = [t_max * (num_steps - i) / num_steps for i in range(num_steps)]
    if shift == 1.0:
        return ts
    return [shift * t / (1 + (shift - 1) * t) for t in ts]


class FewStepStreamer:
    def __init__(self, model, width=640, height=368, max_frames=256,
                 device='cuda', dtype=torch.bfloat16, window_frames=None,
                 time_scale=1000.0, cache_frames=None, block_frames=3,
                 num_steps=4, rope=None, shift=1.0, sampler='renoise',
                 cuda_graphs=False, noisy_context=False):
        self.model = model
        self.device = device
        self.dtype = dtype
        self.time_scale = time_scale
        self.block_frames = block_frames
        self.schedule = few_step_schedule(num_steps, shift=shift)
        self.sampler = sampler
        self.h_lat, self.w_lat, self.hp, self.wp = latent_geometry(width, height)
        self.S = self.hp * self.wp
        self.max_frames = max_frames
        self.window_frames = window_frames
        self.rope = rope if rope is not None else make_rope_table(
            model, self.hp, self.wp, max_frames, device)
        self.cache = make_cache(model, self.S, cache_frames or max_frames,
                                device, dtype)
        self.kv = bc.BufferKV(self.cache)
        self.n_world = 0
        self.n_frames = 0
        # Training-only: pushes the event window this far from the pinned world
        # in RoPE index space, simulating a long stream. At deployment the gap
        # arises naturally from evicted events, so this stays 0.
        self.rope_gap = 0
        # per-channel moment matching towards the world (0 = off); see _renorm
        self.latent_norm = 0.0
        self.ref_stats = None
        self.graphs = SteadyStateGraphs() if cuda_graphs else None
        # "Noisy context": commit the K/V the LAST denoising step already wrote
        # instead of re-running the finished block at t=0. That drops the cost
        # of a unit from N+1 forwards to N -- a third of the latency at N=2,
        # which is the single largest inference-cost lever left in this loop.
        # It is NOT a free inference-time swap: every committed key then comes
        # from a latent at the schedule's lowest noise level rather than a clean
        # one, and Stage 1 and DMD both trained against a clean prefix. Turning
        # it on without retraining under the same convention changes the
        # conditioning the student was distilled for. Exposed here so the
        # latency claim can be measured before that training is paid for.
        self.noisy_context = noisy_context

    # ------------------------------------------------------------------ context
    def set_text(self, ctx_emb, ctx_lens=None):
        self.ctx = ctx_emb
        # Resolve "is a cross-attention mask needed?" ONCE, here, and store None
        # if it is not. Deciding it inside the forward costs a `.item()` on a
        # device tensor -- a host sync, which is illegal during CUDA graph
        # capture. Contexts are padded to the full text length (upstream's own
        # WanModel.forward passes context_lens=None for exactly this reason), so
        # None is not an approximation, it is the same computation.
        full = ctx_emb.shape[1]
        if ctx_lens is not None and int(torch.as_tensor(ctx_lens).min()) < full:
            self.ctx_lens = ctx_lens
        else:
            self.ctx_lens = None

    def _eager(self, z, tbl, e, e0):
        return bc.block_forward(self.model, z, 0.0, 0, self.rope, self.ctx,
                                self.ctx_lens, kv=self.kv, dtype=self.dtype,
                                time_scale=self.time_scale, tbl=tbl, emb=(e, e0))

    def _fwd(self, z, t, t_start):
        tbl = self.rope.span(t_start, z.shape[2])
        e, e0 = bc.time_embed(self.model, t, self.device, self.time_scale)
        if self.graphs is not None and not torch.is_grad_enabled():
            # Key on the cache length: that is the only thing that changes the
            # shapes, and a bounded window makes it constant once saturated.
            g = self.graphs.get(
                (self.cache.length, z.shape[2]),
                lambda: GraphedForward(self._eager, tuple(z.shape),
                                       tuple(tbl.shape), tuple(e.shape),
                                       tuple(e0.shape), self.device))
            if g is not None:
                return g(z, tbl, e, e0)
        return self._eager(z, tbl, e, e0)

    def _commit(self, n):
        self.cache.commit(n * self.S)
        self.n_frames += n
        if self.window_frames is not None:
            self.trim_to_window(self.window_frames)

    def trim_to_window(self, window_frames):
        """Evict the oldest events so only `window_frames` of them remain.

        Deployment calls this from `_commit` after every unit, which is what
        makes per-unit cost flat in stream length. The self-forcing trainer
        cannot: it records denoising steps during a rollout and redoes them
        under gradient afterwards against `buffer[:prefix]`, and an eviction in
        between silently changes what those tokens are (see
        StreamingKVCache.evictions). So the trainer leaves `window_frames` None,
        rolls one iteration's blocks with the buffer growing monotonically, and
        calls this itself at the iteration boundary -- the one point where no
        recorded step is still owed its prefix.

        The world stays pinned either way; only events are dropped.
        """
        protect = self.n_world * self.S
        budget = (self.n_world + window_frames) * self.S
        excess = self.cache.num_tokens - budget
        if excess > 0:
            self.cache.evict_front(excess, protect=protect)

    # -------------------------------------------------------------------- world
    @torch.no_grad()
    def set_world(self, world_latents):
        """Prime the cache from clean world latents, attended bidirectionally."""
        self.cache.reset()
        self.n_frames = 0
        z = world_latents.to(self.device).unsqueeze(0) \
            if world_latents.dim() == 4 else world_latents.to(self.device)
        F = z.shape[2]
        self._fwd(z, 0.0, 0)
        self._commit(F)
        self.n_world = F
        # Per-channel moments of the world: the in-distribution reference for
        # `latent_norm`. Computed over (F, H, W) per channel.
        w = z[0].float()
        self.ref_stats = (w.mean(dim=(1, 2, 3), keepdim=True).unsqueeze(0),
                          w.std(dim=(1, 2, 3), keepdim=True).unsqueeze(0))
        return F

    def _renorm(self, z, strength):
        """Pull a finished block's PER-CHANNEL moments back towards the world's.

        Long autoregressive rollouts drift, and here the drift is measurably a
        per-channel one: the decoded stream develops colour casts and saturated
        blotches long before it loses structure. A distilled student inherits
        this from the CFG-guided real score it was matched to -- CFG moves mass
        towards higher-contrast, more saturated samples, and in a feedback loop
        (output becomes context becomes output) that bias compounds.

        Matching the first two moments per channel against the world is the
        cheapest correction that targets exactly that, and it is applied BEFORE
        the block's K/V is recomputed and committed, so the cached history stays
        in distribution and the correction cannot itself accumulate. It cannot
        fix a wrong *direction*, only a wrong per-channel scale and offset.
        """
        if strength <= 0 or getattr(self, 'ref_stats', None) is None:
            return z
        mu, sd = self.ref_stats
        zm = z.mean(dim=(2, 3, 4), keepdim=True)
        zs = z.std(dim=(2, 3, 4), keepdim=True)
        z_n = (z - zm) / (zs + 1e-5) * sd.to(z.dtype) + mu.to(z.dtype)
        return (1.0 - strength) * z + strength * z_n

    # ------------------------------------------------------------------- events
    def generate_block(self, generator=None, record=None, n_frames=None):
        """Emit one streaming unit. Returns clean latents [1, C, b, H, W]."""
        b = n_frames or self.block_frames
        t0 = bc.event_rope_index(self.n_frames, self.n_world, self.rope_gap)
        z = torch.randn(1, 16, b, self.h_lat, self.w_lat, device=self.device,
                        dtype=torch.float32, generator=generator)
        sched = self.schedule
        for i, t in enumerate(sched):
            if record is not None:
                record.append({'block_start': t0, 'step': i, 't': t,
                               'z': z.detach()})
            v = self._fwd(z, t, t0).float()
            x0 = z - t * v
            tn = sched[i + 1] if i + 1 < len(sched) else 0.0
            if tn == 0.0:
                z = x0
            elif self.sampler == 'renoise':
                # DMD-distilled few-step sampler: re-noise the x0 prediction to
                # the next level with FRESH noise. This is the sampler the
                # student is distilled under, so it is also the one it must be
                # rolled out with.
                eps = torch.randn(z.shape, device=z.device, dtype=z.dtype,
                                  generator=generator)
                z = (1 - tn) * x0 + tn * eps
            else:
                z = z + (tn - t) * v          # deterministic Euler, for the
                #                               undistilled many-step model
        z = self._renorm(z, self.latent_norm)
        if not self.noisy_context:
            self._fwd(z, 0.0, t0)      # clean K/V for the committed block
        # else: the scratch region still holds the K/V written by the final
        # denoising forward, so committing promotes that instead and the extra
        # pass is skipped entirely. (`_renorm` is then not reflected in the
        # cache, which is another reason the two options are not interchangeable
        # without retraining.)
        self._commit(b)
        return z

    @torch.no_grad()
    def stream(self, num_blocks, seed=0, log_every=0, n_frames=None):
        g = torch.Generator(device=self.device).manual_seed(seed)
        lats, times = [], []
        for u in range(num_blocks):
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            z = self.generate_block(generator=g, n_frames=n_frames)
            torch.cuda.synchronize()
            dt = time.perf_counter() - t0
            lats.append(z[0])
            times.append(dt)
            if log_every and (u + 1) % log_every == 0:
                b = n_frames or self.block_frames
                print(f'  unit {u+1}/{num_blocks}: {dt*1000:7.1f} ms '
                      f'({dt/b*1000:6.1f} ms/latent frame, '
                      f'cache {self.cache.num_tokens} tok)', flush=True)
        return torch.cat(lats, dim=1), times

    # --------------------------------------------------- self-forcing rollout
    @torch.no_grad()
    def rollout_record(self, num_blocks, generator=None):
        """Roll out `num_blocks` units and keep what the DMD step needs.

        Returns (x0_blocks, records). `records[i]` holds, for every denoising
        step of block i, the exact input latent and timestep, plus the cache
        length at that point -- enough to recompute any single step under
        gradient without re-running the rollout.
        """
        out, recs = [], []
        for _ in range(num_blocks):
            rec = []
            prefix = self.cache.num_tokens
            ev = self.cache.evictions
            z = self.generate_block(generator=generator, record=rec)
            for r in rec:
                r['prefix'] = prefix
                r['evictions'] = ev
            out.append(z)
            recs.append(rec)
        return out, recs

    def prefix_intact(self, rec_step):
        """Is this recorded step's prefix still literally `buffer[:prefix]`?

        False once anything has been evicted since it was recorded -- the tokens
        at those indices are now different tokens. The trainer must not take a
        gradient through a step for which this is False; the forward would
        succeed and quietly train against the wrong context.
        """
        return rec_step.get('evictions', 0) == self.cache.evictions

    def recompute_step(self, rec_step, prefix_kv):
        """Differentiably redo one recorded denoising step against a detached
        prefix. Returns that step's x0 prediction."""
        z, t = rec_step['z'], rec_step['t']
        v = bc.block_forward(self.model, z, t, rec_step['block_start'],
                             self.rope, self.ctx, self.ctx_lens, kv=prefix_kv,
                             dtype=self.dtype, time_scale=self.time_scale,
                             prefix_upto=rec_step['prefix'],
                             grad_checkpoint=True)
        return z - t * v.float()