SolitaryThinker
9489c6c1dd
The remaining design_v3 probes (all except 8b speculative decoding). All fit with no new runtime primitive (148 tests pass). #6 LTX-2 joint audio+video (§9.11) — LTX-2 declared an audio_vae required_for t2vs but never used it; now a single 2-stage denoise carries a synchronized audio latent (conditioned on video), applies per-modality CFG (guidance_per_modality), and decodes via video VAE + audio VAE → video + audio. Gated on requesting audio, so the T2V path is byte-identical (existing tests untouched). ToyAudioVAE; build_ltx2_av_program. 5 tests. #4 Live weight-sync under in-flight serving (§9.14) — WeightSyncController makes the freeze → drain → transfer → bump version + invalidate → resume lifecycle explicit. Tests: a mid-flight swap corrupts (the hazard); draining first leaves the in-flight request bit-identical to baseline while a post-sync request reflects new weights; transformer-only sync, so the frozen text-encoder cache survives. The RL flywheel's hardest correctness. 3 tests. #5 Reward-model-as-a-served-card (§9.15) — a reward model is a card (scorer + a score loop emitting REWARD_BATCH units); ServedRewardScorer drop-in-replaces the numpy scorer so any RL method becomes RLHF/RLAIF with no method change. ToyRewardModel; models/reward/. 4 tests. #7 Content-adaptive control flow (§9.12) — CacheDiTDenoiseLoop (isolated WanDenoiseLoop subclass) reuses the cached velocity when predictions barely change (cache-dit skip) and early-exits on convergence — variable step count; interleave parity holds across ragged loops. models/adaptive/. 4 tests. #8a Nested workflows (§9.13) — a workflow stage can invoke another workflow (engine.run routes ids); requires/validate recurse; cycles caught at registration + a run-time guard (engine._wf_running). build_t2i_i2v_extend_workflow. 5 tests. designv4 §9.11–§9.15 + falsifier/layout/closing updates (148 tests, 25 files). Also removed a pre-existing unused import in ltx2/loop.py.
82 lines
4.2 KiB
Python
82 lines
4.2 KiB
Python
"""CacheDiTDenoiseLoop — the loop OWNS content-adaptive control flow (design_v3 §2.2).
|
||
|
||
The signature payoff of loop inversion: because the *model* owns control flow, content-adaptive
|
||
decisions are natural. This loop, at each step, decides whether to:
|
||
|
||
* **skip the expensive forward** and reuse the cached velocity when consecutive predictions barely
|
||
change (cache-dit / TeaCache / Δ-DiT — a 1.5–2× inference win), or
|
||
* **early-exit** when the latent has converged.
|
||
|
||
The result is a *variable, content-dependent* step count — impossible in a ``for t in timesteps``
|
||
runtime, trivial in a driven loop whose ``next()`` reads ``LoopState``. It is fully isolated from the
|
||
shared ``WanDenoiseLoop`` (subclass, threshold 0 ⇒ identical behavior), and still interleave-safe:
|
||
different requests skip different steps, and the parity gate still holds (ragged loops don't smear).
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import numpy as np
|
||
|
||
from ...loop.contracts import Done, StepResult
|
||
from ...loop.sampler import flow_match_euler_step
|
||
from ..wan21.loop import WanDenoiseLoop
|
||
|
||
|
||
class CacheDiTDenoiseLoop(WanDenoiseLoop):
|
||
def __init__(self, *, cache_threshold: float = 0.0, exit_threshold: float = 0.0, **kw):
|
||
super().__init__(**kw)
|
||
self.cache_threshold = cache_threshold # rel velocity change below which we reuse (skip the DiT)
|
||
self.exit_threshold = exit_threshold # rel latent change below which we stop early
|
||
|
||
def init(self, req, model, ctx):
|
||
st = super().init(req, model, ctx)
|
||
st.scratch.update(last_v=None, skip_next=False, skipped=0, exited=False)
|
||
return st
|
||
|
||
def next(self, st):
|
||
if st.scratch.get("exited"):
|
||
return Done() # converged early — content-adaptive termination
|
||
plan = super().next(st) # the normal full-forward step plan
|
||
last_v = st.scratch.get("last_v")
|
||
if isinstance(plan, Done) or not self.cache_threshold or not st.scratch.get("skip_next") \
|
||
or last_v is None:
|
||
return plan
|
||
# SKIP: reuse the cached velocity, no DiT forward (the cache-dit win) — wrap the run thunk
|
||
x, prec = st.latents["video"], self.precision
|
||
sigma_t, sigma_next = st.sigmas[st.step_idx], st.sigmas[st.step_idx + 1]
|
||
v = np.asarray(last_v, dtype="float32")
|
||
|
||
def run(model, override=None):
|
||
x_next = flow_match_euler_step(prec.cast(x), v, sigma_t, sigma_next)
|
||
return StepResult(output={"noise_pred": v, "latents": x_next.astype("float32"), "skipped": True})
|
||
|
||
plan.run = run
|
||
plan.label = f"{plan.label}.skip"
|
||
return plan
|
||
|
||
def advance(self, st, result):
|
||
x_prev = np.asarray(st.latents["video"], dtype="float64")
|
||
v = np.asarray(result.output["noise_pred"], dtype="float64")
|
||
prev_v = st.scratch.get("last_v")
|
||
st = super().advance(st, result) # folds latents, step_idx, (rollout) trajectory
|
||
if result.output.get("skipped"):
|
||
st.scratch["skipped"] += 1
|
||
st.scratch["skip_next"] = False # don't chain skips — re-check with a full forward
|
||
elif self.cache_threshold and prev_v is not None:
|
||
rel = float(np.linalg.norm(v - np.asarray(prev_v, dtype="float64"))
|
||
/ (np.linalg.norm(prev_v) + 1e-8))
|
||
st.scratch["skip_next"] = rel < self.cache_threshold # next prediction will barely change
|
||
st.scratch["last_v"] = result.output["noise_pred"]
|
||
half = max(1, (len(st.sigmas) - 1) // 2) # only consider exiting in the low-noise tail
|
||
if self.exit_threshold and st.step_idx >= half:
|
||
x_next = np.asarray(st.latents["video"], dtype="float64")
|
||
dl = float(np.linalg.norm(x_next - x_prev) / (np.linalg.norm(x_prev) + 1e-8))
|
||
if dl < self.exit_threshold:
|
||
st.scratch["exited"] = True # latent converged → stop next tick
|
||
return st
|
||
|
||
def finalize(self, st):
|
||
res = super().finalize(st)
|
||
res.metrics["skipped_steps"] = float(st.scratch.get("skipped", 0))
|
||
res.metrics["early_exited"] = 1.0 if st.scratch.get("exited") else 0.0
|
||
return res
|