Four fixes to the reference implementation of
StreamEdit (arXiv:2605.21466), applied
identically to both backends (Self-Forcing_StreamEdit/, LongLive_StreamEdit/).
Apply against a clean checkout (0002 requires 0001 first):
git -C /path/to/StreamEdit apply /path/to/0001-streamedit-code-fixes.patch
git -C /path/to/StreamEdit apply /path/to/0002-streamedit-multishot-adaptive.patchPatch 0002 adds streaming-causal extensions on top of the fixes; see the 0002 section below.
A literature survey of mechanism-level improvements (localization, SOG guidance,
noise reuse, KV drift) with a prioritized implementation plan lives in
docs/literature-survey-mechanism-improvements.md.
v.to(device) # not in-place; return value discarded.to() returns a new tensor, and the dict entry was never rebound, so the target
KV-cache offload promised by the low_memory path never happened. low_memory
is set automatically when free VRAM < 40 GB, so those users paid the full memory
cost while the code appeared to offload.
Fixed by rebinding (itm[key] = itm[key].to(device)) and moving only k/v —
the bookkeeping tensors stay put so the surrounding .item() calls are unaffected.
Fixing the function alone would crash: _concat_kv_cache consumes kv_cache_trg
at the top of each chunk, while the previous chunk had left it on CPU, and
torch.cat across devices raises. The offload window was therefore added around
the concat, so the target cache is on CPU for the whole spatial denoising loop —
the memory peak, and exactly when the dual cache already holds a copy of it.
torch.cat already allocates, so the trailing .clone() duplicated the dual
cache (30 layers x k/v) for nothing. Removed. Numerically identical.
Never called; the k_src_clean / v_src_clean keys it writes are read nowhere in
wan/. Removed.
_reuse_noise_statistics accepted fg_mask and never used it; the caller computed
it anyway (an F.interpolate per chunk) and threw it away. noise_temporal_mean_fg
and noise_temporal_mean_bg were allocated and never read, as were ema_factor
and alpha_mixed.
The effect: chunk-to-chunk noise correlation (alpha_prog=2, i.e. 2/sqrt(5) ~= 0.894)
was applied uniformly, including inside editing regions — against the paper's own
Sec. 4.4, which argues "randomness in editing-related regions remains beneficial
for effective edits".
Implemented with a single noise state and two mixing coefficients applied
spatially: background keeps alpha_prog, editing regions use noise_alpha_fg.
--noise_alpha_fg 0 # editing regions fully decorrelated across chunksDefault None reproduces the original behaviour exactly. Dead parameters and
dicts removed.
Two places where the code diverges from the paper were left alone, because the code is right and the paper's equation is wrong (the paper's own prose agrees with the code in both cases):
- Eq. (6) writes
M_curr (*) K_src, but the section is titled "Background Preservation" and the code injects the background (~fg_mask). - Eq. (10) and Alg. A.2 L30 write
AMN(...) (*) g_t, but the text says correction is applied to editing-irrelevant regions and the code uses(1 - AMN).
Also left alone: the cross-chunk noise correlation itself is undocumented — Eq. (3)
and Alg. A.1 L14 both specify i.i.d. eps ~ N(0, I). It looks deliberate rather
than accidental, and turning it off would weaken the published baseline, so it is
exposed as a knob rather than removed. Worth stating explicitly in any paper that
uses StreamEdit as a baseline.
Hyper-parameters are untouched and still match Sec. 5.1: fg_boost_factor
(omega) 4.0 for Self Forcing and 2.0 for LongLive, blend_power (rho) 2.0 for both.
- Every edit applied as an anchored string replacement with an exact-count assertion: 7 per pipeline file, all matched.
python -m py_compilepasses on all four modified files.- The two backends' diffs were compared line by line and are equivalent apart from pre-existing signature differences.
- The patch was checked against a pristine tree with
git apply --check.
Nothing was executed. No PyTorch or GPU was available in the environment where these edits were made.
Scope of runtime impact, to make testing cheap:
- On >= 40 GB VRAM,
low_memoryis False and_kv_cache_toreturns immediately. - Without
--noise_alpha_fg, the noise path is bit-identical to before. - Fixes 2 and 3 cannot change numerics.
So on a >= 40 GB GPU with no new flag, output should be bit-identical to the
unpatched code — one run of examples/A_car.mp4 confirms it. The only path with
genuinely new runtime behaviour is low_memory (< 40 GB), which previously did
nothing and now actually moves tensors; sanity-check that separately.
0002-streamedit-multishot-adaptive.patch (applies on top of 0001) adds four
opt-in mechanisms. Every default is off; with no new flag the code paths are
identical to 0001. All mechanisms are streaming-causal: decisions use only
already-arrived source frames and already-emitted output; nothing is retracted.
The reused temporal noise is flipped along the frame axis every chunk
(palindromic pattern). Undocumented in the paper; this flag disables it so the
effect can be measured (see eval/run_ablation.sh).
Per chunk, a latent-domain signal — mean |z0_tgt − x0_src| inside the union
editing mask — is compared to chunk 0. When retention falls below
--adaptive_target (default 0.85), a clamped proportional controller raises
fg_boost_factor (omega) and blend_power (rho) for the next chunk; above
target they relax back. Rationale: emitted frames cannot be fixed afterwards,
so the only defence against edit decay is feedback before the next chunk.
The per-chunk retention trace is printed and exposed as
pipeline.last_edit_retention.
A causal pixel-domain hard-cut detector (frame diff vs rolling median of past
diffs, zero lookahead; dissolves intentionally do not trigger) splits the
source into shots. Each shot is padded to the 12m−3 grid by replicating its
last frame (causal), VAE-encoded separately (frames on the two sides of a
cut never share a latent), edited as an independent inference() call — which
by construction resets all five pieces of cross-chunk state (src KV, tgt KV,
M_prev mask cache, reused-noise statistics, sink) — and trimmed back.
--shot_prompts shots.json optionally supplies per-shot prompts/triggers.
The best-edited chunk of each shot (by the same latent signal) is kept in a small bank with a pooled-latent retrieval key. Two uses:
- Cross-shot (
--use_identity_anchor): at each new shot, the most relevant entry (cosine over pooled latents, MemFlow-style relevance-over-recency, plus an edit-score prior) is injected through the existing visual-prompting prefill path, so the edited subject stays consistent across shots. Honest limit: streaming can only be consistent with history — the first shot in which the subject appears has no anchor. - Mid-shot (
--anchor_refresh_threshold 0.7): when retention drops below the threshold, the attention sink is re-anchored with the best chunk — the anchor latents are re-forwarded atcurrent_start=0into a scratch cache (RoPE therefore matches the sink positions natively, cf. Rolling Forcing), and only foreground tokens overwrite the sink.
Same regime as 0001: python -m py_compile on all four files, patch verified
by round-trip against a pristine 0001 tree, eval self-tests pass. Nothing was
executed on a GPU. First runs to make: (a) no-flag run must be bit-identical
to 0001; (b) eval/run_ablation.sh for the decay baseline and noise ablations;
(c) a two-shot concatenated clip through --multi_shot with and without
--use_identity_anchor, evaluated with eval/edit_retention.py --auto_shots.