Add LingBot-VA Robotwin integration for V2

[jmccaffrey-nv]

Summary

• add a self-contained LingBot-VA RoboTwin model package and V2 application adapter
• port the pinned upstream Robotwin I2AV generator without depending on draft PR Added support for Lingbot-VA #312
• return complete TCHW video, numeric metrics, and denormalized robot actions through the typed tensor-artifact API from Add typed tensor artifact outputs to V2 #524
• preserve the reviewed CFG/KV behavior, strict checkpoint mapping, one-run destructive lifecycle, deferred decode, and resident/offload parity
• incorporate design-review feedback with an explicit offline scope, separate session module, action semantics, and a model-specific action plot/CSV inspector

Stack

This PR is stacked on #524 (dev/jmccaffrey/v2-tensor-artifacts). It is rebased onto its current head, which contains current origin/main. Merge #524 first, then retarget this PR to main.

Relative to #524, this PR has no implementation changes under flashdreams/flashdreams/runtime_v2/.

Integrated scope

This adapter wraps the pinned upstream WanVAServer.generate() path:

• input: one natural-language instruction plus high, left-wrist, and right-wrist RGB PNGs
• output: one complete high-camera video rollout, float32 actions[step, channel], and metrics
• lifecycle: the model generates all chunks internally, tears down denoising state, then decodes and returns one StepResult
• interaction: UserInputEvents are not consumed; no live observation stream, simulator bridge, asynchronous motor execution, or robot actuation is implemented

This is a limitation of the FlashDreams adapter, not of LingBot-VA as a whole. Pinned upstream also has a closed-loop RoboTwin evaluator/server that executes action chunks, captures actual observations, and feeds observations plus executed state back into the KV cache. That path is documented but intentionally not ported here.

The adapter does not use the shared text-to-video session: its inputs include three camera observations, its outputs include robot actions, and its engine does not expose the shared per-block generate/finalize lifecycle. It still reuses V2 interfaces and generic MP4, metrics, and tensor-artifact sinks.

Action artifact decision

Discarding actions would remove the policy half of this video-action model. The adapter therefore declares one typed actions[step, channel] artifact and emits all 16 selected RoboTwin channels:

• left translation delta x/y/z, relative quaternion x/y/z/w, gripper
• right translation delta x/y/z, relative quaternion x/y/z/w, gripper

Upstream composes the relative poses with the initial end-effector poses and normalizes the quaternions before simulator execution. FlashDreams emits the denormalized relative values and does not actuate them.

The generic V2 result already supports multiple independently named tensors via tuple[TensorArtifactOutput, ...]; no list API change is needed. LingBot-VA remains the first model integration exercising the facility, so validation with a second model remains an explicit follow-up rather than a claim in this PR.

A model-specific lingbot-va-visualize-actions command now validates committed artifact manifests and produces a four-panel two-arm trajectory plot plus optional named CSV. It stays outside the generic runtime and is diagnostic only.

Architecture revisions

• split configuration/application ownership into lingbot_va_v2/app.py
• move session, finite model loop, lazy engine ownership, and result validation into lingbot_va_v2/session.py
• update the component, class, use-case, and rollout-sequence diagrams
• document the upstream closed-loop boundary, action consumption semantics, shared-T2V decision, and second-model generalization gate
• update the public model card in docs/source/models/lingbot_va.rst

The earlier review fixes remain intact: steady-state KV rollover, teardown ordering, inactive CFG work elimination, private KV access removal, latent denormalization, prompt validation, cleanup reporting, and stale defaults/constants.

Validation

• combined V2 runtime + LingBot CPU suite: 325 passed, 2 GPU tests deselected
• focused post-edit action-inspector suite: 4 passed
• configured changed-file hooks: passed Ruff fix/format, lock validation, version sync, and ty
• full Sphinx HTML build with -W --keep-going: passed
• uv lock: resolved successfully
• installed lingbot-va-visualize-actions entry point: exercised against legacy and committed-manifest outputs
• fresh real GPU run: 1 passed in 35.82 s on RTX PRO 6000 Blackwell, two chunks, default CFG, component offload, no compile
  • model-reported total: 31.516 s
  • MP4: H.264, 320x256, 10 FPS, 13 frames
  • MP4 SHA-256: 15bcdc4307e080218255e83946c2c2e5dbc30f3b7acd26c8925167017234e586
  • actions: float32 [64, 16], finite, distinct second chunk
  • action SHA-256: 463b307b667c1ca13a47bbbc5a17f68604621dfe3c3a10fc5860077216928d95, byte-identical to prior resident/offload and pre-review runs
  • peak allocation: 39,804,413,440 bytes (37.07 GiB)
  • the inspector consumed the fresh complete manifest and generated a valid 2100x1200 PNG and named CSV

