Files
hao-ai-lab-FastVideo/v2/models/omni/pack.py
T
SolitaryThinker a9e5f6ee7a [refactor] rename package mini_fastvideo → v2
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
2026-07-04 17:15:34 +00:00

46 lines
1.9 KiB
Python

"""Omni program nodes shared by Cosmos3 and the bagel/lance MoT (design_v3 §4.1, §15b).
The ``pack`` node is the seam where the AR (understanding) pathway's tokens become the
conditioning the diffusion (generation) pathway consumes — "the AR reasoner ... used as prompt
upsampling before diffusion in the same request". Both pathways are the SAME resident MoT
module, so this is weight-sharing, not a hand-off between two engines.
"""
from __future__ import annotations
from ..backend import ToyTokenizer
def tokenize_node(instance, slots, request, ctx) -> None:
tok = instance.component("tokenizer") if instance.has_component("tokenizer") else ToyTokenizer()
slots["prompt_tokens"] = tok.encode(request.prompt())
def _tokens_of(slot_value) -> list:
if isinstance(slot_value, dict):
return slot_value.get("tokens", [])
return list(slot_value or [])
def pack_cond_from_tokens(token_slot: str):
"""Build a ComponentNode fn that packs the und-pathway tokens into the gen-pathway conditioning."""
def fn(instance, slots, request, ctx) -> None:
mot = instance.component("transformer") # the SHARED MoT module (und + gen pathways)
toks = _tokens_of(slots.get(token_slot))
slots["text_embeds"] = mot.reasoner_embed(toks)
slots["neg_text_embeds"] = mot.reasoner_embed([0])
return fn
def emit_text_node(token_slot: str, out_slot: str = "text"):
def fn(instance, slots, request, ctx) -> None:
out = slots.get(token_slot, {})
slots[out_slot] = out.get("text", "") if isinstance(out, dict) else str(out)
return fn
def vae_decode_node(latent_slot: str, out_slot: str):
def fn(instance, slots, request, ctx) -> None:
latents = slots[latent_slot]["latents"] if isinstance(slots.get(latent_slot), dict) else None
slots[out_slot] = instance.component("vae").decode(latents) if latents is not None else None
return fn