SolitaryThinker
634f0828a2
Finish the kernel seam across the board so the platform's KernelTable is the universal solver-dispatch path, not just wan21. - Loops: ltx2, wan_causal, adapters, adaptive now resolve the flow-match solver via model.platform.kernels.get(FLOW_MATCH_STEP) instead of importing the numpy sampler directly (wan21 already did). On CPU this is bit-identical (the cpu kernel IS the old function); a GPU/accel backend now overrides every loop. - RL: the FlowGRPO log-prob recompute in unified_rl / joint_multi_rl / workflow_rl dispatches FLOW_SDE_STEP through the platform, pinned to the SAME kernel the rollout used (C2 kernel-pinning — otherwise the PPO ratio biases on a real GPU where rollout and recompute kernels could differ). - accel backend: add a kind-generic AccelComponent wrapper and override the vae component too (text_encoder left unregistered to keep the device->cpu fallback demonstrated), closing the "only dit is overridden" gap. - tests: vae override assertion; a second-loop (wan-causal chunk rollout) parity oracle proving accel == cpu bit-identical beyond wan21. README scope updated. 184 tests pass.
49 lines
2.2 KiB
Python
49 lines
2.2 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 ...platform import FLOW_MATCH_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
|
|
fm = model.platform.kernels.get(FLOW_MATCH_STEP) # solver dispatched per (device, arch)
|
|
res.output["latents"] = fm(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
|