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.
77 lines
4.0 KiB
Python
77 lines
4.0 KiB
Python
"""LTX-2 two-stage distilled program (design_v3 §13, §15f):
|
||
|
||
text_encode → ltx2_base (8-step) → upsample(+noise mix) → ltx2_refine (3-step) → vae_decode
|
||
|
||
The two loop nodes bind the same ``transformer``; the upsample component sits between them and
|
||
mixes noise at σ₀=0.909375 before the refine stage (the repo's stage-2 noise injection).
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import numpy as np
|
||
|
||
from ...program import ComponentNode, ModelLoopNode, Program, ProgramKind
|
||
from ..common import text_encode_node_fn as _text_encode
|
||
from .loop import REFINE_SIGMAS
|
||
|
||
|
||
def _upsample(instance, slots, request, ctx) -> None:
|
||
base = np.asarray(slots["ltx_base_out"]["latents"], dtype="float32")
|
||
up = np.repeat(np.repeat(base, 2, axis=2), 2, axis=3) # 2× spatial (learned upsampler stand-in)
|
||
seed = (request.diffusion.seed if request.diffusion.seed is not None else 0) + 7
|
||
noise = np.random.default_rng(seed).standard_normal(up.shape).astype("float32")
|
||
sigma0 = float(REFINE_SIGMAS[0]) # 0.909375
|
||
slots["ltx_upsampled"] = (noise * sigma0 + up * (1.0 - sigma0)).astype("float32")
|
||
|
||
|
||
def _vae_decode(instance, slots, request, ctx) -> None:
|
||
slots["video"] = instance.component("vae").decode(slots["ltx_refine_out"]["latents"])
|
||
|
||
|
||
def build_ltx2_program() -> Program:
|
||
return Program(
|
||
program_id="ltx2.t2v.2stage", kind=ProgramKind.INLINE,
|
||
nodes=[
|
||
ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")),
|
||
ModelLoopNode("base", loop_id="ltx2_base", output_slot="ltx_base_out",
|
||
reads=("text_embeds",), writes=("ltx_base_out",)),
|
||
ComponentNode("upsample", fn=_upsample, reads=("ltx_base_out",), writes=("ltx_upsampled",)),
|
||
ModelLoopNode("refine", loop_id="ltx2_refine", output_slot="ltx_refine_out",
|
||
reads=("ltx_upsampled",), writes=("ltx_refine_out",)),
|
||
ComponentNode("vae_decode", fn=_vae_decode, reads=("ltx_refine_out",), writes=("video",)),
|
||
],
|
||
output_artifacts={"video": "video", "latents": "ltx_refine_out"},
|
||
).validate()
|
||
|
||
|
||
# --- joint audio+video (T2VS) program (design_v3 §15b; §9.11) ---------------------- #
|
||
def _upsample_av(instance, slots, request, ctx) -> None:
|
||
_upsample(instance, slots, request, ctx) # video latent upsample (reuse)
|
||
base = slots.get("ltx_base_out", {})
|
||
slots["ltx_audio"] = base.get("audio_latents") if isinstance(base, dict) else None # thread audio
|
||
|
||
|
||
def _audio_decode(instance, slots, request, ctx) -> None:
|
||
ro = slots.get("ltx_refine_out", {})
|
||
au = ro.get("audio_latents") if isinstance(ro, dict) else None
|
||
slots["audio"] = instance.component("audio_vae").decode(au) if au is not None else None
|
||
|
||
|
||
def build_ltx2_av_program() -> Program:
|
||
"""T2VS: the two-stage denoise carries a synchronized audio latent through both stages; the video
|
||
VAE and the audio VAE decode the two modalities → video + audio artifacts (per-modality guidance)."""
|
||
return Program(
|
||
program_id="ltx2.t2vs.2stage", kind=ProgramKind.INLINE,
|
||
nodes=[
|
||
ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")),
|
||
ModelLoopNode("base", loop_id="ltx2_base", output_slot="ltx_base_out",
|
||
reads=("text_embeds",), writes=("ltx_base_out",)),
|
||
ComponentNode("upsample", fn=_upsample_av, reads=("ltx_base_out",),
|
||
writes=("ltx_upsampled", "ltx_audio")),
|
||
ModelLoopNode("refine", loop_id="ltx2_refine", output_slot="ltx_refine_out",
|
||
reads=("ltx_upsampled", "ltx_audio"), writes=("ltx_refine_out",)),
|
||
ComponentNode("vae_decode", fn=_vae_decode, reads=("ltx_refine_out",), writes=("video",)),
|
||
ComponentNode("audio_decode", fn=_audio_decode, reads=("ltx_refine_out",), writes=("audio",)),
|
||
],
|
||
output_artifacts={"video": "video", "audio": "audio", "latents": "ltx_refine_out"},
|
||
).validate()
|