Files
hao-ai-lab-FastVideo/v2/models/cosmos3/program.py
T
SolitaryThinker 4e331eb7ff [refactor] v2: use absolute imports (v2.*) everywhere instead of relative
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).
2026-07-04 17:15:34 +00:00

33 lines
1.7 KiB
Python

"""Cosmos3 omni program (design_v3 §15b):
tokenize → reason (ar_decode) → pack(tokens→cond) → diffusion_denoise → vae_decode
The reasoner and the denoiser bind the SAME resident ``transformer``; the scheduler is the mode
multiplexer. This is the workload no DAG-of-engines can express (splitting them doubles the weights
and severs the shared state).
"""
from __future__ import annotations
from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind
from v2.models.omni import emit_text_node, pack_cond_from_tokens, tokenize_node, vae_decode_node
def build_cosmos3_program() -> Program:
return Program(
program_id="cosmos3.omni", kind=ProgramKind.INLINE,
nodes=[
ComponentNode("tokenize", fn=tokenize_node, writes=("prompt_tokens",)),
ModelLoopNode("reason", loop_id="ar_decode", output_slot="reason_out",
reads=("prompt_tokens",), writes=("reason_out",)),
ComponentNode("emit_text", fn=emit_text_node("reason_out", "reason_text"),
reads=("reason_out",), writes=("reason_text",)),
ComponentNode("pack", fn=pack_cond_from_tokens("reason_out"),
reads=("reason_out",), writes=("text_embeds", "neg_text_embeds")),
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", "text": "reason_text", "latents": "denoise_out"},
).validate()