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).
48 lines
2.5 KiB
Python
48 lines
2.5 KiB
Python
"""Unified LM+generator serve program (design_v3 §15b; UniRL/PromptRL inference path):
|
|
|
|
text_encode → refine (ar_decode on llm) → apply_refinement → denoise (diffusion_denoise) → vae_decode
|
|
|
|
The serve path is deterministic: the LM refines greedily (argmax action) and the generator runs the
|
|
ODE Euler sampler. The ``unified_rl`` *training* method drives the same two loops but in the ROLLOUT
|
|
profile — the LM *samples* an action and the generator uses the SDE sampler with per-step log-prob
|
|
capture (design_v3 §9.4: rollout sampler ≠ serve sampler, a controlled divergence). Same loops, two
|
|
profiles — the §10 "train ≡ serve through one loop" invariant, now spanning two loop *types*.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind
|
|
from v2.models.common import text_encode_node_fn
|
|
from v2.models.omni import vae_decode_node
|
|
from v2.models.unified.card import N_REFINE_ACTIONS
|
|
|
|
|
|
def apply_refinement_node(refine_slot: str, n_actions: int = N_REFINE_ACTIONS):
|
|
"""ComponentNode fn: the LM's chosen refinement action reshapes the diffusion conditioning."""
|
|
def fn(instance, slots, request, ctx) -> None:
|
|
llm = instance.component("llm")
|
|
out = slots.get(refine_slot)
|
|
toks = out.get("tokens", []) if isinstance(out, dict) else []
|
|
action = (int(toks[0]) % n_actions) if toks else 0
|
|
slots["text_embeds"] = llm.refined_embed(slots.get("text_embeds"), action)
|
|
slots["refine_action"] = action
|
|
return fn
|
|
|
|
|
|
def build_unified_program() -> Program:
|
|
return Program(
|
|
program_id="unified.unirl", kind=ProgramKind.INLINE,
|
|
nodes=[
|
|
ComponentNode("text_encode", fn=text_encode_node_fn,
|
|
writes=("text_embeds", "neg_text_embeds")),
|
|
ModelLoopNode("refine", loop_id="ar_decode", output_slot="refine_out",
|
|
reads=("text_embeds",), writes=("refine_out",)),
|
|
ComponentNode("apply_refine", fn=apply_refinement_node("refine_out"),
|
|
reads=("refine_out", "text_embeds"), writes=("text_embeds", "refine_action")),
|
|
ModelLoopNode("denoise", loop_id="diffusion_denoise", output_slot="denoise_out",
|
|
reads=("text_embeds",), writes=("denoise_out",)),
|
|
ComponentNode("vae_decode", fn=vae_decode_node("denoise_out", "video"),
|
|
reads=("denoise_out",), writes=("video",)),
|
|
],
|
|
output_artifacts={"video": "video", "latents": "denoise_out"},
|
|
).validate()
|