The model README retains exact input hashes, checkpoint mapping counts, eager/compiled upstream parity results, reproduction commands, memory phases, and safety limitations.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[jmccaffrey-nv]

/ready to test faed607

[jmccaffrey-nv]

/ok to test faed607

[copy-pr-bot]

/ok to test faed607

@jmccaffrey-nv, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[jmccaffrey-nv]

/ok to test d9cdd92

[greptile-apps]

Greptile Summary

The PR adds a self-contained LingBot-VA RoboTwin model package and a V2 application adapter for one-shot joint video/action rollouts.

• Implements checkpoint loading, observation encoding, joint denoising, rolling branch-specific KV caches, deferred video decoding, and action denormalization.
• Registers a V2 session that publishes video, metrics, and typed action tensor artifacts through generic runtime outputs.
• Adds model documentation, artifact inspection tooling, package metadata, dependency locking, and CPU/GPU regression coverage.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains in the eligible follow-up review scope.

No blocking failure remains.

Important Files Changed

Filename	Overview
integrations/lingbot_va/lingbot_va/engine.py	Orchestrates checkpoint-backed prompt and observation encoding, autoregressive generation, destructive teardown, deferred decode, and output aggregation.
integrations/lingbot_va/lingbot_va/transformer/impl/kvcache.py	Implements rolling video/action attention-cache storage and committed-prefix handling for conditional and unconditional branches.
integrations/lingbot_va/lingbot_va/transformer/impl/network.py	Adds the joint video/action transformer network and branch-specific attention flow.
integrations_v2/lingbot_va/lingbot_va_v2/app.py	Defines CLI configuration, fixed output capabilities, input validation, application discovery, and session creation.
integrations_v2/lingbot_va/lingbot_va_v2/session.py	Adapts the one-run engine lifecycle into the V2 session contract and emits video, metrics, and typed action artifacts.
integrations/lingbot_va/lingbot_va/action_artifact.py	Adds validation, plotting, and CSV export for committed action tensor artifacts.
integrations/lingbot_va/pyproject.toml	Defines the standalone model package, dependency ranges, optional visualization support, and action-inspector entry point.
integrations_v2/lingbot_va/pyproject.toml	Packages and registers the LingBot-VA V2 application integration.

Sequence Diagram

sequenceDiagram
  actor User
  participant Runtime as V2 Runtime
  participant App as LingBot-VA Application
  participant Session as Session / Model Loop
  participant Engine as LingBot-VA Engine
  participant Model as Video-Action Model
  participant Sinks as Generic Output Sinks
  User->>Runtime: Launch with prompt, images, and model options
  Runtime->>App: Initialize and validate session contract
  App->>Session: Create session
  Runtime->>Session: step()
  Session->>Engine: Lazily construct and run
  Engine->>Model: Encode prompt and observations
  loop Autoregressive chunks
    Engine->>Model: Denoise video and actions
    Model->>Model: Advance conditional and unconditional KV caches
  end
  Engine->>Model: Release denoising state and decode video
  Model-->>Engine: Video, actions, and metrics
  Engine-->>Session: Complete rollout
  Session-->>Runtime: StepResult with tensor artifact
  Runtime->>Sinks: Write MP4, metrics, and actions

Loading

_{Reviews (7): Last reviewed commit: "Clarify LingBot offline action contract" | Re-trigger Greptile}

[jmccaffrey-nv]

Code review record — #520 (+ dependency #524)

