SolitaryThinker
a9e5f6ee7a
Directory rename (git mv, history preserved) plus rewrite of all references: absolute imports in tests, the zero-dep runner, docstrings, comments, and the README. No behavior change. Run: python3 -m pytest v2/tests/ -q ; python3 v2/run_tests.py ; python3 -m v2.examples
24 lines
1.1 KiB
Python
24 lines
1.1 KiB
Python
"""Wan-causal streaming program: text_encode → chunk_rollout → vae_decode (design_v3 §4, §15e)."""
|
|
from __future__ import annotations
|
|
|
|
from ...program import ComponentNode, ModelLoopNode, Program, ProgramKind
|
|
from ..common import text_encode_node_fn as _text_encode
|
|
|
|
|
|
def _vae_decode(instance, slots, request, ctx) -> None:
|
|
latents = slots["rollout_out"]["latents"]
|
|
slots["video"] = instance.component("vae").decode(latents) if latents is not None else None
|
|
|
|
|
|
def build_wan_causal_program() -> Program:
|
|
return Program(
|
|
program_id="wan_causal.stream", kind=ProgramKind.INLINE,
|
|
nodes=[
|
|
ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")),
|
|
ModelLoopNode("rollout", loop_id="chunk_rollout", output_slot="rollout_out",
|
|
reads=("text_embeds",), writes=("rollout_out",)),
|
|
ComponentNode("vae_decode", fn=_vae_decode, reads=("rollout_out",), writes=("video",)),
|
|
],
|
|
output_artifacts={"video": "video", "latents": "rollout_out"},
|
|
).validate()
|