Sweep all 296 v2 modules: simplify verbose docstrings/comments and remove 376 dangling
"(design_vN §X)" citations to the now-deleted design docs (v2/README.md is the source of
truth). Comment/docstring-only — AST-verified code-identical; the CPU suite holds at 237
passed / 5 pre-existing. Also applies yapf + ruff --fix auto-fixes (import ordering,
forward-ref annotation de-quoting under `from __future__ import annotations`; behavior-
neutral, suite-confirmed) and adds the legitimate domain terms mot/clen/te to the codespell
ignore-list. Remaining ruff (24) + mypy (68 no-untyped-call) findings are pre-existing v2
debt, untouched here.
First net-new architecture ported via the self-contained recipe-package pattern
(card-declared adapter + new loop, no shared-dispatch edit):
* CosmosDenoiseLoop (v2/recipes/cosmos2/loop.py): EDM preconditioning folded into
a flow-match Euler integrator. Faithful port of CosmosDenoisingStage — Karras
sigma schedule (rho=7, sigma_max=80 -> sigma_min=0.002, terminal clamp), latent
init randn*sigma_max, per-step c_in/c_skip/c_out (sigma_data=1) -> x0, CFG in x0
space, x0 -> velocity (x-x0)/sigma, FLOW_MATCH_STEP. video2world frame-replace
conditioning threaded but inert for the t2v preset.
* build_karras_sigmas helper added to v2/loop/sampler.py.
* CosmosDiT + CosmosT5Encoder adapters (v2/platform/backends/torch_cosmos.py),
declared on the card via ComponentSpec.adapter (the ac29750b enabler) — DiT
returns raw EDM output + builds the mandatory zero condition/padding masks + fps;
T5 uses the raw last_hidden_state (no Wan zero-pad). Reuses the WanVAE adapter.
* card/program/registry (HF id nvidia/Cosmos-Predict2-2B-Video2World + arch
fallback on CosmosTransformer3DModel) + COSMOS_NEG prompt + SamplingDefaults
(35 steps, gs 7, 704x1280, 93f, 16fps).
CPU-verified: Karras schedule, card/program build, registry resolve (id + arch),
EDM loop runs end-to-end on the CPU toy backend. GPU load/run is BRINGUP.
v2 mini green (211 passed, 2 skipped); pre-commit clean.
Two more pipelines stress-testing the design, plus a BAGEL-placement note in
designv4. Both fit with no new runtime primitive (109 tests pass).
Pipeline 1 — cross-model T2I→I2V (program/workflow.py, models/image_video/):
- Realizes ProgramKind.WORKFLOW as a thin multi-instance orchestrator ABOVE the
engine (the hot path stays single-instance). A Program composes one model's
loops; a Workflow chains full engine.run calls across distinct cards, threading
artifacts. (LTX-2 already covers same-card multi-stage; cross-model — FLUX→Wan
— is the new capability the single-instance runner can't express.)
- flux-t2i (text→image) and wan-i2v (text+image→video) cards; the I2V program
folds the conditioning image into text_embeds so WanDenoiseLoop is unchanged.
- Each model keeps its own interleave-parity guarantee (crossing instances is a
Workflow boundary, not a loop step). 7 tests incl. video-depends-on-image.
Pipeline 2 — N-way joint RL (training/methods/joint_multi_rl.py, models/multi_expert/):
- JointMultiExpertRL generalizes UnifiedRLMethod (N=2) to N refiner LMs + a
generator: one reward → one group advantage → N token-PG updates + 1 FlowGRPO
PPO update, N+1 independent WeightSyncPlans. Proves the substrate was already
N-ready (card holds N components/loops; per-component weight-sync; dict grad
targets) — only the method body looped over two; now it loops over a list.
- credit="per_expert" learns all N cleanly; credit="shared" (faithful to UniRL)
works but is noisier — the honest multi-agent credit-assignment result, a
reward-shaping choice, not a substrate limit. 6 tests (N=1,3,4; prompt-only).
Fix — flow_sde_ml_velocity (loop/sampler.py): the toy FlowGRPO generator update
targeted the velocity the model already produced (a no-op once guidance_scale=1
was set for the C2 identity; the unified generator moved only on ~1e-7 noise).
The correct PG surrogate targets the max-likelihood velocity of the realized
sample — nonzero at ratio==1. Both UniRL and N-way generators now learn for real;
the C2 ratio==1 identity still holds (measured before the update).
BAGEL: MoT/shared-weight (one transformer on both loops), same row as Cosmos3;
real BAGEL's co-resident experts are expressible via the expert-routing policy
(partial sharing) — captured in designv4 §2.3.
Stress-tests the v2 Card/Loop/Program design with a UniRL/PromptRL-style
joint RL recipe: a prompt-refiner LM expert and a flow generator expert,
two separate experts driven by two loop types in one request, both updated
simultaneously from a single RL reward.
- loop/sampler.py: flow_sde_step_with_logprob — FlowGRPO SDE rollout sampler
(per-step Gaussian log-prob), distinct from the deterministic ODE serve step.
- request/params.py: gated sde_rollout/sde_noise_scale on DiffusionParams so
the serve path stays byte-identical (default ODE).
- models/wan21/loop.py: gated SDE-rollout capture in WanDenoiseLoop.advance.
- models/backend.py: ToyPromptRefiner — a real REINFORCE categorical policy
(the Qwen role), separate weights from the generator.
- models/unified/: the unified card+program — two disjoint experts (llm +
transformer) on ar_decode + diffusion_denoise; the topological opposite of
the Cosmos3 MoT card, same vocabulary.
- training/methods/unified_rl.py: joint GRPO — one reward -> group advantage
-> LM token policy gradient + DiT FlowGRPO PPO; two LRs; prompt-only/joint
flag; reuses the shared diffusion loop for rollout.
- training/weight_sync.py: WeightSyncPlan gains a component scope so the two
experts version + cache-invalidate independently (LM sync never flushes the
frozen text-encoder feature cache).
- tests/test_unified_rl.py: 9 tests incl. likelihood-based C2 identity, the
two-loop interleave parity gate, joint vs prompt-only. Full suite 91 passed.
designv4.md: unified design doc reflecting v2 as built+tested, with the joint
RL stress test as the validating case study (the design held — new card +
new method, no new runtime primitive).
Directory rename (git mv, history preserved) plus rewrite of all references: absolute imports in
tests, the zero-dep runner, docstrings, comments, and the README. No behavior change.
Run: python3 -m pytest v2/tests/ -q ; python3 v2/run_tests.py ; python3 -m v2.examples
SolitaryThinker