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).
48 lines
2.1 KiB
Python
48 lines
2.1 KiB
Python
"""AdapterDenoiseLoop — per-request LoRA / ControlNet over one resident base (design_v3 §9.19).
|
|
|
|
The adapter plane: a request selects adapters (``DiffusionParams.adapters``); the loop applies each
|
|
active adapter's velocity delta on top of the base DiT's prediction, then re-integrates. Many adapters
|
|
are served over ONE resident base; the active set lives in the request/``LoopState`` (never global), so
|
|
two requests using different adapters interleave without smearing. Isolated ``WanDenoiseLoop`` subclass
|
|
(empty adapter set ⇒ identical to the base), so nothing else is touched.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from ...loop.contracts import Done
|
|
from ...loop.sampler import flow_match_euler_step
|
|
from ..wan21.loop import WanDenoiseLoop
|
|
|
|
|
|
class AdapterDenoiseLoop(WanDenoiseLoop):
|
|
def init(self, req, model, ctx):
|
|
st = super().init(req, model, ctx)
|
|
st.scratch["adapters"] = tuple(getattr(req.diffusion, "adapters", ()) or ())
|
|
st.scratch["control"] = ctx.slots.get("control_signal") # ControlNet conditioning (optional)
|
|
return st
|
|
|
|
def next(self, st):
|
|
plan = super().next(st)
|
|
adapters = st.scratch.get("adapters") or ()
|
|
if isinstance(plan, Done) or not adapters:
|
|
return plan
|
|
base_run = plan.run
|
|
x = st.latents["video"]
|
|
control = st.scratch.get("control")
|
|
sigma_t, sigma_next = st.sigmas[st.step_idx], st.sigmas[st.step_idx + 1]
|
|
prec = self.precision
|
|
|
|
def run(model, override=None):
|
|
res = base_run(model, override) # base velocity (CFG combine)
|
|
v = np.asarray(res.output["noise_pred"], dtype="float32")
|
|
for aid in adapters: # apply each active adapter's delta
|
|
v = v + model.component(aid).delta(x, control)
|
|
res.output["noise_pred"] = v
|
|
res.output["latents"] = flow_match_euler_step(prec.cast(x), v, sigma_t, sigma_next).astype("float32")
|
|
return res
|
|
|
|
plan.run = run
|
|
plan.label = f"{plan.label}.adapt[{'+'.join(adapters)}]"
|
|
return plan
|