Skip to content

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Repository files navigation

StreamEdit code fixes

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.patch

Patch 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.

What changed

1. _kv_cache_to was a silent no-op

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.

2. Redundant .clone() in _concat_kv_cache

torch.cat already allocates, so the trailing .clone() duplicated the dual cache (30 layers x k/v) for nothing. Removed. Numerically identical.

3. Dead code: _append_clean_src_kv_cache

Never called; the k_src_clean / v_src_clean keys it writes are read nowhere in wan/. Removed.

4. Foreground-aware temporal noise (was half-implemented)

_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 chunks

Default None reproduces the original behaviour exactly. Dead parameters and dicts removed.

Not changed

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.

Verification status

  • Every edit applied as an anchored string replacement with an exact-count assertion: 7 per pipeline file, all matched.
  • python -m py_compile passes 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_memory is False and _kv_cache_to returns 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.


Patch 0002: closed-loop control, multi-shot, identity anchor

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.

1. Noise-flip ablation (--no_noise_flip, both backends)

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).

2. Closed-loop edit-strength control (--adaptive_edit_control, both backends)

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.

3. Multi-shot streaming editing (--multi_shot, LongLive driver)

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.

4. Identity anchor (--use_identity_anchor, --anchor_refresh_threshold, LongLive)

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 at current_start=0 into a scratch cache (RoPE therefore matches the sink positions natively, cf. Rolling Forcing), and only foreground tokens overwrite the sink.

Verification status (0002)

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.

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages