SolitaryThinker
4e331eb7ff
Mechanical conversion of every relative import under v2/ to an absolute v2.* path (from .x / ..x / ...x -> from v2.<pkg>.x) so imports are unambiguous, grep-able, and stable when code is copied/moved between entrypoints (VideoGenerator, CLI, server). Surgical prefix-only rewrite: only the 'from <dots><module>' prefix changed — import names, parentheses, multi-line formatting, comments, and ordering are byte-for-byte preserved (no collapsing, no reorder, no unrelated reformatting). - 471 imports across 125 files; v2/tests/ was already absolute (untouched). - Validated: all 125 files compile, every 'from v2.* import' target resolves to a real module/package, zero relative imports remain (full sweep), CPU mini suite green (202 passed, 2 skipped).
133 lines
7.1 KiB
Python
133 lines
7.1 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 v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind
|
||
from v2.models.common import text_encode_node_fn as _text_encode
|
||
from v2.models.ltx2.loop import REFINE_SIGMAS
|
||
|
||
|
||
def _upsample(instance, slots, request, ctx) -> None:
|
||
base = np.asarray(slots["ltx_base_out"]["latents"], dtype="float32")
|
||
# Learned 2× spatial latent super-resolution between the base and refine stages. The
|
||
# ``spatial_upsampler`` component is the real LTX2LatentUpsampler on the GPU backend (un_normalize →
|
||
# learned upsample → normalize, via the VAE's per-channel stats) and a nearest-neighbor toy on CPU —
|
||
# same call on both backends (the v2 component-polymorphism seam), no device branch here.
|
||
up = np.asarray(instance.component("spatial_upsampler").upsample(base), dtype="float32")
|
||
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()
|
||
|
||
|
||
# --- single-stage LTX-2 base model (no distill, no upsampler) ---------------------- #
|
||
def _vae_decode_single(instance, slots, request, ctx) -> None:
|
||
slots["video"] = instance.component("vae").decode(slots["ltx_out"]["latents"])
|
||
|
||
|
||
def build_ltx2_base_program() -> Program:
|
||
"""Single-stage LTX-2 base: text_encode → denoise (request-driven many-step, full-res) → vae_decode.
|
||
No upsample/refine (those are the distilled two-stage program above)."""
|
||
return Program(
|
||
program_id="ltx2.t2v.base", kind=ProgramKind.INLINE,
|
||
nodes=[
|
||
ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")),
|
||
ModelLoopNode("denoise", loop_id="ltx2_single", output_slot="ltx_out",
|
||
reads=("text_embeds",), writes=("ltx_out",)),
|
||
ComponentNode("vae_decode", fn=_vae_decode_single, reads=("ltx_out",), writes=("video",)),
|
||
],
|
||
output_artifacts={"video": "video", "latents": "ltx_out"},
|
||
).validate()
|
||
|
||
|
||
# --- single-stage LTX-2.3 joint A/V (T2VS) ----------------------------------------- #
|
||
def _text_encode_av(instance, slots, request, ctx) -> None:
|
||
"""LTX-2.3 connector emits SEPARATE video + audio text projections (one Gemma call)."""
|
||
video_text, audio_text = instance.component("text_encoder").encode_av(request.prompt())
|
||
slots["text_embeds"] = video_text
|
||
slots["audio_text_embeds"] = audio_text
|
||
|
||
|
||
def _audio_decode_single(instance, slots, request, ctx) -> None:
|
||
out = slots.get("ltx_out", {})
|
||
au = out.get("audio_latents") if isinstance(out, dict) else None
|
||
slots["audio"] = instance.component("audio_vae").decode(au) if au is not None else None
|
||
|
||
|
||
def build_ltx2_3_program() -> Program:
|
||
"""LTX-2.3 single-stage T2VS: text_encode (separate video+audio connectors) → JOINT A/V denoise
|
||
(one DiT forward per step cross-attends video<->audio) → video decode + audio decode
|
||
(AudioDecoder→Vocoder). A plain T2V request leaves the audio latent unproduced, so audio_decode
|
||
writes None and the audio components stay unbuilt."""
|
||
return Program(
|
||
program_id="ltx2.3.t2vs", kind=ProgramKind.INLINE,
|
||
nodes=[
|
||
ComponentNode("text_encode", fn=_text_encode_av, writes=("text_embeds", "audio_text_embeds")),
|
||
ModelLoopNode("denoise", loop_id="ltx2_3", output_slot="ltx_out",
|
||
reads=("text_embeds", "audio_text_embeds"), writes=("ltx_out",)),
|
||
ComponentNode("vae_decode", fn=_vae_decode_single, reads=("ltx_out",), writes=("video",)),
|
||
ComponentNode("audio_decode", fn=_audio_decode_single, reads=("ltx_out",), writes=("audio",)),
|
||
],
|
||
output_artifacts={"video": "video", "audio": "audio", "latents": "ltx_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()
|