SolitaryThinker
f1dc587c74
One resident MoT instance runs BOTH an ar_decode loop and a diffusion_denoise loop on shared weights (the §16 claim no DAG-of-engines can express), with both loops runtime-visible: the scheduler prices ar_token AND diffusion_step WorkUnits — unlike vllm-omni's opaque DIFFUSION stage the scheduler never sees inside. - ARDecodeLoop: token decode until EOS/max_tokens, paged text-KV — the omni AR pathway (loop/§5). - ToyMoTDiT: one module exposing an und pathway (ar_forward) AND a gen pathway (denoise __call__); ToyTokenizer. Binding both loops to one instance = shared weights, no duplication. - models/cosmos3/: tokenize → reason(ar_decode) → pack(tokens→conditioning) → diffusion_denoise → vae_decode; sound_vae declared optional_for non-t2vs (the lazy-component P8 fix, not an env-var hack). - models/bagel/: the canonical vllm-omni model — generate_text(ar_decode) → generate_image(diffusion), text+image outputs, both loops step-scheduled. - The diffusion loop is WanDenoiseLoop reused (one loop definition bound to the MoT module). - build_omni_engine() + an omni worked example; 7 omni tests (shared-instance, both-kinds-scheduled, interleave parity across loop types, lazy sound_vae). 61 tests pass (pytest + zero-dep runner).
32 lines
1.7 KiB
Python
32 lines
1.7 KiB
Python
"""BAGEL/lance program (design_v3 §16): one resident MoT instance, two runtime-visible loops.
|
|
|
|
tokenize → generate_text (ar_decode) → pack(text→cond) → generate_image (diffusion) → vae_decode
|
|
|
|
generate_text and generate_image hit the SAME resident weights; their steps are WorkUnits the
|
|
scheduler interleaves — the shared-weight MoT made native, not a monolith bypassing the abstraction.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from ...program import ComponentNode, ModelLoopNode, Program, ProgramKind
|
|
from ..omni import emit_text_node, pack_cond_from_tokens, tokenize_node, vae_decode_node
|
|
|
|
|
|
def build_bagel_program() -> Program:
|
|
return Program(
|
|
program_id="bagel.mot", kind=ProgramKind.INLINE,
|
|
nodes=[
|
|
ComponentNode("tokenize", fn=tokenize_node, writes=("prompt_tokens",)),
|
|
ModelLoopNode("generate_text", loop_id="generate_text", output_slot="gen_text_out",
|
|
reads=("prompt_tokens",), writes=("gen_text_out",)),
|
|
ComponentNode("emit_text", fn=emit_text_node("gen_text_out", "text"),
|
|
reads=("gen_text_out",), writes=("text",)),
|
|
ComponentNode("pack", fn=pack_cond_from_tokens("gen_text_out"),
|
|
reads=("gen_text_out",), writes=("text_embeds", "neg_text_embeds")),
|
|
ModelLoopNode("generate_image", loop_id="generate_image", output_slot="gen_image_out",
|
|
reads=("text_embeds",), writes=("gen_image_out",)),
|
|
ComponentNode("vae_decode", fn=vae_decode_node("gen_image_out", "image"),
|
|
reads=("gen_image_out",), writes=("image",)),
|
|
],
|
|
output_artifacts={"text": "text", "image": "image"},
|
|
).validate()
|