nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[misc] v2: simplify docstrings/comments + drop deleted-design-doc citations

Browse Source 
Sweep all 296 v2 modules: simplify verbose docstrings/comments and remove 376 dangling
"(design_vN §X)" citations to the now-deleted design docs (v2/README.md is the source of
truth). Comment/docstring-only — AST-verified code-identical; the CPU suite holds at 237
passed / 5 pre-existing. Also applies yapf + ruff --fix auto-fixes (import ordering,
forward-ref annotation de-quoting under `from __future__ import annotations`; behavior-
neutral, suite-confirmed) and adds the legitimate domain terms mot/clen/te to the codespell
ignore-list. Remaining ruff (24) + mypy (68 no-untyped-call) findings are pre-existing v2
debt, untouched here.
This commit is contained in:
SolitaryThinker
2026-06-19 20:54:59 +00:00
parent 53d63afa76
commit db866dfb04
234 changed files with 4148 additions and 3080 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pyproject.toml
+3 -2
View File
@@ -223,8 +223,9 @@ follow_imports = "silent"
skip = "./data,./wandb,ui/package-lock.json,performance_dashboard/frontend/package-lock.json,*/_vendored/*"
# "tread" matches daVinci-MagiHuman's acronym "TReAD" (Token Routing and
# Early Drop). codespell lowercases ignore-words entries, so the single
# lowercase form silences all case variants.
ignore-words-list = "tread,passt"
# lowercase form silences all case variants. "mot" = Mixture-of-Transformers
# (MoT); "clen" = a Content-Length local; "te" = a text-embeds local.
ignore-words-list = "tread,passt,mot,clen,te"
[tool.ruff]
# Allow lines to be as long as 120.
v2/__init__.py
+30 -6
View File
@@ -1,4 +1,4 @@
"""v2 — a scoped, CPU-testable realization of the design_v3 model-native runtime.
"""v2 — a scoped, CPU-testable realization of the model-native runtime (see v2/README.md).
> A model card is a (recipe, runtime) pair with a parity obligation.
> The model owns loop semantics; the runtime owns loop lifecycle.
@@ -49,11 +49,35 @@ from v2.runtime import AsyncEngine, Engine
__version__ = "0.2.0"
__all__ = [
"ModelCard", "ComponentSpec", "LoopSpec", "RecipeSpec", "ParitySpec", "CostModel",
"CapabilityMatrix", "ModelInstance", "load_card",
"Engine", "AsyncEngine", "Program", "ProgramKind", "ComponentNode", "ModelLoopNode", "when_task", "when_opt",
"Request", "Session", "Output", "make_request", "TaskType", "SamplingParams", "DiffusionParams",
"LoopKind", "WorkUnitKind", "ConsistencyLevel", "ExecutionProfile", "Capability",
"ModelCard",
"ComponentSpec",
"LoopSpec",
"RecipeSpec",
"ParitySpec",
"CostModel",
"CapabilityMatrix",
"ModelInstance",
"load_card",
"Engine",
"AsyncEngine",
"Program",
"ProgramKind",
"ComponentNode",
"ModelLoopNode",
"when_task",
"when_opt",
"Request",
"Session",
"Output",
"make_request",
"TaskType",
"SamplingParams",
"DiffusionParams",
"LoopKind",
"WorkUnitKind",
"ConsistencyLevel",
"ExecutionProfile",
"Capability",
"VideoGenerator",
"__version__",
]
v2/_enums.py
+28 -29
View File
@@ -1,9 +1,8 @@
"""Shared vocabulary enums (design_v3 §6.1, §9.2, §9.4).
"""Shared vocabulary enums.
These live in a leaf module so both ``card/`` (which references loop kinds and
consistency levels in its specs) and ``loop/``/``runtime/`` can import them
without a circular dependency. design_v3 §18 boundary: ``card/`` imports no
runtime; this is pure vocabulary.
without a circular dependency. ``card/`` imports no runtime; this is pure vocabulary.
"""
from __future__ import annotations
@@ -11,21 +10,21 @@ from enum import Enum
class LoopKind(str, Enum):
"""The kind of iterative computation a LoopSpec describes (design_v3 §4.3)."""
DIFFUSION_DENOISE = "diffusion_denoise" # N solver steps over full clip (Wan, LTX-2)
CHUNK_ROLLOUT = "chunk_rollout" # causal chunk loop × inner denoise (self-forcing, world models)
AR_DECODE = "ar_decode" # token loop until EOS (reasoner/thinker/talker — phase 2)
VAE_TILE = "vae_tile" # tiled VAE encode/decode
ENCODER = "encoder" # one-shot encoder (text/vision) — degenerate single-step loop
AUDIO_DECODE = "audio_decode" # vocoder / codec decode
TRAIN_FORWARD = "train_forward" # grad-enabled forward for a training method
"""The kind of iterative computation a LoopSpec describes."""
DIFFUSION_DENOISE = "diffusion_denoise" # N solver steps over full clip (Wan, LTX-2)
CHUNK_ROLLOUT = "chunk_rollout" # causal chunk loop × inner denoise (self-forcing, world models)
AR_DECODE = "ar_decode" # token loop until EOS (reasoner/thinker/talker — phase 2)
VAE_TILE = "vae_tile" # tiled VAE encode/decode
ENCODER = "encoder" # one-shot encoder (text/vision) — degenerate single-step loop
AUDIO_DECODE = "audio_decode" # vocoder / codec decode
TRAIN_FORWARD = "train_forward" # grad-enabled forward for a training method
class WorkUnitKind(str, Enum):
"""The smallest schedulable action — the scheduler's currency unit (design_v3 §6.1).
"""The smallest schedulable action — the scheduler's currency unit.
Tokens are *one kind*, not the scheduler — the generalization of vLLM's token
scheduler that diffusion forces.
Tokens are one kind among many, not the scheduler itself: this generalizes
vLLM's token scheduler to the work units diffusion needs.
"""
AR_PREFILL = "ar_prefill"
AR_TOKEN = "ar_token"
@@ -35,20 +34,20 @@ class WorkUnitKind(str, Enum):
ENCODER_CHUNK = "encoder_chunk"
VAE_TILE = "vae_tile"
AUDIO_CHUNK = "audio_chunk"
REWARD_BATCH = "reward_batch" # RL
LOGPROB_BATCH = "logprob_batch" # RL (likelihood-based)
REWARD_BATCH = "reward_batch" # RL
LOGPROB_BATCH = "logprob_batch" # RL (likelihood-based)
TRANSFER = "transfer"
CACHE_IO = "cache_io"
GRAPH_CAPTURE = "graph_capture"
class ConsistencyLevel(str, Enum):
"""The consistency ladder (design_v3 §9.2). RL methods declare their required rung."""
C0 = "C0" # component parity (VAE/encoder/block/scheduler-step in isolation)
C1 = "C1" # loop parity (full denoise trajectory / AR logits, fixed seed)
C2 = "C2" # behavioral identity (train-forward ≡ serve-forward on the RL objective's quantity)
C3 = "C3" # distribution parity (rollout distribution under allowed nondeterminism)
C4 = "C4" # artifact quality (SSIM / reward agreement / human preference)
"""The consistency ladder. RL methods declare their required rung."""
C0 = "C0" # component parity (VAE/encoder/block/scheduler-step in isolation)
C1 = "C1" # loop parity (full denoise trajectory / AR logits, fixed seed)
C2 = "C2" # behavioral identity (train-forward ≡ serve-forward on the RL objective's quantity)
C3 = "C3" # distribution parity (rollout distribution under allowed nondeterminism)
C4 = "C4" # artifact quality (SSIM / reward agreement / human preference)
@property
def rank(self) -> int:
@@ -56,25 +55,25 @@ class ConsistencyLevel(str, Enum):
class ExecutionProfile(str, Enum):
"""Three forwards, one loop definition — differ only in grad mode + capture (design_v3 §9.4)."""
SERVE = "serve" # no-grad, graphed, cached, possibly quantized
ROLLOUT = "rollout" # serve profile + behavior capture
TRAIN = "train" # grad, checkpointed, FSDP-gathered
"""Three forwards, one loop definition — differ only in grad mode + capture."""
SERVE = "serve" # no-grad, graphed, cached, possibly quantized
ROLLOUT = "rollout" # serve profile + behavior capture
TRAIN = "train" # grad, checkpointed, FSDP-gathered
class Capability(str, Enum):
"""CapabilityMatrix entries (designv2 model plane)."""
"""CapabilityMatrix entries (model plane)."""
TEXT_TO_VIDEO = "text_to_video"
IMAGE_TO_VIDEO = "image_to_video"
VIDEO_TO_VIDEO = "video_to_video"
TEXT_TO_IMAGE = "text_to_image"
TEXT_TO_VIDEO_SOUND = "text_to_video_sound"
AUDIO_TO_VIDEO = "audio_to_video"
TEXT_TO_SPEECH = "text_to_speech" # thinker→talker→vocoder (Qwen-Omni): reason + speak
TEXT_TO_SPEECH = "text_to_speech" # thinker→talker→vocoder (Qwen-Omni): reason + speak
REASONING_TEXT = "reasoning_text"
ACTION_CONDITIONING = "action_conditioning"
STREAMING_VIDEO_CONTINUATION = "streaming_video_continuation"
VAE_ENCODE = "vae_encode"
VAE_DECODE = "vae_decode"
POLICY_ROLLOUT = "policy_rollout" # can serve as an RL rollout engine
POLICY_ROLLOUT = "policy_rollout" # can serve as an RL rollout engine
LOGPROB_RECOMPUTE = "logprob_recompute"
v2/_types.py
+4 -4
View File
@@ -1,10 +1,10 @@
"""Shared lightweight type aliases for v2.
The core (card/loop/runtime/cache/parity/program/request/training) is numpy-only
and CPU-testable. Tensors are therefore typed structurally as ``TensorLike``:
on the CPU test path they are ``numpy.ndarray``; on a GPU box the same code runs
with ``torch.Tensor``. The core never imports torch — only model-component
*adapters* do, lazily (see ``v2/card/components.py``).
and CPU-testable, so tensors are typed structurally as ``TensorLike``: a
``numpy.ndarray`` on the CPU test path, a ``torch.Tensor`` on GPU. The core never
imports torch — only model-component adapters do, lazily (see
``v2/card/components.py``).
"""
from __future__ import annotations
v2/api.py
+2 -8
View File
@@ -1,12 +1,6 @@
# STUB: re-exports fastvideo until vendored (see memory: v2-vendoring-approach).
"""``api`` facade — the typed config dataclasses the VideoGenerator consumes. Re-exported so v2 code
imports ``v2.api`` instead of ``fastvideo.api``; a vendored cutover replaces these with v2-native configs."""
imports ``v2.api`` instead of ``fastvideo.api``; a vendored cutover will replace these with v2-native configs."""
from fastvideo.api import ( # noqa: F401
EngineConfig,
GenerationRequest,
GenerationResult,
GeneratorConfig,
OffloadConfig,
OutputConfig,
SamplingConfig,
EngineConfig, GenerationRequest, GenerationResult, GeneratorConfig, OffloadConfig, OutputConfig, SamplingConfig,
)
v2/cache/__init__.py
Vendored
+5 -3
View File
@@ -1,9 +1,11 @@
"""Cache plane — correct by key, per-class pools (design_v3 §7)."""
"""Cache plane — correct by key, per-class pools."""
from __future__ import annotations
from v2.cache.classes import FeatureCache, PagedKVCache, ResidualCache, Slab, SlabKVCache, make_pool
from v2.cache.keys import CacheKey, CachePolicy, content_hash
from v2.cache.manager import CacheManager
__all__ = ["CacheKey", "CachePolicy", "content_hash", "CacheManager",
"FeatureCache", "ResidualCache", "SlabKVCache", "PagedKVCache", "Slab", "make_pool"]
__all__ = [
"CacheKey", "CachePolicy", "content_hash", "CacheManager", "FeatureCache", "ResidualCache", "SlabKVCache",
"PagedKVCache", "Slab", "make_pool"
]
v2/cache/classes.py
Vendored
+13 -15
View File
@@ -1,10 +1,7 @@
"""Per-class cache pools (design_v3 §7.2).
"""Per-class cache pools.
> There is **no single unified block pool**, because a unified pool requires uniform
> bytes-per-block and our cache classes differ by 150–500× in natural granularity ...
> Each class gets a statically budgeted pool behind one ``CacheHandle``.
Mini-fastvideo implements four classes (the §7.2 minimal set):
No single unified block pool: cache classes differ by 150-500x in natural granularity, so
each gets its own statically budgeted pool behind one ``CacheHandle``. Four classes:
* ``FeatureCache`` — content-hash keyed, partitioned by adapter+weights (text/vision encoders)
* ``ResidualCache`` — cache-dit residuals, scoped per request AND per CFG branch
* ``SlabKVCache`` — chunk-KV slabs (self-forcing / world models); training mode disables recycle
@@ -15,7 +12,7 @@ KV is the minority case — a pure bidirectional deployment (Wan/LTX T2V) alloca
from __future__ import annotations
from collections import OrderedDict
from dataclasses import dataclass, field
from dataclasses import dataclass
from typing import Any
from v2.cache.keys import CacheKey, CachePolicy
@@ -28,6 +25,7 @@ def _nbytes(value: Any) -> int:
class _Pool:
def __init__(self, policy: CachePolicy):
self.policy = policy
self.used_bytes = 0
@@ -36,7 +34,7 @@ class _Pool:
class FeatureCache(_Pool):
"""content-hash keyed, reference-counted-ish FIFO/LRU, budget-aware (§7.2).
"""Content-hash keyed, budget-aware FIFO/LRU.
Partitioned by ``adapter_versions``/``weights_version`` through the CacheKey, so two
workflows sharing a prompt but differing in te-LoRA stack never serve stale embeddings.
@@ -44,7 +42,7 @@ class FeatureCache(_Pool):
def __init__(self, policy: CachePolicy):
super().__init__(policy)
self._store: "OrderedDict[str, tuple[Any, int, CacheKey]]" = OrderedDict()
self._store: OrderedDict[str, tuple[Any, int, CacheKey]] = OrderedDict()
def get(self, key: CacheKey) -> Any | None:
if not self.policy.reuse_across_requests:
@@ -52,7 +50,7 @@ class FeatureCache(_Pool):
h = key.hash
if h in self._store:
self.hits += 1
self._store.move_to_end(h) # LRU
self._store.move_to_end(h) # LRU
return self._store[h][0]
self.misses += 1
return None
@@ -80,7 +78,7 @@ class FeatureCache(_Pool):
self.used_bytes -= nb
def invalidate_components(self, components: set[str]) -> None:
"""Drop only entries produced by the changed components (design_v3 §7.1 partition-not-flush):
"""Drop only entries produced by the changed components (partition, not flush):
a transformer-only weight sync must NOT evict text-encoder embeddings."""
drop = [h for h, (_v, _nb, k) in self._store.items() if k.component_id in components]
for h in drop:
@@ -89,7 +87,7 @@ class FeatureCache(_Pool):
class ResidualCache(_Pool):
"""cache-dit residual store, scoped per ``LoopState`` AND per CFG branch (§5.1, §11).
"""cache-dit residual store, scoped per ``LoopState`` AND per CFG branch.
Keyed by (namespace, branch, name) where namespace is the request/loop id. This is the
structural fix for the module-global residual state that corrupts cache-dit forks under
@@ -124,10 +122,10 @@ class Slab:
class SlabKVCache(_Pool):
"""Chunk-KV slabs for causal/world-model rollout (§7.2).
"""Chunk-KV slabs for causal/world-model rollout.
``training_mode`` disables mid-rollout recycling and keeps grad-aware index snapshots
so activation-checkpoint recompute doesn't double-advance the cache (self-forcing).
``training_mode`` disables mid-rollout recycling so activation-checkpoint recompute
doesn't double-advance the cache (self-forcing).
"""
def __init__(self, policy: CachePolicy):
v2/cache/keys.py
Vendored
+12 -16
View File
@@ -1,12 +1,9 @@
"""CacheKey — cache correctness is a contract, not hope (design_v3 §7.1).
"""CacheKey — cache correctness is a contract.
> If a field can change output semantics, it is in the key. Incorrect reuse is worse
> than no reuse.
The serving hazard this kills: a workflow-cloud request that shares a prompt but
differs in te-LoRA stack must not serve stale embeddings — so the key is *partitioned*
by ``adapter_versions``, not flushed. An RL ``update_weights`` bumps ``weights_version``
and invalidates wholesale.
If a field can change output semantics, it is in the key (incorrect reuse is worse than no reuse).
The serving hazard this kills: a request that shares a prompt but differs in te-LoRA stack must not
serve stale embeddings — so the key is *partitioned* by ``adapter_versions``, not flushed. An RL
``update_weights`` bumps ``weights_version`` and invalidates wholesale.
"""
from __future__ import annotations
@@ -16,17 +13,16 @@ from typing import Any
def content_hash(obj: Any) -> str:
"""Stable content hash for feature-cache keys (text/vision embeddings, §7.2).
"""Stable content hash for feature-cache keys (text/vision embeddings).
Lets K rollout samples of one prompt reuse a single text encode (design_v3 §10:
24× text-encode reduction).
Lets K rollout samples of one prompt reuse a single text encode.
"""
h = hashlib.sha256()
if isinstance(obj, str):
h.update(obj.encode("utf-8"))
elif isinstance(obj, (bytes, bytearray)):
h.update(obj)
elif hasattr(obj, "tobytes") and hasattr(obj, "shape"): # numpy / torch tensor
elif hasattr(obj, "tobytes") and hasattr(obj, "shape"): # numpy / torch tensor
h.update(str(getattr(obj, "shape", "")).encode())
h.update(str(getattr(obj, "dtype", "")).encode())
try:
@@ -44,7 +40,7 @@ class CacheKey:
component_id: str
loop_id: str | None = None
weights_version: str = "v0"
adapter_versions: tuple[tuple[str, str], ...] = () # sorted (adapter_id, version) pairs
adapter_versions: tuple[tuple[str, str], ...] = () # sorted (adapter_id, version) pairs
precision: str = "float32"
parallel_plan_hash: str = ""
shape_sig: str = ""
@@ -75,11 +71,11 @@ class CacheKey:
@dataclass
class CachePolicy:
"""Runtime config for one cache class pool (design_v3 §7.2; designv2 cache plane)."""
class_name: str # "feature" | "residual" | "slab_kv" | "paged_kv"
"""Runtime config for one cache class pool."""
class_name: str # "feature" | "residual" | "slab_kv" | "paged_kv"
max_bytes: int = 1 << 30
block_bytes: int = 1 << 16
eviction: str = "lru" # "lru" | "fifo" | "none"
eviction: str = "lru" # "lru" | "fifo" | "none"
reuse_across_requests: bool = True
per_component: dict[str, int] = field(default_factory=dict)
training_mode_disables_recycle: bool = False
v2/cache/manager.py
Vendored
+22 -11
View File
@@ -1,4 +1,4 @@
"""CacheManager — per-class pools with static budgets, behind one handle (design_v3 §7.2).
"""CacheManager — per-class pools with static budgets, behind one handle.
Static partitioning makes cross-class fragmentation impossible (jumbo slab traffic cannot
strand text-KV pages and vice versa). Each class gets a budget carved at init from the card's
@@ -13,22 +13,27 @@ from v2.cache.keys import CachePolicy
class CacheManager:
def __init__(self, policies: list[CachePolicy] | None = None):
self._pools: dict[str, _Pool] = {}
for p in (policies or []):
self._pools[p.class_name] = make_pool(p)
@classmethod
def from_card(cls, card) -> "CacheManager":
def from_card(cls, card) -> CacheManager:
"""Build the per-class pools a card declares (KV pools materialize only if declared)."""
policies = []
for cc in card.caches.values():
policies.append(CachePolicy(
class_name=cc.cache_class, max_bytes=cc.max_bytes, block_bytes=cc.block_bytes,
eviction=cc.eviction, reuse_across_requests=cc.reuse_across_requests,
per_component=dict(cc.per_component),
training_mode_disables_recycle=cc.training_mode_disables_recycle,
))
policies.append(
CachePolicy(
class_name=cc.cache_class,
max_bytes=cc.max_bytes,
block_bytes=cc.block_bytes,
eviction=cc.eviction,
reuse_across_requests=cc.reuse_across_requests,
per_component=dict(cc.per_component),
training_mode_disables_recycle=cc.training_mode_disables_recycle,
))
return cls(policies)
def pool(self, class_name: str) -> _Pool:
@@ -45,7 +50,7 @@ class CacheManager:
pool.invalidate_weights(version)
def invalidate_components(self, components) -> None:
"""Component-scoped invalidation (design_v3 §7.1): only drop caches for changed components."""
"""Component-scoped invalidation: only drop caches for changed components."""
comps = set(components)
for pool in self._pools.values():
if hasattr(pool, "invalidate_components"):
@@ -60,5 +65,11 @@ class CacheManager:
pool.free_namespace(namespace)
def stats(self) -> dict[str, Any]:
return {name: {"used_bytes": p.used_bytes, "hits": p.hits, "misses": p.misses}
for name, p in self._pools.items()}
return {
name: {
"used_bytes": p.used_bytes,
"hits": p.hits,
"misses": p.misses
}
for name, p in self._pools.items()
}
v2/card/__init__.py
+23 -6
View File
@@ -1,4 +1,4 @@
"""Model Plane — the center (design_v3 §4). The (recipe, runtime) pair as a typed card."""
"""Model Plane — the (recipe, runtime) pair as a typed card."""
from __future__ import annotations
from v2._enums import Capability, ConsistencyLevel, ExecutionProfile, LoopKind, WorkUnitKind
@@ -22,9 +22,26 @@ from v2.card.specs import (
)
__all__ = [
"ModelCard", "ComponentSpec", "LoopSpec", "RecipeSpec", "ParitySpec", "ParityTestSpec",
"CheckpointManifest", "CapabilityMatrix", "CostModel", "CacheContract", "SamplingDefaults",
"ParallelismContract", "PrecisionContract", "DataRef", "CardValidationError",
"ModelInstance", "load_card",
"Capability", "ConsistencyLevel", "ExecutionProfile", "LoopKind", "WorkUnitKind",
"ModelCard",
"ComponentSpec",
"LoopSpec",
"RecipeSpec",
"ParitySpec",
"ParityTestSpec",
"CheckpointManifest",
"CapabilityMatrix",
"CostModel",
"CacheContract",
"SamplingDefaults",
"ParallelismContract",
"PrecisionContract",
"DataRef",
"CardValidationError",
"ModelInstance",
"load_card",
"Capability",
"ConsistencyLevel",
"ExecutionProfile",
"LoopKind",
"WorkUnitKind",
]
v2/card/instance.py
+30 -24
View File
@@ -1,13 +1,10 @@
"""ModelInstance — a resident, loaded card (design_v3 §4.3).
"""ModelInstance — a resident, loaded card: component instances, model state, caches,
compiled graphs, and a parallel plan. One request may run several of the card's loops
against one ``ModelInstance``.
> A ``ModelInstance`` is a resident, loaded card: component instances, model state,
> caches, compiled graphs, and a parallel plan. **A request may run several of the
> card's loops against one ``ModelInstance``.**
That single sentence is the difference between this design and a stage-only design,
and it is what makes omni native: when two loops bind the same component
(``shared_weight_components``), ``component()`` returns the *exact same* live object
— no duplicated weights, shared live state.
This is what makes omni native: when two loops bind the same component
(``shared_weight_components``), ``component()`` returns the *exact same* live object —
no duplicated weights, shared live state. See v2/README.md.
"""
from __future__ import annotations
@@ -21,23 +18,26 @@ class ModelInstance:
"""The live, resident form of a ModelCard. Also serves as the ``ModelState``
passed to ``Loop.init`` (the loop reads components through it)."""
def __init__(self, card: ModelCard, parallel_plan: Any = None,
cache_manager: Any = None, weights_version: str = "v0",
def __init__(self,
card: ModelCard,
parallel_plan: Any = None,
cache_manager: Any = None,
weights_version: str = "v0",
platform: Any = None):
self.card = card
self.parallel_plan = parallel_plan
self.caches = cache_manager
self.weights_version = weights_version
# The detected (device, arch). Resolves component/kernel implementations through the two
# backend registries; defaults to CPU/numpy. Swapping this to a GPU platform is the whole
# "change only the backend, not the loops/policies/training" story (design_v3 §17).
# backend registries; defaults to CPU/numpy. Swapping this to a GPU platform changes only
# the backend, not the loops/policies/training.
self.platform = platform if platform is not None else Platform.cpu()
# Piecewise CUDA-graph cache for loops declaring graph_capture="breakable_cudagraph". Lazily
# created by the runtime (kept as a plain attr so card/ stays free of any runtime import —
# design_v3 §18 boundary); a captured graph is tied to this instance's resident weights.
# created by the runtime (kept as a plain attr so card/ stays free of any runtime import);
# a captured graph is tied to this instance's resident weights.
self.graphs: Any = None
self.adapter_versions: dict[str, str] = {}
# per-component weight versions (§7.1): a component's version changes only when IT is synced,
# Per-component weight versions: a component's version changes only when IT is synced,
# so a transformer-only RL sync never invalidates the frozen text-encoder's feature cache.
self.component_versions: dict[str, str] = {cid: weights_version for cid in card.components}
self._components: dict[str, Any] = {}
@@ -71,7 +71,7 @@ class ModelInstance:
self._loops[loop_id] = spec.loop_factory()
return self._loops[loop_id]
# --- sleep/wake by component (design_v3 §7.3; CuMem tags = component names) #
# --- sleep/wake by component (CuMem tags = component names) --------------- #
def sleep(self, component_ids: list[str]) -> None:
for cid in component_ids:
self._components.pop(cid, None)
@@ -81,15 +81,15 @@ class ModelInstance:
for cid in component_ids:
self._asleep.discard(cid)
# --- weight sync (design_v3 §10): bump version + invalidate caches -------- #
# --- weight sync: bump version + invalidate caches ----------------------- #
def version_of(self, component_id: str) -> str:
"""The component's own weights version (defaults to the instance version)."""
return self.component_versions.get(component_id, self.weights_version)
def set_weights_version(self, version: str, components: list[str] | None = None) -> None:
"""Publish a new weights version. If ``components`` is given, only those components' versions
bump and only their caches are invalidated (design_v3 §7.1 partition-not-flush) — so a
transformer-only RL weight sync leaves the frozen text-encoder's feature cache intact."""
bump and only their caches are invalidated (partition, not flush) — so a transformer-only RL
weight sync leaves the frozen text-encoder's feature cache intact."""
self.weights_version = version
changed = components if components is not None else list(self.card.components.keys())
for c in changed:
@@ -106,9 +106,13 @@ class ModelInstance:
f"resident={sorted(self._components)})")
def load_card(card: ModelCard, parallel_plan: Any = None, cache_manager: Any = None,
*, validate: bool = True, platform: Any = None) -> ModelInstance:
"""The card-as-factory entrypoint (design_v3 §4.1: card is a runtime factory).
def load_card(card: ModelCard,
parallel_plan: Any = None,
cache_manager: Any = None,
*,
validate: bool = True,
platform: Any = None) -> ModelInstance:
"""The card-as-factory entrypoint: the card is a runtime factory.
``platform`` selects the backend (device, arch); when omitted it is detected (CPU/numpy here,
CUDA on a torch+GPU box). The same card loads on any backend — only the resolved component/kernel
@@ -116,5 +120,7 @@ def load_card(card: ModelCard, parallel_plan: Any = None, cache_manager: Any = N
"""
if validate:
card.validate()
return ModelInstance(card, parallel_plan=parallel_plan, cache_manager=cache_manager,
return ModelInstance(card,
parallel_plan=parallel_plan,
cache_manager=cache_manager,
platform=platform if platform is not None else Platform.detect())
v2/card/specs.py
+66 -73
View File
@@ -1,36 +1,36 @@
"""ModelCard and its sub-specs — the Model Plane, the center (design_v3 §4).
"""ModelCard and its sub-specs — the Model Plane.
> The atomic unit is the (recipe, runtime) pair, owned by a typed ``ModelCard``.
The atomic unit is the (recipe, runtime) pair, owned by a typed ``ModelCard``. The card
is both a declarative contract (strict enough to validate before any GPU touches it — see
``ModelCard.validate``) and a runtime factory (it instantiates components, binds loops, and
resolves caches — see ``card/instance.py``).
The card is both a **declarative contract** (strict enough to validate before any
GPU touches it — see ``ModelCard.validate``) and a **runtime factory** (it knows how
to instantiate components, bind loops, and resolve caches — see ``card/instance.py``).
design_v3 §18 boundary: ``card/`` imports no product/runtime. It depends only on the
shared leaf modules (``_enums``, ``_types``) and references ``ParallelPlan`` under
``TYPE_CHECKING`` so there is no runtime coupling to ``parallel/``.
Boundary: ``card/`` imports no product/runtime. It depends only on the shared leaf modules
(``_enums``, ``_types``) and references ``ParallelPlan`` under ``TYPE_CHECKING`` so there is
no runtime coupling to ``parallel/``.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Callable
from typing import TYPE_CHECKING, Any
from collections.abc import Callable
from v2._enums import Capability, ConsistencyLevel, LoopKind, WorkUnitKind
if TYPE_CHECKING: # avoid runtime card -> parallel coupling
from v2.parallel.plan import ParallelPlan
pass
# --------------------------------------------------------------------------- #
# Cost model — the budget currency, REQUIRED on every LoopSpec (design_v3 §6.1) #
# Cost model — the budget currency, REQUIRED on every LoopSpec #
# --------------------------------------------------------------------------- #
@dataclass
class CostModel:
"""Predicted GPU-time per WorkUnit at (shape, precision, policy).
"The budget currency is predicted GPU-time, not counts." One object, two
consumers: the scheduler's internal budget and the fleet's routing input
(design_v3 §6.1, §14). Calibrated online by the Profiler observer (§11).
The budget currency is predicted GPU-time, not counts. One object, two consumers:
the scheduler's internal budget and the fleet's routing input. Calibrated online by
the Profiler observer.
"""
kind: WorkUnitKind
# synthetic seconds = base + per_unit * work_units(shape) * policy_factor
@@ -39,7 +39,7 @@ class CostModel:
coefficients: dict[str, float] = field(default_factory=dict)
def predict(self, work_units: float, policy_factor: float = 1.0) -> float:
"""Conservative baseline cost (design_v3 §6.1: admission never optimistic)."""
"""Conservative baseline cost: admission is never optimistic."""
return (self.base_seconds + self.per_unit_seconds * float(work_units)) * policy_factor
@@ -50,8 +50,8 @@ class CostModel:
class PrecisionContract:
default_dtype: str = "float32"
component_overrides: dict[str, str] = field(default_factory=dict)
quantization_scheme: str | None = None # "nvfp4" | "int8" | None
training_precision: str = "float32" # distinct from serving precision
quantization_scheme: str | None = None # "nvfp4" | "int8" | None
training_precision: str = "float32" # distinct from serving precision
def dtype_for(self, component_id: str) -> str:
return self.component_overrides.get(component_id, self.default_dtype)
@@ -59,39 +59,39 @@ class PrecisionContract:
@dataclass
class ParallelismContract:
valid_plans: list = field(default_factory=list) # list[ParallelPlan]
default_plan: Any = None # ParallelPlan | None
valid_plans: list = field(default_factory=list) # list[ParallelPlan]
default_plan: Any = None # ParallelPlan | None
@dataclass
class CacheContract:
"""A cache class the card declares it needs (design_v3 §7.2; designv2 cache plane)."""
cache_class: str # "feature" | "residual" | "slab_kv" | "paged_kv"
"""A cache class the card declares it needs."""
cache_class: str # "feature" | "residual" | "slab_kv" | "paged_kv"
max_bytes: int = 1 << 30
block_bytes: int = 1 << 16 # page bytes (paged) / slab granule (slab)
eviction: str = "lru" # "lru" | "fifo" | "none"
reuse_across_requests: bool = True # paged/feature=True; slab depends on mode
block_bytes: int = 1 << 16 # page bytes (paged) / slab granule (slab)
eviction: str = "lru" # "lru" | "fifo" | "none"
reuse_across_requests: bool = True # paged/feature=True; slab depends on mode
per_component: dict[str, int] = field(default_factory=dict)
training_mode_disables_recycle: bool = False # chunk-KV training mode (§7.2)
training_mode_disables_recycle: bool = False # chunk-KV training mode
@dataclass
class ParityTestSpec:
"""A named tap + tolerance + ladder level (design_v3 §9.1)."""
"""A named tap + tolerance + ladder level."""
name: str
level: ConsistencyLevel
tap: str # named activation tap, e.g. "block.0.out"
tap: str # named activation tap, e.g. "block.0.out"
rtol: float = 0.0
atol: float = 0.0
@dataclass
class ParitySpec:
"""The (recipe, runtime) honesty contract (design_v3 §9). Measured, never assumed."""
"""The (recipe, runtime) honesty contract. Measured, never assumed."""
consistency_levels: list[ConsistencyLevel] = field(default_factory=lambda: [ConsistencyLevel.C1])
tests: list[ParityTestSpec] = field(default_factory=list)
tap_tolerances: dict[str, float] = field(default_factory=dict)
interleave_required: bool = True # the batch-of-N gate (§9.3) is non-negotiable
interleave_required: bool = True # the batch-of-N gate is non-negotiable
@property
def max_level(self) -> ConsistencyLevel:
@@ -100,7 +100,7 @@ class ParitySpec:
@dataclass
class DataRef:
"""What a recipe trained on, for governance/reproduction (design_v3 §4.2)."""
"""What a recipe trained on, for governance/reproduction."""
dataset_id: str = ""
revision: str = ""
description: str = ""
@@ -108,41 +108,41 @@ class DataRef:
@dataclass
class RecipeSpec:
"""The provenance half of the (recipe, runtime) pair (design_v3 §4.2, §2.1).
"""The provenance half of the (recipe, runtime) pair.
``assumes_loop`` and ``assumes_precision`` are the teeth: a 4-step distilled
model whose ``assumes_loop = "ddim_4step"`` cannot be served under a 50-step
sampler without a typed mismatch error (enforced in ModelCard.validate).
"""
method: str = "base" # base | dmd2 | self_forcing | diffusion_nft | attn_qat_nvfp4
parents: list[str] = field(default_factory=list) # teacher / base model_ids
method: str = "base" # base | dmd2 | self_forcing | diffusion_nft | attn_qat_nvfp4
parents: list[str] = field(default_factory=list) # teacher / base model_ids
data_contract: DataRef = field(default_factory=DataRef)
assumes_loop: str = "" # loop_id this recipe's weights require
assumes_loop: str = "" # loop_id this recipe's weights require
assumes_precision: str = "float32"
consistency_required: ConsistencyLevel = ConsistencyLevel.C1
@dataclass
class ComponentSpec:
"""A weight-bearing (or processing) component (design_v3 §4.3; designv2 §1.2).
"""A weight-bearing (or processing) component.
Omni-ready fields ``resident_for`` / ``optional_for`` / ``required_for`` turn the
Cosmos3 lazy-sound-VAE problem into a declaration, not an ``if env_var`` inside
``forward`` (design_v3 §4.1).
``forward``.
"""
component_id: str
kind: str # dit | vae | text_encoder | audio_vae | reasoner_tower | ...
load_id: str = "" # "module:Class" for the real (torch) adapter
kind: str # dit | vae | text_encoder | audio_vae | reasoner_tower | ...
load_id: str = "" # "module:Class" for the real (torch) adapter
config_schema: type | None = None
io_schema: tuple[type | None, type | None] = (None, None)
precision_policy: str | None = None
placement_policy: str = "colocated"
parallel_constraints: dict[str, Any] = field(default_factory=dict)
parity_tests: list[ParityTestSpec] = field(default_factory=list)
# omni-ready (design_v3 §4.1; omni-readiness report §4a):
# omni-ready:
resident_for: list[str] = field(default_factory=list) # loop_ids that keep this resident mid-request
optional_for: set[str] = field(default_factory=set) # tasks that don't need it
required_for: set[str] = field(default_factory=set) # tasks that require it
optional_for: set[str] = field(default_factory=set) # tasks that don't need it
required_for: set[str] = field(default_factory=set) # tasks that require it
# v2 wiring: a factory producing the live component (toy numpy or torch adapter)
factory: Callable[..., Any] | None = None
# GPU backend: weights source (HF id or local path) for the real torch adapter resolved from
@@ -157,7 +157,7 @@ class ComponentSpec:
@dataclass
class LoopSpec:
"""Describes one iterative computation the card can run (design_v3 §4.3).
"""Describes one iterative computation the card can run.
The model owns loop *semantics*; the runtime owns loop *lifecycle*. A single
``ModelInstance`` may run several of these against shared components — the MoT
@@ -166,16 +166,16 @@ class LoopSpec:
loop_id: str
kind: LoopKind
work_unit_kind: WorkUnitKind
step_cost_model: CostModel # REQUIRED (design_v3 §6.1)
state_schema: type | None = None # the typed LoopState
step_schema: type | None = None # the typed WorkPlan a step emits
result_schema: type | None = None # the typed StepResult
behavior_schema: type | None = None # what to capture for RL (None if not training-relevant)
extension_schema: type | None = None # per-model LoopState extension (Cosmos3PackedSeq, etc.)
step_cost_model: CostModel # REQUIRED
state_schema: type | None = None # the typed LoopState
step_schema: type | None = None # the typed WorkPlan a step emits
result_schema: type | None = None # the typed StepResult
behavior_schema: type | None = None # what to capture for RL (None if not training-relevant)
extension_schema: type | None = None # per-model LoopState extension (Cosmos3PackedSeq, etc.)
cache_policy: list[str] = field(default_factory=list) # cache class names this loop draws from
valid_parallel_plans: list = field(default_factory=list)
graph_capture: str = "eager" # eager | breakable_cudagraph
# omni-ready (omni-readiness report §4a):
graph_capture: str = "eager" # eager | breakable_cudagraph
# omni-ready:
shared_weight_components: list[str] = field(default_factory=list)
allows_interleaving: bool = True
# v2: the Loop implementation factory (built at bind time)
@@ -184,12 +184,12 @@ class LoopSpec:
@dataclass
class CheckpointManifest:
"""Explicit declared components + key maps — no name-detector guessing (design_v3 §4.1)."""
"""Explicit declared components + key maps — no name-detector guessing."""
upstream_source: str = ""
revision: str = ""
component_ownership: dict[str, list[str]] = field(default_factory=dict) # component_id -> file globs
key_mappings: dict[str, str] = field(default_factory=dict) # ckpt_key -> component.param
required_for: dict[str, set[str]] = field(default_factory=dict) # component_id -> tasks requiring it
key_mappings: dict[str, str] = field(default_factory=dict) # ckpt_key -> component.param
required_for: dict[str, set[str]] = field(default_factory=dict) # component_id -> tasks requiring it
optional_for: dict[str, set[str]] = field(default_factory=dict)
conversion_version: str = "v0"
@@ -202,7 +202,7 @@ class CapabilityMatrix:
return cap in self.capabilities
@classmethod
def of(cls, *caps: Capability) -> "CapabilityMatrix":
def of(cls, *caps: Capability) -> CapabilityMatrix:
return cls(frozenset(caps))
@@ -210,8 +210,7 @@ class CapabilityMatrix:
class SamplingDefaults:
"""Per-model default generation params (the v2 mirror of fastvideo's per-model ``InferencePreset``
defaults). Applied by the entrypoint when the caller didn't specify a value; ``None`` => fall back to
the generic default. ``shift``/``sigmas``/MoE boundary already live in the card's loop policies — these
are the user-facing knobs surfaced at request time."""
the generic default. These are the user-facing knobs surfaced at request time."""
num_steps: int | None = None
guidance_scale: float | None = None
guidance_per_modality: dict[str, float] = field(default_factory=dict) # joint A/V, e.g. {"video":3,"audio":7}
@@ -246,28 +245,23 @@ class ModelCard:
checkpoint: CheckpointManifest = field(default_factory=CheckpointManifest)
sampling_defaults: SamplingDefaults = field(default_factory=SamplingDefaults)
def validate(self) -> "ModelCard":
"""Strict enough to validate before any GPU touches it (design_v3 §4.1).
def validate(self) -> ModelCard:
"""Strict enough to validate before any GPU touches it.
Returns self so it can be chained. Raises CardValidationError on any
contract violation — the (recipe, runtime) binding is enforced here.
"""
errs: list[str] = []
# 1) recipe.assumes_loop must exist (the (recipe, runtime) binding, §2.1/§4.2)
# 1) recipe.assumes_loop must exist (the (recipe, runtime) binding)
if self.recipe.assumes_loop and self.recipe.assumes_loop not in self.loops:
errs.append(
f"recipe.assumes_loop={self.recipe.assumes_loop!r} is not a declared loop "
f"(have {sorted(self.loops)}) — a recipe cannot assume a loop the card does not run"
)
errs.append(f"recipe.assumes_loop={self.recipe.assumes_loop!r} is not a declared loop "
f"(have {sorted(self.loops)}) — a recipe cannot assume a loop the card does not run")
# 2) recipe.assumes_precision must be consistent with the precision contract
if (self.recipe.assumes_precision
and self.recipe.assumes_precision not in (self.precision.default_dtype,
*self.precision.component_overrides.values())):
errs.append(
f"recipe.assumes_precision={self.recipe.assumes_precision!r} not present in precision contract"
)
if (self.recipe.assumes_precision and self.recipe.assumes_precision
not in (self.precision.default_dtype, *self.precision.component_overrides.values())):
errs.append(f"recipe.assumes_precision={self.recipe.assumes_precision!r} not present in precision contract")
# 3) every loop's shared_weight_components must be declared components
for lid, loop in self.loops.items():
@@ -278,7 +272,7 @@ class ModelCard:
if cc not in self.caches:
errs.append(f"loop {lid!r} references cache class {cc!r} not in card.caches")
if loop.step_cost_model is None:
errs.append(f"loop {lid!r} is missing the REQUIRED step_cost_model (design_v3 §6.1)")
errs.append(f"loop {lid!r} is missing the REQUIRED step_cost_model")
# 4) required_for / optional_for tasks must be disjoint per component
for cid, comp in self.components.items():
@@ -287,8 +281,7 @@ class ModelCard:
errs.append(f"component {cid!r} lists tasks {overlap} as both required and optional")
if errs:
raise CardValidationError(
f"ModelCard {self.model_id!r} failed validation:\n - " + "\n - ".join(errs))
raise CardValidationError(f"ModelCard {self.model_id!r} failed validation:\n - " + "\n - ".join(errs))
return self
def loops_sharing(self, component_id: str) -> list[str]:
v2/deploy/__init__.py
+5 -4
View File
@@ -1,4 +1,4 @@
"""Deployment & fleet plane (design_v3 §14) — DeploymentCard, our own LocalFleet, Dynamo adapter.
"""Deployment & fleet plane — DeploymentCard, our own LocalFleet, Dynamo adapter.
The engine exports a DeploymentCard; our LocalFleet routes over it (so we don't rely on Dynamo),
and the DynamoWorkerAdapter exports the same card so Dynamo can front us too — one object, two
@@ -10,6 +10,7 @@ from v2.deploy.card import DeploymentCard, HealthSchema, SLOSchema, build_deploy
from v2.deploy.dynamo import DynamoWorkerAdapter, FakeDynamoRuntime
from v2.deploy.fleet import LocalFleet, NoWorkerAvailable, Worker
__all__ = ["DeploymentCard", "HealthSchema", "SLOSchema", "build_deployment_card",
"LocalFleet", "Worker", "NoWorkerAvailable",
"DynamoWorkerAdapter", "FakeDynamoRuntime"]
__all__ = [
"DeploymentCard", "HealthSchema", "SLOSchema", "build_deployment_card", "LocalFleet", "Worker", "NoWorkerAvailable",
"DynamoWorkerAdapter", "FakeDynamoRuntime"
]
v2/deploy/card.py
+22 -15
View File
@@ -1,7 +1,8 @@
"""DeploymentCard — what an engine exports to a fleet (design_v3 §14; designv2 Deployment plane).
"""DeploymentCard — what an engine exports to a fleet.
> The engine exports a ``DeploymentCard`` and lets a fleet orchestrator route. ... The cost model is
> the SAME §6 object — one object, two consumers (the scheduler's budget and the fleet's routing input).
The engine exports a ``DeploymentCard`` and lets a fleet orchestrator route. The cost model is the
SAME object the scheduler budgets with — one object, two consumers (the scheduler's budget and the
fleet's routing input).
This is the contract our OWN fleet (``deploy/fleet.py``) consumes AND the Dynamo adapter
(``deploy/dynamo.py``) exports — so we are frontable by Dynamo without depending on it.
@@ -10,7 +11,6 @@ from __future__ import annotations
import dataclasses
from dataclasses import dataclass, field
from typing import Any
from v2._enums import Capability
from v2.card import CostModel
@@ -18,14 +18,14 @@ from v2.card import CostModel
@dataclass
class HealthSchema:
status: str = "healthy" # "healthy" | "draining" | "unhealthy"
status: str = "healthy" # "healthy" | "draining" | "unhealthy"
in_flight: int = 0
queue_depth: int = 0
@dataclass
class SLOSchema:
slo_class: str = "standard" # "latency" | "throughput" | "cost"
slo_class: str = "standard" # "latency" | "throughput" | "cost"
max_concurrent: int = 8
@@ -34,9 +34,9 @@ class DeploymentCard:
engine_id: str
model_cards: list[str] = field(default_factory=list)
capabilities: frozenset[Capability] = frozenset()
role_pools: list = field(default_factory=list) # list[RolePoolSpec]
role_pools: list = field(default_factory=list) # list[RolePoolSpec]
supported_programs: list[str] = field(default_factory=list)
cost_model: CostModel | None = None # the SAME §6 cost model
cost_model: CostModel | None = None # the SAME cost model the scheduler uses
health: HealthSchema = field(default_factory=HealthSchema)
slo: SLOSchema = field(default_factory=SLOSchema)
@@ -44,10 +44,14 @@ class DeploymentCard:
return model_id in self.model_cards
def build_deployment_card(engine_id: str, model_cards: list, *, max_concurrent: int = 8,
slo_class: str = "standard", role_pools: list | None = None,
def build_deployment_card(engine_id: str,
model_cards: list,
*,
max_concurrent: int = 8,
slo_class: str = "standard",
role_pools: list | None = None,
supported_programs: list[str] | None = None) -> DeploymentCard:
"""Export a DeploymentCard from the model cards an engine serves (design_v3 §14).
"""Export a DeploymentCard from the model cards an engine serves.
Picks a representative ``step_cost_model`` so the fleet/Dynamo route on the SAME cost object the
scheduler budgets with."""
@@ -64,7 +68,10 @@ def build_deployment_card(engine_id: str, model_cards: list, *, max_concurrent:
cost_model = dataclasses.replace(lp.step_cost_model,
coefficients=dict(lp.step_cost_model.coefficients))
break
return DeploymentCard(
engine_id=engine_id, model_cards=ids, capabilities=frozenset(caps),
role_pools=role_pools or [], supported_programs=supported_programs or [],
cost_model=cost_model, slo=SLOSchema(slo_class=slo_class, max_concurrent=max_concurrent))
return DeploymentCard(engine_id=engine_id,
model_cards=ids,
capabilities=frozenset(caps),
role_pools=role_pools or [],
supported_programs=supported_programs or [],
cost_model=cost_model,
slo=SLOSchema(slo_class=slo_class, max_concurrent=max_concurrent))
v2/deploy/dynamo.py
+20 -16
View File
@@ -1,11 +1,11 @@
"""Dynamo worker adapter (design_v3 §14; design.md §6.3.5) — Dynamo as an *option*, not a dependency.
"""Dynamo worker adapter — Dynamo as an *option*, not a dependency.
> NVIDIA Dynamo is the named first-class partner for the fleet layer ... the engine's job is to be an
> excellent Dynamo citizen: registration, health/drain, cost metrics, affinity events.
NVIDIA Dynamo is the named first-class partner for the fleet layer; the engine's job is to be a good
Dynamo citizen: registration, health/drain, cost metrics, affinity events.
This adapter exposes exactly that contract over an AsyncEngine + DeploymentCard, so Dynamo CAN front
this engine. Crucially, it consumes the **same** DeploymentCard + cost model as our own LocalFleet —
one object, two consumers — so choosing Dynamo vs. our fleet is a deployment decision, not a rewrite.
this engine. It consumes the same DeploymentCard + cost model as our own LocalFleet — one object, two
consumers — so choosing Dynamo vs. our fleet is a deployment decision, not a rewrite.
``FakeDynamoRuntime`` proves the contract is satisfiable end-to-end without importing Dynamo.
"""
from __future__ import annotations
@@ -17,7 +17,7 @@ from v2.deploy.card import DeploymentCard
class DynamoWorkerAdapter:
"""Implements the Dynamo worker surface (design.md §6.3.5 work items 1–4)."""
"""Implements the Dynamo worker surface: registration, metrics, affinity, handle."""
def __init__(self, engine: Any, card: DeploymentCard, *, worker_type: str = "Aggregated"):
self.engine = engine
@@ -30,19 +30,23 @@ class DynamoWorkerAdapter:
def registration(self) -> dict[str, Any]:
self.registered = True
return {
"engine_id": self.card.engine_id,
"engine_id":
self.card.engine_id,
"model_type": ["Videos", "Images"] + (["Chat"] if any("omni" in m or "cosmos" in m or "bagel" in m
for m in self.card.model_cards) else []),
"worker_type": self.worker_type,
"models": list(self.card.model_cards),
"capabilities": sorted(c.value for c in self.card.capabilities),
"supported_programs": list(self.card.supported_programs),
for m in self.card.model_cards) else []),
"worker_type":
self.worker_type,
"models":
list(self.card.model_cards),
"capabilities":
sorted(c.value for c in self.card.capabilities),
"supported_programs":
list(self.card.supported_programs),
}
# 2) metrics for routing + the SLA Planner (the SAME §6 cost model)
# 2) metrics for routing + the SLA Planner (the SAME cost model the scheduler uses)
def metrics(self) -> dict[str, Any]:
return {"in_flight": self.engine.in_flight, "queue_depth": self.engine.queue_depth,
"draining": self.draining}
return {"in_flight": self.engine.in_flight, "queue_depth": self.engine.queue_depth, "draining": self.draining}
def cost_estimate(self, request: Any) -> float:
cm = self.card.cost_model
@@ -57,7 +61,7 @@ class DynamoWorkerAdapter:
def drain(self) -> None:
self.draining = True
# 3) affinity / cache events (KvCacheEventData shape) — checkpoint/session residency (ask A1)
# 3) affinity / cache events (KvCacheEventData shape) — checkpoint/session residency
def cache_event(self, kind: str, key: str) -> dict[str, Any]:
return {"event": kind, "engine_id": self.card.engine_id, "key": key}
v2/deploy/fleet.py
+16 -14
View File
@@ -1,15 +1,15 @@
"""LocalFleet — OUR OWN fleet router (design_v3 §14; design.md §6.3.5/6.3.6 fallbacks).
"""LocalFleet — OUR OWN fleet router.
design.md names Dynamo as the first-class fleet partner, but every Dynamo ask (§6.3.6) has a
*first-class fallback* — and this is it: a self-contained fleet that does discovery, health/drain,
and routing (least-loaded / cost-model / affinity) over multiple engine workers, so we are never
*reliant* on Dynamo. The router's cost input is the SAME §6 cost model the scheduler uses (one
object, two consumers). Affinity routing is the A1 fallback (sticky-by-key for checkpoint/session
residency, "least-loaded + engine redirects").
Dynamo is the first-class fleet partner, but every Dynamo ask has a first-class fallback — and this
is it: a self-contained fleet that does discovery, health/drain, and routing (least-loaded /
cost-model / affinity) over multiple engine workers, so we are never *reliant* on Dynamo. The
router's cost input is the SAME cost model the scheduler uses (one object, two consumers). Affinity
routing is sticky-by-key for checkpoint/session residency (least-loaded + engine redirects).
"""
from __future__ import annotations
from typing import Any, AsyncIterator
from typing import Any
from collections.abc import AsyncIterator
from v2.request.artifacts import Output
from v2.deploy.card import DeploymentCard, HealthSchema
@@ -40,7 +40,7 @@ class Worker:
return self.engine.serves(model_id) or self.card.serves(model_id)
def cost_estimate(self, request: Any) -> float:
"""Predicted GPU-time for this request — the §6 cost model as the fleet's routing input."""
"""Predicted GPU-time for this request — the cost model as the fleet's routing input."""
cm = self.card.cost_model
steps = max(1, int(getattr(request.diffusion, "num_steps", 1) or 1))
work = max(1, int(getattr(request.diffusion, "height", 1)) * int(getattr(request.diffusion, "width", 1)))
@@ -49,16 +49,18 @@ class Worker:
def health(self) -> HealthSchema:
return HealthSchema(status=("draining" if self.draining else "healthy"),
in_flight=self.engine.in_flight, queue_depth=self.engine.queue_depth)
in_flight=self.engine.in_flight,
queue_depth=self.engine.queue_depth)
class LocalFleet:
def __init__(self, policy: str = "least_loaded", *, max_affinity: int = 100_000):
assert policy in ("least_loaded", "cost", "affinity")
self.policy = policy
self.max_affinity = max_affinity
self.workers: dict[str, Worker] = {}
self._affinity: dict[str, str] = {} # affinity key -> worker_id (sticky), FIFO-bounded
self._affinity: dict[str, str] = {} # affinity key -> worker_id (sticky), FIFO-bounded
# --- discovery / health (what Dynamo's registry + planner would do) ------ #
def register(self, worker_id: str, engine: Any, card: DeploymentCard) -> Worker:
@@ -89,14 +91,14 @@ class LocalFleet:
wid = self._affinity.get(key)
if wid in self.workers and self.workers[wid] in cands:
return self.workers[wid]
chosen = min(cands, key=lambda w: w.load) # cold key → least-loaded, then pin
if len(self._affinity) >= self.max_affinity: # FIFO-bound the sticky map
chosen = min(cands, key=lambda w: w.load) # cold key → least-loaded, then pin
if len(self._affinity) >= self.max_affinity: # FIFO-bound the sticky map
self._affinity.pop(next(iter(self._affinity)), None)
self._affinity[key] = chosen.worker_id
return chosen
if self.policy == "cost":
return min(cands, key=lambda w: w.cost_estimate(request) * (1.0 + w.load))
return min(cands, key=lambda w: w.load) # least_loaded (default)
return min(cands, key=lambda w: w.load) # least_loaded (default)
# --- serving (delegates to the chosen worker's engine) ------------------- #
async def generate(self, request: Any, *, affinity_key: str | None = None) -> Output:
v2/examples.py
+24 -18
View File
@@ -1,4 +1,4 @@
"""Worked examples (design_v3 §15), runnable as a demo:
"""Worked examples, runnable as a demo:
python3 -m v2.examples
@@ -6,8 +6,6 @@ Each prints what it demonstrates. This doubles as living documentation of the AP
"""
from __future__ import annotations
import numpy as np
from v2.recipes import build_default_engine, build_omni_engine
from v2.recipes.wan21 import build_wan21_card
from v2.parity import assert_interleave_parity
@@ -16,8 +14,7 @@ from v2.training import build_diffusion_nft
def _t2v(mid, prompt, seed, steps=4, **kw):
return make_request(TaskType.T2V, mid, prompt,
diffusion=DiffusionParams(num_steps=steps, seed=seed), **kw)
return make_request(TaskType.T2V, mid, prompt, diffusion=DiffusionParams(num_steps=steps, seed=seed), **kw)
def example_a_text_to_video(eng) -> None:
@@ -36,8 +33,8 @@ def example_b_ltx2_two_stage(eng) -> None:
def example_c_causal_streaming(eng) -> None:
print("\n(c) Causal streaming (Wan-causal): chunk rollout + slab-KV, streamable by chunk")
out = eng.run(_t2v("wan-causal-sf-1.3b", "a drone flight over mountains", 3,
outputs=OutputSpec(stream={"video": True})))
out = eng.run(
_t2v("wan-causal-sf-1.3b", "a drone flight over mountains", 3, outputs=OutputSpec(stream={"video": True})))
print(f" latents {out.artifacts['latents'].latent.shape} chunks={out.metrics['chunks']:.0f} "
f"streamed_chunks={out.metrics.get('stream_chunks', 0)}")
@@ -63,14 +60,20 @@ def example_d_rl_rollout() -> None:
def example_g_omni_mot() -> None:
print("\n(g) Omni / MoT (§16): ONE resident instance runs AR + diffusion loops on shared weights")
eng = build_omni_engine()
o = eng.run(make_request(TaskType.T2V, "cosmos3-vfm", "a phoenix",
sampling=SamplingParams(max_tokens=6, seed=1),
diffusion=DiffusionParams(num_steps=4, seed=1)))
o = eng.run(
make_request(TaskType.T2V,
"cosmos3-vfm",
"a phoenix",
sampling=SamplingParams(max_tokens=6, seed=1),
diffusion=DiffusionParams(num_steps=4, seed=1)))
print(f" Cosmos3 (reason→joint denoise): text={o.artifacts['text'].text} "
f"video={o.artifacts['video'].frames.shape}")
o2 = eng.run(make_request(TaskType.T2I, "bagel-mot", "a teapot",
sampling=SamplingParams(max_tokens=6, seed=2),
diffusion=DiffusionParams(num_steps=4, seed=2)))
o2 = eng.run(
make_request(TaskType.T2I,
"bagel-mot",
"a teapot",
sampling=SamplingParams(max_tokens=6, seed=2),
diffusion=DiffusionParams(num_steps=4, seed=2)))
print(f" BAGEL (generate_text→generate_image): text={o2.artifacts['text'].text} "
f"image={o2.artifacts['image'].tensor.shape}")
print(f" scheduler priced BOTH WorkUnit kinds (runtime-visible, not one opaque stage): "
@@ -94,8 +97,8 @@ async def _serving_demo() -> None:
pools = PoolSet(wan_t2v_disaggregated(), card)
pools.warmup()
ae.register_disaggregated("wan-disagg", pools, build_wan_t2v_program())
out = await ae.generate(make_request(TaskType.T2V, "wan-disagg", "a wave",
diffusion=DiffusionParams(num_steps=4, seed=1)))
out = await ae.generate(
make_request(TaskType.T2V, "wan-disagg", "a wave", diffusion=DiffusionParams(num_steps=4, seed=1)))
print(f" disaggregated T2V (enc→den→dec): video={out.artifacts['video'].frames.shape} "
f"cross-pool transfers={out.metrics['transfers']:.0f}")
@@ -107,14 +110,17 @@ async def _serving_demo() -> None:
r, w = await asyncio.open_connection(host, port)
w.write(f"{method} {path} HTTP/1.1\r\nHost: x\r\nContent-Length: {len(body)}\r\n\r\n".encode() + body)
await w.drain()
data = await r.read(); w.close()
data = await r.read()
w.close()
return data.decode("utf-8", "replace")
health = await http("GET", "/health")
sse = await http("POST", "/v1/chat/completions",
b'{"model":"cosmos3-vfm","messages":[{"role":"user","content":"a comet"}],"stream":true}')
n_chunks = sse.count("data: ")
print(f" OpenAI server: /health ok={'healthy' in health}; chat SSE streamed {n_chunks} chunks (omni reason→denoise)")
print(
f" OpenAI server: /health ok={'healthy' in health}; chat SSE streamed {n_chunks} chunks (omni reason→denoise)"
)
await server.close()
# our own fleet router + Dynamo adapter (frontable, not relied upon) — same DeploymentCard
@@ -136,7 +142,7 @@ def example_h_serving_and_fleet() -> None:
def main() -> None:
print("=" * 78)
print("v2 — worked examples (design_v3 §6,§13-16). One runtime, many loops.")
print("v2 — worked examples. One runtime, many loops.")
print("=" * 78)
eng = build_default_engine()
print(f"registered (recipe, runtime) cards: {list(eng._registry)}")
v2/extend/__init__.py
+5 -3
View File
@@ -1,4 +1,4 @@
"""Extension plane — observers (read-only) and interceptors (compute-altering) (design_v3 §11)."""
"""Extension plane — observers (read-only) and interceptors (compute-altering)."""
from __future__ import annotations
from v2.extend.base import Interceptor, InterceptorChain, InterceptorConflict, Observer, ObserverBus
@@ -6,5 +6,7 @@ from v2.extend.interceptors import ResidualSkipInterceptor
from v2.extend.observers import NaNWatch, Profiler
from v2.extend.registry import PluginRegistry
__all__ = ["Observer", "Interceptor", "ObserverBus", "InterceptorChain", "InterceptorConflict",
"Profiler", "NaNWatch", "ResidualSkipInterceptor", "PluginRegistry"]
__all__ = [
"Observer", "Interceptor", "ObserverBus", "InterceptorChain", "InterceptorConflict", "Profiler", "NaNWatch",
"ResidualSkipInterceptor", "PluginRegistry"
]
v2/extend/base.py
+14 -10
View File
@@ -1,6 +1,6 @@
"""Observers and interceptors (design_v3 §11) — the optimization/debug/parity surface.
"""Observers and interceptors — the optimization/debug/parity surface.
They compose with §5 cleanly: the hooks wrap ``ctx.execute(plan)``.
They compose with the loop cleanly: the hooks wrap ``ctx.execute(plan)``.
* Observers (read-only): cannot mutate state. An unused hook is *literally absent* from
the hot path (the bus only iterates when observers are attached).
* Interceptors (compute-altering): ``before_step`` may supply a cached prediction (skip);
@@ -18,7 +18,9 @@ from typing import Any, Protocol, runtime_checkable
@runtime_checkable
class Observer(Protocol):
def observe(self, event: str, **kw) -> None: ...
def observe(self, event: str, **kw) -> None:
...
@runtime_checkable
@@ -27,8 +29,11 @@ class Interceptor(Protocol):
distribution_altering: bool
graph_safe: bool
def before_step(self, plan: Any, state: Any) -> Any | None: ... # return override output to skip, or None
def after_step(self, plan: Any, state: Any, result: Any) -> None: ...
def before_step(self, plan: Any, state: Any) -> Any | None:
... # return override output to skip, or None
def after_step(self, plan: Any, state: Any, result: Any) -> None:
...
class ObserverBus:
@@ -45,7 +50,7 @@ class ObserverBus:
return bool(self._observers)
def emit(self, event: str, **kw) -> None:
if not self._observers: # absent from the hot path when off
if not self._observers: # absent from the hot path when off
return
for obs in self._observers:
obs.observe(event, **kw)
@@ -56,7 +61,7 @@ class InterceptorConflict(ValueError):
class InterceptorChain:
"""Ordered interceptor chain; conflicting interceptors rejected pre-flight (§11)."""
"""Ordered interceptor chain; conflicting interceptors rejected pre-flight."""
def __init__(self, interceptors: list[Interceptor] | None = None, *, exact_mode: bool = False):
self._chain = list(interceptors or [])
@@ -65,9 +70,8 @@ class InterceptorChain:
def _validate(self, exact_mode: bool) -> None:
skippers = [i for i in self._chain if getattr(i, "distribution_altering", False)]
if len(skippers) > 1:
raise InterceptorConflict(
f"multiple distribution-altering interceptors {[i.plugin_id for i in skippers]} "
"conflict — only one step-skipper allowed (design_v3 §11)")
raise InterceptorConflict(f"multiple distribution-altering interceptors {[i.plugin_id for i in skippers]} "
"conflict — only one step-skipper allowed")
if exact_mode and skippers:
raise InterceptorConflict(
f"exact-mode rejects distribution_altering interceptors {[i.plugin_id for i in skippers]}")
v2/extend/interceptors.py
+2 -2
View File
@@ -1,9 +1,9 @@
"""cache-dit-style interceptors (design_v3 §11).
"""cache-dit-style interceptors.
``ResidualSkipInterceptor`` is the reference step-skip integration. The load-bearing
correctness property: its per-step state lives in ``LoopState.plugin_state[id][branch]``,
keyed per request AND per CFG branch — NOT a module global. This is exactly why the
interleave gate (§9.3) passes with it on: two interleaved requests have disjoint plugin
interleave gate passes with it on: two interleaved requests have disjoint plugin
state, so neither smears the other's cached prediction.
A 4-step distilled card *rejects* this interceptor (capability negotiation) rather than
v2/extend/observers.py
+11 -14
View File
@@ -1,32 +1,29 @@
"""Built-in read-only observers (design_v3 §11): Profiler, NaNWatch.
"""Built-in read-only observers: Profiler, NaNWatch.
ParityAligner lives in ``parity/`` (it is a first-class package, §18) but also implements
the Observer ``observe`` protocol so it attaches to the same bus.
ParityAligner lives in its own ``parity/`` package but also implements the Observer ``observe``
protocol so it attaches to the same bus.
"""
from __future__ import annotations
from typing import Any
import numpy as np
class Profiler:
"""Per-step wall/CUDA timing; calibrates the cost model (design_v3 §6.1, §11).
"""Per-step wall/CUDA timing that calibrates the cost model.
Here it accumulates (work_units, actual_seconds) samples and can fit a card's
CostModel coefficients — the online calibration the conservative baseline refines.
Accumulates (work_units, actual_seconds) samples and fits a card's CostModel
coefficients — the online calibration that refines the conservative baseline.
"""
def __init__(self) -> None:
self.samples: list[tuple[str, float, float]] = [] # (batch_key_repr, work_units, seconds)
self.samples: list[tuple[str, float, float]] = [] # (batch_key_repr, work_units, seconds)
def observe(self, event: str, **kw) -> None:
if event == "step_complete":
plan, result = kw.get("plan"), kw.get("result")
if plan is not None and result is not None:
self.samples.append((repr(plan.shape_sig.batch_key),
float(plan.shape_sig.work_units),
float(result.actual_seconds)))
self.samples.append(
(repr(plan.shape_sig.batch_key), float(plan.shape_sig.work_units), float(result.actual_seconds)))
def calibrate(self, cost_model) -> None:
"""Fit base + per_unit seconds from observed samples (least squares)."""
@@ -43,10 +40,10 @@ class Profiler:
class NaNWatch:
"""First-NaN/Inf localization (design_v3 §11). Request-fatal → SPMD-consistent abort (§6.4)."""
"""First-NaN/Inf localization. Request-fatal, triggering an SPMD-consistent abort."""
def __init__(self) -> None:
self.first: tuple[str, int] | None = None # (tap/output name, step)
self.first: tuple[str, int] | None = None # (tap/output name, step)
def observe(self, event: str, **kw) -> None:
if event == "step_complete" and self.first is None:
v2/extend/registry.py
+8 -6
View File
@@ -1,21 +1,23 @@
"""Plugin registry + trust boundary (design_v3 §11).
"""Plugin registry + trust boundary.
> Plugins are enabled at deploy scope only (never a per-request ``plugins=[...]`` field
> that would wire third-party code selection into the public API); requests only
> *parameterize* pre-enabled plugins through validated schemas.
Plugins are enabled at deploy scope only (never a per-request ``plugins=[...]`` field that
would wire third-party code selection into the public API); requests only *parameterize*
pre-enabled plugins through validated schemas.
"""
from __future__ import annotations
from typing import Any, Callable
from typing import Any
from collections.abc import Callable
from v2.extend.base import InterceptorChain, Observer
class PluginRegistry:
def __init__(self) -> None:
self._interceptors: dict[str, Callable[..., Any]] = {}
self._observers: dict[str, Callable[..., Observer]] = {}
self._enabled: set[str] = set() # deploy-scope enablement
self._enabled: set[str] = set() # deploy-scope enablement
def register_interceptor(self, plugin_id: str, factory: Callable[..., Any]) -> None:
self._interceptors[plugin_id] = factory
v2/fastvideo_args.py
+1 -1
View File
@@ -1,4 +1,4 @@
# STUB: re-exports fastvideo until vendored (see memory: v2-vendoring-approach).
"""``FastVideoArgs`` facade — the args object the component loaders consume. Re-exported so v2 code imports
``v2.fastvideo_args``; a vendored cutover slims it to ~250 LOC of inference-only args."""
``v2.fastvideo_args``; a vendored cutover slims it to inference-only args."""
from fastvideo.fastvideo_args import FastVideoArgs # noqa: F401
v2/loop/__init__.py
+21 -5
View File
@@ -1,4 +1,4 @@
"""Loop plane — driven loops + policies (design_v3 §5)."""
"""Loop plane — driven loops + policies."""
from __future__ import annotations
from v2._enums import LoopKind, WorkUnitKind
@@ -21,8 +21,24 @@ from v2.loop.driver import LoopRunner
from v2.loop.sampler import add_noise, build_flow_sigmas, flow_match_euler_step, x0_from_velocity
__all__ = [
"Loop", "LoopContext", "LoopState", "LoopResult", "WorkPlan", "Done", "StepResult",
"StepContext", "ShapeSignature", "ResourceRequest", "CachePlan", "CacheOp", "PlacementHint",
"LoopRunner", "LoopKind", "WorkUnitKind",
"build_flow_sigmas", "flow_match_euler_step", "x0_from_velocity", "add_noise",
"Loop",
"LoopContext",
"LoopState",
"LoopResult",
"WorkPlan",
"Done",
"StepResult",
"StepContext",
"ShapeSignature",
"ResourceRequest",
"CachePlan",
"CacheOp",
"PlacementHint",
"LoopRunner",
"LoopKind",
"WorkUnitKind",
"build_flow_sigmas",
"flow_match_euler_step",
"x0_from_velocity",
"add_noise",
]
v2/loop/contracts.py
+48 -36
View File
@@ -1,6 +1,4 @@
"""The driven-loop contract (design_v3 §5) — the keystone.
A loop is a **serializable state machine the runtime drives**:
"""The driven-loop contract — a serializable state machine the runtime drives:
state = loop.init(req, model_state, ctx)
while True:
@@ -11,13 +9,12 @@ A loop is a **serializable state machine the runtime drives**:
for chunk in plan.emits: ctx.emit(chunk)
return loop.finalize(state)
Two properties this contract buys (design_v3 §5.1):
Two properties this contract buys:
* content-adaptive steps are natural — ``next`` reads ``state``, which already folded
in the previous ``StepResult`` via ``advance``;
* cross-request state safety is *structural* — all per-request mutable state lives in
``LoopState`` (incl. ``plugin_state`` per request/CFG-branch), never module globals,
so interleaving requests through one ModelInstance cannot smear state. This is what
the batch-of-N interleave gate (§9.3) verifies.
so interleaving requests through one ModelInstance cannot smear state.
This module is pure stdlib (tensors typed as ``TensorLike``) so it imports no backend.
"""
@@ -36,11 +33,11 @@ from v2.request.streams import StreamChunk
# --------------------------------------------------------------------------- #
@dataclass(frozen=True)
class ShapeSignature:
"""Batch-compatibility + graph-capture key (design_v3 §6.2)."""
"""Batch-compatibility + graph-capture key."""
kind: WorkUnitKind
dims: tuple[int, ...] = ()
dtype: str = "float32"
extra: tuple[tuple[str, Any], ...] = () # e.g. (("cfg","classic"),("expert","e0"))
extra: tuple[tuple[str, Any], ...] = () # e.g. (("cfg","classic"),("expert","e0"))
@property
def work_units(self) -> int:
@@ -51,13 +48,13 @@ class ShapeSignature:
@property
def batch_key(self) -> tuple:
"""Two WorkPlans batch together iff their batch_keys are equal (§6.3 BatchScheduler)."""
"""Two WorkPlans batch together iff their batch_keys are equal."""
return (self.kind, self.dims, self.dtype, self.extra)
@dataclass
class ResourceRequest:
"""Everything admission must reserve (design_v3 §6.2 admission rule)."""
"""Everything admission must reserve."""
compute_seconds: float = 0.0
resident_bytes: int = 0
peak_activation_bytes: int = 0
@@ -69,7 +66,7 @@ class ResourceRequest:
@dataclass
class CacheOp:
cache_class: str
key: Any # CacheKey
key: Any # CacheKey
nbytes: int = 0
@@ -91,7 +88,7 @@ class PlacementHint:
# --------------------------------------------------------------------------- #
@dataclass
class WorkPlan:
"""A typed description of the next step (design_v3 §6.2).
"""A typed description of the next step.
``run`` is the kernel thunk built by ``next()`` but NOT called there (kernel-free
planning). The runtime's ``ctx.execute`` calls it — possibly after batching with
@@ -105,14 +102,14 @@ class WorkPlan:
cache: CachePlan = field(default_factory=CachePlan)
placement: PlacementHint = field(default_factory=PlacementHint)
emits: list[StreamChunk] = field(default_factory=list)
payload: dict[str, Any] = field(default_factory=dict) # inspectable inputs (debug/parity)
payload: dict[str, Any] = field(default_factory=dict) # inspectable inputs (debug/parity)
# the kernel thunk: run(model_instance, override=None) -> StepResult | dict.
# Built by next() but NOT called there (kernel-free planning). ``override`` is an optional
# interceptor-supplied forward result (e.g. a cached prediction); the step body still runs
# the cheap solver step with it.
run: Any = None
label: str = ""
# --- piecewise CUDA-graph capture (design_v3 §6.2; Path A) ----------------------------------- #
# --- piecewise CUDA-graph capture ------------------------------------------------------------ #
# ``capturable``: the model's static declaration that this step is safe to capture/replay — i.e.
# no host RNG / data-dependent control flow inside the captured region. A stochastic step (the
# FlowGRPO SDE rollout) sets this False, forcing the runtime to eager-break it.
@@ -121,7 +118,7 @@ class WorkPlan:
# the captured graph's *shape of computation*. Part of the capture key so a step with a different
# branch set / expert never replays an incompatible graph. Empty ⇒ structure fixed by shape alone.
graph_key: tuple = ()
# The static-buffer capture form (design_v3 §6.2). A capturable step exposes its deterministic op
# The static-buffer capture form. A capturable step exposes its deterministic op
# structure as ``graph_fn(model, workspace) -> StepResult`` reading EVERY per-step-varying input
# (latent, sigmas, conditioning) from ``workspace`` — never from closure — plus ``graph_inputs``,
# the dict of those current values. The runtime allocates address-stable buffers once per key and
@@ -136,22 +133,22 @@ class WorkPlan:
class Done:
"""Sentinel returned by ``next`` when the loop is finished. ``finalize`` produces the
actual LoopResult; ``result`` here is optional (kept for loops that want to carry it)."""
result: "LoopResult | None" = None
result: LoopResult | None = None
@dataclass
class StepResult:
output: dict[str, Any] = field(default_factory=dict) # typed per-loop (e.g. {"noise_pred": ...})
output: dict[str, Any] = field(default_factory=dict) # typed per-loop (e.g. {"noise_pred": ...})
actual_seconds: float = 0.0
cache_writes: list[CacheOp] = field(default_factory=list)
behavior: Any = None # BehaviorRecord slice (rollout profile)
behavior: Any = None # BehaviorRecord slice (rollout profile)
@dataclass
class LoopResult:
outputs: dict[str, Any] = field(default_factory=dict) # final latents / tokens / artifacts
outputs: dict[str, Any] = field(default_factory=dict) # final latents / tokens / artifacts
metrics: dict[str, float] = field(default_factory=dict)
behavior: Any = None # full BehaviorRecord (rollout)
behavior: Any = None # full BehaviorRecord (rollout)
# --------------------------------------------------------------------------- #
@@ -159,23 +156,23 @@ class LoopResult:
# --------------------------------------------------------------------------- #
@dataclass
class LoopState:
"""All per-request mutable state lives here (design_v3 §5.1 structural safety)."""
"""All per-request mutable state lives here (structural cross-request safety)."""
loop_id: str
instance_id: str
request_id: str
profile: ExecutionProfile = ExecutionProfile.SERVE
step_idx: int = 0
done: bool = False
rng: Any = None # seeded numpy Generator (per request)
rng: Any = None # seeded numpy Generator (per request)
seed: int | None = None
# common typed fields (resolved at init)
latents: dict[str, TensorLike] = field(default_factory=dict)
cond: dict[str, Any] = field(default_factory=dict) # conditioning written by ConditioningInjector
cond: dict[str, Any] = field(default_factory=dict) # conditioning written by ConditioningInjector
timesteps: list[float] = field(default_factory=list)
sigmas: list[float] = field(default_factory=list)
# per-model extension (Cosmos3PackedSeq, MatrixGameState, ...) — typed by LoopSpec.extension_schema
extension: Any = None
# interceptor/policy state, keyed per plugin id AND per CFG branch (design_v3 §5.1, §11)
# interceptor/policy state, keyed per plugin id AND per CFG branch
plugin_state: dict[str, Any] = field(default_factory=dict)
cache_handles: dict[str, Any] = field(default_factory=dict)
# capture buffers (rollout): trajectory of per-step records
@@ -184,15 +181,15 @@ class LoopState:
# --------------------------------------------------------------------------- #
# StepContext — what policies read each step (design_v3 §6.2.3 rule 2) #
# StepContext — what policies read each step #
# --------------------------------------------------------------------------- #
@dataclass
class StepContext:
step_idx: int
timestep: float
sigma: float
branch: str = "cond" # current guidance branch
active_expert_id: str | None = None # set by ExpertRouting; observed by AdaptiveGateCFG
branch: str = "cond" # current guidance branch
active_expert_id: str | None = None # set by ExpertRouting; observed by AdaptiveGateCFG
sampler_coeffs: dict[str, float] = field(default_factory=dict)
extra: dict[str, Any] = field(default_factory=dict)
@@ -206,16 +203,31 @@ class LoopContext(Protocol):
scheduler; the scheduler never sees the model's math."""
profile: ExecutionProfile
def execute(self, plan: WorkPlan) -> StepResult: ... # THE INVERSION POINT
def emit(self, chunk: StreamChunk) -> None: ...
def check_cancel(self) -> None: ... # raises request.Cancelled at boundary
def observe(self, event: str, **kw) -> None: ... # observer bus hook
def execute(self, plan: WorkPlan) -> StepResult:
... # THE INVERSION POINT
def emit(self, chunk: StreamChunk) -> None:
...
def check_cancel(self) -> None:
... # raises request.Cancelled at boundary
def observe(self, event: str, **kw) -> None:
... # observer bus hook
@runtime_checkable
class Loop(Protocol):
"""The model-owned control flow (design_v3 §5.1). Four methods; ``next`` is kernel-free."""
def init(self, req: Any, model: Any, ctx: LoopContext) -> LoopState: ...
def next(self, state: LoopState) -> "WorkPlan | Done": ...
def advance(self, state: LoopState, result: StepResult) -> LoopState: ...
def finalize(self, state: LoopState) -> LoopResult: ...
"""The model-owned control flow. Four methods; ``next`` is kernel-free."""
def init(self, req: Any, model: Any, ctx: LoopContext) -> LoopState:
...
def next(self, state: LoopState) -> WorkPlan | Done:
...
def advance(self, state: LoopState, result: StepResult) -> LoopState:
...
def finalize(self, state: LoopState) -> LoopResult:
...
v2/loop/driver.py
+4 -4
View File
@@ -1,10 +1,9 @@
"""LoopRunner — the only place iteration lives (design_v3 §5.1).
"""LoopRunner — the only place iteration lives.
This is the runtime's driver, factored so the engine can either run it to completion
The runtime's driver, factored so the engine can either run it to completion
(``run``) or advance it one step at a time (``peek`` + ``step``) to **interleave** the steps
of concurrent requests. Per-request state is entirely in ``LoopState``; the runner holds no
hidden iteration state beyond a cached pending plan, so interleaving is safe by construction
(the property the §9.3 interleave gate verifies).
hidden iteration state beyond a cached pending plan, so interleaving is safe by construction.
"""
from __future__ import annotations
@@ -14,6 +13,7 @@ from v2.loop.contracts import Done, LoopContext, LoopResult, LoopState, WorkPlan
class LoopRunner:
def __init__(self, loop: Any, ctx: LoopContext, request: Any, model: Any):
self.loop = loop
self.ctx = ctx
v2/loop/policies/__init__.py
+14 -4
View File
@@ -1,4 +1,4 @@
"""Policies — the default step decomposition (design_v3 §5.1, §6.2.3)."""
"""Policies — the default step decomposition."""
from __future__ import annotations
from v2.loop.policies.base import (
@@ -20,7 +20,17 @@ from v2.loop.policies.cfg import (
)
__all__ = [
"CFGPolicy", "ClassicCFG", "BatchedCFG", "EmbeddedGuidance", "AdaptiveGateCFG", "PerModalityCFG",
"FlowShiftPolicy", "PrecisionPolicy", "ExpertRouting", "NoRouting", "BoundaryTimestepRouting",
"ConditioningInjector", "PassthroughConditioning",
"CFGPolicy",
"ClassicCFG",
"BatchedCFG",
"EmbeddedGuidance",
"AdaptiveGateCFG",
"PerModalityCFG",
"FlowShiftPolicy",
"PrecisionPolicy",
"ExpertRouting",
"NoRouting",
"BoundaryTimestepRouting",
"ConditioningInjector",
"PassthroughConditioning",
]
v2/loop/policies/base.py
+20 -15
View File
@@ -1,12 +1,10 @@
"""Policy base classes (design_v3 §5.1, §6.2.3) — the *default* step decomposition.
"""Policy base classes — the *default* step decomposition.
> Policies (CFG, expert routing, precision, flow-shift, conditioning) are the default
> decomposition that deletes duplication for the families that fit — never an admission
> requirement.
A family whose math is braided ships a custom ``next``/``advance`` and uses these as a
library (LTX-2's multi-pass guidance does exactly that). Policy *bindings* resolve at build;
policy *state* is per-request in ``LoopState`` (the adaptive-gate cached delta).
Policies (CFG, expert routing, precision, flow-shift, conditioning) delete duplication for
the families that fit; they are never required. A family whose math is braided ships a custom
``next``/``advance`` and uses these as a library (LTX-2's multi-pass guidance does exactly
that). Policy *bindings* resolve at build; policy *state* is per-request in ``LoopState``
(the adaptive-gate cached delta).
"""
from __future__ import annotations
@@ -23,7 +21,7 @@ from v2.loop.sampler import build_flow_sigmas
# FlowShiftPolicy — resolution-bucket shift + sigma schedule #
# --------------------------------------------------------------------------- #
class FlowShiftPolicy:
"""config-driven flow-shift lookup (design_v3 §6.2.3; Wan 480p shift=3.0, 720p=5.0)."""
"""config-driven flow-shift lookup (e.g. Wan 480p shift=3.0, 720p=5.0)."""
def __init__(self, shift: float = 3.0, bucket_lookup: dict[int, float] | None = None):
self.shift = shift
@@ -32,9 +30,12 @@ class FlowShiftPolicy:
def shift_for(self, height: int = 0, width: int = 0) -> float:
return self.bucket_lookup.get(height * width, self.shift)
def build_schedule(self, num_steps: int, height: int = 0, width: int = 0,
def build_schedule(self,
num_steps: int,
height: int = 0,
width: int = 0,
sigmas: list[float] | None = None) -> np.ndarray:
if sigmas is not None: # explicit distilled schedule (LTX-2)
if sigmas is not None: # explicit distilled schedule (LTX-2)
return np.asarray(sigmas, dtype=np.float64)
return build_flow_sigmas(num_steps, shift=self.shift_for(height, width))
@@ -43,7 +44,7 @@ class FlowShiftPolicy:
# PrecisionPolicy — autocast / scheduler-step dtype control #
# --------------------------------------------------------------------------- #
class PrecisionPolicy:
"""Replaces ``prefix=='Flux'`` autocast hacks + ``scheduler_step_in_fp32`` (§6.2.3)."""
"""Replaces ``prefix=='Flux'`` autocast hacks + ``scheduler_step_in_fp32``."""
def __init__(self, compute_dtype: str = "float32", scheduler_step_in_fp32: bool = True):
self.compute_dtype = compute_dtype
@@ -61,8 +62,10 @@ class PrecisionPolicy:
# ExpertRouting — Wan2.2 boundary-timestep transformer switch #
# --------------------------------------------------------------------------- #
class ExpertRouting(ABC):
@abstractmethod
def expert_for(self, ctx: StepContext) -> str: ...
def expert_for(self, ctx: StepContext) -> str:
...
class NoRouting(ExpertRouting):
@@ -76,7 +79,7 @@ class NoRouting(ExpertRouting):
class BoundaryTimestepRouting(ExpertRouting):
"""Wan2.2 ``boundary_ratio`` switch between two transformers (design_v3 §6.2.3)."""
"""Wan2.2 ``boundary_ratio`` switch between two transformers."""
def __init__(self, high_noise: str, low_noise: str, boundary: float = 0.5):
self.high_noise = high_noise
@@ -91,8 +94,10 @@ class BoundaryTimestepRouting(ExpertRouting):
# ConditioningInjector — writes RequestState.cond; the loop stays agnostic #
# --------------------------------------------------------------------------- #
class ConditioningInjector(ABC):
@abstractmethod
def inject(self, state: Any, request: Any) -> None: ...
def inject(self, state: Any, request: Any) -> None:
...
class PassthroughConditioning(ConditioningInjector):
v2/loop/policies/cfg.py
+12 -14
View File
@@ -1,11 +1,11 @@
"""CFGPolicy — one taxonomy over one shared denoise body (design_v3 §5.3).
"""CFGPolicy — one taxonomy over one shared denoise body.
> Batched-vs-two-forward is a *dispatch detail inside one policy*, not a separate mechanism.
Batched-vs-two-forward is a dispatch detail inside one policy, not a separate mechanism.
The step body asks ``branches_this_step(ctx, state)`` which forwards to run, runs them, then
calls ``combine(preds, scale, ctx, state)``. Per-request mutable state (the adaptive-gate
cached delta) lives in the ``state`` dict, which the step body slices out of
``LoopState.plugin_state`` — never a module global (the §5.1 safety property).
``LoopState.plugin_state`` — never a module global.
cfg-parallel is a *parallelism axis* (not a policy); companions are an *orchestrator pattern*
(not in the loop). Both compose with any policy here.
@@ -13,7 +13,6 @@ cfg-parallel is a *parallelism axis* (not a policy); companions are an *orchestr
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Any
import numpy as np
@@ -21,15 +20,15 @@ from v2.loop.contracts import StepContext
class CFGPolicy(ABC):
batched: bool = False # dispatch detail: stack branches into one forward
branch_vocabulary: list[str] = ["cond", "uncond"] # full declared set (for per-branch plugin state)
batched: bool = False # dispatch detail: stack branches into one forward
branch_vocabulary: list[str] = ["cond", "uncond"] # full declared set (for per-branch plugin state)
def branches_this_step(self, ctx: StepContext, state: dict) -> list[str]:
return list(self.branch_vocabulary)
@abstractmethod
def combine(self, preds: dict[str, np.ndarray], guidance_scale: float,
ctx: StepContext, state: dict) -> np.ndarray: ...
def combine(self, preds: dict[str, np.ndarray], guidance_scale: float, ctx: StepContext, state: dict) -> np.ndarray:
...
class ClassicCFG(CFGPolicy):
@@ -52,7 +51,7 @@ class BatchedCFG(ClassicCFG):
class EmbeddedGuidance(CFGPolicy):
"""Degenerate single-branch identity-combine (Flux): guidance rides in the forward kwarg.
This is *not* 'no CFG' — it is kept inside the same abstraction (design_v3 §5.3)."""
This is *not* 'no CFG' — it is kept inside the same abstraction."""
branch_vocabulary = ["cond"]
def combine(self, preds, guidance_scale, ctx, state):
@@ -60,7 +59,7 @@ class EmbeddedGuidance(CFGPolicy):
class AdaptiveGateCFG(CFGPolicy):
"""Cached-delta reuse with expert-switch self-invalidation (design_v3 §5.1 canonical state).
"""Cached-delta reuse with expert-switch self-invalidation.
On reuse steps it runs ONLY the cond branch and reuses the cached delta:
``out = cond + (s−1)·delta`` (algebraically identical to ``uncond + s·(cond−uncond)``).
@@ -71,20 +70,19 @@ class AdaptiveGateCFG(CFGPolicy):
self.interval = max(1, interval)
def _recompute(self, ctx, state) -> bool:
return (ctx.step_idx % self.interval == 0
or "delta" not in state
return (ctx.step_idx % self.interval == 0 or "delta" not in state
or state.get("expert") != ctx.active_expert_id)
def branches_this_step(self, ctx, state):
return ["cond", "uncond"] if self._recompute(ctx, state) else ["cond"]
def combine(self, preds, guidance_scale, ctx, state):
if "uncond" in preds: # recompute path
if "uncond" in preds: # recompute path
delta = preds["cond"] - preds["uncond"]
state["delta"] = delta
state["expert"] = ctx.active_expert_id
return preds["uncond"] + guidance_scale * delta
delta = state["delta"] # reuse path (skipped the uncond forward)
delta = state["delta"] # reuse path (skipped the uncond forward)
return preds["cond"] + (guidance_scale - 1.0) * delta
v2/loop/sampler.py
+9 -9
View File
@@ -1,12 +1,12 @@
"""Samplers as a small library (design_v3 §5.1: families use samplers as a library).
"""Samplers as a small library; model families compose them.
Flow-match Euler is shared by the Wan and LTX-2 denoise step bodies. On a GPU box the real
UniPC multistep / distilled schedulers are wrapped by the torch component adapters; this
numpy form is the CPU-testable, bit-reproducible reference used by the loop tests.
Flow-match Euler is shared by the Wan and LTX-2 denoise step bodies. On GPU the real UniPC
multistep / distilled schedulers are wrapped by the torch component adapters; this numpy form
is the CPU-testable, bit-reproducible reference used by the loop tests.
Flow-match interpolation: ``x_t = (1-σ_t)·x0 + σ_t·ε``; the model predicts velocity
``v = ε - x0``. Then a deterministic Euler step to σ_next is simply ``x_next = x_t + (σ_next-σ_t)·v``
(this is exactly LTX-2's ``velocity=(latents-denoised)/σ; latents += velocity·dt``).
``v = ε - x0``. A deterministic Euler step to σ_next is ``x_next = x_t + (σ_next-σ_t)·v``
(exactly LTX-2's ``velocity=(latents-denoised)/σ; latents += velocity·dt``).
"""
from __future__ import annotations
@@ -67,7 +67,7 @@ def flow_sde_step_with_logprob(x_t,
noise=None,
prev_sample=None,
noise_scale: float = 0.7):
"""FlowGRPO-style stochastic step + Gaussian log-prob (PromptRL/UniRL §6; FlowGRPO lineage).
"""FlowGRPO-style stochastic step + Gaussian log-prob.
The RL *rollout* sampler (vs the deterministic ODE ``flow_match_euler_step`` used at serve time).
Injects noise for exploration and returns the log-prob of the realized sample under the step's
@@ -95,8 +95,8 @@ def flow_sde_step_with_logprob(x_t,
def flow_sde_ml_velocity(x_t, sample, sigma_t: float, sigma_next: float, *, noise_scale: float = 0.7):
"""The velocity whose deterministic SDE mean lands exactly on ``sample`` (the max-likelihood
velocity for a realized FlowGRPO sample). Moving the policy toward it, advantage-weighted, is the
policy-gradient direction — and is nonzero even at ratio==1, so it is the *correct* FlowGRPO update
surrogate (nudging toward the velocity the model already produced would be a no-op). PromptRL §6.
policy-gradient direction — and is nonzero even at ratio==1, so it is the correct FlowGRPO update
surrogate (nudging toward the velocity the model already produced would be a no-op).
"""
x_t = np.asarray(x_t, dtype=np.float64)
sample = np.asarray(sample, dtype=np.float64)
v2/memory/__init__.py
+1 -1
View File
@@ -1,4 +1,4 @@
"""Memory plane — reservation before admission, sleep/wake by tag (design_v3 §7.3)."""
"""Memory plane — reservation before admission, sleep/wake by tag."""
from __future__ import annotations
from v2.memory.allocator import MemoryManager, OutOfMemory, Reservation
v2/memory/allocator.py
+6 -5
View File
@@ -1,15 +1,15 @@
"""Tagged memory pools with reservation-before-admission (design_v3 §6.2, §7.3).
"""Tagged memory pools with reservation-before-admission.
> Reservation must be pre-flight. A scheduler that admits a diffusion step and then
> discovers it cannot allocate the VAE tile is wrong.
Reservation must be pre-flight: a scheduler that admits a diffusion step and then discovers it
cannot allocate the VAE tile is wrong.
Sleep/wake is component-granular (tags are component names) for RL: drop DiT + caches,
keep VAE/text-encoder resident (design_v3 §7.3; CuMem-style).
keep VAE/text-encoder resident (CuMem-style).
"""
from __future__ import annotations
import itertools
from dataclasses import dataclass, field
from dataclasses import dataclass
class OutOfMemory(Exception):
@@ -28,6 +28,7 @@ class Reservation:
class MemoryManager:
def __init__(self, total_bytes: int = 1 << 40, per_tag_budget: dict[str, int] | None = None):
self.total_bytes = total_bytes
self.per_tag_budget = dict(per_tag_budget or {})
v2/models/dits/lingbotworld.py
+4 -5
View File
@@ -1,11 +1,10 @@
# STUB: re-exports fastvideo until vendored (see memory: v2-vendoring-approach).
"""LingBot-World DiT facade.
The transformer class itself is constructed by the loader from the card's ``load_id``
(``fastvideo.models.dits.lingbotworld.model:LingBotWorldTransformer3DModel``), so it is not imported
here. What the v2 recipe DOES need from the fastvideo source is the camera/Plucker embedding builder
``prepare_camera_embedding`` (poses.npy + intrinsics.npy -> ``c2ws_plucker_emb [B, 6*s^2, F, H, W]``),
re-exported so the program's camera node can build the tensor without reaching into the model package.
The transformer class is constructed by the loader from the card's ``load_id``, so it is not imported
here. This stub only re-exports the camera/Plucker embedding builder ``prepare_camera_embedding``
(poses.npy + intrinsics.npy -> ``c2ws_plucker_emb [B, 6*s^2, F, H, W]``) so the program's camera node
can build the tensor without reaching into the model package.
"""
from fastvideo.models.dits.lingbotworld.cam_utils import ( # noqa: F401
prepare_camera_embedding, )
v2/parallel/__init__.py
+2 -3
View File
@@ -1,9 +1,8 @@
"""Parallelism plane (design_v3 §8) — named axes → validated mesh, part of the cache key."""
"""Parallelism plane — named axes -> validated mesh, part of the cache key."""
from __future__ import annotations
from v2.parallel.mesh import FakeDeviceMesh, build_mesh
from v2.parallel.plan import AXIS_NAMES, ParallelPlan
from v2.parallel.validation import ParallelValidationError, validate_plan
__all__ = ["ParallelPlan", "AXIS_NAMES", "FakeDeviceMesh", "build_mesh",
"validate_plan", "ParallelValidationError"]
__all__ = ["ParallelPlan", "AXIS_NAMES", "FakeDeviceMesh", "build_mesh", "validate_plan", "ParallelValidationError"]
v2/parallel/mesh.py
+3 -5
View File
@@ -1,7 +1,4 @@
"""FakeDeviceMesh — a ParallelDims-style builder, CPU-testable (design_v3 §8; torchtitan).
> Declarative, validated, compiled to a PyTorch ``DeviceMesh`` via a ``ParallelDims``-
> style builder (product-of-degrees validation, cached submeshes).
"""FakeDeviceMesh — a torchtitan-style ParallelDims builder, CPU-testable.
On a GPU box this compiles to ``torch.distributed.device_mesh.DeviceMesh``. Here it is a
pure-Python mesh of rank tuples so topology logic is unit-tested without GPUs.
@@ -16,6 +13,7 @@ from v2.parallel.validation import validate_plan
class FakeDeviceMesh:
def __init__(self, plan: ParallelPlan):
self.plan = plan
self.order = plan.mesh_order or list(plan.axes.keys())
@@ -45,7 +43,7 @@ class FakeDeviceMesh:
def _coord_to_rank(self, coord: tuple[int, ...]) -> int:
rank = 0
for c, d in zip(coord, self.shape):
for c, d in zip(coord, self.shape, strict=False):
rank = rank * d + c
return rank
v2/parallel/plan.py
+6 -6
View File
@@ -1,4 +1,4 @@
"""ParallelPlan — parallelism as a model contract (design_v3 §8).
"""ParallelPlan — parallelism as a model contract.
Parallelism is not a launch flag; it affects cache keys, scheduling, transport,
capture, and parity, so it lives on the card. Declarative, validated, compiled to a
@@ -11,17 +11,17 @@ import hashlib
import json
from dataclasses import dataclass, field
# Canonical axis names (design_v3 §8). cfgp is ≤2.
# Canonical axis names. cfgp is <=2.
AXIS_NAMES = ("dp", "tp", "sp", "cp", "cfgp", "pp_patch", "vae", "ep", "fsdp", "role", "replica")
@dataclass
class ParallelPlan:
axes: dict[str, int] = field(default_factory=dict) # e.g. {"tp": 2, "sp": 4, "cfgp": 2}
axes: dict[str, int] = field(default_factory=dict) # e.g. {"tp": 2, "sp": 4, "cfgp": 2}
mesh_order: list[str] = field(default_factory=list)
placement: str = "colocated"
communication: dict[str, str] = field(default_factory=dict)
# design_v3 §8 + omni-readiness §4b: applicability conditions travel with axes
# applicability conditions travel with axes
applicability: dict[str, dict] = field(default_factory=dict)
per_axis_communication: dict[str, str] = field(default_factory=dict)
@@ -36,11 +36,11 @@ class ParallelPlan:
@property
def hash(self) -> str:
"""Stable hash — part of the CacheKey (design_v3 §7.1: parallel_plan_hash)."""
"""Stable hash — part of the CacheKey (parallel_plan_hash)."""
payload = json.dumps({"axes": self.axes, "order": self.mesh_order}, sort_keys=True)
return hashlib.sha256(payload.encode()).hexdigest()[:16]
@classmethod
def single(cls) -> "ParallelPlan":
def single(cls) -> ParallelPlan:
"""The default deployment: one device, all degree-one trivial groups."""
return cls(axes={"dp": 1}, mesh_order=["dp"])
v2/parallel/validation.py
+11 -10
View File
@@ -1,9 +1,8 @@
"""Pre-flight parallelism validation (design_v3 §8).
"""Pre-flight parallelism validation.
> Pre-flight or it fails at load, never halfway. Ownership conflicts are build errors
> ... Applicability conditions travel with axes: ``pp_patch`` is invalid for causal/AR.
All checks are CPU-testable on a fake mesh (design.md §6.3.4: CPU-testability).
Validate at load or fail, never halfway. Ownership conflicts are build errors, and
applicability conditions travel with axes (e.g. ``pp_patch`` is invalid for causal/AR).
All checks are CPU-testable on a fake mesh.
"""
from __future__ import annotations
@@ -19,16 +18,19 @@ class ParallelValidationError(ValueError):
_CAUSAL_KINDS = {LoopKind.CHUNK_ROLLOUT, LoopKind.AR_DECODE}
def validate_plan(plan: ParallelPlan, card=None, *, world_size: int | None = None,
def validate_plan(plan: ParallelPlan,
card=None,
*,
world_size: int | None = None,
cfg_policy_batched: bool = False) -> ParallelPlan:
"""Validate a plan, optionally against a card. Returns the plan or raises.
Checks (design_v3 §8):
Checks:
1. axis names are known; degrees >= 1
2. cfgp <= 2
3. product of degrees matches world_size (when given)
4. pp_patch invalid for any causal/AR loop on the card
5. ownership conflict: cfgp>1 AND a batched CFGPolicy is rejected (§5.3, §9 build-guard)
5. ownership conflict: cfgp>1 AND a batched CFGPolicy is rejected
"""
errs: list[str] = []
@@ -45,8 +47,7 @@ def validate_plan(plan: ParallelPlan, card=None, *, world_size: int | None = Non
errs.append(f"product of degrees {plan.world_size()} != world_size {world_size}")
if plan.degree("cfgp") > 1 and cfg_policy_batched:
errs.append("ownership conflict: a request owns a BatchedCFG policy OR a cfgp group, never both "
"(design_v3 §5.3 / §9 build-guard)")
errs.append("ownership conflict: a request owns a BatchedCFG policy OR a cfgp group, never both")
if card is not None and plan.degree("pp_patch") > 1:
causal = [lid for lid, lp in card.loops.items() if lp.kind in _CAUSAL_KINDS]
v2/parity/__init__.py
+5 -3
View File
@@ -1,4 +1,4 @@
"""Parity plane — a first-class package, not a test folder (design_v3 §9, §18).
"""Parity plane — a first-class package, not a test folder.
It is how the (recipe, runtime) pair is kept honest: parity is measured by ParityAligner,
the consistency ladder is typed, and the interleave gate is non-negotiable.
@@ -10,5 +10,7 @@ from v2.parity.aligner import ParityAligner
from v2.parity.interleave_gate import assert_interleave_parity, compare_outputs
from v2.parity.ladder import Divergence, array_diff, bit_identical, within
__all__ = ["ConsistencyLevel", "ExecutionProfile", "ParityAligner", "Divergence",
"array_diff", "bit_identical", "within", "assert_interleave_parity", "compare_outputs"]
__all__ = [
"ConsistencyLevel", "ExecutionProfile", "ParityAligner", "Divergence", "array_diff", "bit_identical", "within",
"assert_interleave_parity", "compare_outputs"
]
v2/parity/aligner.py
+7 -7
View File
@@ -1,8 +1,8 @@
"""ParityAligner — parity is measured, never assumed (design_v3 §9.1, §11).
"""ParityAligner — parity is measured, never assumed.
A read-only observer: *record mode* dumps named taps per step/block from a reference
(the official framework, or a pre-change build); *compare mode* replays with fixed seeds
and reports the **first** divergence beyond per-tap tolerance. This is the engine behind
and reports the first divergence beyond per-tap tolerance. This is the engine behind
the "old loop vs new loop, bit-identical" gate and the standing instrument for every
port, precision change, and kernel swap.
"""
@@ -20,7 +20,7 @@ class ParityAligner:
self.name = name
self.default_atol = default_atol
self.default_rtol = default_rtol
self.taps: dict[tuple[int, str], Any] = {} # (step, tap_name) -> value
self.taps: dict[tuple[int, str], Any] = {} # (step, tap_name) -> value
self.tolerances: dict[str, tuple[float, float]] = {} # tap_name -> (atol, rtol)
def set_tolerance(self, tap_name: str, atol: float = 0.0, rtol: float = 0.0) -> None:
@@ -34,19 +34,19 @@ class ParityAligner:
if event == "tap":
self.record_tap(kw.get("step", 0), kw["name"], kw["value"])
def first_divergence(self, reference: "ParityAligner",
def first_divergence(self,
reference: ParityAligner,
level: ConsistencyLevel = ConsistencyLevel.C1) -> Divergence | None:
"""Compare this (current) run against a reference, reporting the FIRST divergence
in step order beyond the per-tap tolerance."""
for (step, tap_name) in sorted(reference.taps, key=lambda k: (k[0], k[1])):
ref_val = reference.taps[(step, tap_name)]
if (step, tap_name) not in self.taps:
return Divergence(f"{tap_name}@{step}", level, float("inf"), float("inf"),
"tap missing in current run")
return Divergence(f"{tap_name}@{step}", level, float("inf"), float("inf"), "tap missing in current run")
cur_val = self.taps[(step, tap_name)]
atol, rtol = self.tolerances.get(tap_name, (self.default_atol, self.default_rtol))
abs_d, rel_d = array_diff(ref_val, cur_val)
if not (abs_d <= atol or rel_d <= rtol): # diverged: outside BOTH tolerances
if not (abs_d <= atol or rel_d <= rtol): # diverged: outside BOTH tolerances
return Divergence(f"{tap_name}@{step}", level, abs_d, rel_d,
f"abs={abs_d:.3e} rel={rel_d:.3e} > (atol={atol},rtol={rtol})")
return None
v2/parity/interleave_gate.py
+21 -24
View File
@@ -1,25 +1,17 @@
"""The batch-of-N interleave parity gate (design_v3 §9.3) — non-negotiable.
"""The batch-of-N interleave parity gate — required, not optional.
> Loop inversion's real hazard is cross-request state smearing under interleaving — and a
> batch-of-1 parity gate is structurally blind to it. So v3 makes a batch-of-N interleave
> parity test a *required* gate: two (or more) concurrent requests, interleaved at step
> granularity, must be bit-identical to the same requests run serially.
Loop inversion's real hazard is cross-request state smearing under interleaving, which a
batch-of-1 parity gate cannot detect. So this gate runs two or more concurrent requests
interleaved at step granularity and requires bit-identical output versus running them serially.
This is the test the whole loop-inversion bet lives or dies on. The gate is pure and
duck-typed: it takes any engine exposing ``run_serial`` and ``run_interleaved``.
The gate is pure and duck-typed: it takes any engine exposing ``run_serial`` and ``run_interleaved``.
"""
from __future__ import annotations
from typing import Any
from v2.request.artifacts import (
AudioArtifact,
LatentArtifact,
Output,
TensorArtifact,
TextArtifact,
VideoArtifact,
)
Output, )
from v2.parity.ladder import ConsistencyLevel, Divergence, bit_identical
@@ -34,27 +26,32 @@ def compare_outputs(serial: dict[str, Output], interleaved: dict[str, Output]) -
"""Bitwise-compare per-request outputs from serial vs interleaved runs."""
divs: list[Divergence] = []
if set(serial) != set(interleaved):
divs.append(Divergence("request_set", ConsistencyLevel.C1, float("inf"), float("inf"),
f"serial reqs {set(serial)} != interleaved {set(interleaved)}"))
divs.append(
Divergence("request_set", ConsistencyLevel.C1, float("inf"), float("inf"),
f"serial reqs {set(serial)} != interleaved {set(interleaved)}"))
return divs
for rid in sorted(serial):
so, io = serial[rid], interleaved[rid]
if set(so.artifacts) != set(io.artifacts):
divs.append(Divergence(f"{rid}:artifacts", ConsistencyLevel.C1, float("inf"), float("inf"),
f"artifact names differ: {set(so.artifacts)} vs {set(io.artifacts)}"))
divs.append(
Divergence(f"{rid}:artifacts", ConsistencyLevel.C1, float("inf"), float("inf"),
f"artifact names differ: {set(so.artifacts)} vs {set(io.artifacts)}"))
continue
for name in sorted(so.artifacts):
a, b = _artifact_payload(so.artifacts[name]), _artifact_payload(io.artifacts[name])
if a is None and b is None:
# both empty is NOT parity — a deferred/aborted request must not pass the gate vacuously
divs.append(Divergence(f"{rid}:{name}", ConsistencyLevel.C1, float("inf"), float("inf"),
"both serial and interleaved produced EMPTY output for a declared "
"artifact (deferred/aborted?) — suspicious, not parity"))
divs.append(
Divergence(
f"{rid}:{name}", ConsistencyLevel.C1, float("inf"), float("inf"),
"both serial and interleaved produced EMPTY output for a declared "
"artifact (deferred/aborted?) — suspicious, not parity"))
continue
if not bit_identical(a, b):
divs.append(Divergence(f"{rid}:{name}", ConsistencyLevel.C1, float("nan"), float("nan"),
"serial vs interleaved output not bit-identical "
"(cross-request state smearing!)"))
divs.append(
Divergence(f"{rid}:{name}", ConsistencyLevel.C1, float("nan"), float("nan"),
"serial vs interleaved output not bit-identical "
"(cross-request state smearing!)"))
return divs
v2/parity/ladder.py
+2 -2
View File
@@ -1,4 +1,4 @@
"""The consistency ladder (design_v3 §9.2) + numeric comparison helpers.
"""The consistency ladder + numeric comparison helpers.
C0 component · C1 loop · C2 behavioral · C3 distribution · C4 artifact.
@@ -18,7 +18,7 @@ from v2._enums import ConsistencyLevel # re-exported for the package
@dataclass
class Divergence:
where: str # tap name / artifact path
where: str # tap name / artifact path
level: ConsistencyLevel
max_abs_diff: float
max_rel_diff: float
v2/platform/__init__.py
+6 -6
View File
@@ -1,12 +1,12 @@
"""The multi-backend dispatch substrate (design_v3 §17).
"""The multi-backend dispatch substrate.
Two tuple-keyed registries (``COMPONENTS``, ``KERNELS``) + a ``Platform`` that detects ``(device,
arch)`` and resolves both — with the numpy reference as the terminal fallback rung and parity
oracle. This package is the membrane that lets CPU, GPU, and other backends coexist: a model's
loops, policies, scheduler, caches, and training code never name a device; they call through the
component/kernel seams and the resolved ``Platform`` decides the implementation.
arch)`` and resolves both, with the numpy reference as the terminal fallback rung and parity oracle.
Lets CPU, GPU, and other backends coexist: a model's loops, policies, scheduler, caches, and
training code never name a device; they call through the component/kernel seams and the resolved
``Platform`` decides the implementation.
Importing this package is cheap and cycle-free (registry + platform classes only). The concrete
Importing this package is cheap and cycle-free (registry + platform classes only); the concrete
backend registrations in ``backends/`` are imported lazily on first ``Platform`` use.
"""
from __future__ import annotations
v2/platform/backends/accel.py
+21 -16
View File
@@ -1,8 +1,8 @@
"""``accel`` — a pure-python stand-in accelerator backend (design_v3 §17 honesty note).
"""``accel`` — a pure-python stand-in accelerator backend.
There is no GPU in this environment, so to *prove* the dispatch substrate is genuinely
device-generic — not a single-backend abstraction with a CPU special case — this backend registers
a second device, ``accel``, entirely in CPU-resident Python. It exercises everything a real GPU
There is no GPU in this environment, so to prove the dispatch substrate is genuinely
device-generic (not a single-backend abstraction with a CPU special case) this backend registers a
second device, ``accel``, entirely in CPU-resident Python. It exercises everything a real GPU
backend would:
* a **component** override — ``AccelDiT`` is built via ``COMPONENTS[("dit", "accel", …)]`` instead
@@ -11,12 +11,12 @@ backend would:
* **kernels at different arch rungs** — ``flow_sde_step`` is registered at ``sm90`` directly, while
``flow_match_step`` is registered only at ``generic``, so resolving it on an ``sm90`` accel
platform must walk the arch fallback ``sm90 → sm80 → generic`` to find it;
* the **parity oracle** — the ODE solver op (``_accel_flow_match_step``) is an *independently
written* kernel that nonetheless matches the numpy reference bit-for-bit, so a denoise run on the
accel platform equals the cpu run. That independent-impl-vs-reference equality is the oracle a
real backend must pass; the stand-in shows the check has teeth. (The SDE op below deliberately
*reuses* the reference impl — a backend may legitimately not specialize every op — so the SDE
test exercises the dispatch path, not an independent implementation.)
* the **parity oracle** — the ODE solver op (``_accel_flow_match_step``) is an independently
written kernel that matches the numpy reference bit-for-bit, so a denoise run on the accel
platform equals the cpu run. That independent-impl-vs-reference equality is the oracle a real
backend must pass. (The SDE op below deliberately reuses the reference impl — a backend may
legitimately not specialize every op — so the SDE test exercises the dispatch path, not an
independent implementation.)
On a real box this file is where a torch/cuda backend would register its adapters and fused kernels
(unified primitives + model-co-located fusion variants) through the same two functions.
@@ -79,8 +79,7 @@ def _build_accel_component(spec: Any, instance: Any, platform: Any) -> AccelComp
# Override more than one component kind on accel (proving the seam is kind-generic, not dit-special).
# text_encoder is deliberately left UNregistered so the device→cpu fallback path stays demonstrated.
for _kind in ("vae", "audio_vae"):
register_component(_kind, _build_accel_component, device="accel",
source="accel(stand-in):AccelComponent")
register_component(_kind, _build_accel_component, device="accel", source="accel(stand-in):AccelComponent")
# --------------------------------------------------------------------------- #
@@ -93,8 +92,11 @@ def _accel_flow_match_step(x_t, velocity, sigma_t: float, sigma_next: float):
# flow_match_step lives only at the `generic` arch rung → resolving on sm90 must fall back through
# the arch chain (sm90 → sm80 → generic). This proves the arch-fallback walk.
register_kernel(FLOW_MATCH_STEP, _accel_flow_match_step,
device="accel", arch="generic", source="accel(stand-in):fused_flow_match")
register_kernel(FLOW_MATCH_STEP,
_accel_flow_match_step,
device="accel",
arch="generic",
source="accel(stand-in):fused_flow_match")
# flow_sde_step is registered at sm90 directly. This backend does NOT specialize the SDE op — it
# reuses the numpy reference, so the accel SDE path is the same function as cpu (the SDE parity test
@@ -102,5 +104,8 @@ register_kernel(FLOW_MATCH_STEP, _accel_flow_match_step,
# independent oracle). A device that DID specialize SDE would be checked the same way the ODE op is.
from v2.loop.sampler import flow_sde_step_with_logprob # noqa: E402
register_kernel(FLOW_SDE_STEP, flow_sde_step_with_logprob,
device="accel", arch="sm90", source="accel(stand-in):sde_logprob")
register_kernel(FLOW_SDE_STEP,
flow_sde_step_with_logprob,
device="accel",
arch="sm90",
source="accel(stand-in):sde_logprob")
v2/platform/backends/cpu.py
+11 -5
View File
@@ -1,4 +1,4 @@
"""CPU / numpy backend — the terminal rung and parity oracle (design_v3 §17).
"""CPU / numpy backend — the terminal rung and parity oracle.
This backend registers the numpy reference kernels. It is the bottom of every device fallback
chain: any op a richer backend hasn't implemented resolves here, and every other backend's output
@@ -16,7 +16,13 @@ from v2.platform.registry import FLOW_MATCH_STEP, FLOW_SDE_STEP, register_kernel
# The numpy solver primitives, registered as the terminal (cpu, numpy) kernels. These are the exact
# same functions the loops used to call directly — so dispatching through the registry on a CPU box
# is bit-for-bit identical to the pre-registry behavior.
register_kernel(FLOW_MATCH_STEP, flow_match_euler_step,
device="cpu", arch="numpy", source="loop.sampler:flow_match_euler_step")
register_kernel(FLOW_SDE_STEP, flow_sde_step_with_logprob,
device="cpu", arch="numpy", source="loop.sampler:flow_sde_step_with_logprob")
register_kernel(FLOW_MATCH_STEP,
flow_match_euler_step,
device="cpu",
arch="numpy",
source="loop.sampler:flow_match_euler_step")
register_kernel(FLOW_SDE_STEP,
flow_sde_step_with_logprob,
device="cpu",
arch="numpy",
source="loop.sampler:flow_sde_step_with_logprob")
v2/platform/backends/torch_backend.py
+11 -12
View File
@@ -1,17 +1,16 @@
"""Real torch component adapters for the GPU backend (design_v3 §17; README "Running the real models").
"""Real torch component adapters for the GPU backend. See v2/README.md ("Running the real models").
ONE cohesive module (was the torch_adapters.py + torch_ltx2.py split): a ``TorchComponent`` base that
centralizes the shared mechanics — ``.to(device,dtype).eval()``, the numpy<->torch marshalling at the
loop boundary, the ``set_forward_context`` wrap every fastvideo forward needs, and the weight surface
(copy_from / blend_from / clone) — plus thin per-model subclasses that carry ONLY the forward semantics
(Wan ``sigma*1000`` -> velocity; LTX-2 per-token sigma, x0 -> ``(x_t-x0)/sigma``, joint A/V; VAE
normalization; T5 padding; Gemma dual-projection; upsampler; audio decode->vocoder). A single
``build_component(spec, instance, platform)`` dispatches by ``spec.kind`` through ``_MAKERS`` (replacing
the six ``build_torch_*`` builders + six lazy trampolines).
A ``TorchComponent`` base centralizes the shared mechanics — ``.to(device,dtype).eval()``, the
numpy<->torch marshalling at the loop boundary, the ``set_forward_context`` wrap every fastvideo
forward needs, and the weight surface (copy_from / blend_from / clone). Thin per-model subclasses
carry only the forward semantics (Wan ``sigma*1000`` -> velocity; LTX-2 per-token sigma,
x0 -> ``(x_t-x0)/sigma``, joint A/V; VAE normalization; T5 padding; Gemma dual-projection; upsampler;
audio decode->vocoder). A single ``build_component(spec, instance, platform)`` dispatches by
``spec.kind`` through ``_MAKERS``.
The loops / CFG / scheduler math stay numpy fp32 (lowest-risk bring-up); the boundary marshals
numpy<->torch in ONE place (``TorchComponent._t``/``_n``) so a torch-native path is a later swap.
Component construction goes through ``load_component`` — the v2-owned loader seam (currently delegates to
Component construction goes through ``load_component``, the v2-owned loader seam (currently delegates to
fastvideo's component_loader; a per-module vendored cutover replaces it later, no caller changes).
Imported ONLY lazily — inside ``torch_cuda.py``'s registered builder — so importing the platform package
@@ -87,7 +86,7 @@ def _fastvideo_args(spec: Any) -> Any:
def load_component(loader_attr: str, path: str, args):
"""v2-owned loader seam (design: ``v2/loader``). Constructs a real component from a checkpoint via the
"""v2-owned loader seam (``v2/loader``). Constructs a real component from a checkpoint via the
component loaders. ``WanTransformer3DModel`` / ``AutoencoderKLWan`` have NO ``from_pretrained``; the
loader reads the checkpoint config, resolves the class (UMT5 vs T5 from config), and loads weights.
Currently delegates to fastvideo's loader; a vendored cutover swaps the body, not the callers."""
@@ -145,7 +144,7 @@ class TorchComponent:
from v2.forward_context import set_forward_context
return set_forward_context(current_timestep=current_timestep, attn_metadata=None)
# weight surface used by serving weight-sync + training (design_v3 §10) ------------------------- #
# weight surface used by serving weight-sync + training ----------------------------------------- #
def copy_from(self, other) -> None:
self.module.load_state_dict(other.module.state_dict())
v2/platform/backends/torch_cosmos25.py
+3 -4
View File
@@ -8,14 +8,13 @@ The three Cosmos2.5 deltas vs Wan (all GPU-path, all carried HERE so the loop st
* ``Cosmos25DiT`` — a flow-match velocity DiT, but the per-step call differs from ``WanDiT`` in three
load-bearing ways (see ``fastvideo/pipelines/stages/denoising.py:Cosmos25DenoisingStage``):
1. timestep is a PER-FRAME 2D tensor ``[B, T]`` filled with the PLAIN sigma (in [0,1]) — NOT Wan's
scalar ``sigma*1000``. (The fastvideo stage feeds ``t*0.001`` of the 0..1000 schedule, i.e. the
shifted sigma directly; ``Cosmos25DenoiseLoop`` already hands us that plain sigma.)
scalar ``sigma*1000``. ``Cosmos25DenoiseLoop`` already hands us that plain sigma.
2. the forward REQUIRES a ``condition_mask[B,1,T,H,W]`` (zeros for t2v) + a ``padding_mask[B,1,H,W]``
(ones for t2v) every step — the model concats them internally (16 -> 17 -> 18ch patch_embed),
so we feed the RAW 16-channel latent (the OPPOSITE of Wan i2v, where the adapter pre-concats).
3. an ``fps`` scalar is required for RoPE temporal scaling (NTK fps-modulation).
The model output is the rectified-flow velocity (a bare tensor); the loop integrates it with the
shared ``FLOW_MATCH_STEP`` Euler solver and combines branches with the standard ``ClassicCFG``.
The output is the rectified-flow velocity (a bare tensor); the loop integrates it with the shared
``FLOW_MATCH_STEP`` Euler solver and combines branches with ``ClassicCFG``.
* ``Cosmos25WanVAE`` — a Wan-style causal 3D VAE (z=16, 8x/4x), BUT unlike ``AutoencoderKLWan`` its
``encode``/``decode`` already apply the (z-mean)/std normalization INTERNALLY (the Cosmos latent
contract: encode -> normalized, decode <- normalized). So this is a thin marshalling wrapper (like
v2/platform/backends/torch_cuda.py
+21 -10
View File
@@ -1,4 +1,4 @@
"""``cuda`` registration — the real torch/GPU backend (design_v3 §17; README "Running the real models").
"""``cuda`` registration — the real torch/GPU backend (see v2/README.md "Running the real models").
Torch-free at import (stdlib only): gates every cell on ``available=_cuda_available`` (a ``find_spec``
check that never imports torch) and registers ONE lazy component builder for all cuda component kinds +
@@ -50,17 +50,28 @@ def _flow_sde_cuda(*args, **kwargs):
# --- components: one generic builder; torch_backend.build_component dispatches by spec.kind ---------- #
_B = "v2.platform.backends.torch_backend"
_KIND_SOURCE = {
"dit": "WanDiT/LTX2DiT", "vae": "WanVAE/LTX2VAE", "text_encoder": "T5Encoder/Gemma",
"image_encoder": "CLIPImageEncoder", "upsampler": "LTX2Upsampler",
"audio_vae": "LTX2AudioVAE", "vocoder": "LTX2Vocoder",
"dit": "WanDiT/LTX2DiT",
"vae": "WanVAE/LTX2VAE",
"text_encoder": "T5Encoder/Gemma",
"image_encoder": "CLIPImageEncoder",
"upsampler": "LTX2Upsampler",
"audio_vae": "LTX2AudioVAE",
"vocoder": "LTX2Vocoder",
}
for _kind, _cls in _KIND_SOURCE.items():
register_component(_kind, _build_component, device="cuda", available=_cuda_available,
source=f"{_B}:{_cls}")
register_component(_kind, _build_component, device="cuda", available=_cuda_available, source=f"{_B}:{_cls}")
del _kind, _cls
# --- solver ops: plain torch elementwise (no fused solver kernel in fastvideo-kernel) ---------------- #
register_kernel(FLOW_MATCH_STEP, _flow_match_cuda, device="cuda", arch="generic",
available=_cuda_available, source="torch elementwise (no fused solver kernel)")
register_kernel(FLOW_SDE_STEP, _flow_sde_cuda, device="cuda", arch="generic",
available=_cuda_available, source="torch elementwise (no fused solver kernel)")
register_kernel(FLOW_MATCH_STEP,
_flow_match_cuda,
device="cuda",
arch="generic",
available=_cuda_available,
source="torch elementwise (no fused solver kernel)")
register_kernel(FLOW_SDE_STEP,
_flow_sde_cuda,
device="cuda",
arch="generic",
available=_cuda_available,
source="torch elementwise (no fused solver kernel)")
v2/platform/backends/torch_gen3c.py
+3 -4
View File
@@ -128,10 +128,9 @@ class Gen3CVAE(TorchComponent):
class Gen3CT5Encoder(CosmosT5Encoder):
"""GEN3C T5-Large (1024-dim): raw last_hidden_state (NaN→0), no Wan zero-pad — identical to Cosmos.
Reuses ``CosmosT5Encoder.encode`` verbatim. (BRINGUP: GEN3C pads with ``max_length`` and zeros rows
past the attention-mask length, which ``CosmosT5Encoder`` already does — verify the max_length /
mask-zeroing matches the converted T5-Large tokenizer on the GPU box.)"""
"""GEN3C T5-Large (1024-dim): raw last_hidden_state (NaN→0), no Wan zero-pad — identical to Cosmos,
so it reuses ``CosmosT5Encoder.encode`` verbatim. (BRINGUP: verify the max_length / mask-zeroing
matches the converted T5-Large tokenizer on the GPU box.)"""
# Re-export so a reader can ``from v2.platform.backends.torch_gen3c import _to_numpy`` if needed.
v2/platform/backends/torch_hunyuan_video15.py
+23 -27
View File
@@ -2,37 +2,33 @@
so the recipe is self-contained (no edit to the shared ``_make_dit``/``_make_vae``/``_make_text_encoder``
dispatch in ``torch_backend.py``). Imported lazily by ``_explicit_adapter`` only on a GPU box.
Architecture deltas vs Wan that force NEW adapters (all faithful to
Architecture deltas vs Wan that force NEW adapters (faithful to
``fastvideo/models/dits/hunyuanvideo15.py`` + ``fastvideo/pipelines/stages/denoising.py``):
* ``HunyuanVideo15DiT`` — the rectified-flow velocity network ``F_θ``. Unlike ``WanDiT`` (single text
embed, single image embed), the real ``forward`` takes ``encoder_hidden_states`` as a **LIST of two**
tensors ``[qwen_hidden(3584-d), byt5_hidden(1472-d)]`` and ``encoder_hidden_states_image`` as a
**LIST of one** image-embed tensor (all-zeros ``[1,729,1152]`` for t2v → the DiT's
``torch.all(image_embeds==0)`` auto-selects the t2v token-ordering branch). The loop hands the two
text embeds as a tuple in ``text_embed`` and the image embeds in ``context``; this adapter marshals
them into the two lists internally and builds the zero image-embed default, so the loop's dit-call
signature stays ``dit(latent, text_embed, sigma, context=, cond=)`` (toy-compatible). For i2v the
loop passes the 33-channel first-frame conditioning latent (32 VAE + 1 mask) in ``cond``; the adapter
concats it onto the noisy latent on the channel dim to reach ``in_channels=65`` BEFORE the forward
(mirrors ``denoising.py`` line 393's ``torch.cat([latent, image_latent], dim=1)``). ``timestep`` is
``sigma*1000`` (the scheduler's discrete timestep; ``scale_model_input`` is identity for flow-match).
The network predicts velocity directly → returned via ``self._n()`` with NO x0→velocity conversion
(unlike ``LTX2DiT``). ``use_meanflow`` SR variants pass ``timestep_r`` — out of scope here (t2v/i2v
ship with ``use_meanflow=False``), so ``timestep_r=None``.
* ``HunyuanVideo15VAE`` — the HunyuanVideo 1.5 causal-3D VAE (z=32, 16× spatial / 4× temporal). Its DiT
latent space is normalized by a single **scalar** ``scaling_factor=1.03682`` (NOT Wan's per-channel
* ``HunyuanVideo15DiT`` — rectified-flow velocity network. Unlike ``WanDiT`` (single text/image embed),
the real ``forward`` takes ``encoder_hidden_states`` as a LIST of two ``[qwen_hidden(3584-d),
byt5_hidden(1472-d)]`` and ``encoder_hidden_states_image`` as a LIST of one image-embed (all-zeros
``[1,729,1152]`` for t2v -> the DiT's ``torch.all(image_embeds==0)`` selects the t2v token-ordering
branch). The loop hands the two text embeds as a tuple in ``text_embed`` and the image embeds in
``context``; this adapter marshals them into the two lists and builds the zero image-embed default, so
the loop's dit-call signature stays ``dit(latent, text_embed, sigma, context=, cond=)``. For i2v the
loop passes the 33-channel first-frame cond latent (32 VAE + 1 mask) in ``cond``; the adapter concats
it onto the noisy latent to reach ``in_channels=65`` before the forward. ``timestep`` is ``sigma*1000``
(``scale_model_input`` is identity for flow-match). Predicts velocity directly, returned via
``self._n()`` with NO x0->velocity conversion. ``timestep_r`` (``use_meanflow`` SR variants) is out of
scope here (t2v/i2v ship with ``use_meanflow=False``), so ``timestep_r=None``.
* ``HunyuanVideo15VAE`` — HunyuanVideo 1.5 causal-3D VAE (z=32, 16x spatial / 4x temporal). DiT latent
space normalized by a single scalar ``scaling_factor=1.03682`` (NOT Wan's per-channel
``latents_mean``/``latents_std`` — the ``WanVAE`` adapter would corrupt these latents, BRINGUP risk D).
Decode: ``module.decode(latent / scaling_factor)``; encode (i2v first-frame cond): ``z * scaling_factor``.
* ``HunyuanVideo15QwenEncoder`` — Qwen2.5-VL primary text encoder. Faithful to
``configs/pipelines/hunyuan15.py:qwen_preprocess_text/qwen_postprocess_text``: wraps the prompt in the
system+chat template, requests ``output_hidden_states=True``, takes ``hidden_states[-3]`` (3rd-from-last
layer, NOT ``last_hidden_state``), and crops the first 108 prompt-template tokens. The generic
``T5Encoder.encode`` (last_hidden_state, pad-to-max) is WRONG for Qwen (BRINGUP risk B).
* ``HunyuanVideo15ByT5Encoder`` — ByT5/Glyph secondary encoder. ``byt5_preprocess_text`` extracts quoted
glyph texts via regex; when the prompt has no quoted text it is ``None`` and the real pipeline emits a
zero-length ``[1, 0, 1472]`` embedding (the DiT always concatenates the txt_in_2 output) — this adapter
reproduces that (BRINGUP risk C). Output is the raw ``last_hidden_state``.
* ``HunyuanVideo15QwenEncoder`` — Qwen2.5-VL primary text encoder. Wraps the prompt in the system+chat
template, requests ``output_hidden_states=True``, takes ``hidden_states[-3]`` (3rd-from-last layer, NOT
``last_hidden_state``), and crops the first 108 template tokens. The generic ``T5Encoder.encode`` is
WRONG for Qwen (BRINGUP risk B).
* ``HunyuanVideo15ByT5Encoder`` — ByT5/Glyph secondary encoder. Extracts quoted glyph texts via regex;
when the prompt has no quoted text the real pipeline emits a zero-length ``[1, 0, 1472]`` embedding (the
DiT always concatenates the txt_in_2 output) — this adapter reproduces that (BRINGUP risk C). Output is
the raw ``last_hidden_state``.
All forwards run inside ``self._ctx()`` (FastVideo attention reads attn_metadata via the forward
context). BRINGUP: the DiT's ``sequence_model_parallel_shard``/all-gather is a no-op at sp=1 (the card
v2/platform/backends/torch_kernels.py
+13 -14
View File
@@ -1,13 +1,12 @@
"""Torch solver-step kernels for the GPU backend (design_v3 §17).
"""Torch solver-step kernels for the GPU backend.
GROUND TRUTH (from the fastvideo-kernel audit): there is NO fused flow-match / SDE *solver* kernel.
fastvideo-kernel ships only primitives — sparse/tiled attention, RMS/LayerNorm, int8 GEMM/quant.
So the GPU solver step is a plain torch elementwise op with the SAME math as ``loop/sampler.py``'s
numpy reference (the consistency oracle that the accel stand-in already proves matches bit-for-bit).
There is NO fused flow-match / SDE *solver* kernel in fastvideo-kernel — it ships only primitives
(sparse/tiled attention, RMS/LayerNorm, int8 GEMM/quant). So the GPU solver step is a plain torch
elementwise op with the SAME math as ``loop/sampler.py``'s numpy reference (which it matches bit-for-bit).
WRITTEN-NOT-RUN: imports torch; loaded only lazily from ``torch_cuda.py``. Under the current numpy
loop surface these marshal numpy↔torch at the boundary (a torch-native surface would keep the latent
on-device through forward→combine→solver — the perf follow-up, Risk G in GPU_BRINGUP.md).
Imports torch; loaded only lazily from ``torch_cuda.py``. Under the current numpy loop surface these
marshal numpy<->torch at the boundary (a torch-native surface would keep the latent on-device through
forward->combine->solver — a perf follow-up).
"""
from __future__ import annotations
@@ -42,16 +41,16 @@ def flow_sde_step(x_t, velocity, sigma_t, sigma_next, *, noise=None, prev_sample
v = v.double()
s = min(float(sigma_t), 0.9999)
dt = float(sigma_next) - float(sigma_t)
std = float((s / (1.0 - s)) ** 0.5 * noise_scale)
std = float((s / (1.0 - s))**0.5 * noise_scale)
denom = 2.0 * max(s, 1e-6)
mean = xt * (1.0 + std ** 2 / denom * dt) + v * (1.0 + std ** 2 * (1.0 - s) / denom) * dt
eff_std = max(std * max(-dt, 1e-12) ** 0.5, 1e-6)
mean = xt * (1.0 + std**2 / denom * dt) + v * (1.0 + std**2 * (1.0 - s) / denom) * dt
eff_std = max(std * max(-dt, 1e-12)**0.5, 1e-6)
if prev_sample is None:
n = _as_torch(noise)[0].double() if noise is not None else torch.zeros_like(xt)
prev = mean + eff_std * n
else:
prev = _as_torch(prev_sample)[0].double()
var = eff_std ** 2
log_prob = float((-((prev - mean) ** 2) / (2.0 * var)
- float(np.log(eff_std)) - 0.5 * float(np.log(2.0 * np.pi))).mean().item())
var = eff_std**2
log_prob = float(
(-((prev - mean)**2) / (2.0 * var) - float(np.log(eff_std)) - 0.5 * float(np.log(2.0 * np.pi))).mean().item())
return (_back(prev.float(), was_np), log_prob, _back(mean.float(), was_np), float(eff_std))
v2/platform/backends/torch_longcat.py
+12 -14
View File
@@ -4,20 +4,18 @@ Imported lazily by ``_explicit_adapter`` only on a GPU box.
* ``LongCatDiT`` — the flow-match velocity predictor. Mirrors ``WanDiT`` (``timestep = sigma*1000``,
scalar ``[B]`` timestep that the DiT internally expands ``[B] -> [B, T_latent]`` for the per-latent-frame
AdaLN time embedding) but with the **critical LongCat sign convention**: the fastvideo
``LongCatDenoisingStage`` does ``noise_pred = -noise_pred`` *before* the flow-match scheduler step
(longcat_denoising.py line 152). The v2 loop integrates ``x + (sigma_next-sigma)*velocity`` with the
velocity this adapter returns, so to reproduce that step EXACTLY this adapter returns ``-velocity``.
Forgetting the negation silently produces diverging samples (spec blocker #1).
AdaLN time embedding) but with the critical LongCat sign convention: the fastvideo
``LongCatDenoisingStage`` does ``noise_pred = -noise_pred`` *before* the flow-match scheduler step. The v2
loop integrates ``x + (sigma_next-sigma)*velocity`` with the velocity this adapter returns, so to reproduce
that step exactly this adapter returns ``-velocity``. Forgetting the negation silently diverges.
LongCat reuses the existing ``WanVAE`` torch adapter unchanged (same ``AutoencoderKLWan`` mean/std
normalization) and the existing ``T5Encoder`` (UMT5, zero-padded to max_length=512 for the CFG concat
uniform-seq contract). Both are wired on the card by ``load_id`` alone — no adapter override needed.
The DiT runs its AdaLN modulation / residual gating + final projection under fp32 and casts the output to
float32 (longcat.py line 1121); we keep the module at its native bf16 dtype (``_native_dtype`` in
``_make_dit``, matching the fastvideo stage's hardcoded bf16 target) and let ``_n`` marshal the fp32 output
back to numpy (spec blocker #6).
float32; we keep the module at its native bf16 dtype (``_native_dtype`` in ``_make_dit``, matching the
fastvideo stage's hardcoded bf16 target) and let ``_n`` marshal the fp32 output back to numpy.
BRINGUP: written-not-run on CPU (no torch/weights here). GPU-verify against a real
``FastVideo/LongCat-Video-T2V-Diffusers`` checkpoint — confirm the bf16 dtype, the scalar-timestep
@@ -40,17 +38,17 @@ class LongCatDiT(TorchComponent):
The loop hands the raw ``sigma`` (1->0); this adapter forms ``timestep = sigma*1000`` (the Wan/diffusers
convention LongCat shares) as a scalar ``[B]`` tensor — the DiT's ``forward`` expands ``[B] -> [B, T_latent]``
internally for the AdaLN time embedding (longcat.py line 1075). It returns the **negated** velocity to
fold in the fastvideo stage's ``noise_pred = -noise_pred`` (spec blocker #1)."""
internally for the AdaLN time embedding. It returns the negated velocity to fold in the fastvideo stage's
``noise_pred = -noise_pred``."""
@torch.no_grad()
def __call__(self, latent, text_embed, sigma, context=None, *, cond=None):
hs = self._t(latent) # [1, C=16, T, H, W]
ehs = self._t(text_embed) # [1, N_text, 4096]
b = hs.shape[0]
# timestep = sigma*1000 (BRINGUP risk B / spec timestep_convention). Scalar [B]; the DiT expands
# [B] -> [B, T_latent] internally for the per-latent-frame AdaLN time embedding. (I2V/VC would build
# an explicit [B, T_latent] with cond frames zeroed — deferred, T2V uses the scalar path.)
# timestep = sigma*1000 (BRINGUP risk B). Scalar [B]; the DiT expands [B] -> [B, T_latent] internally
# for the per-latent-frame AdaLN time embedding. (I2V/VC would build an explicit [B, T_latent] with
# cond frames zeroed — deferred, T2V uses the scalar path.)
ts = float(sigma) * NUM_TRAIN_TIMESTEPS
timestep = torch.full((b, ), ts, device=self.device, dtype=self.dtype)
with self._ctx(current_timestep=ts):
@@ -58,6 +56,6 @@ class LongCatDiT(TorchComponent):
encoder_hidden_states=ehs,
timestep=timestep,
encoder_attention_mask=None)
# CRITICAL (spec blocker #1): negate so the loop's FLOW_MATCH_STEP reproduces the fastvideo stage's
# CRITICAL: negate so the loop's FLOW_MATCH_STEP reproduces the fastvideo stage's
# negate-then-scheduler-step. The DiT already casts to float32; _n marshals fp32 -> numpy.
return self._n(-velocity)
v2/platform/backends/torch_matrixgame2.py
+2 -2
View File
@@ -64,8 +64,8 @@ class MatrixGame2CausalDiT(TorchComponent):
The loop drives the causal block-autoregressive schedule and the DMD epsilon->x0 conversion; this
adapter is a single ``_forward_inference`` call with the cache/i2v/action plumbing. It is NOT the toy
``dit(latent, text_embed, sigma)`` shape — the loop calls ``MatrixGame2CausalDiT.call`` explicitly with
the causal kwargs (the toy backend exercises the SAME loop via its degenerate path; see the recipe
loop). All arch-specifics are built INTERNALLY here so the loop stays backend-agnostic.
the causal kwargs (the toy backend exercises the same loop via its degenerate path). All arch-specifics
are built internally here so the loop stays backend-agnostic.
KV-CACHE OWNERSHIP: ``forward`` dispatches to ``_forward_inference`` ONLY when ``kv_cache is not None``;
otherwise it falls into the train forward (which expects a non-empty text-embed list and uses the
v2/platform/backends/torch_matrixgame3.py
+4 -4
View File
@@ -4,7 +4,7 @@ so the MG3 recipe is self-contained (no edit to the shared ``_make_dit`` dispatc
``MatrixGame3DiT`` is the autoregressive world-model DiT (``MatrixGame3WanModel``, 30 layers, dim 3072,
in/out_channels=48, patch (1,2,2), ``use_memory=True``, ``sigma_theta=0.8``). It departs from ``WanDiT``
in three load-bearing ways (BRINGUP risks 1/2/7 in the spec):
in three load-bearing ways:
1. PER-TOKEN timestep, NOT ``sigma*1000`` scalar. The model's ``timestep`` arg is a flattened
``[1, T*(H/2)*(W/2)]`` LongTensor whose value is the raw FlowUniPC scheduler timestep, with the
@@ -59,9 +59,9 @@ class MatrixGame3DiT(TorchComponent):
def _per_token_timestep(self, hs: torch.Tensor, sigma_value: float, cond_frames: int,
x_memory: torch.Tensor | None) -> torch.Tensor:
"""Build the flattened per-token timestep [1, T*(H/p)*(W/p)] (BRINGUP risk 1). The raw FlowUniPC
timestep value fills every token; ``cond_frames`` latent rows are zeroed (pinned, not denoised);
an all-zeros memory block is prepended when ``x_memory`` is present."""
"""Build the flattened per-token timestep [1, T*(H/p)*(W/p)]. The raw FlowUniPC timestep value fills
every token; ``cond_frames`` latent rows are zeroed (pinned, not denoised); an all-zeros memory block
is prepended when ``x_memory`` is present."""
_b, _c, t, h, w = hs.shape
tokens_per_frame = (h // self.patch_h) * (w // self.patch_w)
tt = hs.new_full((t, tokens_per_frame), float(sigma_value))
v2/platform/backends/toy.py
+60 -45
View File
@@ -1,4 +1,4 @@
"""Toy numpy components — the CPU-testable backend (design_v3 §17 honesty note).
"""Toy numpy components — the CPU-testable backend.
There is no GPU/torch/weights in this environment, so the heavy Wan/LTX neural forwards are
represented by small, deterministic numpy components. They exercise the *real* loop control
@@ -75,17 +75,22 @@ class ToyDiT:
self.s_ctx = 0.3
def _pre_tanh(self, latent: np.ndarray, text_embed, sigma: float, context) -> np.ndarray:
mixed = np.tensordot(self.w_x, latent, axes=([1], [0])) # [C,...] channel mix
mixed = mixed + self.w_t.reshape((self.C,) + (1,) * (latent.ndim - 1)) * float(sigma)
mixed = np.tensordot(self.w_x, latent, axes=([1], [0])) # [C,...] channel mix
mixed = mixed + self.w_t.reshape((self.C, ) + (1, ) * (latent.ndim - 1)) * float(sigma)
cond = float(np.mean(text_embed)) if text_embed is not None else 0.0
mixed = mixed + self.s_text * cond
if context is not None:
mixed = mixed + self.s_ctx * float(np.mean(context))
return mixed
def __call__(self, latent: np.ndarray, text_embed: np.ndarray | None, sigma: float,
context: np.ndarray | None = None, *,
audio_latent: np.ndarray | None = None, audio_text: np.ndarray | None = None,
def __call__(self,
latent: np.ndarray,
text_embed: np.ndarray | None,
sigma: float,
context: np.ndarray | None = None,
*,
audio_latent: np.ndarray | None = None,
audio_text: np.ndarray | None = None,
cond: np.ndarray | None = None):
# ``cond`` (i2v mask+cond latent) is the real WanDiT's 36ch concat; the toy denoises the noise
# channels only (image-conditioning is a GPU-path concern), so it's accepted and ignored here.
@@ -101,22 +106,28 @@ class ToyDiT:
return video, audio
# --- minimal trainable surface (so training methods do real optimizer steps) --------- #
def clone(self) -> "ToyDiT":
def clone(self) -> ToyDiT:
c = ToyDiT.__new__(ToyDiT)
c.C, c.s_text, c.s_ctx = self.C, self.s_text, self.s_ctx
c.w_x, c.w_t = self.w_x.copy(), self.w_t.copy()
return c
def blend_from(self, other: "ToyDiT", decay: float) -> None:
"""EMA / decay-blended-old-policy update: self ← decay·self + (1-decay)·other (design_v3 §10)."""
def blend_from(self, other: ToyDiT, decay: float) -> None:
"""EMA / decay-blended-old-policy update: self ← decay·self + (1-decay)·other."""
self.w_x = (decay * self.w_x + (1.0 - decay) * other.w_x).astype("float32")
self.w_t = (decay * self.w_t + (1.0 - decay) * other.w_t).astype("float32")
def copy_from(self, other: "ToyDiT") -> None:
def copy_from(self, other: ToyDiT) -> None:
self.w_x, self.w_t = other.w_x.copy(), other.w_t.copy()
def mse_grad_step(self, latent: np.ndarray, text_embed, sigma: float, target: np.ndarray,
lr: float, context=None, weight: float = 1.0) -> tuple[float, float]:
def mse_grad_step(self,
latent: np.ndarray,
text_embed,
sigma: float,
target: np.ndarray,
lr: float,
context=None,
weight: float = 1.0) -> tuple[float, float]:
"""One exact SGD step minimizing ``weight·MSE(forward(latent), target)`` w.r.t. ``w_x``.
Returns (loss, grad_norm). The grad_norm is the #1396-style per-method regression signal.
@@ -127,11 +138,11 @@ class ToyDiT:
z = self._pre_tanh(latent, text_embed, sigma, context)
pred = np.tanh(z)
err = pred - target
loss = float(weight * np.mean(err ** 2))
loss = float(weight * np.mean(err**2))
n = err.size
g = weight * (2.0 * err / n) * (1.0 - pred ** 2)
g = weight * (2.0 * err / n) * (1.0 - pred**2)
axes = (list(range(1, g.ndim)), list(range(1, latent.ndim)))
grad_wx = np.tensordot(g, latent, axes=axes) # [C_out, C_in]
grad_wx = np.tensordot(g, latent, axes=axes) # [C_out, C_in]
grad_norm = float(np.linalg.norm(grad_wx))
self.w_x = (self.w_x - lr * grad_wx).astype("float32")
return loss, grad_norm
@@ -152,7 +163,7 @@ class ToyTokenizer:
class ToyMoTDiT(ToyDiT):
"""Mixture-of-Transformers stand-in: ONE resident module that runs BOTH an AR (understanding)
pathway and a diffusion (generation) pathway on shared weights (design_v3 §4.2, §16).
pathway and a diffusion (generation) pathway on shared weights.
``ar_forward`` is the und pathway (next-token); ``__call__`` (inherited) is the gen pathway
(velocity). Binding both the ar_decode and diffusion_denoise loops to one instance of this
@@ -181,7 +192,7 @@ class ToyPromptRefiner:
refined prompt). Sampling an action picks an embedding offset that shifts the diffusion
conditioning; the realized reward then drives the policy by REINFORCE. This is the *second
trainable expert* in the unified-RL stress test: one reward → a token-policy-gradient update
here AND a FlowGRPO update on the DiT, under one advantage (design_v3 §10; PromptRL §6).
here AND a FlowGRPO update on the DiT, under one advantage (PromptRL §6).
Real gradient: loss = −A·logπ(a); ∇_logits = −A·(onehot(a) − softmax). So the toy LM actually
learns to prefer the reward-favored action — the LM-side analogue of ToyDiT.mse_grad_step.
@@ -190,7 +201,7 @@ class ToyPromptRefiner:
def __init__(self, n_actions: int = 8, dim: int = TEXT_DIM, seq: int = TEXT_SEQ, seed: int = 0):
self.n = int(n_actions)
self.dim, self.seq = int(dim), int(seq)
self.logits = np.zeros(self.n, dtype="float64") # the trainable policy params
self.logits = np.zeros(self.n, dtype="float64") # the trainable policy params
rng = np.random.default_rng(seed)
# each action ⇒ a fixed conditioning offset (the "content" of that refined prompt)
self._emb = (rng.standard_normal((self.n, self.seq, self.dim)) * 0.1).astype("float32")
@@ -228,25 +239,25 @@ class ToyPromptRefiner:
p = self._probs()
onehot = np.zeros(self.n, dtype="float64")
onehot[int(action)] = 1.0
grad = -float(advantage) * (onehot - p) # ∇_logits (−A·logπ(a))
self.logits = self.logits - float(lr) * grad # descent on −A·logπ ⇒ ascent on A·logπ
grad = -float(advantage) * (onehot - p) # ∇_logits (−A·logπ(a))
self.logits = self.logits - float(lr) * grad # descent on −A·logπ ⇒ ascent on A·logπ
return float(np.linalg.norm(grad))
def kl_to(self, other: "ToyPromptRefiner") -> float:
def kl_to(self, other: ToyPromptRefiner) -> float:
"""KL(π_self ‖ π_other) over the categorical — the LM-PG reference-KL penalty term."""
p, q = self._probs(), other._probs()
return float(np.sum(p * (np.log(p + 1e-12) - np.log(q + 1e-12))))
def clone(self) -> "ToyPromptRefiner":
def clone(self) -> ToyPromptRefiner:
c = ToyPromptRefiner.__new__(ToyPromptRefiner)
c.n, c.dim, c.seq = self.n, self.dim, self.seq
c.logits, c._emb = self.logits.copy(), self._emb.copy()
return c
def copy_from(self, other: "ToyPromptRefiner") -> None:
def copy_from(self, other: ToyPromptRefiner) -> None:
self.logits = other.logits.copy()
def blend_from(self, other: "ToyPromptRefiner", decay: float) -> None:
def blend_from(self, other: ToyPromptRefiner, decay: float) -> None:
self.logits = decay * self.logits + (1.0 - decay) * other.logits
@@ -265,7 +276,7 @@ class ToyTalker(ToyMoTDiT):
def ar_forward(self, tokens) -> int:
ctx = sum(int(t) for t in tokens)
return int((ctx * 7 + 13 + self.salt) % self.VOCAB) # weight-dependent (salt) + context
return int((ctx * 7 + 13 + self.salt) % self.VOCAB) # weight-dependent (salt) + context
class ToyVocoder:
@@ -279,7 +290,7 @@ class ToyVocoder:
def __init__(self, samples_per_token: int = 16, vocab: int = 256, seed: int = 7):
self.spt = int(samples_per_token)
rng = np.random.default_rng(seed)
self.bank = (rng.standard_normal(vocab) * 0.3).astype("float32") # per-token timbre offset
self.bank = (rng.standard_normal(vocab) * 0.3).astype("float32") # per-token timbre offset
def synthesize(self, token_chunk) -> np.ndarray:
"""One chunk of codec tokens → a waveform segment in [-1, 1], length spt·len(chunk)."""
@@ -296,11 +307,15 @@ class ToyLoRA:
Many of these are served over ONE resident base (the adapter is tiny: a rank-r channel delta);
a request picks which to apply. Swappable/versioned independently (the cache key's
``adapter_versions``). Stand-in for a real LoRA / DoRA / IP-Adapter (design_v3 §9.19)."""
``adapter_versions``). Stand-in for a real LoRA / DoRA / IP-Adapter."""
kind = "lora"
def __init__(self, adapter_id: str, channels: int = LATENT_CHANNELS, scale: float = 0.6,
rank: int = 2, seed: int = 0):
def __init__(self,
adapter_id: str,
channels: int = LATENT_CHANNELS,
scale: float = 0.6,
rank: int = 2,
seed: int = 0):
self.adapter_id = adapter_id
self.scale = float(scale)
rng = np.random.default_rng(seed)
@@ -308,18 +323,18 @@ class ToyLoRA:
self.b = (rng.standard_normal((rank, channels)) * 0.4).astype("float32")
def delta(self, latent, control=None) -> np.ndarray:
w = self.a @ self.b # [C, C] low-rank weight delta
d = np.tensordot(w, np.asarray(latent, dtype="float32"), axes=([1], [0])) # [C, ...]
w = self.a @ self.b # [C, C] low-rank weight delta
d = np.tensordot(w, np.asarray(latent, dtype="float32"), axes=([1], [0])) # [C, ...]
return (self.scale * np.tanh(d)).astype("float32")
def update(self, seed: int) -> None: # hot-swap: replace the adapter's weights
def update(self, seed: int) -> None: # hot-swap: replace the adapter's weights
rng = np.random.default_rng(seed)
self.a = (rng.standard_normal(self.a.shape) * 0.4).astype("float32")
class ToyControlNet:
"""Toy ControlNet adapter — conditions the velocity on a control signal (pose/depth/edge stand-in).
Different control ⇒ different generation; selected per request like a LoRA (design_v3 §9.19)."""
Different control ⇒ different generation; selected per request like a LoRA."""
kind = "controlnet"
def __init__(self, adapter_id: str, channels: int = LATENT_CHANNELS, scale: float = 0.8, seed: int = 1):
@@ -330,8 +345,8 @@ class ToyControlNet:
def delta(self, latent, control=None) -> np.ndarray:
latent = np.asarray(latent, dtype="float32")
c = float(np.mean(control)) if control is not None else 0.0 # the control image's signal
shape = (latent.shape[0],) + (1,) * (latent.ndim - 1)
c = float(np.mean(control)) if control is not None else 0.0 # the control image's signal
shape = (latent.shape[0], ) + (1, ) * (latent.ndim - 1)
return (self.scale * c * self.proj.reshape(shape)).astype("float32")
def update(self, seed: int) -> None:
@@ -366,7 +381,7 @@ class ToyDraftModel:
are accepted) and diverges otherwise (→ a rejection + a target correction). The accept rate varies
*within* a decode, so different rounds accept different lengths (a ragged loop). The demonstration is
exact regardless of draft quality: speculative decoding accepts only tokens the target would have
produced, so the output equals the target's own greedy decode (design_v3 §2.2)."""
produced, so the output equals the target's own greedy decode."""
VOCAB = 256
EOS = 0
@@ -376,9 +391,9 @@ class ToyDraftModel:
def ar_forward(self, tokens) -> int:
target = _spec_target_next(tokens)
s, n = sum(int(t) for t in tokens), len(tokens)
if ((s * 13 + n * 7) % 10) < int(round(self.agree * 10)): # agree on a varying subset
if ((s * 13 + n * 7) % 10) < int(round(self.agree * 10)): # agree on a varying subset
return int(target)
return int((target + 1) % self.VOCAB) # a cheap wrong guess (rejected)
return int((target + 1) % self.VOCAB) # a cheap wrong guess (rejected)
class ToyRewardModel:
@@ -387,7 +402,7 @@ class ToyRewardModel:
Maps a media latent → a scalar reward in [-1, 1] via a fixed projection over per-channel features.
The point isn't the math: it's that the reward is computed by a *model* (served, scheduled as
``REWARD_BATCH`` work units, place-able on its own pool), not a numpy heuristic bolted onto the
method — the reward plane composing with the serving plane (design_v3 §10)."""
method — the reward plane composing with the serving plane."""
def __init__(self, dim: int = LATENT_CHANNELS, seed: int = 5):
rng = np.random.default_rng(seed)
@@ -407,7 +422,7 @@ class ToyAudioVAE:
Decodes a small ``[C, A, 1, 1]`` audio latent (denoised jointly with the video latent) into a 1-D
waveform: channel-collapse via a fixed projection, then upsample along time. Deterministic, so the
joint A/V denoise loop stays interleave-safe. The video VAE and this share nothing — two decoders,
one synchronized latent pair (design_v3 §15b joint A/V)."""
one synchronized latent pair."""
def __init__(self, samples_per_frame: int = 16, seed: int = 2):
self.spf = int(samples_per_frame)
@@ -416,12 +431,12 @@ class ToyAudioVAE:
def decode(self, audio_latent: np.ndarray) -> np.ndarray:
a = np.asarray(audio_latent, dtype="float32")
flat = a.reshape(a.shape[0], -1) # [C, A·...]
flat = a.reshape(a.shape[0], -1) # [C, A·...]
# project over the actual channel count (LTX-2.3 audio latent is 8ch; the 2-stage toy is
# LATENT_CHANNELS) — np.resize is identity when they already match, so existing T2VS is unchanged.
proj = np.resize(self.proj, a.shape[0]).astype("float32")
mono = np.tensordot(proj, flat, axes=([0], [0])) # [A·...]
return np.tanh(np.repeat(mono, self.spf)).astype("float32") # [A·spf] mono waveform
mono = np.tensordot(proj, flat, axes=([0], [0])) # [A·...]
return np.tanh(np.repeat(mono, self.spf)).astype("float32") # [A·spf] mono waveform
class ToyVAE:
@@ -437,13 +452,13 @@ class ToyVAE:
def decode(self, latent: np.ndarray) -> np.ndarray:
"""latent [C,T,H,W] -> video [3, T, H*spatial, W*spatial] (deterministic upsample)."""
latent = np.asarray(latent, dtype=np.float32)
rgb = np.tensordot(self.dec_proj, latent, axes=([1], [0])) # [3,T,H,W]
rgb = np.tensordot(self.dec_proj, latent, axes=([1], [0])) # [3,T,H,W]
up = np.repeat(np.repeat(rgb, self.spatial, axis=2), self.spatial, axis=3)
return np.tanh(up).astype("float32")
def encode(self, video: np.ndarray) -> np.ndarray:
video = np.asarray(video, dtype=np.float32)
pooled = video[:self.C] # crude: take C channels
pooled = video[:self.C] # crude: take C channels
return pooled.astype("float32")
v2/platform/platform.py
+27 -30
View File
@@ -1,4 +1,4 @@
"""Platform — the detected ``(device, arch)`` that resolves both registries (design_v3 §17).
"""Platform — the detected ``(device, arch)`` that resolves both registries.
One ``Platform`` per pool. It owns:
* the **device fallback chain** (e.g. ``accel → cpu``) — so a component/kernel a backend hasn't
@@ -9,14 +9,14 @@ One ``Platform`` per pool. It owns:
``spec.factory`` call); and
* ``kernels`` — a per-platform ``KernelTable`` that resolves an op once and caches the callable.
``Platform.detect()`` is honest: it returns a CPU/numpy platform unless torch+CUDA are actually
present (they are not in this environment). The (device, arch) machinery is fully exercised on the
CPU box by the pure-python ``accel`` stand-in backend — proving the dispatch is generic without
faking a GPU.
``Platform.detect()`` returns a CPU/numpy platform unless torch+CUDA are actually present. The
(device, arch) machinery is exercised on a CPU box by the pure-python ``accel`` stand-in backend,
proving the dispatch is generic without a real GPU.
"""
from __future__ import annotations
from typing import Any, Iterator
from typing import Any
from collections.abc import Iterator
from v2.platform.registry import COMPONENTS, KERNELS
@@ -36,7 +36,7 @@ def ensure_backends_loaded() -> None:
# Arch fallback chains (newest → oldest → portable terminal). Resolution walks these in order.
_CUDA_ARCH_FALLBACK = ("sm100", "sm90", "sm80", "ptx", "generic")
_ACCEL_ARCH_FALLBACK = ("sm90", "sm80", "generic")
_PORTABLE_ARCHS = ("ptx", "generic") # vendor-neutral terminals — valid on any device
_PORTABLE_ARCHS = ("ptx", "generic") # vendor-neutral terminals — valid on any device
def _arch_rank(arch: str) -> int:
@@ -58,9 +58,7 @@ def _suffix_from(chain: tuple[str, ...], arch: str) -> tuple[str, ...]:
if a in tail:
continue
ar = _arch_rank(a)
if ar < 0: # portable terminal (ptx/generic): always valid
tail.append(a)
elif rank >= 0 and ar < rank: # strictly-older real arch: a safe fallback
if ar < 0 or rank >= 0 and ar < rank: # portable terminal (ptx/generic): always valid
tail.append(a)
# else: newer-or-equal real arch → skip (would be binary-incompatible on this device)
return tuple(tail)
@@ -70,7 +68,7 @@ class KernelTable:
"""A platform's resolved op→callable map. Resolves each op once (over the fallback chain),
then caches — so the hot loop pays the registry walk at most once per (op, variant)."""
def __init__(self, platform: "Platform"):
def __init__(self, platform: Platform):
self.platform = platform
self._cache: dict[tuple[str, str], Any] = {}
@@ -79,9 +77,8 @@ class KernelTable:
if ck not in self._cache:
reg = self.platform.resolve_kernel(op, variant)
if reg is None:
raise KeyError(
f"no kernel for op={op!r} variant={variant!r} on device chain "
f"{self.platform.device_chain} (and no numpy terminal registered)")
raise KeyError(f"no kernel for op={op!r} variant={variant!r} on device chain "
f"{self.platform.device_chain} (and no numpy terminal registered)")
self._cache[ck] = reg.fn
return self._cache[ck]
@@ -92,8 +89,8 @@ class KernelTable:
class Platform:
def __init__(self, device: str, arch: str, device_chain: tuple[str, ...],
arch_chains: dict[str, tuple[str, ...]]):
def __init__(self, device: str, arch: str, device_chain: tuple[str, ...], arch_chains: dict[str, tuple[str, ...]]):
self.device = device
self.arch = arch
self.device_chain = tuple(device_chain)
@@ -109,7 +106,7 @@ class Platform:
return self._kernels
def arch_chain(self, device: str) -> tuple[str, ...]:
return self.arch_chains.get(device, ("generic",))
return self.arch_chains.get(device, ("generic", ))
# --- candidate-key generators (the fallback chains, in priority order) --- #
def _component_keys(self, kind: str, variant: str) -> Iterator[tuple]:
@@ -135,10 +132,9 @@ class Platform:
if reg is not None:
return reg.fn(spec, instance, self)
if getattr(spec, "factory", None) is not None:
return spec.factory(instance) # the (kind, "cpu", "numpy") terminal rung
raise RuntimeError(
f"component kind={spec.kind!r} ({getattr(spec, 'component_id', '?')!r}) has no impl on "
f"device chain {self.device_chain} and no factory terminal")
return spec.factory(instance) # the (kind, "cpu", "numpy") terminal rung
raise RuntimeError(f"component kind={spec.kind!r} ({getattr(spec, 'component_id', '?')!r}) has no impl on "
f"device chain {self.device_chain} and no factory terminal")
def resolve_kernel(self, op: str, variant: str = "default"):
ensure_backends_loaded()
@@ -146,25 +142,26 @@ class Platform:
# --- constructors -------------------------------------------------------- #
@classmethod
def cpu(cls) -> "Platform":
return cls("cpu", "numpy", ("cpu",), {"cpu": ("numpy",)})
def cpu(cls) -> Platform:
return cls("cpu", "numpy", ("cpu", ), {"cpu": ("numpy", )})
@classmethod
def accel(cls, arch: str = "sm90") -> "Platform":
def accel(cls, arch: str = "sm90") -> Platform:
"""A pure-python stand-in accelerator (CPU-resident) used to prove cross-device dispatch,
arch fallback, and the parity oracle without a real GPU."""
return cls("accel", arch, ("accel", "cpu"),
{"accel": _suffix_from(_ACCEL_ARCH_FALLBACK, arch), "cpu": ("numpy",)})
return cls("accel", arch, ("accel", "cpu"), {
"accel": _suffix_from(_ACCEL_ARCH_FALLBACK, arch),
"cpu": ("numpy", )
})
@classmethod
def cuda(cls, arch: str = "sm90") -> "Platform":
def cuda(cls, arch: str = "sm90") -> Platform:
"""A real CUDA platform (resolves torch adapters + fastvideo-kernel cells). Only usable on a
box where those cells are ``available``; declared here so detection has a target."""
return cls("cuda", arch, ("cuda", "cpu"),
{"cuda": _suffix_from(_CUDA_ARCH_FALLBACK, arch), "cpu": ("numpy",)})
return cls("cuda", arch, ("cuda", "cpu"), {"cuda": _suffix_from(_CUDA_ARCH_FALLBACK, arch), "cpu": ("numpy", )})
@classmethod
def detect(cls) -> "Platform":
def detect(cls) -> Platform:
"""Honest detection: CUDA iff torch+CUDA are actually importable, else CPU/numpy."""
ensure_backends_loaded()
import importlib.util
v2/platform/registry.py
+41 -20
View File
@@ -1,4 +1,4 @@
"""The two tuple-keyed backend registries — the dispatch membrane (design_v3 §17; multi-backend RFC).
"""The two tuple-keyed backend registries — the dispatch membrane.
Everything device/arch-specific resolves through one of two registries, keyed by a tuple:
@@ -33,13 +33,14 @@ are imported lazily (``Platform.ensure_backends_loaded``) to keep this import-cy
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Callable, Iterable
from typing import Any
from collections.abc import Callable, Iterable
# --------------------------------------------------------------------------- #
# Op names — the stable identifiers a loop dispatches on (not free functions). #
# --------------------------------------------------------------------------- #
FLOW_MATCH_STEP = "flow_match_step" # deterministic flow-match Euler solver step (ODE serve)
FLOW_SDE_STEP = "flow_sde_step" # FlowGRPO stochastic step + log-prob (RL rollout)
FLOW_MATCH_STEP = "flow_match_step" # deterministic flow-match Euler solver step (ODE serve)
FLOW_SDE_STEP = "flow_sde_step" # FlowGRPO stochastic step + log-prob (RL rollout)
def _always_available() -> bool:
@@ -75,14 +76,20 @@ class TupleRegistry:
self.key_fields = key_fields
self._reg: dict[tuple, Registration] = {}
def put(self, key: tuple, fn: Callable[..., Any], *,
available: Callable[[], bool] | None = None, source: str = "",
def put(self,
key: tuple,
fn: Callable[..., Any],
*,
available: Callable[[], bool] | None = None,
source: str = "",
meta: dict[str, Any] | None = None) -> None:
if len(key) != len(self.key_fields):
raise ValueError(f"{self.name}: key {key} does not match fields {self.key_fields}")
self._reg[key] = Registration(key=key, fn=fn,
self._reg[key] = Registration(key=key,
fn=fn,
available=available or _always_available,
source=source, meta=dict(meta or {}))
source=source,
meta=dict(meta or {}))
def lookup(self, key: tuple) -> Registration | None:
"""Exact lookup, ignoring availability (used by ``manifest``/diagnostics)."""
@@ -104,33 +111,47 @@ class TupleRegistry:
"""
rows: list[dict[str, Any]] = []
for key, reg in sorted(self._reg.items(), key=lambda kv: tuple(map(str, kv[0]))):
row = dict(zip(self.key_fields, key))
row = dict(zip(self.key_fields, key, strict=False))
row["available"] = reg.is_available()
row["source"] = reg.source
row.update(reg.meta) # surface declared metadata (e.g. workspace_bytes)
row.update(reg.meta) # surface declared metadata (e.g. workspace_bytes)
rows.append(row)
return rows
# The two registries (design_v3 §17). Populated by ``backends/`` modules (lazily imported).
# The two registries, populated by ``backends/`` modules (lazily imported).
COMPONENTS = TupleRegistry("components", ("kind", "device", "variant"))
KERNELS = TupleRegistry("kernels", ("op", "device", "arch", "variant"))
def register_component(kind: str, fn: Callable[..., Any], *, device: str, variant: str = "default",
available: Callable[[], bool] | None = None, source: str = "") -> None:
def register_component(kind: str,
fn: Callable[..., Any],
*,
device: str,
variant: str = "default",
available: Callable[[], bool] | None = None,
source: str = "") -> None:
"""Register a component builder. ``fn(spec, instance, platform) -> live component``."""
COMPONENTS.put((kind, device, variant), fn, available=available, source=source)
def register_kernel(op: str, fn: Callable[..., Any], *, device: str, arch: str, variant: str = "default",
available: Callable[[], bool] | None = None, source: str = "",
def register_kernel(op: str,
fn: Callable[..., Any],
*,
device: str,
arch: str,
variant: str = "default",
available: Callable[[], bool] | None = None,
source: str = "",
workspace_bytes: int = 0) -> None:
"""Register a stateless kernel. ``fn(*args, **kwargs)`` — same signature as its numpy reference.
``workspace_bytes`` is the kernel's static scratch requirement — the capture-safety contract
(design_v3 §6.2): a kernel used inside a captured CUDA graph must draw its workspace from the
pool-provided static buffer, not malloc per call, or replay corrupts. Declared here so the
requirement is enumerable in the matrix; the numpy reference needs none (0)."""
KERNELS.put((op, device, arch, variant), fn, available=available, source=source,
``workspace_bytes`` is the kernel's static scratch requirement — the capture-safety contract: a
kernel used inside a captured CUDA graph must draw its workspace from the pool-provided static
buffer, not malloc per call, or replay corrupts. Declared here so the requirement is enumerable
in the matrix; the numpy reference needs none (0)."""
KERNELS.put((op, device, arch, variant),
fn,
available=available,
source=source,
meta={"workspace_bytes": int(workspace_bytes)})
v2/program/__init__.py
+5 -4
View File
@@ -1,4 +1,4 @@
"""Program plane — compose a card's loops into a task (design_v3 §13)."""
"""Program plane — compose a card's loops into a task."""
from __future__ import annotations
from v2.program.specs import (
@@ -21,6 +21,7 @@ from v2.program.workflow import (
WorkflowStage,
)
__all__ = ["Program", "ProgramKind", "ProgramNode", "ComponentNode", "ModelLoopNode",
"Edge", "EdgeKind", "always", "when_task", "when_opt",
"Workflow", "WorkflowStage", "WorkflowRegistry", "ParallelWorkflow", "BestOfNWorkflow"]
__all__ = [
"Program", "ProgramKind", "ProgramNode", "ComponentNode", "ModelLoopNode", "Edge", "EdgeKind", "always",
"when_task", "when_opt", "Workflow", "WorkflowStage", "WorkflowRegistry", "ParallelWorkflow", "BestOfNWorkflow"
]
v2/program/specs.py
+10 -9
View File
@@ -1,4 +1,4 @@
"""Programs compose a card's loops into a task (design_v3 §13).
"""Programs compose a card's loops into a task.
> The card says what loops *exist*, the program says how to *run* them for this request.
@@ -11,7 +11,8 @@ from __future__ import annotations
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable
from typing import Any
from collections.abc import Callable
from v2.request.tasks import TaskType
@@ -24,7 +25,7 @@ class ProgramKind(str, Enum):
WORKFLOW = "workflow"
# --- node predicates (replace geometry heuristics; design.md P7) -------------- #
# --- node predicates (declarative, replacing geometry heuristics) ------------- #
def always(_request: Any) -> bool:
return True
@@ -43,8 +44,8 @@ def when_opt(node_id: str, key: str) -> Callable[[Any], bool]:
class ProgramNode:
node_id: str
when: Callable[[Any], bool] = always
reads: tuple[str, ...] = () # input slot names (typed edges in)
writes: tuple[str, ...] = () # output slot names (typed edges out)
reads: tuple[str, ...] = () # input slot names (typed edges in)
writes: tuple[str, ...] = () # output slot names (typed edges out)
@dataclass
@@ -59,7 +60,7 @@ class ComponentNode(ProgramNode):
@dataclass
class ModelLoopNode(ProgramNode):
"""Binds one of the card's loops. The engine drives it via a LoopRunner, one step per tick,
so its steps interleave with other requests' steps (design_v3 §5, §6.3)."""
so its steps interleave with other requests' steps."""
loop_id: str = ""
output_slot: str = "latents"
@@ -76,7 +77,7 @@ class EdgeKind(str, Enum):
@dataclass
class Edge:
src: str # slot or node id
src: str # slot or node id
dst: str
kind: EdgeKind = EdgeKind.TENSOR
@@ -88,13 +89,13 @@ class Program:
kind: ProgramKind
nodes: list[ProgramNode] = field(default_factory=list)
edges: list[Edge] = field(default_factory=list)
# artifact name -> slot that holds its value at the end (design_v3 §12 named artifacts)
# artifact name -> slot that holds its value at the end
output_artifacts: dict[str, str] = field(default_factory=dict)
def active_nodes(self, request: Any) -> list[ProgramNode]:
return [n for n in self.nodes if n.when(request)]
def validate(self) -> "Program":
def validate(self) -> Program:
produced: set[str] = set()
for n in self.nodes:
for r in n.reads:
v2/program/workflow.py
+12 -11
View File
@@ -1,4 +1,4 @@
"""Workflow — multi-MODEL composition above the engine (design_v3 §13; ``ProgramKind.WORKFLOW``).
"""Workflow — multi-MODEL composition above the engine (``ProgramKind.WORKFLOW``).
A ``Program`` composes the loops of ONE resident model instance (the omni/MoT default — every
``ModelLoopNode`` resolves on the program's single instance). A ``Workflow`` composes *across model
@@ -21,7 +21,8 @@ the program runner is what preserves the interleave-parity guarantee for the sin
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Callable
from typing import Any
from collections.abc import Callable
@dataclass
@@ -36,8 +37,8 @@ class WorkflowStage:
def _engine_serves(engine: Any, model_id: str) -> bool:
if hasattr(engine, "serves"):
return engine.serves(model_id) # AsyncEngine
return model_id in getattr(engine, "_registry", {}) # Engine
return engine.serves(model_id) # AsyncEngine
return model_id in getattr(engine, "_registry", {}) # Engine
@dataclass
@@ -46,7 +47,7 @@ class Workflow:
the SAME namespace as a card's ``model_id`` (so the engine, server, and fleet route to it the same
way), and by convention is dotted/namespaced to avoid colliding with a card id (e.g.
``"image_video.t2i_i2v"``). ``requires`` is the set of cards it composes — declared, validated,
never inferred (design.md P7)."""
never inferred."""
workflow_id: str
stages: list[WorkflowStage] = field(default_factory=list)
@@ -62,7 +63,7 @@ class Workflow:
out.append(s.model_id)
return out
def validate(self, engine: Any) -> "Workflow":
def validate(self, engine: Any) -> Workflow:
"""Fail-fast: every required card must already be registered on ``engine`` (no silent skips)."""
if not self.stages:
raise ValueError(f"workflow {self.workflow_id!r} has no stages")
@@ -106,7 +107,7 @@ class ParallelWorkflow(Workflow):
"""Fan-out: run every stage on the SAME initial input (conceptually in parallel — the engine can
interleave their steps), then merge their artifacts into one Output, namespaced by stage label.
The non-linear shape a linear chain can't express: one prompt → N models → a merged result (e.g.
variants, or video+audio+upscale in parallel). Design_v3 §13."""
variants, or video+audio+upscale in parallel)."""
def run(self, engine: Any, **initial: Any) -> Any:
if not self.stages:
@@ -117,13 +118,13 @@ class ParallelWorkflow(Workflow):
rid = self.workflow_id
for stage in self.stages:
label = stage.label or stage.model_id
req = stage.make_request(dict(initial)) # each branch sees the ORIGINAL input
req = stage.make_request(dict(initial)) # each branch sees the ORIGINAL input
if req.model_id != stage.model_id:
raise ValueError(f"{self.workflow_id!r} stage {label!r}: model mismatch")
out = engine.run(req)
rid = out.request_id
for name, art in out.artifacts.items():
merged[f"{label}:{name}"] = art # namespaced — branches don't collide
merged[f"{label}:{name}"] = art # namespaced — branches don't collide
for k, v in out.metrics.items():
metrics[f"{label}:{k}"] = v
return Output(request_id=rid, artifacts=merged, metrics=metrics)
@@ -132,11 +133,11 @@ class ParallelWorkflow(Workflow):
class BestOfNWorkflow(Workflow):
"""Inference-time scaling / rejection sampling: generate N candidates (varying seed), score each with
a reward scorer, return the best. A feedback loop across models — generator + reward (e.g. the served
REWARD_BATCH card, §9.15) — the other non-linear shape. Design_v3 §10/§13."""
REWARD_BATCH card) — the other non-linear shape."""
def __init__(self, workflow_id, generator_stage, *, scorer, n: int = 4, score_key: str = "latents"):
super().__init__(workflow_id, [generator_stage])
self.scorer = scorer # any .score(media, prompts) -> {"avg": ...}
self.scorer = scorer # any .score(media, prompts) -> {"avg": ...}
self.n = int(n)
self.score_key = score_key
v2/recipes/__init__.py
+40 -21
View File
@@ -1,4 +1,4 @@
"""Model definitions — concrete (recipe, runtime) cards (design_v3 §4).
"""Model definitions — concrete (recipe, runtime) cards.
Phase 1: Wan2.1-1.3B (T2V), LTX2.3 (2-stage distilled), Wan-causal (self-forcing student).
``build_default_engine`` loads all three onto one engine (one resident instance per card).
@@ -29,20 +29,40 @@ from v2.recipes.wan21 import build_wan21_card, build_wan_t2v_program
from v2.recipes.wan_causal import build_wan_causal_card, build_wan_causal_program
__all__ = [
"build_wan21_card", "build_wan_t2v_program",
"build_ltx2_card", "build_ltx2_program", "build_ltx2_av_program",
"build_wan_causal_card", "build_wan_causal_program",
"build_cosmos3_card", "build_cosmos3_program",
"build_bagel_card", "build_bagel_program",
"build_qwen_omni_card", "build_qwen_omni_program",
"build_wan21_card",
"build_wan_t2v_program",
"build_ltx2_card",
"build_ltx2_program",
"build_ltx2_av_program",
"build_wan_causal_card",
"build_wan_causal_program",
"build_cosmos3_card",
"build_cosmos3_program",
"build_bagel_card",
"build_bagel_program",
"build_qwen_omni_card",
"build_qwen_omni_program",
"build_reward_card",
"build_speculative_card", "build_speculative_program",
"build_unified_card", "build_unified_program",
"build_flux_t2i_card", "build_flux_t2i_program", "build_wan_i2v_card", "build_wan_i2v_program",
"build_t2i_then_i2v_workflow", "build_t2i_i2v_extend_workflow", "register_workflows",
"build_tiled_card", "build_tiled_program", "build_adaptive_card",
"build_adapter_card", "build_adapter_program",
"build_default_engine", "build_omni_engine", "build_unified_engine", "build_image_video_engine",
"build_speculative_card",
"build_speculative_program",
"build_unified_card",
"build_unified_program",
"build_flux_t2i_card",
"build_flux_t2i_program",
"build_wan_i2v_card",
"build_wan_i2v_program",
"build_t2i_then_i2v_workflow",
"build_t2i_i2v_extend_workflow",
"register_workflows",
"build_tiled_card",
"build_tiled_program",
"build_adaptive_card",
"build_adapter_card",
"build_adapter_program",
"build_default_engine",
"build_omni_engine",
"build_unified_engine",
"build_image_video_engine",
"build_tiled_engine",
]
@@ -78,7 +98,7 @@ def build_omni_engine(engine: Any = None) -> Any:
"""Register the phase-2 omni cards (Cosmos3 + BAGEL/lance) onto one engine.
Each is ONE resident MoT instance whose ``transformer`` runs BOTH an ar_decode loop and a
diffusion_denoise loop (shared weights) — the §16 'true omni/MoT serving' claim made native.
diffusion_denoise loop (shared weights) — true omni/MoT serving.
"""
from v2.cache import CacheManager
from v2.card import load_card
@@ -95,9 +115,8 @@ def build_omni_engine(engine: Any = None) -> Any:
# Adding a custom pipeline is one line here (the cross-model analog of _BUILDERS / _OMNI_BUILDERS;
# cf. vllm-omni's pipeline_registry). ``register_workflows`` registers each whose cards are present.
_WORKFLOWS: dict[str, tuple] = {
"image_video.t2i_i2v": (build_t2i_then_i2v_workflow,
[(build_flux_t2i_card, build_flux_t2i_program),
(build_wan_i2v_card, build_wan_i2v_program)]),
"image_video.t2i_i2v": (build_t2i_then_i2v_workflow, [(build_flux_t2i_card, build_flux_t2i_program),
(build_wan_i2v_card, build_wan_i2v_program)]),
}
@@ -121,7 +140,7 @@ def register_workflows(engine: Any, *, only: list[str] | None = None) -> Any:
def build_image_video_engine(engine: Any = None) -> Any:
"""Register the T2I (``flux-t2i``) and I2V (``wan-i2v``) cards plus the ``image_video.t2i_i2v``
workflow on one engine — two *separate* models chained by a cross-model workflow, addressable by
its workflow_id like any servable (design_v3 §13)."""
its workflow_id like any servable."""
from v2.runtime import Engine
eng = engine if engine is not None else Engine()
return register_workflows(eng, only=["image_video.t2i_i2v"])
@@ -130,7 +149,7 @@ def build_image_video_engine(engine: Any = None) -> Any:
def build_tiled_engine(engine: Any = None) -> Any:
"""Register the tiled-decode card (``wan-tiled``, whose VAE decode emits ``VAE_TILE`` work units)
alongside plain Wan2.1 — so an interleaved batch mixes ``VAE_TILE`` and ``DIFFUSION_STEP`` units
through one admission/scheduler budget (the §17 heterogeneous-co-scheduling probe)."""
through one admission/scheduler budget (the heterogeneous co-scheduling probe)."""
from v2.cache import CacheManager
from v2.card import load_card
from v2.runtime import Engine
@@ -146,7 +165,7 @@ def build_tiled_engine(engine: Any = None) -> Any:
def build_unified_engine(engine: Any = None) -> Any:
"""Register the unified LM+generator card (UniRL/PromptRL): a prompt-refiner ``llm`` expert + a
flow ``transformer`` expert on the SAME engine, each driving its own loop. Two separate experts,
one request, both trainable under one RL reward (the §10 joint-RL stress test)."""
one request, both trainable under one RL reward (the joint-RL stress test)."""
from v2.cache import CacheManager
from v2.card import load_card
from v2.runtime import Engine
v2/recipes/_prompts.py
+2 -2
View File
@@ -1,6 +1,6 @@
"""Default negative prompts per model family — copied verbatim from fastvideo's pipeline presets
(`fastvideo/pipelines/basic/{wan,ltx2}/presets.py`). These are stable recipe DATA (not model code), so
v2 owns its own copy rather than importing private symbols from fastvideo's pipelines package."""
(`fastvideo/pipelines/basic/{wan,ltx2}/presets.py`). These are stable recipe DATA, not model code,
so v2 keeps its own copy rather than importing private symbols from fastvideo's pipelines package."""
from __future__ import annotations
# Wan family — English (Wan2.1 T2V) and Chinese (Wan2.2 / self-forcing) negative prompts.
v2/recipes/adapters/__init__.py
+1 -1
View File
@@ -1,4 +1,4 @@
"""Adapter plane: one base + swappable LoRA/ControlNet adapters, selected per request (design_v3 §9.19)."""
"""Adapter plane: one base + swappable LoRA/ControlNet adapters, selected per request."""
from __future__ import annotations
from v2.recipes.adapters.card import ADAPTERS, build_adapter_card
v2/recipes/adapters/card.py
+53 -28
View File
@@ -1,4 +1,4 @@
"""Adapter-serving card — one base DiT + a set of swappable LoRA/ControlNet adapters (design_v3 §9.19).
"""Adapter-serving card — one base DiT + a set of swappable LoRA/ControlNet adapters.
Many adapters are declared as components over one resident ``transformer``; a request picks which to
apply (``DiffusionParams.adapters``). Adapters are versioned independently (the cache key's
@@ -24,7 +24,7 @@ from v2.parallel import ParallelPlan
from v2.platform.backends.toy import ToyControlNet, ToyDiT, ToyLoRA, ToyTextEncoder, ToyVAE, _seed_from
from v2.recipes.adapters.loop import AdapterDenoiseLoop
ADAPTERS = ("lora_anime", "lora_realistic", "control_pose") # the served adapter library
ADAPTERS = ("lora_anime", "lora_realistic", "control_pose") # the served adapter library
def build_adapter_card(model_id: str = "wan-adapters") -> ModelCard:
@@ -34,40 +34,65 @@ def build_adapter_card(model_id: str = "wan-adapters") -> ModelCard:
precision, expert = PrecisionPolicy(), NoRouting("transformer")
def loop_factory():
return AdapterDenoiseLoop(loop_id="diffusion_denoise", cfg=cfg, flow_shift=flow,
precision=precision, expert=expert, cost=cost)
return AdapterDenoiseLoop(loop_id="diffusion_denoise",
cfg=cfg,
flow_shift=flow,
precision=precision,
expert=expert,
cost=cost)
components = {
"text_encoder": ComponentSpec("text_encoder", kind="text_encoder",
factory=lambda inst: ToyTextEncoder(), required_for={"t2v"}),
"transformer": ComponentSpec("transformer", kind="dit", factory=lambda inst: ToyDiT(seed=seed),
resident_for=["diffusion_denoise"], required_for={"t2v"}),
"vae": ComponentSpec("vae", kind="vae", factory=lambda inst: ToyVAE(), required_for={"t2v"}),
"text_encoder":
ComponentSpec("text_encoder", kind="text_encoder", factory=lambda inst: ToyTextEncoder(), required_for={"t2v"}),
"transformer":
ComponentSpec("transformer",
kind="dit",
factory=lambda inst: ToyDiT(seed=seed),
resident_for=["diffusion_denoise"],
required_for={"t2v"}),
"vae":
ComponentSpec("vae", kind="vae", factory=lambda inst: ToyVAE(), required_for={"t2v"}),
# the adapter library — lightweight, optional, applied per request:
"lora_anime": ComponentSpec("lora_anime", kind="adapter",
factory=lambda inst: ToyLoRA("lora_anime", scale=0.6, seed=seed + 1),
optional_for={"t2v"}),
"lora_realistic": ComponentSpec("lora_realistic", kind="adapter",
factory=lambda inst: ToyLoRA("lora_realistic", scale=0.6, seed=seed + 2),
optional_for={"t2v"}),
"control_pose": ComponentSpec("control_pose", kind="adapter",
factory=lambda inst: ToyControlNet("control_pose", scale=0.8, seed=seed + 3),
optional_for={"t2v"}),
"lora_anime":
ComponentSpec("lora_anime",
kind="adapter",
factory=lambda inst: ToyLoRA("lora_anime", scale=0.6, seed=seed + 1),
optional_for={"t2v"}),
"lora_realistic":
ComponentSpec("lora_realistic",
kind="adapter",
factory=lambda inst: ToyLoRA("lora_realistic", scale=0.6, seed=seed + 2),
optional_for={"t2v"}),
"control_pose":
ComponentSpec("control_pose",
kind="adapter",
factory=lambda inst: ToyControlNet("control_pose", scale=0.8, seed=seed + 3),
optional_for={"t2v"}),
}
loops = {
"diffusion_denoise": LoopSpec("diffusion_denoise", kind=LoopKind.DIFFUSION_DENOISE,
work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=cost,
shared_weight_components=["transformer"], cache_policy=["feature"],
loop_factory=loop_factory),
"diffusion_denoise":
LoopSpec("diffusion_denoise",
kind=LoopKind.DIFFUSION_DENOISE,
work_unit_kind=WorkUnitKind.DIFFUSION_STEP,
step_cost_model=cost,
shared_weight_components=["transformer"],
cache_policy=["feature"],
loop_factory=loop_factory),
}
return ModelCard(
model_id=model_id, family="wan", components=components, loops=loops,
model_id=model_id,
family="wan",
components=components,
loops=loops,
capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.VAE_DECODE),
recipe=RecipeSpec(method="base", assumes_loop="diffusion_denoise",
assumes_precision="float32", consistency_required=ConsistencyLevel.C1),
recipe=RecipeSpec(method="base",
assumes_loop="diffusion_denoise",
assumes_precision="float32",
consistency_required=ConsistencyLevel.C1),
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True),
caches={"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True)},
caches={
"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True)
},
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()],
default_plan=ParallelPlan.single()),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
).validate()
v2/recipes/adapters/loop.py
+11 -10
View File
@@ -1,10 +1,10 @@
"""AdapterDenoiseLoop — per-request LoRA / ControlNet over one resident base (design_v3 §9.19).
"""AdapterDenoiseLoop — per-request LoRA / ControlNet over one resident base.
The adapter plane: a request selects adapters (``DiffusionParams.adapters``); the loop applies each
active adapter's velocity delta on top of the base DiT's prediction, then re-integrates. Many adapters
are served over ONE resident base; the active set lives in the request/``LoopState`` (never global), so
two requests using different adapters interleave without smearing. Isolated ``WanDenoiseLoop`` subclass
(empty adapter set ⇒ identical to the base), so nothing else is touched.
A request selects adapters (``DiffusionParams.adapters``); the loop applies each active adapter's
velocity delta on top of the base DiT's prediction, then re-integrates. Many adapters are served over
ONE resident base; the active set lives in the request/``LoopState`` (never global), so two requests
using different adapters interleave without smearing. An empty adapter set is identical to the base
``WanDenoiseLoop``.
"""
from __future__ import annotations
@@ -16,10 +16,11 @@ from v2.recipes.wan21.loop import WanDenoiseLoop
class AdapterDenoiseLoop(WanDenoiseLoop):
def init(self, req, model, ctx):
st = super().init(req, model, ctx)
st.scratch["adapters"] = tuple(getattr(req.diffusion, "adapters", ()) or ())
st.scratch["control"] = ctx.slots.get("control_signal") # ControlNet conditioning (optional)
st.scratch["control"] = ctx.slots.get("control_signal") # ControlNet conditioning (optional)
return st
def next(self, st):
@@ -34,12 +35,12 @@ class AdapterDenoiseLoop(WanDenoiseLoop):
prec = self.precision
def run(model, override=None):
res = base_run(model, override) # base velocity (CFG combine)
res = base_run(model, override) # base velocity (CFG combine)
v = np.asarray(res.output["noise_pred"], dtype="float32")
for aid in adapters: # apply each active adapter's delta
for aid in adapters: # apply each active adapter's delta
v = v + model.component(aid).delta(x, control)
res.output["noise_pred"] = v
fm = model.platform.kernels.get(FLOW_MATCH_STEP) # solver dispatched per (device, arch)
fm = model.platform.kernels.get(FLOW_MATCH_STEP) # solver dispatched per (device, arch)
res.output["latents"] = fm(prec.cast(x), v, sigma_t, sigma_next).astype("float32")
return res
v2/recipes/adaptive/__init__.py
+1 -1
View File
@@ -1,4 +1,4 @@
"""Adaptive-compute denoise: the loop owns cache-dit skip + early-exit (design_v3 §2.2, §9.12)."""
"""Adaptive-compute denoise: the loop owns cache-dit skip + early-exit."""
from __future__ import annotations
from v2.recipes.adaptive.card import build_adaptive_card
v2/recipes/adaptive/card.py
+43 -21
View File
@@ -1,6 +1,6 @@
"""Adaptive-compute card — a Wan-like T2V model whose denoise loop owns content-adaptive control flow
(cache-dit skip + early-exit; design_v3 §2.2, §9.12). Reuses the Wan T2V program (loop_id
``diffusion_denoise``); only the loop differs."""
(cache-dit skip + early-exit). Reuses the Wan T2V program (loop_id ``diffusion_denoise``); only the
loop differs."""
from __future__ import annotations
from v2._enums import Capability, ConsistencyLevel, LoopKind, WorkUnitKind
@@ -22,7 +22,9 @@ from v2.platform.backends.toy import ToyDiT, ToyTextEncoder, ToyVAE, _seed_from
from v2.recipes.adaptive.loop import CacheDiTDenoiseLoop
def build_adaptive_card(model_id: str = "wan-adaptive", *, cache_threshold: float = 0.02,
def build_adaptive_card(model_id: str = "wan-adaptive",
*,
cache_threshold: float = 0.02,
exit_threshold: float = 0.0) -> ModelCard:
seed = _seed_from(model_id)
cost = CostModel(kind=WorkUnitKind.DIFFUSION_STEP, base_seconds=1e-4, per_unit_seconds=1e-7)
@@ -30,31 +32,51 @@ def build_adaptive_card(model_id: str = "wan-adaptive", *, cache_threshold: floa
precision, expert = PrecisionPolicy(), NoRouting("transformer")
def loop_factory():
return CacheDiTDenoiseLoop(loop_id="diffusion_denoise", cfg=cfg, flow_shift=flow,
precision=precision, expert=expert, cost=cost,
cache_threshold=cache_threshold, exit_threshold=exit_threshold)
return CacheDiTDenoiseLoop(loop_id="diffusion_denoise",
cfg=cfg,
flow_shift=flow,
precision=precision,
expert=expert,
cost=cost,
cache_threshold=cache_threshold,
exit_threshold=exit_threshold)
components = {
"text_encoder": ComponentSpec("text_encoder", kind="text_encoder",
factory=lambda inst: ToyTextEncoder(), required_for={"t2v"}),
"transformer": ComponentSpec("transformer", kind="dit", factory=lambda inst: ToyDiT(seed=seed),
resident_for=["diffusion_denoise"], required_for={"t2v"}),
"vae": ComponentSpec("vae", kind="vae", factory=lambda inst: ToyVAE(), required_for={"t2v"}),
"text_encoder":
ComponentSpec("text_encoder", kind="text_encoder", factory=lambda inst: ToyTextEncoder(), required_for={"t2v"}),
"transformer":
ComponentSpec("transformer",
kind="dit",
factory=lambda inst: ToyDiT(seed=seed),
resident_for=["diffusion_denoise"],
required_for={"t2v"}),
"vae":
ComponentSpec("vae", kind="vae", factory=lambda inst: ToyVAE(), required_for={"t2v"}),
}
loops = {
"diffusion_denoise": LoopSpec("diffusion_denoise", kind=LoopKind.DIFFUSION_DENOISE,
work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=cost,
shared_weight_components=["transformer"], cache_policy=["feature"],
loop_factory=loop_factory),
"diffusion_denoise":
LoopSpec("diffusion_denoise",
kind=LoopKind.DIFFUSION_DENOISE,
work_unit_kind=WorkUnitKind.DIFFUSION_STEP,
step_cost_model=cost,
shared_weight_components=["transformer"],
cache_policy=["feature"],
loop_factory=loop_factory),
}
return ModelCard(
model_id=model_id, family="wan", components=components, loops=loops,
model_id=model_id,
family="wan",
components=components,
loops=loops,
capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.VAE_DECODE),
recipe=RecipeSpec(method="base", assumes_loop="diffusion_denoise",
assumes_precision="float32", consistency_required=ConsistencyLevel.C1),
recipe=RecipeSpec(method="base",
assumes_loop="diffusion_denoise",
assumes_precision="float32",
consistency_required=ConsistencyLevel.C1),
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True),
caches={"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True)},
caches={
"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True)
},
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()],
default_plan=ParallelPlan.single()),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
).validate()
v2/recipes/adaptive/loop.py
+18 -21
View File
@@ -1,16 +1,13 @@
"""CacheDiTDenoiseLoop — the loop OWNS content-adaptive control flow (design_v3 §2.2).
"""CacheDiTDenoiseLoop — the loop owns content-adaptive control flow.
The signature payoff of loop inversion: because the *model* owns control flow, content-adaptive
decisions are natural. This loop, at each step, decides whether to:
Because the loop owns control flow, each step can decide whether to:
* **skip the expensive forward** and reuse the cached velocity when consecutive predictions barely
change (cache-dit / TeaCache / Δ-DiT — a 1.5–2× inference win), or
* **early-exit** when the latent has converged.
* skip the forward and reuse the cached velocity when consecutive predictions barely change
(cache-dit / TeaCache / Δ-DiT), or
* early-exit once the latent has converged.
The result is a *variable, content-dependent* step count — impossible in a ``for t in timesteps``
runtime, trivial in a driven loop whose ``next()`` reads ``LoopState``. It is fully isolated from the
shared ``WanDenoiseLoop`` (subclass, threshold 0 ⇒ identical behavior), and still interleave-safe:
different requests skip different steps, and the parity gate still holds (ragged loops don't smear).
This yields a variable, content-dependent step count. Subclass of ``WanDenoiseLoop`` (threshold 0 ⇒
identical behavior) and interleave-safe: different requests skip different steps without smearing.
"""
from __future__ import annotations
@@ -22,10 +19,11 @@ from v2.recipes.wan21.loop import WanDenoiseLoop
class CacheDiTDenoiseLoop(WanDenoiseLoop):
def __init__(self, *, cache_threshold: float = 0.0, exit_threshold: float = 0.0, **kw):
super().__init__(**kw)
self.cache_threshold = cache_threshold # rel velocity change below which we reuse (skip the DiT)
self.exit_threshold = exit_threshold # rel latent change below which we stop early
self.cache_threshold = cache_threshold # rel velocity change below which we reuse (skip the DiT)
self.exit_threshold = exit_threshold # rel latent change below which we stop early
def init(self, req, model, ctx):
st = super().init(req, model, ctx)
@@ -34,8 +32,8 @@ class CacheDiTDenoiseLoop(WanDenoiseLoop):
def next(self, st):
if st.scratch.get("exited"):
return Done() # converged early — content-adaptive termination
plan = super().next(st) # the normal full-forward step plan
return Done() # converged early — content-adaptive termination
plan = super().next(st) # the normal full-forward step plan
last_v = st.scratch.get("last_v")
if isinstance(plan, Done) or not self.cache_threshold or not st.scratch.get("skip_next") \
or last_v is None:
@@ -57,21 +55,20 @@ class CacheDiTDenoiseLoop(WanDenoiseLoop):
x_prev = np.asarray(st.latents["video"], dtype="float64")
v = np.asarray(result.output["noise_pred"], dtype="float64")
prev_v = st.scratch.get("last_v")
st = super().advance(st, result) # folds latents, step_idx, (rollout) trajectory
st = super().advance(st, result) # folds latents, step_idx, (rollout) trajectory
if result.output.get("skipped"):
st.scratch["skipped"] += 1
st.scratch["skip_next"] = False # don't chain skips — re-check with a full forward
st.scratch["skip_next"] = False # don't chain skips — re-check with a full forward
elif self.cache_threshold and prev_v is not None:
rel = float(np.linalg.norm(v - np.asarray(prev_v, dtype="float64"))
/ (np.linalg.norm(prev_v) + 1e-8))
st.scratch["skip_next"] = rel < self.cache_threshold # next prediction will barely change
rel = float(np.linalg.norm(v - np.asarray(prev_v, dtype="float64")) / (np.linalg.norm(prev_v) + 1e-8))
st.scratch["skip_next"] = rel < self.cache_threshold # next prediction will barely change
st.scratch["last_v"] = result.output["noise_pred"]
half = max(1, (len(st.sigmas) - 1) // 2) # only consider exiting in the low-noise tail
half = max(1, (len(st.sigmas) - 1) // 2) # only consider exiting in the low-noise tail
if self.exit_threshold and st.step_idx >= half:
x_next = np.asarray(st.latents["video"], dtype="float64")
dl = float(np.linalg.norm(x_next - x_prev) / (np.linalg.norm(x_prev) + 1e-8))
if dl < self.exit_threshold:
st.scratch["exited"] = True # latent converged → stop next tick
st.scratch["exited"] = True # latent converged → stop next tick
return st
def finalize(self, st):
v2/recipes/bagel/card.py
+53 -34
View File
@@ -1,13 +1,12 @@
"""BAGEL/lance-style MoT ModelCard — the canonical vllm-omni omni model (design_v3 §1, §16, §19).
"""BAGEL/lance-style MoT ModelCard — canonical vllm-omni omni model.
vllm-omni's ``bagel_single_stage``/``lance`` prove one resident MoT instance can run both AR
``generate_text`` and diffusion ``generate_image`` on co-resident experts in a single request — BUT
they bury that interleaving inside one opaque ``DIFFUSION`` stage the scheduler never sees inside
(``max_num_running_reqs=1``). This card expresses the same shared-weight MoT, but makes both loops
**runtime-visible, step-scheduled, and batchable**: ``generate_text`` (ar_decode) and
``generate_image`` (diffusion_denoise) both bind the one resident ``transformer``, and every AR token
and denoise step is a WorkUnit the scheduler can interleave and price. That visibility is the
differentiation the §1 thesis names.
vllm-omni's ``bagel_single_stage``/``lance`` run one resident MoT instance doing both AR
``generate_text`` and diffusion ``generate_image`` on co-resident experts in a single request, but
bury the interleaving inside one opaque ``DIFFUSION`` stage the scheduler can't see into
(``max_num_running_reqs=1``). This card expresses the same shared-weight MoT with both loops
runtime-visible, step-scheduled, and batchable: ``generate_text`` (ar_decode) and ``generate_image``
(diffusion_denoise) both bind the one resident ``transformer``, and every AR token and denoise step
is a WorkUnit the scheduler can interleave and price.
"""
from __future__ import annotations
@@ -42,44 +41,64 @@ def build_bagel_card(model_id: str = "bagel-mot") -> ModelCard:
return ARDecodeLoop(loop_id="generate_text", transformer_id="transformer", cost=ar_cost, max_tokens=6)
def image_factory():
return WanDenoiseLoop(loop_id="generate_image", cfg=cfg, flow_shift=flow,
precision=precision, expert=expert, cost=dn_cost)
return WanDenoiseLoop(loop_id="generate_image",
cfg=cfg,
flow_shift=flow,
precision=precision,
expert=expert,
cost=dn_cost)
components = {
"tokenizer": ComponentSpec("tokenizer", kind="tokenizer", factory=lambda inst: ToyTokenizer(),
required_for={"reason", "t2i"}),
"transformer": ComponentSpec(
"transformer", kind="dit",
"tokenizer":
ComponentSpec("tokenizer",
kind="tokenizer",
factory=lambda inst: ToyTokenizer(),
required_for={"reason", "t2i"}),
"transformer":
ComponentSpec(
"transformer",
kind="dit",
load_id="vllm_omni.diffusion.models.bagel:BagelTransformer",
factory=lambda inst: ToyMoTDiT(seed=seed),
resident_for=["generate_text", "generate_image"], # one resident copy for BOTH loops
resident_for=["generate_text", "generate_image"], # one resident copy for BOTH loops
required_for={"reason", "t2i"}),
"vae": ComponentSpec("vae", kind="vae", factory=lambda inst: ToyVAE(), required_for={"t2i"}),
"vae":
ComponentSpec("vae", kind="vae", factory=lambda inst: ToyVAE(), required_for={"t2i"}),
}
loops = {
"generate_text": LoopSpec("generate_text", kind=LoopKind.AR_DECODE,
work_unit_kind=WorkUnitKind.AR_TOKEN, step_cost_model=ar_cost,
shared_weight_components=["transformer"], cache_policy=["paged_kv"],
loop_factory=text_factory),
"generate_image": LoopSpec("generate_image", kind=LoopKind.DIFFUSION_DENOISE,
work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=dn_cost,
shared_weight_components=["transformer"], cache_policy=["feature"],
loop_factory=image_factory),
"generate_text":
LoopSpec("generate_text",
kind=LoopKind.AR_DECODE,
work_unit_kind=WorkUnitKind.AR_TOKEN,
step_cost_model=ar_cost,
shared_weight_components=["transformer"],
cache_policy=["paged_kv"],
loop_factory=text_factory),
"generate_image":
LoopSpec("generate_image",
kind=LoopKind.DIFFUSION_DENOISE,
work_unit_kind=WorkUnitKind.DIFFUSION_STEP,
step_cost_model=dn_cost,
shared_weight_components=["transformer"],
cache_policy=["feature"],
loop_factory=image_factory),
}
card = ModelCard(
model_id=model_id, family="bagel", components=components, loops=loops,
capabilities=CapabilityMatrix.of(Capability.TEXT_TO_IMAGE, Capability.REASONING_TEXT,
Capability.VAE_DECODE),
recipe=RecipeSpec(method="base", assumes_loop="generate_image",
assumes_precision="float32", consistency_required=ConsistencyLevel.C1),
model_id=model_id,
family="bagel",
components=components,
loops=loops,
capabilities=CapabilityMatrix.of(Capability.TEXT_TO_IMAGE, Capability.REASONING_TEXT, Capability.VAE_DECODE),
recipe=RecipeSpec(method="base",
assumes_loop="generate_image",
assumes_precision="float32",
consistency_required=ConsistencyLevel.C1),
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True),
caches={
"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True),
"paged_kv": CacheContract("paged_kv", max_bytes=1 << 24, block_bytes=1 << 12,
reuse_across_requests=False),
"paged_kv": CacheContract("paged_kv", max_bytes=1 << 24, block_bytes=1 << 12, reuse_across_requests=False),
},
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()],
default_plan=ParallelPlan.single()),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
)
return card.validate()
v2/recipes/bagel/program.py
+31 -15
View File
@@ -1,9 +1,9 @@
"""BAGEL/lance program (design_v3 §16): one resident MoT instance, two runtime-visible loops.
"""BAGEL/lance program: 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.
scheduler interleaves.
"""
from __future__ import annotations
@@ -13,19 +13,35 @@ from v2.recipes.omni import emit_text_node, pack_cond_from_tokens, tokenize_node
def build_bagel_program() -> Program:
return Program(
program_id="bagel.mot", kind=ProgramKind.INLINE,
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",)),
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"},
output_artifacts={
"text": "text",
"image": "image"
},
).validate()
v2/recipes/common.py
+8 -7
View File
@@ -1,7 +1,7 @@
"""Shared component-node helpers (text encode with feature cache, etc.).
The content-hash feature cache (design_v3 §7.2) is the reuse mechanism behind the §10 claim that
a K-sample RL group encodes its shared prompt once instead of K times.
The content-hash feature cache lets a K-sample RL group encode its shared prompt once instead of
K times.
"""
from __future__ import annotations
@@ -14,14 +14,15 @@ def cached_text_encode(instance: Any, text: str) -> Any:
enc = instance.component("text_encoder")
cache = instance.caches
if cache is not None and cache.has("feature"):
# §7.1: every output-semantic field is in the key — adapter stack + precision partition it,
# and the text-encoder's OWN version (not the instance's) so a transformer weight sync
# does not invalidate the (frozen) text encoder's embeddings.
key = CacheKey(model_id=instance.card.model_id, component_id="text_encoder",
# Every output-semantic field is in the key — adapter stack + precision partition it, and the
# text-encoder's OWN version (not the instance's) so a transformer weight sync does not
# invalidate the (frozen) text encoder's embeddings.
key = CacheKey(model_id=instance.card.model_id,
component_id="text_encoder",
weights_version=instance.version_of("text_encoder"),
adapter_versions=CacheKey.adapters(instance.adapter_versions),
precision=instance.card.precision.dtype_for("text_encoder"),
input_hashes=(("text", content_hash(text)),))
input_hashes=(("text", content_hash(text)), ))
hit = cache.pool("feature").get(key)
if hit is not None:
return hit
v2/recipes/cosmos2/__init__.py
+4 -5
View File
@@ -1,9 +1,8 @@
"""Cosmos-Predict2 (Video2World) — EDM-Karras denoiser ported into the v2 substrate.
"""Cosmos-Predict2 (Video2World) EDM-Karras denoiser recipe.
Self-contained recipe package (the bucket-C pattern): card declares its torch adapters via
``ComponentSpec.adapter`` (``CosmosDiT``/``CosmosT5Encoder`` in ``v2/platform/backends/torch_cosmos.py``)
and a new ``CosmosDenoiseLoop`` (EDM preconditioning folded into flow-match Euler), reusing the Wan VAE
adapter + T5 + ``stamp_wan21_checkpoints``. Registered in ``v2/registry.py``.
Self-contained recipe: the card declares torch adapters (``CosmosDiT``/``CosmosT5Encoder`` in
``v2/platform/backends/torch_cosmos.py``) plus ``CosmosDenoiseLoop`` (EDM preconditioning folded into
flow-match Euler), reusing the Wan VAE adapter + T5 + ``stamp_wan21_checkpoints``.
"""
from __future__ import annotations
v2/recipes/cosmos2/card.py
+7 -8
View File
@@ -1,13 +1,12 @@
"""Cosmos-Predict2-2B-Video2World ModelCard — text→video (the registered preset; video2world
"""Cosmos-Predict2-2B-Video2World ModelCard — text->video (the registered preset; video2world
conditioning is a later capability the loop already threads).
Architecture deltas vs Wan (all GPU-path, declared on the card so the recipe is self-contained):
* DiT ``fastvideo.models.dits.cosmos:CosmosTransformer3DModel`` — an **EDM denoiser** (NOT flow-match);
the adapter (``ComponentSpec.adapter`` -> ``CosmosDiT``) returns the raw network output and
``CosmosDenoiseLoop`` does the EDM ``c_in/c_skip/c_out`` → x0 reconstruction + x0-space CFG + the
Karras (ρ=7, σ_max=80→σ_min=0.002) schedule.
* VAE ``fastvideo.models.vaes.wanvae:AutoencoderKLWan`` — Wan-style (z=16, 8×/4×); reuses the v2
``WanVAE`` adapter unchanged (sigma_data=1.0 makes the Cosmos *σ_data factor a no-op).
Architecture deltas vs Wan (declared on the card so the recipe is self-contained):
* DiT ``fastvideo.models.dits.cosmos:CosmosTransformer3DModel`` — an EDM denoiser (not flow-match); the
``CosmosDiT`` adapter returns the raw network output and ``CosmosDenoiseLoop`` does the EDM
``c_in/c_skip/c_out`` -> x0 reconstruction + x0-space CFG + the Karras (rho=7, sigma 80->0.002) schedule.
* VAE ``fastvideo.models.vaes.wanvae:AutoencoderKLWan`` — Wan-style (z=16, 8x/4x); reuses the v2
``WanVAE`` adapter unchanged (sigma_data=1.0 makes the Cosmos sigma_data factor a no-op).
* Text ``fastvideo.models.encoders.t5:T5EncoderModel`` (T5-Large, 1024-dim) via ``CosmosT5Encoder``
(raw last_hidden_state, no Wan zero-pad).
``stamp_wan21_checkpoints`` applies (diffusers transformer/vae/text_encoder subfolder layout).
v2/recipes/cosmos2/loop.py
+7 -7
View File
@@ -1,16 +1,16 @@
"""CosmosDenoiseLoop — EDM-Karras preconditioning folded into a flow-match Euler integrator.
Cosmos-Predict2 is NOT a plain flow-match model despite the pipeline wrapping a
Cosmos-Predict2 is not a plain flow-match model despite the pipeline wrapping a
``FlowMatchEulerDiscreteScheduler``. The network is an EDM denoiser ``F_θ`` whose preconditioned output
reconstructs ``x0`` (``x0 = c_skip·x + c_out·F_θ(x·c_in)``, ``sigma_data=1``); CFG combines in **x0
space**; then ``x0`` is converted to a flow-match velocity ``(x - x0)/σ`` for the deterministic Euler
update ``x_next = x + (σ_next - σ)·v``. The σ schedule is Karras (ρ=7, σ_max=80 → σ_min=0.002), NOT the
flow-shift linspace, and the latent is initialized at ``randn·σ_max``. Faithful port of
reconstructs ``x0`` (``x0 = c_skip·x + c_out·F_θ(x·c_in)``, ``sigma_data=1``); CFG combines in x0 space;
then ``x0`` is converted to a flow-match velocity ``(x - x0)/σ`` for the Euler update
``x_next = x + (σ_next - σ)·v``. The σ schedule is Karras (ρ=7, σ_max=80 -> σ_min=0.002), not the
flow-shift linspace, and the latent starts at ``randn·σ_max``. Port of
``fastvideo/pipelines/stages/denoising.py:CosmosDenoisingStage`` (see that file for the exact math).
video2world (frame_replace) conditioning is threaded (``conditioning_latents``/``cond_indicator`` from
slots) but is ``None`` for the registered t2v preset → the gated injection branch is inert; the loop
degrades to pure t2v exactly as the fastvideo pipeline does when no image/video is given.
slots) but is ``None`` for the registered t2v preset -> the gated injection is inert; the loop degrades
to pure t2v exactly as the fastvideo pipeline does when no image/video is given.
"""
from __future__ import annotations
v2/recipes/cosmos2/program.py
+3 -4
View File
@@ -1,9 +1,8 @@
"""Cosmos-Predict2 t2v program: text_encode → diffusion_denoise (EDM) → vae_decode.
Same inline shape as the Wan t2v program — the EDM specifics live entirely in ``CosmosDenoiseLoop`` and
the ``CosmosDiT`` adapter, so the node graph is unchanged. (video2world conditioning would add a
VAE-encode node writing ``conditioning_latents``/``cond_indicator`` into slots; the loop already reads
them.)
Same inline shape as the Wan t2v program — the EDM specifics live in ``CosmosDenoiseLoop`` and the
``CosmosDiT`` adapter, so the node graph is unchanged. (video2world conditioning would add a VAE-encode
node writing ``conditioning_latents``/``cond_indicator`` into slots; the loop already reads them.)
"""
from __future__ import annotations
v2/recipes/cosmos25/__init__.py
+6 -8
View File
@@ -1,12 +1,10 @@
"""Cosmos-Predict2.5 (text->video) — a flow-match denoiser ported into the v2 substrate.
"""Cosmos-Predict2.5 (text->video) flow-match denoiser recipe.
Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via
``ComponentSpec.adapter`` (``Cosmos25DiT`` / ``Cosmos25WanVAE`` / ``Cosmos25Reason1Encoder`` in
``v2/platform/backends/torch_cosmos25.py``) and a new ``Cosmos25DenoiseLoop`` (Wan's flow-shift sigma
schedule + ``FLOW_MATCH_STEP`` solver + ``ClassicCFG``, but the PLAIN per-frame sigma timestep ``[B,T]``
Cosmos2.5 needs — NOT Wan's ``sigma*1000``), reusing ``stamp_wan21_checkpoints``. Covers both the 2B and
14B sizes (the card is size-agnostic — size resolves from the checkpoint config). Registered in
``v2/registry.py`` by the orchestrator from the returned metadata.
Self-contained recipe: the card declares torch adapters (``Cosmos25DiT`` / ``Cosmos25WanVAE`` /
``Cosmos25Reason1Encoder`` in ``v2/platform/backends/torch_cosmos25.py``) plus ``Cosmos25DenoiseLoop``,
which reuses Wan's flow-shift sigma schedule + ``FLOW_MATCH_STEP`` + ``ClassicCFG`` but feeds the DiT the
plain per-frame sigma timestep ``[B,T]`` Cosmos2.5 needs (not Wan's ``sigma*1000``). Reuses
``stamp_wan21_checkpoints``. Size-agnostic (2B/14B resolve from the checkpoint config).
"""
from __future__ import annotations
v2/recipes/cosmos25/loop.py
+31 -45
View File
@@ -1,32 +1,20 @@
"""Cosmos25DenoiseLoop — a flow-match Euler denoiser with Cosmos2.5's per-frame sigma timestep.
"""Cosmos25DenoiseLoop — flow-match Euler denoiser with Cosmos2.5's per-frame sigma timestep.
Cosmos-Predict2.5 IS a rectified-flow / flow-match model (unlike Cosmos-Predict2, which is an EDM
denoiser). The sigma schedule and the deterministic Euler solver are EXACTLY Wan's
(``FlowShiftPolicy.build_schedule`` + ``FLOW_MATCH_STEP``), and CFG is the standard ``ClassicCFG``
(``uncond + s*(cond-uncond)``). So this loop is almost identical to ``WanDenoiseLoop``; it forks it for
ONE load-bearing reason (blocker (1)):
Cosmos-Predict2.5 is a flow-match model (unlike Cosmos-Predict2's EDM denoiser). The sigma schedule and
Euler solver are Wan's (``FlowShiftPolicy.build_schedule`` + ``FLOW_MATCH_STEP``, shift=5.0) and CFG is
``ClassicCFG``, so this loop nearly matches ``WanDenoiseLoop``. The one fork: Wan feeds the DiT a scalar
``timestep = sigma*1000``, but Cosmos2.5 feeds a per-frame 2D ``timestep[B,T]`` of the plain sigma in
[0,1]. So we set ``st.timesteps = st.sigmas`` (no *1000) and the ``Cosmos25DiT`` adapter broadcasts the
plain sigma to ``[B, T]`` and builds the mandatory zero ``condition_mask`` / ones ``padding_mask`` +
``fps`` the forward requires (the model concats the masks itself — we feed the raw 16ch latent). See
``v2/platform/backends/torch_cosmos25.py:Cosmos25DiT`` and the source
``fastvideo/pipelines/stages/denoising.py:Cosmos25DenoisingStage``.
* Wan feeds the DiT a SCALAR ``timestep = sigma * num_train_timesteps`` (``sigma*1000``).
* Cosmos2.5 feeds a PER-FRAME 2D ``timestep[B,T]`` of the PLAIN sigma in [0,1] — the fastvideo stage
computes ``t * 0.001`` of the 0..1000 schedule, which is precisely the shifted sigma. So we set
``st.timesteps = st.sigmas`` (NO *1000) and the ``Cosmos25DiT`` adapter broadcasts that plain sigma
to ``[B, T]`` and builds the mandatory zero ``condition_mask`` / ones ``padding_mask`` + ``fps`` the
forward requires (the model concats the masks itself — we feed the RAW 16ch latent). See
``v2/platform/backends/torch_cosmos25.py:Cosmos25DiT`` and the faithful source
``fastvideo/pipelines/stages/denoising.py:Cosmos25DenoisingStage``.
Verification that the math matches the source: the fastvideo ``FlowMatchEulerDiscreteScheduler``
(shift=5.0) builds ``sigma' = shift*sigma/(1+(shift-1)*sigma)`` over a 1->0 base + a terminal 0, sets
``timesteps = sigma'*1000``, feeds the DiT ``timesteps*0.001 = sigma'``, and steps
``prev = sample + (sigma_next - sigma)*velocity`` — i.e. exactly ``build_flow_sigmas(num_steps,
shift=5.0)`` + ``FLOW_MATCH_STEP``.
video2world / image2world (frame_replace) conditioning (cond_indicator GT-x0 clamp + GT-velocity
injection, conditional_frame_timestep=0.1 on conditioned frames) is present in fastvideo
(``Cosmos25V2WDenoisingStage``) but is NOT in the registered HF presets — this loop scopes to t2v.
The hooks are read from ``ctx.slots`` (None for t2v -> inert), mirroring the Wan/Cosmos2 pattern, so a
later capability adds the gated injection without changing the t2v path. (BRINGUP: frame-replace needs a
request-API extension to carry the conditioning latents + indicator.)
video2world / image2world (frame_replace) conditioning exists in fastvideo
(``Cosmos25V2WDenoisingStage``) but not in the registered HF presets, so this loop is t2v-only. The hooks
are read from ``ctx.slots`` (None for t2v -> inert), mirroring Wan/Cosmos2, so a later capability can add
the gated injection without touching the t2v path. (BRINGUP: frame-replace needs a request-API extension
to carry the conditioning latents + indicator.)
"""
from __future__ import annotations
@@ -48,9 +36,8 @@ from v2.loop.contracts import (
from v2.platform import FLOW_MATCH_STEP
from v2.recipes.wan21.loop import latent_shape
# Real Cosmos2.5 Wan-style VAE (Cosmos25WanVAE) compression: z_dim=16, 4x temporal, 8x spatial — the
# same geometry as Wan2.1, so ``latent_shape`` (reused from the Wan recipe) yields
# (16, (F-1)//4 + 1, H//8, W//8) on the real backend.
# Cosmos2.5 Wan-style VAE (Cosmos25WanVAE): z_dim=16, 4x temporal, 8x spatial — same geometry as Wan2.1,
# so the reused ``latent_shape`` yields (16, (F-1)//4 + 1, H//8, W//8) on the real backend.
COSMOS25_LATENT_CHANNELS = 16
COSMOS25_TEMPORAL_RATIO = 4
COSMOS25_SPATIAL_RATIO = 8
@@ -84,8 +71,7 @@ class Cosmos25DenoiseLoop:
def init(self, req, model, ctx) -> LoopState:
seed = req.diffusion.seed if req.diffusion.seed is not None else 0
rng = np.random.default_rng(seed)
# The σ schedule + flow-shift are Wan's verbatim (shift=5.0 for Cosmos2.5): 1 -> 0 over
# num_steps+1 points with the flow-shift warp applied.
# Wan's flow-shift sigma schedule (shift=5.0): 1 -> 0 over num_steps+1 points, warped.
sig = self.flow_shift.build_schedule(req.diffusion.num_steps,
req.diffusion.height,
req.diffusion.width,
@@ -106,16 +92,16 @@ class Cosmos25DenoiseLoop:
seed=seed,
latents={"video": x},
sigmas=[float(s) for s in sig],
# BLOCKER (1): the model timestep is the PLAIN sigma (NOT sigma*1000). The
# Cosmos25DiT adapter broadcasts it to the per-frame [B, T] form the forward needs.
# Model timestep is the plain sigma (not sigma*1000); the Cosmos25DiT adapter
# broadcasts it to the per-frame [B, T] form the forward needs.
timesteps=[float(s) for s in sig])
st.cond["prompt_embeds"] = ctx.slots.get("text_embeds")
st.cond["negative_prompt_embeds"] = ctx.slots.get("neg_text_embeds")
st.scratch["guidance_scale"] = float(req.diffusion.guidance_scale)
st.scratch["stream_video"] = bool(req.outputs.stream.get("video"))
st.scratch["fps"] = int(getattr(req.diffusion, "fps", None) or COSMOS25_DEFAULT_FPS)
# video2world (frame_replace) conditioning: None for the t2v preset -> the gated injection a later
# capability would add is inert (pure t2v, matching the fastvideo pipeline with no image/video).
# video2world (frame_replace) conditioning: None for the t2v preset -> any later gated injection
# is inert (pure t2v, matching the fastvideo pipeline with no image/video).
st.scratch["conditioning_latents"] = ctx.slots.get("conditioning_latents")
st.scratch["cond_indicator"] = ctx.slots.get("cond_indicator")
st.plugin_state["cfg"] = {}
@@ -126,8 +112,8 @@ class Cosmos25DenoiseLoop:
if i >= len(st.sigmas) - 1:
return Done()
sigma_t, sigma_next = st.sigmas[i], st.sigmas[i + 1]
# The model timestep here is the PLAIN sigma (st.timesteps[i] == st.sigmas[i]); the adapter turns
# it into the per-frame [B, T] timestep + masks + fps internally.
# Model timestep is the plain sigma (st.timesteps[i] == st.sigmas[i]); the adapter turns it into
# the per-frame [B, T] timestep + masks + fps internally.
expert_id = self.expert.expert_for(StepContext(i, st.timesteps[i], sigma_t))
sctx = StepContext(step_idx=i, timestep=st.timesteps[i], sigma=sigma_t, active_expert_id=expert_id)
cfg_state = st.plugin_state["cfg"]
@@ -138,16 +124,16 @@ class Cosmos25DenoiseLoop:
cfg, precision = self.cfg, self.precision
def _velocity(model: Any, x_: Any, sigma_t_: float, pe_: Any, ne_: Any, scale_: float) -> np.ndarray:
# The conditioned forward + ClassicCFG combine. The dit-call is toy-compatible:
# Conditioned forward + ClassicCFG combine. The dit-call is toy-compatible:
# ``dit(latent, text_embed, sigma)`` — the Cosmos2.5-specific timestep/masks/fps are built
# INSIDE the Cosmos25DiT adapter (so the CPU ToyDiT, which ignores them, also runs).
# inside the Cosmos25DiT adapter (so the CPU ToyDiT, which ignores them, also runs).
dit = model.component(expert_id)
preds = {b: dit(x_, pe_ if b == "cond" else ne_, sigma_t_) for b in branches}
return precision.cast(cfg.combine(preds, scale_, sctx, cfg_state))
def run(model, override=None):
# The EAGER thunk: handles the override (cached-prediction) path. The deterministic Euler
# solver dispatches through the platform kernel table (numpy on CPU, the device kernel on GPU).
# Eager thunk: handles the override (cached-prediction) path. The Euler solver dispatches
# through the platform kernel table (numpy on CPU, the device kernel on GPU).
kernels = model.platform.kernels
if override is not None and "noise_pred" in override:
velocity = precision.cast(np.asarray(override["noise_pred"], dtype="float32"))
@@ -160,7 +146,7 @@ class Cosmos25DenoiseLoop:
})
def graph_fn(model, ws):
# The CAPTURABLE op-structure: reads EVERY per-step input from the static workspace and writes
# Capturable op-structure: reads every per-step input from the static workspace and writes
# into the static output buffer, so the captured graph replays correctly with rebound buffers.
st_t, st_n = float(ws["sigma_t"]), float(ws["sigma_next"])
velocity = _velocity(model, ws["x"], st_t, ws["pe"], ws["ne"], float(ws["scale"]))
@@ -190,8 +176,8 @@ class Cosmos25DenoiseLoop:
},
run=run,
label=f"cosmos25.denoise.{i}",
# The flow-match Euler step is a clean capturable op (no host RNG / data-dependent branch),
# like Wan t2v: the branch-set + expert + scheduler-precision discriminate the captured graph.
# Flow-match Euler step is a clean capturable op (no host RNG / data-dependent branch), like
# Wan t2v: branch-set + expert + scheduler-precision discriminate the captured graph.
capturable=True,
graph_key=(tuple(sorted(branches)), expert_id, precision.scheduler_step_in_fp32),
graph_fn=graph_fn,
v2/recipes/cosmos25/program.py
+4 -5
View File
@@ -1,11 +1,10 @@
"""Cosmos-Predict2.5 t2v program: text_encode -> diffusion_denoise (flow-match) -> vae_decode.
Same inline shape as the Wan / Cosmos2 t2v programs — the Cosmos2.5 specifics (per-frame plain-sigma
timestep, the mandatory masks + fps, the Reason1 100352-dim embedding) live entirely in
``Cosmos25DenoiseLoop`` and the ``torch_cosmos25`` adapters, so the node graph is unchanged. fps is read
by the loop from the request (``req.diffusion.fps``, default 24) into scratch and threaded to the DiT
adapter. (video2world/image2world conditioning would add a VAE-encode node writing
``conditioning_latents``/``cond_indicator`` into slots; the loop already reads them — BRINGUP.)
timestep, the mandatory masks + fps, the Reason1 100352-dim embedding) live in ``Cosmos25DenoiseLoop``
and the ``torch_cosmos25`` adapters, so the node graph is unchanged. The loop reads fps from the request
(``req.diffusion.fps``, default 24). (video2world/image2world conditioning would add a VAE-encode node
writing ``conditioning_latents``/``cond_indicator`` into slots; the loop already reads them — BRINGUP.)
"""
from __future__ import annotations
v2/recipes/cosmos3/card.py
+60 -38
View File
@@ -1,12 +1,11 @@
"""Cosmos3 omni ModelCard (design_v3 §4.1, §15b) — the forcing function for omni-native serving.
"""Cosmos3 omni ModelCard — one resident MoT transformer serving both und and generation.
ONE resident MoT ``transformer`` is bound to BOTH the ``ar_decode`` loop (the reasoner, und pathway)
and the ``diffusion_denoise`` loop (the joint generation pathway) — ``shared_weight_components`` on
both. A request runs the reasoner (prompt upsampling), packs its tokens into conditioning, then the
joint denoise — all on the same resident weights, with every AR token and denoise step a
runtime-visible WorkUnit (the differentiation vs vllm-omni's opaque DIFFUSION stage). The diffusion
loop is literally ``WanDenoiseLoop`` — the same step body the engine serves for Wan — bound to the
MoT module. ``sound_vae`` is declared ``optional_for`` non-t2vs tasks (the lazy-component fix, P8).
ONE resident MoT ``transformer`` is bound to BOTH the ``ar_decode`` loop (the reasoner) and the
``diffusion_denoise`` loop (joint generation) via ``shared_weight_components``. A request runs the
reasoner (prompt upsampling), packs its tokens into conditioning, then the joint denoise — all on the
same resident weights, with every AR token and denoise step a runtime-visible WorkUnit. The diffusion
loop is ``WanDenoiseLoop`` (the same step body the engine serves for Wan) bound to the MoT module.
``sound_vae`` is ``optional_for`` non-t2vs tasks so it loads lazily.
"""
from __future__ import annotations
@@ -41,49 +40,72 @@ def build_cosmos3_card(model_id: str = "cosmos3-vfm") -> ModelCard:
return ARDecodeLoop(loop_id="ar_decode", transformer_id="transformer", cost=ar_cost, max_tokens=6)
def dn_factory():
return WanDenoiseLoop(loop_id="diffusion_denoise", cfg=cfg, flow_shift=flow,
precision=precision, expert=expert, cost=dn_cost)
return WanDenoiseLoop(loop_id="diffusion_denoise",
cfg=cfg,
flow_shift=flow,
precision=precision,
expert=expert,
cost=dn_cost)
components = {
"tokenizer": ComponentSpec("tokenizer", kind="tokenizer",
factory=lambda inst: ToyTokenizer(),
required_for={"reason", "t2v", "t2vs"}),
"transformer": ComponentSpec(
"transformer", kind="dit",
"tokenizer":
ComponentSpec("tokenizer",
kind="tokenizer",
factory=lambda inst: ToyTokenizer(),
required_for={"reason", "t2v", "t2vs"}),
"transformer":
ComponentSpec(
"transformer",
kind="dit",
load_id="fastvideo.models.dits.cosmos3:Cosmos3VFMTransformer",
factory=lambda inst: ToyMoTDiT(seed=seed),
resident_for=["ar_decode", "diffusion_denoise"], # never unload mid-request
resident_for=["ar_decode", "diffusion_denoise"], # never unload mid-request
required_for={"reason", "t2v", "t2vs"}),
"vae": ComponentSpec("vae", kind="vae", factory=lambda inst: ToyVAE(),
required_for={"t2v", "t2vs"}),
"sound_vae": ComponentSpec("sound_vae", kind="audio_vae",
load_id="fastvideo.models.audio:Cosmos3SoundVAE",
optional_for={"reason", "t2v"}, required_for={"t2vs"}),
"vae":
ComponentSpec("vae", kind="vae", factory=lambda inst: ToyVAE(), required_for={"t2v", "t2vs"}),
"sound_vae":
ComponentSpec("sound_vae",
kind="audio_vae",
load_id="fastvideo.models.audio:Cosmos3SoundVAE",
optional_for={"reason", "t2v"},
required_for={"t2vs"}),
}
loops = {
"ar_decode": LoopSpec("ar_decode", kind=LoopKind.AR_DECODE, work_unit_kind=WorkUnitKind.AR_TOKEN,
step_cost_model=ar_cost, shared_weight_components=["transformer"],
cache_policy=["paged_kv"], loop_factory=ar_factory),
"diffusion_denoise": LoopSpec("diffusion_denoise", kind=LoopKind.DIFFUSION_DENOISE,
work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=dn_cost,
shared_weight_components=["transformer"], cache_policy=["feature"],
loop_factory=dn_factory),
"ar_decode":
LoopSpec("ar_decode",
kind=LoopKind.AR_DECODE,
work_unit_kind=WorkUnitKind.AR_TOKEN,
step_cost_model=ar_cost,
shared_weight_components=["transformer"],
cache_policy=["paged_kv"],
loop_factory=ar_factory),
"diffusion_denoise":
LoopSpec("diffusion_denoise",
kind=LoopKind.DIFFUSION_DENOISE,
work_unit_kind=WorkUnitKind.DIFFUSION_STEP,
step_cost_model=dn_cost,
shared_weight_components=["transformer"],
cache_policy=["feature"],
loop_factory=dn_factory),
}
card = ModelCard(
model_id=model_id, family="cosmos3", components=components, loops=loops,
capabilities=CapabilityMatrix.of(
Capability.TEXT_TO_VIDEO, Capability.REASONING_TEXT, Capability.TEXT_TO_VIDEO_SOUND,
Capability.VAE_DECODE, Capability.POLICY_ROLLOUT),
recipe=RecipeSpec(method="base", assumes_loop="diffusion_denoise",
assumes_precision="float32", consistency_required=ConsistencyLevel.C1),
model_id=model_id,
family="cosmos3",
components=components,
loops=loops,
capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.REASONING_TEXT,
Capability.TEXT_TO_VIDEO_SOUND, Capability.VAE_DECODE,
Capability.POLICY_ROLLOUT),
recipe=RecipeSpec(method="base",
assumes_loop="diffusion_denoise",
assumes_precision="float32",
consistency_required=ConsistencyLevel.C1),
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True),
caches={
"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True),
"paged_kv": CacheContract("paged_kv", max_bytes=1 << 24, block_bytes=1 << 12,
reuse_across_requests=False),
"paged_kv": CacheContract("paged_kv", max_bytes=1 << 24, block_bytes=1 << 12, reuse_across_requests=False),
},
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()],
default_plan=ParallelPlan.single()),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
)
return card.validate()
v2/recipes/cosmos3/program.py
+33 -17
View File
@@ -1,10 +1,9 @@
"""Cosmos3 omni program (design_v3 §15b):
"""Cosmos3 omni program:
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).
The reasoner and the denoiser bind the SAME resident ``transformer``; the scheduler multiplexes the
mode. Splitting them into separate engines would double the weights and sever the shared state.
"""
from __future__ import annotations
@@ -14,19 +13,36 @@ from v2.recipes.omni import emit_text_node, pack_cond_from_tokens, tokenize_node
def build_cosmos3_program() -> Program:
return Program(
program_id="cosmos3.omni", kind=ProgramKind.INLINE,
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",)),
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"},
output_artifacts={
"video": "video",
"text": "reason_text",
"latents": "denoise_out"
},
).validate()
v2/recipes/fastwan/card.py
+27 -28
View File
@@ -1,26 +1,25 @@
"""FastWan (DMD few-step Wan) ModelCard — self-contained recipe package (bucket-C pattern).
"""FastWan (DMD few-step Wan) ModelCard — self-contained recipe package.
FastWan = Distribution-Matching-Distillation few-step Wan. The architecture is the same
``WanTransformer3DModel`` / ``AutoencoderKLWan`` / UMT5 as base Wan — so we REUSE the Wan torch adapters
(no ``adapter=`` needed; the torch backend's ``_make_dit`` already builds ``WanTransformer3DModel`` from
``load_id``) and the Wan VAE / T5 adapters. The ONLY new piece is the DMD few-step schedule, captured by
the in-package ``FastWanDMDLoop`` (predict-x0-then-renoise over a fixed ``dmd_denoising_steps`` list).
FastWan = Distribution-Matching-Distillation few-step Wan. Same architecture as base Wan
(``WanTransformer3DModel`` / ``AutoencoderKLWan`` / UMT5), so the Wan torch adapters are reused via
``load_id`` (no ``adapter=`` needed; the torch backend's ``_make_dit`` already builds the transformer)
along with the Wan VAE / T5 adapters. The only new piece is the DMD few-step schedule in
``FastWanDMDLoop`` (predict-x0-then-renoise over a fixed ``dmd_denoising_steps`` list).
Loadable target (this card's primary, GPU-verifiable id):
* ``FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers`` — FULL ATTENTION (no VSA kernel needed) +
Loadable target (this card's primary id):
* ``FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers`` — full attention (no VSA kernel needed),
TI2V-5B geometry (z_dim=48 VAE, 16x spatial / 4x temporal), DMD steps [1000, 757, 522] (3-step),
flow_shift 5.0 / DMD-scheduler shift 8.0. From ``FastWan2_2_TI2V_5B_Config`` +
``FAST_WAN_2_2_TI2V_5B`` preset (height 704, width 1280, 121 frames, fps 24, guidance 5.0).
flow_shift 5.0 / DMD-scheduler shift 8.0. From ``FAST_WAN_2_2_TI2V_5B`` preset (704x1280, 121
frames, fps 24, guidance 5.0).
BRINGUP variants (registered to a card here, but NOT loadable in this environment — clearly flagged):
* ``FastVideo/FastWan2.2-TI2V-5B-Diffusers`` (the non-FullAttn TI2V-5B) — trained for the **VSA**
(Video Sparse Attention) kernel, which is NOT built here (no nvcc). Same card/geometry; load is
BRINGUP until the VSA kernel is built + selected.
BRINGUP variants (registered here but NOT loadable in this environment):
* ``FastVideo/FastWan2.2-TI2V-5B-Diffusers`` (non-FullAttn TI2V-5B) — trained for the VSA (Video
Sparse Attention) kernel, not built here (no nvcc). Same card/geometry; load is BRINGUP until VSA
is built + selected.
* ``FastVideo/FastWan2.1-T2V-1.3B-Diffusers`` and ``FastVideo/FastWan2.1-T2V-14B-480P-Diffusers``
(``FastWan2_1_T2V_480P_Config``: Wan2.1 16/8/4 geometry, DMD steps [1000, 757, 522], flow_shift 8.0,
448x832x61, fps 16, guidance 3.0) — also VSA-trained, AND the DMD checkpoints carry **gated-attention**
params (``to_gate_compress``) the generic Wan loader cannot map -> they require a v2-side non-strict
loader option. Both are BRINGUP (VSA kernel build + non-strict load).
(Wan2.1 16/8/4 geometry, DMD steps [1000, 757, 522], flow_shift 8.0, 448x832x61, fps 16, guidance
3.0) — also VSA-trained, AND the DMD checkpoints carry gated-attention params (``to_gate_compress``)
the generic Wan loader cannot map, so they need a v2-side non-strict loader option. Both BRINGUP.
``stamp_wan21_checkpoints`` applies (diffusers transformer/vae/text_encoder subfolder layout).
"""
@@ -65,14 +64,14 @@ def _build_fastwan_card(model_id: str,
temporal_ratio: int,
sampling_defaults: SamplingDefaults,
checkpoint_root: str | None = None) -> ModelCard:
"""Shared FastWan card builder. The DMD few-step loop is the only delta vs base Wan; the
WanTransformer3DModel / AutoencoderKLWan / UMT5 torch adapters are reused via ``load_id`` (no
``adapter=``). ``EmbeddedGuidance`` (single-branch) matches DMD's distilled, no-CFG forward.
"""Shared FastWan card builder. The DMD few-step loop is the only delta vs base Wan; the Wan torch
adapters are reused via ``load_id`` (no ``adapter=``). ``EmbeddedGuidance`` (single-branch) matches
DMD's distilled, no-CFG forward.
No ``FlowShiftPolicy`` is needed: DMD does NOT integrate a continuous flow-shift sigma schedule —
the schedule is the fixed ``dmd_denoising_steps`` list, and the DMD re-noise scheduler's shift is
fixed at 8.0 inside ``FastWanDMDLoop`` (the pipeline ``flow_shift`` 5.0/8.0 is doc-level metadata
on the per-variant builders, not a loop input). The negative prompt rides in ``sampling_defaults``."""
No ``FlowShiftPolicy``: DMD does not integrate a continuous flow-shift sigma schedule — the schedule
is the fixed ``dmd_denoising_steps`` list, and the re-noise scheduler's shift is fixed at 8.0 inside
``FastWanDMDLoop`` (the pipeline ``flow_shift`` 5.0/8.0 is doc metadata on the per-variant builders,
not a loop input). The negative prompt rides in ``sampling_defaults``."""
seed = _seed_from(model_id)
cost = CostModel(kind=WorkUnitKind.DIFFUSION_STEP, base_seconds=1e-4, per_unit_seconds=1e-7)
cfg = EmbeddedGuidance() # DMD is distilled single-branch (guidance embedded, NOT CFG)
@@ -149,15 +148,15 @@ def _build_fastwan_card(model_id: str,
def build_fastwan_ti2v_5b_card(model_id: str = "fastwan2.2-ti2v-5b-fullattn",
*,
checkpoint_root: str | None = None) -> ModelCard:
"""FastWan2.2-TI2V-5B-FullAttn (the loadable target). FullAttn = no VSA kernel needed.
"""FastWan2.2-TI2V-5B-FullAttn (the loadable target; FullAttn = no VSA kernel needed).
Geometry: z_dim=48 VAE, 16x spatial / 4x temporal (Wan2.2-TI2V-5B). DMD steps [1000, 757, 522]
(3-step), pipeline flow_shift 5.0 (the DMD re-noise scheduler internally uses shift 8.0). Defaults
from ``FAST_WAN_2_2_TI2V_5B``: 704x1280, 121 frames, fps 24, guidance 5.0. ``num_steps=3`` is the
DMD schedule length (the preset's misleading 50 is ignored by the DMD stage).
DMD schedule length (the preset's 50 is ignored by the DMD stage).
Serves both ``FastWan2.2-TI2V-5B-FullAttn-Diffusers`` (loadable) and ``FastWan2.2-TI2V-5B-Diffusers``
(the non-FullAttn variant -> VSA-trained -> load is BRINGUP; same card/geometry)."""
(non-FullAttn variant, VSA-trained, so load is BRINGUP; same card/geometry)."""
return _build_fastwan_card(model_id,
latent_channels=FASTWAN_TI2V_LATENT_CHANNELS,
spatial_ratio=FASTWAN_TI2V_SPATIAL_RATIO,
v2/recipes/fastwan/loop.py
+2 -2
View File
@@ -1,9 +1,9 @@
"""FastWanDMDLoop — Distribution-Matching-Distillation few-step denoise for FastWan (design_v3 §5).
"""FastWanDMDLoop — Distribution-Matching-Distillation few-step denoise for FastWan.
FastWan is a DMD-distilled Wan: instead of N flow-match Euler steps it runs a handful (3) of
*predict-x0-then-renoise* steps over a fixed list of discrete model timesteps. This is NOT a plain
flow-match Euler integrator (``WanDenoiseLoop``), so it ships its own ``next``/``advance`` and uses the
shared policies as a library (the design_v3 §5.1 escape hatch for braided math).
shared policies as a library.
Faithful port of ``fastvideo/pipelines/stages/denoising.py:DmdDenoisingStage`` (+ ``pred_noise_to_pred_video``
in ``fastvideo/models/utils.py``). The exact math, per discrete model timestep ``t_i`` (e.g. [1000, 757, 522]):
v2/recipes/fastwan/program.py
+3 -5
View File
@@ -1,10 +1,8 @@
"""FastWan (DMD) t2v program: text_encode -> diffusion_denoise (DMD few-step) -> vae_decode.
Same inline node graph as the Wan t2v program — the DMD specifics live entirely in
``FastWanDMDLoop`` (the in-package loop), so the program is unchanged. (i2v conditioning for the TI2V
variant would add image-encode / first-frame-VAE nodes writing ``i2v_img_embeds`` / ``i2v_cond`` into
slots, exactly like ``v2/recipes/wan21/i2v.py``; the loop already reads them. The registered preset is
t2v, so the base program covers it.)
Same inline node graph as the Wan t2v program — the DMD specifics live entirely in ``FastWanDMDLoop``.
The TI2V variant would add image-encode / first-frame-VAE nodes writing ``i2v_img_embeds`` / ``i2v_cond``
(see ``v2/recipes/wan21/i2v.py``; the loop already reads them), but the registered preset is t2v.
"""
from __future__ import annotations
v2/recipes/flux2/sampler.py
+3 -4
View File
@@ -1,10 +1,9 @@
"""FLUX.2 flow-match sigma schedule — the BFL empirical-mu grid (recipe-local; NOT in v2/loop/sampler.py).
FLUX.2 does NOT use ``FlowShiftPolicy.build_schedule`` (the Wan/LTX resolution-bucket shift). It builds a
resolution-dependent ``mu`` from the *packed* image token count and feeds it, together with a custom
resolution-dependent ``mu`` from the *packed* image token count and feeds it, with a custom
``sigmas = linspace(1, 1/N, N)`` grid, to diffusers' ``FlowMatchEulerDiscreteScheduler.set_timesteps``
under ``use_dynamic_shifting=True``. This module reproduces that scheduler's numpy result exactly so the
CPU-toy loop is bit-faithful to the GPU pipeline.
under ``use_dynamic_shifting=True``. This module reproduces that scheduler's numpy result exactly.
Faithful to:
* ``fastvideo/pipelines/basic/flux_2/flux_2_timestep_preparation.py:compute_empirical_mu`` (the BFL
@@ -49,7 +48,7 @@ def flux2_sigmas(num_steps: int, image_seq_len: int) -> np.ndarray:
1. base grid ``sigmas = linspace(1.0, 1/N, N)`` (the custom grid the pipeline passes);
2. resolution-dependent ``mu`` shift ``time_shift(mu, 1.0, sigmas)`` (use_dynamic_shifting=True);
3. append the terminal ``0.0`` (FlowMatchEuler's ``invert_sigmas=False`` tail).
The loop integrates these pairwise via the deterministic flow-match Euler step (FLOW_MATCH_STEP).
The loop integrates these pairwise via the flow-match Euler step (FLOW_MATCH_STEP).
"""
n = max(1, int(num_steps))
base = np.linspace(1.0, 1.0 / n, n, dtype=np.float64)
v2/recipes/gen3c/loop.py
+3 -3
View File
@@ -28,12 +28,12 @@ step math (see ``build_edm_euler_sigmas`` in ``sampler.py``):
Camera / MoGe-depth / 3D-cache conditioning fills ``condition_video_pose`` / ``conditioning_latents`` /
``cond_indicator`` from program slots; ALL are ``None`` for the registered degenerate **t2v** path → the
frame-replace branches are inert and the DiT's internal pose concat is zeroed, exactly as the fastvideo
frame-replace branches are inert and the DiT's internal pose concat is zeroed, matching how the fastvideo
pipeline degrades when no image/camera is given. Faithful port of ``Gen3CDenoisingStage.forward`` +
``Gen3CLatentPreparationStage`` (fastvideo/pipelines/stages/gen3c_stages.py).
The DiT call is ``dit(model_input, text_embed, sigma, cond=bundle)`` — compatible with ``ToyDiT`` (which
accepts and ignores ``cond``), so the loop CPU-verifies on the toy backend.
The DiT call is ``dit(model_input, text_embed, sigma, cond=bundle)`` — ``ToyDiT`` accepts and ignores
``cond``, so the loop CPU-verifies on the toy backend.
"""
from __future__ import annotations
v2/recipes/hunyuan_video/__init__.py
+7 -8
View File
@@ -1,13 +1,12 @@
"""HunyuanVideo — rectified flow-match video DiT ported into the v2 substrate.
Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via
``ComponentSpec.adapter`` (``HunyuanVideoDiT`` / ``HunyuanVideoVAE`` / ``HunyuanVideoLlamaEncoder`` /
``HunyuanVideoCLIPEncoder`` in ``v2/platform/backends/torch_hunyuan_video.py``) and reuses the canonical
``WanDenoiseLoop`` flow-match math via a thin ``HunyuanDenoiseLoop`` (which only threads the CLIP-pooled
vector through Wan's ``context=`` channel). The new work vs Wan is the dual text-encoder marshalling
(LLaMA per-token sequence + CLIP pooled global vector packed into the DiT's 2-element
``encoder_hidden_states``) and the Hunyuan VAE scalar-``scaling_factor`` normalization. Reuses
``stamp_wan21_checkpoints``. Registered in ``v2/registry.py`` by the orchestrator.
Self-contained recipe package: the card declares its torch adapters via ``ComponentSpec.adapter``
(``HunyuanVideoDiT`` / ``HunyuanVideoVAE`` / ``HunyuanVideoLlamaEncoder`` / ``HunyuanVideoCLIPEncoder`` in
``v2/platform/backends/torch_hunyuan_video.py``) and reuses the canonical ``WanDenoiseLoop`` flow-match math
via a thin ``HunyuanDenoiseLoop`` (which only threads the CLIP-pooled vector through Wan's ``context=``
channel). The new work vs Wan: dual text-encoder marshalling (LLaMA per-token sequence + CLIP pooled global
vector packed into the DiT's 2-element ``encoder_hidden_states``) and the Hunyuan VAE scalar-``scaling_factor``
normalization. Reuses ``stamp_wan21_checkpoints``. Registered in ``v2/registry.py``.
"""
from __future__ import annotations
v2/recipes/hunyuan_video/card.py
+4 -4
View File
@@ -13,11 +13,11 @@ Architecture deltas vs Wan (all declared on the card so the recipe is self-conta
``stamp_wan21_checkpoints`` applies (diffusers transformer/vae/text_encoder subfolder layout); HunyuanVideo
additionally ships ``text_encoder_2/`` (CLIP) + ``tokenizer_2/`` — the stamp helper points each declared
component (incl. ``text_encoder_2``) at its subfolder (see the BRINGUP note in ``card.py`` and the CLIP
adapter for the ``tokenizer_2`` wiring caveat).
component (incl. ``text_encoder_2``) at its subfolder (see the BRINGUP note below and the CLIP adapter for
the ``tokenizer_2`` wiring caveat).
Defaults (fastvideo HUNYUAN_T2V preset): 50 steps, guidance_scale 1.0 (CFG effectively off),
720x1280, 125 frames @ 24fps, flow_shift 7.0. A FastHunyuan variant ships shift 17, 6 steps.
Defaults (fastvideo HUNYUAN_T2V preset): 50 steps, guidance_scale 1.0 (CFG effectively off), 720x1280,
125 frames @ 24fps, flow_shift 7.0. A FastHunyuan variant ships shift 17, 6 steps.
"""
from __future__ import annotations
v2/recipes/hunyuan_video/loop.py
+20 -23
View File
@@ -1,27 +1,25 @@
"""HunyuanDenoiseLoop — a thin subclass of the canonical ``WanDenoiseLoop`` (flow-match Euler).
"""HunyuanDenoiseLoop — a thin subclass of ``WanDenoiseLoop`` (flow-match Euler).
HunyuanVideo is a *rectified flow-match* model: the DiT predicts a velocity in unpatchified
HunyuanVideo is a rectified flow-match model: the DiT predicts a velocity in unpatchified
``[B, 16, T, H, W]`` latent space, the model timestep is the scalar ``sigma * 1000`` (a per-batch
``LongTensor`` consumed by the internal ``time_in`` ``TimestepEmbedder``), and the deterministic Euler
update is ``x_next = x + (sigma_next - sigma)*v`` — exactly the math ``WanDenoiseLoop`` already runs. The
latent geometry (16 channels, 8x spatial, 4x temporal) is the Wan default too, so we reuse the loop
verbatim and only carry a Hunyuan-specific docstring + the dual-text-encoder marshalling delta below.
``LongTensor``), and the Euler update is ``x_next = x + (sigma_next - sigma)*v`` — exactly the math
``WanDenoiseLoop`` already runs. The latent geometry (16 channels, 8x spatial, 4x temporal) is the Wan
default too, so we reuse the loop verbatim and only add the dual-text-encoder marshalling delta below.
The ONE delta vs Wan is conditioning shape: the Hunyuan DiT consumes TWO text signals —
The one delta vs Wan is conditioning shape: the Hunyuan DiT consumes TWO text signals —
``encoder_hidden_states = [llama_hidden[B, L, 4096], clip_pooled[B, 768]]`` (a per-token LLaMA sequence
plus a single CLIP-pooled global vector). To keep the loop's ``dit(latent, text_embed, sigma)`` call
identical on the CPU toy and the GPU backend (the cosmos2 discipline: the adapter assembles
arch-specifics INTERNALLY), the program writes the LLaMA sequence into the ``text_embeds`` slot (the
primary ``text_embed`` the loop already threads) and the CLIP-pooled vector into ``text_pooled``. This
loop reuses Wan's existing ``context=`` channel (Wan uses it for the i2v CLIP-vision embeds) to carry
the CLIP-pooled vector to the DiT adapter — so ``WanDiT``-style ``dit(x, text_embed, sigma, context=...)``
is unchanged, ``ToyDiT`` (which means over both ``text_embed`` and ``context``) runs on CPU, and the real
identical on the CPU toy and GPU backends (the adapter assembles arch-specifics internally), the program
writes the LLaMA sequence into the ``text_embeds`` slot and the CLIP-pooled vector into ``text_pooled``.
This loop reuses Wan's existing ``context=`` channel (Wan uses it for i2v CLIP-vision embeds) to carry the
pooled vector to the adapter — so the ``dit(x, text_embed, sigma, context=...)`` call is unchanged,
``ToyDiT`` (which means over both ``text_embed`` and ``context``) runs on CPU, and the real
``HunyuanVideoDiT`` adapter packs ``[text_embed, context]`` into the 2-element ``encoder_hidden_states``.
guidance_scale defaults to 1.0 for base HunyuanVideo → ``ClassicCFG.combine`` collapses to the single
``cond`` branch (``uncond + 1*(cond-uncond) == cond``), i.e. CFG is effectively off, matching the
fastvideo pipeline. (``embedded_cfg_scale``/``guidance`` is a no-op on the base model whose
``guidance_in`` is ``None``; see ``torch_hunyuan_video.HunyuanVideoDiT`` for that BRINGUP note.)
``cond`` branch (CFG effectively off), matching the fastvideo pipeline. (``embedded_cfg_scale``/``guidance``
is a no-op on the base model whose ``guidance_in`` is ``None``; see the BRINGUP note in
``torch_hunyuan_video.HunyuanVideoDiT``.)
Faithful to ``fastvideo/pipelines/basic/hunyuan/hunyuan_pipeline.py`` (flow-match DenoisingStage) and
``fastvideo/models/dits/hunyuanvideo.py:HunyuanVideoTransformer3DModel.forward`` (the list-split of
@@ -41,17 +39,16 @@ HUNYUAN_SPATIAL_RATIO = 8
class HunyuanDenoiseLoop(WanDenoiseLoop):
"""Wan flow-match loop with HunyuanVideo's dual-text-encoder marshalling.
Reuses ``WanDenoiseLoop.next``/``advance``/``finalize`` unchanged (flow-match Euler, velocity output,
timestep ``sigma*1000``, scalar per-batch timestep). Overrides only ``init`` to thread the CLIP-pooled
global vector (slot ``text_pooled``) through Wan's existing ``context=`` channel so the DiT adapter can
reassemble the 2-element ``encoder_hidden_states`` the Hunyuan forward expects.
Reuses ``WanDenoiseLoop.next``/``advance``/``finalize`` unchanged. Overrides only ``init`` to thread the
CLIP-pooled global vector (slot ``text_pooled``) through Wan's existing ``context=`` channel so the DiT
adapter can reassemble the 2-element ``encoder_hidden_states`` the Hunyuan forward expects.
"""
def init(self, req, model, ctx) -> LoopState:
st = super().init(req, model, ctx)
# Wan's ``context`` channel (``i2v_img_embeds`` -> dit ``context=``) carries the CLIP-pooled vector
# for Hunyuan. None for any caller that did not run the dual-encode node -> the DiT adapter falls
# back to a single-tensor encoder_hidden_states (still valid: the Hunyuan forward also accepts a
# bare tensor whose token 0 is the pooled vector — but the dual-slot path is the faithful one).
# for Hunyuan. None when the dual-encode node did not run -> the DiT adapter falls back to a
# single-tensor encoder_hidden_states (still valid: the Hunyuan forward also accepts a bare tensor
# whose token 0 is the pooled vector — but the dual-slot path is the faithful one).
st.scratch["i2v_img_embeds"] = ctx.slots.get("text_pooled")
return st
v2/recipes/hunyuan_video/program.py
+7 -7
View File
@@ -1,15 +1,15 @@
"""HunyuanVideo t2v program: hunyuan_text_encode -> diffusion_denoise (flow-match) -> vae_decode.
Same inline shape as the Wan/cosmos2 t2v programs, with ONE delta: HunyuanVideo has TWO text encoders
(a LLaMA per-token sequence + a CLIP-pooled global vector), so the text-encode node writes BOTH a
Same inline shape as the Wan/cosmos2 t2v programs, with one delta: HunyuanVideo has TWO text encoders
(a LLaMA per-token sequence + a CLIP-pooled global vector), so the text-encode node writes both a
``text_embeds`` slot (the LLaMA sequence, the loop's primary ``text_embed``) and a ``text_pooled`` slot
(the CLIP vector, threaded to the DiT via the loop's ``context=`` channel). The dual-encode node uses the
content-hash feature cache (like ``recipes.common.text_encode_node_fn``) so a repeated prompt encodes once.
content-hash feature cache so a repeated prompt encodes once.
On the CPU toy backend both components are ``ToyTextEncoder`` (single arrays), so ``ToyDiT`` — which means
over both ``text_embed`` and ``context`` — runs end-to-end; on the GPU backend the real LLaMA/CLIP adapters
produce the sequence + pooled pair the ``HunyuanVideoDiT`` adapter reassembles into the 2-element
``encoder_hidden_states`` the Hunyuan forward expects.
``encoder_hidden_states``.
"""
from __future__ import annotations
@@ -23,9 +23,9 @@ def _ensure_text_encoder_2_stamped(instance: Any) -> None:
"""The shared ``stamp_wan21_checkpoints`` (called by the entrypoint after ``build_card``) only knows
Wan's subfolder superset, which does NOT include ``text_encoder_2`` (HunyuanVideo's CLIP secondary
encoder). Stamp it here — right before the component is materialized — from the already-stamped
``transformer`` checkpoint's parent (the model root). Idempotent / no-op once set; this keeps the fix
inside the recipe (no edit to the shared stamp). Build-time stamping via ``build_*_card(checkpoint_root)``
already covers the local-path path; this covers the registry+entrypoint path."""
``transformer`` checkpoint's parent (the model root). Idempotent / no-op once set; keeps the fix in the
recipe. Build-time stamping via ``build_*_card(checkpoint_root)`` covers the local-path case; this covers
the registry+entrypoint case."""
import os
te2 = instance.card.components.get("text_encoder_2")
if te2 is None or te2.checkpoint:
v2/recipes/hunyuan_video15/__init__.py
+7 -7
View File
@@ -1,12 +1,12 @@
"""HunyuanVideo 1.5 (T2V) — rectified-flow video DiT ported into the v2 substrate.
Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via
``ComponentSpec.adapter`` (``HunyuanVideo15DiT``/``HunyuanVideo15VAE``/``HunyuanVideo15QwenEncoder``/
``HunyuanVideo15ByT5Encoder`` in ``v2/platform/backends/torch_hunyuan_video15.py``) and a
``HunyuanVideo15DenoiseLoop`` (a thin ``WanDenoiseLoop`` subclass with z=32/16×/4× geometry). The
flow-match loop math is unchanged from Wan; the arch deltas (two text embeds, an image-embed list, the
33-channel i2v cond concat, the scalar-scaling VAE) live in the adapters. Registered in ``v2/registry.py``
(by the orchestrator) for ``hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v``.
Self-contained recipe package: the card declares its torch adapters via ``ComponentSpec.adapter``
(``HunyuanVideo15DiT``/``HunyuanVideo15VAE``/``HunyuanVideo15QwenEncoder``/``HunyuanVideo15ByT5Encoder`` in
``v2/platform/backends/torch_hunyuan_video15.py``) plus a ``HunyuanVideo15DenoiseLoop`` (a thin
``WanDenoiseLoop`` subclass with z=32/16×/4× geometry). The flow-match loop math is unchanged from Wan; the
arch deltas (two text embeds, an image-embed list, the 33-channel i2v cond concat, the scalar-scaling VAE)
live in the adapters. Registered in ``v2/registry.py`` (by the orchestrator) for
``hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v``.
"""
from __future__ import annotations
v2/recipes/hunyuan_video15/card.py
+14 -14
View File
@@ -1,14 +1,14 @@
"""HunyuanVideo 1.5 ModelCard — text→video (the registered 480p t2v preset; i2v + SR are later
capabilities the loop/adapter thread but the t2v program does not yet exercise on CPU).
"""HunyuanVideo 1.5 ModelCard — text→video (the registered 480p t2v preset). i2v + SR are later
capabilities the loop/adapter thread but the t2v program does not yet exercise on CPU.
Architecture deltas vs Wan (all declared on the card so the recipe is self-contained):
* DiT ``fastvideo.models.dits.hunyuanvideo15:HunyuanVideo15Transformer3DModel`` — a rectified-flow
velocity network (54 dual-stream MMDoubleStream blocks, hidden 2048, in/out 65/32). Its forward takes
a LIST of two text embeds (Qwen2.5-VL + ByT5) and a LIST image-embed; the new ``HunyuanVideo15DiT``
adapter marshals those, so ``HunyuanVideo15DenoiseLoop`` (a thin ``WanDenoiseLoop`` subclass with
z=32/16×/4× geometry) drives it with the unchanged flow-match math.
* VAE ``fastvideo.models.vaes.hunyuan15vae`` — causal-3D, z=32, 16× spatial / 4× temporal, normalized
by a SCALAR ``scaling_factor=1.03682`` -> the ``HunyuanVideo15VAE`` adapter (NOT Wan's mean/std).
* DiT ``HunyuanVideo15Transformer3DModel`` — a rectified-flow velocity network (54 dual-stream
MMDoubleStream blocks, hidden 2048, in/out 65/32). Its forward takes a list of two text embeds
(Qwen2.5-VL + ByT5) and a list image-embed; the ``HunyuanVideo15DiT`` adapter marshals those, so
``HunyuanVideo15DenoiseLoop`` (a thin ``WanDenoiseLoop`` subclass with z=32/16×/4× geometry) drives it
with the unchanged flow-match math.
* VAE ``hunyuan15vae`` — causal-3D, z=32, 16× spatial / 4× temporal, normalized by a SCALAR
``scaling_factor=1.03682`` via the ``HunyuanVideo15VAE`` adapter (not Wan's mean/std).
* Text encoders: ``text_encoder`` = Qwen2.5-VL (3584-d, ``hidden_states[-3]`` cropped of 108 template
tokens) via ``HunyuanVideo15QwenEncoder``; ``text_encoder_2`` = ByT5/Glyph (1472-d, quoted-glyph text,
may be empty) via ``HunyuanVideo15ByT5Encoder``.
@@ -60,7 +60,7 @@ _DEFAULT_STEPS = 50
def _linspace_sigmas(num_steps: int) -> tuple[float, ...]:
"""Faithful to ``presets.py:_sigmas`` — ``np.linspace(1,0,n+1)[:-1]`` (drops the terminal 0). The
loop's ``build_schedule`` consumes this verbatim when present (FlowShiftPolicy passes it through)."""
loop's ``build_schedule`` consumes it verbatim when present (FlowShiftPolicy passes it through)."""
return tuple((1.0 - i / num_steps) for i in range(num_steps))
@@ -166,10 +166,10 @@ def build_hunyuan_video15_card(model_id: str = "hunyuan-video-1.5-t2v-480p",
card.validate()
if checkpoint_root:
stamp_wan21_checkpoints(card, checkpoint_root)
# stamp_wan21_checkpoints only knows the Wan superset of subfolders; HunyuanVideo 1.5 adds a
# second text encoder (ByT5 in ``text_encoder_2/``) that the shared stamp does not cover. Stamp it
# here so the recipe stays self-contained (no edit to the shared stamp). The model root is the
# parent of any already-stamped subfolder (stamp resolves HF ids -> a local snapshot).
# stamp_wan21_checkpoints does not cover HunyuanVideo 1.5's second text encoder (ByT5 in
# ``text_encoder_2/``), so stamp it here to keep the recipe self-contained (no edit to the shared
# stamp). The model root is the parent of any already-stamped subfolder (stamp resolves HF ids -> a
# local snapshot).
import os
if "text_encoder_2" in card.components and not card.components["text_encoder_2"].checkpoint:
model_root = os.path.dirname(os.path.normpath(card.components["transformer"].checkpoint))
v2/recipes/hunyuan_video15/loop.py
+18 -21
View File
@@ -1,39 +1,36 @@
"""HunyuanVideo15DenoiseLoop — the canonical flow-match denoise loop, parameterized for HunyuanVideo 1.5.
"""HunyuanVideo15DenoiseLoop — the flow-match denoise loop parameterized for HunyuanVideo 1.5.
The loop *math* is identical to Wan's (``WanDenoiseLoop``): N deterministic flow-match Euler steps over
the full latent, ``ClassicCFG`` combine, ``timestep = sigma*1000``, ``x_next = x + (σ_next-σ)·v``. This is
faithful to ``fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py`` — ``scale_model_input``
is the identity for flow-match, ``step`` is ``prev_sample = sample + (sigma_next - sigma)*model_output``,
and HunyuanVideo 1.5's DiT predicts velocity directly (no x0 reconstruction). The latent is initialized at
``randn·σ[0]`` with ``σ[0]=1.0`` (the preset's ``linspace(1,0,n+1)[:-1]`` schedule).
The loop math is identical to ``WanDenoiseLoop``: N deterministic flow-match Euler steps over the full
latent, ``ClassicCFG`` combine, ``timestep = sigma*1000``, ``x_next = x + (σ_next-σ)·v``. HunyuanVideo 1.5's
DiT predicts velocity directly (no x0 reconstruction); the latent starts at ``randn·σ[0]`` with ``σ[0]=1.0``.
The ONLY architectural delta lives in the ``HunyuanVideo15DiT`` torch adapter (two text embeds + an
image-embed list + the 33-channel i2v cond concat). Because ``WanDenoiseLoop`` already threads
``context=`` (image embeds) and ``cond=`` (the conditioning latent) into the adapter, the cleanest port
is to subclass it: change only the latent geometry (z=32, 16× spatial, 4× temporal) and the slot names
the loop reads for the i2v signal. For t2v those slots are absent (``None``) → the dit-call degenerates
to ``dit(latent, text_embed, sigma)`` exactly like Wan t2v, and the toy backend's ``ToyDiT`` runs it.
The only architectural delta lives in the ``HunyuanVideo15DiT`` torch adapter (two text embeds + an
image-embed list + the 33-channel i2v cond concat). Because ``WanDenoiseLoop`` already threads ``context=``
(image embeds) and ``cond=`` (the conditioning latent) into the adapter, the port subclasses it: change only
the latent geometry (z=32, 16× spatial, 4× temporal) and the slot names the loop reads for the i2v signal.
For t2v those slots are ``None`` → the dit-call degenerates to ``dit(latent, text_embed, sigma)`` like Wan
t2v, and the toy backend's ``ToyDiT`` runs it.
BRINGUP: i2v conditioning (``context``=zero image embeds, ``cond``=33ch first-frame latent) is threaded
but the program currently writes neither slot for the registered t2v preset, so the loop is pure t2v on
CPU; the i2v VAE-encode node + image-latent slot are documented hooks (see ``program.py``).
BRINGUP: i2v conditioning (``context``=zero image embeds, ``cond``=33ch first-frame latent) is threaded but
the program writes neither slot for the registered t2v preset, so the loop is pure t2v on CPU; the i2v
VAE-encode node + image-latent slot are documented hooks (see ``program.py``).
"""
from __future__ import annotations
from v2.recipes.wan21.loop import WanDenoiseLoop
# HunyuanVideo 1.5 causal-3D VAE compression (configs/models/vaes/hunyuan15vae.py): z=32, 16× spatial,
# 4× temporal. The DiT's noisy-latent in_channels is 32 (out_channels); i2v reaches 65 by concatenating
# the 33ch first-frame cond latent inside the adapter, so the loop's latent stays 32ch.
# 4× temporal. The DiT's noisy-latent in_channels is 32; i2v reaches 65 by concatenating the 33ch
# first-frame cond latent inside the adapter, so the loop's latent stays 32ch.
HUNYUAN15_LATENT_CHANNELS = 32
HUNYUAN15_SPATIAL_RATIO = 16
HUNYUAN15_TEMPORAL_RATIO = 4
class HunyuanVideo15DenoiseLoop(WanDenoiseLoop):
"""Wan flow-match loop with HunyuanVideo 1.5 latent geometry. The i2v conditioning the base loop reads
from ``ctx.slots`` is keyed the same way (``i2v_cond``/``i2v_img_embeds``), so a future i2v program can
write them with no loop change; the t2v program leaves them unset (pure t2v)."""
"""Wan flow-match loop with HunyuanVideo 1.5 latent geometry. The i2v slots the base loop reads from
``ctx.slots`` are keyed identically (``i2v_cond``/``i2v_img_embeds``), so a future i2v program can write
them with no loop change; the t2v program leaves them unset (pure t2v)."""
def __init__(self,
*,
v2/recipes/hunyuan_video15/program.py
+16 -19
View File
@@ -1,8 +1,8 @@
"""HunyuanVideo 1.5 t2v program: text_encode → diffusion_denoise → vae_decode.
Same inline shape as the Wan/Cosmos t2v programs — the arch specifics (two text embeds, the image-embed
list, the 33ch i2v concat) live entirely in the ``HunyuanVideo15DiT`` adapter, so the node graph is
unchanged except for the text-encode node, which must produce BOTH text embeddings.
list, the 33ch i2v concat) live entirely in the ``HunyuanVideo15DiT`` adapter, so the node graph differs
only in the text-encode node, which must produce BOTH text embeddings.
HunyuanVideo 1.5 conditions on two encoders (faithful to ``HunyuanVideo15Pipeline``'s
``TextEncodingStage(text_encoders=[text_encoder, text_encoder_2])``):
@@ -11,10 +11,10 @@ HunyuanVideo 1.5 conditions on two encoders (faithful to ``HunyuanVideo15Pipelin
zero-length embedding when the prompt carries no quoted text.
The text-encode node packs them as a ``(qwen, byt5)`` tuple in the slot; ``HunyuanVideo15DiT`` unpacks
that into the two-element ``encoder_hidden_states`` list. On the CPU **toy** backend there is a single
that into the two-element ``encoder_hidden_states`` list. On the CPU toy backend there is a single
``text_encoder`` component (``ToyTextEncoder``) and no ``text_encoder_2``: the node then writes the plain
toy embedding (a single numpy array), which ``ToyDiT`` consumes directly — so the program CPU-verifies
unchanged. The feature cache (design_v3 §7.2) still applies to the Qwen path via ``cached_text_encode``.
unchanged. The feature cache still applies to the Qwen path via ``cached_text_encode``.
BRINGUP (i2v): the i2v path adds a VAE-encode node that writes the 33-channel first-frame conditioning
latent into ``i2v_cond`` and zero image embeds into ``i2v_img_embeds`` (the slots the loop already reads);
@@ -33,11 +33,10 @@ from v2.recipes.common import cached_text_encode
def _ensure_byt5_checkpoint(instance: Any) -> None:
"""Stamp ``text_encoder_2``'s weights subfolder if the shared stamp left it empty.
``VideoGenerator.from_pretrained`` stamps the card via ``stamp_wan21_checkpoints``, whose subfolder
superset is Wan-centric and does NOT cover HunyuanVideo 1.5's second text encoder (ByT5 lives in
``text_encoder_2/``). Rather than edit the shared stamp, the recipe stamps it here at first use,
deriving the diffusers model root from the already-stamped primary ``text_encoder`` subfolder
(sibling layout). Idempotent: a no-op once ``checkpoint`` is set (e.g. by the card's own stamp path)."""
``stamp_wan21_checkpoints`` (run by ``VideoGenerator.from_pretrained``) is Wan-centric and does NOT cover
HunyuanVideo 1.5's second text encoder (ByT5 lives in ``text_encoder_2/``). Rather than edit the shared
stamp, derive the model root from the already-stamped primary ``text_encoder`` (sibling layout) and stamp
here at first use. Idempotent: a no-op once ``checkpoint`` is set."""
import os
spec2 = instance.card.components.get("text_encoder_2")
if spec2 is None or spec2.checkpoint:
@@ -51,16 +50,14 @@ def _text_encode(instance: Any, slots: dict, request: Any, ctx: Any) -> None:
"""Encode prompt + negative prompt. When a second text encoder (ByT5) exists and the prompt carries
quoted glyph text, pack the two embeds; otherwise write the single Qwen embedding.
The shared ``WanDenoiseLoop`` does numpy bookkeeping on the conditioning slot (``np.asarray(e).nbytes``
for the resident-bytes estimate, and ``np.array(e)`` to allocate the CUDA-graph static workspace). A
plain Python tuple of the (qwen[seq,3584], byt5[bseq,1472]) embeds is INHOMOGENEOUS, so np.asarray /
np.array raise. For the common t2v case the ByT5 stream is a zero-length ``[0,1472]`` embedding (no
quoted text), which the DiT concatenates as 0 tokens — i.e. a no-op. So we pass ONLY the Qwen embedding
then; the ``HunyuanVideo15DiT`` adapter rebuilds the empty ByT5 default internally (its ``byt5_e is
None`` branch). When glyph text IS present we pack the pair as a numpy OBJECT array (length 2): an
object array is a single ndarray, so ``np.asarray(...).nbytes`` and ``np.array(...)`` both succeed
(no stacking), and the adapter unpacks it. This keeps the t2v path capturable (single ndarray) while
still threading real ByT5 embeddings for glyph prompts (those steps eager-break on the object array)."""
Why not a plain tuple: ``WanDenoiseLoop`` does numpy bookkeeping on the conditioning slot
(``np.asarray(e).nbytes`` for resident-bytes, ``np.array(e)`` to allocate the CUDA-graph workspace), and
a tuple of (qwen[seq,3584], byt5[bseq,1472]) is INHOMOGENEOUS so both raise. For the common t2v case the
ByT5 stream is a zero-length ``[0,1472]`` embedding (no quoted text), a 0-token DiT no-op, so we pass ONLY
the Qwen embedding and the adapter rebuilds the empty ByT5 default (its ``byt5_e is None`` branch). When
glyph text IS present we pack the pair as a numpy OBJECT array (a single ndarray, so the np.asarray /
np.array calls succeed without stacking) and the adapter unpacks it. This keeps the t2v path capturable
(single ndarray) while still threading real ByT5 embeddings for glyph prompts (which eager-break)."""
has_byt5 = "text_encoder_2" in instance.card.components
# The dual-encoder packing is a GPU-path concern: the CPU toy backend's ToyDiT takes a single
# [seq,dim] float embedding (and its ToyTextEncoder has no glyph/ByT5 stream), so on the toy we pass
v2/recipes/hunyuangamecraft/__init__.py
+6 -7
View File
@@ -1,12 +1,11 @@
"""HunyuanGameCraft (interactive camera/action-conditioned i2v) ported into the v2 substrate.
Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via
``ComponentSpec.adapter`` (``GameCraftDiT``/``GameCraftVAE``/``GameCraftLlamaEncoder``/
``GameCraftClipEncoder`` in ``v2/platform/backends/torch_hunyuangamecraft.py``) and a new
``GameCraftDenoiseLoop`` (flow-match Euler + 33ch concat + per-step clean-ref injection). The registered
path is the t2v/degenerate denoise; the camera/action (CameraNet Plücker) conditioning is BRINGUP — it
needs a request-API camera-input channel v2 lacks today. Registered in ``v2/registry.py`` by the
orchestrator from this package's returned metadata.
Self-contained recipe package: the card declares its torch adapters via ``ComponentSpec.adapter``
(``GameCraftDiT``/``GameCraftVAE``/``GameCraftLlamaEncoder``/``GameCraftClipEncoder`` in
``v2/platform/backends/torch_hunyuangamecraft.py``) plus a ``GameCraftDenoiseLoop`` (flow-match Euler +
33ch concat + per-step clean-ref injection). The registered path is the t2v/degenerate denoise; the
camera/action (CameraNet Plücker) conditioning is BRINGUP — it needs a request-API camera-input channel v2
lacks today. Registered in ``v2/registry.py`` by the orchestrator from this package's returned metadata.
"""
from __future__ import annotations
v2/recipes/hunyuangamecraft/card.py
+19 -20
View File
@@ -1,20 +1,19 @@
"""HunyuanGameCraft ModelCard — camera/action-conditioned interactive i2v (the registered preset is the
t2v/degenerate path; the action/camera conditioning is BRINGUP — see the recipe ``__init__`` docstring).
"""HunyuanGameCraft ModelCard — camera/action-conditioned interactive i2v. The registered preset is the
t2v/degenerate path; the action/camera conditioning is BRINGUP (see the recipe ``__init__`` docstring).
Architecture deltas vs Wan (all declared on the card so the recipe is self-contained):
* DiT ``fastvideo.models.dits.hunyuangamecraft:HunyuanGameCraftTransformer3DModel`` — a flow-match DiT
whose forward takes a **33-channel** input ``cat([latent16|gt_latent16|mask1])`` (assembled by the
loop/adapter, NOT the DiT), a **list** of text states ``[LLaMA(4096d), CLIP-pooled(768d)]``, an
optional ``camera_states`` (CameraNet Plücker), and ``guidance=None`` (plain ClassicCFG, NO embedded
guidance). Adapter -> ``GameCraftDiT``; the new ``GameCraftDenoiseLoop`` does the 33ch concat
(via the adapter), per-step clean-ref injection, and flow-match Euler.
* VAE ``fastvideo.models.vaes.gamecraftvae:GameCraftVAE`` — SCALAR ``scaling_factor=0.476986``
normalization (NOT Wan per-channel stats). Adapter -> ``GameCraftVAE`` (cannot reuse ``WanVAE``).
* Text ``fastvideo.models.encoders.llama:LlamaModel`` (LLaVA-LLaMA-3-8B, 4096d -> ``text_states``) and
``fastvideo.models.encoders.clip:CLIPTextModel`` (CLIP ViT-L/14, 768d pooled -> ``text_states_2``).
Adapters -> ``GameCraftLlamaEncoder`` / ``GameCraftClipEncoder``.
The diffusers checkpoint layout has extra ``text_encoder_2/`` + ``tokenizer_2/`` subfolders the Wan stamp
lacks, so this file extends ``stamp_wan21_checkpoints`` with a GameCraft-specific stamp.
* DiT ``HunyuanGameCraftTransformer3DModel`` — a flow-match DiT whose forward takes a 33-channel input
``cat([latent16|gt_latent16|mask1])`` (assembled by the adapter, not the DiT), a list of text states
``[LLaMA(4096d), CLIP-pooled(768d)]``, an optional ``camera_states`` (CameraNet Plücker), and
``guidance=None`` (plain ClassicCFG, no embedded guidance). Adapter ``GameCraftDiT``; the
``GameCraftDenoiseLoop`` does the 33ch concat (via the adapter), per-step clean-ref injection, and
flow-match Euler.
* VAE ``GameCraftVAE`` — SCALAR ``scaling_factor=0.476986`` normalization (not Wan per-channel stats);
cannot reuse ``WanVAE``.
* Text ``LlamaModel`` (LLaVA-LLaMA-3-8B, 4096d -> ``text_states``) and ``CLIPTextModel`` (CLIP ViT-L/14,
768d pooled -> ``text_states_2``). Adapters ``GameCraftLlamaEncoder`` / ``GameCraftClipEncoder``.
The diffusers layout adds ``text_encoder_2/`` + ``tokenizer_2/`` subfolders the Wan stamp lacks, so this
file extends ``stamp_wan21_checkpoints`` with a GameCraft-specific stamp.
"""
from __future__ import annotations
@@ -159,11 +158,11 @@ _GAMECRAFT_SUBFOLDERS = {
def stamp_gamecraft_checkpoints(card: ModelCard, model_root: str) -> ModelCard:
"""GPU deploy-time: point each component's ``ComponentSpec.checkpoint`` at its weights subfolder under
``model_root`` (a local diffusers dir or an HF id resolved via ``snapshot_download``). Mirrors
``stamp_wan21_checkpoints`` but covers GameCraft's ``text_encoder_2`` (the Wan stamp lacks it). The
sibling ``tokenizer/`` (LLaMA) and ``tokenizer_2/`` (CLIP) + ``scheduler/`` subfolders are loaded by
the torch adapters from ``dirname(checkpoint)``. Mutates and returns the card. BRINGUP risk A."""
"""Point each component's ``ComponentSpec.checkpoint`` at its weights subfolder under ``model_root`` (a
local diffusers dir or an HF id resolved via ``snapshot_download``). Like ``stamp_wan21_checkpoints``
but also covers GameCraft's ``text_encoder_2`` (the Wan stamp lacks it). The torch adapters load the
sibling ``tokenizer/`` (LLaMA), ``tokenizer_2/`` (CLIP), and ``scheduler/`` subfolders from
``dirname(checkpoint)``. Mutates and returns the card. BRINGUP risk A."""
import os
if not os.path.isdir(model_root):
from huggingface_hub import snapshot_download
v2/recipes/hunyuangamecraft/loop.py
+33 -36
View File
@@ -1,37 +1,35 @@
"""GameCraftDenoiseLoop — a flow-match Euler denoise loop with GameCraft i2v + camera conditioning.
"""GameCraftDenoiseLoop — flow-match Euler denoise with GameCraft i2v + camera conditioning.
GameCraft (HunyuanGameCraft) is an *interactive* world model: a camera/action-conditioned i2v diffusion
DiT. The sampler is plain flow-match (``FlowMatchEulerDiscreteScheduler``, ``flow_shift=5.0``,
``timestep = σ·1000``) with standard ClassicCFG (``guidance_scale=6.0``, ``guidance=None`` to the DiT — NO
embedded guidance). So the *integrator* is identical to ``WanDenoiseLoop`` (``x_next = x + (σ_next−σ)·v``
via ``FLOW_MATCH_STEP``). The deltas this loop carries (faithful to
GameCraft (HunyuanGameCraft) is an interactive world model: a camera/action-conditioned i2v diffusion DiT.
The sampler is plain flow-match (``FlowMatchEulerDiscreteScheduler``, ``flow_shift=5.0``,
``timestep = σ·1000``) with standard ClassicCFG (``guidance_scale=6.0``, ``guidance=None`` to the DiT, no
embedded guidance), so the integrator is identical to ``WanDenoiseLoop`` (``x_next = x + (σ_next−σ)·v`` via
``FLOW_MATCH_STEP``). The deltas this loop carries (faithful to
``fastvideo/pipelines/stages/gamecraft_denoising.py``):
1. **33-channel input is assembled by the DiT adapter, not here.** The DiT forward wants
``cat([latent16 | gt_latent16 | mask1], dim=1)``; the loop hands the bare 16ch latent + a ``cond``
pack ``{"gt_latents","mask","camera_states"}`` to the ``GameCraftDiT`` adapter, which concats. The
toy ``ToyDiT`` accepts (and ignores) the ``cond`` kwarg, so this loop CPU-verifies unchanged.
pack ``{"gt_latents","mask","camera_states"}`` to the adapter, which concats. The toy ``ToyDiT``
accepts and ignores the ``cond`` kwarg, so this loop CPU-verifies unchanged.
2. **Per-step clean-ref injection (frame_replace).** Official GameCraft overwrites the *conditioned*
latent frames (``mask>0.5``) with the clean reference latent at EVERY step (not just step 0), so the
conditioned frames stay noise-free throughout denoising. ``next()`` mutates the working latent before
the forward when a reference latent is present.
2. **Per-step clean-ref injection (frame_replace).** Official GameCraft overwrites the conditioned latent
frames (``mask>0.5``) with the clean reference latent at every step (not just step 0) to keep them
noise-free. ``next()`` mutates the working latent before the forward when a reference latent is present.
3. **Dual text encoders.** ``encoder_hidden_states = [LLaMA_states(4096d), CLIP_pooled(768d)]``. The
loop threads the LLaMA states as the usual ``text_embed`` (cond vs uncond) and the CLIP pooled embed
3. **Dual text encoders.** ``encoder_hidden_states = [LLaMA_states(4096d), CLIP_pooled(768d)]``. The loop
threads the LLaMA states as the usual ``text_embed`` (cond vs uncond) and the CLIP pooled embed
through the adapter's ``context=`` slot per branch, so ClassicCFG combines both correctly.
4. **BRINGUP — camera/action conditioning.** GameCraft is fundamentally CameraNet (Plücker)
conditioned, with the forward special-casing the latent temporal length (18/9/10) for autoregressive
multi-chunk generation. v2 has no camera/action request field or Plücker-conversion node yet, so the
registered path passes ``camera_states=None`` -> the DiT's ``if camera_states is not None`` branch is
skipped (the standard / degenerate denoise). For pure t2v (no reference image) gt_latents+mask are
zero -> the 33ch concat reduces to ``[noise | 0 | 0]``, exactly the ``GameCraftDenoisingStage``
fallback. Wiring real camera/action + the 18/9/10 length special-casing needs a request-API
extension (see ``program.py`` blockers).
4. **BRINGUP — camera/action conditioning.** GameCraft is CameraNet (Plücker) conditioned, with the
forward special-casing the latent temporal length (18/9/10) for autoregressive multi-chunk generation.
v2 has no camera/action request field or Plücker-conversion node yet, so the registered path passes
``camera_states=None`` -> the DiT's ``if camera_states is not None`` branch is skipped (the degenerate
denoise). For pure t2v (no reference image) gt_latents+mask are zero -> the 33ch concat reduces to
``[noise | 0 | 0]``, exactly the ``GameCraftDenoisingStage`` fallback. Wiring real camera/action + the
18/9/10 length special-casing needs a request-API extension (see ``program.py`` blockers).
All per-request state lives in ``LoopState`` (interleave-safe). ``next`` is kernel-free; ``advance`` folds
the result + (under ROLLOUT) captures a trajectory slice — same shape as ``WanDenoiseLoop``.
the result and, under ROLLOUT, captures a trajectory slice — same shape as ``WanDenoiseLoop``.
"""
from __future__ import annotations
@@ -115,7 +113,7 @@ class GameCraftDenoiseLoop:
st.scratch["guidance_scale"] = float(req.diffusion.guidance_scale)
st.scratch["stream_video"] = bool(req.outputs.stream.get("video"))
# i2v conditioning slots written by the image-cond program node. All None for pure t2v -> the 33ch
# concat degenerates to [noise | 0 | 0] and the ref-injection branch is inert (the fastvideo
# concat degenerates to [noise | 0 | 0] and the ref-injection branch is inert (the
# GameCraftDenoisingStage fallback). camera_states is a BRINGUP slot (no request plumbing yet).
st.scratch["gt_latents"] = ctx.slots.get("gt_latents")
st.scratch["conditioning_mask"] = ctx.slots.get("conditioning_mask")
@@ -128,9 +126,9 @@ class GameCraftDenoiseLoop:
"""Per-step frame_replace: overwrite conditioned frames (mask>0.5) with the clean reference latent.
Faithful to ``GameCraftDenoisingStage`` (``latents[:,:,t_idx] = ref[:,:,0]`` for each conditioned
``t``). Loop latents here are unbatched ``[C,T,h,w]`` (the toy/CPU geometry), so the mask's temporal
axis is read at ``mask[0,0,:,0,0]`` when 5D (matching the fastvideo ``[B,1,T,H,W]`` mask), or as a
1D ``[T]`` vector otherwise. ``ref`` is the clean ref latent ``[C,1,h,w]`` (or ``[C,T,h,w]``)."""
``t``). Loop latents here are unbatched ``[C,T,h,w]`` (toy/CPU geometry), so the mask's temporal axis
is read at ``mask[0,0,:,0,0]`` when 5D (the fastvideo ``[B,1,T,H,W]`` mask) or as a 1D ``[T]`` vector
otherwise. ``ref`` is the clean ref latent ``[C,1,h,w]`` (or ``[C,T,h,w]``)."""
m = np.asarray(mask)
if m.ndim == 5: # fastvideo-shaped [B,1,T,H,W]: read the temporal slice
frame_mask = m[0, 0, :, 0, 0]
@@ -143,10 +141,9 @@ class GameCraftDenoiseLoop:
r = np.asarray(ref, dtype="float32")
ref_frame = r[:, 0] if r.ndim >= 4 else r # [C,h,w]
out = np.array(x, dtype="float32", copy=True)
# The conditioned-frame replacement requires the ref frame to match the working latent's per-frame
# geometry (it always does on the real VAE geometry). If a backend's encode produced a mismatched
# shape (e.g. the crude CPU toy VAE that does not spatially downsample), skip injection rather than
# crash — the degenerate denoise still proceeds.
# Frame replacement needs the ref frame to match the working latent's per-frame geometry (always
# true on the real VAE). If a backend's encode produced a mismatched shape (e.g. the toy VAE that
# does not spatially downsample), skip injection rather than crash — the degenerate denoise proceeds.
if ref_frame.shape != out[:, 0].shape:
return out
n_t = min(out.shape[1], int(frame_mask.shape[0]))
@@ -180,14 +177,14 @@ class GameCraftDenoiseLoop:
scale = st.scratch["guidance_scale"]
cfg, precision = self.cfg, self.precision
# One ``cond`` pack reused across CFG branches (gt/mask/camera are branch-independent); the adapter
# assembles the 33ch concat from it. None entries -> zero gt/mask + no camera inside the adapter.
# assembles the 33ch concat from it. None entries -> zero gt/mask and no camera inside the adapter.
cond_pack = {"gt_latents": gt, "mask": mask, "camera_states": camera}
def _velocity(model: Any, x_: np.ndarray, sigma_t_: float, pe_: Any, clip_: Any, ne_: Any, clip_n_: Any,
scale_: float) -> np.ndarray:
# Conditioned forward + ClassicCFG combine. The CLIP pooled embed rides ``context=`` per branch
# so text_states_2 differs cond vs uncond; ``cond`` carries the i2v/camera pack. The toy DiT
# accepts+ignores both kwargs, so this is CPU-toy compatible.
# accepts and ignores both kwargs, so this is CPU-toy compatible.
dit = model.component(expert_id)
preds = {
b:
@@ -230,9 +227,9 @@ class GameCraftDenoiseLoop:
},
run=run,
label=f"gamecraft.denoise.{i}",
# i2v ref-injection mutates host state + threads non-workspace conditioning each step -> the
# static-buffer capture form would miss it, so this step takes the eager path (like cosmos2's
# gated frame-injection). Pure-t2v could capture, but keeping one path is simpler + correct.
# i2v ref-injection mutates host state and threads non-workspace conditioning each step, which
# static-buffer capture would miss, so this step takes the eager path (like cosmos2's gated
# frame-injection). Pure-t2v could capture, but keeping one path is simpler and correct.
capturable=False)
def advance(self, st: LoopState, result: StepResult) -> LoopState:
v2/recipes/hunyuangamecraft/program.py
+21 -21
View File
@@ -3,26 +3,26 @@
Same backbone as the Wan/Cosmos t2v programs, with two GameCraft-specific deltas:
* **Dual text encode.** ``_text_encode_dual`` runs BOTH encoders: the LLaMA encoder (``text_encoder``)
writes the DiT's ``text_states`` into ``text_embeds``/``neg_text_embeds`` (the usual slots the loop
reads), and the CLIP encoder (``text_encoder_2``) writes the pooled ``text_states_2`` into
* **Dual text encode.** ``_text_encode_dual`` runs both encoders: LLaMA (``text_encoder``) writes the
DiT's ``text_states`` into ``text_embeds``/``neg_text_embeds`` (the usual slots the loop reads), and
CLIP (``text_encoder_2``) writes the pooled ``text_states_2`` into
``clip_text_embeds``/``clip_neg_text_embeds`` (threaded to the adapter's ``context=`` per CFG branch).
On the CPU toy both are ``ToyTextEncoder`` (cached via the content-hash feature cache).
* **Image conditioning node (i2v).** ``_image_cond_encode`` mirrors ``GameCraftImageVAEEncodingStage``:
VAE-encode the reference image -> scaled ``ref_latent`` (the ``GameCraftVAE`` adapter applies the
scalar ``scaling_factor`` already), repeat to ``T_lat`` frames -> ``gt_latents``, build the binary
VAE-encode the reference image -> scaled ``ref_latent`` (the ``GameCraftVAE`` adapter already applies
the scalar ``scaling_factor``), repeat to ``T_lat`` frames -> ``gt_latents``, build the binary
``conditioning_mask`` (latent_frames<=10: only frame 0; else first half), zero ``gt_latents`` on
non-conditioned frames, and expose ``ref_latent_for_injection = gt_latents[:, 0:1]`` for the loop's
per-step frame_replace. No request image -> a no-op (slots stay unset) -> the loop runs pure t2v
(zero gt/mask), exactly the fastvideo fallback.
per-step frame_replace. No request image -> no-op (slots stay unset) -> the loop runs pure t2v (zero
gt/mask), exactly the fastvideo fallback.
BRINGUP — the camera/action conditioning node is intentionally omitted: GameCraft's CameraNet needs a
Plücker-coordinate ``camera_states`` tensor derived from an action/camera sequence the v2 request API does
not carry yet, and the DiT's autoregressive latent-length special-casing (18/9/10) for multi-chunk
generation has no request surface. Adding a ``camera_encode`` node that writes the ``camera_states`` slot
(which the loop already reads) is the follow-up once the request API is extended. T_lat helper below
matches the GameCraft geometry (4x temporal compression).
not carry yet, and the DiT's autoregressive latent-length special-casing (18/9/10) has no request surface.
Adding a ``camera_encode`` node that writes the ``camera_states`` slot (which the loop already reads) is the
follow-up once the request API is extended. The T_lat helper below matches GameCraft's 4x temporal
compression.
"""
from __future__ import annotations
@@ -38,12 +38,12 @@ GAMECRAFT_SPATIAL_RATIO = 8
def _stamp_text_encoder_2(instance: Any) -> None:
"""The shared ``from_config`` stamps the card via ``stamp_wan21_checkpoints``, whose subfolder map
knows only ``transformer/vae/text_encoder`` — NOT GameCraft's second text encoder. So ``text_encoder_2``
arrives at GPU build time with an empty ``checkpoint`` and ``_require_checkpoint`` fails. Derive the
model root from the already-stamped ``text_encoder`` (LLaMA) checkpoint and point ``text_encoder_2`` at
the sibling ``text_encoder_2/`` subfolder, here in the recipe's own node (no edit to the shared stamp).
No-op once stamped or on the CPU toy (empty LLaMA checkpoint -> the toy factory needs no weights)."""
"""``stamp_wan21_checkpoints`` knows only ``transformer/vae/text_encoder``, not GameCraft's second text
encoder, so ``text_encoder_2`` reaches GPU build time with an empty ``checkpoint`` and
``_require_checkpoint`` fails. Derive the model root from the already-stamped ``text_encoder`` (LLaMA)
checkpoint and point ``text_encoder_2`` at the sibling ``text_encoder_2/`` subfolder, here in the
recipe's own node (no edit to the shared stamp). No-op once stamped or on the CPU toy (empty LLaMA
checkpoint -> the toy factory needs no weights)."""
import os
comps = instance.card.components
spec2 = comps.get("text_encoder_2")
@@ -84,10 +84,10 @@ def _image_cond_encode(instance: Any, slots: dict, request: Any, ctx: Any) -> No
return # T2V mode: leave the i2v slots unset (loop falls back to zero gt/mask)
d = request.diffusion
# The GameCraftVAE adapter applies the scalar scaling_factor in ``encode`` (matches the fastvideo
# stage's ``ref_latents.mul_(scaling_factor)``), so we use the returned latent directly. The real
# adapter also runs the resize/center-crop/normalize preprocessing on the PIL image (BRINGUP); the
# CPU toy VAE consumes the pixel array directly.
# The GameCraftVAE adapter applies the scalar scaling_factor in ``encode`` (matching the fastvideo
# stage's ``ref_latents.mul_(scaling_factor)``), so use the returned latent directly. The real adapter
# also runs resize/center-crop/normalize on the PIL image (BRINGUP); the toy VAE consumes the pixel
# array directly.
ref_latent = np.asarray(instance.component("vae").encode(pixels), dtype="float32") # [C,1,h,w] (or [C,h,w])
if ref_latent.ndim == 3: # [C, h, w] -> add the singleton temporal axis
ref_latent = ref_latent[:, None]
v2/recipes/hyworld/card.py
+15 -15
View File
@@ -1,28 +1,28 @@
"""HY-WorldPlay-Bidirectional ModelCard — interactive world model ported into the v2 substrate.
"""HY-WorldPlay-Bidirectional ModelCard: interactive world model on the v2 substrate.
HY-WorldPlay is a chunk-rollout *world model*: it generates video in temporal chunks where each
chunk after the first retrieves camera-aligned history "memory" frames as frozen context, and it
conditions on THREE text/image streams (Qwen2.5-VL mllm + ByT5 glyph + SigLIP image) plus per-frame
camera (viewmats/Ks, injected via ProPE attention) and per-frame action ids. The registered preset
here is the **bidirectional t2v / degenerate (no-action, single straight-trajectory pose) path** that
CPU-verifies; the full action/camera/memory conditioning is BRINGUP (needs a request-API extension to
carry the pose string -> viewmats/Ks/action expansion + the point-cloud memory retrieval). See
``loop.py`` and ``torch_hyworld.py`` for the BRINGUP markers.
HY-WorldPlay is a chunk-rollout world model: it generates video in temporal chunks where each chunk
after the first retrieves camera-aligned history "memory" frames as frozen context, and conditions on
three text/image streams (Qwen2.5-VL mllm + ByT5 glyph + SigLIP image) plus per-frame camera
(viewmats/Ks, injected via ProPE attention) and per-frame action ids. The registered preset is the
bidirectional t2v / degenerate (no-action, single straight-trajectory pose) path that CPU-verifies;
the full action/camera/memory conditioning is BRINGUP (needs a request-API extension to carry the pose
string -> viewmats/Ks/action expansion + point-cloud memory retrieval). See ``loop.py`` and
``torch_hyworld.py`` for the BRINGUP markers.
Architecture deltas vs Wan/Cosmos (all declared on the card so the recipe is self-contained):
* DiT ``fastvideo.models.dits.hyworld.hyworld:HYWorldTransformer3DModel`` (subclasses
HunyuanVideo15Transformer3DModel) — a **rectified-flow velocity** predictor whose forward takes a
65ch pre-concatenated latent ``[noise(32) | cond(32) + mask(1)]``, a LIST of text embeds
``[qwen, byt5]``, a LIST image embed ``[siglip]``, a PER-LATENT-FRAME LongTensor ``timestep`` (+ a
HunyuanVideo15Transformer3DModel) — a rectified-flow velocity predictor whose forward takes a
65ch pre-concatenated latent ``[noise(32) | cond(32) + mask(1)]``, a list of text embeds
``[qwen, byt5]``, a list image embed ``[siglip]``, a per-latent-frame LongTensor ``timestep`` (+ a
scalar ``timestep_txt``), and ``viewmats/Ks/action`` per-frame camera/action conditioning. The
adapter (``HYWorldDiT``) marshals all of this INTERNALLY so the loop only hands it
adapter (``HYWorldDiT``) marshals all of this internally so the loop only hands it
``dit(latent, text_embed, sigma)`` (the cosmos2/ToyDiT-compatible signature).
* VAE ``fastvideo.models.vaes.hyworldvae:AutoencoderKLHYWorld`` (z=32, 16x spatial, 4x temporal,
scaling_factor=1.03682) — encode -> ``.mode() * scaling_factor`` (no shift/mean), decode inverts.
* Encoders: Qwen2.5-VL (mllm primary text) + ByT5/T5 (glyph) + SigLIP (image). For the t2v preset the
glyph + image streams are zeroed (the DiT detects an all-zero image and masks it).
* Sampler: flow-match (``FlowMatchEulerDiscreteScheduler``) but the preset OVERRIDES the schedule with
explicit sigmas ``linspace(1.0, 0.0, 51)[:-1]`` (50 steps), flow_shift 5.0, guidance 6.0.
* Sampler: flow-match (``FlowMatchEulerDiscreteScheduler``), but the preset overrides the schedule
with explicit sigmas ``linspace(1.0, 0.0, 51)[:-1]`` (50 steps), flow_shift 5.0, guidance 6.0.
``stamp_wan21_checkpoints`` applies (diffusers transformer/vae/text_encoder/text_encoder_2/image_encoder
subfolder layout); the superset stamp helper already covers ``text_encoder_2`` + ``image_encoder``.
"""
v2/recipes/hyworld/loop.py
+7 -7
View File
@@ -17,15 +17,15 @@ the 65ch latent concat, the 3-stream conditioning list, the per-frame heterogene
camera/action tensors INTERNALLY, so this loop only calls the cosmos2/ToyDiT-compatible
``dit(latent, text_embed, sigma)`` (with the per-chunk camera context threaded via ``context=``).
BRINGUP (documented honestly — needs a request-API extension to be real):
* The default preset (num_frames=125 -> ~32 latent frames, chunk 16 -> 2 chunks) WILL execute chunk>0.
BRINGUP (needs a request-API extension to be real):
* The default preset (num_frames=125 -> ~32 latent frames, chunk 16 -> 2 chunks) will execute chunk>0.
Camera-aligned memory retrieval, ProPE camera attention, and per-frame action conditioning all need
``viewmats``/``Ks``/``action`` carried from a pose string (``pose="w-31"`` -> ``pose_to_input`` +
``compute_latent_num``). The v2 request/program has no slot for the pose string yet, so on the CPU
toy path (and absent a pose) this loop runs the **degenerate single-chunk t2v path**: it clamps
``chunk_num`` to the number of whole chunks but the camera/memory selection is skipped (the adapter
builds zero camera/action internally). On a GPU box with a pose-carrying request, the chunk>0 branch
(memory retrieval + frozen context) activates — that path is written-not-run here.
``compute_latent_num``). The v2 request/program has no slot for the pose string yet, so absent a
pose (the CPU toy path) this loop runs the degenerate per-chunk t2v path: it clamps ``chunk_num`` to
the number of whole chunks but skips camera/memory selection (the adapter builds zero camera/action
internally). On a GPU box with a pose-carrying request the chunk>0 branch (memory retrieval + frozen
context) activates; that path is written-not-run here.
"""
from __future__ import annotations
v2/recipes/image_video/__init__.py
+7 -5
View File
@@ -1,6 +1,6 @@
"""Text-to-image → image-to-video: two distinct generation models chained by a cross-model
``Workflow`` (design_v3 §13). The realistic FLUX→Wan pipeline — two cards, two weight sets — proving
the design composes across model *instances*, not just loops within one (which LTX-2 already covers)."""
"""Text-to-image -> image-to-video: two distinct generation models chained by a cross-model
``Workflow``. The realistic FLUX->Wan pipeline (two cards, two weight sets) proving the design
composes across model instances, not just loops within one (which LTX-2 already covers)."""
from __future__ import annotations
from v2.recipes.image_video.card import build_flux_t2i_card, build_wan_i2v_card
@@ -11,5 +11,7 @@ from v2.recipes.image_video.program import (
build_wan_i2v_program,
)
__all__ = ["build_flux_t2i_card", "build_wan_i2v_card", "build_flux_t2i_program",
"build_wan_i2v_program", "build_t2i_then_i2v_workflow", "build_t2i_i2v_extend_workflow"]
__all__ = [
"build_flux_t2i_card", "build_wan_i2v_card", "build_flux_t2i_program", "build_wan_i2v_program",
"build_t2i_then_i2v_workflow", "build_t2i_i2v_extend_workflow"
]
v2/recipes/image_video/card.py
+48 -29
View File
@@ -1,13 +1,13 @@
"""Two single-purpose generation cards for the T2I→I2V workflow (design_v3 §4, §13).
"""Two single-purpose generation cards for the T2I->I2V workflow.
These are deliberately *separate models* — distinct weights, distinct cards — because the realistic
"text-to-image followed by image-to-video" pipeline is FLUX→Wan, not one model's two stages (that
same-card case is already covered by LTX-2's base→refine program). Chaining them is a ``Workflow``
These are deliberately separate models (distinct weights, distinct cards) because the realistic
"text-to-image followed by image-to-video" pipeline is FLUX->Wan, not one model's two stages (that
same-card case is already covered by LTX-2's base->refine program). Chaining them is a ``Workflow``
(``program/workflow.py``), not a single-instance ``Program``.
Both reuse the *unchanged* ``WanDenoiseLoop`` — the image-conditioning for I2V lives in the I2V
*program* (it folds the conditioning image into ``text_embeds``), so no loop change is needed. That
is the point: a new pipeline is cards + a program + a workflow, never a new loop or runtime.
Both reuse the unchanged ``WanDenoiseLoop``: the I2V image-conditioning lives in the I2V program (it
folds the conditioning image into ``text_embeds``), so no loop change is needed. A new pipeline is
cards + a program + a workflow, never a new loop or runtime.
"""
from __future__ import annotations
@@ -37,45 +37,64 @@ def _diffusion_card(model_id, family, loop_id, caps, tasks, *, shift):
precision, expert = PrecisionPolicy(), NoRouting("transformer")
def loop_factory():
return WanDenoiseLoop(loop_id=loop_id, cfg=cfg, flow_shift=flow,
precision=precision, expert=expert, cost=cost)
return WanDenoiseLoop(loop_id=loop_id, cfg=cfg, flow_shift=flow, precision=precision, expert=expert, cost=cost)
components = {
"text_encoder": ComponentSpec("text_encoder", kind="text_encoder",
factory=lambda inst: ToyTextEncoder(), required_for=tasks),
"transformer": ComponentSpec("transformer", kind="dit",
factory=lambda inst: ToyDiT(seed=seed),
resident_for=[loop_id], required_for=tasks),
"vae": ComponentSpec("vae", kind="vae", factory=lambda inst: ToyVAE(), required_for=tasks),
"text_encoder":
ComponentSpec("text_encoder", kind="text_encoder", factory=lambda inst: ToyTextEncoder(), required_for=tasks),
"transformer":
ComponentSpec("transformer",
kind="dit",
factory=lambda inst: ToyDiT(seed=seed),
resident_for=[loop_id],
required_for=tasks),
"vae":
ComponentSpec("vae", kind="vae", factory=lambda inst: ToyVAE(), required_for=tasks),
}
loops = {
loop_id: LoopSpec(loop_id, kind=LoopKind.DIFFUSION_DENOISE,
work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=cost,
shared_weight_components=["transformer"], cache_policy=["feature"],
loop_factory=loop_factory),
loop_id:
LoopSpec(loop_id,
kind=LoopKind.DIFFUSION_DENOISE,
work_unit_kind=WorkUnitKind.DIFFUSION_STEP,
step_cost_model=cost,
shared_weight_components=["transformer"],
cache_policy=["feature"],
loop_factory=loop_factory),
}
return ModelCard(
model_id=model_id, family=family, components=components, loops=loops,
model_id=model_id,
family=family,
components=components,
loops=loops,
capabilities=CapabilityMatrix.of(*caps),
recipe=RecipeSpec(method="base", assumes_loop=loop_id,
assumes_precision="float32", consistency_required=ConsistencyLevel.C1),
recipe=RecipeSpec(method="base",
assumes_loop=loop_id,
assumes_precision="float32",
consistency_required=ConsistencyLevel.C1),
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True),
caches={"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True)},
caches={
"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True)
},
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()],
default_plan=ParallelPlan.single()),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
).validate()
def build_flux_t2i_card(model_id: str = "flux-t2i") -> ModelCard:
"""Stage 1 — text → image (a single-frame diffusion). Stand-in for FLUX/SD."""
return _diffusion_card(model_id, "flux", "t2i_denoise",
return _diffusion_card(model_id,
"flux",
"t2i_denoise",
caps=(Capability.TEXT_TO_IMAGE, Capability.VAE_DECODE),
tasks={"t2i"}, shift=3.0)
tasks={"t2i"},
shift=3.0)
def build_wan_i2v_card(model_id: str = "wan-i2v") -> ModelCard:
"""Stage 2 — text + conditioning image → video. Stand-in for Wan-I2V / SVD."""
return _diffusion_card(model_id, "wan", "i2v_denoise",
return _diffusion_card(model_id,
"wan",
"i2v_denoise",
caps=(Capability.IMAGE_TO_VIDEO, Capability.VAE_DECODE),
tasks={"i2v"}, shift=5.0)
tasks={"i2v"},
shift=5.0)

Some files were not shown because too many files have changed in this diff Show More