SolitaryThinker
431f4daddb
Three more capabilities on distinct untested surfaces (167 tests pass, no new runtime primitive). A — Adapter plane (§9.19): one base + swappable LoRA/ControlNet adapters, selected per request (DiffusionParams.adapters); AdapterDenoiseLoop applies each active adapter's velocity delta. Per-request selection changes output, multi-LoRA composes, ControlNet conditions on a control image, mixed-adapter requests interleave without smearing, hot-swap changes generation, cache key partitions by adapter stack (the adapter_versions field, previously declared-only). ToyLoRA/ToyControlNet; models/adapters/. 6 tests. B — Non-linear workflows (§9.17): ParallelWorkflow (fan-out: one input → N models → merged) and BestOfNWorkflow (generate N → score with the served reward card → return best; inference-time scaling). The shapes a linear chain can't express. 4 tests. D — RL→distill flywheel (§9.18): run_flywheel RL-improves the base (NFT), then distills FROM the RL'd model (DMD2 teacher = RL'd policy) into a faster card, recording the base→rl→distilled provenance chain in RecipeSpec.parents. The distilled student is measurably closer to the RL'd teacher than the base; the distilled card serves few-step. training/flywheel.py. 4 tests. designv4 §9.17–§9.19 + layout/closing/counts (167 tests, 29 files).
74 lines
3.4 KiB
Python
74 lines
3.4 KiB
Python
"""Non-linear workflow shapes — fan-out + best-of-N feedback (design_v3 §13; §9.17).
|
|
|
|
Every workflow so far was a linear chain. These exercise the two non-linear shapes:
|
|
|
|
* **Fan-out** (`ParallelWorkflow`): one input → N models in parallel → merged result (namespaced by
|
|
branch). The engine can interleave the branches' steps.
|
|
* **Best-of-N** (`BestOfNWorkflow`): generate N candidates, score each with a reward scorer, return
|
|
the best — inference-time scaling / rejection sampling, a feedback loop composing a generator with
|
|
the served REWARD_BATCH card (§9.15).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from v2.models import build_default_engine, build_reward_card
|
|
from v2.program import BestOfNWorkflow, ParallelWorkflow, WorkflowStage
|
|
from v2.request import DiffusionParams, TaskType, make_request
|
|
from v2.training.methods.base import new_instance
|
|
from v2.training.rewards import ServedRewardScorer
|
|
|
|
|
|
def _t2v(model_id):
|
|
return lambda s: make_request(TaskType.T2V, model_id, s["prompt"],
|
|
diffusion=DiffusionParams(num_steps=4, seed=s.get("seed", 0)))
|
|
|
|
|
|
def test_fan_out_runs_branches_and_merges():
|
|
eng = build_default_engine()
|
|
wf = ParallelWorkflow("variants.fanout", [
|
|
WorkflowStage("wan2.1-1.3b", _t2v("wan2.1-1.3b"), label="wan"),
|
|
WorkflowStage("ltx2.3-distilled", _t2v("ltx2.3-distilled"), label="ltx"),
|
|
])
|
|
out = wf.run(eng, prompt="a comet over the sea", seed=1)
|
|
# both branches' artifacts present, namespaced by label
|
|
assert "wan:video" in out.artifacts and "ltx:video" in out.artifacts
|
|
wan_v = np.asarray(out.artifacts["wan:video"].frames)
|
|
ltx_v = np.asarray(out.artifacts["ltx:video"].frames)
|
|
assert wan_v.ndim == 4 and ltx_v.ndim == 4
|
|
# two different models on the same prompt → different videos (genuine fan-out, not duplication)
|
|
assert wan_v.shape != ltx_v.shape or not np.array_equal(wan_v, ltx_v)
|
|
assert wf.requires == ["wan2.1-1.3b", "ltx2.3-distilled"]
|
|
|
|
|
|
def _best_of_n(n=4):
|
|
eng = build_default_engine()
|
|
scorer = ServedRewardScorer(new_instance(build_reward_card(batch=2)))
|
|
wf = BestOfNWorkflow("wan.best_of_n", WorkflowStage("wan2.1-1.3b", _t2v("wan2.1-1.3b")),
|
|
scorer=scorer, n=n, score_key="latents")
|
|
return eng, scorer, wf
|
|
|
|
|
|
def test_best_of_n_returns_the_highest_scored_candidate():
|
|
eng, scorer, wf = _best_of_n(n=4)
|
|
out = wf.run(eng, prompt="a fox", seed=2)
|
|
assert out.metrics["best_of_n"] == 4.0
|
|
# re-score the n candidates independently and confirm the returned one is the argmax
|
|
cands = [eng.run(make_request(TaskType.T2V, "wan2.1-1.3b", "a fox",
|
|
diffusion=DiffusionParams(num_steps=4, seed=2 * 100 + i))) for i in range(4)]
|
|
scores = scorer.score([np.asarray(c.artifacts["latents"].latent) for c in cands], ["a fox"] * 4)["avg"]
|
|
assert out.metrics["best_score"] == float(np.max(scores))
|
|
assert out.metrics["best_index"] == float(int(np.argmax(scores)))
|
|
|
|
|
|
def test_best_of_n_is_deterministic():
|
|
eng, _, wf = _best_of_n(n=3)
|
|
a = wf.run(eng, prompt="a city", seed=5).metrics["best_index"]
|
|
b = wf.run(eng, prompt="a city", seed=5).metrics["best_index"]
|
|
assert a == b
|
|
|
|
|
|
def test_best_of_n_requires_only_the_generator():
|
|
_, _, wf = _best_of_n()
|
|
assert wf.requires == ["wan2.1-1.3b"] # the reward scorer is a held instance, not a chained stage
|