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.
80 lines
3.3 KiB
Python
80 lines
3.3 KiB
Python
"""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. ``card/`` imports no runtime; this is pure vocabulary.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from enum import Enum
|
||
|
||
|
||
class LoopKind(str, Enum):
|
||
"""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.
|
||
|
||
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"
|
||
DIFFUSION_STEP = "diffusion_step"
|
||
DIFFUSION_WINDOW = "diffusion_window"
|
||
CHUNK_STEP = "chunk_step"
|
||
ENCODER_CHUNK = "encoder_chunk"
|
||
VAE_TILE = "vae_tile"
|
||
AUDIO_CHUNK = "audio_chunk"
|
||
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. 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:
|
||
return {"C0": 0, "C1": 1, "C2": 2, "C3": 3, "C4": 4}[self.value]
|
||
|
||
|
||
class ExecutionProfile(str, Enum):
|
||
"""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 (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
|
||
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
|
||
LOGPROB_RECOMPUTE = "logprob_recompute"
|