From 94b84b03b9a4eb8ec02a4e034f7fcc48d7a8fe47 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Tue, 16 Jun 2026 22:13:43 -0700 Subject: [PATCH] =?UTF-8?q?[examples]=20Add=20v2=5Fexamples/{training,omni?= =?UTF-8?q?,workflows}/=20=E2=80=94=20runnable=20examples?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three more example folders alongside inference/, all CPU/numpy, self-contained (sys.path bootstrap), public API only, every script verified to run green. training/ (7) — one per method, all via the uniform method.train_step seam: 01 finetune · 02 dmd2 distillation · 03 diffusion_nft (likelihood-free RL, samples from the old policy, feature-cache reuse) · 04 self-forcing (causal chunk_rollout) · 05 joint LM+generator RL (UniRL; joint + prompt-only) · 06 N-way joint RL (per_expert vs shared credit) · 07 end-to-end workflow RL (T2I+I2V from one final-video reward). omni/ (4) — 01 Cosmos3 (reason→joint denoise, shared MoT) · 02 BAGEL (text→image, shared MoT; scheduler prices both WorkUnit kinds) · 03 Qwen-Omni (thinker→talker→vocoder, three separate experts, text+audio) · 04 interleave parity across AR + diffusion loop types. workflows/ (2) — 01 cross-model T2I→I2V workflow (image provably conditions the video) · 02 workflow as a first-class servable (requires/validate, address by id, register_workflows catalog, WorkflowRegistry). Each folder has a README indexing its scripts. --- v2_examples/omni/01_cosmos3.py | 39 ++++++++++++++ v2_examples/omni/02_bagel.py | 40 ++++++++++++++ v2_examples/omni/03_qwen_omni.py | 42 +++++++++++++++ .../omni/04_interleave_across_loop_types.py | 39 ++++++++++++++ v2_examples/omni/README.md | 25 +++++++++ v2_examples/training/01_finetune.py | 36 +++++++++++++ v2_examples/training/02_distillation_dmd2.py | 37 +++++++++++++ v2_examples/training/03_rl_diffusion_nft.py | 44 +++++++++++++++ v2_examples/training/04_self_forcing.py | 37 +++++++++++++ .../training/05_joint_lm_generator_rl.py | 53 ++++++++++++++++++ v2_examples/training/06_nway_joint_rl.py | 49 +++++++++++++++++ .../training/07_end_to_end_workflow_rl.py | 49 +++++++++++++++++ v2_examples/training/README.md | 26 +++++++++ v2_examples/workflows/01_t2i_then_i2v.py | 42 +++++++++++++++ .../workflows/02_register_as_servable.py | 54 +++++++++++++++++++ v2_examples/workflows/README.md | 18 +++++++ 16 files changed, 630 insertions(+) create mode 100644 v2_examples/omni/01_cosmos3.py create mode 100644 v2_examples/omni/02_bagel.py create mode 100644 v2_examples/omni/03_qwen_omni.py create mode 100644 v2_examples/omni/04_interleave_across_loop_types.py create mode 100644 v2_examples/omni/README.md create mode 100644 v2_examples/training/01_finetune.py create mode 100644 v2_examples/training/02_distillation_dmd2.py create mode 100644 v2_examples/training/03_rl_diffusion_nft.py create mode 100644 v2_examples/training/04_self_forcing.py create mode 100644 v2_examples/training/05_joint_lm_generator_rl.py create mode 100644 v2_examples/training/06_nway_joint_rl.py create mode 100644 v2_examples/training/07_end_to_end_workflow_rl.py create mode 100644 v2_examples/training/README.md create mode 100644 v2_examples/workflows/01_t2i_then_i2v.py create mode 100644 v2_examples/workflows/02_register_as_servable.py create mode 100644 v2_examples/workflows/README.md diff --git a/v2_examples/omni/01_cosmos3.py b/v2_examples/omni/01_cosmos3.py new file mode 100644 index 000000000..7bb55e161 --- /dev/null +++ b/v2_examples/omni/01_cosmos3.py @@ -0,0 +1,39 @@ +#!/usr/bin/env python3 +"""01 — Cosmos3: reason → joint denoise on ONE resident MoT instance. + +A single request runs the reasoner (`ar_decode`, the und pathway) to upsample the prompt, packs its +tokens into conditioning, then the joint generation (`diffusion_denoise`) — **both loops bound to the +same `transformer`** (shared weights). The §16 claim no DAG-of-engines can express, made native: both +the AR tokens and the denoise steps are WorkUnits the scheduler prices and interleaves. + +Run: python3 v2_examples/omni/01_cosmos3.py +""" +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +import numpy as np + +from v2.models import build_omni_engine +from v2.request import DiffusionParams, SamplingParams, TaskType, make_request + + +def main() -> None: + eng = build_omni_engine() # registers cosmos3-vfm, bagel-mot, qwen-omni-tts + out = eng.run(make_request( + TaskType.T2V, "cosmos3-vfm", "a phoenix rising over a volcano", + sampling=SamplingParams(max_tokens=6, seed=1), # the reasoner (ar_decode) + diffusion=DiffusionParams(num_steps=4, seed=1))) # the joint denoise (diffusion_denoise) + + print("Cosmos3 (one MoT instance, two loop types on shared weights):") + print(f" reasoner text : {out.artifacts['text'].text}") + print(f" video : {np.asarray(out.artifacts['video'].frames).shape}") + print(f" metrics : reasoner tokens + denoise steps both ran as WorkUnits " + f"(gpu_seconds={out.metrics['gpu_seconds']:.2e})") + + +if __name__ == "__main__": + main() diff --git a/v2_examples/omni/02_bagel.py b/v2_examples/omni/02_bagel.py new file mode 100644 index 000000000..da11dd461 --- /dev/null +++ b/v2_examples/omni/02_bagel.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python3 +"""02 — BAGEL: generate_text → generate_image on ONE resident MoT instance. + +The canonical vllm-omni unified AR+diffusion model. `generate_text` (`ar_decode`) and `generate_image` +(`diffusion_denoise`) both bind the one resident `transformer`. Unlike vllm-omni's opaque `DIFFUSION` +stage, BOTH loops are runtime-visible here — the scheduler prices `ar_token` AND `diffusion_step` +WorkUnits. (BAGEL is MoT/shared-weight, the same topology row as Cosmos3 — see designv4 §2.3.) + +Run: python3 v2_examples/omni/02_bagel.py +""" +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +import numpy as np + +from v2.models import build_omni_engine +from v2.request import DiffusionParams, SamplingParams, TaskType, make_request + + +def main() -> None: + eng = build_omni_engine() + out = eng.run(make_request( + TaskType.T2I, "bagel-mot", "a porcelain teapot on a windowsill", + sampling=SamplingParams(max_tokens=6, seed=2), + diffusion=DiffusionParams(num_steps=4, seed=2))) + + print("BAGEL (generate_text → generate_image, shared MoT weights):") + print(f" text : {out.artifacts['text'].text}") + print(f" image : {np.asarray(out.artifacts['image'].tensor).shape}") + # the scheduler saw BOTH WorkUnit kinds (not one opaque stage) + by_kind = dict(eng.admission.metrics.by_kind) + print(f" WorkUnit kinds priced by the scheduler: {by_kind}") + + +if __name__ == "__main__": + main() diff --git a/v2_examples/omni/03_qwen_omni.py b/v2_examples/omni/03_qwen_omni.py new file mode 100644 index 000000000..1a956bd04 --- /dev/null +++ b/v2_examples/omni/03_qwen_omni.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python3 +"""03 — Qwen-Omni: thinker → talker → vocoder (three separate experts). + +The cascade topology (designv4 §9.5): three *distinct* experts on three loop types in one request — +the `thinker` (`ar_decode` → text), the `talker` (`ar_decode` → speech codec tokens, conditioned on the +thinker's tokens + hidden state), and the `vocoder` (`audio_decode` → waveform). Produces both text and +audio; the vocoder streams `AUDIO_CHUNK` units. This is the topological opposite of MoT: no shared +weights, but the same Card/Loop/Program vocabulary. + +Run: python3 v2_examples/omni/03_qwen_omni.py +""" +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +import numpy as np + +from v2.models import build_omni_engine +from v2.request import OutputSpec, TaskType, make_request + + +def main() -> None: + eng = build_omni_engine() + out = eng.run(make_request( + TaskType.T2A, "qwen-omni-tts", "Explain audio generation in fifteen words.", + outputs=OutputSpec(stream={"audio": True}))) + + samples = np.asarray(out.artifacts["audio"].samples) + print("Qwen-Omni (thinker → talker → vocoder, three separate experts / three loops):") + print(f" thinker text : {out.artifacts['text'].text}") + print(f" speech tokens : {out.artifacts['speech_tokens'].tensor}") + print(f" audio waveform : shape={samples.shape} sample_rate={out.artifacts['audio'].sample_rate} " + f"range=[{samples.min():.3f}, {samples.max():.3f}]") + print(f" vocoder chunks : {out.metrics['audio_chunks']:.0f} AUDIO_CHUNK units " + f"(total streamed events incl. AR text: {out.metrics.get('stream_chunks', 0):.0f})") + + +if __name__ == "__main__": + main() diff --git a/v2_examples/omni/04_interleave_across_loop_types.py b/v2_examples/omni/04_interleave_across_loop_types.py new file mode 100644 index 000000000..1614283b1 --- /dev/null +++ b/v2_examples/omni/04_interleave_across_loop_types.py @@ -0,0 +1,39 @@ +#!/usr/bin/env python3 +"""04 — The interleave parity gate holds across loop *types* (omni). + +The loop-inversion safety guarantee isn't limited to homogeneous denoise batches: concurrent omni +requests — each running an `ar_decode` loop THEN a `diffusion_denoise` loop — interleaved at step +granularity are still bit-identical to serial. Per-request state lives in `LoopState`, so mixing AR +tokens and denoise steps from different requests in one schedule cannot smear them (designv4 §6, §9.5). + +Run: python3 v2_examples/omni/04_interleave_across_loop_types.py +""" +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +from v2.models import build_omni_engine +from v2.parity import assert_interleave_parity +from v2.request import DiffusionParams, SamplingParams, TaskType, make_request + + +def _cosmos(prompt, seed): + return make_request(TaskType.T2V, "cosmos3-vfm", prompt, + sampling=SamplingParams(max_tokens=5, seed=seed), + diffusion=DiffusionParams(num_steps=4, seed=seed)) + + +def main() -> None: + eng = build_omni_engine() + reqs = [_cosmos("a comet", 11), _cosmos("a glacier", 22), _cosmos("a comet", 11)] + divs = assert_interleave_parity(eng, reqs) + print("3 omni requests (each ar_decode → diffusion_denoise), serial vs step-interleaved:") + print(f" interleave parity : {'PASS — bit-identical across AR + diffusion steps ✓' if not divs else divs}") + print(" (mixing AR tokens and denoise steps from different requests cannot smear state)") + + +if __name__ == "__main__": + main() diff --git a/v2_examples/omni/README.md b/v2_examples/omni/README.md new file mode 100644 index 000000000..336aa02bc --- /dev/null +++ b/v2_examples/omni/README.md @@ -0,0 +1,25 @@ +# v2_examples/omni — MoT & cascade models + +Runnable examples of the omni models: one request that spans **multiple loop types**. Two topologies +(designv4 §2.3): + + * **MoT / shared-weight** (Cosmos3, BAGEL) — ONE resident instance whose `transformer` runs *both* an + `ar_decode` loop and a `diffusion_denoise` loop on the same weights. Every AR token and denoise step + is a runtime-visible WorkUnit (the differentiation vs vllm-omni's opaque DIFFUSION stage). + * **Cascade / separate experts** (Qwen-Omni) — three *distinct* experts on three loop types + (`thinker → talker → vocoder`), cross-stage conditioning, streaming codec→waveform. + +Run from anywhere: + +```bash +python3 v2_examples/omni/01_cosmos3.py +``` + +| Script | Model | Loops in one request | +|---|---|---| +| `01_cosmos3.py` | Cosmos3 (shared MoT) | `ar_decode` (reason) → `diffusion_denoise` (joint video) | +| `02_bagel.py` | BAGEL (shared MoT) | `ar_decode` (generate_text) → `diffusion_denoise` (generate_image) | +| `03_qwen_omni.py` | Qwen-Omni (3 separate experts) | `ar_decode` → `ar_decode` → `audio_decode` (text + speech) | +| `04_interleave_across_loop_types.py` | — | the interleave parity gate holds across AR + diffusion steps | + +`build_omni_engine()` registers all three; see designv4 §9.5 (Qwen-Omni cascade). diff --git a/v2_examples/training/01_finetune.py b/v2_examples/training/01_finetune.py new file mode 100644 index 000000000..2699ab262 --- /dev/null +++ b/v2_examples/training/01_finetune.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python3 +"""01 — Flow-match finetuning (Wan2.1-1.3B). + +The simplest method: regress the student's velocity prediction toward the flow-match target on noised +latents. It drives no rollout (it trains on given/sampled latents), but it goes through the same +`TrainingMethod` seam every other method does: `train_step(batch, iteration) -> (loss_map, metrics)`. + +Run: python3 v2_examples/training/01_finetune.py +""" +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +from v2.models.wan21 import build_wan21_card +from v2.training.methods import build_finetune + + +def main() -> None: + method = build_finetune(build_wan21_card(), lr=0.05) + batch = {"prompts": ["a red sports car", "a sailboat at dawn"], "seeds": [1, 2]} + + print("method:", method.name, "| manages_optimization:", method.manages_optimization()) + first = last = None + for it in range(20): + loss, metrics = method.train_step(batch, it) + first = first or metrics + last = metrics + print(f"loss : {first['loss']:.4f} -> {last['loss']:.4f} (decreasing = the student fits)") + print(f"grad_norm : {last['grad_norm/student']:.4f} (the #1396-style per-method regression signal)") + + +if __name__ == "__main__": + main() diff --git a/v2_examples/training/02_distillation_dmd2.py b/v2_examples/training/02_distillation_dmd2.py new file mode 100644 index 000000000..e339bf779 --- /dev/null +++ b/v2_examples/training/02_distillation_dmd2.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python3 +"""02 — DMD2 distillation (Wan2.1-1.3B). + +Distribution-matching distillation: three roles on the same card — a **student** (the few-step model +being trained), a trainable **fake-score critic**, and a frozen **teacher**. The generator loss is the +DMD score difference (teacher − critic) on the student's own rollout latents; the critic is trained to +score the student. Two grad streams (student + critic) ⇒ two grad-norm signals. + +Run: python3 v2_examples/training/02_distillation_dmd2.py +""" +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +from v2.models.wan21 import build_wan21_card +from v2.training.methods import build_dmd2 + + +def main() -> None: + method = build_dmd2(build_wan21_card(), rollout_steps=3) + batch = {"prompts": ["a waterfall in a canyon", "a lantern festival"], "seeds": [1, 2]} + + print("method:", method.name, "| roles: student + fake-score critic + frozen teacher") + last = None + for it in range(15): + loss, last = method.train_step(batch, it) + print(f"dmd_loss : {last['dmd_loss']:.4f} (generator: teacher − critic score difference)") + print(f"critic_loss : {last['critic_loss']:.4f} (critic learns to score the student's samples)") + print(f"grad_norm : student={last.get('grad_norm/student', 0.0):.4f} " + f"critic={last['grad_norm/critic']:.4f}") + + +if __name__ == "__main__": + main() diff --git a/v2_examples/training/03_rl_diffusion_nft.py b/v2_examples/training/03_rl_diffusion_nft.py new file mode 100644 index 000000000..f2c87f8f7 --- /dev/null +++ b/v2_examples/training/03_rl_diffusion_nft.py @@ -0,0 +1,44 @@ +#!/usr/bin/env python3 +"""03 — DiffusionNFT: likelihood-free RL (Wan2.1-1.3B). + +The landed RL method, and the reason for the C2 *split*: it is **likelihood-free** — no log-probs to +match — so consistency is the C2 *behavioral* rung (seeded sample + prediction-space identity). Key +properties this shows: + + * **Rollout == serve** — K samples of a prompt are produced by the SAME denoise loop the engine + serves (in the ROLLOUT profile), captured as behavior. + * **Samples from the *old* (decay-blended) policy**, not the student (the WeightSyncPlan carries that). + * **Group-relative advantages** drive the update; the shared prompt encodes ONCE per K-sample group + via the feature cache (the "24×" text-encode reduction). + +Run: python3 v2_examples/training/03_rl_diffusion_nft.py +""" +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +from v2.models.wan21 import build_wan21_card +from v2.training.methods import build_diffusion_nft + + +def main() -> None: + nft = build_diffusion_nft(build_wan21_card(), num_video_per_prompt=4, num_inner_timesteps=2) + batch = {"prompts": ["a red car", "a blue boat"], "seeds": [1, 2]} + + print("method:", nft.name, "| consistency:", nft.consistency_level().value, "(likelihood-free)") + print("| samples from role:", nft.old_sync.role.value, "(NOT the student)") + loss, m = nft.train_step(batch, 0) + print(f"policy_loss : {loss['policy_loss']:.3f} kl_div_loss : {loss['kl_div_loss']:.5f}") + print(f"reward_mean : {m['reward_mean']:.3f} advantage_std : {m['advantage_std']:.3f}") + print(f"old_decay : {m['old_decay']:.3f} old_weights_version : {m['old_weights_version']}") + + fc = nft.old.caches.stats()["feature"] + print(f"feature cache : {fc['hits']} hits / {fc['misses']} misses " + f"(each K-sample group encodes its prompt once)") + + +if __name__ == "__main__": + main() diff --git a/v2_examples/training/04_self_forcing.py b/v2_examples/training/04_self_forcing.py new file mode 100644 index 000000000..578be6b4a --- /dev/null +++ b/v2_examples/training/04_self_forcing.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python3 +"""04 — Self-forcing: causal video distillation (Wan-causal). + +"Train against your own KV-cached causal rollout." Self-forcing extends DMD2, but the student rollout +is the SHARED `chunk_rollout` loop (with the slab-KV cache) — exactly the causal loop the engine +*streams* at serve time. So the causal student is distilled under the runtime it will be served on +(the (recipe, runtime) pair). The loss is DMD2's distribution-matching, applied to the rollout latents. + +Run: python3 v2_examples/training/04_self_forcing.py +""" +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +from v2.models.wan_causal import build_wan_causal_card +from v2.training.methods import build_self_forcing + + +def main() -> None: + method = build_self_forcing(build_wan_causal_card()) + batch = {"prompts": ["a drone flight over dunes", "a river through a forest"], "seeds": [1, 2]} + + print("method:", method.name, "| student_loop_id:", method.student_loop_id, "(the causal, KV-cached loop)") + last = None + for it in range(12): + loss, last = method.train_step(batch, it) + print(f"dmd_loss : {last['dmd_loss']:.4f} critic_loss : {last['critic_loss']:.4f}") + print(f"grad_norm : student={last.get('grad_norm/student', 0.0):.4f} critic={last['grad_norm/critic']:.4f}") + print("\nThe student trains on its OWN causal rollout (chunk_rollout + slab-KV) — same loop the " + "engine streams.") + + +if __name__ == "__main__": + main() diff --git a/v2_examples/training/05_joint_lm_generator_rl.py b/v2_examples/training/05_joint_lm_generator_rl.py new file mode 100644 index 000000000..f147d3bb4 --- /dev/null +++ b/v2_examples/training/05_joint_lm_generator_rl.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +"""05 — Joint LM + generator RL (UniRL/PromptRL). + +One reward, two trainable experts: an `llm` prompt-refiner and a `transformer` flow generator. The +reward → a group-relative advantage → BOTH a **token policy gradient** on the LM and a **FlowGRPO PPO** +on the DiT, with two independent WeightSyncPlans. Likelihood-**based** C2: the per-step log-prob +identity (recomputed logp == rollout logp ⇒ PPO ratio ≈ 1). `joint=False` is PromptRL's prompt-only +mode (generator frozen; only the LM learns). + +Run: python3 v2_examples/training/05_joint_lm_generator_rl.py +""" +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +from v2.models.unified import build_unified_card +from v2.training.methods import build_unified_rl + +TARGET = 3 + + +def _run(joint: bool, iters: int = 40): + m = build_unified_rl(build_unified_card(), joint=joint, target_action=TARGET, + num_samples_per_prompt=4, num_skip_refinement=1) + batch = {"prompts": ["a fox", "blue sky", "a city at night"], "seeds": [1, 2, 3]} + p0 = float(m.llm._probs()[TARGET]) + last = None + for it in range(iters): + _, last = m.train_step(batch, it) + return m, p0, last + + +def main() -> None: + print("== joint (LM + generator both trained) ==") + m, p0, last = _run(joint=True) + print(f" LM P(target action): {p0:.3f} -> {last['refine_target_prob']:.3f} (the LM learns to refine)") + print(f" lm_pg_loss={last['lm_pg_loss']:.4f} dit_pg_loss={last['dit_pg_loss']:.4f} " + f"ppo_ratio={last['ppo_ratio_mean']:.4f}") + print(f" grad_norm: llm={last['grad_norm/llm']:.4f} transformer={last['grad_norm/transformer']:.4f}") + print(f" independent versions: llm={last['llm_weights_version']} dit={last['transformer_weights_version']}") + + print("\n== prompt-only (generator frozen; PromptRL ablation) ==") + _, p0b, lastb = _run(joint=False) + print(f" LM P(target action): {p0b:.3f} -> {lastb['refine_target_prob']:.3f}") + print(f" generator version : {lastb['transformer_weights_version']} (frozen; no dit_pg_loss key: " + f"{'dit_pg_loss' not in lastb})") + + +if __name__ == "__main__": + main() diff --git a/v2_examples/training/06_nway_joint_rl.py b/v2_examples/training/06_nway_joint_rl.py new file mode 100644 index 000000000..f8661b5b6 --- /dev/null +++ b/v2_examples/training/06_nway_joint_rl.py @@ -0,0 +1,49 @@ +#!/usr/bin/env python3 +"""06 — N-way joint RL over arbitrary experts. + +Generalizes joint RL from 2 experts to **N refiner LMs + a generator**, from one reward. The substrate +was already N-ready (a card holds N components/loops; WeightSyncPlan is per-component; grad-targets are +a dict) — only the method body loops over a list. Two credit modes: + + * `per_expert` — each refiner gets an advantage from its own reward term ⇒ all N learn cleanly. + * `shared` — one reward to everyone (faithful to UniRL) ⇒ works but noisier (multi-agent variance). + +Run: python3 v2_examples/training/06_nway_joint_rl.py +""" +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +from v2.models.multi_expert import build_multi_expert_card, refiner_ids +from v2.training.methods import build_joint_multi_rl + + +def _run(credit: str, n: int = 3, iters: int = 60): + rids = refiner_ids(n) + targets = {rid: (i * 2 + 1) % 6 for i, rid in enumerate(rids)} + m = build_joint_multi_rl(build_multi_expert_card(n_refiners=n), refiner_ids=rids, + target_actions=targets, joint_generator=True, credit=credit) + p0 = {r: float(m.refiner(r)._probs()[targets[r]]) for r in rids} + batch = {"prompts": ["a fox", "blue sky", "a city at night"], "seeds": [1, 2, 3]} + last = None + for it in range(iters): + _, last = m.train_step(batch, it) + return m, rids, targets, p0, last + + +def main() -> None: + for credit in ("per_expert", "shared"): + m, rids, targets, p0, last = _run(credit) + print(f"== {int(last['n_experts'])} experts ({len(rids)} refiners + generator), credit={credit} ==") + for r in rids: + print(f" {r} P(target): {p0[r]:.3f} -> {float(m.refiner(r)._probs()[targets[r]]):.3f}") + vers = [last[f"weights_version/{r}"] for r in rids] + [last["weights_version/transformer"]] + print(f" generator grad_norm={last['grad_norm/transformer']:.3f} " + f"distinct expert versions: {len(set(vers))}/{len(vers)}\n") + + +if __name__ == "__main__": + main() diff --git a/v2_examples/training/07_end_to_end_workflow_rl.py b/v2_examples/training/07_end_to_end_workflow_rl.py new file mode 100644 index 000000000..6ec28747d --- /dev/null +++ b/v2_examples/training/07_end_to_end_workflow_rl.py @@ -0,0 +1,49 @@ +#!/usr/bin/env python3 +"""07 — End-to-end RL over a cross-model workflow (T2I → I2V). + +Train BOTH stages of the T2I→I2V workflow (`flux-t2i` and `wan-i2v` — two *separate* models) from ONE +final-video reward. The rollout spans both instances with SDE capture; the same final-video advantage +drives a FlowGRPO PPO update on each stage's transformer, with two WeightSyncPlans on two instances. + +The point: the **earlier model (T2I) is trained by a reward on the final video** — end-to-end credit +across a model boundary. So "rollout == serve + capture" holds for a *workflow*, not just one card. + +Run: python3 v2_examples/training/07_end_to_end_workflow_rl.py +""" +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +import numpy as np + +from v2.models import build_flux_t2i_card, build_wan_i2v_card +from v2.training.methods import build_workflow_rl + + +def main() -> None: + m = build_workflow_rl(build_flux_t2i_card(), build_wan_i2v_card(), + num_samples_per_prompt=4, rollout_steps=4, t2i_lr=0.03, i2v_lr=0.03) + t2i_w0 = m.t2i.component("transformer").w_x.copy() + i2v_w0 = m.i2v.component("transformer").w_x.copy() + + batch = {"prompts": ["a fox", "blue sky", "a city"], "seeds": [1, 2, 3]} + last = None + for it in range(20): + _, last = m.train_step(batch, it) + + moved_t2i = not np.array_equal(t2i_w0, m.t2i.component("transformer").w_x) + moved_i2v = not np.array_equal(i2v_w0, m.i2v.component("transformer").w_x) + print("method:", m.name, "| two separate models trained from one final-video reward") + print(f" reward_mean : {last['reward_mean']:.4f}") + print(f" T2I generator moved: {moved_t2i} <- trained by a reward on the FINAL video (end-to-end credit)") + print(f" I2V generator moved: {moved_i2v}") + print(f" grad_norm : t2i={last['grad_norm/t2i']:.3f} i2v={last['grad_norm/i2v']:.3f} " + f"ppo_ratio={last['ppo_ratio_mean']:.4f}") + print(f" independent versions: t2i={last['t2i_weights_version']} i2v={last['i2v_weights_version']}") + + +if __name__ == "__main__": + main() diff --git a/v2_examples/training/README.md b/v2_examples/training/README.md new file mode 100644 index 000000000..c15380782 --- /dev/null +++ b/v2_examples/training/README.md @@ -0,0 +1,26 @@ +# v2_examples/training — training & RL on the *same* loops + +Runnable examples of every training method in `v2`. The load-bearing idea (designv4 §8): **a method's +rollout drives the same loop the engine serves** (in the ROLLOUT profile, with behavior capture) — there +is one numerics surface, not a second sampler. Each method differs only in loss math, roles, and capture. + +CPU + numpy only; the toy `ToyDiT.mse_grad_step` / `ToyPromptRefiner.pg_step` are *real* gradient steps, +so the toy students actually learn — the structural stand-in for an FSDP/optimizer step. Run from anywhere: + +```bash +python3 v2_examples/training/03_rl_diffusion_nft.py +``` + +| Script | Method | Kind | +|---|---|---| +| `01_finetune.py` | flow-match finetuning | supervised | +| `02_distillation_dmd2.py` | DMD2 (student + fake-score critic + teacher) | distillation | +| `03_rl_diffusion_nft.py` | DiffusionNFT (samples from the decay-blended *old* policy) | RL, likelihood-**free** (C2) | +| `04_self_forcing.py` | self-forcing on the causal `chunk_rollout` loop | causal distillation | +| `05_joint_lm_generator_rl.py` | UniRL/PromptRL: LM refiner + flow generator, one reward | joint RL, likelihood-**based** (C2) | +| `06_nway_joint_rl.py` | N refiner experts + generator, one reward | N-way joint RL | +| `07_end_to_end_workflow_rl.py` | T2I + I2V trained from one final-video reward | end-to-end RL across a workflow | + +Every method exposes the same entry point: `method.train_step(batch, iteration) -> (loss_map, metrics)`. +RL methods (`manages_optimization() == True`) own their sample→score→update cadence inside that call. +See designv4 §9.3 (joint RL), §9.7 (N-way), §9.9 (workflow RL). diff --git a/v2_examples/workflows/01_t2i_then_i2v.py b/v2_examples/workflows/01_t2i_then_i2v.py new file mode 100644 index 000000000..f1163cfbf --- /dev/null +++ b/v2_examples/workflows/01_t2i_then_i2v.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python3 +"""01 — Cross-model T2I → I2V workflow (FLUX→Wan-style). + +Two *separate* models chained by a `Workflow`: `flux-t2i` (text → image) → `wan-i2v` (text + image → +video). The workflow runs a full `engine.run` per stage and threads the stage-1 image artifact into the +stage-2 request as an `ImagePart`; the I2V program folds the image into its conditioning. Each model is +a distinct instance with its own weights — composition across model *instances*, not just loops within +one (designv4 §9.6). + +Run: python3 v2_examples/workflows/01_t2i_then_i2v.py +""" +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +import numpy as np + +from v2.models import build_image_video_engine, build_t2i_then_i2v_workflow + + +def main() -> None: + eng = build_image_video_engine() # registers flux-t2i + wan-i2v (+ the workflow) + wf = build_t2i_then_i2v_workflow() + print("workflow:", wf.workflow_id, "| requires cards:", wf.requires) + + out = wf.run(eng, prompt="a red fox in snow", seed=3) + video = np.asarray(out.artifacts["video"].frames) + print(f" T2I → I2V video: {video.shape}") + + # the I2V stage genuinely consumes the generated image: a different prompt → different image → + # different video (the hand-off is real, not cosmetic) + a = np.asarray(build_t2i_then_i2v_workflow().run(eng, prompt="alpha", seed=1).artifacts["video"].frames) + b = np.asarray(build_t2i_then_i2v_workflow().run(eng, prompt="beta gamma", seed=1).artifacts["video"].frames) + print(f" different prompt → different video: {not np.array_equal(a, b)}") + print(" (stage 2 is conditioned on stage 1's image — composition across two model instances)") + + +if __name__ == "__main__": + main() diff --git a/v2_examples/workflows/02_register_as_servable.py b/v2_examples/workflows/02_register_as_servable.py new file mode 100644 index 000000000..9565b59de --- /dev/null +++ b/v2_examples/workflows/02_register_as_servable.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +"""02 — A workflow is a first-class servable (naming & registration). + +A `Workflow` is named and registered with the same discipline as a card: a namespaced `workflow_id` +(`image_video.t2i_i2v`) in the SAME servable namespace as a `model_id`, with declared+validated +`requires`. Once registered it is addressable exactly like a model — `engine.serves(id)` reports it and +`engine.run(request)` routes to it (designv4 §9.6). This shows the per-engine registry, the fail-fast +validation, the declarative catalog (`register_workflows`), and the `WorkflowRegistry`. + +Run: python3 v2_examples/workflows/02_register_as_servable.py +""" +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +import numpy as np + +from v2.models import build_image_video_engine, build_t2i_then_i2v_workflow, register_workflows +from v2.program import WorkflowRegistry +from v2.request import DiffusionParams, TaskType, make_request +from v2.runtime import Engine + + +def main() -> None: + # build_image_video_engine registers the cards AND the workflow as a servable + eng = build_image_video_engine() + print("serves('image_video.t2i_i2v'):", eng.serves("image_video.t2i_i2v"), "(a workflow, addressable like a model)") + print("serves('flux-t2i'):", eng.serves("flux-t2i"), "(its required card)") + + # address the workflow BY ID — same call shape as any model + out = eng.run(make_request(TaskType.T2V, "image_video.t2i_i2v", "a lighthouse", + diffusion=DiffusionParams(seed=3))) + print("run-by-id video:", np.asarray(out.artifacts["video"].frames).shape) + + # declared + validated dependencies: registering on a bare engine fails fast + try: + Engine().register_workflow(build_t2i_then_i2v_workflow()) + print("FAIL: expected validation error") + except ValueError as e: + print("fail-fast validation:", "requires unregistered" in str(e)) + + # the declarative catalog + the standalone registry + fresh = register_workflows(Engine(), only=["image_video.t2i_i2v"]) # brings up cards + workflow + print("register_workflows → serves:", fresh.serves("image_video.t2i_i2v")) + reg = WorkflowRegistry() + reg.register("image_video.t2i_i2v", build_t2i_then_i2v_workflow) + print("WorkflowRegistry names:", reg.names(), "| requires:", reg.build("image_video.t2i_i2v").requires) + + +if __name__ == "__main__": + main() diff --git a/v2_examples/workflows/README.md b/v2_examples/workflows/README.md new file mode 100644 index 000000000..9f1ce3bd5 --- /dev/null +++ b/v2_examples/workflows/README.md @@ -0,0 +1,18 @@ +# v2_examples/workflows — cross-model composition + +Runnable examples of **cross-model `Workflow`s** — chaining two *distinct* models, threading typed +artifacts between them (designv4 §9.6). A `Program` composes the loops of ONE resident model +(step-interleaved, parity-gated); a `Workflow` composes *across* model instances, above the engine — so +the per-step hot path is unchanged and each model keeps its own interleave-parity guarantee. + +The shipped example is **T2I → I2V** (the realistic FLUX→Wan pipeline): a text-to-image model produces +an image, which conditions an image-to-video model. Run from anywhere: + +```bash +python3 v2_examples/workflows/01_t2i_then_i2v.py +``` + +| Script | Shows | +|---|---| +| `01_t2i_then_i2v.py` | run the cross-model workflow; the I2V stage provably consumes the generated image | +| `02_register_as_servable.py` | a workflow is a first-class servable — declared `requires`, validated, addressable by `workflow_id` like a model; the declarative catalog + `WorkflowRegistry` |