Reviewed at 94bee383 (#520) and 213fecdb (#524), both matching the current PR heads. Read the full diff of both branches plus the flashdreams components they build on (runtime_v2, core/attention, core/checkpoint, core/io, infra/diffusion/scheduler, infra/encoder, recipes/wan, t2v_v2).

Reproduced the stated CPU validation: pytest flashdreams/test_v2 integrations/lingbot_va integrations_v2/lingbot_va -m ci_cpu -q → 272 passed, 2 deselected. Matches the PR body.

Two findings were confirmed by execution rather than by reading; both are marked below.

---

#524 — typed tensor artifacts

Well-scoped and model-neutral as claimed. The generation semantics (latest-wins, buffered older output cleared, late stale batches ignored), the per-file NamedTemporaryFile → fsync → os.replace write, the pre-acceptance schema/rank/dtype validation, and the _PerResultModelOutputSink shim that keeps the original per-result contract intact are all right. The api_v2 and integrations_v2 doc additions are genuinely good — the "one name at most once across the whole result batch" rule is the kind of thing that is invisible until someone violates it.

1. A failed run still writes its artifacts (medium). run_session's finally closes every attempted sink regardless of outcome, so TensorArtifactOutputSink.close() persists whatever chunks it buffered before the failure. A truncated actions.npy from a run that died at chunk 6 of 10 is byte-for-byte indistinguishable from a complete one — same name, same dtype, plausible shape. The CLI exits non-zero, but anything scanning the output directory afterwards can't tell. MetricsOutputSink has the same shape, but a partial benchmark JSON is self-evidently partial; a tensor isn't. Worth either skipping the write when the run failed, or recording completeness somewhere the consumer sees.

2. --tensor-artifact-dir is registered in the wrong module (low). It's added in client_window_factory.add_client_window_arguments, but it's a model-output concern with no window involvement — cli._add_session_arguments is the cohesive home. --stats-path is misplaced the same way already, so this follows precedent rather than setting it; flagging so the precedent doesn't harden.

3. Output directory isn't validated until close() (low). __init__ stores the path and open() ignores it; the mkdir happens at teardown. A path that exists as a file, or an unwritable directory, fails after the entire rollout. open() has the session and could mkdir(parents=True, exist_ok=True) there. This mirrors MetricsOutputSink, but the stakes differ: the LingBot GPU run in the evidence is ~39s for two chunks and minutes-to-tens-of-minutes at realistic chunk counts, and #524 already fails fast on the schema half of the same flag.

4. Sinks opened after session.init() are unreachable by either cleanup path (low). run_session appends to attempted_output_sinks inside the try, after session.init(). If init() raises, model_output_sinks never get closed there — and ApplicationRunner won't close them either, because run_started is already True. Harmless for TensorArtifactOutputSink (nothing buffered, close() no-ops), but the ModelOutputSink protocol invites implementations that acquire resources in __init__.

5. Nit. close()'s finally unlinks every staged temp path, including the ones os.replace already consumed. It's correct (missing_ok=True), but it reads as unconditional cleanup when the intent is cleanup-on-failure.

---

#520 — LingBot-VA integration

1. The destructive teardown frees none of the KV it targets — confirmed by execution

engine.py:509-513, inside _release_denoising_state:

block_cache.self_attn.reset()
block_cache.cross_attn.text.k = torch.empty(0)
block_cache.cross_attn.text.v = torch.empty(0)

Both lines are no-ops:

• BlockKVCache stores _k / _v. It has no k / v fields and no __slots__, so those two assignments silently create two brand-new attributes and leave the buffers fully allocated.
• BlockKVCache.reset() is documented as "Reset bookkeeping while preserving the allocated tensor storage." It resets _prev_chunk_idx / _curr_chunk_idx / _n_cached and nothing else.

Verified directly:

fields: [... '_n_cached', '_k', '_v']       has slots: False
after assigning .k/.v, _k numel still: 24   same buffer: True
after reset(),          _k numel: 24

Scale under the Robotwin config: window_slots 36 × slot_size 272 = 9,792 tokens × 24 heads × 128 head_dim × 2 bytes × 2 (k,v) × 30 blocks × 2 CFG branches ≈ 7.2 GiB of self-attention KV, plus ~0.4 GiB of cross-attention KV.

The buffers are eventually reclaimed by refcount when self._pipeline_cache = None and the local cache goes out of scope — but that happens after the gc.collect() and _empty_cuda_cache() at the end of the same function, so the caching allocator is never trimmed once they're actually released. The run passes today because 37 GiB peak on a 97 GiB card leaves room and the allocator reuses freed blocks, so this reads as latent rather than currently-breaking. But the README's stated justification for the whole destructive design is "VAE decoding cannot fit beside the complete DiT, text encoder, and KV footprint on the supported capacity path" — this is precisely the release that isn't happening, and no test covers _release_denoising_state.

Suggested fix, which doubles as a reuse win: add a public release() / free_storage() to BlockKVCache in flashdreams.core.attention.kvcache — every streaming integration wants it — and call that. At minimum, del cache before the gc.collect().

2. Rolling-window steady state reports a duplicated chunk as committed — confirmed by execution

kvcache.py:86-88 returns BlockKVCache._n_cached as n_committed_tokens, and network.py:298-308 slices _k[:, :n_committed_tokens]. _n_cached saturates at the buffer size. Once before_update() rolls the window left, the trailing chunk_size slots hold a stale copy until update() overwrites them — but _extract_cache_tensors runs at the top of forward_video / forward_action, before that write.

Simulated with a 3-slot cache, one distinct value per chunk:

chunk 2: n_committed=6  committed= [1,1,1, 2,2,2]              # correct
chunk 3: n_committed=9  committed= [2,2,2, 3,3,3, 3,3,3]       # chunk 2 duplicated
chunk 4: n_committed=9  committed= [3,3,3, 4,4,4, 4,4,4]

From the first steady-state chunk onward, every denoising step attends to the previous chunk's KV twice and one chunk too few of history. Shapes don't change, nothing raises, and the outputs stay finite — it just quietly stops matching upstream.

Reachable at --num-chunks 37 (ROBOTWIN_ATTENTION_WINDOW 72 → 36 slots). --num-chunks has no upper bound in _parse_args or LingbotVAEngineConfig.__post_init__, the README's "Remaining limits" section doesn't mention it, and the deepest test is 2 chunks. The correct value is already computed inside BlockKVCache: _current_write_bounds()[0] = min(_n_cached, total_size - chunk_size). Either expose it, or cap num_chunks at window_slots until there's a test past the window.

3. The unconditional branch is computed and discarded when its scale is 1.0 (efficiency)

transformer/__init__.py:299-315 — predict_action_flow always runs a full forward_action on network_cache_uncond whenever that cache exists, then returns flow_cond unless action_guidance_scale > 1.0. Robotwin's default action_guidance_scale is exactly 1.0, so every one of the ~50 action denoise steps per chunk pays a full 30-block unconditional forward whose result is thrown away.

The uncond cache exists because guidance_scale (5.0) needs it for video, and the persist=True step must still run to advance the uncond KV. But the ~50 non-persist steps don't. Guarding with if persist or scale > 1.0 roughly halves the action denoise loop at zero output change. predict_flow is symmetric for guidance_scale <= 1.0.

4. VAKVCache.committed_kv_plus_fresh is dead code (low)

Never called — the compiled path takes pre-extracted tensors from _extract_cache_tensors instead. ~22 lines, and it carries the same n_committed_tokens defect as #2, so it would be wrong if revived. VASelfAttention's docstring still points at it.

5. Private-attribute reach-through into flashdreams core (low)

_extract_cache_tensors uses bc.self_attn.kv_cache._k/._v and bc.cross_attn.text._k/._v/._n_cached; VAKVCache.n_committed_tokens uses _n_cached. (Engine-side: vae._feat_map, vae._conv_idx, transformer._network, and diffusers.models.autoencoders.autoencoder_kl_wan.unpatchify — the diffusers half is acknowledged in the README and pinned to <0.39.) The flashdreams-side reach-through is the half that's cheap to fix, and the fix is the same public accessor #2 needs.

6. Decode denormalizes through a bf16 reciprocal (low)

engine.py:545-556 builds latent_inverse_std = 1.0 / std in config.dtype (bf16) and then computes latent / latent_inverse_std — two bf16 roundings where latent * std needs one. _encode_observation does the corresponding multiply in fp32, so the two paths aren't symmetric.

7. network.to("cpu") during teardown (low)

_release_denoising_state moves the 5.09B-parameter DiT to host memory purely in order to free the GPU allocation — a ~10 GiB host spike — and the local network reference keeps it alive through the gc.collect() at the end of the function. Dropping the reference frees the GPU memory by refcount without the host copy. Same pattern for self._text_encoder.to("cpu").

8. Nits

• _prompt_clean guards import ftfy with except ImportError: pass. ftfy>=6.0 is a hard flashdreams dependency and lingbot_va depends on flashdreams, so the branch is unreachable — and if it ever were reachable, prompts would be cleaned differently from the recorded parity evidence, silently.
• resolve_prompt validates user input with bare assert (stripped under -O); the rest of the package raises ValueError.
• LingbotVAEngine.run() swallows cleanup failures with except Exception: pass — no log line. The README's "cleanup errors do not replace an earlier inference failure" is satisfied, but the error vanishes entirely.
• ROBOTWIN_ENV_TYPE is unused.
• config.py hardcodes guidance_scale=5.0 / action_guidance_scale=1.0 where ROBOTWIN_GUIDANCE_SCALE / ROBOTWIN_ACTION_GUIDANCE_SCALE already exist.
• attn_window: int = 64 is the default in both LingbotVAInferencePipelineConfig and LingbotVATransformerConfig while the real value is 72. Always overridden, so dead — but two wrong numbers sitting next to the right one.

Test coverage gaps

No test exercises past the attention window (#2), _release_denoising_state (#1), or the discarded-uncond path (#3). All three are CPU-testable with the existing fake-engine machinery.

---

SLoC and reuse analysis

Where the 5,532 lines go: 3,140 model src · 561 model tests · 425 compare_upstream.py · 444 V2 adapter · 552 V2 tests · 270 docs · ~140 packaging.

For calibration: the five T2V V2 integrations that build on flashdreams.t2v_v2 are 44–66 lines of app code each (t2v_wan21/app.py is 55). This adapter is 444. Most of that gap is real — I2AV instead of T2V, tensor artifacts, a one-shot destructive engine that can't use T2VApplication's per-chunk model loop — but not all of it.

Concrete reductions, highest payoff first:

(a) UMT5 text encoder — ~70 lines, the component already exists. flashdreams.infra.encoder.text.umt5 provides UMT5TextEncoder and prompt_clean. engine._prompt_clean (11 lines), engine._encode_prompt (30), and _loaders.load_text_encoder + load_tokenizer (32) are near line-for-line copies — same 512-token max-length contract, same mask.gt(0).sum(dim=1) trim, same zero-pad-and-stack. Two blockers, both one-line upstream config changes: UMT5TextEncoderConfig.model_id_or_local_path is Literal[<3 Wan repo ids>] so a Robotwin snapshot path doesn't fit, and it pins T5Tokenizer where LingBot needs T5TokenizerFast. Widening that field to str and adding a tokenizer knob deletes ~70 lines here and stops prompt cleaning from drifting between the two copies.

(b) Wan VAE — ~100 lines and a documented limitation. flashdreams.recipes.wan.autoencoder.vae.WanVAE already implements the streaming causal Wan 2.2 VAE at z_dim=48, patch_size=2, is_residual=True — the same family Robotwin uses — with prepare_cache() / encode(x, cache) / decode(z, cache), _patchify/_unpatchify, and wan22_ti2v_5b_vae_state_dict_transform for diffusers-layout keys. That subsumes load_vae + WanVAEStreamingWrapper + patchify (~100 lines in _loaders.py) and the hand-rolled per-frame vae._feat_map / _conv_idx decode loop in engine.py. It also retires the README's "Diffusers 0.38 private VAE streaming state" limitation, the diffusers>=0.38,<0.39 pin, and the two-WanVAEStreamingWrappers-over-one-module trick — prepare_cache() returns independent cache objects, which is exactly what those two wrappers are faking. Caveat: it's a reimplementation, so the recorded action SHA-256 and the eager/compiled parity thresholds would need re-validating. Better as a tracked follow-up than as churn in this PR.

(c) Checkpoint remap — ~15 lines. flashdreams.core.checkpoint.remap.remap_checkpoint_keys is the established shared helper — hy_worldplay, fastvideo_causal_wan22, recipes/cosmos, and recipes/wan all use it. transformer/checkpoint.state_dict_transform reimplements the same first-match-wins re.match / re.sub loop. The drop-prefix filter and the strict unmapped-key error are genuine improvements over the common version — which is the argument for pushing them into it as optional drop_prefixes / strict parameters, so the other four callers get the strictness too.

(d) VABlock — ~50 lines. VABlock.__init__ and update_parameters_after_loading_checkpoint are copies of recipes/wan/transformer/impl/modules.Block's, and VABlock.forward's cross-attention arm re-derives CrossAttention.forward's non-i2v body inline. A _self_attention_cls class attribute on Block (one line) lets VABlock subclass it; a tensor-level CrossAttention.forward_kv(x, k, v) that forward delegates to (six lines) removes the duplicated math without putting a cache object inside the compiled graph. The duplicated __init__ is the riskier half — the checkpoint remap rules target those exact submodule names, so a rename in the common Block breaks LingBot silently rather than loudly.

(e) RoPE — ~35 lines. compute_rope_freqs_from_grid re-derives what flashdreams.core.attention.rope already has: the head-dim split (head_dim//3 vs head_dim//6*2 — identical at 128), the base frequencies (_compute_freqs with extrapolation_ratio=1.0), and the interleave (_cat_freqs(interleaved=True)). The real difference is arbitrary per-token grid ids instead of a regular t/h/w grid, plus float64 base accumulation. Promoting _compute_freqs with a dtype argument leaves ~15 lines of grid indexing here.

(f) Scheduler — ~90 lines. LingbotVAFlowMatchScheduler differs from FlowMatchEulerDiscreteScheduler in exactly two ways: sigma_max=1.0 rather than the train-grid 1 - 1/1000, and an N-entry schedule plus a padded_timesteps property rather than the common N+1 with a terminal 0.0. Those are the same thing — fixed_timesteps = tuple(warp(linspace(1,0,N+1)[:-1], shift) * 1000) + (0.0,) reproduces padded_timesteps exactly. Since the pipeline drives its own loop and only ever touches step() and padded_timesteps (never sample()), reuse needs a public step() on the common Euler scheduler — which its other users would also benefit from. That leaves a ~25-line schedule-builder subclass.

(g) Config duplication — ~25 lines. LingbotVAInferencePipelineConfig and LingbotVATransformerConfig each carry latent_height, latent_width, frame_chunk_size, action_per_frame, and attn_window, with literal defaults that config.py always overrides from constants.py, and nothing checks that the two copies agree. Have the pipeline read them from self.transformer.config and drop the defaults.

(h) V2 adapter — ~30 lines, and reusable beyond this PR. _validate_requested_session is the generic "my model has one non-negotiable output contract" guard. T2VApplication._validate_layout / _validate_frame_size is a weaker version of the same idea. One require_session_fields(requested, canonical, fields) helper in flashdreams.runtime_v2 removes it here and lets the T2V apps drop theirs.

(i) Headers — ~55 lines. All four integrations_v2/lingbot_va files carry the 15-line Apache header; sibling V2 integrations (t2v_wan21/app.py) use the 2-line SPDX form.

Net: roughly 400–500 lines off the 3,584 non-test, non-doc source — about 12–15% — of which (a), (c), (g), (h), and (i) are safe now, (d)/(e)/(f) need small additive changes to shared flashdreams code, and (b) is a follow-up gated on re-validating the parity evidence. Most of that isn't really about line count: (a), (c), (d), and (e) each replace a silently drifting copy of a shared component with a call into it.

What this analysis does not suggest shrinking: WanVADiTNetwork's forward paths, the CFG/KV lifecycle, LingbotVAActionProcessor, and the engine's phase state machine. That is the actual model, and its size is honest.

One structural note: compare_upstream.py is 425 lines, 8% of the diff. [tool.setuptools.packages.find] include = ["lingbot_va*"] correctly keeps it out of the wheel, so it's review surface rather than shipped surface — fine to keep, worth knowing when reading the diff stat.

---

Summary

#524 is ready modulo the failed-run-writes-artifacts question (#524.1), which is a design call rather than a defect.

For #520, the two confirmed findings — the teardown that frees nothing (#1) and the steady-state cache duplication (#2) — are the ones I'd want resolved before merge, along with the discarded uncond forward (#3), which is a free ~2× on the action denoise loop. Everything else is cleanup. The reuse work in (a), (c), (g), (h), (i) is safe to do in this PR; (b), (d), (e), (f) are better as follow-ups since they touch shared flashdreams code that other integrations depend on.

The evidence discipline in the README — exact hashes, thresholds, checkpoint key counts, reproduction commands, and an explicit limits section — is the right standard for an integration of this size.

[jmccaffrey-nv]

Addressed the actionable PR 520 review findings in 20b8d1c6 (evidence correction in dc92b215):

• fixed steady-state KV rollover so the stale trailing write region is not exposed as committed context; added a CPU regression that crosses the window boundary
• dropped cache/DiT/text owners before GC and CUDA allocator trimming, removed the ineffective cache-field assignments and host model copies, and added teardown-order coverage
• skipped the inactive action unconditional branch on non-persistent steps while retaining its persistent cache-advance step; the same guard is applied symmetrically to video
• removed the dead cache read path and FlashDreams-private KV reads; fixed latent denormalization, cleanup logging, prompt validation, stale defaults/constants, and scale constant reuse

Final validation: 275 CPU tests passed (2 GPU tests deselected), all pre-commit hooks passed including ty, and the real two-chunk GPU test passed in 34.02 s. The action artifact remains byte-identical (463b307b…8d95), with 39,804,413,440-byte peak allocation.

The suggested shared UMT5/Wan-VAE/scheduler/RoPE/block/checkpoint/session-validation refactors are reasonable follow-ups, but are intentionally not mixed into this model-only stacked PR because they cross the #524 boundary or require renewed parity validation. The #524 sink-lifecycle observations likewise belong on that dependency PR rather than in this diff.

[jarcherNV]

Should the 2 in cfg.attn_window // 2 be cfg.frame_chunk_size? A slot holds one chunk of latent frames, so those look like they should be the same number. They match right now, so nothing's broken, but frame_chunk_size defaults to 4 on line 107 of this file, so if anyone ever leans on that default the window silently doubles. Should line 107's default be 2, to match the pipeline config and ROBOTWIN_FRAME_CHUNK_SIZE?

[jmccaffrey-nv]

(GPT Sol) Good catch on the stale default: frame_chunk_size now defaults to 2, matching Robotwin and the pipeline, with focused CPU coverage. I kept attn_window // 2 intentionally: pinned upstream create_empty_cache() uses that divisor for both its two-frame Robotwin config and its four-frame LIBERO/demo/Franka configs. The window counts paired video/action regions, not frames per chunk. I documented that invariant in code and the architecture README. Addressed in 830019e6.

[jarcherNV]

Is action_text_embedding meant to be used somewhere? I can't find a caller. initialize_cache on line 264 builds the cross-attention text context once from self.text_embedding (the video one), and both the video and action forwards read that same cache.

[jmccaffrey-nv]

(GPT Sol) It is intentionally not used in inference. The pinned upstream model deep-copies the full condition embedder, but _input_embed(..., input_type="text") always calls condition_embedder.text_embedder; action mode only selects the copied time embedding/projection. Wiring the copied action text projection here would therefore diverge from upstream and invalidate the recorded parity. I retained it solely so the upstream checkpoint loads strictly and documented that invariant in 830019e6.

[jarcherNV]

Left a couple of comments but I don't think they should block anything. Up to you if you want to address them or not. The PR LGTM.

[jmccaffrey-nv]

/ok to test bd95bf5

[jmccaffrey-nv]

Design-review follow-up: the three questions are valid and are addressed in 7a15f43. The V2 session/model-loop now lives separately from app.py; the documentation explains why this fixed-input video-plus-action lifecycle does not map to the shared text-only T2V session; and the model package now includes a validated action manifest loader with four-panel plot and optional named CSV output. The docs also distinguish this offline I2AV adapter from the pinned upstream closed-loop RoboTwin client/server path.

[jmccaffrey-nv]

/ok to test 7a15f43

[jmccaffrey-nv]

/ok to test 7a15f43