From db866dfb04a139d29051fe233b36c9dd13fee2e5 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Fri, 19 Jun 2026 20:54:59 +0000 Subject: [PATCH] [misc] v2: simplify docstrings/comments + drop deleted-design-doc citations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- pyproject.toml | 5 +- v2/__init__.py | 36 ++- v2/_enums.py | 57 ++-- v2/_types.py | 8 +- v2/api.py | 10 +- v2/cache/__init__.py | 8 +- v2/cache/classes.py | 28 +- v2/cache/keys.py | 28 +- v2/cache/manager.py | 33 ++- v2/card/__init__.py | 29 +- v2/card/instance.py | 54 ++-- v2/card/specs.py | 139 +++++----- v2/deploy/__init__.py | 9 +- v2/deploy/card.py | 37 +-- v2/deploy/dynamo.py | 36 +-- v2/deploy/fleet.py | 30 ++- v2/examples.py | 42 +-- v2/extend/__init__.py | 8 +- v2/extend/base.py | 24 +- v2/extend/interceptors.py | 4 +- v2/extend/observers.py | 25 +- v2/extend/registry.py | 14 +- v2/fastvideo_args.py | 2 +- v2/loop/__init__.py | 26 +- v2/loop/contracts.py | 84 +++--- v2/loop/driver.py | 8 +- v2/loop/policies/__init__.py | 18 +- v2/loop/policies/base.py | 35 +-- v2/loop/policies/cfg.py | 26 +- v2/loop/sampler.py | 18 +- v2/memory/__init__.py | 2 +- v2/memory/allocator.py | 11 +- v2/models/dits/lingbotworld.py | 9 +- v2/parallel/__init__.py | 5 +- v2/parallel/mesh.py | 8 +- v2/parallel/plan.py | 12 +- v2/parallel/validation.py | 21 +- v2/parity/__init__.py | 8 +- v2/parity/aligner.py | 14 +- v2/parity/interleave_gate.py | 45 ++-- v2/parity/ladder.py | 4 +- v2/platform/__init__.py | 12 +- v2/platform/backends/accel.py | 37 +-- v2/platform/backends/cpu.py | 16 +- v2/platform/backends/torch_backend.py | 23 +- v2/platform/backends/torch_cosmos25.py | 7 +- v2/platform/backends/torch_cuda.py | 31 ++- v2/platform/backends/torch_gen3c.py | 7 +- v2/platform/backends/torch_hunyuan_video15.py | 50 ++-- v2/platform/backends/torch_kernels.py | 27 +- v2/platform/backends/torch_longcat.py | 26 +- v2/platform/backends/torch_matrixgame2.py | 4 +- v2/platform/backends/torch_matrixgame3.py | 8 +- v2/platform/backends/toy.py | 105 ++++---- v2/platform/platform.py | 57 ++-- v2/platform/registry.py | 61 +++-- v2/program/__init__.py | 9 +- v2/program/specs.py | 19 +- v2/program/workflow.py | 23 +- v2/recipes/__init__.py | 61 +++-- v2/recipes/_prompts.py | 4 +- v2/recipes/adapters/__init__.py | 2 +- v2/recipes/adapters/card.py | 81 ++++-- v2/recipes/adapters/loop.py | 21 +- v2/recipes/adaptive/__init__.py | 2 +- v2/recipes/adaptive/card.py | 64 +++-- v2/recipes/adaptive/loop.py | 39 ++- v2/recipes/bagel/card.py | 87 +++--- v2/recipes/bagel/program.py | 46 ++-- v2/recipes/common.py | 15 +- v2/recipes/cosmos2/__init__.py | 9 +- v2/recipes/cosmos2/card.py | 15 +- v2/recipes/cosmos2/loop.py | 14 +- v2/recipes/cosmos2/program.py | 7 +- v2/recipes/cosmos25/__init__.py | 14 +- v2/recipes/cosmos25/loop.py | 76 +++--- v2/recipes/cosmos25/program.py | 9 +- v2/recipes/cosmos3/card.py | 98 ++++--- v2/recipes/cosmos3/program.py | 50 ++-- v2/recipes/fastwan/card.py | 55 ++-- v2/recipes/fastwan/loop.py | 4 +- v2/recipes/fastwan/program.py | 8 +- v2/recipes/flux2/sampler.py | 7 +- v2/recipes/gen3c/loop.py | 6 +- v2/recipes/hunyuan_video/__init__.py | 15 +- v2/recipes/hunyuan_video/card.py | 8 +- v2/recipes/hunyuan_video/loop.py | 43 ++- v2/recipes/hunyuan_video/program.py | 14 +- v2/recipes/hunyuan_video15/__init__.py | 14 +- v2/recipes/hunyuan_video15/card.py | 28 +- v2/recipes/hunyuan_video15/loop.py | 39 ++- v2/recipes/hunyuan_video15/program.py | 35 ++- v2/recipes/hunyuangamecraft/__init__.py | 13 +- v2/recipes/hunyuangamecraft/card.py | 39 ++- v2/recipes/hunyuangamecraft/loop.py | 69 +++-- v2/recipes/hunyuangamecraft/program.py | 42 +-- v2/recipes/hyworld/card.py | 30 +-- v2/recipes/hyworld/loop.py | 14 +- v2/recipes/image_video/__init__.py | 12 +- v2/recipes/image_video/card.py | 77 ++++-- v2/recipes/image_video/program.py | 113 +++++--- v2/recipes/lingbotworld/__init__.py | 8 +- v2/recipes/lingbotworld/card.py | 39 ++- v2/recipes/lingbotworld/loop.py | 26 +- v2/recipes/lingbotworld/program.py | 34 +-- v2/recipes/longcat/__init__.py | 12 +- v2/recipes/longcat/card.py | 20 +- v2/recipes/longcat/loop.py | 7 +- v2/recipes/ltx2/card.py | 4 +- v2/recipes/ltx2/loop.py | 189 ++++++++----- v2/recipes/ltx2/program.py | 103 ++++--- v2/recipes/matrixgame2/__init__.py | 8 +- v2/recipes/matrixgame2/loop.py | 45 ++-- v2/recipes/matrixgame2/program.py | 16 +- v2/recipes/matrixgame2/sampler.py | 21 +- v2/recipes/matrixgame3/__init__.py | 12 +- v2/recipes/matrixgame3/card.py | 18 +- v2/recipes/multi_expert/__init__.py | 2 +- v2/recipes/multi_expert/card.py | 101 ++++--- v2/recipes/omni/__init__.py | 5 +- v2/recipes/omni/ar_loop.py | 50 ++-- v2/recipes/omni/pack.py | 15 +- v2/recipes/omni/vocoder_loop.py | 47 ++-- v2/recipes/qwen_omni/__init__.py | 2 +- v2/recipes/qwen_omni/card.py | 131 +++++---- v2/recipes/qwen_omni/program.py | 71 +++-- v2/recipes/reward/__init__.py | 2 +- v2/recipes/reward/card.py | 43 +-- v2/recipes/reward/loop.py | 34 ++- v2/recipes/sd35/__init__.py | 12 +- v2/recipes/sd35/card.py | 9 +- v2/recipes/sd35/loop.py | 6 +- v2/recipes/sd35/program.py | 12 +- v2/recipes/sfwan22/card.py | 6 +- v2/recipes/sfwan22/loop.py | 4 +- v2/recipes/sfwan22/program.py | 2 +- v2/recipes/speculative/__init__.py | 2 +- v2/recipes/speculative/card.py | 71 +++-- v2/recipes/speculative/loop.py | 88 +++--- v2/recipes/speculative/program.py | 22 +- v2/recipes/stable_audio/card.py | 28 +- v2/recipes/stable_audio/loop.py | 25 +- v2/recipes/stable_audio/program.py | 18 +- v2/recipes/stable_audio/sampler.py | 2 +- v2/recipes/tiled/__init__.py | 2 +- v2/recipes/tiled/card.py | 80 ++++-- v2/recipes/tiled/loop.py | 48 ++-- v2/recipes/tiled/program.py | 24 +- v2/recipes/turbowan/card.py | 40 ++- v2/recipes/turbowan/loop.py | 33 ++- v2/recipes/turbowan/program.py | 6 +- v2/recipes/turbowan/sampler.py | 37 ++- v2/recipes/unified/__init__.py | 4 +- v2/recipes/unified/card.py | 110 +++++--- v2/recipes/unified/program.py | 45 ++-- v2/recipes/wan21/card.py | 254 +++++++++++------- v2/recipes/wan21/loop.py | 156 +++++++---- v2/recipes/wan21/program.py | 24 +- v2/recipes/wan_causal/card.py | 10 +- v2/recipes/wan_causal/loop.py | 80 ++++-- v2/recipes/wan_causal/program.py | 19 +- v2/recipes/wan_fun_control/__init__.py | 13 +- v2/recipes/wan_fun_control/loop.py | 25 +- v2/recipes/wan_fun_control/program.py | 53 ++-- v2/registry.py | 8 +- v2/request/__init__.py | 40 ++- v2/request/artifacts.py | 20 +- v2/request/cancel.py | 9 +- v2/request/modalpart.py | 12 +- v2/request/params.py | 20 +- v2/request/requests.py | 27 +- v2/request/streams.py | 24 +- v2/request/tasks.py | 26 +- v2/runtime/__init__.py | 11 +- v2/runtime/async_engine.py | 45 ++-- v2/runtime/context.py | 24 +- v2/runtime/cudagraph.py | 38 ++- v2/runtime/disaggregated.py | 63 +++-- v2/runtime/engine.py | 77 +++--- v2/runtime/pools.py | 52 ++-- v2/runtime/scheduler.py | 14 +- v2/runtime/session.py | 52 ++-- v2/serving/http.py | 48 ++-- v2/serving/protocol.py | 62 +++-- v2/serving/server.py | 144 +++++++--- v2/tests/test_adapters.py | 9 +- v2/tests/test_adaptive_compute.py | 2 +- v2/tests/test_admission.py | 4 +- v2/tests/test_cache.py | 2 +- v2/tests/test_contracts.py | 2 +- v2/tests/test_cudagraph.py | 2 +- v2/tests/test_extensions.py | 6 +- v2/tests/test_flywheel.py | 9 +- v2/tests/test_inference.py | 2 +- v2/tests/test_interleave_gate.py | 2 +- v2/tests/test_joint_multi_rl.py | 4 +- v2/tests/test_ltx2_av.py | 2 +- v2/tests/test_nested_workflow.py | 2 +- v2/tests/test_omni.py | 6 +- v2/tests/test_parity.py | 2 +- v2/tests/test_platform.py | 15 +- v2/tests/test_policies.py | 2 +- v2/tests/test_review_fixes.py | 4 +- v2/tests/test_served_reward.py | 6 +- v2/tests/test_speculative.py | 2 +- v2/tests/test_t2i_i2v_workflow.py | 2 +- v2/tests/test_thinker_talker.py | 2 +- v2/tests/test_tiled_scheduling.py | 8 +- v2/tests/test_torch_backend.py | 2 +- v2/tests/test_training.py | 4 +- v2/tests/test_unified_rl.py | 6 +- v2/tests/test_weight_sync_live.py | 4 +- v2/tests/test_workflow_rl.py | 2 +- v2/tests/test_workflow_shapes.py | 4 +- v2/tests/test_world_session.py | 8 +- v2/training/__init__.py | 40 ++- v2/training/behavior.py | 12 +- v2/training/flywheel.py | 42 +-- v2/training/methods/__init__.py | 27 +- v2/training/methods/base.py | 14 +- v2/training/methods/diffusion_nft.py | 70 +++-- v2/training/methods/dmd2.py | 33 ++- v2/training/methods/finetune.py | 18 +- v2/training/methods/joint_multi_rl.py | 128 +++++---- v2/training/methods/self_forcing.py | 26 +- v2/training/methods/unified_rl.py | 122 +++++---- v2/training/methods/workflow_rl.py | 161 ++++++----- v2/training/rewards.py | 21 +- v2/training/rollout.py | 28 +- v2/training/weight_sync.py | 59 ++-- v2/transport/__init__.py | 8 +- v2/transport/connector.py | 63 +++-- v2/transport/manifest.py | 8 +- v2/video_generator.py | 29 +- 234 files changed, 4148 insertions(+), 3080 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 4aba0d350..e06d5ca84 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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. diff --git a/v2/__init__.py b/v2/__init__.py index c13d356dd..410506ce5 100644 --- a/v2/__init__.py +++ b/v2/__init__.py @@ -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__", ] diff --git a/v2/_enums.py b/v2/_enums.py index 6d7f3c669..82e938747 100644 --- a/v2/_enums.py +++ b/v2/_enums.py @@ -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" diff --git a/v2/_types.py b/v2/_types.py index 3573d9eee..da9fb8be9 100644 --- a/v2/_types.py +++ b/v2/_types.py @@ -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 diff --git a/v2/api.py b/v2/api.py index e0672eb0b..084e7f941 100644 --- a/v2/api.py +++ b/v2/api.py @@ -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, ) diff --git a/v2/cache/__init__.py b/v2/cache/__init__.py index c83f13db8..a4de163e7 100644 --- a/v2/cache/__init__.py +++ b/v2/cache/__init__.py @@ -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" +] diff --git a/v2/cache/classes.py b/v2/cache/classes.py index c4bbaae32..ccadfb3ba 100644 --- a/v2/cache/classes.py +++ b/v2/cache/classes.py @@ -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): diff --git a/v2/cache/keys.py b/v2/cache/keys.py index f69964f23..083d6a7d5 100644 --- a/v2/cache/keys.py +++ b/v2/cache/keys.py @@ -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 diff --git a/v2/cache/manager.py b/v2/cache/manager.py index 783334cf7..3d20db96e 100644 --- a/v2/cache/manager.py +++ b/v2/cache/manager.py @@ -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() + } diff --git a/v2/card/__init__.py b/v2/card/__init__.py index 78d388ec2..c9993e3f6 100644 --- a/v2/card/__init__.py +++ b/v2/card/__init__.py @@ -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", ] diff --git a/v2/card/instance.py b/v2/card/instance.py index afec31da3..e948eb00f 100644 --- a/v2/card/instance.py +++ b/v2/card/instance.py @@ -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()) diff --git a/v2/card/specs.py b/v2/card/specs.py index 924a3a751..0797d8146 100644 --- a/v2/card/specs.py +++ b/v2/card/specs.py @@ -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]: diff --git a/v2/deploy/__init__.py b/v2/deploy/__init__.py index 39d150edd..11e1c64cb 100644 --- a/v2/deploy/__init__.py +++ b/v2/deploy/__init__.py @@ -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" +] diff --git a/v2/deploy/card.py b/v2/deploy/card.py index 985676f05..43257f8e7 100644 --- a/v2/deploy/card.py +++ b/v2/deploy/card.py @@ -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)) diff --git a/v2/deploy/dynamo.py b/v2/deploy/dynamo.py index 7380ab039..fdfce9764 100644 --- a/v2/deploy/dynamo.py +++ b/v2/deploy/dynamo.py @@ -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} diff --git a/v2/deploy/fleet.py b/v2/deploy/fleet.py index 72636fa58..ba3c31ea8 100644 --- a/v2/deploy/fleet.py +++ b/v2/deploy/fleet.py @@ -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: diff --git a/v2/examples.py b/v2/examples.py index a1822efcd..2c23d20a4 100644 --- a/v2/examples.py +++ b/v2/examples.py @@ -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)}") diff --git a/v2/extend/__init__.py b/v2/extend/__init__.py index c70f2e444..489dfa68a 100644 --- a/v2/extend/__init__.py +++ b/v2/extend/__init__.py @@ -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" +] diff --git a/v2/extend/base.py b/v2/extend/base.py index 13981de48..8de66fa28 100644 --- a/v2/extend/base.py +++ b/v2/extend/base.py @@ -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]}") diff --git a/v2/extend/interceptors.py b/v2/extend/interceptors.py index de7797c0f..8ddb8f540 100644 --- a/v2/extend/interceptors.py +++ b/v2/extend/interceptors.py @@ -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 diff --git a/v2/extend/observers.py b/v2/extend/observers.py index 0fa45a00c..520c223df 100644 --- a/v2/extend/observers.py +++ b/v2/extend/observers.py @@ -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: diff --git a/v2/extend/registry.py b/v2/extend/registry.py index 9f8f680bf..e0e097059 100644 --- a/v2/extend/registry.py +++ b/v2/extend/registry.py @@ -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 diff --git a/v2/fastvideo_args.py b/v2/fastvideo_args.py index cdb71911e..e61d4d790 100644 --- a/v2/fastvideo_args.py +++ b/v2/fastvideo_args.py @@ -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 diff --git a/v2/loop/__init__.py b/v2/loop/__init__.py index 8bd50b2d7..0ee4d1f6c 100644 --- a/v2/loop/__init__.py +++ b/v2/loop/__init__.py @@ -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", ] diff --git a/v2/loop/contracts.py b/v2/loop/contracts.py index 8456883e8..5676c53f3 100644 --- a/v2/loop/contracts.py +++ b/v2/loop/contracts.py @@ -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: + ... diff --git a/v2/loop/driver.py b/v2/loop/driver.py index 1486839b8..11cdc3e96 100644 --- a/v2/loop/driver.py +++ b/v2/loop/driver.py @@ -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 diff --git a/v2/loop/policies/__init__.py b/v2/loop/policies/__init__.py index 7f2d9e4ab..8cdc6b9e8 100644 --- a/v2/loop/policies/__init__.py +++ b/v2/loop/policies/__init__.py @@ -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", ] diff --git a/v2/loop/policies/base.py b/v2/loop/policies/base.py index 378e01389..9437a2026 100644 --- a/v2/loop/policies/base.py +++ b/v2/loop/policies/base.py @@ -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): diff --git a/v2/loop/policies/cfg.py b/v2/loop/policies/cfg.py index 60df4e258..d67efa386 100644 --- a/v2/loop/policies/cfg.py +++ b/v2/loop/policies/cfg.py @@ -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 diff --git a/v2/loop/sampler.py b/v2/loop/sampler.py index d1b1b5d17..285bfafb5 100644 --- a/v2/loop/sampler.py +++ b/v2/loop/sampler.py @@ -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) diff --git a/v2/memory/__init__.py b/v2/memory/__init__.py index 4aad30e07..85e9a52da 100644 --- a/v2/memory/__init__.py +++ b/v2/memory/__init__.py @@ -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 diff --git a/v2/memory/allocator.py b/v2/memory/allocator.py index 80f1cdd9f..b3828b8b8 100644 --- a/v2/memory/allocator.py +++ b/v2/memory/allocator.py @@ -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 {}) diff --git a/v2/models/dits/lingbotworld.py b/v2/models/dits/lingbotworld.py index 46385745b..78a1de829 100644 --- a/v2/models/dits/lingbotworld.py +++ b/v2/models/dits/lingbotworld.py @@ -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, ) diff --git a/v2/parallel/__init__.py b/v2/parallel/__init__.py index 9e30c6f23..1b3d35379 100644 --- a/v2/parallel/__init__.py +++ b/v2/parallel/__init__.py @@ -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"] diff --git a/v2/parallel/mesh.py b/v2/parallel/mesh.py index d41df249d..72f853d39 100644 --- a/v2/parallel/mesh.py +++ b/v2/parallel/mesh.py @@ -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 diff --git a/v2/parallel/plan.py b/v2/parallel/plan.py index 16dae9593..3cf5edcea 100644 --- a/v2/parallel/plan.py +++ b/v2/parallel/plan.py @@ -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"]) diff --git a/v2/parallel/validation.py b/v2/parallel/validation.py index 083435943..0b2484f23 100644 --- a/v2/parallel/validation.py +++ b/v2/parallel/validation.py @@ -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] diff --git a/v2/parity/__init__.py b/v2/parity/__init__.py index 284fd4057..43f0bd6da 100644 --- a/v2/parity/__init__.py +++ b/v2/parity/__init__.py @@ -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" +] diff --git a/v2/parity/aligner.py b/v2/parity/aligner.py index 28d9cd4d6..ac1b937d7 100644 --- a/v2/parity/aligner.py +++ b/v2/parity/aligner.py @@ -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 diff --git a/v2/parity/interleave_gate.py b/v2/parity/interleave_gate.py index dd66cfdf5..33c84e272 100644 --- a/v2/parity/interleave_gate.py +++ b/v2/parity/interleave_gate.py @@ -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 diff --git a/v2/parity/ladder.py b/v2/parity/ladder.py index b82ef3a20..efe5255ed 100644 --- a/v2/parity/ladder.py +++ b/v2/parity/ladder.py @@ -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 diff --git a/v2/platform/__init__.py b/v2/platform/__init__.py index 18adfa90a..4c431fafc 100644 --- a/v2/platform/__init__.py +++ b/v2/platform/__init__.py @@ -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 diff --git a/v2/platform/backends/accel.py b/v2/platform/backends/accel.py index d754a2e82..10befde1e 100644 --- a/v2/platform/backends/accel.py +++ b/v2/platform/backends/accel.py @@ -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") diff --git a/v2/platform/backends/cpu.py b/v2/platform/backends/cpu.py index 9400bed98..0199fd2e8 100644 --- a/v2/platform/backends/cpu.py +++ b/v2/platform/backends/cpu.py @@ -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") diff --git a/v2/platform/backends/torch_backend.py b/v2/platform/backends/torch_backend.py index 22c4bf9d3..6aa2563a1 100644 --- a/v2/platform/backends/torch_backend.py +++ b/v2/platform/backends/torch_backend.py @@ -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()) diff --git a/v2/platform/backends/torch_cosmos25.py b/v2/platform/backends/torch_cosmos25.py index de87563f2..ffaa2069d 100644 --- a/v2/platform/backends/torch_cosmos25.py +++ b/v2/platform/backends/torch_cosmos25.py @@ -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 diff --git a/v2/platform/backends/torch_cuda.py b/v2/platform/backends/torch_cuda.py index 5492fc33c..e5b476d1d 100644 --- a/v2/platform/backends/torch_cuda.py +++ b/v2/platform/backends/torch_cuda.py @@ -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)") diff --git a/v2/platform/backends/torch_gen3c.py b/v2/platform/backends/torch_gen3c.py index 2f385deae..66452a828 100644 --- a/v2/platform/backends/torch_gen3c.py +++ b/v2/platform/backends/torch_gen3c.py @@ -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. diff --git a/v2/platform/backends/torch_hunyuan_video15.py b/v2/platform/backends/torch_hunyuan_video15.py index 6a75a0c90..e4f533fcb 100644 --- a/v2/platform/backends/torch_hunyuan_video15.py +++ b/v2/platform/backends/torch_hunyuan_video15.py @@ -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 diff --git a/v2/platform/backends/torch_kernels.py b/v2/platform/backends/torch_kernels.py index 60f4ef3b6..4371735e3 100644 --- a/v2/platform/backends/torch_kernels.py +++ b/v2/platform/backends/torch_kernels.py @@ -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)) diff --git a/v2/platform/backends/torch_longcat.py b/v2/platform/backends/torch_longcat.py index a8229ed47..f1d16870b 100644 --- a/v2/platform/backends/torch_longcat.py +++ b/v2/platform/backends/torch_longcat.py @@ -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) diff --git a/v2/platform/backends/torch_matrixgame2.py b/v2/platform/backends/torch_matrixgame2.py index 8980807b4..a1f50b498 100644 --- a/v2/platform/backends/torch_matrixgame2.py +++ b/v2/platform/backends/torch_matrixgame2.py @@ -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 diff --git a/v2/platform/backends/torch_matrixgame3.py b/v2/platform/backends/torch_matrixgame3.py index 12c0e0464..ed9dbcfee 100644 --- a/v2/platform/backends/torch_matrixgame3.py +++ b/v2/platform/backends/torch_matrixgame3.py @@ -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)) diff --git a/v2/platform/backends/toy.py b/v2/platform/backends/toy.py index 730776068..91aca30e7 100644 --- a/v2/platform/backends/toy.py +++ b/v2/platform/backends/toy.py @@ -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") diff --git a/v2/platform/platform.py b/v2/platform/platform.py index f574cf574..8bd507ef1 100644 --- a/v2/platform/platform.py +++ b/v2/platform/platform.py @@ -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 diff --git a/v2/platform/registry.py b/v2/platform/registry.py index d551ba0d3..e71a9d0cf 100644 --- a/v2/platform/registry.py +++ b/v2/platform/registry.py @@ -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)}) diff --git a/v2/program/__init__.py b/v2/program/__init__.py index dc976567c..0876f5832 100644 --- a/v2/program/__init__.py +++ b/v2/program/__init__.py @@ -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" +] diff --git a/v2/program/specs.py b/v2/program/specs.py index 51c6cd1a3..cf3a34d29 100644 --- a/v2/program/specs.py +++ b/v2/program/specs.py @@ -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: diff --git a/v2/program/workflow.py b/v2/program/workflow.py index 7133ef15b..a8b1e0cfe 100644 --- a/v2/program/workflow.py +++ b/v2/program/workflow.py @@ -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 diff --git a/v2/recipes/__init__.py b/v2/recipes/__init__.py index de36a3bbd..d2b4d199f 100644 --- a/v2/recipes/__init__.py +++ b/v2/recipes/__init__.py @@ -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 diff --git a/v2/recipes/_prompts.py b/v2/recipes/_prompts.py index ec1b360ad..6e071f238 100644 --- a/v2/recipes/_prompts.py +++ b/v2/recipes/_prompts.py @@ -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. diff --git a/v2/recipes/adapters/__init__.py b/v2/recipes/adapters/__init__.py index 7f454b640..e43e88469 100644 --- a/v2/recipes/adapters/__init__.py +++ b/v2/recipes/adapters/__init__.py @@ -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 diff --git a/v2/recipes/adapters/card.py b/v2/recipes/adapters/card.py index 749553cb5..dc3182ca2 100644 --- a/v2/recipes/adapters/card.py +++ b/v2/recipes/adapters/card.py @@ -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() diff --git a/v2/recipes/adapters/loop.py b/v2/recipes/adapters/loop.py index 144fff27e..7ed3d1010 100644 --- a/v2/recipes/adapters/loop.py +++ b/v2/recipes/adapters/loop.py @@ -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 diff --git a/v2/recipes/adaptive/__init__.py b/v2/recipes/adaptive/__init__.py index 4f242075e..dfab584a4 100644 --- a/v2/recipes/adaptive/__init__.py +++ b/v2/recipes/adaptive/__init__.py @@ -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 diff --git a/v2/recipes/adaptive/card.py b/v2/recipes/adaptive/card.py index 9d37a668f..1c4ca2411 100644 --- a/v2/recipes/adaptive/card.py +++ b/v2/recipes/adaptive/card.py @@ -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() diff --git a/v2/recipes/adaptive/loop.py b/v2/recipes/adaptive/loop.py index 7e05f2a99..e57fab29c 100644 --- a/v2/recipes/adaptive/loop.py +++ b/v2/recipes/adaptive/loop.py @@ -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): diff --git a/v2/recipes/bagel/card.py b/v2/recipes/bagel/card.py index ea22c942c..6307c76d4 100644 --- a/v2/recipes/bagel/card.py +++ b/v2/recipes/bagel/card.py @@ -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() diff --git a/v2/recipes/bagel/program.py b/v2/recipes/bagel/program.py index fa00f3a6c..b0fabf826 100644 --- a/v2/recipes/bagel/program.py +++ b/v2/recipes/bagel/program.py @@ -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() diff --git a/v2/recipes/common.py b/v2/recipes/common.py index 8a29b0c22..5ec259a91 100644 --- a/v2/recipes/common.py +++ b/v2/recipes/common.py @@ -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 diff --git a/v2/recipes/cosmos2/__init__.py b/v2/recipes/cosmos2/__init__.py index 846303504..9e65d6f0a 100644 --- a/v2/recipes/cosmos2/__init__.py +++ b/v2/recipes/cosmos2/__init__.py @@ -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 diff --git a/v2/recipes/cosmos2/card.py b/v2/recipes/cosmos2/card.py index 1b4dc6ade..d15d1c6a8 100644 --- a/v2/recipes/cosmos2/card.py +++ b/v2/recipes/cosmos2/card.py @@ -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). diff --git a/v2/recipes/cosmos2/loop.py b/v2/recipes/cosmos2/loop.py index 1f9551d76..340caef1f 100644 --- a/v2/recipes/cosmos2/loop.py +++ b/v2/recipes/cosmos2/loop.py @@ -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 diff --git a/v2/recipes/cosmos2/program.py b/v2/recipes/cosmos2/program.py index 139390efa..d47b7ef74 100644 --- a/v2/recipes/cosmos2/program.py +++ b/v2/recipes/cosmos2/program.py @@ -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 diff --git a/v2/recipes/cosmos25/__init__.py b/v2/recipes/cosmos25/__init__.py index 4fd4ccacc..2426033d0 100644 --- a/v2/recipes/cosmos25/__init__.py +++ b/v2/recipes/cosmos25/__init__.py @@ -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 diff --git a/v2/recipes/cosmos25/loop.py b/v2/recipes/cosmos25/loop.py index 9883db82a..42ee8caa2 100644 --- a/v2/recipes/cosmos25/loop.py +++ b/v2/recipes/cosmos25/loop.py @@ -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, diff --git a/v2/recipes/cosmos25/program.py b/v2/recipes/cosmos25/program.py index c8b956cc1..ee76f6f64 100644 --- a/v2/recipes/cosmos25/program.py +++ b/v2/recipes/cosmos25/program.py @@ -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 diff --git a/v2/recipes/cosmos3/card.py b/v2/recipes/cosmos3/card.py index 7d462de3a..f12813be0 100644 --- a/v2/recipes/cosmos3/card.py +++ b/v2/recipes/cosmos3/card.py @@ -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() diff --git a/v2/recipes/cosmos3/program.py b/v2/recipes/cosmos3/program.py index c57a74346..ed5c373c9 100644 --- a/v2/recipes/cosmos3/program.py +++ b/v2/recipes/cosmos3/program.py @@ -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() diff --git a/v2/recipes/fastwan/card.py b/v2/recipes/fastwan/card.py index 719144aa0..309ed25a0 100644 --- a/v2/recipes/fastwan/card.py +++ b/v2/recipes/fastwan/card.py @@ -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, diff --git a/v2/recipes/fastwan/loop.py b/v2/recipes/fastwan/loop.py index 36e1c672a..0c8824fe8 100644 --- a/v2/recipes/fastwan/loop.py +++ b/v2/recipes/fastwan/loop.py @@ -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]): diff --git a/v2/recipes/fastwan/program.py b/v2/recipes/fastwan/program.py index eff890e93..9a0bb086a 100644 --- a/v2/recipes/fastwan/program.py +++ b/v2/recipes/fastwan/program.py @@ -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 diff --git a/v2/recipes/flux2/sampler.py b/v2/recipes/flux2/sampler.py index 9c5375b0c..7bc32cf70 100644 --- a/v2/recipes/flux2/sampler.py +++ b/v2/recipes/flux2/sampler.py @@ -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) diff --git a/v2/recipes/gen3c/loop.py b/v2/recipes/gen3c/loop.py index fef800843..bc4427f7d 100644 --- a/v2/recipes/gen3c/loop.py +++ b/v2/recipes/gen3c/loop.py @@ -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 diff --git a/v2/recipes/hunyuan_video/__init__.py b/v2/recipes/hunyuan_video/__init__.py index 9f7744c1f..b85794aba 100644 --- a/v2/recipes/hunyuan_video/__init__.py +++ b/v2/recipes/hunyuan_video/__init__.py @@ -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 diff --git a/v2/recipes/hunyuan_video/card.py b/v2/recipes/hunyuan_video/card.py index 75d599195..8c263b559 100644 --- a/v2/recipes/hunyuan_video/card.py +++ b/v2/recipes/hunyuan_video/card.py @@ -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 diff --git a/v2/recipes/hunyuan_video/loop.py b/v2/recipes/hunyuan_video/loop.py index 41067794d..7840544ab 100644 --- a/v2/recipes/hunyuan_video/loop.py +++ b/v2/recipes/hunyuan_video/loop.py @@ -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 diff --git a/v2/recipes/hunyuan_video/program.py b/v2/recipes/hunyuan_video/program.py index 4b8b00ba7..0b5aaa61b 100644 --- a/v2/recipes/hunyuan_video/program.py +++ b/v2/recipes/hunyuan_video/program.py @@ -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: diff --git a/v2/recipes/hunyuan_video15/__init__.py b/v2/recipes/hunyuan_video15/__init__.py index 787f7b392..21512d916 100644 --- a/v2/recipes/hunyuan_video15/__init__.py +++ b/v2/recipes/hunyuan_video15/__init__.py @@ -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 diff --git a/v2/recipes/hunyuan_video15/card.py b/v2/recipes/hunyuan_video15/card.py index c22eedfdd..0e987cf5f 100644 --- a/v2/recipes/hunyuan_video15/card.py +++ b/v2/recipes/hunyuan_video15/card.py @@ -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)) diff --git a/v2/recipes/hunyuan_video15/loop.py b/v2/recipes/hunyuan_video15/loop.py index fba740d6a..9e174d5b1 100644 --- a/v2/recipes/hunyuan_video15/loop.py +++ b/v2/recipes/hunyuan_video15/loop.py @@ -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, *, diff --git a/v2/recipes/hunyuan_video15/program.py b/v2/recipes/hunyuan_video15/program.py index d8aa52fe1..535a60c07 100644 --- a/v2/recipes/hunyuan_video15/program.py +++ b/v2/recipes/hunyuan_video15/program.py @@ -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 diff --git a/v2/recipes/hunyuangamecraft/__init__.py b/v2/recipes/hunyuangamecraft/__init__.py index 15e26e905..9e020bb70 100644 --- a/v2/recipes/hunyuangamecraft/__init__.py +++ b/v2/recipes/hunyuangamecraft/__init__.py @@ -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 diff --git a/v2/recipes/hunyuangamecraft/card.py b/v2/recipes/hunyuangamecraft/card.py index 1128fa4ef..852a5b4bc 100644 --- a/v2/recipes/hunyuangamecraft/card.py +++ b/v2/recipes/hunyuangamecraft/card.py @@ -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 diff --git a/v2/recipes/hunyuangamecraft/loop.py b/v2/recipes/hunyuangamecraft/loop.py index 927a07077..671858d68 100644 --- a/v2/recipes/hunyuangamecraft/loop.py +++ b/v2/recipes/hunyuangamecraft/loop.py @@ -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: diff --git a/v2/recipes/hunyuangamecraft/program.py b/v2/recipes/hunyuangamecraft/program.py index 833da5586..77ea75a57 100644 --- a/v2/recipes/hunyuangamecraft/program.py +++ b/v2/recipes/hunyuangamecraft/program.py @@ -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] diff --git a/v2/recipes/hyworld/card.py b/v2/recipes/hyworld/card.py index bdc3d84de..3611a7287 100644 --- a/v2/recipes/hyworld/card.py +++ b/v2/recipes/hyworld/card.py @@ -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``. """ diff --git a/v2/recipes/hyworld/loop.py b/v2/recipes/hyworld/loop.py index 173e42cb2..e62b46886 100644 --- a/v2/recipes/hyworld/loop.py +++ b/v2/recipes/hyworld/loop.py @@ -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 diff --git a/v2/recipes/image_video/__init__.py b/v2/recipes/image_video/__init__.py index 0fdf42c78..35148e682 100644 --- a/v2/recipes/image_video/__init__.py +++ b/v2/recipes/image_video/__init__.py @@ -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" +] diff --git a/v2/recipes/image_video/card.py b/v2/recipes/image_video/card.py index eafd07624..e2d324f43 100644 --- a/v2/recipes/image_video/card.py +++ b/v2/recipes/image_video/card.py @@ -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) diff --git a/v2/recipes/image_video/program.py b/v2/recipes/image_video/program.py index 70ca71908..8399b0ff6 100644 --- a/v2/recipes/image_video/program.py +++ b/v2/recipes/image_video/program.py @@ -1,4 +1,4 @@ -"""T2I and I2V programs + the cross-model workflow that chains them (design_v3 §13, §15). +"""T2I and I2V programs + the cross-model workflow that chains them. Workflow "t2i_then_i2v": stage 1 (flux-t2i): text_encode → t2i_denoise → vae_decode → image @@ -28,14 +28,21 @@ def _t2i_vae_decode(instance, slots, request, ctx) -> None: def build_flux_t2i_program() -> Program: return Program( - program_id="flux.t2i", kind=ProgramKind.INLINE, + program_id="flux.t2i", + kind=ProgramKind.INLINE, nodes=[ ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")), - ModelLoopNode("denoise", loop_id="t2i_denoise", output_slot="t2i_out", - reads=("text_embeds",), writes=("t2i_out",)), - ComponentNode("vae_decode", fn=_t2i_vae_decode, reads=("t2i_out",), writes=("image",)), + ModelLoopNode("denoise", + loop_id="t2i_denoise", + output_slot="t2i_out", + reads=("text_embeds", ), + writes=("t2i_out", )), + ComponentNode("vae_decode", fn=_t2i_vae_decode, reads=("t2i_out", ), writes=("image", )), ], - output_artifacts={"image": "image", "latents": "t2i_out"}, + output_artifacts={ + "image": "image", + "latents": "t2i_out" + }, ).validate() @@ -55,7 +62,7 @@ def _condition_on_image(instance, slots, request, ctx) -> None: cond = slots.get("cond_latent") if cond is None: return - shift = float(np.tanh(np.mean(np.asarray(cond, dtype="float64")))) # image-derived conditioning + shift = float(np.tanh(np.mean(np.asarray(cond, dtype="float64")))) # image-derived conditioning te = np.asarray(slots["text_embeds"], dtype="float32") slots["text_embeds"] = (te + shift).astype("float32") slots["image_shift"] = shift @@ -67,57 +74,81 @@ def _i2v_vae_decode(instance, slots, request, ctx) -> None: def build_wan_i2v_program() -> Program: return Program( - program_id="wan.i2v", kind=ProgramKind.INLINE, + program_id="wan.i2v", + kind=ProgramKind.INLINE, nodes=[ - ComponentNode("encode_cond_image", fn=_encode_cond_image, writes=("cond_latent",)), + ComponentNode("encode_cond_image", fn=_encode_cond_image, writes=("cond_latent", )), ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")), - ComponentNode("condition_on_image", fn=_condition_on_image, - reads=("cond_latent", "text_embeds"), writes=("text_embeds", "image_shift")), - ModelLoopNode("denoise", loop_id="i2v_denoise", output_slot="i2v_out", - reads=("text_embeds",), writes=("i2v_out",)), - ComponentNode("vae_decode", fn=_i2v_vae_decode, reads=("i2v_out",), writes=("video",)), + ComponentNode("condition_on_image", + fn=_condition_on_image, + reads=("cond_latent", "text_embeds"), + writes=("text_embeds", "image_shift")), + ModelLoopNode("denoise", + loop_id="i2v_denoise", + output_slot="i2v_out", + reads=("text_embeds", ), + writes=("i2v_out", )), + ComponentNode("vae_decode", fn=_i2v_vae_decode, reads=("i2v_out", ), writes=("video", )), ], - output_artifacts={"video": "video", "latents": "i2v_out"}, + output_artifacts={ + "video": "video", + "latents": "i2v_out" + }, ).validate() # --- the cross-model workflow ----------------------------------------------------- # def build_t2i_then_i2v_workflow(t2i_id: str = "flux-t2i", i2v_id: str = "wan-i2v") -> Workflow: + def t2i_stage(state): - return make_request(TaskType.T2I, t2i_id, state["prompt"], - diffusion=DiffusionParams(num_steps=4, num_frames=1, - seed=state.get("seed", 0))) + return make_request(TaskType.T2I, + t2i_id, + state["prompt"], + diffusion=DiffusionParams(num_steps=4, num_frames=1, seed=state.get("seed", 0))) def i2v_stage(state): - image = state["t2i:image"].tensor # the decoded image from stage 1 - return make_request(TaskType.I2V, i2v_id, state["prompt"], - inputs=(TextPart(state["prompt"]), ImagePart(pixels=image)), - diffusion=DiffusionParams(num_steps=4, num_frames=81, # → multi-frame video - seed=state.get("seed", 0))) + image = state["t2i:image"].tensor # the decoded image from stage 1 + return make_request( + TaskType.I2V, + i2v_id, + state["prompt"], + inputs=(TextPart(state["prompt"]), ImagePart(pixels=image)), + diffusion=DiffusionParams( + num_steps=4, + num_frames=81, # → multi-frame video + seed=state.get("seed", 0))) - return Workflow("image_video.t2i_i2v", [ # namespaced id: . (a servable) - WorkflowStage(t2i_id, t2i_stage, label="t2i"), - WorkflowStage(i2v_id, i2v_stage, label="i2v"), - ]) + return Workflow( + "image_video.t2i_i2v", + [ # namespaced id: . (a servable) + WorkflowStage(t2i_id, t2i_stage, label="t2i"), + WorkflowStage(i2v_id, i2v_stage, label="i2v"), + ]) def build_t2i_i2v_extend_workflow(inner_id: str = "image_video.t2i_i2v", i2v_id: str = "wan-i2v") -> Workflow: - """A NESTED workflow (design_v3 §13; §9.13): its first stage is *another workflow* (the T2I→I2V - pipeline), whose video is then extended by an I2V pass. Proves composition is recursive — a - workflow id is a servable, so a stage can invoke a workflow exactly as it invokes a model.""" - def shot1(state): # stage 1 IS a workflow (engine.run routes it) - return make_request(TaskType.T2V, inner_id, state["prompt"], + """A NESTED workflow: its first stage is *another workflow* (the T2I→I2V pipeline), whose video + is then extended by an I2V pass. Composition is recursive — a workflow id is a servable, so a + stage can invoke a workflow exactly as it invokes a model.""" + + def shot1(state): # stage 1 IS a workflow (engine.run routes it) + return make_request(TaskType.T2V, + inner_id, + state["prompt"], diffusion=DiffusionParams(num_steps=4, num_frames=1, seed=state.get("seed", 0))) def shot2(state): - video = np.asarray(state["shot1:video"].frames) # the inner workflow's output - first_frame = video[:, :1] # extend from its first frame - return make_request(TaskType.I2V, i2v_id, state["prompt"], + video = np.asarray(state["shot1:video"].frames) # the inner workflow's output + first_frame = video[:, :1] # extend from its first frame + return make_request(TaskType.I2V, + i2v_id, + state["prompt"], inputs=(TextPart(state["prompt"]), ImagePart(pixels=first_frame)), - diffusion=DiffusionParams(num_steps=4, num_frames=81, - seed=state.get("seed", 0) + 1)) + diffusion=DiffusionParams(num_steps=4, num_frames=81, seed=state.get("seed", 0) + 1)) - return Workflow("image_video.t2i_i2v_extend", [ - WorkflowStage(inner_id, shot1, label="shot1"), # ← a workflow stage - WorkflowStage(i2v_id, shot2, label="shot2"), - ]) + return Workflow( + "image_video.t2i_i2v_extend", + [ + WorkflowStage(inner_id, shot1, label="shot1"), # ← a workflow stage + WorkflowStage(i2v_id, shot2, label="shot2"), + ]) diff --git a/v2/recipes/lingbotworld/__init__.py b/v2/recipes/lingbotworld/__init__.py index c21a97846..ce61dbcf7 100644 --- a/v2/recipes/lingbotworld/__init__.py +++ b/v2/recipes/lingbotworld/__init__.py @@ -1,9 +1,9 @@ """LingBot-World-Base-Cam (camera-conditioned dual-guidance MoE i2v) ported into the v2 substrate. -Self-contained recipe package (the bucket-C pattern): the card declares its torch adapter via -``ComponentSpec.adapter`` (``LingBotWorldDiT`` in ``v2/platform/backends/torch_lingbotworld.py``) and a -new ``LingBotWorldDenoiseLoop`` (Wan i2v flow-match + boundary-routed MoE + DUAL guidance + camera/Plucker -conditioning), reusing the Wan VAE / T5 / CLIP adapters + ``stamp_wan21_checkpoints``. The program adds a +Self-contained recipe package: the card declares its torch adapter via ``ComponentSpec.adapter`` +(``LingBotWorldDiT`` in ``v2/platform/backends/torch_lingbotworld.py``) and a new +``LingBotWorldDenoiseLoop`` (Wan i2v flow-match + boundary-routed MoE + dual guidance + camera/Plucker +conditioning), reusing the Wan VAE / T5 adapters + ``stamp_wan21_checkpoints``. The program adds a ``camera_encode`` node (Plucker tensor; BRINGUP — needs a per-request camera input). Registered in ``v2/registry.py`` by the orchestrator from the returned metadata. """ diff --git a/v2/recipes/lingbotworld/card.py b/v2/recipes/lingbotworld/card.py index 3abbb8a11..fd261d4bd 100644 --- a/v2/recipes/lingbotworld/card.py +++ b/v2/recipes/lingbotworld/card.py @@ -1,28 +1,27 @@ """LingBot-World-Base-Cam ModelCard — camera-conditioned, dual-guidance MoE image-to-video. -LingBot-World-Base-Cam (``FastVideo/LingBot-World-Base-Cam-Diffusers``) is Wan2.2-I2V-A14B with a NEW -camera/Plucker conditioning input. Architecture (all declared on the card so the recipe is self-contained -— no edit to the shared ``_make_dit`` dispatch): +LingBot-World-Base-Cam (``FastVideo/LingBot-World-Base-Cam-Diffusers``) is Wan2.2-I2V-A14B plus a +camera/Plucker conditioning input. Everything is declared on the card so the recipe is self-contained: * DiT (x2) ``fastvideo.models.dits.lingbotworld.model:LingBotWorldTransformer3DModel`` — a 2x14B boundary-routed MoE (``transformer`` high-noise, ``transformer_2`` low-noise; ``boundary_ratio=0.947``). - The forward gains ``c2ws_plucker_emb`` (FiLM camera injection per block) which a vanilla WanDiT reuse - would drop, so both experts use the ``LingBotWorldDiT`` adapter (``ComponentSpec.adapter``) that threads - the Plucker tensor + keeps the Wan i2v 36ch ``[noise|mask+cond]`` concat + CLIP image embeds. + The forward takes ``c2ws_plucker_emb`` (per-block FiLM camera injection) that a vanilla WanDiT would + drop, so both experts use the ``LingBotWorldDiT`` adapter, which threads the Plucker tensor and keeps + the Wan i2v 36ch ``[noise|mask+cond]`` concat. * VAE ``fastvideo.models.vaes.wanvae:AutoencoderKLWan`` (z=16, 8x spatial / 4x temporal, mean/std - normalized latent space) — reuses the v2 ``WanVAE`` adapter unchanged. ``spatial_scale=8`` for the - Plucker downsample matches this 8x spatial compression. + normalized latent space) — reuses the v2 ``WanVAE`` adapter. ``spatial_scale=8`` for the Plucker + downsample matches the 8x spatial compression. * Text ``fastvideo.models.encoders.t5:T5EncoderModel`` (UMT5, text_len=512) — reuses ``T5Encoder``. - * Image ``fastvideo.models.encoders.clip:CLIPVisionModel`` — first-frame CLIP conditioning (i2v), reuses - ``CLIPImageEncoder``. + * No CLIP image encoder: ``image_dim=null``, so first-frame conditioning is carried entirely by the + 36ch ``[noise|mask+cond]`` latent concat (see the ``cond_encode`` node). -Sampler: flow-match (``FlowUniPCMultistepScheduler``, ``prediction_type='flow_prediction'``) -> the v2 -``FlowShiftPolicy`` + ``FLOW_MATCH_STEP``, exactly as Wan. NOTE the unusual defaults: ``flow_shift=10.0``, -``num_inference_steps=70``, dual ``guidance_scale``/``guidance_scale_2`` (both 5.0), and the Chinese -negative prompt — these must NOT inherit the Wan defaults. +Sampler: flow-match (``FlowUniPCMultistepScheduler``, ``prediction_type='flow_prediction'``) via the v2 +``FlowShiftPolicy`` + ``FLOW_MATCH_STEP``, as in Wan. The defaults are unusual and must NOT inherit Wan's: +``flow_shift=10.0``, ``num_inference_steps=70``, dual ``guidance_scale``/``guidance_scale_2`` (both 5.0), +and a Chinese negative prompt. -``stamp_wan21_checkpoints`` applies (diffusers ``transformer``/``transformer_2``/``vae``/``text_encoder``/ -``image_encoder`` subfolder layout, all in ``_WAN21_SUBFOLDERS``). +``stamp_wan21_checkpoints`` applies (diffusers ``transformer``/``transformer_2``/``vae``/``text_encoder`` +subfolder layout, all in ``_WAN21_SUBFOLDERS``). """ from __future__ import annotations @@ -65,10 +64,10 @@ def build_lingbotworld_card(model_id: str = "lingbot-world-base-cam", sampling_defaults: SamplingDefaults | None = None) -> ModelCard: """LingBot-World-Base-Cam card — 2x14B camera-conditioned dual-guidance MoE i2v. - Mirrors ``build_wan22_i2v_a14b_card`` but: ``flow_shift=10.0``; ``boundary_ratio=0.947``; both DiT - experts use the ``LingBotWorldDiT`` adapter (threads ``c2ws_plucker_emb`` + Wan i2v cond); the - ``LingBotWorldDenoiseLoop`` (dual guidance + camera); and the LingBot defaults (70 steps, gs 5.0, - 480x832, Chinese negative prompt). 2x14B bf16 won't co-reside on 80GB -> the adapter keeps the WanDiT + Mirrors ``build_wan22_i2v_a14b_card`` but with ``flow_shift=10.0``, ``boundary_ratio=0.947``, both DiT + experts on the ``LingBotWorldDiT`` adapter (threads ``c2ws_plucker_emb`` + Wan i2v cond), the + ``LingBotWorldDenoiseLoop`` (dual guidance + camera), and the LingBot defaults (70 steps, gs 5.0, + 480x832, Chinese negative prompt). 2x14B bf16 won't co-reside on 80GB, so the adapter keeps the WanDiT CPU-offload swap (GPU-pending).""" seed = _seed_from(model_id) cost = CostModel(kind=WorkUnitKind.DIFFUSION_STEP, base_seconds=1e-4, per_unit_seconds=1e-7) diff --git a/v2/recipes/lingbotworld/loop.py b/v2/recipes/lingbotworld/loop.py index 68bab8a90..cc58321e5 100644 --- a/v2/recipes/lingbotworld/loop.py +++ b/v2/recipes/lingbotworld/loop.py @@ -1,19 +1,17 @@ -"""LingBotWorldDenoiseLoop — Wan flow-match denoise + boundary-routed MoE, extended with (a) the camera -(Plucker) conditioning slot threaded to the DiT and (b) DUAL classifier-free guidance. +"""LingBotWorldDenoiseLoop — Wan flow-match denoise + boundary-routed MoE, extended with (a) a camera +(Plucker) conditioning slot threaded to the DiT and (b) dual classifier-free guidance. -LingBot-World-Base-Cam is Wan2.2-I2V-A14B i2v conditioning (CLIP + first-frame ``[mask|cond]`` latent) -plus a camera/Plucker tensor. The denoise math is otherwise the canonical Wan flow-match Euler step with +LingBot-World-Base-Cam is Wan2.2-I2V-A14B i2v conditioning (first-frame ``[mask|cond]`` latent) plus a +camera/Plucker tensor. The denoise math is otherwise the canonical Wan flow-match Euler step with ``BoundaryTimestepRouting`` between the high-noise ``transformer`` and the low-noise ``transformer_2``. Two -deltas vs ``WanDenoiseLoop`` (faithful to ``LingBotWorldImageToVideoPipeline`` == the Wan i2v denoise -stage, denoising.py:352-375): +deltas vs ``WanDenoiseLoop`` (matching the fastvideo Wan i2v denoise stage): - * DUAL guidance — fastvideo switches ``guidance_scale -> guidance_scale_2`` at the SAME expert boundary: - ``t >= boundary_timestep`` (the high-noise expert) uses ``guidance_scale``; below it (the low-noise - expert) uses ``guidance_scale_2``. The boundary is in timestep space ``boundary_ratio*1000``, i.e. - sigma-space ``ctx.sigma >= boundary`` (preset ``boundary_ratio=0.947``). The preset sets both scales - to 5.0, so the registered path is numerically a single scale — but the loop honors a distinct - ``guidance_scale_2`` whenever the request carries one. - * CAMERA — the program's camera node writes ``c2ws_plucker_emb`` into slots; the loop publishes it onto + * Dual guidance — fastvideo switches ``guidance_scale -> guidance_scale_2`` at the SAME expert boundary: + the high-noise expert (``ctx.sigma >= boundary``) uses ``guidance_scale``; the low-noise expert uses + ``guidance_scale_2``. The boundary is ``boundary_ratio*1000`` in timestep space, i.e. sigma-space + ``boundary`` (preset ``boundary_ratio=0.947``). The preset sets both scales to 5.0, so the registered + path is numerically a single scale, but the loop honors a distinct ``guidance_scale_2`` when one is set. + * Camera — the program's camera node writes ``c2ws_plucker_emb`` into slots; the loop publishes it onto the active DiT adapter (``dit.c2ws_plucker_emb = ...``) before the forward, because the v2 dit-call surface (and ``ToyDiT``) has no such kwarg. On the CPU toy backend the slot is ``None`` -> the degenerate no-camera path (a plain Wan i2v step), so the loop CPU-verifies end-to-end. @@ -21,7 +19,7 @@ stage, denoising.py:352-375): BRINGUP (needs request-API extension): the Plucker tensor is built from a per-request camera trajectory (``poses.npy``/``intrinsics.npy``); v2's ``Request`` has no camera slot yet, so the program node writes ``None`` unless an ``action_path`` override is supplied. GPU-verify (a) the 2x14B MoE CPU-offload swap, -(b) the camera FiLM injection actually steers the generation, (c) the dual-scale boundary parity. +(b) that the camera FiLM injection steers the generation, (c) the dual-scale boundary parity. """ from __future__ import annotations diff --git a/v2/recipes/lingbotworld/program.py b/v2/recipes/lingbotworld/program.py index 844441591..204feaebb 100644 --- a/v2/recipes/lingbotworld/program.py +++ b/v2/recipes/lingbotworld/program.py @@ -1,20 +1,20 @@ -"""LingBot-World i2v program: image_encode (CLIP) -> cond_encode (first-frame VAE [mask|cond]) -> +"""LingBot-World i2v program: image_encode -> cond_encode (first-frame VAE [mask|cond]) -> camera_encode (Plucker) -> text_encode -> i2v_denoise (dual-guidance MoE) -> vae_decode. -Same inline shape as the Wan i2v program (the i2v conditioning is identical to Wan2.2-I2V-A14B), with one -NEW node: ``camera_encode`` builds the camera/Plucker tensor ``c2ws_plucker_emb`` the LingBot-World DiT -FiLM-injects per block. The ``LingBotWorldDenoiseLoop`` reads that slot and publishes it onto the active -expert adapter; the dual-guidance / MoE specifics live in the loop + adapter, so the node graph is a thin -superset of Wan i2v. +Same inline shape as the Wan i2v program (the i2v conditioning matches Wan2.2-I2V-A14B), plus one new +node: ``camera_encode`` builds the camera/Plucker tensor ``c2ws_plucker_emb`` that the DiT FiLM-injects +per block. The ``LingBotWorldDenoiseLoop`` reads that slot and publishes it onto the active expert adapter; +the dual-guidance / MoE specifics live in the loop + adapter, so the node graph is a thin superset of Wan +i2v. -BRINGUP (blocker #2 — needs a request-API extension): v2's ``Request`` has no camera-trajectory input -slot. The camera node therefore looks for an ``action_path`` in the node override +BRINGUP (needs a request-API extension): v2's ``Request`` has no camera-trajectory input slot. The camera +node therefore looks for an ``action_path`` in the node override (``request.node_override("camera_encode")["action_path"]``, pointing at a dir with ``poses.npy`` + ``intrinsics.npy``) and, if present, builds the Plucker tensor via the fastvideo ``prepare_camera_embedding`` (``spatial_scale=8`` to match the VAE 8x). Absent (the default, and the CPU -toy path) -> it writes ``None`` and the loop runs the degenerate no-camera i2v step. ``prepare_camera_embedding`` -also re-clamps ``num_frames`` to the trajectory length (4k+1); the v2 frame count is taken from the -request, so on GPU the camera node and the request num_frames must agree (verify at bring-up). +toy path) -> it writes ``None`` and the loop runs the degenerate no-camera i2v step. +``prepare_camera_embedding`` also re-clamps ``num_frames`` to the trajectory length (4k+1); the v2 frame +count comes from the request, so on GPU the camera node and the request num_frames must agree. """ from __future__ import annotations @@ -38,12 +38,12 @@ def _i2v_cond_encode(instance, slots, request, ctx) -> None: """First-frame VAE conditioning + mask (mirrors fastvideo's ImageVAEEncodingStage): encode the conditioning image as frame 0 (rest zeros), prepend a 4-channel mask (first latent frame = 1). - This model's DiT has ``in_channels=36`` (16 noise + 4 mask + 16 cond latent), so the i2v cond slot is - REQUIRED for a well-formed forward — a None cond would feed only 16ch into a 36ch patch embedding. The - v2 convenience API (``generate_video(prompt=...)``) supplies no image, so for the bring-up we synthesize - a BLANK first frame (zeros): the latent concat shape is correct and the run produces a finite output - (degenerate no-image i2v, the goal being loads+runs+finite, not quality). A real conditioning image - (when ``request.image()`` is present) takes the normal path.""" + The DiT has ``in_channels=36`` (16 noise + 4 mask + 16 cond latent), so the i2v cond slot is REQUIRED + for a well-formed forward — a None cond would feed only 16ch into a 36ch patch embedding. The v2 + convenience API (``generate_video(prompt=...)``) supplies no image, so for bring-up we synthesize a + blank (zeros) first frame: the latent concat shape is correct and the run produces a finite output + (degenerate no-image i2v). A real conditioning image (``request.image()`` present) takes the normal + path.""" img = request.image() nf = int(request.diffusion.num_frames) h_px, w_px = int(request.diffusion.height), int(request.diffusion.width) diff --git a/v2/recipes/longcat/__init__.py b/v2/recipes/longcat/__init__.py index 5d0fb35d4..5748d398c 100644 --- a/v2/recipes/longcat/__init__.py +++ b/v2/recipes/longcat/__init__.py @@ -1,11 +1,11 @@ """LongCat-Video (T2V) — flow-match denoiser with CFG-zero guidance, ported into the v2 substrate. -Self-contained recipe package (the bucket-C pattern): the card declares its torch adapter via -``ComponentSpec.adapter`` (``LongCatDiT`` in ``v2/platform/backends/torch_longcat.py``) and a new -``LongCatDenoiseLoop`` (explicit ``linspace(1.0, 0.001, num_steps)`` sigma schedule + CFG-zero -optimized-scale combine), reusing the Wan ``WanVAE`` + ``T5Encoder`` adapters and ``stamp_wan21_checkpoints``. -The DiT adapter negates the velocity to fold in the fastvideo stage's ``noise_pred = -noise_pred``. -Registered in ``v2/registry.py`` by the orchestrator (transformer ``LongCatTransformer3DModel``). +Self-contained recipe package: the card declares its torch adapter via ``ComponentSpec.adapter`` +(``LongCatDiT`` in ``v2/platform/backends/torch_longcat.py``) and a new ``LongCatDenoiseLoop`` (explicit +``linspace(1.0, 0.001, num_steps)`` sigma schedule + CFG-zero optimized-scale combine), reusing the Wan +``WanVAE`` + ``T5Encoder`` adapters and ``stamp_wan21_checkpoints``. The DiT adapter negates the velocity +to fold in the fastvideo stage's ``noise_pred = -noise_pred``. Registered in ``v2/registry.py`` by the +orchestrator (transformer ``LongCatTransformer3DModel``). """ from __future__ import annotations diff --git a/v2/recipes/longcat/card.py b/v2/recipes/longcat/card.py index fdf04357c..bbd0531a9 100644 --- a/v2/recipes/longcat/card.py +++ b/v2/recipes/longcat/card.py @@ -2,20 +2,20 @@ Architecture deltas vs Wan (all declared on the card so the recipe is self-contained): * DiT ``fastvideo.models.dits.longcat:LongCatTransformer3DModel`` (depth 48, hidden 4096, 32 heads, - in/out 16ch, patch [1,2,2], caption_channels 4096). The adapter (``ComponentSpec.adapter`` -> ``LongCatDiT`` - in ``v2/platform/backends/torch_longcat.py``) returns the **negated** velocity (the fastvideo stage's - ``noise_pred = -noise_pred`` before the scheduler step) so the loop's flow-match Euler reproduces it. - * VAE ``fastvideo.models.vaes.wanvae:AutoencoderKLWan`` — Wan-style (z=16, 8x/4x) in NORMALIZED latent - space (mean/std). Reuses the v2 ``WanVAE`` torch adapter UNCHANGED (no adapter override). + in/out 16ch, patch [1,2,2], caption_channels 4096). The adapter (``LongCatDiT`` in + ``v2/platform/backends/torch_longcat.py``) returns the **negated** velocity (folding the fastvideo + stage's ``noise_pred = -noise_pred`` before the scheduler step) so the loop's flow-match Euler matches. + * VAE ``fastvideo.models.vaes.wanvae:AutoencoderKLWan`` — Wan-style (z=16, 8x/4x) in normalized latent + space (mean/std). Reuses the v2 ``WanVAE`` torch adapter (no adapter override). * Text ``fastvideo.models.encoders.t5:T5EncoderModel`` (UMT5, 4096-dim) via the v2 ``T5Encoder`` (zero-pads to max_length=512 for the CFG-concat uniform-seq contract — no adapter override needed). - * Sampler: ``FlowMatchEulerDiscreteScheduler`` with an EXPLICIT ``linspace(1.0, 0.001, num_steps)`` sigma + * Sampler: ``FlowMatchEulerDiscreteScheduler`` with an explicit ``linspace(1.0, 0.001, num_steps)`` sigma schedule (NOT flow-shift) + CFG-zero optimized-scale guidance — both live in ``LongCatDenoiseLoop``. -``stamp_wan21_checkpoints`` applies (the diffusers transformer/vae/text_encoder/tokenizer subfolder layout is -identical to Wan). i2v/VC (first-frame latent replacement + per-frame timestep masking + num_cond_latents), -the refine stages (t_thresh crop, which is the latent-prep branch that SKIPS the init-noise re-scale), KV-cache, -and BSA are deferred (spec blocker #8) — the registered preset is base T2V. +``stamp_wan21_checkpoints`` applies (the diffusers transformer/vae/text_encoder/tokenizer subfolder layout +is identical to Wan). i2v/VC (first-frame latent replacement + per-frame timestep masking + num_cond_latents), +the refine stages (t_thresh crop, the latent-prep branch that SKIPS the init-noise re-scale), KV-cache, and +BSA are deferred — the registered preset is base T2V. """ from __future__ import annotations diff --git a/v2/recipes/longcat/loop.py b/v2/recipes/longcat/loop.py index 9e847cd5e..eea0a5e33 100644 --- a/v2/recipes/longcat/loop.py +++ b/v2/recipes/longcat/loop.py @@ -58,10 +58,9 @@ class CFGZeroPolicy(CFGPolicy): Per sample, project cond onto uncond to get the optimized scale ``st_star = / (||v_uncond||^2 + 1e-8)``, then combine - ``out = v_uncond*st_star + scale*(v_cond - v_uncond*st_star)``. This is the exact algebra of - ``LongCatDenoisingStage.optimized_scale`` + the CFG application (longcat_denoising.py lines 28-149), with - the dot-products taken per leading "batch" group. In the v2 loop a step works on ONE unbatched sample - ([C,T,H,W]), so ``st_star`` is a single scalar over the whole tensor (B=1 in fastvideo's reshape(B,-1)).""" + ``out = v_uncond*st_star + scale*(v_cond - v_uncond*st_star)``. This matches + ``LongCatDenoisingStage.optimized_scale`` + its CFG application (longcat_denoising.py). A v2 step works on + ONE unbatched sample ([C,T,H,W]), so ``st_star`` is a single scalar over the whole tensor (B=1).""" def combine(self, preds, guidance_scale, ctx, state): cond = np.asarray(preds["cond"], dtype=np.float64) diff --git a/v2/recipes/ltx2/card.py b/v2/recipes/ltx2/card.py index 86b775451..55edd5139 100644 --- a/v2/recipes/ltx2/card.py +++ b/v2/recipes/ltx2/card.py @@ -1,4 +1,4 @@ -"""LTX-2 ModelCards — split by ARCHITECTURE, not by version label (design_v3 §4, §15). +"""LTX-2 ModelCards — split by ARCHITECTURE, not by version label. * ``build_ltx2_card`` — the **two-stage distilled** pipeline: ``ltx2_base`` (8-step) → learned spatial upsampler → ``ltx2_refine`` (3-step), both binding the same ``transformer``. Serves the distilled @@ -92,7 +92,7 @@ def build_ltx2_card(model_id: str = "ltx2-2stage-distilled") -> ModelCard: factory=lambda inst: ToyDiT(seed=seed), resident_for=["ltx2_base", "ltx2_refine"], required_for={"t2v"}), - # audio branch: lazy for T2V (not loaded), required for T2VS (joint audio+video, §9.11) + # audio branch: lazy for T2V (not loaded), required for T2VS (joint audio+video) "audio_vae": ComponentSpec(component_id="audio_vae", kind="audio_vae", diff --git a/v2/recipes/ltx2/loop.py b/v2/recipes/ltx2/loop.py index 641c39d31..c431207c4 100644 --- a/v2/recipes/ltx2/loop.py +++ b/v2/recipes/ltx2/loop.py @@ -1,11 +1,10 @@ -"""LTX2DenoiseLoop — a custom-step driven loop (design_v3 §5.1 escape hatch; §15 LTX-2 2-stage). +"""LTX2DenoiseLoop — a custom-step driven loop for LTX-2's 2-stage pipeline. LTX-2's guidance is braided (multi-pass: positive / negative / STG-perturbed, combined per the -LTX formula), so the step body is hand-written using samplers + the DiT as a library — exactly -the legitimate custom-step pattern. It is still a Loop (scheduled, admitted, observed, parity- -gated). One ``LTX2DenoiseLoop`` is instantiated per stage (base / refine) with that stage's -distilled sigma schedule; both stages bind the SAME ``transformer`` component (shared by -reference — the MoT 'one instance, many loops' property, demonstrated even here). +LTX formula), so the step body is hand-written using samplers + the DiT as a library. It is +still a Loop (scheduled, admitted, observed, parity-gated). One ``LTX2DenoiseLoop`` is +instantiated per stage (base / refine) with that stage's distilled sigma schedule; both stages +bind the SAME ``transformer`` component (shared by reference — one instance, many loops). Real LTX-2 distilled schedules (from the repo): base : [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0] (8 steps) @@ -34,7 +33,6 @@ from v2.platform.backends.toy import LATENT_CHANNELS BASE_SIGMAS = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0] REFINE_SIGMAS = [0.909375, 0.725, 0.421875, 0.0] - # Real LTX-2 latent: 128 channels; CausalVideoAutoencoder ~32x spatial / 8x temporal. The base stage # runs at HALF the full latent spatial res (stage-2 upsamples 2x), so divide by 64 (=32*2). LTX2_LATENT_CHANNELS = 128 @@ -67,17 +65,29 @@ def base_flow_sigmas(num_steps: int, shift: float = 1.0) -> list[float]: class LTX2DenoiseLoop: - def __init__(self, *, loop_id, stage, sigmas, cfg_scale, stg_scale, cost, - input_slot=None, seed_offset=0, audio_input_slot=None, - full_res=False, request_steps=False, shift=1.0): + + def __init__(self, + *, + loop_id, + stage, + sigmas, + cfg_scale, + stg_scale, + cost, + input_slot=None, + seed_offset=0, + audio_input_slot=None, + full_res=False, + request_steps=False, + shift=1.0): self.loop_id = loop_id - self.stage = stage # "base" | "refine" | "single" (base model) + self.stage = stage # "base" | "refine" | "single" (base model) self.sigmas = list(sigmas) self.cfg_scale = cfg_scale self.stg_scale = stg_scale self.cost = cost - self.input_slot = input_slot # None => fresh noise (base); slot name => read latents (refine) - self.audio_input_slot = audio_input_slot # joint A/V: where to read the threaded audio latent + self.input_slot = input_slot # None => fresh noise (base); slot name => read latents (refine) + self.audio_input_slot = audio_input_slot # joint A/V: where to read the threaded audio latent self.seed_offset = seed_offset # Single-stage LTX-2 base model: full-res latent + a request-driven many-step schedule (distilled # stages keep full_res=False + the fixed few-step self.sigmas). @@ -90,13 +100,18 @@ class LTX2DenoiseLoop: rng = np.random.default_rng(seed) sig = base_flow_sigmas(req.diffusion.num_steps, self.shift) if self.request_steps else self.sigmas if self.input_slot is None: - x = (rng.standard_normal(ltx_base_latent_shape(req, model, full_res=self.full_res)) - * float(sig[0])).astype("float32") + x = (rng.standard_normal(ltx_base_latent_shape(req, model, full_res=self.full_res)) * + float(sig[0])).astype("float32") else: x = np.asarray(ctx.slots[self.input_slot], dtype="float32") - st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id, - request_id=req.request_id, profile=ctx.profile, rng=rng, seed=seed, - latents={"video": x}, sigmas=[float(s) for s in sig], + st = LoopState(loop_id=self.loop_id, + instance_id=model.card.model_id, + request_id=req.request_id, + profile=ctx.profile, + rng=rng, + seed=seed, + latents={"video": x}, + sigmas=[float(s) for s in sig], timesteps=[float(s) * 1000.0 for s in sig]) st.cond["prompt_embeds"] = ctx.slots.get("text_embeds") st.cond["negative_prompt_embeds"] = ctx.slots.get("neg_text_embeds") @@ -105,7 +120,7 @@ class LTX2DenoiseLoop: want_audio = "audio" in getattr(req.outputs, "modalities", frozenset()) gpm = req.diffusion.guidance_per_modality or {} st.scratch["want_audio"] = want_audio - st.scratch["v_guidance"] = float(gpm.get("video", self.cfg_scale)) # == cfg_scale for T2V + st.scratch["v_guidance"] = float(gpm.get("video", self.cfg_scale)) # == cfg_scale for T2V st.scratch["a_guidance"] = float(gpm.get("audio", self.cfg_scale)) if want_audio: if self.audio_input_slot is None: @@ -123,58 +138,69 @@ class LTX2DenoiseLoop: sigma_t, sigma_next = st.sigmas[i], st.sigmas[i + 1] x = st.latents["video"] pe, ne = st.cond["prompt_embeds"], st.cond["negative_prompt_embeds"] - v_guidance, stg_scale = st.scratch["v_guidance"], self.stg_scale # per-modality video CFG + v_guidance, stg_scale = st.scratch["v_guidance"], self.stg_scale # per-modality video CFG want_audio, a_guidance = st.scratch.get("want_audio", False), st.scratch.get("a_guidance", 0.0) au = st.latents.get("audio") def run(model, override=None): - 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) if override is not None and "noise_pred" in override: velocity = np.asarray(override["noise_pred"], dtype="float32") else: dit = model.component("transformer") - pos = dit(x, pe, sigma_t) # full conditioned pass - neg = dit(x, ne, sigma_t) # negative (CFG) pass + pos = dit(x, pe, sigma_t) # full conditioned pass + neg = dit(x, ne, sigma_t) # negative (CFG) pass # STG-perturbed pass (here: drop text conditioning) — the multi-pass braid ptb = dit(x, np.zeros_like(pe) if pe is not None else None, sigma_t) - velocity = (pos - + (v_guidance - 1.0) * (pos - neg) - + stg_scale * (pos - ptb)) + velocity = (pos + (v_guidance - 1.0) * (pos - neg) + stg_scale * (pos - ptb)) x_next = fm(x, np.asarray(velocity, dtype="float32"), sigma_t, sigma_next) out = {"noise_pred": np.asarray(velocity, dtype="float32"), "latents": x_next.astype("float32")} - if want_audio and au is not None: # joint audio denoise, conditioned on video + if want_audio and au is not None: # joint audio denoise, conditioned on video dit = model.component("transformer") - a_pos = dit(au, pe, sigma_t, context=x) # context=x ⇒ audio synced to the video + a_pos = dit(au, pe, sigma_t, context=x) # context=x ⇒ audio synced to the video a_neg = dit(au, ne, sigma_t, context=x) - a_vel = a_pos + (a_guidance - 1.0) * (a_pos - a_neg) # audio's OWN guidance scale + a_vel = a_pos + (a_guidance - 1.0) * (a_pos - a_neg) # audio's OWN guidance scale out["audio_latents"] = fm(au, a_vel, sigma_t, sigma_next).astype("float32") return StepResult(output=out) res = ResourceRequest( compute_seconds=self.cost.predict(int(np.prod(x.shape)), 3.0), # 3 passes - resident_bytes=int(x.nbytes), peak_activation_bytes=int(x.nbytes * 3)) - return WorkPlan( - loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.DIFFUSION_STEP, - shape_sig=ShapeSignature(WorkUnitKind.DIFFUSION_STEP, dims=tuple(x.shape), - extra=(("stage", self.stage),)), - resources=res, payload={"branch": "combined", "step": i, "stage": self.stage}, run=run, - label=f"ltx2.{self.stage}.{i}") + resident_bytes=int(x.nbytes), + peak_activation_bytes=int(x.nbytes * 3)) + return WorkPlan(loop_id=self.loop_id, + instance_id=st.instance_id, + kind=WorkUnitKind.DIFFUSION_STEP, + shape_sig=ShapeSignature(WorkUnitKind.DIFFUSION_STEP, + dims=tuple(x.shape), + extra=(("stage", self.stage), )), + resources=res, + payload={ + "branch": "combined", + "step": i, + "stage": self.stage + }, + run=run, + label=f"ltx2.{self.stage}.{i}") def advance(self, st: LoopState, result: StepResult) -> LoopState: st.latents["video"] = result.output["latents"] if "audio_latents" in result.output: st.latents["audio"] = result.output["audio_latents"] if st.profile == ExecutionProfile.ROLLOUT: - st.trajectory.append({"step": st.step_idx, "stage": self.stage, - "latents": np.asarray(st.latents["video"]).copy()}) + st.trajectory.append({ + "step": st.step_idx, + "stage": self.stage, + "latents": np.asarray(st.latents["video"]).copy() + }) st.step_idx += 1 return st def finalize(self, st: LoopState) -> LoopResult: outs = {"latents": st.latents["video"]} if st.scratch.get("want_audio") and "audio" in st.latents: - outs["audio_latents"] = st.latents["audio"] # threaded to the next stage / decoded - return LoopResult(outputs=outs, metrics={f"{self.stage}_steps": float(st.step_idx)}, + outs["audio_latents"] = st.latents["audio"] # threaded to the next stage / decoded + return LoopResult(outputs=outs, + metrics={f"{self.stage}_steps": float(st.step_idx)}, behavior=st.trajectory or None) @@ -190,11 +216,18 @@ class LTX23DenoiseLoop: so both latents step together each step. Single-pass (distilled guidance≈1.0). Full-res video latent (//32) + a ``[8, T_aud, 16]`` audio latent. For T2V (no audio requested) it runs video-only.""" - def __init__(self, *, loop_id="ltx2_3", sigmas=None, cfg_scale=1.0, stg_scale=0.0, cost, - audio_channels=LTX2_AUDIO_CHANNELS, audio_mel_bins=LTX2_AUDIO_MEL_BINS): + def __init__(self, + *, + loop_id="ltx2_3", + sigmas=None, + cfg_scale=1.0, + stg_scale=0.0, + cost, + audio_channels=LTX2_AUDIO_CHANNELS, + audio_mel_bins=LTX2_AUDIO_MEL_BINS): self.loop_id = loop_id self.stage = "single" - self.sigmas = list(sigmas) if sigmas else list(BASE_SIGMAS) # tuned distilled few-step schedule + self.sigmas = list(sigmas) if sigmas else list(BASE_SIGMAS) # tuned distilled few-step schedule self.cfg_scale = cfg_scale self.stg_scale = stg_scale self.cost = cost @@ -205,19 +238,24 @@ class LTX23DenoiseLoop: seed = req.diffusion.seed if req.diffusion.seed is not None else 0 rng = np.random.default_rng(seed) sig = self.sigmas - x = (rng.standard_normal(ltx_base_latent_shape(req, model, full_res=True)) - * float(sig[0])).astype("float32") - st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id, request_id=req.request_id, - profile=ctx.profile, rng=rng, seed=seed, latents={"video": x}, - sigmas=[float(s) for s in sig], timesteps=[float(s) for s in sig]) + x = (rng.standard_normal(ltx_base_latent_shape(req, model, full_res=True)) * float(sig[0])).astype("float32") + st = LoopState(loop_id=self.loop_id, + instance_id=model.card.model_id, + request_id=req.request_id, + profile=ctx.profile, + rng=rng, + seed=seed, + latents={"video": x}, + sigmas=[float(s) for s in sig], + timesteps=[float(s) for s in sig]) st.cond["prompt_embeds"] = ctx.slots.get("text_embeds") st.cond["audio_prompt_embeds"] = ctx.slots.get("audio_text_embeds") want_audio = "audio" in getattr(req.outputs, "modalities", frozenset()) st.scratch["want_audio"] = want_audio if want_audio: - t_aud = max(2, int(req.diffusion.num_frames)) # ~one audio latent frame per video frame - au = (rng.standard_normal((self.audio_channels, t_aud, self.audio_mel_bins)) - * float(sig[0])).astype("float32") + t_aud = max(2, int(req.diffusion.num_frames)) # ~one audio latent frame per video frame + au = (rng.standard_normal( + (self.audio_channels, t_aud, self.audio_mel_bins)) * float(sig[0])).astype("float32") st.latents["audio"] = au return st @@ -234,32 +272,48 @@ class LTX23DenoiseLoop: def run(model, override=None): fm = model.platform.kernels.get(FLOW_MATCH_STEP) dit = model.component("transformer") - if au is not None: # joint A/V: one forward -> both velocities + if au is not None: # joint A/V: one forward -> both velocities vvel, avel = dit(x, pe, sigma_t, audio_latent=au, audio_text=apa) - out = {"latents": fm(x, np.asarray(vvel, "float32"), sigma_t, sigma_next).astype("float32"), - "audio_latents": fm(au, np.asarray(avel, "float32"), sigma_t, sigma_next).astype("float32"), - "noise_pred": np.asarray(vvel, dtype="float32")} - else: # video-only (T2V) + out = { + "latents": fm(x, np.asarray(vvel, "float32"), sigma_t, sigma_next).astype("float32"), + "audio_latents": fm(au, np.asarray(avel, "float32"), sigma_t, sigma_next).astype("float32"), + "noise_pred": np.asarray(vvel, dtype="float32") + } + else: # video-only (T2V) vvel = dit(x, pe, sigma_t) - out = {"latents": fm(x, np.asarray(vvel, "float32"), sigma_t, sigma_next).astype("float32"), - "noise_pred": np.asarray(vvel, dtype="float32")} + out = { + "latents": fm(x, np.asarray(vvel, "float32"), sigma_t, sigma_next).astype("float32"), + "noise_pred": np.asarray(vvel, dtype="float32") + } return StepResult(output=out) res = ResourceRequest(compute_seconds=self.cost.predict(int(np.prod(x.shape)), 1.0), - resident_bytes=int(x.nbytes), peak_activation_bytes=int(x.nbytes * 2)) - return WorkPlan( - loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.DIFFUSION_STEP, - shape_sig=ShapeSignature(WorkUnitKind.DIFFUSION_STEP, dims=tuple(x.shape), - extra=(("stage", self.stage),)), - resources=res, payload={"step": i, "stage": self.stage}, run=run, label=f"ltx2_3.{i}") + resident_bytes=int(x.nbytes), + peak_activation_bytes=int(x.nbytes * 2)) + return WorkPlan(loop_id=self.loop_id, + instance_id=st.instance_id, + kind=WorkUnitKind.DIFFUSION_STEP, + shape_sig=ShapeSignature(WorkUnitKind.DIFFUSION_STEP, + dims=tuple(x.shape), + extra=(("stage", self.stage), )), + resources=res, + payload={ + "step": i, + "stage": self.stage + }, + run=run, + label=f"ltx2_3.{i}") def advance(self, st: LoopState, result: StepResult) -> LoopState: st.latents["video"] = result.output["latents"] if "audio_latents" in result.output: st.latents["audio"] = result.output["audio_latents"] if st.profile == ExecutionProfile.ROLLOUT: - st.trajectory.append({"step": st.step_idx, "stage": self.stage, - "latents": np.asarray(st.latents["video"]).copy()}) + st.trajectory.append({ + "step": st.step_idx, + "stage": self.stage, + "latents": np.asarray(st.latents["video"]).copy() + }) st.step_idx += 1 return st @@ -267,5 +321,4 @@ class LTX23DenoiseLoop: outs = {"latents": st.latents["video"]} if st.scratch.get("want_audio") and "audio" in st.latents: outs["audio_latents"] = st.latents["audio"] - return LoopResult(outputs=outs, metrics={"single_steps": float(st.step_idx)}, - behavior=st.trajectory or None) + return LoopResult(outputs=outs, metrics={"single_steps": float(st.step_idx)}, behavior=st.trajectory or None) diff --git a/v2/recipes/ltx2/program.py b/v2/recipes/ltx2/program.py index c567354cb..191733722 100644 --- a/v2/recipes/ltx2/program.py +++ b/v2/recipes/ltx2/program.py @@ -1,4 +1,4 @@ -"""LTX-2 two-stage distilled program (design_v3 §13, §15f): +"""LTX-2 two-stage distilled program: text_encode → ltx2_base (8-step) → upsample(+noise mix) → ltx2_refine (3-step) → vae_decode @@ -19,11 +19,11 @@ def _upsample(instance, slots, request, ctx) -> None: # Learned 2× spatial latent super-resolution between the base and refine stages. The # ``spatial_upsampler`` component is the real LTX2LatentUpsampler on the GPU backend (un_normalize → # learned upsample → normalize, via the VAE's per-channel stats) and a nearest-neighbor toy on CPU — - # same call on both backends (the v2 component-polymorphism seam), no device branch here. + # same call on both backends, no device branch here. up = np.asarray(instance.component("spatial_upsampler").upsample(base), dtype="float32") seed = (request.diffusion.seed if request.diffusion.seed is not None else 0) + 7 noise = np.random.default_rng(seed).standard_normal(up.shape).astype("float32") - sigma0 = float(REFINE_SIGMAS[0]) # 0.909375 + sigma0 = float(REFINE_SIGMAS[0]) # 0.909375 slots["ltx_upsampled"] = (noise * sigma0 + up * (1.0 - sigma0)).astype("float32") @@ -33,17 +33,27 @@ def _vae_decode(instance, slots, request, ctx) -> None: def build_ltx2_program() -> Program: return Program( - program_id="ltx2.t2v.2stage", kind=ProgramKind.INLINE, + program_id="ltx2.t2v.2stage", + kind=ProgramKind.INLINE, nodes=[ ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")), - ModelLoopNode("base", loop_id="ltx2_base", output_slot="ltx_base_out", - reads=("text_embeds",), writes=("ltx_base_out",)), - ComponentNode("upsample", fn=_upsample, reads=("ltx_base_out",), writes=("ltx_upsampled",)), - ModelLoopNode("refine", loop_id="ltx2_refine", output_slot="ltx_refine_out", - reads=("ltx_upsampled",), writes=("ltx_refine_out",)), - ComponentNode("vae_decode", fn=_vae_decode, reads=("ltx_refine_out",), writes=("video",)), + ModelLoopNode("base", + loop_id="ltx2_base", + output_slot="ltx_base_out", + reads=("text_embeds", ), + writes=("ltx_base_out", )), + ComponentNode("upsample", fn=_upsample, reads=("ltx_base_out", ), writes=("ltx_upsampled", )), + ModelLoopNode("refine", + loop_id="ltx2_refine", + output_slot="ltx_refine_out", + reads=("ltx_upsampled", ), + writes=("ltx_refine_out", )), + ComponentNode("vae_decode", fn=_vae_decode, reads=("ltx_refine_out", ), writes=("video", )), ], - output_artifacts={"video": "video", "latents": "ltx_refine_out"}, + output_artifacts={ + "video": "video", + "latents": "ltx_refine_out" + }, ).validate() @@ -56,14 +66,21 @@ def build_ltx2_base_program() -> Program: """Single-stage LTX-2 base: text_encode → denoise (request-driven many-step, full-res) → vae_decode. No upsample/refine (those are the distilled two-stage program above).""" return Program( - program_id="ltx2.t2v.base", kind=ProgramKind.INLINE, + program_id="ltx2.t2v.base", + kind=ProgramKind.INLINE, nodes=[ ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")), - ModelLoopNode("denoise", loop_id="ltx2_single", output_slot="ltx_out", - reads=("text_embeds",), writes=("ltx_out",)), - ComponentNode("vae_decode", fn=_vae_decode_single, reads=("ltx_out",), writes=("video",)), + ModelLoopNode("denoise", + loop_id="ltx2_single", + output_slot="ltx_out", + reads=("text_embeds", ), + writes=("ltx_out", )), + ComponentNode("vae_decode", fn=_vae_decode_single, reads=("ltx_out", ), writes=("video", )), ], - output_artifacts={"video": "video", "latents": "ltx_out"}, + output_artifacts={ + "video": "video", + "latents": "ltx_out" + }, ).validate() @@ -87,21 +104,29 @@ def build_ltx2_3_program() -> Program: (AudioDecoder→Vocoder). A plain T2V request leaves the audio latent unproduced, so audio_decode writes None and the audio components stay unbuilt.""" return Program( - program_id="ltx2.3.t2vs", kind=ProgramKind.INLINE, + program_id="ltx2.3.t2vs", + kind=ProgramKind.INLINE, nodes=[ ComponentNode("text_encode", fn=_text_encode_av, writes=("text_embeds", "audio_text_embeds")), - ModelLoopNode("denoise", loop_id="ltx2_3", output_slot="ltx_out", - reads=("text_embeds", "audio_text_embeds"), writes=("ltx_out",)), - ComponentNode("vae_decode", fn=_vae_decode_single, reads=("ltx_out",), writes=("video",)), - ComponentNode("audio_decode", fn=_audio_decode_single, reads=("ltx_out",), writes=("audio",)), + ModelLoopNode("denoise", + loop_id="ltx2_3", + output_slot="ltx_out", + reads=("text_embeds", "audio_text_embeds"), + writes=("ltx_out", )), + ComponentNode("vae_decode", fn=_vae_decode_single, reads=("ltx_out", ), writes=("video", )), + ComponentNode("audio_decode", fn=_audio_decode_single, reads=("ltx_out", ), writes=("audio", )), ], - output_artifacts={"video": "video", "audio": "audio", "latents": "ltx_out"}, + output_artifacts={ + "video": "video", + "audio": "audio", + "latents": "ltx_out" + }, ).validate() -# --- joint audio+video (T2VS) program (design_v3 §15b; §9.11) ---------------------- # +# --- joint audio+video (T2VS) program ---------------------------------------------- # def _upsample_av(instance, slots, request, ctx) -> None: - _upsample(instance, slots, request, ctx) # video latent upsample (reuse) + _upsample(instance, slots, request, ctx) # video latent upsample (reuse) base = slots.get("ltx_base_out", {}) slots["ltx_audio"] = base.get("audio_latents") if isinstance(base, dict) else None # thread audio @@ -116,17 +141,27 @@ def build_ltx2_av_program() -> Program: """T2VS: the two-stage denoise carries a synchronized audio latent through both stages; the video VAE and the audio VAE decode the two modalities → video + audio artifacts (per-modality guidance).""" return Program( - program_id="ltx2.t2vs.2stage", kind=ProgramKind.INLINE, + program_id="ltx2.t2vs.2stage", + kind=ProgramKind.INLINE, nodes=[ ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")), - ModelLoopNode("base", loop_id="ltx2_base", output_slot="ltx_base_out", - reads=("text_embeds",), writes=("ltx_base_out",)), - ComponentNode("upsample", fn=_upsample_av, reads=("ltx_base_out",), - writes=("ltx_upsampled", "ltx_audio")), - ModelLoopNode("refine", loop_id="ltx2_refine", output_slot="ltx_refine_out", - reads=("ltx_upsampled", "ltx_audio"), writes=("ltx_refine_out",)), - ComponentNode("vae_decode", fn=_vae_decode, reads=("ltx_refine_out",), writes=("video",)), - ComponentNode("audio_decode", fn=_audio_decode, reads=("ltx_refine_out",), writes=("audio",)), + ModelLoopNode("base", + loop_id="ltx2_base", + output_slot="ltx_base_out", + reads=("text_embeds", ), + writes=("ltx_base_out", )), + ComponentNode("upsample", fn=_upsample_av, reads=("ltx_base_out", ), writes=("ltx_upsampled", "ltx_audio")), + ModelLoopNode("refine", + loop_id="ltx2_refine", + output_slot="ltx_refine_out", + reads=("ltx_upsampled", "ltx_audio"), + writes=("ltx_refine_out", )), + ComponentNode("vae_decode", fn=_vae_decode, reads=("ltx_refine_out", ), writes=("video", )), + ComponentNode("audio_decode", fn=_audio_decode, reads=("ltx_refine_out", ), writes=("audio", )), ], - output_artifacts={"video": "video", "audio": "audio", "latents": "ltx_refine_out"}, + output_artifacts={ + "video": "video", + "audio": "audio", + "latents": "ltx_refine_out" + }, ).validate() diff --git a/v2/recipes/matrixgame2/__init__.py b/v2/recipes/matrixgame2/__init__.py index 94bd5a480..5e13b3597 100644 --- a/v2/recipes/matrixgame2/__init__.py +++ b/v2/recipes/matrixgame2/__init__.py @@ -1,11 +1,11 @@ -"""Matrix-Game 2.0 (Base, distilled) — the INTERACTIVE mouse/keyboard world model ported into the v2 -substrate as a self-contained recipe package (bucket-C pattern). +"""Matrix-Game 2.0 (Base, distilled) — interactive mouse/keyboard world model, as a self-contained v2 +recipe package (bucket-C pattern). The card declares its torch adapters via ``ComponentSpec.adapter`` (``MatrixGame2CausalDiT`` / -``MatrixGame2CLIPImageEncoder`` in ``v2/platform/backends/torch_matrixgame2.py``) and a new +``MatrixGame2CLIPImageEncoder`` in ``v2/platform/backends/torch_matrixgame2.py``) plus ``MatrixGame2CausalDMDLoop`` (few-step DMD: epsilon->x0 via the FlowUniPC sigma table + re-add_noise, causal block-autoregressive with a sliding-window KV cache, action-conditioned). Reuses the Wan VAE adapter + -``stamp_wan21_checkpoints``. Registered in ``v2/registry.py`` (the orchestrator adds the entry). +``stamp_wan21_checkpoints``. Registered in ``v2/registry.py``. The registered path is the i2v world-rollout (first-frame + CLIP context); live mouse/keyboard action routing is BRINGUP (needs a request-API extension — the loop + program already thread the action slots). diff --git a/v2/recipes/matrixgame2/loop.py b/v2/recipes/matrixgame2/loop.py index b3bd5c419..91eac60cd 100644 --- a/v2/recipes/matrixgame2/loop.py +++ b/v2/recipes/matrixgame2/loop.py @@ -1,32 +1,29 @@ -"""MatrixGame2CausalDMDLoop — the few-step, causal/autoregressive, action-conditioned DMD denoise loop for -Matrix-Game 2.0 (the INTERACTIVE world model). Cannot reuse ``WanDenoiseLoop``/``CosmosDenoiseLoop``: +"""MatrixGame2CausalDMDLoop — few-step, causal/autoregressive, action-conditioned DMD denoise loop for +Matrix-Game 2.0. Cannot reuse ``WanDenoiseLoop``/``CosmosDenoiseLoop`` because: (a) Few-step DMD, NOT a flow-match velocity ODE. The model output is EPSILON; each DMD step converts - ``eps -> x0`` via ``pred_video = x - sigma_t*eps`` (the scheduler sigma table), then RE-ADDS noise to - the next DMD timestep. Fixed schedule ``dmd_denoising_steps=[1000,666,333]`` (3 steps), warped through - the FlowUniPC timestep grid (``warp_denoising_step``). CFG is OFF (guidance_scale=1.0). + ``eps -> x0`` via ``pred_video = x - sigma_t*eps`` (scheduler sigma table), then RE-ADDS noise to + the next DMD timestep. Fixed 3-step schedule ``[1000,666,333]``, warped through the FlowUniPC + timestep grid (``warp_denoising_step``). CFG is OFF (guidance_scale=1.0). (b) Causal block-autoregressive: the video latent is split into blocks of ``num_frames_per_block`` (=3); each block is fully DMD-denoised, then a context pass at ``timestep=context_noise(=0)`` writes its CLEAN K/V into a sliding-window KV cache before the next block. Per-frame timestep ``[B, num_frames]``. (c) Action-conditioned: per-frame mouse[F,2] / keyboard[F,4] feed ActionModule blocks with their own KV - caches. The DEGENERATE (no-action) path is what CPU-verifies here; live action routing is BRINGUP - (needs a request-API extension — the loop exposes the ``mouse_cond``/``keyboard_cond`` slots and a - ``streaming_reset``/``streaming_step`` surface for interactive play). + caches. The no-action degenerate path is what CPU-verifies here; live action routing is BRINGUP + (needs a request-API extension — the loop exposes ``mouse_cond``/``keyboard_cond`` slots). -Faithful port of ``fastvideo/pipelines/stages/matrixgame2_denoising.py:MatrixGame2CausalDenoisingStage`` -(``forward`` / ``_process_single_block`` / ``_update_context_cache``). The DiT call is dispatched two ways -so the SAME loop runs on the CPU toy AND the GPU adapter: +Faithful port of ``fastvideo/pipelines/stages/matrixgame2_denoising.py:MatrixGame2CausalDenoisingStage``. +The DiT call is dispatched two ways so the SAME loop runs on the CPU toy AND the GPU adapter: * GPU: ``model.component('transformer')`` is a ``MatrixGame2CausalDiT`` exposing ``.call(...)`` with the causal kv_cache / i2v / action plumbing built INTERNALLY (the loop just passes numpy + the cache lists). - * CPU toy: ``ToyDiT`` has only ``dit(latent, text_embed, sigma)``; the loop's ``_eps`` falls back to that - (no kv_cache — the toy is a degenerate stand-in exercising the DMD + block + epsilon->x0 control flow - with real numbers). The KV caches are still allocated + threaded so the structure is identical. + * CPU toy: ``ToyDiT`` has only ``dit(latent, text_embed, sigma)``; ``_eps`` falls back to that (no + kv_cache — a degenerate stand-in exercising the DMD + block + epsilon->x0 control flow with real + numbers). The KV caches are still allocated + threaded so the structure is identical. -BRINGUP: the GPU kv_cache / crossattn_cache / action-cache shapes (``local_attn_size`` window, mouse-cache -batch dim ``B*frame_seq_len``) come from the loaded arch_config and are allocated by the GPU adapter's -companion init; on CPU we allocate lightweight placeholder dict lists so the threading is faithful but the -toy ignores them. The MoE high/low-noise boundary path (``pred_noise_to_x_bound`` + ``add_noise_high``) -exists for GTA/TempleRun variants (``boundary_ratio`` set); the Base distilled checkpoint is single-expert. +The GPU kv_cache / crossattn_cache / action-cache shapes come from the loaded arch_config and are +allocated by the GPU adapter's companion init; on CPU we allocate placeholder dict lists the toy ignores. +The MoE high/low-noise boundary path (``pred_noise_to_x_bound`` + ``add_noise_high``) exists for +GTA/TempleRun variants (``boundary_ratio`` set); the Base distilled checkpoint is single-expert. """ from __future__ import annotations @@ -57,7 +54,7 @@ from v2.recipes.matrixgame2.sampler import ( ) from v2.recipes.wan21.loop import latent_shape -# Real Wan2.1 VAE (AutoencoderKLWan): z=16, 4x temporal, 8x spatial. The DiT patch_embedding is 32-in +# Wan2.1 VAE (AutoencoderKLWan): z=16, 4x temporal, 8x spatial. The DiT patch_embedding is 32-in # (16 noise + 16 i2v cond_concat) but the LOOP's latent is the 16ch noise; the adapter does the concat. MATRIXGAME2_LATENT_CHANNELS = 16 MATRIXGAME2_TEMPORAL_RATIO = 4 @@ -291,10 +288,10 @@ class MatrixGame2CausalDMDLoop: # advance # # --------------------------------------------------------------------- # def advance(self, st: LoopState, result: StepResult) -> LoopState: - # Monotonic progress counter: this loop drives control flow off block_idx/dmd_idx/phase (in - # scratch), but the runtime's no-progress watchdog (engine ProgramRunner._progress) keys on - # st.step_idx. Bump it on EVERY executed work unit (each DMD step AND each clean-context pass) - # so a multi-block causal rollout is seen to advance — every other recipe loop does the same. + # Monotonic progress counter: control flow runs off block_idx/dmd_idx/phase in scratch, but the + # runtime's no-progress watchdog (engine ProgramRunner._progress) keys on st.step_idx. Bump it on + # EVERY executed work unit (each DMD step AND each clean-context pass) so a multi-block causal + # rollout is seen to advance. st.step_idx += 1 out = result.output start, num = out["start"], out["num"] diff --git a/v2/recipes/matrixgame2/program.py b/v2/recipes/matrixgame2/program.py index 5b8879946..7a44e3c45 100644 --- a/v2/recipes/matrixgame2/program.py +++ b/v2/recipes/matrixgame2/program.py @@ -5,9 +5,8 @@ Deltas vs the Wan2.1 i2v program (faithful to ``matrixgame2_causal_dmd_pipeline. ``MatrixGame2ImageVAEEncodingStage``): * NO text_encode node — Matrix-Game 2.0 ignores text; the CLIP image embeds are the sole cross-attn context (written to ``i2v_img_embeds``). - * cond_encode writes the RAW first-frame VAE latent to ``i2v_cond`` with NO 4-channel mask (unlike Wan2.1 - i2v's ``[mask|cond]``); the DiT channel-concats this raw cond onto the noise latent internally (the - patch_embedding is the 32-in checkpoint). + * cond_encode writes a 20-channel ``cond_concat`` (4 first-frame mask channels + 16 VAE-latent channels) + to ``i2v_cond``; the DiT channel-concats this onto the 16ch noise latent internally. * action_prepare is a HOOK that routes per-frame mouse/keyboard arrays into ``mouse_cond``/``keyboard_cond`` when the request carries them. BRINGUP: the request API does not yet surface game-action arrays, so this is a no-op on the registered path (the world model degrades to a first-frame-conditioned rollout). When @@ -28,12 +27,11 @@ _MG2_TEMPORAL_RATIO = 4 def _first_frame_pixels(request) -> np.ndarray: - """Return the first-frame conditioning image as ``[3, H, W]`` float32 in [-1, 1]. Matrix-Game 2.0 is an - i2v world model: a first frame is MANDATORY (it is the only spatial conditioning + drives the cond_concat - that makes the DiT's 36 in-channels). When the request carries no image (degenerate world-rollout / a - plain text smoke request), synthesize a neutral mid-gray frame at the requested resolution -- the world - model then extends from a blank first frame. Faithful to fastvideo's ``create_default_image`` fallback - in ``MatrixGame2ImageEncodingStage`` (which uses a default image when ``batch.pil_image is None``).""" + """Return the first-frame conditioning image as ``[3, H, W]`` float32 in [-1, 1]. A first frame is + MANDATORY: it is the only spatial conditioning and drives the cond_concat for the DiT's 36 in-channels. + When the request carries no image (degenerate world-rollout / plain text smoke request), synthesize a + neutral mid-gray frame at the requested resolution and extend from there. Mirrors fastvideo's + ``create_default_image`` fallback in ``MatrixGame2ImageEncodingStage`` (``batch.pil_image is None``).""" img = request.image() h = int(request.diffusion.height) w = int(request.diffusion.width) diff --git a/v2/recipes/matrixgame2/sampler.py b/v2/recipes/matrixgame2/sampler.py index e38c411eb..a9e900a4a 100644 --- a/v2/recipes/matrixgame2/sampler.py +++ b/v2/recipes/matrixgame2/sampler.py @@ -1,18 +1,17 @@ -"""Matrix-Game 2.0 DMD sampler helpers — the few-step distilled epsilon->x0 math, ported faithfully from +"""Matrix-Game 2.0 DMD sampler helpers — the few-step distilled epsilon->x0 math, ported from ``fastvideo/models/utils.py`` (``pred_noise_to_pred_video`` / ``pred_noise_to_x_bound``) and the -``FlowUniPCMultistepScheduler`` sigma/timestep table it indexes. Kept HERE (a recipe-local sibling), NOT -in ``v2/loop/sampler.py``, because the DMD epsilon->x0 + re-add_noise schedule is specific to this arch. +``FlowUniPCMultistepScheduler`` sigma/timestep table it indexes. Kept recipe-local (not in +``v2/loop/sampler.py``) because the DMD epsilon->x0 + re-add_noise schedule is specific to this arch. -The shipped distilled checkpoints run a FIXED 3-step DMD schedule ``dmd_denoising_steps=[1000,666,333]``. -The scheduler (``FlowUniPCMultistepScheduler(shift=5.0)``) is used ONLY as a sigma/timestep lookup table + -``add_noise``, never as a UniPC multistep corrector. With ``warp_denoising_step=True`` the nominal integer -steps are remapped through the scheduler's timestep grid: -``timesteps = cat(scheduler.timesteps, [0])[1000 - dmd_steps]``. +The distilled checkpoints run a FIXED 3-step DMD schedule ``[1000,666,333]``. The scheduler +(``FlowUniPCMultistepScheduler(shift=5.0)``) is used ONLY as a sigma/timestep lookup table + ``add_noise``, +never as a UniPC multistep corrector. With ``warp_denoising_step=True`` the nominal integer steps are +remapped through its timestep grid: ``timesteps = cat(scheduler.timesteps, [0])[1000 - dmd_steps]``. This numpy form is the CPU-testable, bit-reproducible reference. ``build_flow_unipc_table`` reproduces the -scheduler's ``sigmas``/``timesteps`` (flow-shift over a 1000-step alpha ramp; identical formula to -``FlowShiftPolicy``) so the loop can run on the CPU toy without instantiating the real scheduler. On a GPU -box the torch scheduler is built from the card's ``scheduler`` load_id and the SAME math applies. +scheduler's ``sigmas``/``timesteps`` (flow-shift over a 1000-step alpha ramp; same formula as +``FlowShiftPolicy``) so the loop runs on the CPU toy without the real scheduler. On GPU the torch scheduler +is built from the card's ``scheduler`` load_id and the SAME math applies. """ from __future__ import annotations diff --git a/v2/recipes/matrixgame3/__init__.py b/v2/recipes/matrixgame3/__init__.py index 68b1f14b8..2a54c9208 100644 --- a/v2/recipes/matrixgame3/__init__.py +++ b/v2/recipes/matrixgame3/__init__.py @@ -1,10 +1,10 @@ -"""Matrix-Game 3.0 (Base-Distilled) — autoregressive multi-clip world model ported into the v2 substrate. +"""Matrix-Game 3.0 (Base-Distilled) — autoregressive multi-clip world model, as a self-contained v2 recipe. -Self-contained recipe package (the bucket-C pattern): the card declares its torch adapter via -``ComponentSpec.adapter`` (``MatrixGame3DiT`` in ``v2/platform/backends/torch_matrixgame3.py``) and a new -``MatrixGame3DenoiseLoop`` (the autoregressive multi-clip loop: per-token timesteps + first-frame pinning; -action/camera/KV-memory are BRINGUP), reusing the Wan VAE adapter + T5 + ``stamp_wan21_checkpoints``. The -orchestrator registers it in ``v2/registry.py`` (the ``MatrixGame3WanModel`` transformer fallback). +Bucket-C pattern: the card declares its torch adapter via ``ComponentSpec.adapter`` (``MatrixGame3DiT`` in +``v2/platform/backends/torch_matrixgame3.py``) plus ``MatrixGame3DenoiseLoop`` (autoregressive multi-clip +loop: per-token timesteps + first-frame pinning; action/camera/KV-memory are BRINGUP), reusing the Wan VAE +adapter + T5 + ``stamp_wan21_checkpoints``. Registered in ``v2/registry.py`` (the ``MatrixGame3WanModel`` +transformer fallback). """ from __future__ import annotations diff --git a/v2/recipes/matrixgame3/card.py b/v2/recipes/matrixgame3/card.py index b2edfa008..19a8ce72d 100644 --- a/v2/recipes/matrixgame3/card.py +++ b/v2/recipes/matrixgame3/card.py @@ -1,20 +1,20 @@ """Matrix-Game 3.0 ModelCard — image+action+camera -> video autoregressive world model. -Self-contained recipe package (the bucket-C pattern, mirroring cosmos2): the card declares its torch -adapter via ``ComponentSpec.adapter`` (``MatrixGame3DiT`` in ``v2/platform/backends/torch_matrixgame3.py``) -and a new ``MatrixGame3DenoiseLoop`` (the autoregressive multi-clip loop), reusing the Wan VAE adapter + -T5 + ``stamp_wan21_checkpoints``. +Self-contained recipe package (bucket-C pattern, mirroring cosmos2): the card declares its torch adapter +via ``ComponentSpec.adapter`` (``MatrixGame3DiT`` in ``v2/platform/backends/torch_matrixgame3.py``) plus +``MatrixGame3DenoiseLoop`` (the autoregressive multi-clip loop), reusing the Wan VAE adapter + T5 + +``stamp_wan21_checkpoints``. -Architecture deltas vs Wan (all GPU-path, declared on the card so the recipe is self-contained): +Architecture deltas vs Wan (all GPU-path, declared on the card): * DiT ``fastvideo.models.dits.matrixgame3.model:MatrixGame3WanModel`` — a 5B autoregressive world-model DiT (30 layers, dim 24*128=3072, ffn 14336, patch (1,2,2), in/out_channels=48, ``use_memory=True``, ``sigma_theta=0.8``, action_config on the first 15 blocks). Predicts a velocity - (``flow_prediction``). The adapter (``ComponentSpec.adapter`` -> ``MatrixGame3DiT``) takes a PER-TOKEN - timestep (cond-frame rows zeroed) + the action/camera/KV-memory bundle — built INTERNALLY so the - loop's dit-call stays ``dit(latent, text_embed, sigma)``-shaped. + (``flow_prediction``). The adapter takes a PER-TOKEN timestep (cond-frame rows zeroed) + the + action/camera/KV-memory bundle, built INTERNALLY so the loop's dit-call stays + ``dit(latent, text_embed, sigma)``-shaped. * VAE ``fastvideo.models.vaes.wanvae:AutoencoderKLWan`` — the ``light_vae`` variant (z_dim=48, 16x spatial / 4x temporal, Wan2.2-TI2V geometry); reuses the v2 ``WanVAE`` adapter unchanged (its - ``_mean_invstd`` normalization path is the MG3 normalized-latent convention). + ``_mean_invstd`` normalization is the MG3 normalized-latent convention). * Text ``fastvideo.models.encoders.t5:T5EncoderModel`` (T5, fp32, text_len 512) via the default ``T5Encoder`` adapter (Wan zero-pad-to-512 convention). NO image/CLIP encoder — the conditioning image goes only through the VAE (BRINGUP risk 6). diff --git a/v2/recipes/multi_expert/__init__.py b/v2/recipes/multi_expert/__init__.py index 5bbf74c7f..b16d7bbfd 100644 --- a/v2/recipes/multi_expert/__init__.py +++ b/v2/recipes/multi_expert/__init__.py @@ -1,5 +1,5 @@ """Multi-expert card: N refiner LMs + one generator (N+1 trainable experts, N+1 loops) — the -substrate for N-way joint RL (design_v3 §4, §10). Generalizes the two-expert ``unified`` card.""" +substrate for N-way joint RL. Generalizes the two-expert ``unified`` card.""" from __future__ import annotations from v2.recipes.multi_expert.card import N_REFINE_ACTIONS, build_multi_expert_card, refiner_ids diff --git a/v2/recipes/multi_expert/card.py b/v2/recipes/multi_expert/card.py index d04009fed..c4c32d6ae 100644 --- a/v2/recipes/multi_expert/card.py +++ b/v2/recipes/multi_expert/card.py @@ -1,15 +1,13 @@ -"""Multi-expert ModelCard — N refiner LMs + one generator, for N-way joint RL (design_v3 §4, §10). +"""Multi-expert ModelCard — N refiner LMs + one generator, for N-way joint RL. -The generalization test of the UniRL stress test: where ``unified`` has *two* trainable experts -(``llm`` + ``transformer``), this card has ``n_refiners`` prompt-refiner LMs **plus** a generator — -N+1 disjoint experts on N+1 loops (``ar_decode_0…n-1`` + ``diffusion_denoise``). It exists to show -that "joint RL over more than two experts" needs **no new card primitive**: a card already holds an -arbitrary number of components and loops (Qwen-Omni proved three). The only question is whether the -*method* can drive N updates from one reward — answered by ``JointMultiExpertRL`` (no, it needs no -rewrite; it loops over an expert list). +Generalizes the two-expert ``unified`` card: instead of ``llm`` + ``transformer``, this card has +``n_refiners`` prompt-refiner LMs plus a generator — N+1 disjoint experts on N+1 loops +(``ar_decode_0…n-1`` + ``diffusion_denoise``). It shows that joint RL over more than two experts needs +no new card primitive (a card already holds an arbitrary number of components and loops); +``JointMultiExpertRL`` drives the N updates from one reward by looping over an expert list. -The N refiners' refinements *compose* into the diffusion conditioning (each adds its embedding -offset), so the generator's output depends on every refiner — and one reward trains them all. +The N refiners' refinements compose into the diffusion conditioning (each adds its embedding offset), +so the generator's output depends on every refiner and one reward trains them all. """ from __future__ import annotations @@ -48,52 +46,71 @@ def build_multi_expert_card(model_id: str = "multi-expert-rl", *, n_refiners: in rids = refiner_ids(n_refiners) def denoise_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) def ar_factory(rid): - return lambda: ARDecodeLoop(loop_id=f"ar_decode_{rid}", transformer_id=rid, cost=ar_cost, - max_tokens=1, prompt_slot="prompt_tokens") + return lambda: ARDecodeLoop( + loop_id=f"ar_decode_{rid}", transformer_id=rid, cost=ar_cost, max_tokens=1, prompt_slot="prompt_tokens") 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=dn_cost, - shared_weight_components=["transformer"], cache_policy=["feature"], - loop_factory=denoise_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=denoise_factory), } - for i, rid in enumerate(rids): # N refiner experts + N ar loops + for i, rid in enumerate(rids): # N refiner experts + N ar loops components[rid] = ComponentSpec( - rid, kind="dit", + rid, + kind="dit", factory=(lambda s: (lambda inst: ToyPromptRefiner(n_actions=N_REFINE_ACTIONS, seed=s)))(seed + 11 + i), - resident_for=[f"ar_decode_{rid}"], required_for={"t2v"}) - loops[f"ar_decode_{rid}"] = LoopSpec( - f"ar_decode_{rid}", kind=LoopKind.AR_DECODE, work_unit_kind=WorkUnitKind.AR_TOKEN, - step_cost_model=ar_cost, shared_weight_components=[rid], cache_policy=["paged_kv"], - loop_factory=ar_factory(rid)) + resident_for=[f"ar_decode_{rid}"], + required_for={"t2v"}) + loops[f"ar_decode_{rid}"] = LoopSpec(f"ar_decode_{rid}", + kind=LoopKind.AR_DECODE, + work_unit_kind=WorkUnitKind.AR_TOKEN, + step_cost_model=ar_cost, + shared_weight_components=[rid], + cache_policy=["paged_kv"], + loop_factory=ar_factory(rid)) card = ModelCard( - model_id=model_id, family="multi_expert", components=components, loops=loops, - capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.REASONING_TEXT, - Capability.VAE_DECODE, Capability.POLICY_ROLLOUT), - recipe=RecipeSpec(method="base", assumes_loop="diffusion_denoise", - assumes_precision="float32", consistency_required=ConsistencyLevel.C2), - parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1, ConsistencyLevel.C2], - interleave_required=True), + model_id=model_id, + family="multi_expert", + components=components, + loops=loops, + capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.REASONING_TEXT, Capability.VAE_DECODE, + Capability.POLICY_ROLLOUT), + recipe=RecipeSpec(method="base", + assumes_loop="diffusion_denoise", + assumes_precision="float32", + consistency_required=ConsistencyLevel.C2), + parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1, ConsistencyLevel.C2], 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() diff --git a/v2/recipes/omni/__init__.py b/v2/recipes/omni/__init__.py index 1090fa191..2aba947aa 100644 --- a/v2/recipes/omni/__init__.py +++ b/v2/recipes/omni/__init__.py @@ -1,9 +1,8 @@ -"""Shared omni building blocks (design_v3 §4): the AR decode loop + pack/emit program nodes.""" +"""Shared omni building blocks: the AR decode loop + pack/emit program nodes.""" from __future__ import annotations from v2.recipes.omni.ar_loop import ARDecodeLoop from v2.recipes.omni.pack import emit_text_node, pack_cond_from_tokens, tokenize_node, vae_decode_node from v2.recipes.omni.vocoder_loop import VocoderLoop -__all__ = ["ARDecodeLoop", "VocoderLoop", "tokenize_node", "pack_cond_from_tokens", - "emit_text_node", "vae_decode_node"] +__all__ = ["ARDecodeLoop", "VocoderLoop", "tokenize_node", "pack_cond_from_tokens", "emit_text_node", "vae_decode_node"] diff --git a/v2/recipes/omni/ar_loop.py b/v2/recipes/omni/ar_loop.py index a01ab2e99..077097682 100644 --- a/v2/recipes/omni/ar_loop.py +++ b/v2/recipes/omni/ar_loop.py @@ -1,10 +1,10 @@ -"""ARDecodeLoop — token decode until EOS/max_tokens (design_v3 §4.3, §5; §15b Cosmos3 reasoner). +"""ARDecodeLoop — token decode until EOS/max_tokens. The autoregressive pathway: ``init`` prefills the prompt tokens (and allocates paged text-KV); ``next`` emits one ``ar_token`` WorkUnit (kernel-free); ``advance`` appends the sampled token and streams it; ``done`` on EOS/max_tokens. It binds the SAME ``transformer`` component as the omni -model's ``diffusion_denoise`` loop — one resident MoT instance running both loop types (the §16 -claim). All per-request state (the token list, paged-KV namespace) lives in ``LoopState``. +model's ``diffusion_denoise`` loop — one resident MoT instance running both loop types. All +per-request state (the token list, paged-KV namespace) lives in ``LoopState``. """ from __future__ import annotations @@ -22,24 +22,27 @@ from v2.request.streams import StreamChunk class ARDecodeLoop: - def __init__(self, *, loop_id, transformer_id="transformer", cost, max_tokens=8, - prompt_slot="prompt_tokens"): + + def __init__(self, *, loop_id, transformer_id="transformer", cost, max_tokens=8, prompt_slot="prompt_tokens"): self.loop_id = loop_id self.transformer_id = transformer_id self.cost = cost self.max_tokens = max_tokens - self.prompt_slot = prompt_slot # which slot holds the prefill (chained AR loops differ) + self.prompt_slot = prompt_slot # which slot holds the prefill (chained AR loops differ) def init(self, req, model, ctx) -> LoopState: seed = req.sampling.seed if req.sampling.seed is not None else 0 - st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id, - request_id=req.request_id, profile=ctx.profile, seed=seed) + st = LoopState(loop_id=self.loop_id, + instance_id=model.card.model_id, + request_id=req.request_id, + profile=ctx.profile, + seed=seed) prompt_tokens = ctx.slots.get(self.prompt_slot) or [1] - st.scratch["tokens"] = list(prompt_tokens) # prefill (folded into init for the toy) + st.scratch["tokens"] = list(prompt_tokens) # prefill (folded into init for the toy) st.scratch["generated"] = [] st.scratch["max_tokens"] = min(req.sampling.max_tokens, self.max_tokens) caches = getattr(model, "caches", None) - if caches is not None and caches.has("paged_kv"): # reasoner paged text-KV (minority case) + if caches is not None and caches.has("paged_kv"): # reasoner paged text-KV (minority case) caches.pool("paged_kv").allocate(req.request_id, n_blocks=1) return st @@ -50,17 +53,25 @@ class ARDecodeLoop: tid = self.transformer_id def run(model, override=None): - nxt = model.component(tid).ar_forward(tokens) # und pathway of the shared MoT weights + nxt = model.component(tid).ar_forward(tokens) # und pathway of the shared MoT weights return StepResult(output={"token": int(nxt)}) pos = len(st.scratch["generated"]) emits = [StreamChunk(stream_id=st.request_id, modality="text", seq=pos, preview=True)] - return WorkPlan( - loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.AR_TOKEN, - shape_sig=ShapeSignature(WorkUnitKind.AR_TOKEN, dims=(len(tokens),)), - resources=ResourceRequest(compute_seconds=self.cost.predict(len(tokens)), - resident_bytes=8 * len(tokens), peak_activation_bytes=8), - payload={"branch": "ar", "pos": pos}, run=run, emits=emits, label=f"ar.tok{pos}") + return WorkPlan(loop_id=self.loop_id, + instance_id=st.instance_id, + kind=WorkUnitKind.AR_TOKEN, + shape_sig=ShapeSignature(WorkUnitKind.AR_TOKEN, dims=(len(tokens), )), + resources=ResourceRequest(compute_seconds=self.cost.predict(len(tokens)), + resident_bytes=8 * len(tokens), + peak_activation_bytes=8), + payload={ + "branch": "ar", + "pos": pos + }, + run=run, + emits=emits, + label=f"ar.tok{pos}") def advance(self, st: LoopState, result: StepResult) -> LoopState: tok = int(result.output["token"]) @@ -73,6 +84,9 @@ class ARDecodeLoop: def finalize(self, st: LoopState) -> LoopResult: toks = st.scratch["generated"] - return LoopResult(outputs={"tokens": toks, "text": "tok:" + ",".join(map(str, toks))}, + return LoopResult(outputs={ + "tokens": toks, + "text": "tok:" + ",".join(map(str, toks)) + }, metrics={"tokens": float(len(toks))}, behavior=st.trajectory or None) diff --git a/v2/recipes/omni/pack.py b/v2/recipes/omni/pack.py index 4971dba28..afa6fcc19 100644 --- a/v2/recipes/omni/pack.py +++ b/v2/recipes/omni/pack.py @@ -1,9 +1,8 @@ -"""Omni program nodes shared by Cosmos3 and the bagel/lance MoT (design_v3 §4.1, §15b). +"""Omni program nodes shared by Cosmos3 and the bagel/lance MoT. The ``pack`` node is the seam where the AR (understanding) pathway's tokens become the -conditioning the diffusion (generation) pathway consumes — "the AR reasoner ... used as prompt -upsampling before diffusion in the same request". Both pathways are the SAME resident MoT -module, so this is weight-sharing, not a hand-off between two engines. +conditioning the diffusion (generation) pathway consumes. Both pathways are the SAME resident +MoT module, so this is weight-sharing within a request, not a hand-off between two engines. """ from __future__ import annotations @@ -23,23 +22,29 @@ def _tokens_of(slot_value) -> list: def pack_cond_from_tokens(token_slot: str): """Build a ComponentNode fn that packs the und-pathway tokens into the gen-pathway conditioning.""" + def fn(instance, slots, request, ctx) -> None: - mot = instance.component("transformer") # the SHARED MoT module (und + gen pathways) + mot = instance.component("transformer") # the SHARED MoT module (und + gen pathways) toks = _tokens_of(slots.get(token_slot)) slots["text_embeds"] = mot.reasoner_embed(toks) slots["neg_text_embeds"] = mot.reasoner_embed([0]) + return fn def emit_text_node(token_slot: str, out_slot: str = "text"): + def fn(instance, slots, request, ctx) -> None: out = slots.get(token_slot, {}) slots[out_slot] = out.get("text", "") if isinstance(out, dict) else str(out) + return fn def vae_decode_node(latent_slot: str, out_slot: str): + def fn(instance, slots, request, ctx) -> None: latents = slots[latent_slot]["latents"] if isinstance(slots.get(latent_slot), dict) else None slots[out_slot] = instance.component("vae").decode(latents) if latents is not None else None + return fn diff --git a/v2/recipes/omni/vocoder_loop.py b/v2/recipes/omni/vocoder_loop.py index 0b266e9d3..0645e664c 100644 --- a/v2/recipes/omni/vocoder_loop.py +++ b/v2/recipes/omni/vocoder_loop.py @@ -1,11 +1,10 @@ -"""VocoderLoop — streaming code2wav (design_v3 §4.3 ``LoopKind.AUDIO_DECODE``; Qwen-Omni stage 2). +"""VocoderLoop — streaming code2wav (``LoopKind.AUDIO_DECODE``; Qwen-Omni stage 2). The third loop type in the thinker→talker→vocoder cascade. It consumes the talker's speech-codec -tokens in fixed-size chunks; **each chunk is one ``AUDIO_CHUNK`` WorkUnit** that synthesizes a -waveform segment and streams it. This is the differentiation vllm-omni's own docs name: their -code2wav is one opaque ``LLM_GENERATION`` stage the scheduler never sees inside — here every chunk -is a runtime-visible, step-scheduled, cancellable WorkUnit, on the *same* driven-loop contract the -engine uses for denoise steps and AR tokens. All per-request state lives in ``LoopState``. +tokens in fixed-size chunks; each chunk is one ``AUDIO_CHUNK`` WorkUnit that synthesizes and streams +a waveform segment. Unlike vllm-omni's opaque code2wav stage, every chunk here is a runtime-visible, +step-scheduled, cancellable WorkUnit on the same driven-loop contract used for denoise steps and AR +tokens. All per-request state lives in ``LoopState``. """ from __future__ import annotations @@ -25,6 +24,7 @@ from v2.request.streams import StreamChunk class VocoderLoop: + def __init__(self, *, loop_id, vocoder_id="vocoder", cost, chunk_tokens=2, speech_slot="speech_tokens"): self.loop_id = loop_id self.vocoder_id = vocoder_id @@ -34,11 +34,12 @@ class VocoderLoop: def init(self, req, model, ctx) -> LoopState: toks = list(ctx.slots.get(self.speech_slot) or [1]) - st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id, - request_id=req.request_id, profile=ctx.profile) + st = LoopState(loop_id=self.loop_id, + instance_id=model.card.model_id, + request_id=req.request_id, + profile=ctx.profile) st.scratch["tokens"] = toks - st.scratch["chunks"] = [toks[i:i + self.chunk_tokens] - for i in range(0, len(toks), self.chunk_tokens)] or [[1]] + st.scratch["chunks"] = [toks[i:i + self.chunk_tokens] for i in range(0, len(toks), self.chunk_tokens)] or [[1]] st.scratch["wave"] = [] st.scratch["stream_audio"] = bool(req.outputs.stream.get("audio")) return st @@ -52,19 +53,23 @@ class VocoderLoop: vid = self.vocoder_id def run(model, override=None): - seg = model.component(vid).synthesize(chunk) # codec tokens → waveform segment + seg = model.component(vid).synthesize(chunk) # codec tokens → waveform segment return StepResult(output={"segment": np.asarray(seg, dtype="float32")}) emits = [] if st.scratch.get("stream_audio"): emits.append(StreamChunk(stream_id=st.request_id, modality="audio", seq=i, preview=False)) - return WorkPlan( - loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.AUDIO_CHUNK, - shape_sig=ShapeSignature(WorkUnitKind.AUDIO_CHUNK, dims=(len(chunk),)), - resources=ResourceRequest(compute_seconds=self.cost.predict(len(chunk)), - resident_bytes=8 * len(st.scratch["tokens"]), - peak_activation_bytes=64), - payload={"chunk": i}, run=run, emits=emits, label=f"vocoder.chunk{i}") + return WorkPlan(loop_id=self.loop_id, + instance_id=st.instance_id, + kind=WorkUnitKind.AUDIO_CHUNK, + shape_sig=ShapeSignature(WorkUnitKind.AUDIO_CHUNK, dims=(len(chunk), )), + resources=ResourceRequest(compute_seconds=self.cost.predict(len(chunk)), + resident_bytes=8 * len(st.scratch["tokens"]), + peak_activation_bytes=64), + payload={"chunk": i}, + run=run, + emits=emits, + label=f"vocoder.chunk{i}") def advance(self, st: LoopState, result: StepResult) -> LoopState: st.scratch["wave"].append(np.asarray(result.output["segment"], dtype="float32")) @@ -76,6 +81,8 @@ class VocoderLoop: def finalize(self, st: LoopState) -> LoopResult: wave = np.concatenate(st.scratch["wave"]) if st.scratch["wave"] else np.zeros(0, dtype="float32") return LoopResult(outputs={"samples": wave.astype("float32")}, - metrics={"audio_chunks": float(st.step_idx), - "audio_samples": float(wave.size)}, + metrics={ + "audio_chunks": float(st.step_idx), + "audio_samples": float(wave.size) + }, behavior=st.trajectory or None) diff --git a/v2/recipes/qwen_omni/__init__.py b/v2/recipes/qwen_omni/__init__.py index 586d7fef2..9db0b7efd 100644 --- a/v2/recipes/qwen_omni/__init__.py +++ b/v2/recipes/qwen_omni/__init__.py @@ -1,6 +1,6 @@ """Qwen-Omni thinker→talker→vocoder model (vllm-omni ``qwen2_5_omni``): three disjoint experts, three loops (ar_decode → ar_decode → audio_decode), cascaded conditioning + streaming codec→waveform. The -third weight-sharing topology proving the Card/Loop/Program vocabulary holds (design_v3 §4).""" +third weight-sharing topology in the Card/Loop/Program vocabulary.""" from __future__ import annotations from v2.recipes.qwen_omni.card import build_qwen_omni_card diff --git a/v2/recipes/qwen_omni/card.py b/v2/recipes/qwen_omni/card.py index 0bbfa0cd2..ce8124f28 100644 --- a/v2/recipes/qwen_omni/card.py +++ b/v2/recipes/qwen_omni/card.py @@ -1,21 +1,20 @@ -"""Qwen-Omni thinker→talker→vocoder ModelCard (design_v3 §4; vllm-omni ``qwen2_5_omni`` pipeline). +"""Qwen-Omni thinker→talker→vocoder ModelCard (vllm-omni ``qwen2_5_omni`` pipeline). -The canonical vllm-omni 3-stage omni-speech model, made model-native. vllm-omni expresses it as three -opaque stages (``thinker``/``talker``/``code2wav``) with ``custom_process_input_func`` hand-offs, each -request-scheduled, the scheduler never seeing inside the AR generation or the code2wav synthesis. v2 -expresses the SAME cascade as three driven loops on three components — and every thinker token, talker -token, and vocoder chunk is a runtime-visible WorkUnit (the differentiation §1 names). +vllm-omni expresses this as three opaque stages (``thinker``/``talker``/``code2wav``) with +``custom_process_input_func`` hand-offs, request-scheduled, the scheduler never seeing inside. v2 +expresses the same cascade as three driven loops on three components, where every thinker token, +talker token, and vocoder chunk is a runtime-visible WorkUnit. -It is a *third* weight-sharing topology, in the same Card/Loop/Program vocabulary: +A third weight-sharing topology, in the same Card/Loop/Program vocabulary: * Cosmos3/BAGEL — ONE shared MoT module, two loops (weight sharing within a request). * unified/UniRL — TWO disjoint experts, two loops, jointly RL-trained. - * Qwen-Omni — THREE disjoint experts (thinker, talker, vocoder), three loops, **cascaded**: + * Qwen-Omni — THREE disjoint experts (thinker, talker, vocoder), three loops, cascaded: the talker conditions on the thinker's tokens+hidden state, the vocoder on the talker's speech tokens; streaming codec→waveform. -The loop kinds were anticipated: ``LoopKind.AR_DECODE`` (its docstring names "thinker/talker") for the -two language stages, ``LoopKind.AUDIO_DECODE`` for the vocoder — no new primitive. +Reuses existing loop kinds: ``LoopKind.AR_DECODE`` for the two language stages, +``LoopKind.AUDIO_DECODE`` for the vocoder. """ from __future__ import annotations @@ -43,63 +42,99 @@ def build_qwen_omni_card(model_id: str = "qwen-omni-tts") -> ModelCard: wav_cost = CostModel(kind=WorkUnitKind.AUDIO_CHUNK, base_seconds=3e-5, per_unit_seconds=2e-7) def thinker_factory(): - return ARDecodeLoop(loop_id="thinker_decode", transformer_id="thinker", cost=ar_cost, - max_tokens=4, prompt_slot="prompt_tokens") + return ARDecodeLoop(loop_id="thinker_decode", + transformer_id="thinker", + cost=ar_cost, + max_tokens=4, + prompt_slot="prompt_tokens") def talker_factory(): - return ARDecodeLoop(loop_id="talker_decode", transformer_id="talker", cost=ar_cost, - max_tokens=4, prompt_slot="talker_prompt_tokens") + return ARDecodeLoop(loop_id="talker_decode", + transformer_id="talker", + cost=ar_cost, + max_tokens=4, + prompt_slot="talker_prompt_tokens") def vocoder_factory(): - return VocoderLoop(loop_id="vocoder", vocoder_id="vocoder", cost=wav_cost, - chunk_tokens=2, speech_slot="speech_tokens") + return VocoderLoop(loop_id="vocoder", + vocoder_id="vocoder", + cost=wav_cost, + chunk_tokens=2, + speech_slot="speech_tokens") components = { - "tokenizer": ComponentSpec("tokenizer", kind="tokenizer", factory=lambda inst: ToyTokenizer(), - required_for={"reason", "t2a"}), - "thinker": ComponentSpec( # stage 0: multimodal understanding + text - "thinker", kind="dit", + "tokenizer": + ComponentSpec("tokenizer", + kind="tokenizer", + factory=lambda inst: ToyTokenizer(), + required_for={"reason", "t2a"}), + "thinker": + ComponentSpec( # stage 0: multimodal understanding + text + "thinker", + kind="dit", load_id="vllm_omni.model_executor.models.qwen2_5_omni:Thinker", factory=lambda inst: ToyMoTDiT(seed=seed), - resident_for=["thinker_decode"], required_for={"reason", "t2a"}), - "talker": ComponentSpec( # stage 1: text+hidden → speech tokens - "talker", kind="dit", + resident_for=["thinker_decode"], + required_for={"reason", "t2a"}), + "talker": + ComponentSpec( # stage 1: text+hidden → speech tokens + "talker", + kind="dit", load_id="vllm_omni.model_executor.models.qwen2_5_omni:Talker", factory=lambda inst: ToyTalker(seed=seed + 1), - resident_for=["talker_decode"], required_for={"t2a"}), - "vocoder": ComponentSpec( # stage 2: speech tokens → waveform - "vocoder", kind="audio_vae", + resident_for=["talker_decode"], + required_for={"t2a"}), + "vocoder": + ComponentSpec( # stage 2: speech tokens → waveform + "vocoder", + kind="audio_vae", load_id="vllm_omni.model_executor.models.qwen2_5_omni:Code2Wav", factory=lambda inst: ToyVocoder(seed=seed + 2), - resident_for=["vocoder"], required_for={"t2a"}), + resident_for=["vocoder"], + required_for={"t2a"}), } loops = { - "thinker_decode": LoopSpec("thinker_decode", kind=LoopKind.AR_DECODE, - work_unit_kind=WorkUnitKind.AR_TOKEN, step_cost_model=ar_cost, - shared_weight_components=["thinker"], cache_policy=["paged_kv"], - loop_factory=thinker_factory), - "talker_decode": LoopSpec("talker_decode", kind=LoopKind.AR_DECODE, - work_unit_kind=WorkUnitKind.AR_TOKEN, step_cost_model=ar_cost, - shared_weight_components=["talker"], cache_policy=["paged_kv"], - loop_factory=talker_factory), - "vocoder": LoopSpec("vocoder", kind=LoopKind.AUDIO_DECODE, - work_unit_kind=WorkUnitKind.AUDIO_CHUNK, step_cost_model=wav_cost, - shared_weight_components=["vocoder"], cache_policy=[], - loop_factory=vocoder_factory), + "thinker_decode": + LoopSpec("thinker_decode", + kind=LoopKind.AR_DECODE, + work_unit_kind=WorkUnitKind.AR_TOKEN, + step_cost_model=ar_cost, + shared_weight_components=["thinker"], + cache_policy=["paged_kv"], + loop_factory=thinker_factory), + "talker_decode": + LoopSpec("talker_decode", + kind=LoopKind.AR_DECODE, + work_unit_kind=WorkUnitKind.AR_TOKEN, + step_cost_model=ar_cost, + shared_weight_components=["talker"], + cache_policy=["paged_kv"], + loop_factory=talker_factory), + "vocoder": + LoopSpec("vocoder", + kind=LoopKind.AUDIO_DECODE, + work_unit_kind=WorkUnitKind.AUDIO_CHUNK, + step_cost_model=wav_cost, + shared_weight_components=["vocoder"], + cache_policy=[], + loop_factory=vocoder_factory), } card = ModelCard( - model_id=model_id, family="qwen_omni", components=components, loops=loops, - capabilities=CapabilityMatrix.of( - Capability.REASONING_TEXT, Capability.TEXT_TO_SPEECH, Capability.POLICY_ROLLOUT), - recipe=RecipeSpec(method="base", assumes_loop="vocoder", - assumes_precision="float32", consistency_required=ConsistencyLevel.C1), + model_id=model_id, + family="qwen_omni", + components=components, + loops=loops, + capabilities=CapabilityMatrix.of(Capability.REASONING_TEXT, Capability.TEXT_TO_SPEECH, + Capability.POLICY_ROLLOUT), + recipe=RecipeSpec(method="base", + assumes_loop="vocoder", + assumes_precision="float32", + consistency_required=ConsistencyLevel.C1), parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True), caches={ - "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() diff --git a/v2/recipes/qwen_omni/program.py b/v2/recipes/qwen_omni/program.py index 811358b8c..6d156f406 100644 --- a/v2/recipes/qwen_omni/program.py +++ b/v2/recipes/qwen_omni/program.py @@ -7,11 +7,10 @@ → vocoder (audio_decode) → emit_audio The hand-off nodes are the model-native analogue of vllm-omni's ``custom_process_input_func`` -(``thinker2talker_full_payload``, ``talker2code2wav_full_payload``) — but here they are explicit -slot writes between three driven loops, not opaque stage-boundary callbacks. The talker conditions on -the thinker's **hidden state** (``reasoner_embed``), not only its tokens — the "full payload" path; -the streaming/sync path would pass tokens only (noted inline). Final outputs: ``text`` (thinker) and -``audio`` (vocoder waveform). +(``thinker2talker_full_payload``, ``talker2code2wav_full_payload``), but here they are explicit slot +writes between three driven loops, not opaque stage-boundary callbacks. The talker conditions on the +thinker's hidden state (``reasoner_embed``), not only its tokens — the "full payload" path; the +streaming/sync path would pass tokens only. Final outputs: ``text`` (thinker) and ``audio`` (vocoder). """ from __future__ import annotations @@ -27,8 +26,8 @@ def thinker_to_talker_node(instance, slots, request, ctx) -> None: would write only the tokens.)""" out = slots.get("thinker_out", {}) toks = out.get("tokens", []) if isinstance(out, dict) else [] - hidden = instance.component("thinker").reasoner_embed(toks or [1]) # thinker hidden state - h = int(abs(float(np.sum(hidden))) * 1000.0) % 251 + 1 # one hidden-derived token + hidden = instance.component("thinker").reasoner_embed(toks or [1]) # thinker hidden state + h = int(abs(float(np.sum(hidden))) * 1000.0) % 251 + 1 # one hidden-derived token slots["talker_prompt_tokens"] = list(toks) + [h] @@ -39,31 +38,55 @@ def talker_to_vocoder_node(instance, slots, request, ctx) -> None: def emit_audio_node(loop_slot: str, out_slot: str = "audio"): + def fn(instance, slots, request, ctx) -> None: out = slots.get(loop_slot, {}) slots[out_slot] = out.get("samples") if isinstance(out, dict) else out + return fn def build_qwen_omni_program() -> Program: return Program( - program_id="qwen_omni.tts", kind=ProgramKind.INLINE, + program_id="qwen_omni.tts", + kind=ProgramKind.INLINE, nodes=[ - ComponentNode("tokenize", fn=tokenize_node, writes=("prompt_tokens",)), - ModelLoopNode("thinker", loop_id="thinker_decode", output_slot="thinker_out", - reads=("prompt_tokens",), writes=("thinker_out",)), - ComponentNode("emit_text", fn=emit_text_node("thinker_out", "text"), - reads=("thinker_out",), writes=("text",)), - ComponentNode("thinker_to_talker", fn=thinker_to_talker_node, - reads=("thinker_out",), writes=("talker_prompt_tokens",)), - ModelLoopNode("talker", loop_id="talker_decode", output_slot="talker_out", - reads=("talker_prompt_tokens",), writes=("talker_out",)), - ComponentNode("talker_to_vocoder", fn=talker_to_vocoder_node, - reads=("talker_out",), writes=("speech_tokens",)), - ModelLoopNode("vocoder", loop_id="vocoder", output_slot="vocoder_out", - reads=("speech_tokens",), writes=("vocoder_out",)), - ComponentNode("emit_audio", fn=emit_audio_node("vocoder_out", "audio"), - reads=("vocoder_out",), writes=("audio",)), + ComponentNode("tokenize", fn=tokenize_node, writes=("prompt_tokens", )), + ModelLoopNode("thinker", + loop_id="thinker_decode", + output_slot="thinker_out", + reads=("prompt_tokens", ), + writes=("thinker_out", )), + ComponentNode("emit_text", + fn=emit_text_node("thinker_out", "text"), + reads=("thinker_out", ), + writes=("text", )), + ComponentNode("thinker_to_talker", + fn=thinker_to_talker_node, + reads=("thinker_out", ), + writes=("talker_prompt_tokens", )), + ModelLoopNode("talker", + loop_id="talker_decode", + output_slot="talker_out", + reads=("talker_prompt_tokens", ), + writes=("talker_out", )), + ComponentNode("talker_to_vocoder", + fn=talker_to_vocoder_node, + reads=("talker_out", ), + writes=("speech_tokens", )), + ModelLoopNode("vocoder", + loop_id="vocoder", + output_slot="vocoder_out", + reads=("speech_tokens", ), + writes=("vocoder_out", )), + ComponentNode("emit_audio", + fn=emit_audio_node("vocoder_out", "audio"), + reads=("vocoder_out", ), + writes=("audio", )), ], - output_artifacts={"text": "text", "audio": "audio", "speech_tokens": "talker_out"}, + output_artifacts={ + "text": "text", + "audio": "audio", + "speech_tokens": "talker_out" + }, ).validate() diff --git a/v2/recipes/reward/__init__.py b/v2/recipes/reward/__init__.py index 64e7e8230..384890b78 100644 --- a/v2/recipes/reward/__init__.py +++ b/v2/recipes/reward/__init__.py @@ -1,4 +1,4 @@ -"""Reward-model card: a served scorer/verifier emitting REWARD_BATCH work units (design_v3 §10).""" +"""Reward-model card: a served scorer/verifier emitting REWARD_BATCH work units.""" from __future__ import annotations from v2.recipes.reward.card import build_reward_card diff --git a/v2/recipes/reward/card.py b/v2/recipes/reward/card.py index 7859ae7b4..05f64dbb1 100644 --- a/v2/recipes/reward/card.py +++ b/v2/recipes/reward/card.py @@ -1,9 +1,8 @@ -"""Reward-model ModelCard — a served scorer/verifier for RL (design_v3 §10). +"""Reward-model ModelCard — a served scorer/verifier for RL. A reward model is just another card: a `scorer` component + a `score` loop emitting `REWARD_BATCH` -work units. So a learned reward (PickScore / CLIP / a VLM judge) is loaded, placed, priced, and -scheduled like any model — enabling RLHF/RLAIF where the reward is a model, not a heuristic, and can -run on its own role-pool (reward disaggregation). +work units. A learned reward (PickScore / CLIP / a VLM judge) is loaded, placed, priced, and scheduled +like any model, and can run on its own role-pool (reward disaggregation). """ from __future__ import annotations @@ -32,24 +31,36 @@ def build_reward_card(model_id: str = "pickscore-reward", *, batch: int = 4) -> return RewardLoop(loop_id="score", scorer_id="scorer", cost=cost, media_slot="media", batch=batch) components = { - "scorer": ComponentSpec("scorer", kind="reward_model", - load_id="fastvideo.rl.rewards:PickScore", - factory=lambda inst: ToyRewardModel(seed=seed), - resident_for=["score"], required_for={"reason"}), + "scorer": + ComponentSpec("scorer", + kind="reward_model", + load_id="fastvideo.rl.rewards:PickScore", + factory=lambda inst: ToyRewardModel(seed=seed), + resident_for=["score"], + required_for={"reason"}), } loops = { - "score": LoopSpec("score", kind=LoopKind.ENCODER, work_unit_kind=WorkUnitKind.REWARD_BATCH, - step_cost_model=cost, shared_weight_components=["scorer"], cache_policy=[], - loop_factory=score_factory), + "score": + LoopSpec("score", + kind=LoopKind.ENCODER, + work_unit_kind=WorkUnitKind.REWARD_BATCH, + step_cost_model=cost, + shared_weight_components=["scorer"], + cache_policy=[], + loop_factory=score_factory), } return ModelCard( - model_id=model_id, family="reward", components=components, loops=loops, + model_id=model_id, + family="reward", + components=components, + loops=loops, capabilities=CapabilityMatrix.of(Capability.POLICY_ROLLOUT), - recipe=RecipeSpec(method="base", assumes_loop="score", - assumes_precision="float32", consistency_required=ConsistencyLevel.C1), + recipe=RecipeSpec(method="base", + assumes_loop="score", + assumes_precision="float32", + consistency_required=ConsistencyLevel.C1), parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True), caches={}, 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() diff --git a/v2/recipes/reward/loop.py b/v2/recipes/reward/loop.py index 3777652b5..70a584d82 100644 --- a/v2/recipes/reward/loop.py +++ b/v2/recipes/reward/loop.py @@ -1,10 +1,9 @@ -"""RewardLoop — a served reward model scoring samples as ``REWARD_BATCH`` WorkUnits (design_v3 §10). +"""RewardLoop — a served reward model scoring samples as ``REWARD_BATCH`` WorkUnits. -The reward in RL is computed by a *model* (PickScore/CLIP/a VLM verifier), not a heuristic. This loop -makes that a first-class part of the runtime: it reads a batch of K rollout media from a slot and emits -``REWARD_BATCH`` work units (chunked) that score them through the engine — so reward compute is admitted, -priced, and interleavable exactly like a denoise step or an AR token. A GRPO group's K samples score in -one (or a few) batched units — the homogeneous-batch win the scheduler likes. +The RL reward is computed by a *model* (PickScore/CLIP/a VLM verifier), not a heuristic. This loop reads +a batch of K rollout media from a slot and emits chunked ``REWARD_BATCH`` work units that score them +through the engine, so reward compute is admitted, priced, and interleaved like a denoise step or an AR +token. A GRPO group's K samples score in one (or a few) batched units. """ from __future__ import annotations @@ -21,6 +20,7 @@ from v2.loop.contracts import ( class RewardLoop: + def __init__(self, *, loop_id, scorer_id="scorer", cost, media_slot="media", batch=4): self.loop_id = loop_id self.scorer_id = scorer_id @@ -30,8 +30,10 @@ class RewardLoop: def init(self, req, model, ctx) -> LoopState: media = list(ctx.slots.get(self.media_slot) or []) - st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id, - request_id=req.request_id, profile=ctx.profile) + st = LoopState(loop_id=self.loop_id, + instance_id=model.card.model_id, + request_id=req.request_id, + profile=ctx.profile) st.scratch["batches"] = [media[i:i + self.batch] for i in range(0, len(media), self.batch)] or [[]] st.scratch["scores"] = [] return st @@ -47,12 +49,16 @@ class RewardLoop: scorer = model.component(sid) return StepResult(output={"scores": [float(scorer.score(m)) for m in chunk], "i": i}) - return WorkPlan( - loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.REWARD_BATCH, - shape_sig=ShapeSignature(WorkUnitKind.REWARD_BATCH, dims=(len(chunk),)), - resources=ResourceRequest(compute_seconds=self.cost.predict(len(chunk)), - resident_bytes=64 * max(1, len(chunk)), peak_activation_bytes=64), - payload={"reward_batch": i}, run=run, label=f"reward.batch{i}") + return WorkPlan(loop_id=self.loop_id, + instance_id=st.instance_id, + kind=WorkUnitKind.REWARD_BATCH, + shape_sig=ShapeSignature(WorkUnitKind.REWARD_BATCH, dims=(len(chunk), )), + resources=ResourceRequest(compute_seconds=self.cost.predict(len(chunk)), + resident_bytes=64 * max(1, len(chunk)), + peak_activation_bytes=64), + payload={"reward_batch": i}, + run=run, + label=f"reward.batch{i}") def advance(self, st: LoopState, result: StepResult) -> LoopState: st.scratch["scores"].extend(result.output["scores"]) diff --git a/v2/recipes/sd35/__init__.py b/v2/recipes/sd35/__init__.py index c38b7cea3..64cad5a0f 100644 --- a/v2/recipes/sd35/__init__.py +++ b/v2/recipes/sd35/__init__.py @@ -1,11 +1,11 @@ """Stable Diffusion 3.5 Medium (text→image) — a FLOW-MATCH MMDiT ported into the v2 substrate. -Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via -``ComponentSpec.adapter`` (``SD3DiT``/``SD3VAE``/``SD3ClipEncoder``/``SD3T5Encoder`` in -``v2/platform/backends/torch_sd35.py``) and a new ``SD3DenoiseLoop`` (flow-match Euler over 4D image -latents with dual text conditioning: the triple-encoder joint embed + the dual-CLIP pooled vector). It -reuses the ``FLOW_MATCH_STEP`` kernel + ``FlowShiftPolicy`` + ``ClassicCFG``. Registered in ``v2/registry.py`` -by the orchestrator (transformer ``_class_name`` = ``SD3Transformer2DModel`` for the arch fallback). +Self-contained recipe package: the card declares its torch adapters via ``ComponentSpec.adapter`` +(``SD3DiT``/``SD3VAE``/``SD3ClipEncoder``/``SD3T5Encoder`` in ``v2/platform/backends/torch_sd35.py``) +and ``SD3DenoiseLoop`` (flow-match Euler over 4D image latents with dual text conditioning: the +triple-encoder joint embed + the dual-CLIP pooled vector). Reuses ``FLOW_MATCH_STEP`` + ``FlowShiftPolicy`` ++ ``ClassicCFG``. Registered in ``v2/registry.py`` (transformer ``_class_name`` = ``SD3Transformer2DModel`` +for the arch fallback). """ from __future__ import annotations diff --git a/v2/recipes/sd35/card.py b/v2/recipes/sd35/card.py index 1c273f7b4..13948b8a4 100644 --- a/v2/recipes/sd35/card.py +++ b/v2/recipes/sd35/card.py @@ -2,11 +2,10 @@ 1-frame video (num_frames=1, 4D latents) throughout the v2 substrate. Architecture deltas vs Wan (all declared on the card so the recipe is self-contained): - * DiT ``fastvideo.models.dits.sd3:SD3Transformer2DModel`` — a FLOW-MATCH MMDiT (reuses the - flow-match step), but the forward needs TWO text conditioners: the assembled triple-encoder joint - embed (``encoder_hidden_states``) AND a separate ``pooled_projections`` vector. The adapter - (``ComponentSpec.adapter`` -> ``SD3DiT``) takes both (joint embed as ``text_embed``, pooled as the - positional ``context`` arg) and ``SD3DenoiseLoop`` threads them per CFG branch. + * DiT ``fastvideo.models.dits.sd3:SD3Transformer2DModel`` — a FLOW-MATCH MMDiT, but the forward needs + TWO text conditioners: the assembled triple-encoder joint embed (``encoder_hidden_states``) AND a + separate ``pooled_projections`` vector. The ``SD3DiT`` adapter takes both (joint embed as + ``text_embed``, pooled as the positional ``context`` arg) and ``SD3DenoiseLoop`` threads them per CFG branch. * VAE ``fastvideo.models.vaes.autoencoder_kl:AutoencoderKL`` — SD3.5 shift/scale normalization (``scaling_factor=1.5305``, ``shift_factor=0.0609`` read from the checkpoint vae config); the ``SD3VAE`` adapter normalizes/denormalizes around the raw AutoencoderKL space (4D image, not 5D video). diff --git a/v2/recipes/sd35/loop.py b/v2/recipes/sd35/loop.py index 459aa5d76..54e9cdbca 100644 --- a/v2/recipes/sd35/loop.py +++ b/v2/recipes/sd35/loop.py @@ -1,7 +1,7 @@ """SD3DenoiseLoop — a thin flow-match Euler denoise loop for Stable Diffusion 3.5 (image / 4D latents). SD3.5 is a plain FLOW-MATCH MMDiT, so this loop reuses the canonical pieces (``FLOW_MATCH_STEP`` kernel, -``FlowShiftPolicy.build_schedule``, ``ClassicCFG``) — the only deltas vs ``WanDenoiseLoop`` are: +``FlowShiftPolicy.build_schedule``, ``ClassicCFG``). Deltas vs ``WanDenoiseLoop``: * IMAGE modality / 4D latents: the latent is ``[16, h//8, w//8]`` (num_frames=1), with NO temporal dim. The fastvideo stage carries a fake T dim in its 5D ForwardBatch (unsqueeze(2)/squeeze(2)) but @@ -15,8 +15,8 @@ SD3.5 is a plain FLOW-MATCH MMDiT, so this loop reuses the canonical pieces (``F Euler update is the plain ``x_next = x + (sigma_next - sigma)*v``. Faithful to ``fastvideo/pipelines/stages/sd35_conditioning.py:SD35DenoisingStage`` (uncond + s·(cond − -uncond) CFG over the FlowMatchEuler schedule). The σ schedule is the flow-shift linspace; ``shift`` comes -from the diffusers ``FlowMatchEulerDiscreteScheduler`` config (3.0 for SD3.5-medium) — see ``card.py``. +uncond) CFG over the FlowMatchEuler schedule). The σ schedule is the flow-shift linspace; ``shift`` is +the diffusers ``FlowMatchEulerDiscreteScheduler`` config (3.0 for SD3.5-medium) — see ``card.py``. """ from __future__ import annotations diff --git a/v2/recipes/sd35/program.py b/v2/recipes/sd35/program.py index d655224bd..66cac8991 100644 --- a/v2/recipes/sd35/program.py +++ b/v2/recipes/sd35/program.py @@ -1,16 +1,16 @@ -"""SD3.5 t2i program: sd3_text_encode → diffusion_denoise (flow-match) → vae_decode. +"""SD3.5 t2i program: sd3_text_encode -> diffusion_denoise (flow-match) -> vae_decode. -Single-stage INLINE program (like ``build_ltx2_base_program``). The SD3-specific work lives in the -custom text-encode node (the triple-encoder joint-embed + dual-CLIP pooled assembly that -``SD35ConditioningStage`` does) and the VAE decode node (the [-1,1] -> [0,1] remap of -``SD35DecodingStage``); the denoise node is the generic ``ModelLoopNode`` binding ``SD3DenoiseLoop``. +Single-stage INLINE program. The SD3-specific work lives in the custom text-encode node (the +triple-encoder joint-embed + dual-CLIP pooled assembly of ``SD35ConditioningStage``) and the VAE decode +node (the [-1,1] -> [0,1] remap of ``SD35DecodingStage``); the denoise node is the generic +``ModelLoopNode`` binding ``SD3DenoiseLoop``. The text-encode node is backend-agnostic: on the GPU torch backend the CLIP adapters (``SD3ClipEncoder``) return ``{"hidden": penultimate[-2], "pooled": pooler_output}`` and ``SD3T5Encoder`` returns ``last_hidden_state``, which ``assemble_sd3_conditioning`` combines into (joint_embed, pooled). On the CPU toy backend the encoders return plain arrays (``ToyTextEncoder.encode``), so the node falls back to a degenerate-but-finite assembly (joint = the array, pooled = its channel-mean broadcast) — enough -to exercise the dual-conditioning plumbing end-to-end without any GPU/torch. +to exercise the dual-conditioning plumbing without any GPU/torch. """ from __future__ import annotations diff --git a/v2/recipes/sfwan22/card.py b/v2/recipes/sfwan22/card.py index 2fa962e6a..3ccca25c0 100644 --- a/v2/recipes/sfwan22/card.py +++ b/v2/recipes/sfwan22/card.py @@ -3,9 +3,9 @@ This is the causal *student* of the Wan2.2 MoE produced by self-forcing distillation: * ``recipe.method='self_forcing'`` / ``assumes_loop='chunk_rollout'`` (causal/AR; ``pp_patch`` rejected); * DiT ``fastvideo.models.dits.causal_wanvideo:CausalWanTransformer3DModel`` — TWO experts - (``transformer`` high-noise + ``transformer_2`` low-noise) switched by ``boundary_ratio`` (design_v3 - §6.2.3); the GPU adapter is the existing CausalWan/Wan torch path (pure Wan arch -> ``load_id`` only, - no ``adapter=``, exactly like the wan_causal + wan21 cards); + (``transformer`` high-noise + ``transformer_2`` low-noise) switched by ``boundary_ratio``; the GPU + adapter is the existing CausalWan/Wan torch path (pure Wan arch -> ``load_id`` only, no ``adapter=``, + exactly like the wan_causal + wan21 cards); * VAE ``fastvideo.models.vaes.wanvae:AutoencoderKLWan`` (z=16, 8x/4x) and UMT5 text encoder — reused; * a ``slab_kv`` (chunk-KV) cache class in addition to the text feature cache. diff --git a/v2/recipes/sfwan22/loop.py b/v2/recipes/sfwan22/loop.py index 9ea989713..278abfabb 100644 --- a/v2/recipes/sfwan22/loop.py +++ b/v2/recipes/sfwan22/loop.py @@ -1,4 +1,4 @@ -"""SFWan22ChunkRolloutLoop — Self-Forcing Wan2.2-A14B causal DMD rollout (design_v3 §4 causal, §5, §10). +"""SFWan22ChunkRolloutLoop — Self-Forcing Wan2.2-A14B causal DMD rollout. Self-Forcing Wan2.2-A14B is the CAUSAL student of the Wan2.2 MoE: a few-step DMD denoiser run block-by-block (causal chunks, carrying per-chunk KV context) over TWO transformer experts switched by @@ -177,7 +177,7 @@ class SFWan22ChunkRolloutLoop: st.cond["prompt_embeds"] = ctx.slots.get("text_embeds") st.cond["negative_prompt_embeds"] = ctx.slots.get("neg_text_embeds") boundary_t = self.boundary_ratio * NUM_TRAIN_TIMESTEPS - # world-model continuation: seed the chunk context from prior chunks (interactive sessions, §16). + # world-model continuation: seed the chunk context from prior chunks (interactive sessions). prior = [np.asarray(c, dtype="float32") for c in (ctx.slots.get("world_context") or [])] st.scratch.update(sched_sigmas=sched_sigmas, dmd_timesteps=dmd_t, diff --git a/v2/recipes/sfwan22/program.py b/v2/recipes/sfwan22/program.py index d1b634444..08664cd0b 100644 --- a/v2/recipes/sfwan22/program.py +++ b/v2/recipes/sfwan22/program.py @@ -1,4 +1,4 @@ -"""Self-Forcing Wan2.2-A14B programs (design_v3 §4 causal, §15e). +"""Self-Forcing Wan2.2-A14B programs. * t2v: text_encode -> chunk_rollout -> vae_decode (the wan_causal streaming shape, with the SF Wan2.2 boundary-routed DMD loop bound to ``chunk_rollout``). diff --git a/v2/recipes/speculative/__init__.py b/v2/recipes/speculative/__init__.py index cc0e1cf5f..9e19fe7dc 100644 --- a/v2/recipes/speculative/__init__.py +++ b/v2/recipes/speculative/__init__.py @@ -1,4 +1,4 @@ -"""Speculative (draft-verify) decoding: a draft + target AR pair, exact + lower-latency (design_v3 §9.16).""" +"""Speculative (draft-verify) decoding: a draft + target AR pair, exact + lower-latency.""" from __future__ import annotations from v2.recipes.speculative.card import build_speculative_card diff --git a/v2/recipes/speculative/card.py b/v2/recipes/speculative/card.py index 5ab0452ff..347411a1c 100644 --- a/v2/recipes/speculative/card.py +++ b/v2/recipes/speculative/card.py @@ -1,4 +1,4 @@ -"""Speculative-decoding card — a draft + target AR pair on one instance (design_v3 §2.2, §9.16).""" +"""Speculative-decoding card — a draft + target AR pair on one instance.""" from __future__ import annotations from v2._enums import Capability, ConsistencyLevel, LoopKind, WorkUnitKind @@ -19,40 +19,61 @@ from v2.platform.backends.toy import ToyDraftModel, ToyTargetModel, ToyTokenizer from v2.recipes.speculative.loop import SpeculativeARLoop -def build_speculative_card(model_id: str = "spec-decode", *, spec_len: int = 4, - draft_agree: float = 0.7) -> ModelCard: +def build_speculative_card(model_id: str = "spec-decode", *, spec_len: int = 4, draft_agree: float = 0.7) -> ModelCard: seed = _seed_from(model_id) ar_cost = CostModel(kind=WorkUnitKind.AR_TOKEN, base_seconds=5e-5, per_unit_seconds=1e-7) def spec_factory(): - return SpeculativeARLoop(loop_id="spec_decode", draft_id="draft", target_id="target", - cost=ar_cost, spec_len=spec_len, max_tokens=12) + return SpeculativeARLoop(loop_id="spec_decode", + draft_id="draft", + target_id="target", + cost=ar_cost, + spec_len=spec_len, + max_tokens=12) components = { - "tokenizer": ComponentSpec("tokenizer", kind="tokenizer", factory=lambda inst: ToyTokenizer(), - required_for={"reason"}), - "draft": ComponentSpec("draft", kind="dit", load_id="fastvideo.models.llm:DraftModel", - factory=lambda inst: ToyDraftModel(agree=draft_agree), - resident_for=["spec_decode"], required_for={"reason"}), - "target": ComponentSpec("target", kind="dit", load_id="fastvideo.models.llm:TargetModel", - factory=lambda inst: ToyTargetModel(seed=seed), # the ground-truth AR - resident_for=["spec_decode"], required_for={"reason"}), + "tokenizer": + ComponentSpec("tokenizer", kind="tokenizer", factory=lambda inst: ToyTokenizer(), required_for={"reason"}), + "draft": + ComponentSpec("draft", + kind="dit", + load_id="fastvideo.models.llm:DraftModel", + factory=lambda inst: ToyDraftModel(agree=draft_agree), + resident_for=["spec_decode"], + required_for={"reason"}), + "target": + ComponentSpec( + "target", + kind="dit", + load_id="fastvideo.models.llm:TargetModel", + factory=lambda inst: ToyTargetModel(seed=seed), # the ground-truth AR + resident_for=["spec_decode"], + required_for={"reason"}), } loops = { - "spec_decode": LoopSpec("spec_decode", kind=LoopKind.AR_DECODE, work_unit_kind=WorkUnitKind.AR_TOKEN, - step_cost_model=ar_cost, shared_weight_components=["draft", "target"], - cache_policy=["paged_kv"], loop_factory=spec_factory), + "spec_decode": + LoopSpec("spec_decode", + kind=LoopKind.AR_DECODE, + work_unit_kind=WorkUnitKind.AR_TOKEN, + step_cost_model=ar_cost, + shared_weight_components=["draft", "target"], + cache_policy=["paged_kv"], + loop_factory=spec_factory), } return ModelCard( - model_id=model_id, family="speculative", components=components, loops=loops, + model_id=model_id, + family="speculative", + components=components, + loops=loops, capabilities=CapabilityMatrix.of(Capability.REASONING_TEXT), - recipe=RecipeSpec(method="base", assumes_loop="spec_decode", - assumes_precision="float32", consistency_required=ConsistencyLevel.C1), - parity=ParitySpec(consistency_levels=[ConsistencyLevel.C0, ConsistencyLevel.C1], - interleave_required=True), - caches={"paged_kv": CacheContract("paged_kv", max_bytes=1 << 24, block_bytes=1 << 12, - reuse_across_requests=False)}, + recipe=RecipeSpec(method="base", + assumes_loop="spec_decode", + assumes_precision="float32", + consistency_required=ConsistencyLevel.C1), + parity=ParitySpec(consistency_levels=[ConsistencyLevel.C0, ConsistencyLevel.C1], interleave_required=True), + caches={ + "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()), ).validate() diff --git a/v2/recipes/speculative/loop.py b/v2/recipes/speculative/loop.py index 41726f03d..d8e8a1f77 100644 --- a/v2/recipes/speculative/loop.py +++ b/v2/recipes/speculative/loop.py @@ -1,18 +1,17 @@ -"""SpeculativeARLoop — draft-verify (speculative) decoding (design_v3 §2.2; §9.16). +"""SpeculativeARLoop — draft-verify (speculative) decoding. -The AR-side analog of content-adaptive compute (§9.12): a cheap **draft** model proposes K tokens, the -**target** model verifies them, and the loop accepts the longest matching prefix plus one target -correction — a *variable* number of tokens per round (a ragged loop the model owns). Two properties: +A cheap draft model proposes K tokens, the target model verifies them, and the loop accepts the longest +matching prefix plus one target correction — a variable number of tokens per round. Two properties: - * **Exact.** Every accepted token is one the target would have produced greedily, and the correction - is the target's own token — so the emitted sequence is *identical* to the target's standalone greedy - decode. The speedup is free (same output, fewer sequential target steps). - * **Latency win.** The target's per-round forwards are one batched verify step (here counted as one - ``verify_round``); standalone decoding would take one *sequential* target step per token. So - ``verify_rounds`` (the expensive model's latency steps) is well below the token count. + * Exact. Every accepted token is one the target would have produced greedily, and the correction is the + target's own token, so the emitted sequence is identical to the target's standalone greedy decode. The + speedup is free (same output, fewer sequential target steps). + * Latency win. The target's per-round forwards are one batched verify step (counted as one + ``verify_round``); standalone decoding would take one sequential target step per token, so + ``verify_rounds`` is well below the token count. -It binds TWO components (``draft`` + ``target``) on one resident instance — two models co-scheduled in -one request — and every round is an ``AR_TOKEN`` WorkUnit. +It binds two components (``draft`` + ``target``) on one resident instance, and every round is an +``AR_TOKEN`` WorkUnit. """ from __future__ import annotations @@ -30,19 +29,23 @@ from v2.request.streams import StreamChunk class SpeculativeARLoop: + def __init__(self, *, loop_id, draft_id="draft", target_id="target", cost, spec_len=4, max_tokens=12): self.loop_id = loop_id self.draft_id = draft_id self.target_id = target_id self.cost = cost - self.spec_len = spec_len # K: candidate tokens drafted per round + self.spec_len = spec_len # K: candidate tokens drafted per round self.max_tokens = max_tokens def init(self, req, model, ctx) -> LoopState: seed = req.sampling.seed if req.sampling.seed is not None else 0 - st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id, - request_id=req.request_id, profile=ctx.profile, seed=seed) - st.scratch["tokens"] = list(ctx.slots.get("prompt_tokens") or [1]) # prefill + st = LoopState(loop_id=self.loop_id, + instance_id=model.card.model_id, + request_id=req.request_id, + profile=ctx.profile, + seed=seed) + st.scratch["tokens"] = list(ctx.slots.get("prompt_tokens") or [1]) # prefill st.scratch["generated"] = [] st.scratch["max_tokens"] = min(req.sampling.max_tokens, self.max_tokens) st.scratch.update(verify_rounds=0, target_forwards=0, draft_forwards=0, accepted=0, drafted=0) @@ -59,33 +62,41 @@ class SpeculativeARLoop: # 1) DRAFT proposes k candidate tokens autoregressively (cheap) cand, c = [], list(ctx_tokens) for _ in range(k): - t = int(draft.ar_forward(c)); cand.append(t); c.append(t) + t = int(draft.ar_forward(c)) + cand.append(t) + c.append(t) # 2) TARGET verifies (one batched step): accept the matching prefix, then 1 correction accepted, tf = [], 0 for j in range(k): - tj = int(target.ar_forward(ctx_tokens + cand[:j])); tf += 1 + tj = int(target.ar_forward(ctx_tokens + cand[:j])) + tf += 1 if tj == cand[j]: - accepted.append(cand[j]) # draft agreed with target → accept (free) + accepted.append(cand[j]) # draft agreed with target → accept (free) else: - accepted.append(tj) # mismatch → take the target's token, stop + accepted.append(tj) # mismatch → take the target's token, stop break else: - accepted.append(int(target.ar_forward(ctx_tokens + cand))); tf += 1 # all k accepted + 1 + accepted.append(int(target.ar_forward(ctx_tokens + cand))) + tf += 1 # all k accepted + 1 return StepResult(output={"accepted": accepted, "target_forwards": tf, "drafted": k}) pos = len(st.scratch["generated"]) emits = [StreamChunk(stream_id=st.request_id, modality="text", seq=pos, preview=True)] - return WorkPlan( - loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.AR_TOKEN, - shape_sig=ShapeSignature(WorkUnitKind.AR_TOKEN, dims=(len(ctx_tokens), k)), - resources=ResourceRequest(compute_seconds=self.cost.predict(len(ctx_tokens) + k), - resident_bytes=8 * len(ctx_tokens), peak_activation_bytes=8 * k), - payload={"round": st.scratch["verify_rounds"]}, run=run, emits=emits, - label=f"spec.round{st.scratch['verify_rounds']}") + return WorkPlan(loop_id=self.loop_id, + instance_id=st.instance_id, + kind=WorkUnitKind.AR_TOKEN, + shape_sig=ShapeSignature(WorkUnitKind.AR_TOKEN, dims=(len(ctx_tokens), k)), + resources=ResourceRequest(compute_seconds=self.cost.predict(len(ctx_tokens) + k), + resident_bytes=8 * len(ctx_tokens), + peak_activation_bytes=8 * k), + payload={"round": st.scratch["verify_rounds"]}, + run=run, + emits=emits, + label=f"spec.round{st.scratch['verify_rounds']}") def advance(self, st: LoopState, result: StepResult) -> LoopState: room = st.scratch["max_tokens"] - len(st.scratch["generated"]) - accepted = list(result.output["accepted"])[:room] # don't overshoot max_tokens + accepted = list(result.output["accepted"])[:room] # don't overshoot max_tokens st.scratch["tokens"].extend(accepted) st.scratch["generated"].extend(accepted) st.scratch["verify_rounds"] += 1 @@ -100,10 +111,15 @@ class SpeculativeARLoop: def finalize(self, st: LoopState) -> LoopResult: toks = st.scratch["generated"] rounds = max(1, st.scratch["verify_rounds"]) - return LoopResult( - outputs={"tokens": toks, "text": "tok:" + ",".join(map(str, toks))}, - metrics={"tokens": float(len(toks)), "verify_rounds": float(st.scratch["verify_rounds"]), - "target_forwards": float(st.scratch["target_forwards"]), - "draft_forwards": float(st.scratch["drafted"]), - "tokens_per_round": float(len(toks)) / rounds}, - behavior=st.trajectory or None) + return LoopResult(outputs={ + "tokens": toks, + "text": "tok:" + ",".join(map(str, toks)) + }, + metrics={ + "tokens": float(len(toks)), + "verify_rounds": float(st.scratch["verify_rounds"]), + "target_forwards": float(st.scratch["target_forwards"]), + "draft_forwards": float(st.scratch["drafted"]), + "tokens_per_round": float(len(toks)) / rounds + }, + behavior=st.trajectory or None) diff --git a/v2/recipes/speculative/program.py b/v2/recipes/speculative/program.py index 18c5dd32f..2eeec3512 100644 --- a/v2/recipes/speculative/program.py +++ b/v2/recipes/speculative/program.py @@ -1,4 +1,4 @@ -"""Speculative-decoding program: tokenize → spec_decode (draft+verify) → emit_text (design_v3 §9.16).""" +"""Speculative-decoding program: tokenize -> spec_decode (draft+verify) -> emit_text.""" from __future__ import annotations from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind @@ -7,13 +7,19 @@ from v2.recipes.omni import emit_text_node, tokenize_node def build_speculative_program() -> Program: return Program( - program_id="speculative.decode", kind=ProgramKind.INLINE, + program_id="speculative.decode", + kind=ProgramKind.INLINE, nodes=[ - ComponentNode("tokenize", fn=tokenize_node, writes=("prompt_tokens",)), - ModelLoopNode("spec", loop_id="spec_decode", output_slot="spec_out", - reads=("prompt_tokens",), writes=("spec_out",)), - ComponentNode("emit_text", fn=emit_text_node("spec_out", "text"), - reads=("spec_out",), writes=("text",)), + ComponentNode("tokenize", fn=tokenize_node, writes=("prompt_tokens", )), + ModelLoopNode("spec", + loop_id="spec_decode", + output_slot="spec_out", + reads=("prompt_tokens", ), + writes=("spec_out", )), + ComponentNode("emit_text", fn=emit_text_node("spec_out", "text"), reads=("spec_out", ), writes=("text", )), ], - output_artifacts={"text": "text", "tokens": "spec_out"}, + output_artifacts={ + "text": "text", + "tokens": "spec_out" + }, ).validate() diff --git a/v2/recipes/stable_audio/card.py b/v2/recipes/stable_audio/card.py index 63fdbe0fa..224b6aad1 100644 --- a/v2/recipes/stable_audio/card.py +++ b/v2/recipes/stable_audio/card.py @@ -8,14 +8,14 @@ Two variants from the converted Diffusers repos: loader from the checkpoint config; the card carries the per-variant sample_size + audio defaults.) Architecture deltas vs the video recipes (all declared on the card so the recipe is self-contained): - * DiT ``fastvideo.models.dits.stable_audio:StableAudioDiT`` — a v-PREDICTION network (EDM-v / VDenoiser), - NOT flow-match and NOT EDM-Karras. The adapter (``ComponentSpec.adapter`` -> ``StableAudioDiT``) packs - the conditioner payload into (cross_attn_cond, global_embed), passes the RAW continuous sigma as the - timestep, and returns the raw v output; ``StableAudioDenoiseLoop`` does the VDenoiser v->x0 + x0-space - CFG + the polyexponential schedule (sigma_min 0.3, sigma_max 500, rho 1) and a DPM++ multistep step. - * VAE ``fastvideo.models.vaes.oobleck:OobleckVAE`` — 1-D conv audio VAE (hop_length 2048). The - ``OobleckVAE`` adapter is RAW latent space (NO mean/std normalization, unlike WanVAE) and exposes - ``sampling_rate`` for the decode-slice. + * DiT ``fastvideo.models.dits.stable_audio:StableAudioDiT`` — a v-prediction network (EDM-v / VDenoiser), + not flow-match and not EDM-Karras. The adapter packs the conditioner payload into (cross_attn_cond, + global_embed), passes the raw continuous sigma as the timestep, and returns the raw v output; + ``StableAudioDenoiseLoop`` does the VDenoiser v->x0 + x0-space CFG + the polyexponential schedule + (sigma_min 0.3, sigma_max 500, rho 1) and a DPM++ multistep step. + * VAE ``fastvideo.models.vaes.oobleck:OobleckVAE`` — 1-D conv audio VAE (hop_length 2048). The adapter + is raw latent space (no mean/std normalization, unlike WanVAE) and exposes ``sampling_rate`` for the + decode-slice. * Conditioner ``fastvideo.models.encoders.stable_audio_conditioner:StableAudioMultiConditioner`` — T5-base (max_length 128) + duration NumberConditioners. Declared with ``kind="text_encoder"`` (kind reuse — the shared ``_MAKERS`` has no ``conditioner`` kind and editing it is out of scope) + the explicit @@ -186,10 +186,10 @@ class _ToyAudioDiT: """CPU toy v-prediction stand-in for ``StableAudioDiT`` (1-D audio latent [64, L]). The shared ``ToyDiT`` is channel-fixed (4 channels) and shaped for video; SA's latent is 64-channel - 1-D. This tiny deterministic toy honors the same contract the loop needs — ``__call__(latent, - text_embed, sigma) -> v`` with the loop-supplied VDenoiser-scaled input — so the loop runs end-to-end - on CPU. ``text_embed`` is the packed conditioning payload (the toy just means-pools it as a scalar - conditioning signal, matching ToyDiT's text handling). Deterministic given (seed, inputs).""" + 1-D. This deterministic toy honors the contract the loop needs — ``__call__(latent, text_embed, sigma) + -> v`` with the loop-supplied VDenoiser-scaled input — so the loop runs end-to-end on CPU. + ``text_embed`` is the packed conditioning payload (the toy means-pools it as a scalar conditioning + signal, matching ToyDiT's text handling). Deterministic given (seed, inputs).""" def __init__(self, seed: int = 0): import numpy as np @@ -207,8 +207,8 @@ class _ToyAudioDiT: class _ToyOobleckVAE: - """CPU toy 1-D audio VAE stand-in for ``OobleckVAE`` (the shared ``ToyVAE`` is a video VAE that decodes - a 4-channel ``[C,T,H,W]`` latent). Decodes the 64-channel ``[64, L]`` audio latent to a + """CPU toy 1-D audio VAE stand-in for ``OobleckVAE`` (the shared ``ToyVAE`` is a video VAE decoding a + 4-channel ``[C,T,H,W]`` latent). Decodes the 64-channel ``[64, L]`` audio latent to a ``[channels, samples]`` waveform (channel-collapse projection + hop_length upsample), deterministic so the program node is interleave-safe. Exposes ``hop_length`` / ``sampling_rate`` like the real adapter.""" diff --git a/v2/recipes/stable_audio/loop.py b/v2/recipes/stable_audio/loop.py index 306905318..71b4a87dc 100644 --- a/v2/recipes/stable_audio/loop.py +++ b/v2/recipes/stable_audio/loop.py @@ -1,10 +1,9 @@ """StableAudioDenoiseLoop — a v-prediction (EDM-v / VDenoiser) DPM++ denoise loop for Stable Audio Open. -Stable Audio is NOT a flow-match model and NOT an EDM-Karras model: it is a v-prediction network sampled -by k-diffusion's ``dpmpp-3m-sde`` (a 3rd-order multistep stochastic DPM-Solver++) over a POLYEXPONENTIAL +Stable Audio is not a flow-match model and not an EDM-Karras model: it is a v-prediction network sampled +by k-diffusion's ``dpmpp-3m-sde`` (a 3rd-order multistep stochastic DPM-Solver++) over a polyexponential sigma schedule (sigma_min=0.3, sigma_max=500, rho=1.0). So neither ``WanDenoiseLoop`` (flow-match) nor -``CosmosDenoiseLoop`` (EDM-Karras x0) can be reused — this is a NEW loop (the spec's -``NEEDS_NEW:StableAudioDenoiseLoop``). +``CosmosDenoiseLoop`` (EDM-Karras x0) can be reused — this needs a new loop. Faithfulness vs ``fastvideo/pipelines/basic/stable_audio/stages/denoising.py``: * The schedule is polyexponential (``sampler.build_polyexponential_sigmas``), seeded noise scaled by @@ -15,15 +14,15 @@ Faithfulness vs ``fastvideo/pipelines/basic/stable_audio/stages/denoising.py``: * CFG: ``denoised = uncond + (cond - uncond)·scale``, combined in x0 space (the ``_DiTAdapter`` convention; the negative cross-attn is null/zero, applied by the conditioning node). -GPU PATH (BRINGUP — k_diffusion is an external dep, NOT installed in this env): the faithful, lowest-risk -GPU route is to NOT decompose ``dpmpp-3m-sde`` into per-step v2 WorkPlans but to run the whole -k-diffusion sampler inside ONE thunk: ``import k_diffusion as K; K.sampling.sample_dpmpp_3m_sde( -K.external.VDenoiser(adapter), x, sigmas, ...)`` — a direct lift of the fastvideo stage. That requires a -per-loop "single-shot" execution mode the v2 ``ctx.execute`` per-step contract does not yet expose -(BRINGUP: needs a loop-as-one-thunk hook). For CPU-verification + per-step serving granularity, this loop -ships a DETERMINISTIC DPM-Solver++(2M) recurrence in numpy over the SAME VDenoiser x0 — a faithful -v-prediction multistep integrator that runs end-to-end on the CPU toy. The SDE (stochastic) higher-order -``dpmpp-3m-sde`` terms are the GPU-parity refinement (BRINGUP) layered on top of this deterministic spine. +GPU PATH (BRINGUP — k_diffusion is an external dep, not installed in this env): the faithful, lowest-risk +GPU route is to run the whole k-diffusion sampler inside one thunk rather than decomposing +``dpmpp-3m-sde`` into per-step v2 WorkPlans: ``K.sampling.sample_dpmpp_3m_sde(K.external.VDenoiser(adapter), +x, sigmas, ...)`` — a direct lift of the fastvideo stage. That needs a per-loop "single-shot" execution +mode the v2 ``ctx.execute`` per-step contract does not yet expose (BRINGUP: needs a loop-as-one-thunk hook). +For CPU-verification + per-step serving granularity, this loop ships a deterministic DPM-Solver++(2M) +recurrence in numpy over the same VDenoiser x0 — a faithful v-prediction multistep integrator that runs +end-to-end on the CPU toy. The stochastic higher-order ``dpmpp-3m-sde`` terms are the GPU-parity refinement +(BRINGUP) layered on top of this deterministic spine. A2A (init_audio -> lowered sigma_max) and RePaint inpainting (per-step blend callback) are threaded as inert hooks (None for the registered T2A preset) and documented BRINGUP — base port is pure text->audio. diff --git a/v2/recipes/stable_audio/program.py b/v2/recipes/stable_audio/program.py index 9b4b2a3ef..2abb1ed0a 100644 --- a/v2/recipes/stable_audio/program.py +++ b/v2/recipes/stable_audio/program.py @@ -1,8 +1,8 @@ """Stable Audio Open T2A program: conditioning -> stable_audio_denoise (v-prediction DPM++) -> audio_decode. -Mirrors the cosmos2 t2v inline program shape — the SA specifics (the conditioner's (cross_attn_cond, -global_embed) packing; the v-prediction sampler) live entirely in the conditioning node, the -``StableAudioDenoiseLoop``, and the torch adapters, so the node graph stays a clean 3-node line. +The SA specifics (the conditioner's (cross_attn_cond, global_embed) packing; the v-prediction sampler) +live in the conditioning node, the ``StableAudioDenoiseLoop``, and the torch adapters, so the node graph +stays a clean 3-node line. Deltas vs the video programs: * The first node is ``conditioning`` (NOT the generic ``text_encode``): it runs the SA multi-conditioner @@ -10,15 +10,13 @@ Deltas vs the video programs: payload (so the loop->dit call stays modality-agnostic). It degrades to the generic text encoder on the CPU toy (which exposes only ``.encode(text)``), so the loop still runs end-to-end. * The last node is ``audio_decode`` (OobleckVAE.decode + slice [start,end] seconds) writing an ``audio`` - artifact (NOT ``video``). BRINGUP: the v2 output-artifact surface is video-shaped (cards/programs - declare a ``video`` artifact); the registered ``output_artifacts`` here declares ``audio`` (+ latents), - which an audio-aware saver/output path must consume — that audio output node is the request-API - extension flagged in the spec (blocker #6). The slot value is the decoded waveform numpy array. + artifact (NOT ``video``). BRINGUP: the v2 output-artifact surface is video-shaped, so the ``audio`` + artifact declared here needs an audio-aware saver/output path to consume it. The slot value is the + decoded waveform numpy array. Audio duration knobs: ``DiffusionParams`` has no ``audio_start_in_s``/``audio_end_in_s`` (the v2 request -surface is video-shaped — spec blocker #6). They are read from ``request.node_params['conditioning']`` -(the per-node override surface) with the card's defaults as fallback. A first-class audio request field is -the request-API extension (BRINGUP). +surface is video-shaped). They are read from ``request.node_params['conditioning']`` (the per-node +override surface) with the card's defaults as fallback. A first-class audio request field is TODO (BRINGUP). """ from __future__ import annotations diff --git a/v2/recipes/stable_audio/sampler.py b/v2/recipes/stable_audio/sampler.py index e4ccb32b6..d3f9dc7f7 100644 --- a/v2/recipes/stable_audio/sampler.py +++ b/v2/recipes/stable_audio/sampler.py @@ -1,6 +1,6 @@ """Stable Audio sampler helpers (recipe-local — NOT in v2/loop/sampler.py, which is flow-match/EDM-Karras). -Three numpy primitives the ``StableAudioDenoiseLoop`` composes, faithful to k-diffusion's behavior for the +Three numpy primitives the ``StableAudioDenoiseLoop`` composes, matching k-diffusion's behavior for the Stable Audio ``dpmpp-3m-sde`` sampler over a v-prediction (VDenoiser) network: * ``build_polyexponential_sigmas`` — k-diffusion ``get_sigmas_polyexponential(n, sigma_min, sigma_max, diff --git a/v2/recipes/tiled/__init__.py b/v2/recipes/tiled/__init__.py index b4d424a10..4a52e90d5 100644 --- a/v2/recipes/tiled/__init__.py +++ b/v2/recipes/tiled/__init__.py @@ -1,4 +1,4 @@ -"""Tiled VAE decode: the VAE_TILE WorkUnit kind co-scheduled with denoise steps (design_v3 §4, §17).""" +"""Tiled VAE decode: the VAE_TILE WorkUnit kind co-scheduled with denoise steps.""" from __future__ import annotations from v2.recipes.tiled.card import build_tiled_card diff --git a/v2/recipes/tiled/card.py b/v2/recipes/tiled/card.py index 12d375ad2..5ed58cdf5 100644 --- a/v2/recipes/tiled/card.py +++ b/v2/recipes/tiled/card.py @@ -1,5 +1,5 @@ -"""Tiled-decode card (design_v3 §4, §17) — a diffusion model whose VAE decode is a scheduled -``vae_tile`` loop, so its tiles co-schedule with denoise steps under the one WorkUnit budget.""" +"""Tiled-decode card — a diffusion model whose VAE decode is a scheduled ``vae_tile`` loop, so its +tiles co-schedule with denoise steps under the one WorkUnit budget.""" from __future__ import annotations from v2._enums import Capability, ConsistencyLevel, LoopKind, WorkUnitKind @@ -30,39 +30,65 @@ def build_tiled_card(model_id: str = "wan-tiled", *, tile_rows: int = 1) -> Mode precision, expert = PrecisionPolicy(), NoRouting("transformer") def denoise_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) def tile_factory(): - return VAETileLoop(loop_id="vae_tile", vae_id="vae", cost=tile_cost, - tile_rows=tile_rows, latent_slot="denoise_out") + return VAETileLoop(loop_id="vae_tile", + vae_id="vae", + cost=tile_cost, + tile_rows=tile_rows, + latent_slot="denoise_out") 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(), - resident_for=["vae_tile"], 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(), resident_for=["vae_tile"], + required_for={"t2v"}), } loops = { - "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=denoise_factory), - "vae_tile": LoopSpec("vae_tile", kind=LoopKind.VAE_TILE, work_unit_kind=WorkUnitKind.VAE_TILE, - step_cost_model=tile_cost, shared_weight_components=["vae"], - cache_policy=[], loop_factory=tile_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=denoise_factory), + "vae_tile": + LoopSpec("vae_tile", + kind=LoopKind.VAE_TILE, + work_unit_kind=WorkUnitKind.VAE_TILE, + step_cost_model=tile_cost, + shared_weight_components=["vae"], + cache_policy=[], + loop_factory=tile_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), - parity=ParitySpec(consistency_levels=[ConsistencyLevel.C0, ConsistencyLevel.C1], - interleave_required=True), - caches={"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True)}, + recipe=RecipeSpec(method="base", + assumes_loop="diffusion_denoise", + assumes_precision="float32", + consistency_required=ConsistencyLevel.C1), + parity=ParitySpec(consistency_levels=[ConsistencyLevel.C0, ConsistencyLevel.C1], interleave_required=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() diff --git a/v2/recipes/tiled/loop.py b/v2/recipes/tiled/loop.py index 8af204717..849a468f8 100644 --- a/v2/recipes/tiled/loop.py +++ b/v2/recipes/tiled/loop.py @@ -1,13 +1,11 @@ -"""VAETileLoop — tiled VAE decode (design_v3 §4.3 ``LoopKind.VAE_TILE``; §17 falsifier probe). +"""VAETileLoop — tiled VAE decode (``LoopKind.VAE_TILE``). -Decodes a latent in spatial row-tiles, **each tile a ``VAE_TILE`` WorkUnit** that flows through the -exact same admission + step-scheduling + interleave machinery as a denoise step. This is the §17 -falsifier made testable: *does the general WorkUnit scheduler earn its generality on non-step kinds, or -should tiles collapse to an in-loop call?* At the mechanism level it composes cleanly — VAE_TILE units -interleave bit-identically with DIFFUSION_STEP units across requests (the §9.3 gate over mixed kinds). -Whether it *pays* on a real GPU duty-cycle is the empirical half of the falsifier (deferred to the port). +Decodes a latent in spatial row-tiles, each tile a ``VAE_TILE`` WorkUnit flowing through the same +admission + step-scheduling + interleave machinery as a denoise step. So VAE_TILE units interleave with +DIFFUSION_STEP units across requests through one scheduler; this exercises whether that general scheduler +earns its generality on non-step kinds (whether it pays on a real GPU duty-cycle is deferred to the port). -The toy VAE is spatially local (channel-mix + repeat-upsample), so tiled decode is **exactly** equal to +The toy VAE is spatially local (channel-mix + repeat-upsample), so tiled decode is exactly equal to one-shot decode — a C0 component parity the tests assert. """ from __future__ import annotations @@ -27,6 +25,7 @@ from v2.loop.contracts import ( class VAETileLoop: + def __init__(self, *, loop_id, vae_id="vae", cost, tile_rows=1, latent_slot="denoise_out"): self.loop_id = loop_id self.vae_id = vae_id @@ -38,8 +37,10 @@ class VAETileLoop: out = ctx.slots.get(self.latent_slot) latent = out["latents"] if isinstance(out, dict) else out latent = np.asarray(latent, dtype="float32") - st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id, - request_id=req.request_id, profile=ctx.profile) + st = LoopState(loop_id=self.loop_id, + instance_id=model.card.model_id, + request_id=req.request_id, + profile=ctx.profile) h = latent.shape[2] st.scratch["latent"] = latent st.scratch["tiles"] = [(r, min(r + self.tile_rows, h)) for r in range(0, h, self.tile_rows)] or [(0, h)] @@ -55,18 +56,22 @@ class VAETileLoop: latent, vid = st.scratch["latent"], self.vae_id def run(model, override=None): - tile = latent[:, :, r0:r1, :] # a spatial row-band of the latent + tile = latent[:, :, r0:r1, :] # a spatial row-band of the latent dec = model.component(vid).decode(tile) return StepResult(output={"tile": np.asarray(dec, dtype="float32"), "i": i}) rows = r1 - r0 - res = ResourceRequest( - compute_seconds=self.cost.predict(int(rows * latent.shape[1] * latent.shape[3])), - resident_bytes=int(latent.nbytes), peak_activation_bytes=int(latent[:, :, r0:r1, :].nbytes)) - return WorkPlan( - loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.VAE_TILE, - shape_sig=ShapeSignature(WorkUnitKind.VAE_TILE, dims=(rows, latent.shape[3])), - resources=res, payload={"tile": i}, run=run, label=f"vae.tile{i}") + res = ResourceRequest(compute_seconds=self.cost.predict(int(rows * latent.shape[1] * latent.shape[3])), + resident_bytes=int(latent.nbytes), + peak_activation_bytes=int(latent[:, :, r0:r1, :].nbytes)) + return WorkPlan(loop_id=self.loop_id, + instance_id=st.instance_id, + kind=WorkUnitKind.VAE_TILE, + shape_sig=ShapeSignature(WorkUnitKind.VAE_TILE, dims=(rows, latent.shape[3])), + resources=res, + payload={"tile": i}, + run=run, + label=f"vae.tile{i}") def advance(self, st: LoopState, result: StepResult) -> LoopState: st.scratch["decoded"][int(result.output["i"])] = result.output["tile"] @@ -75,6 +80,9 @@ class VAETileLoop: def finalize(self, st: LoopState) -> LoopResult: tiles = [t for t in st.scratch["decoded"] if t is not None] - video = np.concatenate(tiles, axis=2) if tiles else None # stitch row-bands back together - return LoopResult(outputs={"video": video, "latents": st.scratch["latent"]}, + video = np.concatenate(tiles, axis=2) if tiles else None # stitch row-bands back together + return LoopResult(outputs={ + "video": video, + "latents": st.scratch["latent"] + }, metrics={"vae_tiles": float(st.step_idx)}) diff --git a/v2/recipes/tiled/program.py b/v2/recipes/tiled/program.py index b02e91f52..41cca23c4 100644 --- a/v2/recipes/tiled/program.py +++ b/v2/recipes/tiled/program.py @@ -1,4 +1,4 @@ -"""Tiled-decode program: text_encode → diffusion_denoise → vae_tile (design_v3 §4, §17). +"""Tiled-decode program: text_encode → diffusion_denoise → vae_tile. The VAE decode is a ``ModelLoopNode`` (the ``vae_tile`` loop), not a one-shot ``ComponentNode`` — so its tiles become scheduled ``VAE_TILE`` WorkUnits that interleave with other requests' ``DIFFUSION_STEP`` @@ -12,13 +12,23 @@ from v2.recipes.common import text_encode_node_fn as _text_encode def build_tiled_program() -> Program: return Program( - program_id="wan.t2v.tiled", kind=ProgramKind.INLINE, + program_id="wan.t2v.tiled", + kind=ProgramKind.INLINE, nodes=[ ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")), - ModelLoopNode("denoise", loop_id="diffusion_denoise", output_slot="denoise_out", - reads=("text_embeds",), writes=("denoise_out",)), - ModelLoopNode("vae_tile", loop_id="vae_tile", output_slot="tiled_out", - reads=("denoise_out",), writes=("tiled_out",)), + ModelLoopNode("denoise", + loop_id="diffusion_denoise", + output_slot="denoise_out", + reads=("text_embeds", ), + writes=("denoise_out", )), + ModelLoopNode("vae_tile", + loop_id="vae_tile", + output_slot="tiled_out", + reads=("denoise_out", ), + writes=("tiled_out", )), ], - output_artifacts={"video": "tiled_out", "latents": "denoise_out"}, + output_artifacts={ + "video": "tiled_out", + "latents": "denoise_out" + }, ).validate() diff --git a/v2/recipes/turbowan/card.py b/v2/recipes/turbowan/card.py index 1acfafc51..72abd30ee 100644 --- a/v2/recipes/turbowan/card.py +++ b/v2/recipes/turbowan/card.py @@ -1,24 +1,22 @@ -"""TurboWan ModelCards — rCM few-step distilled Wan (the bucket-C self-contained pattern). +"""TurboWan ModelCards — rCM few-step distilled Wan (self-contained recipe package). -TurboWan (= Reparameterized/recurrent Consistency Model distilled Wan, TurboDiffusion arXiv:2512.16093) -REUSES the Wan architecture end-to-end: ``WanTransformer3DModel`` velocity DiT, ``AutoencoderKLWan`` VAE, -UMT5 text encoder, (+ a CLIP vision encoder for the I2V variant). So there is NO new torch adapter — the -components carry the same ``load_id`` strings as ``v2/recipes/wan21`` and the existing ``_make_dit`` / -``WanDiT`` dispatch in ``torch_backend.py`` builds them as-is. The ONLY new work is the sampler/loop: the -rCM consistency schedule + step (``v2/recipes/turbowan/{sampler,loop}.py``). +TurboWan (Reparameterized Consistency Model distilled Wan, TurboDiffusion arXiv:2512.16093) reuses the Wan +architecture end-to-end: ``WanTransformer3DModel`` velocity DiT, ``AutoencoderKLWan`` VAE, UMT5 text encoder +(+ CLIP vision encoder for I2V). There is no new torch adapter — components carry the same ``load_id`` strings +as ``v2/recipes/wan21``, built by the existing ``WanDiT`` dispatch in ``torch_backend.py``. The only new work +is the rCM sampler/loop (``v2/recipes/turbowan/{sampler,loop}.py``). Three HF ids, two builders: - * ``build_turbowan_card`` (size-agnostic T2V) serves ``TurboWan2.1-T2V-1.3B`` (480p, flow geometry - notwithstanding — rCM has no flow-shift schedule) and ``TurboWan2.1-T2V-14B`` (720p). The size is - resolved from the checkpoint by the loader; the per-id deltas are just the default h/w. - * ``build_turbowan_i2v_a14b_card`` serves ``TurboWan2.2-I2V-A14B`` — a MoE (two WanTransformer3DModel - experts, boundary-routed) + i2v conditioning (CLIP image encoder + first-frame [mask|cond]), reusing - ``BoundaryTimestepRouting`` and the Wan i2v hooks. sigma_max=200 (vs 80 for T2V). + * ``build_turbowan_card`` (size-agnostic T2V) serves ``TurboWan2.1-T2V-1.3B`` (480p) and + ``TurboWan2.1-T2V-14B`` (720p); rCM has no flow-shift schedule. Size is resolved from the checkpoint by + the loader, so the per-id delta is just the default h/w. + * ``build_turbowan_i2v_a14b_card`` serves ``TurboWan2.2-I2V-A14B`` — a boundary-routed MoE (two + WanTransformer3DModel experts) + i2v conditioning (CLIP image encoder + first-frame [mask|cond]), reusing + ``BoundaryTimestepRouting`` and the Wan i2v hooks. sigma_max=200 vs 80 for T2V. -SamplingDefaults are taken from ``fastvideo/pipelines/basic/turbodiffusion/presets.py``: 4 inference steps, -guidance_scale 1.0 (CFG OFF — the few-step recipe), EMPTY negative prompt (a local constant here, since the -shared WAN_NEG_* are not what the TurboWan preset uses). ``stamp_wan21_checkpoints`` applies (diffusers -transformer/transformer_2/vae/text_encoder/image_encoder subfolder layout). +SamplingDefaults match ``fastvideo/pipelines/basic/turbodiffusion/presets.py``: 4 steps, guidance_scale 1.0 +(CFG off), and an empty negative prompt (a local constant — the shared WAN_NEG_* are not what the preset uses). +``stamp_wan21_checkpoints`` applies the diffusers subfolder layout. """ from __future__ import annotations @@ -43,8 +41,8 @@ from v2.platform.backends.toy import ToyDiT, ToyImageEncoder, ToyTextEncoder, To from v2.recipes.turbowan.loop import TurboWanDenoiseLoop from v2.recipes.wan21.card import stamp_wan21_checkpoints -# The registered TurboWan presets use an EMPTY negative prompt (guidance is off at scale 1.0), so the -# shared Wan negatives do not apply. Local constant per the self-contained-package rule. +# TurboWan presets use an empty negative prompt (guidance off at scale 1.0), so the shared Wan negatives +# do not apply. Local constant per the self-contained-package rule. TURBOWAN_NEG = "" _WAN_DIT = "fastvideo.models.dits.wanvideo:WanTransformer3DModel" @@ -151,8 +149,8 @@ def build_turbowan_i2v_a14b_card(model_id: str = "turbowan2.2-i2v-a14b", checkpoint_root: str | None = None) -> ModelCard: """TurboWan2.2-I2V-A14B — rCM few-step MoE i2v. Two WanTransformer3DModel experts boundary-routed (``boundary_ratio`` 0.9, compared in raw-sigma space) + i2v conditioning (CLIP + first-frame - [mask|cond]). sigma_max=200 (the I2V pipeline's value vs 80 for T2V). Reuses the Wan adapter, the - boundary policy, and the i2v hooks; only the rCM sampler/loop is new.""" + [mask|cond]). sigma_max=200 (the I2V pipeline's value vs 80 for T2V). Reuses the Wan adapter, boundary + policy, and i2v hooks; only the rCM sampler/loop is new.""" seed = _seed_from(model_id) cost = CostModel(kind=WorkUnitKind.DIFFUSION_STEP, base_seconds=1e-4, per_unit_seconds=1e-7) cfg = ClassicCFG() diff --git a/v2/recipes/turbowan/loop.py b/v2/recipes/turbowan/loop.py index 82511055a..22fb2448a 100644 --- a/v2/recipes/turbowan/loop.py +++ b/v2/recipes/turbowan/loop.py @@ -1,25 +1,22 @@ """TurboWanDenoiseLoop — the rCM (Reparameterized Consistency Model) few-step denoise loop. -TurboWan reuses the Wan architecture (``WanTransformer3DModel`` velocity predictor, ``AutoencoderKLWan``, -UMT5) but replaces the multistep flow-match sampler with the rCM consistency sampler. So this loop is the -ONLY genuinely new piece vs ``v2/recipes/wan21``: the standard FLOW_MATCH schedule + Euler step are -swapped for the TrigFlow->RectifiedFlow RCM schedule + the stochastic consistency step in -``v2/recipes/turbowan/sampler.py`` (a faithful port of ``RCMScheduler``). +TurboWan reuses the Wan architecture but swaps the multistep flow-match sampler for the rCM consistency +sampler, so this loop is the only genuinely new piece vs ``v2/recipes/wan21``: the FLOW_MATCH schedule + +Euler step become the TrigFlow->RectifiedFlow schedule + stochastic consistency step in ``sampler.py`` +(a faithful port of ``RCMScheduler``). -Per step (mirroring ``fastvideo/pipelines/stages/denoising.py`` driving ``RCMScheduler``): - * the model timestep is the SCALED sigma ``t = sigma * 1000`` (the DiT's training convention); - * the DiT predicts a velocity (CFG-combined when a guidance scale != 1 and a negative branch exists — - the registered TurboWan presets use guidance_scale 1.0, so CFG is OFF, matching the few-step recipe); - * the rCM SDE step ``x_next = (1 - t_next)*(x - t_cur*v) + t_next*noise`` injects FRESH per-step noise. - The host RNG (and the consistency math, which is not the platform's flow-match kernel) means every - step is eager — ``capturable=False`` — exactly like the cosmos2 EDM loop and the FlowGRPO SDE rollout. +Per step (mirroring ``fastvideo/pipelines/stages/denoising.py``): + * the model timestep is the scaled sigma ``t = sigma * 1000`` (the DiT's training convention); + * the DiT predicts a velocity (CFG-combined when guidance_scale != 1 and a negative branch exists — the + presets use guidance_scale 1.0, so CFG is off); + * the rCM SDE step ``x_next = (1 - t_next)*(x - t_cur*v) + t_next*noise`` injects fresh per-step noise. + The host RNG makes every step eager (``capturable=False``), like the cosmos2 EDM and FlowGRPO SDE loops. -MoE + i2v (Wan2.2-I2V-A14B): expert routing and the i2v conditioning hooks are reused unchanged from the -Wan loop. ``self.expert`` (a ``BoundaryTimestepRouting`` for the MoE variant) picks the high/low-noise -transformer per step by comparing the RAW sigma against the boundary (0.9 == ``boundary_ratio`` 0.9, since -the fastvideo boundary ``0.9*num_train_timesteps`` is compared against the scaled timestep ``sigma*1000``). -The i2v ``[mask|cond]`` latent + CLIP image embeds flow from slots to the WanDiT adapter exactly as in -``v2/recipes/wan21/i2v.py``; ``None`` for T2V -> the plain forward. The toy DiT accepts (and ignores) both. +MoE + i2v (Wan2.2-I2V-A14B): expert routing and i2v conditioning hooks are reused unchanged from the Wan +loop. ``self.expert`` (a ``BoundaryTimestepRouting`` for the MoE variant) picks the high/low-noise transformer +per step by comparing the raw sigma against the boundary (fastvideo's ``0.9*num_train_timesteps`` vs the +scaled timestep ``sigma*1000``). The i2v ``[mask|cond]`` latent + CLIP image embeds flow from slots to the +WanDiT adapter as in ``v2/recipes/wan21/i2v.py``; ``None`` for T2V -> the plain forward. """ from __future__ import annotations diff --git a/v2/recipes/turbowan/program.py b/v2/recipes/turbowan/program.py index 237a8c253..578991c3b 100644 --- a/v2/recipes/turbowan/program.py +++ b/v2/recipes/turbowan/program.py @@ -1,9 +1,9 @@ """TurboWan programs (node graphs identical to Wan — the rCM specifics live in the loop/sampler). T2V: text_encode -> diffusion_denoise (rCM) -> vae_decode (same inline shape as the Wan/cosmos2 t2v -programs). I2V: the Wan2.2-I2V-A14B variant reuses the ``v2/recipes/wan21/i2v.py`` conditioning nodes -verbatim (CLIP image encode + first-frame [mask|cond] VAE encode), just bound to TurboWan's ``i2v_denoise`` -loop id. No new conditioning math — the only TurboWan delta is the loop's rCM sampler. +programs). I2V (Wan2.2-I2V-A14B): reuses the ``v2/recipes/wan21/i2v.py`` conditioning nodes verbatim (CLIP +image encode + first-frame [mask|cond] VAE encode), bound to TurboWan's ``i2v_denoise`` loop. No new +conditioning math — the only TurboWan delta is the loop's rCM sampler. """ from __future__ import annotations diff --git a/v2/recipes/turbowan/sampler.py b/v2/recipes/turbowan/sampler.py index 8dfc265e9..29bb31642 100644 --- a/v2/recipes/turbowan/sampler.py +++ b/v2/recipes/turbowan/sampler.py @@ -1,30 +1,27 @@ -"""RCM (Reparameterized / recurrent Consistency Model) sampler helpers — TurboWan's few-step distilled -schedule + SDE consistency step. +"""RCM (Reparameterized Consistency Model) sampler helpers — TurboWan's few-step schedule + SDE step. -Faithful CPU-testable port of ``fastvideo/models/schedulers/scheduling_rcm.py:RCMScheduler`` (the -TurboDiffusion rCM sampler, arXiv:2512.16093). The math has two pieces that this module mirrors exactly: +Faithful CPU-testable port of ``fastvideo/models/schedulers/scheduling_rcm.py:RCMScheduler`` (TurboDiffusion +rCM sampler, arXiv:2512.16093). Two pieces mirrored exactly: - 1. ``build_rcm_sigmas`` — the TrigFlow -> RectifiedFlow timestep schedule. Starting from - ``[atan(sigma_max), *mid_timesteps, 0]`` in TrigFlow angle space, each entry is converted with - ``t = sin(t) / (cos(t) + sin(t))`` to a RectifiedFlow sigma in ``[0, 1]``. ``mid_timesteps`` defaults - to ``[1.5, 1.4, 1.0]`` (the scheduler's visual-quality default), sliced to ``num_steps - 1`` so the - final schedule has ``num_steps + 1`` points (the appended terminal ``0``). This matches the loop - convention everywhere else in v2 (``i`` ranges over ``len(sigmas) - 1`` pairwise intervals). + 1. ``build_rcm_sigmas`` — the TrigFlow -> RectifiedFlow timestep schedule. From the TrigFlow angles + ``[atan(sigma_max), *mid_timesteps, 0]``, each entry is converted with ``t = sin(t)/(cos(t)+sin(t))`` to a + RectifiedFlow sigma in ``[0, 1]``. ``mid_timesteps`` defaults to ``[1.5, 1.4, 1.0]`` (the scheduler's + quality default), sliced to ``num_steps - 1`` so the schedule has ``num_steps + 1`` points (terminal + ``0``), matching the v2 loop convention (``i`` ranges over ``len(sigmas) - 1`` intervals). - 2. ``rcm_step`` — the rCM SDE consistency update. The model predicts a velocity ``v_pred`` at the - CURRENT (raw) sigma ``t_cur``; the scheduler reconstructs the clean estimate - ``x_denoised = x - t_cur * v_pred`` and re-noises it to the NEXT sigma with FRESH Gaussian noise: + 2. ``rcm_step`` — the rCM SDE consistency update. The model predicts a velocity ``v_pred`` at the current + (raw) sigma ``t_cur``; the scheduler reconstructs the clean estimate ``x_denoised = x - t_cur * v_pred`` + and re-noises it to the next sigma with fresh Gaussian noise: x_next = (1 - t_next) * x_denoised + t_next * noise - This is NOT a deterministic flow-match Euler step — it is a stochastic consistency step that - injects fresh noise every step (which is why TurboWan reaches quality in 1-4 steps). The fresh - host RNG is exactly why the loop must run this on the eager path (like the FlowGRPO SDE rollout), - never inside a captured CUDA graph. + This stochastic consistency step (not a deterministic Euler step) injects fresh noise every step, which + is why TurboWan reaches quality in 1-4 steps — and why the loop must run it on the eager path (like the + FlowGRPO SDE rollout), never inside a captured CUDA graph. -The real scheduler scales the model-input timestep by 1000 (``timesteps = sigmas * 1000``) — the loop -passes that scaled timestep to the DiT; the step formula here uses the RAW sigmas. ``init_noise_sigma`` -is ``sigmas[0]`` (the loop scales the initial latent by it, matching ``scale_noise``). +The real scheduler scales the model-input timestep by 1000 (``timesteps = sigmas * 1000``); the loop passes +that scaled timestep to the DiT while the step formula here uses the raw sigmas. ``init_noise_sigma`` is +``sigmas[0]`` (the loop scales the initial latent by it, matching ``scale_noise``). """ from __future__ import annotations diff --git a/v2/recipes/unified/__init__.py b/v2/recipes/unified/__init__.py index a18367c8b..3e431a492 100644 --- a/v2/recipes/unified/__init__.py +++ b/v2/recipes/unified/__init__.py @@ -1,6 +1,6 @@ """Unified LM+generator model (UniRL/PromptRL): a prompt-refiner LM expert + a flow generator expert, -both trainable under one RL reward. The stress test that the Card/Loop/Program split holds for -joint multi-expert RL (design_v3 §4, §10).""" +both trainable under one RL reward. Stress-tests that the Card/Loop/Program split holds for +joint multi-expert RL.""" from __future__ import annotations from v2.recipes.unified.card import N_REFINE_ACTIONS, build_unified_card diff --git a/v2/recipes/unified/card.py b/v2/recipes/unified/card.py index 514e4312e..518829612 100644 --- a/v2/recipes/unified/card.py +++ b/v2/recipes/unified/card.py @@ -1,18 +1,18 @@ -"""Unified LM+generator ModelCard — the UniRL/PromptRL stress test (design_v3 §4, §10). +"""Unified LM+generator ModelCard — the UniRL/PromptRL stress test. -This is the topological *opposite* of the Cosmos3 omni card, expressed in the SAME vocabulary: +This is the topological opposite of the Cosmos3 omni card, expressed in the same vocabulary: - * Cosmos3 MoT: ONE ``transformer`` bound to BOTH ``ar_decode`` and ``diffusion_denoise`` - (shared weights — splitting them severs shared state). - * Unified UniRL: TWO separate experts — an ``llm`` prompt-refiner (bound to ``ar_decode``) and a - ``transformer`` flow generator (bound to ``diffusion_denoise``). The LM refines the prompt; its - refinement conditions the generator. They do NOT share weights (UniRL keeps Qwen and FLUX as - distinct experts), yet BOTH are trainable policies updated under one RL reward. + * Cosmos3 MoT: one ``transformer`` bound to both ``ar_decode`` and ``diffusion_denoise`` (shared weights — + splitting them severs shared state). + * Unified UniRL: two separate experts — an ``llm`` prompt-refiner (bound to ``ar_decode``) and a + ``transformer`` flow generator (bound to ``diffusion_denoise``). The LM refines the prompt, which + conditions the generator. They do not share weights (Qwen and FLUX stay distinct), yet both are trainable + policies updated under one RL reward. -The stress-test claim: the Card/Loop/Program split that expresses MoT weight-sharing *also* expresses -two-separate-experts-joint-RL with no new primitive — only a second ``shared_weight_components`` target -and a method (``unified_rl``) that emits two WeightSyncPlans. The loops, caches, parity gate, rollout -capture, and serve path are unchanged. +The stress-test claim: the Card/Loop/Program split that expresses MoT weight-sharing also expresses +two-separate-experts-joint-RL with no new primitive — only a second ``shared_weight_components`` target and a +method (``unified_rl``) that emits two WeightSyncPlans. Loops, caches, parity gate, rollout capture, and serve +path are unchanged. """ from __future__ import annotations @@ -50,51 +50,73 @@ def build_unified_card(model_id: str = "unirl-qwenflux") -> ModelCard: return ARDecodeLoop(loop_id="ar_decode", transformer_id="llm", cost=ar_cost, max_tokens=1) def denoise_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 = { - "text_encoder": ComponentSpec("text_encoder", kind="text_encoder", - factory=lambda inst: ToyTextEncoder(), - required_for={"reason", "t2v"}), - "llm": ComponentSpec( # the prompt-refiner expert (Qwen role) - "llm", kind="dit", + "text_encoder": + ComponentSpec("text_encoder", + kind="text_encoder", + factory=lambda inst: ToyTextEncoder(), + required_for={"reason", "t2v"}), + "llm": + ComponentSpec( # the prompt-refiner expert (Qwen role) + "llm", + kind="dit", load_id="fastvideo.models.llm:QwenPromptRefiner", factory=lambda inst: ToyPromptRefiner(n_actions=N_REFINE_ACTIONS, seed=seed + 1), - resident_for=["ar_decode"], required_for={"reason", "t2v"}), - "transformer": ComponentSpec( # the flow generator expert (FLUX role) - "transformer", kind="dit", + resident_for=["ar_decode"], + required_for={"reason", "t2v"}), + "transformer": + ComponentSpec( # the flow generator expert (FLUX role) + "transformer", + kind="dit", load_id="fastvideo.models.dits.flux:FluxTransformer", 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"}), + resident_for=["diffusion_denoise"], + required_for={"t2v"}), + "vae": + ComponentSpec("vae", kind="vae", factory=lambda inst: ToyVAE(), required_for={"t2v"}), } loops = { - "ar_decode": LoopSpec("ar_decode", kind=LoopKind.AR_DECODE, work_unit_kind=WorkUnitKind.AR_TOKEN, - step_cost_model=ar_cost, shared_weight_components=["llm"], - cache_policy=["paged_kv"], loop_factory=refine_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=denoise_factory), + "ar_decode": + LoopSpec("ar_decode", + kind=LoopKind.AR_DECODE, + work_unit_kind=WorkUnitKind.AR_TOKEN, + step_cost_model=ar_cost, + shared_weight_components=["llm"], + cache_policy=["paged_kv"], + loop_factory=refine_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=denoise_factory), } card = ModelCard( - model_id=model_id, family="unified", components=components, loops=loops, - capabilities=CapabilityMatrix.of( - Capability.TEXT_TO_VIDEO, Capability.REASONING_TEXT, - Capability.VAE_DECODE, Capability.POLICY_ROLLOUT), - recipe=RecipeSpec(method="base", assumes_loop="diffusion_denoise", - assumes_precision="float32", consistency_required=ConsistencyLevel.C2), - parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1, ConsistencyLevel.C2], - interleave_required=True), + model_id=model_id, + family="unified", + components=components, + loops=loops, + capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.REASONING_TEXT, Capability.VAE_DECODE, + Capability.POLICY_ROLLOUT), + recipe=RecipeSpec(method="base", + assumes_loop="diffusion_denoise", + assumes_precision="float32", + consistency_required=ConsistencyLevel.C2), + parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1, ConsistencyLevel.C2], 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() diff --git a/v2/recipes/unified/program.py b/v2/recipes/unified/program.py index e9c47abdd..825d8bcd0 100644 --- a/v2/recipes/unified/program.py +++ b/v2/recipes/unified/program.py @@ -1,12 +1,12 @@ -"""Unified LM+generator serve program (design_v3 §15b; UniRL/PromptRL inference path): +"""Unified LM+generator serve program (UniRL/PromptRL inference path): text_encode → refine (ar_decode on llm) → apply_refinement → denoise (diffusion_denoise) → vae_decode The serve path is deterministic: the LM refines greedily (argmax action) and the generator runs the ODE Euler sampler. The ``unified_rl`` *training* method drives the same two loops but in the ROLLOUT profile — the LM *samples* an action and the generator uses the SDE sampler with per-step log-prob -capture (design_v3 §9.4: rollout sampler ≠ serve sampler, a controlled divergence). Same loops, two -profiles — the §10 "train ≡ serve through one loop" invariant, now spanning two loop *types*. +capture (rollout sampler ≠ serve sampler, a controlled divergence). Same loops, two profiles — the +"train ≡ serve through one loop" invariant, now spanning two loop *types*. """ from __future__ import annotations @@ -18,6 +18,7 @@ from v2.recipes.unified.card import N_REFINE_ACTIONS def apply_refinement_node(refine_slot: str, n_actions: int = N_REFINE_ACTIONS): """ComponentNode fn: the LM's chosen refinement action reshapes the diffusion conditioning.""" + def fn(instance, slots, request, ctx) -> None: llm = instance.component("llm") out = slots.get(refine_slot) @@ -25,23 +26,37 @@ def apply_refinement_node(refine_slot: str, n_actions: int = N_REFINE_ACTIONS): action = (int(toks[0]) % n_actions) if toks else 0 slots["text_embeds"] = llm.refined_embed(slots.get("text_embeds"), action) slots["refine_action"] = action + return fn def build_unified_program() -> Program: return Program( - program_id="unified.unirl", kind=ProgramKind.INLINE, + program_id="unified.unirl", + kind=ProgramKind.INLINE, nodes=[ - ComponentNode("text_encode", fn=text_encode_node_fn, - writes=("text_embeds", "neg_text_embeds")), - ModelLoopNode("refine", loop_id="ar_decode", output_slot="refine_out", - reads=("text_embeds",), writes=("refine_out",)), - ComponentNode("apply_refine", fn=apply_refinement_node("refine_out"), - reads=("refine_out", "text_embeds"), writes=("text_embeds", "refine_action")), - 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("text_encode", fn=text_encode_node_fn, writes=("text_embeds", "neg_text_embeds")), + ModelLoopNode("refine", + loop_id="ar_decode", + output_slot="refine_out", + reads=("text_embeds", ), + writes=("refine_out", )), + ComponentNode("apply_refine", + fn=apply_refinement_node("refine_out"), + reads=("refine_out", "text_embeds"), + writes=("text_embeds", "refine_action")), + 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", "latents": "denoise_out"}, + output_artifacts={ + "video": "video", + "latents": "denoise_out" + }, ).validate() diff --git a/v2/recipes/wan21/card.py b/v2/recipes/wan21/card.py index df047f6f5..eae41384b 100644 --- a/v2/recipes/wan21/card.py +++ b/v2/recipes/wan21/card.py @@ -1,4 +1,4 @@ -"""Wan2.1-1.3B ModelCard (design_v3 §4) — the (recipe, runtime) pair for T2V. +"""Wan2.1-1.3B ModelCard — the (recipe, runtime) pair for T2V. Mirrors the real Wan2.1-1.3B wiring discovered in the repo: * DiT ``fastvideo.models.dits.wanvideo:WanTransformer3DModel`` (40 layers, dim 5120) @@ -37,9 +37,14 @@ from v2.recipes._prompts import WAN_NEG_CN, WAN_NEG_EN from v2.recipes.wan21.loop import WanDenoiseLoop -def build_wan21_card(model_id: str = "wan2.1-1.3b", *, cfg_policy=None, flow_shift: float = 3.0, - checkpoint_root: str | None = None, latent_channels: int = 16, - spatial_ratio: int = 8, temporal_ratio: int = 4, +def build_wan21_card(model_id: str = "wan2.1-1.3b", + *, + cfg_policy=None, + flow_shift: float = 3.0, + checkpoint_root: str | None = None, + latent_channels: int = 16, + spatial_ratio: int = 8, + temporal_ratio: int = 4, sampling_defaults: SamplingDefaults | None = None) -> ModelCard: seed = _seed_from(model_id) cost = CostModel(kind=WorkUnitKind.DIFFUSION_STEP, base_seconds=1e-4, per_unit_seconds=1e-7) @@ -49,54 +54,75 @@ def build_wan21_card(model_id: str = "wan2.1-1.3b", *, cfg_policy=None, flow_shi expert = NoRouting("transformer") def loop_factory(): - return WanDenoiseLoop(loop_id="diffusion_denoise", cfg=cfg, flow_shift=flow, - precision=precision, expert=expert, cost=cost, - latent_channels=latent_channels, spatial_ratio=spatial_ratio, + return WanDenoiseLoop(loop_id="diffusion_denoise", + cfg=cfg, + flow_shift=flow, + precision=precision, + expert=expert, + cost=cost, + latent_channels=latent_channels, + spatial_ratio=spatial_ratio, temporal_ratio=temporal_ratio) components = { - "text_encoder": ComponentSpec( - component_id="text_encoder", kind="text_encoder", - load_id="fastvideo.models.encoders.t5:T5EncoderModel", - factory=lambda inst: ToyTextEncoder(), required_for={"t2v"}), - "vae": ComponentSpec( - component_id="vae", kind="vae", - load_id="fastvideo.models.vaes.wanvae:AutoencoderKLWan", - factory=lambda inst: ToyVAE(), required_for={"t2v"}), - "transformer": ComponentSpec( - component_id="transformer", kind="dit", - load_id="fastvideo.models.dits.wanvideo:WanTransformer3DModel", - factory=lambda inst: ToyDiT(seed=seed), - resident_for=["diffusion_denoise"], required_for={"t2v"}), + "text_encoder": + ComponentSpec(component_id="text_encoder", + kind="text_encoder", + load_id="fastvideo.models.encoders.t5:T5EncoderModel", + factory=lambda inst: ToyTextEncoder(), + required_for={"t2v"}), + "vae": + ComponentSpec(component_id="vae", + kind="vae", + load_id="fastvideo.models.vaes.wanvae:AutoencoderKLWan", + factory=lambda inst: ToyVAE(), + required_for={"t2v"}), + "transformer": + ComponentSpec(component_id="transformer", + kind="dit", + load_id="fastvideo.models.dits.wanvideo:WanTransformer3DModel", + factory=lambda inst: ToyDiT(seed=seed), + resident_for=["diffusion_denoise"], + required_for={"t2v"}), } loops = { - "diffusion_denoise": LoopSpec( - loop_id="diffusion_denoise", kind=LoopKind.DIFFUSION_DENOISE, - work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=cost, - shared_weight_components=["transformer"], cache_policy=["feature"], - graph_capture="breakable_cudagraph", loop_factory=loop_factory), + "diffusion_denoise": + LoopSpec(loop_id="diffusion_denoise", + kind=LoopKind.DIFFUSION_DENOISE, + work_unit_kind=WorkUnitKind.DIFFUSION_STEP, + step_cost_model=cost, + shared_weight_components=["transformer"], + cache_policy=["feature"], + graph_capture="breakable_cudagraph", + loop_factory=loop_factory), } card = 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, # base Wan2.1 is T2V-only; i2v is the separate InP variant - 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, - tests=[ParityTestSpec(name="denoise_trajectory", level=ConsistencyLevel.C1, - tap="latents")]), - caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, - reuse_across_requests=True)}, + Capability.TEXT_TO_VIDEO, # base Wan2.1 is T2V-only; i2v is the separate InP variant + 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, + tests=[ParityTestSpec(name="denoise_trajectory", level=ConsistencyLevel.C1, tap="latents")]), + caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)}, precision=PrecisionContract(default_dtype="float32", training_precision="float32"), parallelism=ParallelismContract( valid_plans=[ParallelPlan.single(), - ParallelPlan(axes={"sp": 2, "cfgp": 2}, mesh_order=["cfgp", "sp"])], + ParallelPlan(axes={ + "sp": 2, + "cfgp": 2 + }, mesh_order=["cfgp", "sp"])], default_plan=ParallelPlan.single()), sampling_defaults=sampling_defaults or SamplingDefaults( - num_steps=50, guidance_scale=3.0, height=480, width=832, num_frames=81, fps=16, - negative_prompt=WAN_NEG_EN), + num_steps=50, guidance_scale=3.0, height=480, width=832, num_frames=81, fps=16, negative_prompt=WAN_NEG_EN), ) card.validate() if checkpoint_root: @@ -104,31 +130,46 @@ def build_wan21_card(model_id: str = "wan2.1-1.3b", *, cfg_policy=None, flow_shi return card -def build_wan22_ti2v_card(model_id: str = "wan2.2-ti2v-5b", *, - checkpoint_root: str | None = None) -> ModelCard: +def build_wan22_ti2v_card(model_id: str = "wan2.2-ti2v-5b", *, checkpoint_root: str | None = None) -> ModelCard: """Wan2.2-TI2V-5B card. Same component classes as Wan2.1 (WanTransformer3DModel / AutoencoderKLWan / UMT5), so the torch adapters are reused as-is — the deltas are the higher-compression VAE geometry (z_dim=48, 16x spatial, 4x temporal) and the 480p flow shift 5.0. The DiT forward accepts a scalar timestep (its 1D path), so no per-frame ``expand_timesteps`` is needed for pure t2v. Uses the same ``transformer/vae/text_encoder`` subfolder layout, so ``stamp_wan21_checkpoints`` applies.""" - return build_wan21_card(model_id=model_id, flow_shift=5.0, checkpoint_root=checkpoint_root, - latent_channels=48, spatial_ratio=16, temporal_ratio=4, - sampling_defaults=SamplingDefaults( - num_steps=50, guidance_scale=5.0, height=704, width=1280, - num_frames=121, fps=24, negative_prompt=WAN_NEG_CN)) + return build_wan21_card(model_id=model_id, + flow_shift=5.0, + checkpoint_root=checkpoint_root, + latent_channels=48, + spatial_ratio=16, + temporal_ratio=4, + sampling_defaults=SamplingDefaults(num_steps=50, + guidance_scale=5.0, + height=704, + width=1280, + num_frames=121, + fps=24, + negative_prompt=WAN_NEG_CN)) def build_wan_t2v_14b_card(model_id: str = "wan2.1-t2v-14b", *, checkpoint_root: str | None = None) -> ModelCard: """Wan2.1-T2V-14B (720p). Same WanTransformer3DModel / AutoencoderKLWan / UMT5 + 16/8/4 VAE geometry as the 1.3B — reuses the Wan recipe + torch adapter unchanged; only the size (resolved from the checkpoint by the loader), the 720p flow shift 5.0, and the per-model defaults differ.""" - return build_wan21_card(model_id=model_id, flow_shift=5.0, checkpoint_root=checkpoint_root, - sampling_defaults=SamplingDefaults( - num_steps=50, guidance_scale=5.0, height=720, width=1280, - num_frames=81, fps=16, negative_prompt=WAN_NEG_EN)) + return build_wan21_card(model_id=model_id, + flow_shift=5.0, + checkpoint_root=checkpoint_root, + sampling_defaults=SamplingDefaults(num_steps=50, + guidance_scale=5.0, + height=720, + width=1280, + num_frames=81, + fps=16, + negative_prompt=WAN_NEG_EN)) -def build_wan22_a14b_card(model_id: str = "wan2.2-t2v-a14b", *, checkpoint_root: str | None = None, +def build_wan22_a14b_card(model_id: str = "wan2.2-t2v-a14b", + *, + checkpoint_root: str | None = None, boundary: float = 0.875) -> ModelCard: """Wan2.2-T2V-A14B card — MoE: two WanTransformer3DModel experts (in_ch=16, Wan2.1 16/8/4 geometry) with a boundary-timestep switch (``boundary_ratio`` 0.875). ``transformer`` is the high-noise expert @@ -145,53 +186,73 @@ def build_wan22_a14b_card(model_id: str = "wan2.2-t2v-a14b", *, checkpoint_root: expert = BoundaryTimestepRouting(high_noise="transformer", low_noise="transformer_2", boundary=boundary) def loop_factory(): - return WanDenoiseLoop(loop_id="diffusion_denoise", cfg=cfg, flow_shift=flow, - precision=precision, expert=expert, cost=cost) # 16/8/4 default geometry + return WanDenoiseLoop(loop_id="diffusion_denoise", + cfg=cfg, + flow_shift=flow, + precision=precision, + expert=expert, + cost=cost) # 16/8/4 default geometry def _dit_spec(cid): - return ComponentSpec( - component_id=cid, kind="dit", - load_id="fastvideo.models.dits.wanvideo:WanTransformer3DModel", - factory=lambda inst: ToyDiT(seed=seed), - resident_for=["diffusion_denoise"], required_for={"t2v"}) + return ComponentSpec(component_id=cid, + kind="dit", + load_id="fastvideo.models.dits.wanvideo:WanTransformer3DModel", + factory=lambda inst: ToyDiT(seed=seed), + resident_for=["diffusion_denoise"], + required_for={"t2v"}) components = { - "text_encoder": ComponentSpec( - component_id="text_encoder", kind="text_encoder", - load_id="fastvideo.models.encoders.t5:T5EncoderModel", - factory=lambda inst: ToyTextEncoder(), required_for={"t2v"}), - "vae": ComponentSpec( - component_id="vae", kind="vae", - load_id="fastvideo.models.vaes.wanvae:AutoencoderKLWan", - factory=lambda inst: ToyVAE(), required_for={"t2v"}), - "transformer": _dit_spec("transformer"), - "transformer_2": _dit_spec("transformer_2"), + "text_encoder": + ComponentSpec(component_id="text_encoder", + kind="text_encoder", + load_id="fastvideo.models.encoders.t5:T5EncoderModel", + factory=lambda inst: ToyTextEncoder(), + required_for={"t2v"}), + "vae": + ComponentSpec(component_id="vae", + kind="vae", + load_id="fastvideo.models.vaes.wanvae:AutoencoderKLWan", + factory=lambda inst: ToyVAE(), + required_for={"t2v"}), + "transformer": + _dit_spec("transformer"), + "transformer_2": + _dit_spec("transformer_2"), } loops = { - "diffusion_denoise": LoopSpec( - loop_id="diffusion_denoise", kind=LoopKind.DIFFUSION_DENOISE, - work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=cost, - shared_weight_components=["transformer", "transformer_2"], cache_policy=["feature"], - graph_capture="breakable_cudagraph", loop_factory=loop_factory), + "diffusion_denoise": + LoopSpec(loop_id="diffusion_denoise", + kind=LoopKind.DIFFUSION_DENOISE, + work_unit_kind=WorkUnitKind.DIFFUSION_STEP, + step_cost_model=cost, + shared_weight_components=["transformer", "transformer_2"], + cache_policy=["feature"], + graph_capture="breakable_cudagraph", + loop_factory=loop_factory), } card = ModelCard( - model_id=model_id, family="wan", - components=components, loops=loops, - capabilities=CapabilityMatrix.of( - Capability.TEXT_TO_VIDEO, 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, - tests=[ParityTestSpec(name="denoise_trajectory", level=ConsistencyLevel.C1, - tap="latents")]), - caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, - reuse_across_requests=True)}, + model_id=model_id, + family="wan", + components=components, + loops=loops, + capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, 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, + tests=[ParityTestSpec(name="denoise_trajectory", level=ConsistencyLevel.C1, tap="latents")]), + caches={"feature": CacheContract(cache_class="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()), - sampling_defaults=SamplingDefaults( - num_steps=40, guidance_scale=4.0, height=480, width=832, num_frames=81, fps=16, - negative_prompt=WAN_NEG_CN), + parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()), + sampling_defaults=SamplingDefaults(num_steps=40, + guidance_scale=4.0, + height=480, + width=832, + num_frames=81, + fps=16, + negative_prompt=WAN_NEG_CN), ) card.validate() if checkpoint_root: @@ -202,11 +263,16 @@ def build_wan22_a14b_card(model_id: str = "wan2.2-t2v-a14b", *, checkpoint_root: # Diffusers checkpoint layout: each component's weights live in a subfolder of the model root. The stamp # loop skips any subfolder a given card doesn't declare, so this superset covers Wan2.1, Wan2.2 MoE # (``transformer_2``), and LTX-2 (``spatial_upsampler``) alike. -_WAN21_SUBFOLDERS = {"transformer": "transformer", "transformer_2": "transformer_2", - "vae": "vae", "text_encoder": "text_encoder", - "image_encoder": "image_encoder", # Wan i2v CLIP vision - "spatial_upsampler": "spatial_upsampler", - "audio_vae": "audio_vae", "vocoder": "vocoder"} # LTX-2.3 T2VS audio branch +_WAN21_SUBFOLDERS = { + "transformer": "transformer", + "transformer_2": "transformer_2", + "vae": "vae", + "text_encoder": "text_encoder", + "image_encoder": "image_encoder", # Wan i2v CLIP vision + "spatial_upsampler": "spatial_upsampler", + "audio_vae": "audio_vae", + "vocoder": "vocoder" +} # LTX-2.3 T2VS audio branch def stamp_wan21_checkpoints(card: ModelCard, model_root: str) -> ModelCard: diff --git a/v2/recipes/wan21/loop.py b/v2/recipes/wan21/loop.py index 1783d3b05..d81a95059 100644 --- a/v2/recipes/wan21/loop.py +++ b/v2/recipes/wan21/loop.py @@ -1,10 +1,10 @@ -"""WanDenoiseLoop — the canonical driven denoise loop (design_v3 §5, §15a). +"""WanDenoiseLoop — the canonical driven denoise loop. Bidirectional video diffusion: N flow-match steps over the full latent. Composed from policies -(CFG / flow-shift / precision / expert routing) — the default decomposition. ``next`` is -kernel-free (it builds the forward+combine+solver thunk); ``advance`` folds the result and, -under the ROLLOUT profile, captures a behavior slice — so the *same* loop serves and rolls out -for RL (design_v3 §10). All per-request state lives in ``LoopState`` (interleave-safe). +(CFG / flow-shift / precision / expert routing). ``next`` is kernel-free (it builds the +forward+combine+solver thunk); ``advance`` folds the result and, under the ROLLOUT profile, +captures a behavior slice — so the *same* loop serves and rolls out for RL. All per-request +state lives in ``LoopState`` (interleave-safe). """ from __future__ import annotations @@ -25,14 +25,16 @@ from v2.platform import FLOW_MATCH_STEP, FLOW_SDE_STEP from v2.request.streams import StreamChunk from v2.platform.backends.toy import LATENT_CHANNELS - # Real Wan2.1 VAE (AutoencoderKLWan) compression: z_dim=16, 4x temporal, 8x spatial. WAN_LATENT_CHANNELS = 16 WAN_TEMPORAL_RATIO = 4 WAN_SPATIAL_RATIO = 8 -def latent_shape(req, model=None, *, channels=WAN_LATENT_CHANNELS, +def latent_shape(req, + model=None, + *, + channels=WAN_LATENT_CHANNELS, spatial_ratio=WAN_SPATIAL_RATIO, temporal_ratio=WAN_TEMPORAL_RATIO) -> tuple[int, int, int, int]: """Latent geometry for the denoise loop. On the real (GPU) backend the DiT/VAE require the true Wan @@ -42,8 +44,7 @@ def latent_shape(req, model=None, *, channels=WAN_LATENT_CHANNELS, d = req.diffusion if model is not None and getattr(getattr(model, "platform", None), "device", "cpu") == "cuda": t = (max(1, d.num_frames) - 1) // temporal_ratio + 1 - return (channels, max(1, t), - max(1, d.height // spatial_ratio), max(1, d.width // spatial_ratio)) + return (channels, max(1, t), max(1, d.height // spatial_ratio), max(1, d.width // spatial_ratio)) t = max(1, d.num_frames // 40) h = max(2, d.height // 120) w = max(2, d.width // 120) @@ -51,8 +52,17 @@ def latent_shape(req, model=None, *, channels=WAN_LATENT_CHANNELS, class WanDenoiseLoop: - def __init__(self, *, loop_id, cfg, flow_shift, precision, expert, cost, - latent_channels=WAN_LATENT_CHANNELS, spatial_ratio=WAN_SPATIAL_RATIO, + + def __init__(self, + *, + loop_id, + cfg, + flow_shift, + precision, + expert, + cost, + latent_channels=WAN_LATENT_CHANNELS, + spatial_ratio=WAN_SPATIAL_RATIO, temporal_ratio=WAN_TEMPORAL_RATIO): self.loop_id = loop_id self.cfg = cfg @@ -67,14 +77,24 @@ class WanDenoiseLoop: 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) - sig = self.flow_shift.build_schedule(req.diffusion.num_steps, req.diffusion.height, - req.diffusion.width, sigmas=req.diffusion.sigmas or None) - shape = latent_shape(req, model, channels=self.latent_channels, - spatial_ratio=self.spatial_ratio, temporal_ratio=self.temporal_ratio) + sig = self.flow_shift.build_schedule(req.diffusion.num_steps, + req.diffusion.height, + req.diffusion.width, + sigmas=req.diffusion.sigmas or None) + shape = latent_shape(req, + model, + channels=self.latent_channels, + spatial_ratio=self.spatial_ratio, + temporal_ratio=self.temporal_ratio) x = (rng.standard_normal(shape) * float(sig[0])).astype("float32") - st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id, - request_id=req.request_id, profile=ctx.profile, rng=rng, seed=seed, - latents={"video": x}, sigmas=[float(s) for s in sig], + st = LoopState(loop_id=self.loop_id, + instance_id=model.card.model_id, + request_id=req.request_id, + profile=ctx.profile, + rng=rng, + seed=seed, + latents={"video": x}, + sigmas=[float(s) for s in sig], timesteps=[float(s) * 1000.0 for s in sig]) st.cond["prompt_embeds"] = ctx.slots.get("text_embeds") st.cond["negative_prompt_embeds"] = ctx.slots.get("neg_text_embeds") @@ -110,12 +130,14 @@ class WanDenoiseLoop: def _velocity(model, x_, sigma_t_, pe_, ne_, scale_): # The conditioned forward + CFG combine. Solver/forward dispatch through the platform's - # kernel table (numpy on CPU, the device kernel on a GPU/accel backend — design_v3 §17). + # kernel table (numpy on CPU, the device kernel on a GPU/accel backend). # i2v_ctx/i2v_cond are None for T2V (the plain forward); for i2v the adapter concats the # conditioning latent and passes the CLIP embeds as encoder_hidden_states_image. dit = model.component(expert_id) - preds = {b: dit(x_, pe_ if b == "cond" else ne_, sigma_t_, context=i2v_ctx, cond=i2v_cond) - for b in branches} + preds = { + b: dit(x_, pe_ if b == "cond" else ne_, sigma_t_, context=i2v_ctx, cond=i2v_cond) + for b in branches + } return precision.cast(cfg.combine(preds, scale_, sctx, cfg_state)) def run(model, override=None): @@ -126,44 +148,66 @@ class WanDenoiseLoop: velocity = precision.cast(np.asarray(override["noise_pred"], dtype="float32")) else: velocity = _velocity(model, x, sigma_t, pe, ne, scale) - if sde: # FlowGRPO rollout: stochastic + log-prob + if sde: # FlowGRPO rollout: stochastic + log-prob noise = rng.standard_normal(x.shape) - x_next, logp, _m, _s = kernels.get(FLOW_SDE_STEP)( - precision.cast(x), velocity, sigma_t, sigma_next, noise=noise, noise_scale=noise_scale) - return StepResult(output={"noise_pred": np.asarray(velocity, dtype="float32"), - "latents": x_next.astype("float32"), - "sde_logprob": logp, "prev": np.asarray(x, dtype="float32")}) + x_next, logp, _m, _s = kernels.get(FLOW_SDE_STEP)(precision.cast(x), + velocity, + sigma_t, + sigma_next, + noise=noise, + noise_scale=noise_scale) + return StepResult( + output={ + "noise_pred": np.asarray(velocity, dtype="float32"), + "latents": x_next.astype("float32"), + "sde_logprob": logp, + "prev": np.asarray(x, dtype="float32") + }) x_next = kernels.get(FLOW_MATCH_STEP)(precision.cast(x), velocity, sigma_t, sigma_next) - return StepResult(output={"noise_pred": np.asarray(velocity, dtype="float32"), - "latents": x_next.astype("float32")}) + return StepResult(output={ + "noise_pred": np.asarray(velocity, dtype="float32"), + "latents": x_next.astype("float32") + }) def graph_fn(model, ws): # The CAPTURABLE op-structure: reads EVERY per-step input from the static workspace (never # from closure over per-step data) and writes into the static output buffer — so the same - # captured graph replays correctly with rebound buffer contents (design_v3 §6.2). + # captured graph replays correctly with rebound buffer contents. 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"])) x_next = model.platform.kernels.get(FLOW_MATCH_STEP)(precision.cast(ws["x"]), velocity, st_t, st_n) - np.copyto(ws["out"], x_next.astype("float32")) # write the static output buffer - return StepResult(output={"noise_pred": np.asarray(velocity, dtype="float32"), - "latents": np.array(ws["out"], copy=True)}) + np.copyto(ws["out"], x_next.astype("float32")) # write the static output buffer + return StepResult(output={ + "noise_pred": np.asarray(velocity, dtype="float32"), + "latents": np.array(ws["out"], copy=True) + }) cond_bytes = sum(int(np.asarray(e).nbytes) for e in (pe, ne) if e is not None) res = ResourceRequest( compute_seconds=self.cost.predict(int(np.prod(x.shape)), float(len(branches))), - resident_bytes=int(x.nbytes) + cond_bytes, # latents + conditioning held for the loop - peak_activation_bytes=int(x.nbytes)) # one step's transient working buffer + resident_bytes=int(x.nbytes) + cond_bytes, # latents + conditioning held for the loop + peak_activation_bytes=int(x.nbytes)) # one step's transient working buffer emits = [] if st.scratch.get("stream_video"): - emits.append(StreamChunk(stream_id=st.request_id, modality="video", seq=i, - data=x, preview=True)) # carry the latent as a preview payload + emits.append(StreamChunk(stream_id=st.request_id, modality="video", seq=i, data=x, + preview=True)) # carry the latent as a preview payload return WorkPlan( - loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.DIFFUSION_STEP, - shape_sig=ShapeSignature(WorkUnitKind.DIFFUSION_STEP, dims=tuple(x.shape), - dtype=precision.compute_dtype, # compute dtype is part of the key - extra=(("cfg", type(cfg).__name__),)), - resources=res, payload={"branch": "combined", "step": i}, run=run, - label=f"wan.denoise.{i}", emits=emits, + loop_id=self.loop_id, + instance_id=st.instance_id, + kind=WorkUnitKind.DIFFUSION_STEP, + shape_sig=ShapeSignature( + WorkUnitKind.DIFFUSION_STEP, + dims=tuple(x.shape), + dtype=precision.compute_dtype, # compute dtype is part of the key + extra=(("cfg", type(cfg).__name__), )), + resources=res, + payload={ + "branch": "combined", + "step": i + }, + run=run, + label=f"wan.denoise.{i}", + emits=emits, # CUDA-graph capture (Path A): the deterministic ODE step is capturable; the stochastic # SDE rollout step has host RNG / a data-dependent branch, so it must eager-break. The # op-structure key carries the CFG branch set, active expert, and the scheduler-precision @@ -171,25 +215,35 @@ class WanDenoiseLoop: # incompatible captured graph. (Compute dtype rides in shape_sig.dtype above.) NOTE: this # is sound for branch-set/expert-determined policies (ClassicCFG); a policy whose op # structure forks on other state (PerModalityCFG modality) would need a graph_key hook. - capturable=not sde and i2v_cond is None, # i2v threads non-workspace conditioning -> eager + capturable=not sde and i2v_cond is None, # i2v threads non-workspace conditioning -> eager graph_key=(tuple(sorted(branches)), expert_id, precision.scheduler_step_in_fp32), # the static-buffer capture form: graph_fn reads all per-step inputs from the workspace graph_fn=graph_fn, - graph_inputs={"x": x, "sigma_t": sigma_t, "sigma_next": sigma_next, - "pe": pe, "ne": ne, "scale": scale}) + graph_inputs={ + "x": x, + "sigma_t": sigma_t, + "sigma_next": sigma_next, + "pe": pe, + "ne": ne, + "scale": scale + }) def advance(self, st: LoopState, result: StepResult) -> LoopState: st.latents["video"] = result.output["latents"] if st.profile == ExecutionProfile.ROLLOUT: i = st.step_idx - rec = {"step": i, "sigma": st.sigmas[i], - "velocity": np.asarray(result.output["noise_pred"]).copy(), - "latents": np.asarray(st.latents["video"]).copy()} - if "sde_logprob" in result.output: # FlowGRPO: capture the PPO log-prob slice + rec = { + "step": i, + "sigma": st.sigmas[i], + "velocity": np.asarray(result.output["noise_pred"]).copy(), + "latents": np.asarray(st.latents["video"]).copy() + } + if "sde_logprob" in result.output: # FlowGRPO: capture the PPO log-prob slice rec.update(sde_logprob=result.output["sde_logprob"], prev=np.asarray(result.output["prev"]).copy(), sample=np.asarray(result.output["latents"]).copy(), - sigma_t=st.sigmas[i], sigma_next=st.sigmas[i + 1]) + sigma_t=st.sigmas[i], + sigma_next=st.sigmas[i + 1]) st.trajectory.append(rec) st.step_idx += 1 return st diff --git a/v2/recipes/wan21/program.py b/v2/recipes/wan21/program.py index 78f9b2df5..1ef73eb3b 100644 --- a/v2/recipes/wan21/program.py +++ b/v2/recipes/wan21/program.py @@ -1,8 +1,7 @@ -"""Wan2.1 T2V program (design_v3 §13, §15a): text_encode → diffusion_denoise → vae_decode. +"""Wan2.1 T2V program: text_encode → diffusion_denoise → vae_decode. -The text-encode node demonstrates the content-hash feature cache (design_v3 §7.2, §10): the -same prompt across requests (e.g. a K=24 RL group) encodes once and reuses — the 24× text-encode -reduction the vendored RL sampler cannot express. +The text-encode node uses the content-hash feature cache: the same prompt across requests (e.g. a +K=24 RL group) encodes once and reuses. """ from __future__ import annotations @@ -17,12 +16,19 @@ def _vae_decode(instance, slots, request, ctx) -> None: def build_wan_t2v_program() -> Program: return Program( - program_id="wan.t2v.inline", kind=ProgramKind.INLINE, + program_id="wan.t2v.inline", + kind=ProgramKind.INLINE, nodes=[ ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")), - ModelLoopNode("denoise", loop_id="diffusion_denoise", output_slot="denoise_out", - reads=("text_embeds", "neg_text_embeds"), writes=("denoise_out",)), - ComponentNode("vae_decode", fn=_vae_decode, reads=("denoise_out",), writes=("video",)), + ModelLoopNode("denoise", + loop_id="diffusion_denoise", + output_slot="denoise_out", + reads=("text_embeds", "neg_text_embeds"), + writes=("denoise_out", )), + ComponentNode("vae_decode", fn=_vae_decode, reads=("denoise_out", ), writes=("video", )), ], - output_artifacts={"video": "video", "latents": "denoise_out"}, + output_artifacts={ + "video": "video", + "latents": "denoise_out" + }, ).validate() diff --git a/v2/recipes/wan_causal/card.py b/v2/recipes/wan_causal/card.py index 0cef45f46..5d27e3647 100644 --- a/v2/recipes/wan_causal/card.py +++ b/v2/recipes/wan_causal/card.py @@ -1,9 +1,9 @@ -"""Wan-causal (self-forcing) ModelCard — causal/streaming video (design_v3 §4, §10). +"""Wan-causal (self-forcing) ModelCard — causal/streaming video. -This is the causal student *created by* self-forcing distillation: ``recipe.method='self_forcing'``, -``assumes_loop='chunk_rollout'``. It declares a ``slab_kv`` cache class (chunk-KV) in addition to the -text feature cache. The chunk_rollout loop is causal/AR, so ``pp_patch`` parallelism is rejected for -this card by the validator (stale KV breaks causality). +The causal student *created by* self-forcing distillation: ``recipe.method='self_forcing'``, +``assumes_loop='chunk_rollout'``. Declares a ``slab_kv`` cache class (chunk-KV) alongside the text +feature cache. The chunk_rollout loop is causal/AR, so the validator rejects ``pp_patch`` parallelism +for this card (stale KV breaks causality). """ from __future__ import annotations diff --git a/v2/recipes/wan_causal/loop.py b/v2/recipes/wan_causal/loop.py index ca5c614d6..555d5e0bb 100644 --- a/v2/recipes/wan_causal/loop.py +++ b/v2/recipes/wan_causal/loop.py @@ -1,9 +1,9 @@ -"""ChunkRolloutLoop — causal/AR video (design_v3 §4 causal, §5; design.md taxonomy). +"""ChunkRolloutLoop — causal/AR video. Loop shape: outer loop over latent-frame chunks × inner denoise per chunk, carrying a per-chunk KV/context via the slab-KV cache (sink/local window in inference; full history in self-forcing training mode). Each completed chunk is streamable (``StepResult.emit``), and the rollout profile -captures per-chunk behavior for the self-forcing distillation method (design_v3 §10). +captures per-chunk behavior for the self-forcing distillation method. State is entirely in ``LoopState`` (the slab list + chunk position), so interleaving causal rollouts of two sessions cannot smear context — and the slab-KV pool is namespaced per request. @@ -30,8 +30,8 @@ from v2.platform.backends.toy import LATENT_CHANNELS class ChunkRolloutLoop: - def __init__(self, *, loop_id, num_chunks, chunk_size, steps_per_chunk, cfg, flow_shift, - precision, cost): + + def __init__(self, *, loop_id, num_chunks, chunk_size, steps_per_chunk, cfg, flow_shift, precision, cost): self.loop_id = loop_id self.num_chunks = num_chunks self.chunk_size = chunk_size @@ -54,15 +54,23 @@ class ChunkRolloutLoop: seed = req.diffusion.seed if req.diffusion.seed is not None else 0 rng = np.random.default_rng(seed) sig = self.flow_shift.build_schedule(self.steps_per_chunk, req.diffusion.height, req.diffusion.width) - st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id, - request_id=req.request_id, profile=ctx.profile, rng=rng, seed=seed, - sigmas=[float(s) for s in sig], timesteps=[float(s) * 1000.0 for s in sig]) + st = LoopState(loop_id=self.loop_id, + instance_id=model.card.model_id, + request_id=req.request_id, + profile=ctx.profile, + rng=rng, + seed=seed, + sigmas=[float(s) for s in sig], + timesteps=[float(s) * 1000.0 for s in sig]) st.cond["prompt_embeds"] = ctx.slots.get("text_embeds") st.cond["negative_prompt_embeds"] = ctx.slots.get("neg_text_embeds") # world-model continuation: seed the chunk context from prior chunks (an interactive session's - # persistent world state; design_v3 §16). Default [] ⇒ a fresh rollout (unchanged one-shot path). + # persistent world state). Default [] ⇒ a fresh rollout (unchanged one-shot path). prior = [np.asarray(c, dtype="float32") for c in (ctx.slots.get("world_context") or [])] - st.scratch.update(chunk_idx=0, step_in_chunk=0, slabs=prior, chunks_out=[], + st.scratch.update(chunk_idx=0, + step_in_chunk=0, + slabs=prior, + chunks_out=[], guidance_scale=float(req.diffusion.guidance_scale), caches=getattr(model, "caches", None)) st.latents["chunk"] = (rng.standard_normal(self._chunk_shape(req, model)) * float(sig[0])).astype("float32") @@ -84,7 +92,7 @@ class ChunkRolloutLoop: cfg, precision = self.cfg, self.precision def run(model, override=None): - 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) if override is not None and "noise_pred" in override: velocity = np.asarray(override["noise_pred"], dtype="float32") else: @@ -92,22 +100,37 @@ class ChunkRolloutLoop: preds = {b: dit(x, pe if b == "cond" else ne, sigma_t, context=context) for b in branches} velocity = cfg.combine(preds, scale, sctx, cfg_state) x_next = fm(precision.cast(x), precision.cast(velocity), sigma_t, sigma_next) - return StepResult(output={"noise_pred": np.asarray(velocity, dtype="float32"), - "latents": x_next.astype("float32")}) + return StepResult(output={ + "noise_pred": np.asarray(velocity, dtype="float32"), + "latents": x_next.astype("float32") + }) emits = [] - if i == self.steps_per_chunk - 1: # last step of this chunk → streamable - emits.append(StreamChunk(stream_id=st.request_id, modality="video", - seq=st.scratch["chunk_idx"], data=x, preview=False)) - res = ResourceRequest( - compute_seconds=self.cost.predict(int(np.prod(x.shape)), float(len(branches))), - resident_bytes=int(x.nbytes), peak_activation_bytes=int(x.nbytes * len(branches))) - return WorkPlan( - loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.CHUNK_STEP, - shape_sig=ShapeSignature(WorkUnitKind.CHUNK_STEP, dims=tuple(x.shape), - extra=(("chunk", st.scratch["chunk_idx"]),)), - resources=res, payload={"branch": "combined", "chunk": st.scratch["chunk_idx"], "step": i}, - run=run, label=f"causal.c{st.scratch['chunk_idx']}.s{i}", emits=emits) + if i == self.steps_per_chunk - 1: # last step of this chunk → streamable + emits.append( + StreamChunk(stream_id=st.request_id, + modality="video", + seq=st.scratch["chunk_idx"], + data=x, + preview=False)) + res = ResourceRequest(compute_seconds=self.cost.predict(int(np.prod(x.shape)), float(len(branches))), + resident_bytes=int(x.nbytes), + peak_activation_bytes=int(x.nbytes * len(branches))) + return WorkPlan(loop_id=self.loop_id, + instance_id=st.instance_id, + kind=WorkUnitKind.CHUNK_STEP, + shape_sig=ShapeSignature(WorkUnitKind.CHUNK_STEP, + dims=tuple(x.shape), + extra=(("chunk", st.scratch["chunk_idx"]), )), + resources=res, + payload={ + "branch": "combined", + "chunk": st.scratch["chunk_idx"], + "step": i + }, + run=run, + label=f"causal.c{st.scratch['chunk_idx']}.s{i}", + emits=emits) def advance(self, st: LoopState, result: StepResult) -> LoopState: st.latents["chunk"] = result.output["latents"] @@ -115,10 +138,10 @@ class ChunkRolloutLoop: st.step_idx += 1 if st.scratch["step_in_chunk"] >= self.steps_per_chunk: chunk = np.asarray(st.latents["chunk"]).copy() - st.scratch["slabs"].append(chunk) # context for the next chunk + st.scratch["slabs"].append(chunk) # context for the next chunk st.scratch["chunks_out"].append(chunk) caches = st.scratch.get("caches") - if caches is not None and caches.has("slab_kv"): # demonstrate the slab-KV class + if caches is not None and caches.has("slab_kv"): # demonstrate the slab-KV class caches.pool("slab_kv").append(st.request_id, Slab(st.scratch["chunk_idx"], k=chunk, v=None)) if st.profile == ExecutionProfile.ROLLOUT: st.trajectory.append({"chunk": st.scratch["chunk_idx"], "latents": chunk}) @@ -130,6 +153,9 @@ class ChunkRolloutLoop: def finalize(self, st: LoopState) -> LoopResult: video = np.concatenate(st.scratch["chunks_out"], axis=1) if st.scratch["chunks_out"] else None - return LoopResult(outputs={"latents": video, "chunks": list(st.scratch["chunks_out"])}, + return LoopResult(outputs={ + "latents": video, + "chunks": list(st.scratch["chunks_out"]) + }, metrics={"chunks": float(st.scratch["chunk_idx"])}, behavior=st.trajectory or None) diff --git a/v2/recipes/wan_causal/program.py b/v2/recipes/wan_causal/program.py index 264a7c401..9a5504a2d 100644 --- a/v2/recipes/wan_causal/program.py +++ b/v2/recipes/wan_causal/program.py @@ -1,4 +1,4 @@ -"""Wan-causal streaming program: text_encode → chunk_rollout → vae_decode (design_v3 §4, §15e).""" +"""Wan-causal streaming program: text_encode → chunk_rollout → vae_decode.""" from __future__ import annotations from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind @@ -12,12 +12,19 @@ def _vae_decode(instance, slots, request, ctx) -> None: def build_wan_causal_program() -> Program: return Program( - program_id="wan_causal.stream", kind=ProgramKind.INLINE, + program_id="wan_causal.stream", + kind=ProgramKind.INLINE, nodes=[ ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")), - ModelLoopNode("rollout", loop_id="chunk_rollout", output_slot="rollout_out", - reads=("text_embeds",), writes=("rollout_out",)), - ComponentNode("vae_decode", fn=_vae_decode, reads=("rollout_out",), writes=("video",)), + ModelLoopNode("rollout", + loop_id="chunk_rollout", + output_slot="rollout_out", + reads=("text_embeds", ), + writes=("rollout_out", )), + ComponentNode("vae_decode", fn=_vae_decode, reads=("rollout_out", ), writes=("video", )), ], - output_artifacts={"video": "video", "latents": "rollout_out"}, + output_artifacts={ + "video": "video", + "latents": "rollout_out" + }, ).validate() diff --git a/v2/recipes/wan_fun_control/__init__.py b/v2/recipes/wan_fun_control/__init__.py index fa2819294..4627c9e10 100644 --- a/v2/recipes/wan_fun_control/__init__.py +++ b/v2/recipes/wan_fun_control/__init__.py @@ -1,11 +1,10 @@ -"""Wan2.1-Fun-Control — control-video-conditioned Wan2.1 (V2V) ported into the v2 substrate. +"""Wan2.1-Fun-Control — control-video-conditioned Wan2.1 (V2V). -Self-contained recipe package (the bucket-C pattern): a thin recipe over the SHARED Wan architecture. -The control video is VAE-encoded and concatenated onto the noise latent -(``cat([noise(16), control(16), zero_pad(16)])`` → the 48ch Fun-Control DiT input), carried through the -existing i2v ``cond=`` thread of the shared ``WanDenoiseLoop`` / ``WanDiT`` adapter — so NO new DiT/VAE/T5 -adapter and NO new sampler are needed. ``build_wan_fun_control_card`` + ``build_wan_fun_control_program`` -are registered in ``v2/registry.py`` for ``IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers``. +Self-contained recipe over the shared Wan architecture. The control video is VAE-encoded and concatenated +onto the noise latent (``cat([noise(16), control(16), zero_pad(16)])`` → the 48ch Fun-Control DiT input), +carried through the existing i2v ``cond=`` thread of the shared ``WanDenoiseLoop`` / ``WanDiT`` adapter, so +no new DiT/VAE/T5 adapter and no new sampler are needed. Registered in ``v2/registry.py`` for +``IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers``. """ from __future__ import annotations diff --git a/v2/recipes/wan_fun_control/loop.py b/v2/recipes/wan_fun_control/loop.py index c816aca46..9df781c30 100644 --- a/v2/recipes/wan_fun_control/loop.py +++ b/v2/recipes/wan_fun_control/loop.py @@ -1,22 +1,19 @@ -"""WanFunControlDenoiseLoop — Wan2.1-Fun-Control reuses the canonical Wan denoise loop verbatim. +"""Wan2.1-Fun-Control reuses the canonical Wan denoise loop verbatim. -Wan2.1-Fun-Control is *architecturally* a plain Wan2.1 flow-match DiT: the only delta from T2V is the -DiT input channel count (the control latent is concatenated onto the noise latent, see the program + -the ``WanDiT`` adapter). The denoise control flow — flow-shift schedule, classifier-free guidance, -precision policy, the deterministic flow-match Euler step — is identical to ``v2.recipes.wan21``. - -So there is NO new sampler and NO new loop math here: we re-export the shared ``WanDenoiseLoop``. It -already threads an optional ``i2v_cond`` slot (the i2v ``[mask|cond]`` concat) straight through to the -adapter's ``cond=`` kwarg; the Fun-Control program reuses that exact slot to carry the control-video -``[control_latent | zero_pad]`` concat (same mechanism, different conditioning tensor). Keeping a module -here (rather than importing ``WanDenoiseLoop`` directly in the card) keeps the recipe self-contained and -gives one obvious place to fork the loop should Fun-Control ever need bespoke step math. +Fun-Control is architecturally a plain Wan2.1 flow-match DiT; the only delta from T2V is the DiT input +channel count (the control latent is concatenated onto the noise latent, see the program + the ``WanDiT`` +adapter). The denoise control flow — flow-shift schedule, CFG, precision policy, the flow-match Euler step +— is identical to ``v2.recipes.wan21``, so we just re-export the shared ``WanDenoiseLoop``. It already +threads an optional ``i2v_cond`` slot through to the adapter's ``cond=`` kwarg; the Fun-Control program +reuses that exact slot to carry the control-video ``[control_latent | zero_pad]`` concat. Keeping a module +here (rather than importing ``WanDenoiseLoop`` in the card) keeps the recipe self-contained and gives one +place to fork the loop should Fun-Control ever need bespoke step math. """ from __future__ import annotations -# Re-export the canonical Wan loop + latent-geometry helper. ``WanDenoiseLoop`` is reused unchanged: the +# Re-export the canonical Wan loop + latent-geometry helper. ``WanDenoiseLoop`` is reused unchanged; the # Fun-Control conditioning rides on its existing i2v ``cond=`` thread (program writes the control concat -# into the ``i2v_cond`` slot). ``latent_shape`` gives the Wan2.1 16/8/4 geometry on GPU + the toy shape. +# into the ``i2v_cond`` slot). ``latent_shape`` gives the Wan2.1 16/8/4 geometry. from v2.recipes.wan21.loop import WanDenoiseLoop, latent_shape # Wan2.1 VAE (AutoencoderKLWan): z_dim=16, 4x temporal, 8x spatial — the geometry the control latent and diff --git a/v2/recipes/wan_fun_control/program.py b/v2/recipes/wan_fun_control/program.py index 0bc8c8a6b..2e7bda51b 100644 --- a/v2/recipes/wan_fun_control/program.py +++ b/v2/recipes/wan_fun_control/program.py @@ -1,27 +1,24 @@ """Wan2.1-Fun-Control program: (ref-image CLIP encode) + control-video VAE encode → denoise → decode. -Faithful port of ``fastvideo/pipelines/basic/wan/wan_v2v_pipeline.py`` + its stages: +Port of ``fastvideo/pipelines/basic/wan/wan_v2v_pipeline.py`` + its stages: - * ``RefImageEncodingStage`` (image_encoding.py): a CLIP-vision reference image → ``image_embeds`` passed - to the DiT as ``encoder_hidden_states_image``. Optional — defaults to a zero/None embed when absent. - * ``VideoVAEEncodingStage`` (image_encoding.py): the CONTROL VIDEO is VAE-encoded to a latent - ``control_latent`` (scaled into the DiT's normalized latent space by the ``WanVAE`` adapter, exactly - like the noise latent). - * ``DenoisingStage`` (denoising.py L380-390): for the V2V/Control path the DiT input is - ``cat([noise_latent, control_latent, v2v_zero_pad], dim=channels)`` — i.e. noise(16) + control(16) + - a 16-channel zero pad = the 48-channel Fun-Control DiT input (``WanVideoArchConfig`` in_channels=48). + * ``RefImageEncodingStage``: a CLIP-vision reference image → ``image_embeds`` passed to the DiT as + ``encoder_hidden_states_image``. Optional — defaults to a zero/None embed when absent. + * ``VideoVAEEncodingStage``: the control video is VAE-encoded to ``control_latent`` (scaled into the + DiT's normalized latent space by the ``WanVAE`` adapter, like the noise latent). + * ``DenoisingStage``: the V2V/Control DiT input is ``cat([noise(16), control(16), zero_pad(16)])`` = + the 48-channel Fun-Control input (``WanVideoArchConfig`` in_channels=48). -We carry the ``[control_latent | zero_pad]`` (32ch) concat through the *existing* ``i2v_cond`` slot the -shared ``WanDenoiseLoop`` already threads to the ``WanDiT`` adapter's ``cond=`` kwarg. The adapter then -does ``cat([noise(16ch), cond(32ch)], dim=channels)`` → the 48ch DiT input, byte-for-byte the V2V concat -order ``[latent, video_latent, zero_pad]``. So the Wan recipe/adapter/loop are reused unchanged. +We carry the ``[control_latent | zero_pad]`` (32ch) concat through the existing ``i2v_cond`` slot the +shared ``WanDenoiseLoop`` threads to the ``WanDiT`` adapter's ``cond=`` kwarg. The adapter then does +``cat([noise(16ch), cond(32ch)])`` → the 48ch DiT input, matching the V2V concat order +``[latent, video_latent, zero_pad]``, so the Wan recipe/adapter/loop are reused unchanged. BRINGUP (written-not-run): the control-video request input. v2's ``Request`` has no ``video()`` accessor -yet, so we scan ``request.inputs`` for a ``VideoPart`` ourselves. When NO control video is present this -DEGRADES gracefully: ``i2v_cond`` is ``None`` → ``WanDiT`` runs the plain 16ch forward → pure T2V (and a -ref image, if any, still flows as i2v-style CLIP context). GPU-verify: (a) Wan2.1-Fun-1.3B-Control loads -via the generic Wan loader with in_channels=48, (b) the control video drives the motion, (c) the -``WAN2_1ControlCLIPVisionConfig`` reference-image encoder subfolder/dtype. +yet, so we scan ``request.inputs`` for a ``VideoPart`` ourselves. With no control video this degrades +gracefully (a zero control video is encoded; see ``_control_video_encode``). GPU-verify: (a) +Wan2.1-Fun-1.3B-Control loads via the generic Wan loader with in_channels=48, (b) the control video drives +the motion, (c) the ``WAN2_1ControlCLIPVisionConfig`` reference-image encoder subfolder/dtype. """ from __future__ import annotations @@ -34,10 +31,10 @@ from v2.recipes.common import text_encode_node_fn as _text_encode def _control_video_part(request: Any) -> Any: - """Find the control-video input (a ``VideoPart``) on the request, or ``None`` for t2v degrade. + """Find the control-video input (a ``VideoPart``) on the request, or ``None``. - ``Request`` exposes ``image()`` but not ``video()``, and we must not edit shared files, so we scan - ``request.inputs`` for the VIDEO modality directly (duck-typed on ``.frames``).""" + ``Request`` exposes ``image()`` but not ``video()``, so we scan ``request.inputs`` for the video + modality directly (duck-typed on ``.frames``).""" for part in getattr(request, "inputs", ()) or (): if getattr(part, "modality", None) == "video" or getattr(part, "frames", None) is not None: return part @@ -46,8 +43,8 @@ def _control_video_part(request: Any) -> Any: def _ref_image_encode(instance: Any, slots: dict, request: Any, ctx: Any) -> None: """Optional reference image → CLIP-vision context (``RefImageEncodingStage``). Carried through the - i2v ``i2v_img_embeds`` slot the WanDiT adapter passes as ``encoder_hidden_states_image``. Absent -> - None (the adapter passes no image context; Fun-Control's primary signal is the control video).""" + ``i2v_img_embeds`` slot the WanDiT adapter passes as ``encoder_hidden_states_image``. Absent -> None + (no image context; Fun-Control's primary signal is the control video).""" img = request.image() slots["i2v_img_embeds"] = (instance.component("image_encoder").encode_image(img.pixels) if img is not None and getattr(img, "pixels", None) is not None else None) @@ -60,12 +57,10 @@ def _control_video_encode(instance: Any, slots: dict, request: Any, ctx: Any) -> (``cat([latent, video_latent, zero_pad])``). We pre-build the ``[video_latent | zero_pad]`` (32ch) half so the loop/adapter only has to ``cat([noise(16ch), this(32ch)])`` → the 48ch Fun-Control DiT input. - BRINGUP: the Fun-Control DiT's ``patch_embedding`` is a Conv3d with ``in_channels=48`` — it ALWAYS - requires the 48ch input (no pure-16ch T2V path). The real pipeline reflects this: when no control - video is supplied, ``RefImageEncodingStage`` still sets ``video_latent = zeros`` and - ``VideoVAEEncodingStage`` VAE-encodes it (an asserted, non-optional stage). So we NEVER emit - ``i2v_cond=None``; with no control video we VAE-encode a ZERO control video (faithful to the real - zeros-``video_latent`` path) and concat the zero pad — the 48ch input is always well-formed.""" + BRINGUP: the Fun-Control DiT's ``patch_embedding`` Conv3d has ``in_channels=48`` and always requires + the 48ch input (no pure-16ch T2V path). The real pipeline reflects this: with no control video it still + VAE-encodes a zeros ``video_latent``. So we never emit ``i2v_cond=None`` — with no control video we + VAE-encode a zero control video and concat the zero pad, keeping the 48ch input well-formed.""" nf = int(request.diffusion.num_frames) part = _control_video_part(request) if part is None or getattr(part, "frames", None) is None: diff --git a/v2/registry.py b/v2/registry.py index 591707b5c..d03083ff7 100644 --- a/v2/registry.py +++ b/v2/registry.py @@ -34,12 +34,12 @@ class ModelEntry: # Bucket-C ports — each is a self-contained recipe package (card-declared torch adapter via -# ``ComponentSpec.adapter`` + a new/forked loop, no shared-backend edit). ONE table drives BOTH the +# ``ComponentSpec.adapter`` + a new/forked loop, no shared-backend edit). One table drives both the # explicit HF-id registry (PRIMARY) and the ``select_by_architecture`` fallback, so adding a port is one # row. Rows: (hf_ids, package, card_builder, program_builder, transformer_cls). ``transformer_cls`` is the # diffusers/EntryClass name for the fallback ("" = explicit-id-only; confirm against the real checkpoint's -# transformer/config.json _class_name). Builders are resolved lazily (``import v2.registry`` stays cheap / -# torch-free; the recipe packages are CPU-clean but importing 14 of them eagerly on every resolve is waste). +# transformer/config.json _class_name). Builders are resolved lazily so ``import v2.registry`` stays cheap +# and torch-free (the recipe packages are CPU-clean, but importing them all eagerly on every resolve is waste). _BUCKET_C: tuple[tuple, ...] = ( (("KyleShao/Cosmos-Predict2.5-2B-Diffusers", "nvidia/Cosmos-Predict2.5-14B"), "cosmos25", "build_cosmos25_card", "build_cosmos25_program", "Cosmos25Transformer3DModel"), @@ -217,7 +217,7 @@ def select_by_architecture(sig: dict): """FALLBACK: map an architecture signature -> (build_card, build_program) by class names, so a local path / renamed repo / new distilled variant of a known arch still resolves with no registry entry.""" tr, pipe = sig.get("transformer_cls"), sig.get("pipeline") - for (_hf, pkg, cb, pb, cls) in _BUCKET_C: # the 14 bucket-C ports, keyed by transformer cls + for (_hf, pkg, cb, pb, cls) in _BUCKET_C: # the bucket-C ports, keyed by transformer cls if cls and tr == cls: import importlib mod = importlib.import_module(f"v2.recipes.{pkg}") diff --git a/v2/request/__init__.py b/v2/request/__init__.py index e965642f6..1257c6f8b 100644 --- a/v2/request/__init__.py +++ b/v2/request/__init__.py @@ -1,4 +1,4 @@ -"""Request plane — typed runtime objects crossing the product boundary (design_v3 §12).""" +"""Request plane — typed runtime objects crossing the product boundary.""" from __future__ import annotations from v2.request.artifacts import ( @@ -27,12 +27,34 @@ from v2.request.streams import OmniEvent, Stream, StreamChunk from v2.request.tasks import TaskType __all__ = [ - "TaskType", "Modality", "ModalPart", "TextPart", "ImagePart", "VideoPart", - "AudioPart", "ActionPart", "LatentPart", - "SamplingParams", "DiffusionParams", "OutputSpec", "CaptureMode", - "Request", "Session", "make_request", "new_request_id", - "Artifact", "VideoArtifact", "AudioArtifact", "TextArtifact", - "TensorArtifact", "LatentArtifact", "Output", - "Stream", "OmniEvent", "StreamChunk", - "CancelScope", "CancelKind", "Cancelled", + "TaskType", + "Modality", + "ModalPart", + "TextPart", + "ImagePart", + "VideoPart", + "AudioPart", + "ActionPart", + "LatentPart", + "SamplingParams", + "DiffusionParams", + "OutputSpec", + "CaptureMode", + "Request", + "Session", + "make_request", + "new_request_id", + "Artifact", + "VideoArtifact", + "AudioArtifact", + "TextArtifact", + "TensorArtifact", + "LatentArtifact", + "Output", + "Stream", + "OmniEvent", + "StreamChunk", + "CancelScope", + "CancelKind", + "Cancelled", ] diff --git a/v2/request/artifacts.py b/v2/request/artifacts.py index e59656960..a06305582 100644 --- a/v2/request/artifacts.py +++ b/v2/request/artifacts.py @@ -1,8 +1,8 @@ -"""Artifact — a named, typed output with provenance (design_v3 §12). +"""Artifact — a named, typed output with provenance. -"This kills the ``extra['audio']`` pattern — audio carries its sample rate as a -first-class artifact, not a dict passenger." Every artifact records which program -node produced it (``producer``), so multi-modal fan-out outputs are traceable. +Each modality is its own typed artifact (e.g. audio carries its sample rate as a +field, not via an ``extra`` dict). Every artifact records which program node +produced it (``producer``), so multi-modal fan-out outputs are traceable. """ from __future__ import annotations @@ -14,21 +14,21 @@ from v2.request.modalpart import Modality @dataclass class Artifact: - name: str # output slot name, e.g. "video", "audio", "text" - producer: str = "" # program node id that produced it (provenance) + name: str # output slot name, e.g. "video", "audio", "text" + producer: str = "" # program node id that produced it (provenance) meta: dict[str, object] = field(default_factory=dict) @dataclass class VideoArtifact(Artifact): - frames: TensorLike = None # T x H x W x C + frames: TensorLike = None # T x H x W x C fps: int = 16 @dataclass class AudioArtifact(Artifact): samples: TensorLike = None - sample_rate: int = 44100 # first-class, not an extra-dict passenger + sample_rate: int = 44100 # first-class, not an extra-dict passenger @dataclass @@ -39,7 +39,7 @@ class TextArtifact(Artifact): @dataclass class TensorArtifact(Artifact): - tensor: TensorLike = None # generic (e.g. action stream) + tensor: TensorLike = None # generic (e.g. action stream) @dataclass @@ -54,7 +54,7 @@ class Output: request_id: str artifacts: dict[str, Artifact] = field(default_factory=dict) metrics: dict[str, float] = field(default_factory=dict) - error: str | None = None # set on partial/aborted delivery (failure isolation) + error: str | None = None # set on partial/aborted delivery (failure isolation) def get(self, name: str) -> Artifact | None: return self.artifacts.get(name) diff --git a/v2/request/cancel.py b/v2/request/cancel.py index 03d737a3e..698f0e46b 100644 --- a/v2/request/cancel.py +++ b/v2/request/cancel.py @@ -1,9 +1,8 @@ -"""Structured cancellation (design_v3 §6.4, §12). +"""Structured cancellation. -"Cancellation is common-path, not exceptional" — vibe-directing makes abandoning -in-flight work the *normal* user action. Cancel takes effect at the next loop step -boundary: the driver checks the scope between steps, drops queued work, releases -LoopState + cache handles, and reports ``cancelled``. +Cancellation is a common-path action, not exceptional. Cancel takes effect at the +next loop step boundary: the driver checks the scope between steps, drops queued +work, releases LoopState + cache handles, and reports ``cancelled``. """ from __future__ import annotations diff --git a/v2/request/modalpart.py b/v2/request/modalpart.py index c9d0bc0d5..2bc22ee04 100644 --- a/v2/request/modalpart.py +++ b/v2/request/modalpart.py @@ -1,9 +1,9 @@ -"""ModalPart — typed multimodal request inputs (design_v3 §12; design.md §6.1 G2). +"""ModalPart — typed multimodal request inputs. Requests carry a list of typed modality parts instead of a god-batch with -per-model fields. This kills the ``ForwardBatch`` blackboard / ``extra[...]`` -pattern (design.md P3) at the input boundary; artifacts (outputs) do the same -on the way out (see ``artifacts.py``). +per-model fields, avoiding a ``ForwardBatch`` blackboard / ``extra[...]`` pattern +at the input boundary; artifacts (outputs) do the same on the way out (see +``artifacts.py``). """ from __future__ import annotations @@ -38,7 +38,7 @@ class TextPart(ModalPart): @dataclass(frozen=True) class ImagePart(ModalPart): - pixels: TensorLike = None # HxWxC or CxHxW (backend decides), or a path + pixels: TensorLike = None # HxWxC or CxHxW (backend decides), or a path path: str | None = None def __init__(self, pixels: TensorLike = None, path: str | None = None): @@ -49,7 +49,7 @@ class ImagePart(ModalPart): @dataclass(frozen=True) class VideoPart(ModalPart): - frames: TensorLike = None # T x H x W x C + frames: TensorLike = None # T x H x W x C fps: float = 16.0 path: str | None = None diff --git a/v2/request/params.py b/v2/request/params.py index 92d15fa49..10f761842 100644 --- a/v2/request/params.py +++ b/v2/request/params.py @@ -1,9 +1,7 @@ """Request parameters: AR sampling vs diffusion params + output/capture spec. -design_v3 §12: "AR ``sampling`` vs ``diffusion`` params, an ``OutputSpec`` -(requested modalities + streaming + capture flags), and per-node overrides." -design.md §6.1: ``sampling``/``diffusion`` are *defaults*; ``node_params`` override -per graph node (validated against the PipelineSpec's node-parameter schema). +``sampling``/``diffusion`` are defaults; ``node_params`` override per graph node +(validated against the PipelineSpec's node-parameter schema). """ from __future__ import annotations @@ -29,7 +27,7 @@ class DiffusionParams: guidance_scale: float = 5.0 # Per-modality guidance for joint denoise (LTX-2 A/V, Cosmos3 t2vs). Empty => use scalar. guidance_per_modality: dict[str, float] = field(default_factory=dict) - shift: float = 5.0 # flow-match sigma shift + shift: float = 5.0 # flow-match sigma shift sigmas: tuple[float, ...] | None = None # explicit schedule (distilled few-step) negative_prompt: str = "" height: int = 480 @@ -38,18 +36,18 @@ class DiffusionParams: fps: int = 16 seed: int | None = None # FlowGRPO-style RL rollout: use the SDE sampler (with per-step log-prob capture) instead of the - # deterministic ODE serve sampler (PromptRL/UniRL §6). Default False ⇒ the serve path is unchanged. + # deterministic ODE serve sampler. Default False ⇒ the serve path is unchanged. sde_rollout: bool = False sde_noise_scale: float = 0.7 - # Active adapters for this request (LoRA / ControlNet ids on the base) — the adapter plane (§9.19). + # Active adapters for this request (LoRA / ControlNet ids on the base). # Empty ⇒ the base model, unmodified. Part of the cache key (adapter_versions) so adapted ≠ base. adapters: tuple[str, ...] = () class CaptureMode(str, Enum): - """What the loop captures while running (design_v3 §9.4, §10).""" - NONE = "none" # serve forward: no-grad, no capture - BEHAVIOR = "behavior" # rollout forward: emit BehaviorRecord slices for RL/distill + """What the loop captures while running.""" + NONE = "none" # serve forward: no-grad, no capture + BEHAVIOR = "behavior" # rollout forward: emit BehaviorRecord slices for RL/distill @dataclass(frozen=True) @@ -58,7 +56,7 @@ class OutputSpec: # streaming flags per modality: True | chunk size (ms / frames). {} => no streaming. stream: dict[str, object] = field(default_factory=dict) return_latents: bool = False - return_trajectory: bool = False # full denoise trajectory (parity / debug) + return_trajectory: bool = False # full denoise trajectory (parity / debug) capture: CaptureMode = CaptureMode.NONE @property diff --git a/v2/request/requests.py b/v2/request/requests.py index f92ccc57a..9f72c3115 100644 --- a/v2/request/requests.py +++ b/v2/request/requests.py @@ -1,8 +1,8 @@ -"""Request + Session — the only currency crossing the product boundary (design_v3 §12). +"""Request + Session — the only currency crossing the product boundary. -"A typed ``Request`` carries ``task`` (declared, never inferred), ``inputs`` +A typed ``Request`` carries ``task`` (declared, never inferred), ``inputs`` (ModalParts), AR ``sampling`` vs ``diffusion`` params, an ``OutputSpec``, and -per-node overrides." Products construct Requests; they never reach into the model. +per-node overrides. Products construct Requests; they never reach into the model. """ from __future__ import annotations @@ -26,15 +26,15 @@ def new_request_id() -> str: class Request: request_id: str task: TaskType - model_id: str # which ModelCard to run + model_id: str # which ModelCard to run inputs: tuple[ModalPart, ...] = () sampling: SamplingParams = field(default_factory=SamplingParams) diffusion: DiffusionParams = field(default_factory=DiffusionParams) outputs: OutputSpec = field(default_factory=OutputSpec) - node_params: dict[str, dict] = field(default_factory=dict) # per graph-node overrides + node_params: dict[str, dict] = field(default_factory=dict) # per graph-node overrides priority: int = 0 - # --- convenience accessors (do not infer task from these; §12) --- + # --- convenience accessors (do not infer task from these) --- def prompt(self) -> str: for p in self.inputs: if isinstance(p, TextPart): @@ -51,7 +51,10 @@ class Request: return self.node_params.get(node_id, {}) -def make_request(task: TaskType, model_id: str, prompt: str = "", *, +def make_request(task: TaskType, + model_id: str, + prompt: str = "", + *, inputs: tuple[ModalPart, ...] | None = None, sampling: SamplingParams | None = None, diffusion: DiffusionParams | None = None, @@ -60,10 +63,12 @@ def make_request(task: TaskType, model_id: str, prompt: str = "", *, priority: int = 0, request_id: str | None = None) -> Request: """Ergonomic constructor used by the offline API and tests.""" - parts: tuple[ModalPart, ...] = inputs if inputs is not None else ((TextPart(prompt),) if prompt else ()) + parts: tuple[ModalPart, ...] = inputs if inputs is not None else ((TextPart(prompt), ) if prompt else ()) return Request( request_id=request_id or new_request_id(), - task=task, model_id=model_id, inputs=parts, + task=task, + model_id=model_id, + inputs=parts, sampling=sampling or SamplingParams(), diffusion=diffusion or DiffusionParams(), outputs=outputs or OutputSpec(), @@ -74,7 +79,7 @@ def make_request(task: TaskType, model_id: str, prompt: str = "", *, @dataclass class Session: - """A long-lived interactive context (design_v3 §12, §15e Dreamverse). + """A long-lived interactive context. Holds prompt memory, media streams, a cancel scope, and a handle to cross-request chunk-KV that persists for a game/scene session (self-forcing / @@ -84,7 +89,7 @@ class Session: prompt_memory: list[str] = field(default_factory=list) streams: dict[str, Stream] = field(default_factory=dict) cancel_scope: CancelScope | None = None - kv_handle: object | None = None # persistent chunk-KV CacheHandle (see cache/) + kv_handle: object | None = None # persistent chunk-KV CacheHandle (see cache/) def __post_init__(self): if self.cancel_scope is None: diff --git a/v2/request/streams.py b/v2/request/streams.py index d55689c74..2e857a108 100644 --- a/v2/request/streams.py +++ b/v2/request/streams.py @@ -1,10 +1,10 @@ -"""Streams + typed event taxonomy (design_v3 §12). +"""Streams + typed event taxonomy. -"A ``Stream`` is one ordered event channel for previews, media chunks, progress, -logs, finals." Event taxonomy: ``request.*``, ``session.*``, ``artifact.*``, -``media.{init,chunk,complete}``, ``trace.*``. A ``media.chunk`` must know its +A ``Stream`` is one ordered event channel for previews, media chunks, progress, +logs, and finals. Event taxonomy: ``request.*``, ``session.*``, ``artifact.*``, +``media.{init,chunk,complete}``, ``trace.*``. A ``media.chunk`` carries its stream, byte-range/buffer-ref, codec/container, timestamp range, and -preview-vs-final — invalid combinations are unrepresentable. +preview-vs-final flag, so invalid combinations are unrepresentable. For v2 the Stream is a simple synchronous ordered channel (the core is numpy/sync and CPU-testable); a GPU/server deployment swaps in an async channel @@ -19,7 +19,7 @@ from v2._types import TensorLike @dataclass(frozen=True) class OmniEvent: - type: str # e.g. "request.start", "media.chunk", "artifact.ready", "trace.step" + type: str # e.g. "request.start", "media.chunk", "artifact.ready", "trace.step" request_id: str seq: int = 0 payload: dict[str, object] = field(default_factory=dict) @@ -29,14 +29,14 @@ class OmniEvent: class StreamChunk: """A media.chunk payload — invalid field combinations are unrepresentable.""" stream_id: str - modality: str # "video" | "audio" | "text" + modality: str # "video" | "audio" | "text" seq: int - data: TensorLike = None # in-proc buffer (frames / samples / token ids) - codec: str = "raw" # raw | h264 | opus | ... - ts_start: float = 0.0 # presentation timestamp range (seconds) + data: TensorLike = None # in-proc buffer (frames / samples / token ids) + codec: str = "raw" # raw | h264 | opus | ... + ts_start: float = 0.0 # presentation timestamp range (seconds) ts_end: float = 0.0 - preview: bool = False # preview vs final - final: bool = False # last chunk of the stream + preview: bool = False # preview vs final + final: bool = False # last chunk of the stream class Stream: diff --git a/v2/request/tasks.py b/v2/request/tasks.py index f7de45c7f..b97783815 100644 --- a/v2/request/tasks.py +++ b/v2/request/tasks.py @@ -1,7 +1,7 @@ -"""TaskType — declared, never inferred (design_v3 §12; design.md §6.1 P7). +"""TaskType — declared, never inferred. The task is part of the typed request and the program branches on it. Heuristics -may only *suggest* a default at the product boundary; the runtime never infers it. +may only suggest a default at the product boundary; the runtime never infers it. """ from __future__ import annotations @@ -10,18 +10,18 @@ from enum import Enum class TaskType(str, Enum): # --- wired in phase 1 (Wan2.1 / LTX2) --- - T2V = "t2v" # text -> video - I2V = "i2v" # image -> video - TI2V = "ti2v" # text+image -> video - V2V = "v2v" # video -> video - T2I = "t2i" # text -> image (batchable diffusion) + T2V = "t2v" # text -> video + I2V = "i2v" # image -> video + TI2V = "ti2v" # text+image -> video + V2V = "v2v" # video -> video + T2I = "t2i" # text -> image (batchable diffusion) # --- declared now for omni-readiness; implemented in phase 2 --- - T2A = "t2a" # text -> audio (diffusion) - T2VS = "t2vs" # text -> video+sound (joint denoise) - V2W = "v2w" # video -> world rollout (chunked, causal) - A2W = "a2w" # action -> world rollout - REASON = "reason" # AR text reasoning (Cosmos3 reasoner / omni thinker) + T2A = "t2a" # text -> audio (diffusion) + T2VS = "t2vs" # text -> video+sound (joint denoise) + V2W = "v2w" # video -> world rollout (chunked, causal) + A2W = "a2w" # action -> world rollout + REASON = "reason" # AR text reasoning (Cosmos3 reasoner / omni thinker) @property def is_world_rollout(self) -> bool: @@ -29,4 +29,4 @@ class TaskType(str, Enum): @property def is_autoregressive(self) -> bool: - return self in (TaskType.REASON,) + return self in (TaskType.REASON, ) diff --git a/v2/runtime/__init__.py b/v2/runtime/__init__.py index 141df4259..c01b15227 100644 --- a/v2/runtime/__init__.py +++ b/v2/runtime/__init__.py @@ -1,4 +1,4 @@ -"""Execution plane — engine + scheduler (design_v3 §6).""" +"""Execution plane — engine + scheduler.""" from __future__ import annotations from v2.runtime.async_engine import AsyncEngine, RequestState @@ -16,7 +16,8 @@ from v2.runtime.scheduler import ( WorkUnit, ) -__all__ = ["Engine", "ProgramRunner", "RuntimeLoopContext", "AdmissionController", - "AdmissionInfeasible", "BatchScheduler", "SchedulerMetrics", "StepTicket", "WorkUnit", - "AsyncEngine", "RequestState", "DisaggregatedRunner", "WorldModelSession", - "DeployConfig", "PoolSet", "RolePool", "RolePoolSpec", "wan_t2v_disaggregated"] +__all__ = [ + "Engine", "ProgramRunner", "RuntimeLoopContext", "AdmissionController", "AdmissionInfeasible", "BatchScheduler", + "SchedulerMetrics", "StepTicket", "WorkUnit", "AsyncEngine", "RequestState", "DisaggregatedRunner", + "WorldModelSession", "DeployConfig", "PoolSet", "RolePool", "RolePoolSpec", "wan_t2v_disaggregated" +] diff --git a/v2/runtime/async_engine.py b/v2/runtime/async_engine.py index 06da97360..464a216b2 100644 --- a/v2/runtime/async_engine.py +++ b/v2/runtime/async_engine.py @@ -1,18 +1,19 @@ -"""AsyncEngine — the serving runtime (design_v3 §6; plan M3). +"""AsyncEngine — the serving runtime. -A real request queue, lifecycle state machine, live streaming, and step-level concurrency over the -same step scheduler. Each request runs as an asyncio task that ticks its runner one step at a time -and yields the event loop (`await asyncio.sleep(0)`), so concurrent requests interleave at step +A request queue, lifecycle state machine, live streaming, and step-level concurrency over the same +step scheduler. Each request runs as an asyncio task that ticks its runner one step at a time and +yields the event loop (`await asyncio.sleep(0)`), so concurrent requests interleave at step granularity — and a per-request stall backs off instead of busy-spinning. Drives BOTH inline ``ProgramRunner`` and ``DisaggregatedRunner`` (capacity-aware pool dispatch) through one driver. -Cancellation is common-path (design_v3 §6.4): ``cancel(request_id)`` trips the runner's CancelScope, -which raises at the next step boundary and delivers partial artifacts + a structured ``cancelled``. +Cancellation is common-path: ``cancel(request_id)`` trips the runner's CancelScope, which raises at +the next step boundary and delivers partial artifacts + a structured ``cancelled``. """ from __future__ import annotations import asyncio -from typing import Any, AsyncIterator +from typing import Any +from collections.abc import AsyncIterator from v2.request.artifacts import Output from v2.request.cancel import Cancelled @@ -23,7 +24,7 @@ from v2.runtime.engine import Engine _SENTINEL = object() -def _safe_put(q: "asyncio.Queue", item: Any) -> None: +def _safe_put(q: asyncio.Queue, item: Any) -> None: """Non-blocking put on an unbounded queue — safe to call from finally/cancellation paths.""" try: q.put_nowait(item) @@ -40,15 +41,15 @@ class RequestState: class AsyncEngine: - def __init__(self, engine: Engine | None = None, *, max_stall_ticks: int = 200_000, - max_history: int = 512): + + def __init__(self, engine: Engine | None = None, *, max_stall_ticks: int = 200_000, max_history: int = 512): self.engine = engine if engine is not None else Engine() - self._disagg: dict[str, tuple] = {} # model_id -> (PoolSet, Program) + self._disagg: dict[str, tuple] = {} # model_id -> (PoolSet, Program) self._events: dict[str, asyncio.Queue] = {} self._states: dict[str, str] = {} self._results: dict[str, Output] = {} self._runners: dict[str, Any] = {} - self._order: list[str] = [] # submission order, for bounded eviction + self._order: list[str] = [] # submission order, for bounded eviction self.max_stall_ticks = max_stall_ticks self.max_history = max_history @@ -83,8 +84,11 @@ class AsyncEngine: def _make_runner(self, request: Any) -> Any: if request.model_id in self._disagg: pools, program = self._disagg[request.model_id] - return DisaggregatedRunner(pools, program, request, - observers=self.engine.observers, interceptors=self.engine.interceptors) + return DisaggregatedRunner(pools, + program, + request, + observers=self.engine.observers, + interceptors=self.engine.interceptors) return self.engine._make_runner(request) # --- lifecycle / control ------------------------------------------------ # @@ -110,8 +114,7 @@ class AsyncEngine: return sum(1 for s in self._states.values() if s == RequestState.WAITING) def serves(self, model_id: str) -> bool: - return (model_id in self.engine._registry or model_id in self._disagg - or model_id in self.engine._workflows) + return (model_id in self.engine._registry or model_id in self._disagg or model_id in self.engine._workflows) # --- the driver: one asyncio task per request --------------------------- # async def _run(self, request: Any) -> None: @@ -126,7 +129,7 @@ class AsyncEngine: self._states[rid] = RequestState.RUNNING await evq.put(OmniEvent("request.start", rid, payload={"task": request.task.value})) while not runner.done: - runner.cancel_scope.check() # common-path cancellation (step boundary) + runner.cancel_scope.check() # common-path cancellation (step boundary) p0 = runner._progress runner.tick() # surface any newly emitted media chunks as live events @@ -138,10 +141,10 @@ class AsyncEngine: stall += 1 if stall > self.max_stall_ticks: raise RuntimeError(f"request {rid} stalled (state={getattr(runner,'state','?')})") - await asyncio.sleep(0.0005) # backpressure backoff (capacity wait) + await asyncio.sleep(0.0005) # backpressure backoff (capacity wait) else: stall = 0 - await asyncio.sleep(0) # cooperative yield → step interleaving + await asyncio.sleep(0) # cooperative yield → step interleaving out = runner.output() self._results[rid] = out for name in out.artifacts: @@ -158,7 +161,7 @@ class AsyncEngine: except Exception: pass _safe_put(evq, OmniEvent("request.cancelled", rid)) - except asyncio.CancelledError: # driver task cancelled (e.g. client gone) + except asyncio.CancelledError: # driver task cancelled (e.g. client gone) self._states[rid] = RequestState.CANCELLED raise except Exception as e: # noqa: BLE001 — request-fatal isolation: one request's error @@ -198,7 +201,7 @@ class AsyncEngine: task.cancel() try: await task - except (asyncio.CancelledError, Exception): # noqa: BLE001 + except (asyncio.CancelledError, Exception): # noqa: BLE001 pass async def generate(self, request: Any) -> Output: diff --git a/v2/runtime/context.py b/v2/runtime/context.py index d825ba980..22e0d473e 100644 --- a/v2/runtime/context.py +++ b/v2/runtime/context.py @@ -1,4 +1,4 @@ -"""RuntimeLoopContext — the concrete inversion point (design_v3 §5.1, §11). +"""RuntimeLoopContext — the concrete inversion point. This is the single seam: ``execute(plan)`` runs the kernel thunk built by the model's ``next()``, wrapped by interceptors (compute-altering) and observers (read-only). The model @@ -16,19 +16,19 @@ from v2.request.streams import OmniEvent, Stream, StreamChunk class RuntimeLoopContext: - def __init__(self, instance: Any, *, observers, interceptors, slots: dict[str, Any], - stream: Stream, cancel_scope, profile: ExecutionProfile, metrics: dict, - request_id: str): + + def __init__(self, instance: Any, *, observers, interceptors, slots: dict[str, Any], stream: Stream, cancel_scope, + profile: ExecutionProfile, metrics: dict, request_id: str): self.instance = instance self.observers = observers self.interceptors = interceptors - self.slots = slots # the program's typed named edges + self.slots = slots # the program's typed named edges self.stream = stream self.cancel_scope = cancel_scope self.profile = profile self.metrics = metrics self.request_id = request_id - self.state = None # bound per step by the LoopRunner + self.state = None # bound per step by the LoopRunner def bind_state(self, state) -> None: self.state = state @@ -68,9 +68,9 @@ class RuntimeLoopContext: def _run_step(self, plan: WorkPlan, override, eager): """Execute the step body, under piecewise CUDA-graph capture/replay when the loop declares - graph_capture="breakable_cudagraph" (design_v3 §6.2; Path A). Otherwise plain eager — the - unchanged path. Replay still runs the current thunk (CPU models the lifecycle, not the - speedup); see runtime/cudagraph.py.""" + graph_capture="breakable_cudagraph". Otherwise plain eager — the unchanged path. Replay + still runs the current thunk (CPU models the lifecycle, not the speedup); see + runtime/cudagraph.py.""" if plan.run is None: return {} loop = self.instance.card.loops.get(plan.loop_id) if self.instance.card else None @@ -82,13 +82,13 @@ class RuntimeLoopContext: return self.instance.graphs.dispatch(plan, self.instance, override, eager) def emit(self, chunk: StreamChunk) -> None: - self.stream.emit(OmniEvent(type="media.chunk", request_id=self.request_id, - seq=chunk.seq, payload={"chunk": chunk})) + self.stream.emit( + OmniEvent(type="media.chunk", request_id=self.request_id, seq=chunk.seq, payload={"chunk": chunk})) self.metrics["stream_chunks"] = self.metrics.get("stream_chunks", 0) + 1 def check_cancel(self) -> None: if self.cancel_scope is not None: - self.cancel_scope.check() # raises request.Cancelled at the step boundary + self.cancel_scope.check() # raises request.Cancelled at the step boundary def observe(self, event: str, **kw) -> None: self.observers.emit(event, **kw) diff --git a/v2/runtime/cudagraph.py b/v2/runtime/cudagraph.py index cdc13e828..62de6ec0a 100644 --- a/v2/runtime/cudagraph.py +++ b/v2/runtime/cudagraph.py @@ -1,7 +1,7 @@ -"""Piecewise CUDA-graph capture/replay at the driven-loop step boundary (Path A; design_v3 §6.2). +"""Piecewise CUDA-graph capture/replay at the driven-loop step boundary. -The chosen optimization path (Path A) is hand-fused kernels behind the registry + piecewise CUDA -graphs captured at the step boundary — NOT a compiler. This module is the capture/replay LIFECYCLE: +The chosen optimization path is hand-fused kernels behind the registry + piecewise CUDA graphs +captured at the step boundary — NOT a compiler. This module is the capture/replay LIFECYCLE: when to capture a step into a graph, when to replay it, when to eager-break (a data-dependent step a static graph cannot hold), and when to invalidate (shape or resident-weights change). @@ -54,7 +54,8 @@ Honest scope — what this mini still does NOT model (bites a real GPU, not the from __future__ import annotations from dataclasses import dataclass -from typing import Any, Callable +from typing import Any +from collections.abc import Callable import numpy as np @@ -69,14 +70,14 @@ class StaticWorkspace: @staticmethod def _mk(v: Any): - return None if v is None else np.array(v) # scalars become 0-d arrays; arrays are copied + return None if v is None else np.array(v) # scalars become 0-d arrays; arrays are copied @classmethod - def alloc(cls, inputs: dict[str, Any]) -> "StaticWorkspace": + def alloc(cls, inputs: dict[str, Any]) -> StaticWorkspace: b = {k: cls._mk(v) for k, v in inputs.items()} x = inputs.get("x") if x is not None: - b["out"] = np.array(x) # static output buffer, latent-shaped + b["out"] = np.array(x) # static output buffer, latent-shaped return cls(b) def bind(self, inputs: dict[str, Any]) -> None: @@ -86,7 +87,7 @@ class StaticWorkspace: if v is not None: raise ValueError(f"workspace[{k!r}] was None at capture but has a value on replay") continue - np.copyto(cur, v) # in place; raises if shapes/dtypes don't fit + np.copyto(cur, v) # in place; raises if shapes/dtypes don't fit def __getitem__(self, k: str): return self._b[k] @@ -96,7 +97,7 @@ class StaticWorkspace: class CapturedGraph: """A keyed record standing in for a captured ``torch.cuda.CUDAGraph`` + its static buffers.""" key: tuple - workspace: Any # StaticWorkspace (the reused fixed buffers), or None + workspace: Any # StaticWorkspace (the reused fixed buffers), or None workspace_bytes: int replays: int = 0 @@ -106,7 +107,7 @@ class GraphCapturer: captured graph is valid only for that instance's resident weights + shapes.""" def __init__(self, enabled: bool = True) -> None: - self.enabled = enabled # off ⇒ always eager (the no-cudagraph baseline) + self.enabled = enabled # off ⇒ always eager (the no-cudagraph baseline) self._graphs: dict[tuple, CapturedGraph] = {} self.stats = {"captures": 0, "replays": 0, "eager_breaks": 0} @@ -116,14 +117,12 @@ class GraphCapturer: backend (device, arch) · loop · batch-shape signature · resident-weight versions · the op-structure discriminators (CFG branch set / expert) the plan declares.""" loop = instance.card.loops.get(plan.loop_id) if instance.card else None - resident = tuple(sorted((cid, instance.version_of(cid)) - for cid in (loop.shared_weight_components if loop else ()))) + resident = tuple( + sorted((cid, instance.version_of(cid)) for cid in (loop.shared_weight_components if loop else ()))) p = instance.platform - return (p.device, p.arch, plan.loop_id, plan.shape_sig.batch_key, - resident, tuple(plan.graph_key)) + return (p.device, p.arch, plan.loop_id, plan.shape_sig.batch_key, resident, tuple(plan.graph_key)) - def dispatch(self, plan: Any, instance: Any, override: Any, - eager: Callable[[Any], Any]) -> Any: + def dispatch(self, plan: Any, instance: Any, override: Any, eager: Callable[[Any], Any]) -> Any: """Run one step under the capture lifecycle. ``eager(override)`` runs the eager step thunk (override / stochastic / no-capture paths). A capturable step runs its ``graph_fn`` against the keyed ``StaticWorkspace``; replay rebinds the current inputs into the SAME buffers and @@ -150,10 +149,9 @@ class GraphCapturer: # HIT → replay. peak-workspace proxy guard (declared static workspace must match)... if g.workspace_bytes != ws_bytes: - raise RuntimeError( - f"cudagraph key collision: key {key} captured a {g.workspace_bytes}-byte workspace " - f"but this step needs {ws_bytes} — the capture key is insufficient and would corrupt " - f"a real graph. This is a keying bug, not a runtime condition.") + raise RuntimeError(f"cudagraph key collision: key {key} captured a {g.workspace_bytes}-byte workspace " + f"but this step needs {ws_bytes} — the capture key is insufficient and would corrupt " + f"a real graph. This is a keying bug, not a runtime condition.") # ...and the real backstop: rebinding into the fixed buffers raises if shapes don't fit. g.workspace.bind(plan.graph_inputs) g.replays += 1 diff --git a/v2/runtime/disaggregated.py b/v2/runtime/disaggregated.py index ca7df81c0..5013fbb6a 100644 --- a/v2/runtime/disaggregated.py +++ b/v2/runtime/disaggregated.py @@ -1,4 +1,4 @@ -"""Disaggregated program execution (design_v3 §13; design.md §6.3.4; sglang multimodal_gen). +"""Disaggregated program execution (cf. sglang multimodal_gen). Runs a Program across role pools: each node executes on its assigned pool's resident instance, and cross-pool edges move slots through the producing pool's connector (with credit flow-control + @@ -25,6 +25,7 @@ from v2.runtime.engine import _to_artifact class DisaggregatedRunner: + def __init__(self, pools, program: Program, request: Any, *, observers, interceptors): self.pools = pools self.program = program @@ -32,7 +33,7 @@ class DisaggregatedRunner: self.observers = observers self.interceptors = interceptors self.slots: dict[str, Any] = {} - self.slot_pool: dict[str, str] = {} # slot name -> pool_id currently holding it + self.slot_pool: dict[str, str] = {} # slot name -> pool_id currently holding it self.metrics: dict[str, float] = {} self.transfers = 0 self.node_idx = 0 @@ -53,10 +54,15 @@ class DisaggregatedRunner: def _ctx_for(self, pool) -> RuntimeLoopContext: if pool.pool_id not in self._ctxs: - self._ctxs[pool.pool_id] = RuntimeLoopContext( - pool.instance, observers=self.observers, interceptors=self.interceptors, - slots=self.slots, stream=self.stream, cancel_scope=self.cancel_scope, - profile=ExecutionProfile.SERVE, metrics=self.metrics, request_id=self.request.request_id) + self._ctxs[pool.pool_id] = RuntimeLoopContext(pool.instance, + observers=self.observers, + interceptors=self.interceptors, + slots=self.slots, + stream=self.stream, + cancel_scope=self.cancel_scope, + profile=ExecutionProfile.SERVE, + metrics=self.metrics, + request_id=self.request.request_id) return self._ctxs[pool.pool_id] def _leave_current(self) -> None: @@ -70,19 +76,22 @@ class DisaggregatedRunner: holder_id = self.slot_pool.get(r) if holder_id and holder_id != pool.pool_id and r in self.slots: src = self.pools.by_id[holder_id] - got = src.connector.acquire_credit() # credit-based flow control + got = src.connector.acquire_credit() # credit-based flow control try: - manifest = TransferManifest(keys=(r,), producer_id=holder_id, consumer_id=pool.pool_id, - src_location=src.connector.name, dst_location=pool.connector.name, + manifest = TransferManifest(keys=(r, ), + producer_id=holder_id, + consumer_id=pool.pool_id, + src_location=src.connector.name, + dst_location=pool.connector.name, nbytes=payload_nbytes(self.slots.get(r))) src.connector.put(r, self.slots.get(r), manifest) # producer publishes (chunk_ready) - if src.connector.chunk_ready(r): # readiness signal - self.slots[r] = src.connector.take(r) # consumer fetches - self.slot_pool[r] = pool.pool_id # re-home ONLY on a successful fetch + if src.connector.chunk_ready(r): # readiness signal + self.slots[r] = src.connector.take(r) # consumer fetches + self.slot_pool[r] = pool.pool_id # re-home ONLY on a successful fetch self.transfers += 1 self.metrics["transfers"] = self.metrics.get("transfers", 0) + 1 finally: - if got: # never leak a credit on an error path + if got: # never leak a credit on an error path src.connector.release_credit() def _commit_node_writes(self, node, pool) -> None: @@ -92,7 +101,7 @@ class DisaggregatedRunner: def tick(self) -> bool: if self.done: return True - self.cancel_scope.check() # common-path cancellation (online + offline, all nodes) + self.cancel_scope.check() # common-path cancellation (online + offline, all nodes) nodes = self.program.active_nodes(self.request) while True: if self.node_idx >= len(nodes): @@ -103,10 +112,10 @@ class DisaggregatedRunner: node = nodes[self.node_idx] pool = self.pools.pool_for(node.node_id) - if pool is not self._cur_pool: # switching pools: capacity-aware dispatch + if pool is not self._cur_pool: # switching pools: capacity-aware dispatch if not pool.can_admit(): self.state = f"waiting:{pool.role}" - return False # backpressure — pool at capacity + return False # backpressure — pool at capacity self._leave_current() pool.enter() self._cur_pool = pool @@ -123,8 +132,8 @@ class DisaggregatedRunner: if isinstance(node, ModelLoopNode): if self.loop_runner is None: - self.loop_runner = LoopRunner(pool.instance.loop(node.loop_id), - self._ctx_for(pool), self.request, pool.instance) + self.loop_runner = LoopRunner(pool.instance.loop(node.loop_id), self._ctx_for(pool), self.request, + pool.instance) plan = self.loop_runner.peek() if plan is None: self._finish_loop(node, pool) @@ -133,7 +142,7 @@ class DisaggregatedRunner: self.metrics["stepped_units"] = self.metrics.get("stepped_units", 0) + 1 if self.loop_runner.done: self._finish_loop(node, pool) - return False # yield one loop step (for interleaving) + return False # yield one loop step (for interleaving) self.node_idx += 1 @@ -153,7 +162,7 @@ class DisaggregatedRunner: self.tick() if not self.done and self._progress == p0: stuck += 1 - if stuck > 3: # single-runner: a stall means real OOM/capacity-0 + if stuck > 3: # single-runner: a stall means real OOM/capacity-0 raise RuntimeError(f"disaggregated request {self.request.request_id} stalled " f"(state={self.state}); a pool cannot admit it") else: @@ -161,11 +170,15 @@ class DisaggregatedRunner: def close(self) -> None: """Release any pool this request is occupying (idempotent). Called on terminal exit so a - cancel/error mid-loop never leaks pool capacity (the §6.4 failure-isolation guarantee).""" + cancel/error mid-loop never leaks pool capacity (failure-isolation guarantee).""" self._leave_current() def output(self) -> Output: - artifacts = {name: _to_artifact(name, self.slots.get(slot), producer=slot) - for name, slot in self.program.output_artifacts.items()} - return Output(request_id=self.request.request_id, artifacts=artifacts, - metrics=dict(self.metrics, transfers=float(self.transfers)), error=self.error) + artifacts = { + name: _to_artifact(name, self.slots.get(slot), producer=slot) + for name, slot in self.program.output_artifacts.items() + } + return Output(request_id=self.request.request_id, + artifacts=artifacts, + metrics=dict(self.metrics, transfers=float(self.transfers)), + error=self.error) diff --git a/v2/runtime/engine.py b/v2/runtime/engine.py index b184d8d28..f030db686 100644 --- a/v2/runtime/engine.py +++ b/v2/runtime/engine.py @@ -1,13 +1,13 @@ -"""The engine — drives Programs over resident ModelInstances (design_v3 §6). +"""The engine — drives Programs over resident ModelInstances. Two execution modes share one ProgramRunner: * ``run_serial`` — each request's program runs to completion before the next. * ``run_interleaved`` — round-robin one *unit* per request per tick (a unit = one loop step, or one one-shot node), so the denoise steps of concurrent requests interleave. This is the - mode the §9.3 interleave parity gate exercises: serial and interleaved must be bit-identical. + mode the interleave parity gate exercises: serial and interleaved must be bit-identical. The engine owns admission (reservation before stepping), the observer bus, and the interceptor -chain (deploy scope). It never imports ``training`` (design_v3 §10 dependency rule). +chain (deploy scope). It never imports ``training`` (dependency rule). """ from __future__ import annotations @@ -56,7 +56,7 @@ def _to_artifact(name: str, value: Any, producer: str = "") -> Any: class ProgramRunner: """Drives one Program for one request. Advances one *unit* per ``tick`` for interleaving.""" - def __init__(self, engine: "Engine", program: Program, instance: Any, request: Any): + def __init__(self, engine: Engine, program: Program, instance: Any, request: Any): self.engine = engine self.program = program self.instance = instance @@ -67,10 +67,15 @@ class ProgramRunner: self.stream = Stream(request.request_id) profile = ExecutionProfile.ROLLOUT if getattr(request.outputs, "capture", None) and \ request.outputs.capture.value == "behavior" else engine.profile - self.ctx = RuntimeLoopContext( - instance, observers=engine.observers, interceptors=engine.interceptors, - slots=self.slots, stream=self.stream, cancel_scope=self.cancel_scope, - profile=profile, metrics=self.metrics, request_id=request.request_id) + self.ctx = RuntimeLoopContext(instance, + observers=engine.observers, + interceptors=engine.interceptors, + slots=self.slots, + stream=self.stream, + cancel_scope=self.cancel_scope, + profile=profile, + metrics=self.metrics, + request_id=request.request_id) self.node_idx = 0 self.loop_runner: LoopRunner | None = None self._resident_res: Any = None @@ -84,7 +89,7 @@ class ProgramRunner: def tick(self) -> bool: if self.done: return True - self.cancel_scope.check() # common-path cancellation (offline + all-component programs) + self.cancel_scope.check() # common-path cancellation (offline + all-component programs) nodes = self.program.nodes while True: if self.node_idx >= len(nodes): @@ -99,7 +104,7 @@ class ProgramRunner: if node.fn is not None: node.fn(self.instance, self.slots, self.request, self.ctx) self.node_idx += 1 - continue # one-shots run within this tick + continue # one-shots run within this tick if isinstance(node, ModelLoopNode): if self.loop_runner is None: @@ -116,13 +121,12 @@ class ProgramRunner: need = plan0.resources.resident_bytes if need > 0: if not self.engine.admission.feasible_resident(need): - raise AdmissionInfeasible( - f"{node.loop_id}: resident {need}B exceeds pool capacity " - f"{self.engine.admission.memory.total_bytes}B") + raise AdmissionInfeasible(f"{node.loop_id}: resident {need}B exceeds pool capacity " + f"{self.engine.admission.memory.total_bytes}B") res = self.engine.admission.reserve_resident(node.loop_id, need) if res is None: self.engine.metrics.deferred += 1 - return False # defer: cannot fit resident yet (jointly-OOM guard) + return False # defer: cannot fit resident yet (jointly-OOM guard) self._resident_res = res else: self._resident_res = True @@ -137,13 +141,13 @@ class ProgramRunner: f"{plan.resources.compute_seconds:.4g}s — exceeds pool/budget capacity") step_res = self.engine.admission.admit_step(plan) if step_res is None: - return False # deferred this tick (will retry; refundable) + return False # deferred this tick (will retry; refundable) self.loop_runner.step() self.engine.admission.release(step_res) self.engine.metrics.stepped_units += 1 if self.loop_runner.done: self._finish_loop(node) - return False # yielded one loop step → return for interleaving + return False # yielded one loop step → return for interleaving # unknown node kind self.node_idx += 1 @@ -152,7 +156,7 @@ class ProgramRunner: self.slots[node.output_slot] = self.loop_runner.result.outputs self._merge_metrics(self.loop_runner.result.metrics) if self._resident_res not in (None, True): - self.engine.admission.memory.release(self._resident_res) # resident is a plain Reservation + self.engine.admission.memory.release(self._resident_res) # resident is a plain Reservation self._resident_res = None self.loop_runner = None self.node_idx += 1 @@ -170,10 +174,9 @@ class ProgramRunner: self.tick() if not self.done and self._progress == p0: stuck += 1 - if stuck > 3: # fail fast, not a 10M busy-spin - raise AdmissionInfeasible( - f"program {self.program.program_id} made no progress for several ticks " - "(unsatisfiable reservation or deadlock)") + if stuck > 3: # fail fast, not a 10M busy-spin + raise AdmissionInfeasible(f"program {self.program.program_id} made no progress for several ticks " + "(unsatisfiable reservation or deadlock)") else: stuck = 0 @@ -181,21 +184,23 @@ class ProgramRunner: artifacts: dict[str, Any] = {} for name, slot in self.program.output_artifacts.items(): artifacts[name] = _to_artifact(name, self.slots.get(slot), producer=slot) - return Output(request_id=self.request.request_id, artifacts=artifacts, - metrics=dict(self.metrics), error=self.error) + return Output(request_id=self.request.request_id, + artifacts=artifacts, + metrics=dict(self.metrics), + error=self.error) @dataclass class Engine: - """A single-pool engine (design_v3 §6). Pools are single-node; fleet scale is multiple pools.""" + """A single-pool engine. Pools are single-node; fleet scale is multiple pools.""" profile: ExecutionProfile = ExecutionProfile.SERVE observers: ObserverBus = field(default_factory=ObserverBus) interceptors: InterceptorChain = field(default_factory=InterceptorChain) admission: AdmissionController = field(default_factory=AdmissionController) metrics: SchedulerMetrics = field(default_factory=SchedulerMetrics) _registry: dict[str, tuple[Any, Program]] = field(default_factory=dict) - _workflows: dict[str, Any] = field(default_factory=dict) # workflow_id -> validated Workflow - _wf_running: set = field(default_factory=set) # workflow ids in-flight (nested cycle guard) + _workflows: dict[str, Any] = field(default_factory=dict) # workflow_id -> validated Workflow + _wf_running: set = field(default_factory=set) # workflow ids in-flight (nested cycle guard) def register(self, model_id: str, instance: Any, program: Program) -> None: self._registry[model_id] = (instance, program) @@ -223,9 +228,9 @@ class Engine: # --- offline single-request (the VideoGenerator path) -------------------- # def run(self, request: Any) -> Output: - if request.model_id in self._workflows: # a workflow servable → orchestrate stages + if request.model_id in self._workflows: # a workflow servable → orchestrate stages wf = self._workflows[request.model_id] - if request.model_id in self._wf_running: # nested workflows: cycle guard (§9.13) + if request.model_id in self._wf_running: # nested workflows: cycle guard raise ValueError(f"workflow cycle: {request.model_id!r} is already running " f"(a workflow cannot transitively invoke itself)") seed = request.diffusion.seed if request.diffusion.seed is not None else 0 @@ -250,10 +255,11 @@ class Engine: def run_interleaved(self, requests: list[Any]) -> dict[str, Output]: runners = [self._make_runner(req) for req in requests] while not all(r.done for r in runners): - # accounting: how many ready WorkPlans would co-batch this round (BatchScheduler, §6.3) - units = [WorkUnit(r.request.request_id, r.loop_runner.peek()) - for r in runners if not r.done and r.loop_runner is not None - and r.loop_runner.peek() is not None] + # accounting: how many ready WorkPlans would co-batch this round (BatchScheduler) + units = [ + WorkUnit(r.request.request_id, r.loop_runner.peek()) for r in runners + if not r.done and r.loop_runner is not None and r.loop_runner.peek() is not None + ] if units: self.metrics.batches += len(BatchScheduler.group(units)) progressed = False @@ -263,8 +269,7 @@ class Engine: r.tick() if r.done or r._progress != p0: progressed = True - if not progressed: # nobody could be admitted → real deadlock, fail fast - raise AdmissionInfeasible( - "interleaved scheduling deadlocked: no request could be admitted this round " - "(jointly infeasible — would require step-boundary preemption, §6.4)") + if not progressed: # nobody could be admitted → real deadlock, fail fast + raise AdmissionInfeasible("interleaved scheduling deadlocked: no request could be admitted this round " + "(jointly infeasible — would require step-boundary preemption, §6.4)") return {r.request.request_id: r.output() for r in runners} diff --git a/v2/runtime/pools.py b/v2/runtime/pools.py index f3415e9f7..663c1e55a 100644 --- a/v2/runtime/pools.py +++ b/v2/runtime/pools.py @@ -1,10 +1,10 @@ -"""Role / stage pools (design_v3 §13, design.md §6.3.4; port of sglang multimodal_gen disagg). +"""Role / stage pools (port of sglang multimodal_gen disagg). A ``RolePool`` is a single-node worker pool for one role (encoder / denoiser / decoder / …) holding its OWN resident ModelInstance (a component subset), its OWN cache manager, an outbound connector, and a capacity (max concurrent requests → capacity-aware dispatch). ``DeployConfig`` maps program -nodes onto pools. Pools are single-node; multi-node scale is *multiple pools* fronted by the fleet -(§14) — never a cross-node mesh inside one pool. +nodes onto pools. Pools are single-node; multi-node scale is *multiple pools* fronted by the fleet — +never a cross-node mesh inside one pool. """ from __future__ import annotations @@ -18,26 +18,27 @@ from v2.transport import Connector, make_connector @dataclass class RolePoolSpec: - role: str # "encoder" | "denoiser" | "decoder" | ... + role: str # "encoder" | "denoiser" | "decoder" | ... pool_id: str - components: tuple[str, ...] = () # components this pool is responsible for - capacity: int = 2 # max concurrent requests occupying this pool + components: tuple[str, ...] = () # components this pool is responsible for + capacity: int = 2 # max concurrent requests occupying this pool replicas: int = 1 - connector: str = "in_proc" # outbound transport backend - slo_class: str = "throughput" # "latency" | "throughput" | "cost" + connector: str = "in_proc" # outbound transport backend + slo_class: str = "throughput" # "latency" | "throughput" | "cost" @dataclass class DeployConfig: - """Maps a Program's nodes onto role pools (the vllm-omni deploy-YAML analog, §6.3.4).""" + """Maps a Program's nodes onto role pools (the vllm-omni deploy-YAML analog).""" pools: list[RolePoolSpec] = field(default_factory=list) - placement: dict[str, str] = field(default_factory=dict) # node_id -> pool_id + placement: dict[str, str] = field(default_factory=dict) # node_id -> pool_id def pool_id_for(self, node_id: str, default: str | None = None) -> str: return self.placement.get(node_id, default or (self.pools[0].pool_id if self.pools else "")) class RolePool: + def __init__(self, spec: RolePoolSpec, card: Any): self.spec = spec self.role = spec.role @@ -66,9 +67,14 @@ class RolePool: return self.in_flight / max(1, self.capacity) def stats(self) -> dict[str, Any]: - return {"role": self.role, "in_flight": self.in_flight, "capacity": self.capacity, - "served": self.served, "transfers_out": self.connector.transfers, - "credits_available": self.connector.credits.available} + return { + "role": self.role, + "in_flight": self.in_flight, + "capacity": self.capacity, + "served": self.served, + "transfers_out": self.connector.transfers, + "credits_available": self.connector.credits.available + } class PoolSet: @@ -85,7 +91,7 @@ class PoolSet: return self.by_id[self.deploy.pool_id_for(node_id)] def warmup(self) -> dict[str, int]: - """Instantiate each pool's components ahead of first request (design.md §6.3.5 warmup).""" + """Instantiate each pool's components ahead of first request.""" built: dict[str, int] = {} for pool in self.by_id.values(): for cid in pool.components: @@ -102,11 +108,17 @@ def wan_t2v_disaggregated() -> DeployConfig: """The canonical N:M:K split for Wan T2V (encoder → denoiser → decoder).""" return DeployConfig( pools=[ - RolePoolSpec(role="encoder", pool_id="enc", components=("text_encoder",), capacity=4, - slo_class="latency"), - RolePoolSpec(role="denoiser", pool_id="den", components=("transformer",), capacity=1, - slo_class="throughput"), # jumbo video step stays batch-of-1 - RolePoolSpec(role="decoder", pool_id="dec", components=("vae",), capacity=2), + RolePoolSpec(role="encoder", pool_id="enc", components=("text_encoder", ), capacity=4, slo_class="latency"), + RolePoolSpec(role="denoiser", + pool_id="den", + components=("transformer", ), + capacity=1, + slo_class="throughput"), # jumbo video step stays batch-of-1 + RolePoolSpec(role="decoder", pool_id="dec", components=("vae", ), capacity=2), ], - placement={"text_encode": "enc", "denoise": "den", "vae_decode": "dec"}, + placement={ + "text_encode": "enc", + "denoise": "den", + "vae_decode": "dec" + }, ) diff --git a/v2/runtime/scheduler.py b/v2/runtime/scheduler.py index 8fd0d1df4..94a4d1894 100644 --- a/v2/runtime/scheduler.py +++ b/v2/runtime/scheduler.py @@ -1,11 +1,11 @@ -"""The scheduler, in layers (design_v3 §6.3) — each testable on a fake pool, no GPU. +"""The scheduler, in layers — each testable on a fake pool, no GPU. Mini-fastvideo implements the load-bearing layers: - * ``AdmissionController`` — the reservation gate (§6.2): do not admit a WorkUnit unless its + * ``AdmissionController`` — the reservation gate: do not admit a WorkUnit unless its compute budget AND memory (resident + worst-case peak) can be reserved. Reservations are *refundable* (compute is a concurrency gate, not a lifetime cap), and infeasible reservations (need > pool capacity) fail fast with ``AdmissionInfeasible`` instead of spinning. - * ``BatchScheduler`` — groups compatible WorkPlans by ``shape_sig.batch_key`` (§6.3). NOTE: in this + * ``BatchScheduler`` — groups compatible WorkPlans by ``shape_sig.batch_key``. NOTE: in this single-process mini, grouping drives *accounting* (how many units would co-batch); kernels still execute per-plan. Real cross-request batched execution is a GPU-path concern (documented, not faked). """ @@ -25,7 +25,7 @@ class AdmissionInfeasible(RuntimeError): @dataclass class WorkUnit: - """The smallest schedulable action (design_v3 §6.1).""" + """The smallest schedulable action.""" request_id: str plan: WorkPlan priority: int = 0 @@ -49,7 +49,7 @@ class SchedulerMetrics: class AdmissionController: - """Reservation before admission (design_v3 §6.2). Compute budget + memory both gate; both refund.""" + """Reservation before admission. Compute budget + memory both gate; both refund.""" def __init__(self, memory: MemoryManager | None = None, compute_budget_seconds: float = float("inf")): self.memory = memory or MemoryManager() @@ -96,11 +96,11 @@ class AdmissionController: return if ticket.memory_res is not None: self.memory.release(ticket.memory_res) - self.compute_spent = max(0.0, self.compute_spent - ticket.compute_seconds) # refund + self.compute_spent = max(0.0, self.compute_spent - ticket.compute_seconds) # refund class BatchScheduler: - """Group compatible WorkPlans across requests (design_v3 §6.3 layer 3).""" + """Group compatible WorkPlans across requests.""" @staticmethod def group(units: list[WorkUnit]) -> list[list[WorkUnit]]: diff --git a/v2/runtime/session.py b/v2/runtime/session.py index 60d3aaaf5..bb55969e6 100644 --- a/v2/runtime/session.py +++ b/v2/runtime/session.py @@ -1,4 +1,4 @@ -"""WorldModelSession — a long-lived interactive world-model rollout (design_v3 §16, §12, §15e). +"""WorldModelSession — a long-lived interactive world-model rollout. The realtime/interactive plane, made concrete. A one-shot ``engine.run`` generates a clip and forgets; a *session* continues the world: each ``.act(action)`` resumes the causal ``chunk_rollout`` loop from @@ -13,7 +13,8 @@ cancelled rollout leaves the session resumable from the last good frame. """ from __future__ import annotations -from typing import Any, Callable +from typing import Any +from collections.abc import Callable from v2._enums import ExecutionProfile from v2.loop.driver import LoopRunner @@ -25,16 +26,22 @@ from v2.runtime.context import RuntimeLoopContext class WorldModelSession: - def __init__(self, engine: Any, model_id: str, *, session_id: str = "world", - window: int = 2, loop_id: str = "chunk_rollout"): + + def __init__(self, + engine: Any, + model_id: str, + *, + session_id: str = "world", + window: int = 2, + loop_id: str = "chunk_rollout"): self.engine = engine self.model_id = model_id - self.window = window # how many prior chunks of world state to carry forward + self.window = window # how many prior chunks of world state to carry forward self.loop_id = loop_id self.session = Session(session_id=session_id) self.stream = Stream(session_id) self.session.streams["video"] = self.stream - self.world_context: list = [] # the persistent world chunks (the §16 cross-request KV) + self.world_context: list = [] # persistent world chunks (the cross-request KV) self.session.kv_handle = self.world_context self.step = 0 @@ -44,33 +51,42 @@ class WorldModelSession: raise KeyError(f"no model {self.model_id!r} registered on engine") return self.engine._registry[self.model_id][0] - def act(self, action: str, *, seed: int | None = None, - step_hook: Callable[["WorldModelSession", int], None] | None = None) -> Any: + def act(self, + action: str, + *, + seed: int | None = None, + step_hook: Callable[[WorldModelSession, int], None] | None = None) -> Any: """Advance the world by one action, continuing from persistent context; stream frames. Raises ``request.Cancelled`` if cancelled at a step boundary — and because the world state is only committed after the loop completes, a cancelled act leaves the session resumable.""" inst = self.instance - req = make_request(TaskType.V2W, self.model_id, action, + req = make_request(TaskType.V2W, + self.model_id, + action, diffusion=DiffusionParams(seed=self.step if seed is None else seed)) slots: dict[str, Any] = {} text_encode_node_fn(inst, slots, req, None) - slots["world_context"] = list(self.world_context) # thread the persistent state in - ctx = RuntimeLoopContext( - inst, observers=self.engine.observers, interceptors=self.engine.interceptors, - slots=slots, stream=self.stream, cancel_scope=self.session.cancel_scope, - profile=ExecutionProfile.SERVE, metrics={}, - request_id=f"{self.session.session_id}.{self.step}") + slots["world_context"] = list(self.world_context) # thread the persistent state in + ctx = RuntimeLoopContext(inst, + observers=self.engine.observers, + interceptors=self.engine.interceptors, + slots=slots, + stream=self.stream, + cancel_scope=self.session.cancel_scope, + profile=ExecutionProfile.SERVE, + metrics={}, + request_id=f"{self.session.session_id}.{self.step}") runner = LoopRunner(inst.loop(self.loop_id), ctx, req, inst) n = 0 while not runner.done: if step_hook is not None: - step_hook(self, n) # test seam: may call self.cancel() at a boundary - runner.step() # peek() raises request.Cancelled if cancelled + step_hook(self, n) # test seam: may call self.cancel() at a boundary + runner.step() # peek() raises request.Cancelled if cancelled n += 1 result = runner.result new_chunks = list(result.outputs.get("chunks") or []) - self.world_context = (self.world_context + new_chunks)[-self.window:] # slide window; commit + self.world_context = (self.world_context + new_chunks)[-self.window:] # slide window; commit self.session.kv_handle = self.world_context self.step += 1 return result diff --git a/v2/serving/http.py b/v2/serving/http.py index d6e688bd6..5c3adef69 100644 --- a/v2/serving/http.py +++ b/v2/serving/http.py @@ -1,16 +1,16 @@ """A tiny framework-free async HTTP/1.1 server (stdlib asyncio only). -No FastAPI/uvicorn dependency — the design says the production OpenAI *fleet* frontend is Dynamo's -job; this is the per-engine worker server (our own version). Supports JSON responses and streamed -(SSE) responses. Deliberately minimal: one request per connection (Connection: close), enough to be -real and curl-able, not a hardened web server. +No FastAPI/uvicorn dependency — the production OpenAI *fleet* frontend is Dynamo's job; this is the +per-engine worker server. Supports JSON and streamed (SSE) responses. Deliberately minimal: one +request per connection (Connection: close), enough to be real and curl-able, not a hardened server. """ from __future__ import annotations import asyncio import json from dataclasses import dataclass, field -from typing import Any, AsyncIterator, Awaitable, Callable +from typing import Any +from collections.abc import AsyncIterator, Awaitable, Callable @dataclass @@ -30,20 +30,27 @@ class Response: status: int = 200 body: bytes = b"" content_type: str = "application/json" - stream: AsyncIterator | None = None # if set, body is ignored and chunks are streamed + stream: AsyncIterator | None = None # if set, body is ignored and chunks are streamed @classmethod - def json(cls, obj: Any, status: int = 200) -> "Response": + def json(cls, obj: Any, status: int = 200) -> Response: return cls(status=status, body=json.dumps(obj).encode(), content_type="application/json") @classmethod - def sse(cls, gen: AsyncIterator) -> "Response": + def sse(cls, gen: AsyncIterator) -> Response: return cls(status=200, content_type="text/event-stream", stream=gen) -_STATUS = {200: "OK", 201: "Created", 202: "Accepted", 400: "Bad Request", - 404: "Not Found", 408: "Request Timeout", 413: "Payload Too Large", - 500: "Internal Server Error"} +_STATUS = { + 200: "OK", + 201: "Created", + 202: "Accepted", + 400: "Bad Request", + 404: "Not Found", + 408: "Request Timeout", + 413: "Payload Too Large", + 500: "Internal Server Error" +} Handler = Callable[[Request], Awaitable[Response]] @@ -102,12 +109,17 @@ async def _write(writer: asyncio.StreamWriter, resp: Response) -> None: class HttpServer: - def __init__(self, dispatch: Handler, *, read_timeout: float = 30.0, - max_body_bytes: int = 64 << 20, header_limit: int = 1 << 20): + + def __init__(self, + dispatch: Handler, + *, + read_timeout: float = 30.0, + max_body_bytes: int = 64 << 20, + header_limit: int = 1 << 20): self.dispatch = dispatch - self.read_timeout = read_timeout # slowloris guard - self.max_body_bytes = max_body_bytes # body-size cap → 413 - self.header_limit = header_limit # StreamReader buffer for the header block + self.read_timeout = read_timeout # slowloris guard + self.max_body_bytes = max_body_bytes # body-size cap → 413 + self.header_limit = header_limit # StreamReader buffer for the header block self._server: asyncio.AbstractServer | None = None async def _conn(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: @@ -127,9 +139,9 @@ class HttpServer: try: await _write(writer, resp) except Exception: - if resp.stream is not None: # client gone mid-SSE → close the generator + if resp.stream is not None: # client gone mid-SSE → close the generator try: - await resp.stream.aclose() # → GeneratorExit → submit() cancels the driver + await resp.stream.aclose() # → GeneratorExit → submit() cancels the driver except Exception: pass try: diff --git a/v2/serving/protocol.py b/v2/serving/protocol.py index 9cbf0d7c1..74e9f2aed 100644 --- a/v2/serving/protocol.py +++ b/v2/serving/protocol.py @@ -1,12 +1,11 @@ -"""OpenAI-shaped request schemas (design_v3 §12; the vllm-omni protocol, framework-free). +"""OpenAI-shaped request schemas (vllm-omni protocol, framework-free). Plain dataclasses parsed from JSON (no pydantic dep). Each translates to a typed ``OmniRequest`` at -the server boundary — "the request is the only currency crossing the product boundary" (§12). +the server boundary — the request is the only currency crossing the product boundary. """ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any from v2.request import DiffusionParams, OutputSpec, SamplingParams, TaskType, make_request @@ -45,17 +44,25 @@ class ImageGenerationRequest: seed: int = 0 @classmethod - def from_json(cls, d: dict) -> "ImageGenerationRequest": - return cls(prompt=d.get("prompt", ""), model=d.get("model", ""), n=int(d.get("n", 1)), - size=d.get("size", "832x480"), num_inference_steps=int(d.get("num_inference_steps", 4)), - guidance_scale=float(d.get("guidance_scale", 5.0)), seed=int(d.get("seed", 0))) + def from_json(cls, d: dict) -> ImageGenerationRequest: + return cls(prompt=d.get("prompt", ""), + model=d.get("model", ""), + n=int(d.get("n", 1)), + size=d.get("size", "832x480"), + num_inference_steps=int(d.get("num_inference_steps", 4)), + guidance_scale=float(d.get("guidance_scale", 5.0)), + seed=int(d.get("seed", 0))) def to_omni(self): w, h = _wh(self.size) - return make_request(_task_for(self.model, "image"), self.model, self.prompt, + return make_request(_task_for(self.model, "image"), + self.model, + self.prompt, diffusion=DiffusionParams(num_steps=self.num_inference_steps, guidance_scale=self.guidance_scale, - height=h, width=w, seed=self.seed)) + height=h, + width=w, + seed=self.seed)) @dataclass @@ -71,20 +78,28 @@ class VideoGenerationRequest: stream: dict = field(default_factory=dict) @classmethod - def from_json(cls, d: dict) -> "VideoGenerationRequest": - return cls(prompt=d.get("prompt", ""), model=d.get("model", ""), - seconds=float(d.get("seconds", 2.0)), size=d.get("size", "832x480"), + def from_json(cls, d: dict) -> VideoGenerationRequest: + return cls(prompt=d.get("prompt", ""), + model=d.get("model", ""), + seconds=float(d.get("seconds", 2.0)), + size=d.get("size", "832x480"), num_frames=int(d.get("num_frames", 81)), num_inference_steps=int(d.get("num_inference_steps", 4)), - guidance_scale=float(d.get("guidance_scale", 5.0)), seed=int(d.get("seed", 0)), + guidance_scale=float(d.get("guidance_scale", 5.0)), + seed=int(d.get("seed", 0)), stream={"video": True} if d.get("stream") else {}) def to_omni(self): w, h = _wh(self.size) - return make_request(TaskType.T2V, self.model, self.prompt, + return make_request(TaskType.T2V, + self.model, + self.prompt, diffusion=DiffusionParams(num_steps=self.num_inference_steps, - guidance_scale=self.guidance_scale, height=h, width=w, - num_frames=self.num_frames, seed=self.seed), + guidance_scale=self.guidance_scale, + height=h, + width=w, + num_frames=self.num_frames, + seed=self.seed), outputs=OutputSpec(modalities=frozenset({"video"}), stream=self.stream)) @@ -98,10 +113,13 @@ class ChatCompletionRequest: seed: int = 0 @classmethod - def from_json(cls, d: dict) -> "ChatCompletionRequest": - return cls(model=d.get("model", ""), messages=d.get("messages", []), - stream=bool(d.get("stream", False)), max_tokens=int(d.get("max_tokens", 6)), - num_inference_steps=int(d.get("num_inference_steps", 4)), seed=int(d.get("seed", 0))) + def from_json(cls, d: dict) -> ChatCompletionRequest: + return cls(model=d.get("model", ""), + messages=d.get("messages", []), + stream=bool(d.get("stream", False)), + max_tokens=int(d.get("max_tokens", 6)), + num_inference_steps=int(d.get("num_inference_steps", 4)), + seed=int(d.get("seed", 0))) def prompt(self) -> str: for m in reversed(self.messages): @@ -111,7 +129,9 @@ class ChatCompletionRequest: return "" def to_omni(self): - return make_request(_task_for(self.model, "chat"), self.model, self.prompt(), + return make_request(_task_for(self.model, "chat"), + self.model, + self.prompt(), sampling=SamplingParams(max_tokens=self.max_tokens, seed=self.seed), diffusion=DiffusionParams(num_steps=self.num_inference_steps, seed=self.seed), outputs=OutputSpec(stream={"video": True} if self.stream else {})) diff --git a/v2/serving/server.py b/v2/serving/server.py index c9fc5727d..012d2dc09 100644 --- a/v2/serving/server.py +++ b/v2/serving/server.py @@ -1,12 +1,12 @@ -"""OmniOpenAIServer — our own OpenAI-compatible worker server (design.md §6.3.5; vllm-omni pattern). +"""OmniOpenAIServer — OpenAI-compatible worker server (vllm-omni pattern). Fronts ONE AsyncEngine (the per-pool worker surface). Routes OpenAI-shaped endpoints to typed OmniRequests, streams chat via SSE, and serves video as an async job (POST returns a job id; GET polls) plus a /sync variant — mirroring vllm-omni's `/v1/videos` + `/v1/videos/sync`. The fleet (deploy/fleet.py) or Dynamo routes across multiple such servers/engines above this layer. -Unlike vllm-omni's request-scheduled opaque DIFFUSION stage, the engine underneath is step-scheduled -— the endpoint is a thin shim; the runtime-visible loop scheduler is the substance. +Unlike vllm-omni's request-scheduled opaque DIFFUSION stage, the engine underneath is step-scheduled; +the endpoint is a thin shim over the loop scheduler. """ from __future__ import annotations @@ -27,11 +27,19 @@ def _json_default(o: Any) -> Any: """Make SSE event payloads JSON-safe (media.chunk carries a StreamChunk with an array).""" if isinstance(o, StreamChunk): shape = list(np.asarray(o.data).shape) if o.data is not None else None - return {"stream_id": o.stream_id, "modality": o.modality, "seq": o.seq, - "preview": o.preview, "final": o.final, "shape": shape} + return { + "stream_id": o.stream_id, + "modality": o.modality, + "seq": o.seq, + "preview": o.preview, + "final": o.final, + "shape": shape + } if isinstance(o, np.ndarray): return {"shape": list(o.shape), "dtype": str(o.dtype)} return str(o) + + from v2.serving.protocol import ChatCompletionRequest, ImageGenerationRequest, VideoGenerationRequest _job_ctr = itertools.count(1) @@ -55,20 +63,25 @@ def _ser_artifacts(out: Output) -> dict: res[name] = {"type": "empty"} else: a = np.asarray(p) - res[name] = {"type": "tensor", "shape": list(a.shape), "dtype": str(a.dtype), - "b64_json": base64.b64encode(np.ascontiguousarray(a).tobytes()).decode()} + res[name] = { + "type": "tensor", + "shape": list(a.shape), + "dtype": str(a.dtype), + "b64_json": base64.b64encode(np.ascontiguousarray(a).tobytes()).decode() + } return res class OmniOpenAIServer: + def __init__(self, engine: Any, *, engine_id: str = "v2", max_jobs: int = 512): - self.engine = engine # an AsyncEngine + self.engine = engine # an AsyncEngine self.engine_id = engine_id - self.jobs: dict[str, dict] = {} # video job store + self.jobs: dict[str, dict] = {} # video job store self.http = HttpServer(self.dispatch) self.served = 0 self.max_jobs = max_jobs - self._tasks: set = set() # tracked video-job tasks (no fire-and-forget leak) + self._tasks: set = set() # tracked video-job tasks (no fire-and-forget leak) def _evict_jobs(self) -> None: terminal = ("completed", "failed") @@ -101,20 +114,30 @@ class OmniOpenAIServer: # --- metadata endpoints ------------------------------------------------- # def _health(self) -> dict: - return {"status": "healthy", "engine_id": self.engine_id, - "in_flight": self.engine.in_flight, "queue_depth": self.engine.queue_depth} + return { + "status": "healthy", + "engine_id": self.engine_id, + "in_flight": self.engine.in_flight, + "queue_depth": self.engine.queue_depth + } def _metrics(self) -> dict: sm = self.engine.engine.metrics adm = self.engine.engine.admission.metrics - return {"served": self.served, "in_flight": self.engine.in_flight, - "stepped_units": sm.stepped_units, "batches": sm.batches, - "admitted": adm.admitted, "deferred": adm.deferred, - "gpu_seconds": round(adm.gpu_seconds, 6), "by_kind": dict(adm.by_kind)} + return { + "served": self.served, + "in_flight": self.engine.in_flight, + "stepped_units": sm.stepped_units, + "batches": sm.batches, + "admitted": adm.admitted, + "deferred": adm.deferred, + "gpu_seconds": round(adm.gpu_seconds, 6), + "by_kind": dict(adm.by_kind) + } def _models(self) -> dict: - ids = (list(self.engine.engine._registry) + list(self.engine._disagg) - + list(self.engine.engine._workflows)) # workflows are servables too + ids = (list(self.engine.engine._registry) + list(self.engine._disagg) + list(self.engine.engine._workflows) + ) # workflows are servables too return {"object": "list", "data": [{"id": m, "object": "model"} for m in ids]} # --- images ------------------------------------------------------------- # @@ -124,9 +147,16 @@ class OmniOpenAIServer: self.served += 1 arts = _ser_artifacts(out) primary = arts.get("image") or arts.get("video") or arts.get("latents") - return Response.json({"created": int(time.time()), "model": ig.model, - "data": [primary] * max(1, ig.n), "omni": arts, - "metrics": {k: round(v, 6) for k, v in out.metrics.items()}}) + return Response.json({ + "created": int(time.time()), + "model": ig.model, + "data": [primary] * max(1, ig.n), + "omni": arts, + "metrics": { + k: round(v, 6) + for k, v in out.metrics.items() + } + }) # --- video: async job + poll + sync ------------------------------------- # async def _video_create(self, req: Request) -> Response: @@ -135,7 +165,7 @@ class OmniOpenAIServer: self.jobs[job_id] = {"id": job_id, "status": "queued", "model": vg.model} t = asyncio.create_task(self._run_video_job(job_id, vg)) self._tasks.add(t) - t.add_done_callback(self._tasks.discard) # tracked, not fire-and-forget + t.add_done_callback(self._tasks.discard) # tracked, not fire-and-forget self._evict_jobs() return Response.json({"id": job_id, "status": "queued", "model": vg.model}, status=202) @@ -144,8 +174,12 @@ class OmniOpenAIServer: try: out = await self.engine.generate(vg.to_omni()) self.served += 1 - self.jobs[job_id].update(status="completed", result=_ser_artifacts(out), - metrics={k: round(v, 6) for k, v in out.metrics.items()}) + self.jobs[job_id].update(status="completed", + result=_ser_artifacts(out), + metrics={ + k: round(v, 6) + for k, v in out.metrics.items() + }) except Exception as e: # noqa: BLE001 self.jobs[job_id].update(status="failed", error=str(e)) @@ -159,8 +193,14 @@ class OmniOpenAIServer: vg = VideoGenerationRequest.from_json(req.json()) out = await self.engine.generate(vg.to_omni()) self.served += 1 - return Response.json({"model": vg.model, "data": _ser_artifacts(out), - "metrics": {k: round(v, 6) for k, v in out.metrics.items()}}) + return Response.json({ + "model": vg.model, + "data": _ser_artifacts(out), + "metrics": { + k: round(v, 6) + for k, v in out.metrics.items() + } + }) # --- chat: SSE stream or JSON ------------------------------------------- # async def _chat(self, req: Request) -> Response: @@ -173,26 +213,54 @@ class OmniOpenAIServer: arts = _ser_artifacts(out) content = arts.get("text", {}).get("text", "") if "text" in arts else "" return Response.json({ - "id": f"chatcmpl-{next(_chat_ctr)}", "object": "chat.completion", - "created": int(time.time()), "model": cc.model, - "choices": [{"index": 0, "finish_reason": "stop", - "message": {"role": "assistant", "content": content}}], - "omni": arts, "metrics": {k: round(v, 6) for k, v in out.metrics.items()}}) + "id": + f"chatcmpl-{next(_chat_ctr)}", + "object": + "chat.completion", + "created": + int(time.time()), + "model": + cc.model, + "choices": [{ + "index": 0, + "finish_reason": "stop", + "message": { + "role": "assistant", + "content": content + } + }], + "omni": + arts, + "metrics": { + k: round(v, 6) + for k, v in out.metrics.items() + } + }) async def _chat_sse(self, omni: Any, model: str): import json cid = f"chatcmpl-{next(_chat_ctr)}" - async for ev in self.engine.submit(omni): # live OmniEvents → SSE chunks - chunk = {"id": cid, "object": "chat.completion.chunk", "model": model, - "event": ev.type, "data": ev.payload} + async for ev in self.engine.submit(omni): # live OmniEvents → SSE chunks + chunk = {"id": cid, "object": "chat.completion.chunk", "model": model, "event": ev.type, "data": ev.payload} yield f"data: {json.dumps(chunk, default=_json_default)}\n\n" self.served += 1 out = self.engine.result(omni.request_id) arts = _ser_artifacts(out) if out is not None else {} content = arts.get("text", {}).get("text", "") if "text" in arts else "" - final = {"id": cid, "object": "chat.completion.chunk", "model": model, - "choices": [{"index": 0, "delta": {"role": "assistant", "content": content}, - "finish_reason": "stop"}], "omni": arts} + final = { + "id": cid, + "object": "chat.completion.chunk", + "model": model, + "choices": [{ + "index": 0, + "delta": { + "role": "assistant", + "content": content + }, + "finish_reason": "stop" + }], + "omni": arts + } yield f"data: {json.dumps(final, default=_json_default)}\n\n" yield "data: [DONE]\n\n" @@ -201,7 +269,7 @@ class OmniOpenAIServer: return await self.http.start(host, port) async def close(self) -> None: - for t in list(self._tasks): # cancel + drain in-flight video jobs on shutdown + for t in list(self._tasks): # cancel + drain in-flight video jobs on shutdown t.cancel() for t in list(self._tasks): try: diff --git a/v2/tests/test_adapters.py b/v2/tests/test_adapters.py index ed9fedffa..f2ded5401 100644 --- a/v2/tests/test_adapters.py +++ b/v2/tests/test_adapters.py @@ -1,10 +1,9 @@ -"""The adapter plane — per-request LoRA / ControlNet over one base (design_v3 §7.1, §9.19). +"""The adapter plane — per-request LoRA / ControlNet over one base. Many adapters served over ONE resident base; a request picks which to apply (`DiffusionParams.adapters`). -The `adapter_versions` cache-key field existed but the serving capability didn't. These assert: -per-request selection changes output, multi-LoRA composes, ControlNet conditions on a control image, -mixed-adapter requests interleave without smearing, hot-swap changes generation, and the cache key -partitions by the active adapter stack (adapted ≠ base, and version A ≠ version B). +These assert: per-request selection changes output, multi-LoRA composes, ControlNet conditions on a +control image, mixed-adapter requests interleave without smearing, hot-swap changes generation, and the +cache key partitions by the active adapter stack (adapted ≠ base, and version A ≠ version B). """ from __future__ import annotations diff --git a/v2/tests/test_adaptive_compute.py b/v2/tests/test_adaptive_compute.py index cb0233b91..bcadfb6b0 100644 --- a/v2/tests/test_adaptive_compute.py +++ b/v2/tests/test_adaptive_compute.py @@ -1,4 +1,4 @@ -"""Content-adaptive control flow — the loop owns it (design_v3 §2.2, §9.12). +"""Content-adaptive control flow — the loop owns it. Because the model owns control flow (the driven loop), content-adaptive compute is natural: the `CacheDiTDenoiseLoop`'s `next()` reuses the cached velocity when consecutive predictions barely change diff --git a/v2/tests/test_admission.py b/v2/tests/test_admission.py index 7bdbb6d67..abd79946a 100644 --- a/v2/tests/test_admission.py +++ b/v2/tests/test_admission.py @@ -1,6 +1,6 @@ -"""Admission = reservation before stepping (design_v3 §6.2). +"""Admission = reservation before stepping. -> Two requests that fit individually but jointly OOM are rejected at admission, not at step 37. +Two requests that fit individually but jointly OOM are rejected at admission, not at step 37. And crucially: admission-forced serialization does NOT change outputs (still bit-identical). """ from __future__ import annotations diff --git a/v2/tests/test_cache.py b/v2/tests/test_cache.py index a2c936744..9ff6a18ba 100644 --- a/v2/tests/test_cache.py +++ b/v2/tests/test_cache.py @@ -1,4 +1,4 @@ -"""Per-class cache pools (design_v3 §7.2): feature partition/invalidation, slab-KV window/training.""" +"""Per-class cache pools: feature partition/invalidation, slab-KV window/training.""" from __future__ import annotations import numpy as np diff --git a/v2/tests/test_contracts.py b/v2/tests/test_contracts.py index 936a204a7..d03eb41b1 100644 --- a/v2/tests/test_contracts.py +++ b/v2/tests/test_contracts.py @@ -1,4 +1,4 @@ -"""Card / recipe / parallelism / cache-key contracts (design_v3 §4, §7.1, §8, §9.2).""" +"""Card / recipe / parallelism / cache-key contracts.""" from __future__ import annotations import numpy as np diff --git a/v2/tests/test_cudagraph.py b/v2/tests/test_cudagraph.py index 24dd41dc2..e72ae6324 100644 --- a/v2/tests/test_cudagraph.py +++ b/v2/tests/test_cudagraph.py @@ -1,4 +1,4 @@ -"""Piecewise CUDA-graph capture/replay at the step boundary (Path A; design_v3 §6.2). +"""Piecewise CUDA-graph capture/replay at the step boundary. These tests exercise the capture LIFECYCLE that is correctness-critical regardless of device: the capture key (capture-once / replay-many for compatible steps), eager-break for data-dependent steps diff --git a/v2/tests/test_extensions.py b/v2/tests/test_extensions.py index e5516244c..15c5883bf 100644 --- a/v2/tests/test_extensions.py +++ b/v2/tests/test_extensions.py @@ -1,6 +1,6 @@ -"""Observers + interceptors (design_v3 §11). The headline: per-request interceptor state is -interleave-safe; a module-global version provably is NOT — which is *why* §11 scopes state to -``LoopState.plugin_state`` per request and per CFG branch. +"""Observers + interceptors. Per-request interceptor state is interleave-safe; a module-global +version provably is NOT — which is why state is scoped to ``LoopState.plugin_state`` per request +and per CFG branch. """ from __future__ import annotations diff --git a/v2/tests/test_flywheel.py b/v2/tests/test_flywheel.py index 7e7a8cccb..ea647afa3 100644 --- a/v2/tests/test_flywheel.py +++ b/v2/tests/test_flywheel.py @@ -1,9 +1,8 @@ -"""The (recipe, runtime) flywheel — RL → distill → faster card (design_v3 §16; §9.18). +"""The (recipe, runtime) flywheel — RL → distill → faster card. -design_v3's central product claim, made concrete: RL the base (DiffusionNFT), then distill from the -RL'd model (a DMD2 student whose teacher IS the RL'd policy), producing a faster card whose recipe -records the provenance chain base → rl → distilled. Both stages drive the same denoise loop the engine -serves — the flywheel is real, not aspirational. +RL the base (DiffusionNFT), then distill from the RL'd model (a DMD2 student whose teacher IS the +RL'd policy), producing a faster card whose recipe records the provenance chain base → rl → +distilled. Both stages drive the same denoise loop the engine serves. """ from __future__ import annotations diff --git a/v2/tests/test_inference.py b/v2/tests/test_inference.py index 3dfe2b0ea..38cee1b30 100644 --- a/v2/tests/test_inference.py +++ b/v2/tests/test_inference.py @@ -1,4 +1,4 @@ -"""Inference paths: Wan2.1 T2V, LTX2.3 two-stage distilled, Wan-causal streaming (design_v3 §15).""" +"""Inference paths: Wan2.1 T2V, LTX2.3 two-stage distilled, Wan-causal streaming.""" from __future__ import annotations import numpy as np diff --git a/v2/tests/test_interleave_gate.py b/v2/tests/test_interleave_gate.py index e66b29c15..ae591331f 100644 --- a/v2/tests/test_interleave_gate.py +++ b/v2/tests/test_interleave_gate.py @@ -1,4 +1,4 @@ -"""The batch-of-N interleave parity gate (design_v3 §9.3) — the bet loop-inversion lives/dies on. +"""The batch-of-N interleave parity gate — the bet loop-inversion lives or dies on. Two+ concurrent requests, interleaved at step granularity, MUST be bit-identical to the same requests run serially. A batch-of-1 gate is structurally blind to cross-request state smearing; diff --git a/v2/tests/test_joint_multi_rl.py b/v2/tests/test_joint_multi_rl.py index c229a634f..a8d7dd232 100644 --- a/v2/tests/test_joint_multi_rl.py +++ b/v2/tests/test_joint_multi_rl.py @@ -1,4 +1,4 @@ -"""N-way joint RL over arbitrary experts (design_v3 §10) — generalizing the UniRL stress test. +"""N-way joint RL over arbitrary experts — generalizing the UniRL stress test. Answers "can we do joint RL over MORE than two experts, or does it need a rewrite?": no rewrite. The substrate is already N-ready (a card holds N components/loops — Qwen-Omni shipped three; WeightSyncPlan @@ -97,7 +97,7 @@ def test_each_expert_versions_independently_text_encoder_cache_preserved(): m, rids, _ = _method(2) _train(m, iters=10) inst = m.student - assert inst.version_of("text_encoder") == "v0" # frozen encoder cache survives (§7.1) + assert inst.version_of("text_encoder") == "v0" # frozen encoder cache survives expert_vers = {inst.version_of(r) for r in rids} | {inst.version_of("transformer")} assert "v0" not in expert_vers # every trained expert bumped assert len(expert_vers) == 3 # and each independently diff --git a/v2/tests/test_ltx2_av.py b/v2/tests/test_ltx2_av.py index e6dfe2158..499a8984d 100644 --- a/v2/tests/test_ltx2_av.py +++ b/v2/tests/test_ltx2_av.py @@ -1,4 +1,4 @@ -"""LTX-2 joint audio+video denoise — T2VS with per-modality guidance (design_v3 §15b; §9.11). +"""LTX-2 joint audio+video denoise — T2VS with per-modality guidance. LTX-2 declared an `audio_vae` (`required_for={"t2vs"}`) but never used it. These tests light it up: a single two-stage denoise carries a **synchronized audio latent** alongside the video (audio conditioned diff --git a/v2/tests/test_nested_workflow.py b/v2/tests/test_nested_workflow.py index bbfa29c8a..9a4623162 100644 --- a/v2/tests/test_nested_workflow.py +++ b/v2/tests/test_nested_workflow.py @@ -1,4 +1,4 @@ -"""Nested workflows — recursive cross-model composition (design_v3 §13; §9.13). +"""Nested workflows — recursive cross-model composition. A workflow id is a servable, so a `WorkflowStage` can invoke a *workflow* exactly as it invokes a model. These tests build a workflow whose first stage IS the T2I→I2V workflow (then extends its video), diff --git a/v2/tests/test_omni.py b/v2/tests/test_omni.py index 7ceb91151..441c5b4a0 100644 --- a/v2/tests/test_omni.py +++ b/v2/tests/test_omni.py @@ -1,8 +1,8 @@ """Phase-2 omni: one resident MoT instance runs AR + diffusion loops on shared weights. -This is the §16 claim — true omni/MoT serving: one resident model, many loop types, scheduled at -step granularity, in one request — made native (not a DAG doubling weights, not a monolith bypassing -the abstraction). Cosmos3 = reason → joint denoise; BAGEL/lance = generate_text → generate_image. +True omni/MoT serving: one resident model, many loop types, scheduled at step granularity, in one +request — made native (not a DAG doubling weights, not a monolith bypassing the abstraction). +Cosmos3 = reason → joint denoise; BAGEL/lance = generate_text → generate_image. """ from __future__ import annotations diff --git a/v2/tests/test_parity.py b/v2/tests/test_parity.py index 6ce608ade..f5143fa47 100644 --- a/v2/tests/test_parity.py +++ b/v2/tests/test_parity.py @@ -1,4 +1,4 @@ -"""ParityAligner + consistency ladder (design_v3 §9.1, §9.2).""" +"""ParityAligner + consistency ladder.""" from __future__ import annotations import numpy as np diff --git a/v2/tests/test_platform.py b/v2/tests/test_platform.py index 254e606a2..61d121941 100644 --- a/v2/tests/test_platform.py +++ b/v2/tests/test_platform.py @@ -1,10 +1,9 @@ -"""Multi-backend dispatch substrate: device/arch resolution, fallback, and the parity oracle (§17). +"""Multi-backend dispatch substrate: device/arch resolution, fallback, and the parity oracle. -These tests prove the two tuple-keyed registries (COMPONENTS, KERNELS) + the Platform resolve -implementations generically across devices, with the numpy reference as the terminal rung — and -that a second backend (the pure-python ``accel`` stand-in) is bit-for-bit identical to it. The -``cuda`` cells are declared-but-unavailable here, so the matrix is enumerable without a GPU and -without importing torch. +The two tuple-keyed registries (COMPONENTS, KERNELS) plus Platform resolve implementations +generically across devices, with the numpy reference as the terminal rung; a second backend (the +pure-python ``accel`` stand-in) is bit-for-bit identical to it. The ``cuda`` cells are +declared-but-unavailable here, so the matrix is enumerable without a GPU and without importing torch. """ from __future__ import annotations @@ -153,8 +152,8 @@ def test_cuda_component_falls_back_to_factory_when_unavailable(): # (checked with the typed C1 instrument `parity.bit_identical`, not a bare ==) # # --------------------------------------------------------------------------- # def test_parity_oracle_accel_equals_cpu_ode(): - """accel's ODE solver kernel is an INDEPENDENT impl (`_accel_flow_match_step`) — this is the real - oracle: a separately-written backend kernel must match the numpy reference, C1 bit-identical.""" + """accel's ODE solver kernel is an independent impl (`_accel_flow_match_step`): a + separately-written backend kernel must match the numpy reference, C1 bit-identical.""" cpu_latent = _drive_denoise(_instance(Platform.cpu()), seed=7, steps=8) accel_latent = _drive_denoise(_instance(Platform.accel("sm90")), seed=7, steps=8) assert bit_identical(cpu_latent, accel_latent) diff --git a/v2/tests/test_policies.py b/v2/tests/test_policies.py index dbfcd8136..1b2b1c932 100644 --- a/v2/tests/test_policies.py +++ b/v2/tests/test_policies.py @@ -1,4 +1,4 @@ -"""Policy semantics (design_v3 §5.3, §6.2.3): CFG taxonomy, flow-shift, expert routing.""" +"""Policy semantics: CFG taxonomy, flow-shift, expert routing.""" from __future__ import annotations import numpy as np diff --git a/v2/tests/test_review_fixes.py b/v2/tests/test_review_fixes.py index 8f9cc151c..ce4f9316e 100644 --- a/v2/tests/test_review_fixes.py +++ b/v2/tests/test_review_fixes.py @@ -1,7 +1,7 @@ """Regression tests for the adversarial-review findings (untested paths now covered). Locks in: fail-fast on infeasible/deadlocked admission (no busy-spin), refundable compute budget -(concurrency gate, not lifetime cap), §7.1 feature-key partitioning by adapter stack + precision, +(concurrency gate, not lifetime cap), feature-key partitioning by adapter stack + precision, component-scoped weight invalidation (transformer sync keeps the text-encoder cache), and the interleave gate flagging symmetric-empty output. """ @@ -55,7 +55,7 @@ def test_feature_key_partitions_by_adapter_stack(): assert inst.caches.stats()["feature"]["hits"] == h + 1 inst.adapter_versions = {"te_lora": "2"} # swap te-LoRA stack on the resident instance m = inst.caches.stats()["feature"]["misses"] - cached_text_encode(inst, "p") # different stack ⇒ MISS, never a stale hit (§7.1) + cached_text_encode(inst, "p") # different stack ⇒ MISS, never a stale hit assert inst.caches.stats()["feature"]["misses"] == m + 1 diff --git a/v2/tests/test_served_reward.py b/v2/tests/test_served_reward.py index 3cbfe2c54..05b7ebc73 100644 --- a/v2/tests/test_served_reward.py +++ b/v2/tests/test_served_reward.py @@ -1,9 +1,9 @@ -"""Reward-model-as-a-served-card — REWARD_BATCH work units (design_v3 §10). +"""Reward-model-as-a-served-card — REWARD_BATCH work units. -In real RL the reward is a *model* (PickScore/CLIP/a VLM judge), not a heuristic. These tests make the +In real RL the reward is a model (PickScore/CLIP/a VLM judge), not a heuristic. These tests make the reward a first-class card: a `scorer` component + a `score` loop emitting `REWARD_BATCH` work units, and a `ServedRewardScorer` that drop-in replaces the numpy scorer so any RL method becomes RLHF/RLAIF with a -learned reward — no method change. Exercises the `REWARD_BATCH` WorkUnit kind (previously zero coverage). +learned reward, no method change. Exercises the `REWARD_BATCH` WorkUnit kind. """ from __future__ import annotations diff --git a/v2/tests/test_speculative.py b/v2/tests/test_speculative.py index 1a9752a78..f47b5f38d 100644 --- a/v2/tests/test_speculative.py +++ b/v2/tests/test_speculative.py @@ -1,4 +1,4 @@ -"""Speculative (draft-verify) decoding — exact + lower-latency AR (design_v3 §2.2; §9.16). +"""Speculative (draft-verify) decoding — exact + lower-latency AR. A cheap draft model proposes K tokens, the target verifies them, and the loop accepts the matching prefix + one target correction — a variable accepted-length per round (a ragged AR loop). The two diff --git a/v2/tests/test_t2i_i2v_workflow.py b/v2/tests/test_t2i_i2v_workflow.py index 292b0dbd9..c482987d2 100644 --- a/v2/tests/test_t2i_i2v_workflow.py +++ b/v2/tests/test_t2i_i2v_workflow.py @@ -1,4 +1,4 @@ -"""Cross-model T2I→I2V workflow — composition across model *instances* (design_v3 §13). +"""Cross-model T2I→I2V workflow — composition across model *instances*. LTX-2 already proves same-card multi-stage diffusion (base→refine). This proves the harder case the single-instance ``Program`` runner cannot express: chaining two *different* models (a T2I model and an diff --git a/v2/tests/test_thinker_talker.py b/v2/tests/test_thinker_talker.py index bd4986484..ffaa3dd47 100644 --- a/v2/tests/test_thinker_talker.py +++ b/v2/tests/test_thinker_talker.py @@ -1,4 +1,4 @@ -"""Qwen-Omni thinker→talker→vocoder — the third weight-sharing topology (design_v3 §4; vllm-omni +"""Qwen-Omni thinker→talker→vocoder — the third weight-sharing topology (vllm-omni ``qwen2_5_omni`` pipeline). Asserts the Card/Loop/Program vocabulary holds for a THREE-expert, THREE-loop cascade — the topology diff --git a/v2/tests/test_tiled_scheduling.py b/v2/tests/test_tiled_scheduling.py index 71407685d..13160087b 100644 --- a/v2/tests/test_tiled_scheduling.py +++ b/v2/tests/test_tiled_scheduling.py @@ -1,4 +1,4 @@ -"""Heterogeneous WorkUnit co-scheduling (design_v3 §6.1, §9.3, §17 falsifier). +"""Heterogeneous WorkUnit co-scheduling. The `VAE_TILE` WorkUnit kind had zero coverage. These tests put tiled VAE decode (one `VAE_TILE` unit per latent row-band) through the same admission/scheduler/interleave machinery as denoise steps, and @@ -7,11 +7,11 @@ assert: * **tiling is exact** — tiled decode == one-shot decode (a C0 component parity); * **the kind is real WorkUnits** — the tiled run emits one `VAE_TILE` unit per row; * **mixed kinds interleave bit-identically** — a `VAE_TILE` pipeline and a `DIFFUSION_STEP` pipeline, - interleaved, equal their serial run (the §9.3 gate over *heterogeneous* units); + interleaved, equal their serial run (the interleave gate over *heterogeneous* units); * **the scheduler co-runs both kinds** in one interleaved batch. -This validates the *mechanism* the §17 falsifier questions (non-step kinds flow through one budget). Whether -it *pays* on a real GPU duty-cycle is the empirical half of the falsifier — deferred to the port. +This validates the *mechanism*: non-step kinds flow through one budget. Whether it *pays* on a real GPU +duty-cycle is the empirical half — deferred to the port. """ from __future__ import annotations diff --git a/v2/tests/test_torch_backend.py b/v2/tests/test_torch_backend.py index f8b724daa..aa5d699ac 100644 --- a/v2/tests/test_torch_backend.py +++ b/v2/tests/test_torch_backend.py @@ -1,4 +1,4 @@ -"""GPU torch backend — the CPU-verifiable tier (design_v3 §17). +"""GPU torch backend — the CPU-verifiable tier. The torch adapters/kernels themselves need a GPU box (see GPU_BRINGUP.md). What IS verifiable on a torchless laptop is the wiring that must be right for the GPU path to even resolve correctly: the diff --git a/v2/tests/test_training.py b/v2/tests/test_training.py index 4e0960c28..31352bf40 100644 --- a/v2/tests/test_training.py +++ b/v2/tests/test_training.py @@ -1,4 +1,4 @@ -"""Training methods on Wan2.1-1.3B (design_v3 §10) — all over the SHARED loops, train ≡ serve. +"""Training methods on Wan2.1-1.3B — all over the SHARED loops, train ≡ serve. Verifies: finetune learns (loss ↓), DMD2 real/fake scores, NFT likelihood-free C2 + group reuse + old-as-behavior-policy, self-forcing on the causal chunk loop + slab-KV. @@ -40,7 +40,7 @@ def test_dmd2_real_and_fake_scores_finite_with_two_optimizers(): def test_diffusion_nft_is_likelihood_free_c2_with_group_reuse(): nft = build_diffusion_nft(build_wan21_card(), num_video_per_prompt=4, num_inner_timesteps=2) assert nft.manages_optimization() and nft.consistency_level() is ConsistencyLevel.C2 - assert nft.old_sync.role is WeightRole.OLD_POLICY # sample from OLD, not student (§8.4) + assert nft.old_sync.role is WeightRole.OLD_POLICY # sample from OLD, not student lm, m = nft.managed_train_step(BATCH, 0) assert np.isfinite(lm["policy_loss"]) and np.isfinite(m["old_deviate"]) assert m["advantage_std"] >= 0.0 diff --git a/v2/tests/test_unified_rl.py b/v2/tests/test_unified_rl.py index 5094b8183..ad7ae9a95 100644 --- a/v2/tests/test_unified_rl.py +++ b/v2/tests/test_unified_rl.py @@ -1,4 +1,4 @@ -"""UniRL/PromptRL joint LM+generator RL — the design stress test (design_v3 §4, §10; arXiv 2510.17937). +"""UniRL/PromptRL joint LM+generator RL — the design stress test (arXiv 2510.17937). These tests assert the *design holds*: a joint two-expert RL recipe fits inside the existing Card/Loop/Program + training-plane vocabulary with no new primitive. Specifically: @@ -67,7 +67,7 @@ def test_two_separate_experts_not_a_shared_mot(): assert inst.component("llm") is not inst.component("transformer") # genuinely separate weights -# --- serve ODE vs rollout SDE (gated, §9.4) --------------------------------------- # +# --- serve ODE vs rollout SDE (gated) --------------------------------------------- # def test_serve_is_ode_rollout_is_sde(): eng = build_unified_engine() @@ -136,7 +136,7 @@ def test_joint_rl_updates_both_experts(): # 4) two experts versioned INDEPENDENTLY; the frozen text-encoder cache is NOT invalidated inst = m.student assert last["llm_weights_version"] != last["transformer_weights_version"] - assert inst.version_of("text_encoder") == "v0" # §7.1 cache scope preserved + assert inst.version_of("text_encoder") == "v0" # cache scope preserved assert inst.version_of("llm") != "v0" and inst.version_of("transformer") != "v0" diff --git a/v2/tests/test_weight_sync_live.py b/v2/tests/test_weight_sync_live.py index 4714e6ee7..5eb6e7af4 100644 --- a/v2/tests/test_weight_sync_live.py +++ b/v2/tests/test_weight_sync_live.py @@ -1,4 +1,4 @@ -"""Live weight-sync under in-flight serving — the RL flywheel's hardest correctness (designv4 §9.14). +"""Live weight-sync under in-flight serving — the RL flywheel's hardest correctness. Collocated RL serves rollouts and receives weight updates on the same instance. The hazard: if new weights are swapped in *while a denoise loop is mid-flight*, that request's trajectory becomes a @@ -84,6 +84,6 @@ def test_sync_versions_and_invalidates_only_the_synced_component(): hits0 = inst.caches.stats()["feature"]["hits"] WeightSyncController(inst).sync(ToyDiT(seed=7)) # transformer-only sync assert inst.version_of("transformer") != "v0" # synced component bumped - assert inst.version_of("text_encoder") == "v0" # frozen encoder untouched (§7.1) + assert inst.version_of("text_encoder") == "v0" # frozen encoder untouched cached_text_encode(inst, "a shared prompt") # same prompt after the sync assert inst.caches.stats()["feature"]["hits"] == hits0 + 1 # text cache SURVIVED the transformer sync diff --git a/v2/tests/test_workflow_rl.py b/v2/tests/test_workflow_rl.py index 616bcee24..344e47884 100644 --- a/v2/tests/test_workflow_rl.py +++ b/v2/tests/test_workflow_rl.py @@ -1,4 +1,4 @@ -"""End-to-end RL over a cross-model workflow (design_v3 §10 + §13). +"""End-to-end RL over a cross-model workflow. Probes the training-plane / Workflow boundary: can "rollout == serve + capture" span *two model instances*, and can one final reward train an *earlier* model? These tests train both stages of the diff --git a/v2/tests/test_workflow_shapes.py b/v2/tests/test_workflow_shapes.py index 1c1ba0ad4..7f90614ad 100644 --- a/v2/tests/test_workflow_shapes.py +++ b/v2/tests/test_workflow_shapes.py @@ -1,4 +1,4 @@ -"""Non-linear workflow shapes — fan-out + best-of-N feedback (design_v3 §13; §9.17). +"""Non-linear workflow shapes — fan-out + best-of-N feedback. Every workflow so far was a linear chain. These exercise the two non-linear shapes: @@ -6,7 +6,7 @@ Every workflow so far was a linear chain. These exercise the two non-linear shap branch). The engine can interleave the branches' steps. * **Best-of-N** (`BestOfNWorkflow`): generate N candidates, score each with a reward scorer, return the best — inference-time scaling / rejection sampling, a feedback loop composing a generator with - the served REWARD_BATCH card (§9.15). + the served REWARD_BATCH card. """ from __future__ import annotations diff --git a/v2/tests/test_world_session.py b/v2/tests/test_world_session.py index 96d1b0c89..2a761f25f 100644 --- a/v2/tests/test_world_session.py +++ b/v2/tests/test_world_session.py @@ -1,14 +1,14 @@ -"""Interactive world-model session (design_v3 §16, §12) — the realtime/interactive plane. +"""Interactive world-model session — the realtime/interactive plane. The Session plane had zero coverage. These tests drive the causal ``chunk_rollout`` loop as a long-lived session and assert the seams a one-shot ``generate()`` can't reach: - * **persistent cross-request world state** — each ``act`` continues from the last (the §16 KV); + * **persistent cross-request world state** — each ``act`` continues from the last (the carried KV); * **history-dependence** — the same action with a different history yields a different world; * **transactional step-boundary cancellation** — a cancelled act leaves the session resumable; * **frame streaming** — chunks stream as they complete; * **no cross-session smearing** — interleaving two sessions is bit-identical to running them apart - (the §9.3 interleave guarantee, now for sessions). + (the interleave guarantee, now for sessions). """ from __future__ import annotations @@ -74,7 +74,7 @@ def test_cancel_mid_rollout_is_transactional(): def test_two_sessions_do_not_smear(): """Interleaving a second session's acts between this session's acts must not change this session's - output — per-session state isolation (the §9.3 interleave guarantee for the session plane).""" + output — per-session state isolation (the interleave guarantee for the session plane).""" eng = build_default_engine() a = WorldModelSession(eng, MID) a.act("x", seed=1) diff --git a/v2/training/__init__.py b/v2/training/__init__.py index 8c93fe853..92325d164 100644 --- a/v2/training/__init__.py +++ b/v2/training/__init__.py @@ -1,7 +1,7 @@ -"""Training plane — rollout records behavior on the same loops it serves (design_v3 §10). +"""Training plane — rollout records behavior on the same loops it serves. Dependency rule (enforced): ``training`` imports ``card``/``loop``/``runtime``/``models``; the -engine never imports ``training``. The rollout forward IS the serve forward plus capture. +engine never imports ``training``. The rollout forward is the serve forward plus capture. """ from __future__ import annotations @@ -29,11 +29,33 @@ from v2.training.rollout import RolloutContext, Trajectory, TrajectoryBuffer, ro from v2.training.weight_sync import WeightRole, WeightSyncController, WeightSyncPlan __all__ = [ - "TrainingMethod", "FineTuneMethod", "DMD2Method", "DiffusionNFTMethod", "SelfForcingMethod", - "UnifiedRLMethod", "JointMultiExpertRL", "WorkflowRLMethod", - "build_finetune", "build_dmd2", "build_diffusion_nft", "build_self_forcing", "build_unified_rl", - "build_joint_multi_rl", "build_workflow_rl", "run_flywheel", "derive_card", - "BehaviorRecord", "TrajectoryStatus", "Trajectory", "TrajectoryBuffer", "RolloutContext", - "rollout_loop", "WeightSyncPlan", "WeightSyncController", "WeightRole", - "MultiRewardScorer", "ServedRewardScorer", "build_multi_reward_scorer", + "TrainingMethod", + "FineTuneMethod", + "DMD2Method", + "DiffusionNFTMethod", + "SelfForcingMethod", + "UnifiedRLMethod", + "JointMultiExpertRL", + "WorkflowRLMethod", + "build_finetune", + "build_dmd2", + "build_diffusion_nft", + "build_self_forcing", + "build_unified_rl", + "build_joint_multi_rl", + "build_workflow_rl", + "run_flywheel", + "derive_card", + "BehaviorRecord", + "TrajectoryStatus", + "Trajectory", + "TrajectoryBuffer", + "RolloutContext", + "rollout_loop", + "WeightSyncPlan", + "WeightSyncController", + "WeightRole", + "MultiRewardScorer", + "ServedRewardScorer", + "build_multi_reward_scorer", ] diff --git a/v2/training/behavior.py b/v2/training/behavior.py index 71547a781..ef1558220 100644 --- a/v2/training/behavior.py +++ b/v2/training/behavior.py @@ -1,9 +1,7 @@ -"""BehaviorRecord — captured at generation time on the serving loop (design_v3 §10). +"""BehaviorRecord — captured at generation time on the serving loop. -> The rollout forward *is* the serve forward plus capture. - -Sized honestly (design_v3 §10): full routing capture is GB/sample for omni, so it is an opt-in -instrument. ``logprobs`` is None for likelihood-free methods (DiffusionNFT) — the C2 split. +The rollout forward is the serve forward plus capture. Full routing capture is GB/sample for omni, +so it is opt-in. ``logprobs`` is None for likelihood-free methods (DiffusionNFT). """ from __future__ import annotations @@ -30,8 +28,8 @@ class BehaviorRecord: seeds: dict[str, int] = field(default_factory=dict) timesteps: list[float] = field(default_factory=list) sigmas: list[float] = field(default_factory=list) - latents: Any = None # final clean latents (or refs if too large) - logprobs: list[float] | None = None # None for likelihood-free methods (NFT) + latents: Any = None # final clean latents (or refs if too large) + logprobs: list[float] | None = None # None for likelihood-free methods (NFT) sampled_tokens: list[int] = field(default_factory=list) guidance: dict[str, float] = field(default_factory=dict) reward_inputs: dict[str, Any] = field(default_factory=dict) diff --git a/v2/training/flywheel.py b/v2/training/flywheel.py index f95f9b84b..acda9becb 100644 --- a/v2/training/flywheel.py +++ b/v2/training/flywheel.py @@ -1,17 +1,14 @@ -"""The (recipe, runtime) flywheel — RL → distill → a faster card (design_v3 §16; §9.18). +"""The (recipe, runtime) flywheel — RL then distill into a faster card. -design_v3's central product claim: *preference/reward data → RL → a faster distilled card → a better -product*, with the recipe carrying provenance so the chain is auditable. This wires it end-to-end on the -shared loops: +Reward data → RL → a faster distilled card, with the recipe carrying provenance so the chain is +auditable. Wired end-to-end on the shared loops: - 1. **RL** the base model (DiffusionNFT) → improved policy weights. - 2. **Distill** from the RL'd model: a DMD2 student with the **RL'd model as its teacher** learns to - match it in fewer steps. - 3. The distilled card's ``RecipeSpec.parents`` records the chain ``base → rl → distilled`` — the - (recipe, runtime) provenance that makes "this fast model came from that RL run" a typed fact. + 1. RL the base model (DiffusionNFT) → improved policy weights. + 2. Distill from the RL'd model: a DMD2 student with the RL'd model as its teacher learns to match + it in fewer steps. + 3. The distilled card's ``RecipeSpec.parents`` records the chain ``base → rl → distilled``. -Both stages drive the SAME ``diffusion_denoise`` loop the engine serves (no second sampler) — so the -flywheel is real, not aspirational. +Both stages drive the same ``diffusion_denoise`` loop the engine serves (no second sampler). """ from __future__ import annotations @@ -27,8 +24,7 @@ def derive_card(base_card: Any, model_id: str, *, method: str, parents: list[str return dataclasses.replace(base_card, model_id=model_id, recipe=recipe).validate() -def run_flywheel(base_card: Any, *, rl_iters: int = 8, distill_iters: int = 15, - batch: dict | None = None) -> dict: +def run_flywheel(base_card: Any, *, rl_iters: int = 8, distill_iters: int = 15, batch: dict | None = None) -> dict: """RL the base, then distill from the RL'd model. Returns the instances + the provenance cards.""" batch = batch or {"prompts": ["a red car", "a blue boat"], "seeds": [1, 2]} @@ -36,17 +32,23 @@ def run_flywheel(base_card: Any, *, rl_iters: int = 8, distill_iters: int = 15, nft = build_diffusion_nft(base_card, num_video_per_prompt=4, num_inner_timesteps=2) for it in range(rl_iters): nft.train_step(batch, it) - rl_card = derive_card(base_card, base_card.model_id + "-rl", - method="diffusion_nft", parents=[base_card.model_id]) + rl_card = derive_card(base_card, base_card.model_id + "-rl", method="diffusion_nft", parents=[base_card.model_id]) # 2) Distill FROM the RL'd model: the DMD2 teacher IS the RL'd policy dmd = build_dmd2(base_card, rollout_steps=4) dmd.teacher.component("transformer").copy_from(nft.student.component("transformer")) for it in range(distill_iters): dmd.train_step(batch, it) - distilled_card = derive_card(base_card, base_card.model_id + "-rl-distilled", - method="dmd2", parents=[rl_card.model_id]) + distilled_card = derive_card(base_card, + base_card.model_id + "-rl-distilled", + method="dmd2", + parents=[rl_card.model_id]) - return {"rl": nft.student, "rl_card": rl_card, - "distilled": dmd.student, "distilled_card": distilled_card, - "teacher": dmd.teacher, "base_card": base_card} + return { + "rl": nft.student, + "rl_card": rl_card, + "distilled": dmd.student, + "distilled_card": distilled_card, + "teacher": dmd.teacher, + "base_card": base_card + } diff --git a/v2/training/methods/__init__.py b/v2/training/methods/__init__.py index 4c658c85b..fe984d31f 100644 --- a/v2/training/methods/__init__.py +++ b/v2/training/methods/__init__.py @@ -1,4 +1,4 @@ -"""Training methods — all on Wan2.1-1.3B, all over the SHARED loops (design_v3 §10).""" +"""Training methods — all on Wan2.1-1.3B, all over the SHARED loops.""" from __future__ import annotations from v2.training.methods.base import TrainingMethod, new_instance, predict_x0 @@ -11,12 +11,21 @@ from v2.training.methods.unified_rl import UnifiedRLMethod, build_unified_rl from v2.training.methods.workflow_rl import WorkflowRLMethod, build_workflow_rl __all__ = [ - "TrainingMethod", "new_instance", "predict_x0", - "FineTuneMethod", "build_finetune", - "DMD2Method", "build_dmd2", - "DiffusionNFTMethod", "build_diffusion_nft", - "SelfForcingMethod", "build_self_forcing", - "UnifiedRLMethod", "build_unified_rl", - "JointMultiExpertRL", "build_joint_multi_rl", - "WorkflowRLMethod", "build_workflow_rl", + "TrainingMethod", + "new_instance", + "predict_x0", + "FineTuneMethod", + "build_finetune", + "DMD2Method", + "build_dmd2", + "DiffusionNFTMethod", + "build_diffusion_nft", + "SelfForcingMethod", + "build_self_forcing", + "UnifiedRLMethod", + "build_unified_rl", + "JointMultiExpertRL", + "build_joint_multi_rl", + "WorkflowRLMethod", + "build_workflow_rl", ] diff --git a/v2/training/methods/base.py b/v2/training/methods/base.py index 61f5f01c9..d0bfc6010 100644 --- a/v2/training/methods/base.py +++ b/v2/training/methods/base.py @@ -1,9 +1,8 @@ -"""TrainingMethod base (design_v3 §10; mirrors the repo's train/methods contract). +"""TrainingMethod base (mirrors the repo's train/methods contract). -The load-bearing principle: every method's *rollout* drives the SAME loop the engine serves -(via ``_rollout`` → ``rollout_loop``), so there is one numerics surface. Methods differ only in -loss math, roles, and capture — not in a second sampler. ``manages_optimization`` matches the -landed ``TrainingMethod.manages_optimization()`` seam (RL owns its sample→score→inner-train cadence). +Every method's rollout drives the same loop the engine serves (via ``_rollout`` → ``rollout_loop``), +so there is one numerics surface. Methods differ only in loss math, roles, and capture — not in a +second sampler. ``manages_optimization`` lets RL own its sample→score→inner-train cadence. """ from __future__ import annotations @@ -61,7 +60,10 @@ class TrainingMethod(ABC): def _rollout(self, request, *, instance: Any = None, loop_id: str | None = None): inst = instance or self.student slots = self._encode_slots(request, inst) - return rollout_loop(inst, loop_id or self.student_loop_id, request, slots=slots, + return rollout_loop(inst, + loop_id or self.student_loop_id, + request, + slots=slots, profile=ExecutionProfile.ROLLOUT) @abstractmethod diff --git a/v2/training/methods/diffusion_nft.py b/v2/training/methods/diffusion_nft.py index 1441c4129..75556fd29 100644 --- a/v2/training/methods/diffusion_nft.py +++ b/v2/training/methods/diffusion_nft.py @@ -1,11 +1,10 @@ -"""DiffusionNFT — likelihood-free RL (design_v3 §10, §9.2 C2; repo: train/methods/rl/diffusion_nft). +"""DiffusionNFT — likelihood-free RL (repo: train/methods/rl/diffusion_nft). -The landed RL method, the C2-split's reason for being: it is **likelihood-free** — no log-probs to -match — so consistency is the C2 *behavioral* rung (seeded final sample + prediction-space identity: -``old_deviate`` / ref-MSE), declared on the recipe. Sampling is from the **old** (decay-blended -behavior) policy, NOT the student (the §8.4 caveat the WeightSyncPlan must carry). K samples of one -prompt form a homogeneous group; group-relative advantages drive the update; the shared prompt -encodes once via the feature cache (the §10 ``24×`` text-encode reduction). +Likelihood-free (no log-probs to match), so consistency is the C2 behavioral rung: seeded final +sample plus prediction-space identity (``old_deviate`` / ref-MSE), declared on the recipe. Sampling +is from the old (decay-blended behavior) policy, NOT the student — a caveat the WeightSyncPlan must +carry. K samples of one prompt form a homogeneous group; group-relative advantages drive the update; +the shared prompt encodes once via the feature cache. NFT objective (repo diffusion_nft.py): positive = β·forward + (1−β)·old ; implicit_neg = (1+β)·old − β·forward @@ -28,15 +27,26 @@ from v2.training.methods.base import TrainingMethod, new_instance class DiffusionNFTMethod(TrainingMethod): name = "diffusion_nft" - consistency = ConsistencyLevel.C2 # likelihood-free behavioral identity + consistency = ConsistencyLevel.C2 # likelihood-free behavioral identity - def __init__(self, student_instance, old_instance, reference_instance, *, - reward_fn: dict | None = None, num_video_per_prompt: int = 4, rollout_steps: int = 3, - beta: float = 0.1, kl_beta: float = 1e-4, adv_clip_max: float = 5.0, - num_inner_timesteps: int = 2, lr: float = 0.05, decay_type: int = 1, **kw): + def __init__(self, + student_instance, + old_instance, + reference_instance, + *, + reward_fn: dict | None = None, + num_video_per_prompt: int = 4, + rollout_steps: int = 3, + beta: float = 0.1, + kl_beta: float = 1e-4, + adv_clip_max: float = 5.0, + num_inner_timesteps: int = 2, + lr: float = 0.05, + decay_type: int = 1, + **kw): super().__init__(student_instance, lr=lr, **kw) - self.old = old_instance # behavior policy (sampled from!) - self.reference = reference_instance # frozen KL anchor + self.old = old_instance # behavior policy (sampled from!) + self.reference = reference_instance # frozen KL anchor self.K = num_video_per_prompt self.rollout_steps = rollout_steps self.beta = beta @@ -50,11 +60,11 @@ class DiffusionNFTMethod(TrainingMethod): self.old.component("transformer").copy_from(self.student_dit) self.reference.component("transformer").copy_from(self.student_dit) - def manages_optimization(self) -> bool: # RL owns its sample→score→train cadence + def manages_optimization(self) -> bool: # RL owns its sample→score→train cadence return True def get_grad_clip_targets(self, iteration: int = 0) -> dict: - return {"student": self.student_dit} # NFT clips internally (matches repo) + return {"student": self.student_dit} # NFT clips internally (matches repo) @staticmethod def _decay(step: int, decay_type: int) -> float: @@ -69,12 +79,18 @@ class DiffusionNFTMethod(TrainingMethod): buf = TrajectoryBuffer() for k in range(self.K): seed = base_seed * 1000 + k - req = make_request(TaskType.T2V, self.old.card.model_id, prompt, + req = make_request(TaskType.T2V, + self.old.card.model_id, + prompt, diffusion=DiffusionParams(num_steps=self.rollout_steps, seed=seed)) - res = self._rollout(req, instance=self.old) # behavior policy = OLD - buf.add(Trajectory(request_id=req.request_id, prompt=prompt, seed=seed, - latents=np.asarray(res.outputs["latents"], dtype="float32"), - behavior=res.behavior, weights_version=self.old.weights_version)) + res = self._rollout(req, instance=self.old) # behavior policy = OLD + buf.add( + Trajectory(request_id=req.request_id, + prompt=prompt, + seed=seed, + latents=np.asarray(res.outputs["latents"], dtype="float32"), + behavior=res.behavior, + weights_version=self.old.weights_version)) return buf def managed_train_step(self, batch: dict, iteration: int) -> tuple[dict, dict]: @@ -88,12 +104,12 @@ class DiffusionNFTMethod(TrainingMethod): media = [t.latents for t in group] scored = self.scorer.score(media, [prompt] * len(group)) rewards = scored["avg"] - adv = (rewards - rewards.mean()) / (rewards.std() + 1e-4) # group-relative + adv = (rewards - rewards.mean()) / (rewards.std() + 1e-4) # group-relative adv = np.clip(adv, -self.adv_clip_max, self.adv_clip_max) all_reward.extend(rewards.tolist()) all_adv.extend(adv.tolist()) - emb = cached_text_encode(self.student, prompt) # encoded ONCE for the group + emb = cached_text_encode(self.student, prompt) # encoded ONCE for the group old_dit, ref_dit = self.old.component("transformer"), self.reference.component("transformer") for k, traj in enumerate(group): x0 = traj.latents @@ -112,13 +128,13 @@ class DiffusionNFTMethod(TrainingMethod): x0n = xt - t * neg wfp = max(float(np.mean(np.abs(x0p - x0))), 1e-5) wfn = max(float(np.mean(np.abs(x0n - x0))), 1e-5) - pos_loss = float(np.mean((x0p - x0) ** 2) / wfp) - neg_loss = float(np.mean((x0n - x0) ** 2) / wfn) + pos_loss = float(np.mean((x0p - x0)**2) / wfp) + neg_loss = float(np.mean((x0n - x0)**2) / wfn) policy_loss = (r * pos_loss / self.beta + (1 - r) * neg_loss / self.beta) * self.adv_clip_max - kl = float(np.mean((fwd - ref_p) ** 2)) + kl = float(np.mean((fwd - ref_p)**2)) policy_losses.append(policy_loss) kl_losses.append(kl) - old_devs.append(float(np.mean((fwd - old_p) ** 2))) + old_devs.append(float(np.mean((fwd - old_p)**2))) # update student toward x0 along the (reward-weighted) positive direction v_target = ((xt - x0) / max(t, 1e-3)).astype("float32") _, gn = self.student_dit.mse_grad_step(xt, emb, t, v_target, self.lr * max(r, 0.05)) diff --git a/v2/training/methods/dmd2.py b/v2/training/methods/dmd2.py index 1373a767c..1f36e9610 100644 --- a/v2/training/methods/dmd2.py +++ b/v2/training/methods/dmd2.py @@ -1,8 +1,8 @@ -"""DMD2 distribution-matching distillation (design_v3 §10; repo: train/methods/distribution_matching). +"""DMD2 distribution-matching distillation (repo: train/methods/distribution_matching). Roles: student (trainable) · teacher (frozen real score) · critic (trainable fake score). -The student rollout is driven through the SHARED denoise loop (not a vendored sampler) — the §10 -collocation moat. Loss math (repo dmd2.py): +The student rollout is driven through the shared denoise loop (not a vendored sampler). Loss math +(repo dmd2.py): real_cfg_x0 = teacher_uncond_x0 + s·(teacher_cond_x0 − teacher_uncond_x0) faker_x0 = critic_x0 @@ -22,11 +22,19 @@ from v2.training.methods.base import TrainingMethod, new_instance, predict_x0 class DMD2Method(TrainingMethod): name = "dmd2" - consistency = ConsistencyLevel.C1 # kernel-pinned: rollout under trainer kernels + consistency = ConsistencyLevel.C1 # kernel-pinned: rollout under trainer kernels - def __init__(self, student_instance, teacher_instance, critic_instance, *, - lr: float = 0.05, critic_lr: float = 0.05, guidance_scale: float = 4.5, - generator_update_interval: int = 5, rollout_steps: int = 3, **kw): + def __init__(self, + student_instance, + teacher_instance, + critic_instance, + *, + lr: float = 0.05, + critic_lr: float = 0.05, + guidance_scale: float = 4.5, + generator_update_interval: int = 5, + rollout_steps: int = 3, + **kw): super().__init__(student_instance, lr=lr, **kw) self.teacher = teacher_instance self.critic = critic_instance @@ -42,7 +50,9 @@ class DMD2Method(TrainingMethod): return t def _rollout_sample(self, prompt: str, seed: int) -> np.ndarray: - req = make_request(TaskType.T2V, self.student.card.model_id, prompt, + req = make_request(TaskType.T2V, + self.student.card.model_id, + prompt, diffusion=DiffusionParams(num_steps=self.rollout_steps, seed=seed)) return np.asarray(self._rollout(req).outputs["latents"], dtype="float32") @@ -55,7 +65,7 @@ class DMD2Method(TrainingMethod): seed = batch.get("seeds", list(range(len(batch["prompts"]))))[i] emb = cached_text_encode(self.student, p) neg = cached_text_encode(self.student, "") - x0g = self._rollout_sample(p, seed) # student rollout via SHARED loop + x0g = self._rollout_sample(p, seed) # student rollout via SHARED loop rng = np.random.default_rng(seed ^ 0xD2) sigma = float(rng.uniform(0.1, 0.9)) @@ -69,7 +79,7 @@ class DMD2Method(TrainingMethod): faker = predict_x0(critic(noised, emb, sigma), noised, sigma) denom = float(np.mean(np.abs(x0g - real_cfg))) + 1e-6 grad = np.nan_to_num((faker - real_cfg) / denom).astype("float32") - gen_losses.append(0.5 * float(np.mean(grad ** 2))) # DMD generator loss value + gen_losses.append(0.5 * float(np.mean(grad**2))) # DMD generator loss value # generator (student) update toward (gen_x0 − grad), converted to velocity space if iteration % self.generator_update_interval == 0: @@ -82,8 +92,7 @@ class DMD2Method(TrainingMethod): sigma2 = float(rng.uniform(0.1, 0.9)) noise2 = rng.standard_normal(x0g.shape).astype("float32") noised2 = ((1.0 - sigma2) * x0g + sigma2 * noise2).astype("float32") - c_loss, gc = critic.mse_grad_step(noised2, emb, sigma2, (noise2 - x0g).astype("float32"), - self.critic_lr) + c_loss, gc = critic.mse_grad_step(noised2, emb, sigma2, (noise2 - x0g).astype("float32"), self.critic_lr) critic_losses.append(c_loss) gnorm_c.append(gc) diff --git a/v2/training/methods/finetune.py b/v2/training/methods/finetune.py index 80eb8311a..d43dd1f55 100644 --- a/v2/training/methods/finetune.py +++ b/v2/training/methods/finetune.py @@ -1,7 +1,7 @@ -"""FineTuneMethod — plain flow-matching finetune (design_v3 §10; repo: train/methods/fine_tuning). +"""FineTuneMethod — plain flow-matching finetune. -Loss (repo finetune.py): flow-matching target ``v = ε − x0``; ``L = MSE(student(x_σ), v)`` where -``x_σ = (1−σ)·x0 + σ·ε``. Data-centric (no rollout) — the simplest (recipe, runtime) consumer. +Flow-matching target ``v = ε − x0``; ``L = MSE(student(x_σ), v)`` where ``x_σ = (1−σ)·x0 + σ·ε``. +Data-centric (no rollout) — the simplest (recipe, runtime) consumer. """ from __future__ import annotations @@ -26,20 +26,24 @@ class FineTuneMethod(TrainingMethod): seeds = batch.get("seeds", list(range(len(prompts)))) latents = batch.get("latents") losses, gnorms = [], [] - for i, (p, s) in enumerate(zip(prompts, seeds)): + for i, (p, s) in enumerate(zip(prompts, seeds, strict=False)): x0 = np.asarray(latents[i], dtype="float32") if latents is not None else self._x0(s) emb = cached_text_encode(self.student, p) rng = np.random.default_rng(s) sigma = float(rng.uniform(0.05, 0.95)) noise = rng.standard_normal(x0.shape).astype("float32") noised = ((1.0 - sigma) * x0 + sigma * noise).astype("float32") - target = (noise - x0).astype("float32") # flow-matching velocity target + target = (noise - x0).astype("float32") # flow-matching velocity target loss, gnorm = self.student_dit.mse_grad_step(noised, emb, sigma, target, self.lr) losses.append(loss) gnorms.append(gnorm) self.iteration = iteration - return ({"loss": float(np.mean(losses))}, - {"loss": float(np.mean(losses)), "grad_norm/student": float(np.mean(gnorms))}) + return ({ + "loss": float(np.mean(losses)) + }, { + "loss": float(np.mean(losses)), + "grad_norm/student": float(np.mean(gnorms)) + }) def build_finetune(card, **kw) -> FineTuneMethod: diff --git a/v2/training/methods/joint_multi_rl.py b/v2/training/methods/joint_multi_rl.py index 0bca25f43..d49c17d5a 100644 --- a/v2/training/methods/joint_multi_rl.py +++ b/v2/training/methods/joint_multi_rl.py @@ -1,26 +1,17 @@ -"""JointMultiExpertRL — N-way joint RL over an arbitrary set of experts (design_v3 §10). +"""JointMultiExpertRL — N-way joint RL over an arbitrary set of experts. -The generalization of ``UnifiedRLMethod`` (which is the N=2 special case: one refiner LM + one -generator) to **N refiner LMs + one generator**. It answers the question "can we do joint RL over more -than two experts?" — *yes, with no substrate rewrite*. The proof is structural: +Generalizes ``UnifiedRLMethod`` (the N=2 case: one refiner LM + one generator) to N refiner LMs + one +generator, with no substrate change: the card already holds N components/loops, WeightSyncPlan is +per-component (one plan per expert), and ``get_grad_clip_targets`` already returns a dict. Only this +method's body changes — a loop over the expert list. One reward → one group-relative advantage → N +policy-gradient updates (N token-PG refiners + one FlowGRPO PPO generator). The refiners' refinements +compose into the diffusion conditioning, so the single reward depends on all of them. - * the **card** already holds N components and N loops (Qwen-Omni shipped three) — no change; - * **WeightSyncPlan** is already per-component, so N experts version + cache-scope independently — - this method just builds a ``dict`` of N plans instead of two; - * ``get_grad_clip_targets`` already returns a ``dict`` — N entries instead of two; - * rollout / behavior capture / the consistency ladder are per-loop, agnostic to N. - -So the ONLY thing that was "two" was this method's body. Here it is a loop over an expert list. One -reward → one group-relative advantage → N policy-gradient updates (N token-PG refiners + one FlowGRPO -PPO generator). The N refiners' refinements **compose** into the diffusion conditioning, so one reward -genuinely depends on all of them. - -Credit assignment (the real research subtlety, not an architectural one): - * ``credit="shared"`` — one reward, one advantage applied to every expert (faithful to UniRL). - Works, but each refiner sees the others as noise (shared-reward multi-agent variance). - * ``credit="per_expert"`` — each refiner gets an advantage from its OWN reward term (a clean - decomposition). Also supported with no substrate change — it is a reward-shaping choice. The - generator always uses the total reward's advantage. +Credit assignment (a reward-shaping choice, not an architectural one): + * ``credit="shared"`` — one advantage applied to every expert (faithful to UniRL); each refiner sees + the others as noise (shared-reward multi-agent variance). + * ``credit="per_expert"`` — each refiner gets an advantage from its OWN reward term (clean + decomposition). The generator always uses the total reward's advantage. """ from __future__ import annotations @@ -41,12 +32,26 @@ class JointMultiExpertRL(TrainingMethod): consistency = ConsistencyLevel.C2 student_loop_id = "diffusion_denoise" - def __init__(self, student_instance, reference_instance, *, refiner_ids, generator_id="transformer", - target_actions, joint_generator: bool = True, credit: str = "shared", - num_samples_per_prompt: int = 4, num_skip_refinement: int = 1, rollout_steps: int = 4, - refiner_lr: float = 0.4, dit_lr: float = 0.02, beta: float = 0.01, - adv_clip_max: float = 5.0, ppo_clip: float = 1e-2, sde_noise_scale: float = 0.7, - action_bonus: float = 2.0, **kw): + def __init__(self, + student_instance, + reference_instance, + *, + refiner_ids, + generator_id="transformer", + target_actions, + joint_generator: bool = True, + credit: str = "shared", + num_samples_per_prompt: int = 4, + num_skip_refinement: int = 1, + rollout_steps: int = 4, + refiner_lr: float = 0.4, + dit_lr: float = 0.02, + beta: float = 0.01, + adv_clip_max: float = 5.0, + ppo_clip: float = 1e-2, + sde_noise_scale: float = 0.7, + action_bonus: float = 2.0, + **kw): super().__init__(student_instance, lr=dit_lr, **kw) self.reference = reference_instance self.refiner_ids = list(refiner_ids) @@ -60,12 +65,11 @@ class JointMultiExpertRL(TrainingMethod): self.refiner_lr, self.dit_lr = refiner_lr, dit_lr self.beta, self.adv_clip_max, self.ppo_clip = beta, adv_clip_max, ppo_clip self.sde_noise_scale, self.action_bonus = sde_noise_scale, action_bonus - # one WeightSyncPlan PER trainable expert — independent versions + cache scopes (§7.1, §10) - self.sync = {rid: WeightSyncPlan(role=WeightRole.STUDENT, components=(rid,)) - for rid in self.refiner_ids} + # one WeightSyncPlan PER trainable expert — independent versions + cache scopes + self.sync = {rid: WeightSyncPlan(role=WeightRole.STUDENT, components=(rid, )) for rid in self.refiner_ids} if joint_generator: - self.sync[generator_id] = WeightSyncPlan(role=WeightRole.STUDENT, components=(generator_id,)) - for rid in self.refiner_ids: # frozen reference per refiner + self.sync[generator_id] = WeightSyncPlan(role=WeightRole.STUDENT, components=(generator_id, )) + for rid in self.refiner_ids: # frozen reference per refiner self.reference.component(rid).copy_from(self.student.component(rid)) self.reference.component(generator_id).copy_from(self.student.component(generator_id)) @@ -98,21 +102,25 @@ class JointMultiExpertRL(TrainingMethod): diff = float(np.tanh(-np.std(np.asarray(latent, dtype="float64")))) per, total = {}, diff for rid in self.refiner_ids: - hit = self.action_bonus if (actions is not None - and actions.get(rid) == self.targets.get(rid)) else 0.0 + hit = self.action_bonus if (actions is not None and actions.get(rid) == self.targets.get(rid)) else 0.0 total += hit per[rid] = diff + hit return total, per def _rollout_sample(self, prompt, emb, neg_emb, seed): - req = make_request(TaskType.T2V, self.student.card.model_id, prompt, - diffusion=DiffusionParams(num_steps=self.rollout_steps, seed=seed, - guidance_scale=1.0, sde_rollout=self.joint_generator, + req = make_request(TaskType.T2V, + self.student.card.model_id, + prompt, + diffusion=DiffusionParams(num_steps=self.rollout_steps, + seed=seed, + guidance_scale=1.0, + sde_rollout=self.joint_generator, sde_noise_scale=self.sde_noise_scale)) - slots = {"text_embeds": np.asarray(emb, dtype="float32"), - "neg_text_embeds": np.asarray(neg_emb, dtype="float32")} - return rollout_loop(self.student, self.student_loop_id, req, slots=slots, - profile=ExecutionProfile.ROLLOUT) + slots = { + "text_embeds": np.asarray(emb, dtype="float32"), + "neg_text_embeds": np.asarray(neg_emb, dtype="float32") + } + return rollout_loop(self.student, self.student_loop_id, req, slots=slots, profile=ExecutionProfile.ROLLOUT) def _dit_ppo(self, behavior, emb, advantage): ref_dit = self.reference.component(self.generator_id) @@ -125,21 +133,23 @@ class JointMultiExpertRL(TrainingMethod): prev, sample = rec["prev"], rec["sample"] st, sn = float(rec["sigma_t"]), float(rec["sigma_next"]) v_cur = self.dit(prev, emb, st) - _, logp_cur, mean_cur, eff_std = sde_step( - prev, v_cur, st, sn, prev_sample=sample, noise_scale=self.sde_noise_scale) + _, logp_cur, mean_cur, eff_std = sde_step(prev, + v_cur, + st, + sn, + prev_sample=sample, + noise_scale=self.sde_noise_scale) v_ref = ref_dit(prev, emb, st) - _, _, mean_ref, _ = sde_step( - prev, v_ref, st, sn, prev_sample=sample, noise_scale=self.sde_noise_scale) + _, _, mean_ref, _ = sde_step(prev, v_ref, st, sn, prev_sample=sample, noise_scale=self.sde_noise_scale) ratio = float(np.exp(np.clip(logp_cur - float(rec["sde_logprob"]), -10.0, 10.0))) clipped = float(np.clip(ratio, 1.0 - self.ppo_clip, 1.0 + self.ppo_clip)) ppo_losses.append(max(-advantage * ratio, -advantage * clipped)) - kls.append(float(np.mean((mean_cur - mean_ref) ** 2) / (2.0 * max(eff_std ** 2, 1e-8)))) + kls.append(float(np.mean((mean_cur - mean_ref)**2) / (2.0 * max(eff_std**2, 1e-8)))) ratios.append(ratio) v_target = _ml_velocity(prev, sample, st, sn, noise_scale=self.sde_noise_scale) # PG direction - _, gn = self.dit.mse_grad_step(prev, emb, st, v_target, - self.dit_lr * float(np.clip(advantage, -1.0, 1.0))) + _, gn = self.dit.mse_grad_step(prev, emb, st, v_target, self.dit_lr * float(np.clip(advantage, -1.0, 1.0))) gnorms.append(gn) - m = lambda xs: float(np.mean(xs)) if xs else 0.0 # noqa: E731 + m = lambda xs: float(np.mean(xs)) if xs else 0.0 # noqa: E731 return m(ppo_losses), m(kls), m(ratios), m(gnorms) @staticmethod @@ -160,7 +170,7 @@ class JointMultiExpertRL(TrainingMethod): base = cached_text_encode(self.student, prompt) rng = np.random.default_rng((seeds[pi] + 1) * 7919 + iteration) actions, logps = {}, {} - for rid in self.refiner_ids: # each refiner samples its action + for rid in self.refiner_ids: # each refiner samples its action actions[rid], logps[rid] = self.refiner(rid).sample_refinement(rng) composed = self._compose(base, actions) @@ -194,20 +204,21 @@ class JointMultiExpertRL(TrainingMethod): for k in range(self.K): if not samples[k][2]: continue - ppo, kl, ratio, gn = self._dit_ppo(samples[k][0].behavior, samples[k][1], - float(adv_total[k])) + ppo, kl, ratio, gn = self._dit_ppo(samples[k][0].behavior, samples[k][1], float(adv_total[k])) dit_losses.append(ppo + self.beta * kl) kls.append(kl) ratios.append(ratio) dit_gn.append(gn) # publish every expert independently (N+? distinct versions + cache scopes) - versions = {rid: self.sync[rid].apply(self.refiner(rid), self.refiner(rid), self.student) - for rid in self.refiner_ids} + versions = { + rid: self.sync[rid].apply(self.refiner(rid), self.refiner(rid), self.student) + for rid in self.refiner_ids + } if self.joint_generator: versions[self.generator_id] = self.sync[self.generator_id].apply(self.dit, self.dit, self.student) - m = lambda xs: float(np.mean(xs)) if xs else 0.0 # noqa: E731 + m = lambda xs: float(np.mean(xs)) if xs else 0.0 # noqa: E731 metrics = { "reward_mean": m(all_reward), "n_experts": float(len(self.refiner_ids) + (1 if self.joint_generator else 0)), @@ -233,5 +244,8 @@ class JointMultiExpertRL(TrainingMethod): def build_joint_multi_rl(card, *, refiner_ids, target_actions, **kw) -> JointMultiExpertRL: """Two roles: a trainable student (N refiners + generator) and a frozen reference (all experts).""" - return JointMultiExpertRL(new_instance(card), new_instance(card), - refiner_ids=refiner_ids, target_actions=target_actions, **kw) + return JointMultiExpertRL(new_instance(card), + new_instance(card), + refiner_ids=refiner_ids, + target_actions=target_actions, + **kw) diff --git a/v2/training/methods/self_forcing.py b/v2/training/methods/self_forcing.py index 3fbf00517..674c057cc 100644 --- a/v2/training/methods/self_forcing.py +++ b/v2/training/methods/self_forcing.py @@ -1,13 +1,11 @@ -"""Self-forcing — causal video distillation (design_v3 §10; repo: self_forcing.py extends DMD2). +"""Self-forcing — causal video distillation (extends DMD2): train against your own KV-cached causal +rollout. -> Self-forcing is "train against your own KV-cached causal rollout." - -The student rollout is the SHARED ``chunk_rollout`` loop (with the slab-KV cache) — exactly the -loop the engine streams at serve time — so the causal student is distilled under the runtime it -will be served on (the (recipe, runtime) pair, §2.1). The distribution-matching loss is DMD2's, -applied to the causal rollout latents. (The repo uses a dedicated SelfForcingFlowMatchScheduler; -the mini reuses the shared flow-match sampler — the structure being demonstrated is the causal -rollout + KV + DMD loss.) +The student rollout is the SHARED ``chunk_rollout`` loop (with the slab-KV cache) — exactly the loop +the engine streams at serve time — so the causal student is distilled under the runtime it will be +served on. The loss is DMD2's distribution-matching loss applied to the causal rollout latents. (The +demonstrated structure is the causal rollout + KV + DMD loss; this mini reuses the shared flow-match +sampler in place of the repo's dedicated SelfForcingFlowMatchScheduler.) """ from __future__ import annotations @@ -22,17 +20,15 @@ from v2.training.methods.dmd2 import DMD2Method class SelfForcingMethod(DMD2Method): name = "self_forcing" consistency = ConsistencyLevel.C1 - student_loop_id = "chunk_rollout" # the SHARED causal loop (KV-cached) + student_loop_id = "chunk_rollout" # the SHARED causal loop (KV-cached) def _rollout_sample(self, prompt: str, seed: int) -> np.ndarray: - req = make_request(TaskType.T2V, self.student.card.model_id, prompt, - diffusion=DiffusionParams(seed=seed)) - res = self._rollout(req) # drives chunk_rollout + slab-KV + req = make_request(TaskType.T2V, self.student.card.model_id, prompt, diffusion=DiffusionParams(seed=seed)) + res = self._rollout(req) # drives chunk_rollout + slab-KV latents = res.outputs["latents"] return np.asarray(latents, dtype="float32") def build_self_forcing(causal_card, **kw) -> SelfForcingMethod: """Build a self-forcing method over a Wan-causal card (chunk_rollout loop).""" - return SelfForcingMethod(new_instance(causal_card), new_instance(causal_card), - new_instance(causal_card), **kw) + return SelfForcingMethod(new_instance(causal_card), new_instance(causal_card), new_instance(causal_card), **kw) diff --git a/v2/training/methods/unified_rl.py b/v2/training/methods/unified_rl.py index 197adb5d8..0e54fe526 100644 --- a/v2/training/methods/unified_rl.py +++ b/v2/training/methods/unified_rl.py @@ -1,30 +1,21 @@ -"""UnifiedRLMethod — joint LM + generator RL (UniRL / PromptRL §6; design_v3 §10). +"""UnifiedRLMethod — joint LM + generator RL (UniRL / PromptRL). -The stress test for the whole design. One reward, computed once per sample, drives TWO trainable -experts under one group-relative advantage: +One reward, computed once per sample, drives TWO trainable experts under one group-relative advantage: - * the ``llm`` prompt-refiner — a **token policy gradient** (REINFORCE on the refinement action, + * the ``llm`` prompt-refiner — a token policy gradient (REINFORCE on the refinement action, ``−A·logπ(a) + β·KL(π‖π_ref)``); and - * the ``transformer`` flow generator — **FlowGRPO PPO** over the SDE rollout's per-step log-probs + * the ``transformer`` flow generator — FlowGRPO PPO over the SDE rollout's per-step log-probs (``max(−A·ρ, −A·clip(ρ, 1±ε)) + β·KL`` with ``ρ = exp(logp_cur − logp_rollout)``). -What it exercises that nothing else in v2 does: - 1. a training rollout that spans **two loop types** — the ``ar_decode`` refine *and* the - ``diffusion_denoise`` rollout — not one; - 2. **dual log-prob capture** — a categorical action log-prob (LM) and a Gaussian per-step - log-prob (SDE diffusion) — feeding one advantage; - 3. **likelihood-based C2** — per-step log-prob identity (the OTHER half of the C2 split; NFT is - likelihood-*free*); - 4. a **multi-component joint update** — two WeightSyncPlans (one per expert), versioned and - cache-invalidated independently (the LM sync must not flush the text-encoder feature cache); - 5. the rollout SDE sampler ≠ the serve ODE sampler — a controlled §9.4 divergence, gated by - ``DiffusionParams.sde_rollout`` so the serve path is untouched. +Notable: the training rollout spans two loop types (``ar_decode`` refine + ``diffusion_denoise``); +dual log-prob capture (categorical LM action + Gaussian per-step SDE) feeds one advantage; the two +experts use independent WeightSyncPlans versioned/cache-invalidated separately (the LM sync must not +flush the text-encoder feature cache); the rollout SDE sampler differs from the serve ODE sampler, a +controlled divergence gated by ``DiffusionParams.sde_rollout`` so the serve path is untouched. -``joint=False`` is PromptRL's *prompt-only* mode: the generator is frozen (ODE rollout, no DiT -update); only the LM learns. ``num_skip_refinement`` samples per group use the *original* prompt -(partial refinement) so the advantage contrasts refined vs unrefined — the LM's learning signal. - -This is a stress test, not a redesign: it adds a method and a card, no new loop/runtime primitive. +``joint=False`` is PromptRL's prompt-only mode: the generator is frozen (ODE rollout, no DiT update), +only the LM learns. ``num_skip_refinement`` samples per group use the original prompt (partial +refinement) so the advantage contrasts refined vs unrefined — the LM's learning signal. """ from __future__ import annotations @@ -42,14 +33,26 @@ from v2.training.methods.base import TrainingMethod, new_instance class UnifiedRLMethod(TrainingMethod): name = "unified_rl" - consistency = ConsistencyLevel.C2 # likelihood-BASED (per-step log-prob identity) + consistency = ConsistencyLevel.C2 # likelihood-BASED (per-step log-prob identity) student_loop_id = "diffusion_denoise" - def __init__(self, student_instance, reference_instance, *, joint: bool = True, - num_samples_per_prompt: int = 4, num_skip_refinement: int = 1, rollout_steps: int = 4, - llm_lr: float = 0.4, dit_lr: float = 0.02, beta: float = 0.01, adv_clip_max: float = 5.0, - ppo_clip: float = 1e-2, sde_noise_scale: float = 0.7, target_action: int = 3, - action_bonus: float = 2.0, **kw): + def __init__(self, + student_instance, + reference_instance, + *, + joint: bool = True, + num_samples_per_prompt: int = 4, + num_skip_refinement: int = 1, + rollout_steps: int = 4, + llm_lr: float = 0.4, + dit_lr: float = 0.02, + beta: float = 0.01, + adv_clip_max: float = 5.0, + ppo_clip: float = 1e-2, + sde_noise_scale: float = 0.7, + target_action: int = 3, + action_bonus: float = 2.0, + **kw): super().__init__(student_instance, lr=dit_lr, **kw) self.reference = reference_instance self.joint = joint @@ -60,9 +63,9 @@ class UnifiedRLMethod(TrainingMethod): self.beta, self.adv_clip_max, self.ppo_clip = beta, adv_clip_max, ppo_clip self.sde_noise_scale = sde_noise_scale self.target_action, self.action_bonus = target_action, action_bonus - # one plan PER trainable expert — versioned + cache-scoped independently (design_v3 §7.1, §10) - self.llm_sync = WeightSyncPlan(role=WeightRole.STUDENT, components=("llm",)) - self.dit_sync = WeightSyncPlan(role=WeightRole.STUDENT, components=("transformer",)) + # one plan PER trainable expert — versioned + cache-scoped independently + self.llm_sync = WeightSyncPlan(role=WeightRole.STUDENT, components=("llm", )) + self.dit_sync = WeightSyncPlan(role=WeightRole.STUDENT, components=("transformer", )) # reference = frozen copy of BOTH experts at init (the two KL anchors) self.reference.component("llm").copy_from(self.llm) self.reference.component("transformer").copy_from(self.dit) @@ -99,14 +102,19 @@ class UnifiedRLMethod(TrainingMethod): # guidance_scale=1 ⇒ the CFG combine collapses to the cond branch, so the PPO recompute # (a single cond forward) replays the rollout velocity exactly — the likelihood-based C2 # identity (FLUX is guidance-distilled: one effective forward, no CFG at RL time). - req = make_request(TaskType.T2V, self.student.card.model_id, prompt, - diffusion=DiffusionParams(num_steps=self.rollout_steps, seed=seed, - guidance_scale=1.0, sde_rollout=self.joint, + req = make_request(TaskType.T2V, + self.student.card.model_id, + prompt, + diffusion=DiffusionParams(num_steps=self.rollout_steps, + seed=seed, + guidance_scale=1.0, + sde_rollout=self.joint, sde_noise_scale=self.sde_noise_scale)) - slots = {"text_embeds": np.asarray(emb, dtype="float32"), - "neg_text_embeds": np.asarray(neg_emb, dtype="float32")} - return rollout_loop(self.student, self.student_loop_id, req, slots=slots, - profile=ExecutionProfile.ROLLOUT) + slots = { + "text_embeds": np.asarray(emb, dtype="float32"), + "neg_text_embeds": np.asarray(neg_emb, dtype="float32") + } + return rollout_loop(self.student, self.student_loop_id, req, slots=slots, profile=ExecutionProfile.ROLLOUT) def _dit_ppo(self, behavior, emb, advantage): """FlowGRPO PPO over one rollout's SDE trajectory: per-step ratio + KL, then a toy DiT step. @@ -122,42 +130,44 @@ class UnifiedRLMethod(TrainingMethod): prev, sample = rec["prev"], rec["sample"] st, sn = float(rec["sigma_t"]), float(rec["sigma_next"]) v_cur = self.dit(prev, emb, st) - _, logp_cur, mean_cur, eff_std = sde_step( - prev, v_cur, st, sn, prev_sample=sample, noise_scale=self.sde_noise_scale) + _, logp_cur, mean_cur, eff_std = sde_step(prev, + v_cur, + st, + sn, + prev_sample=sample, + noise_scale=self.sde_noise_scale) v_ref = ref_dit(prev, emb, st) - _, _, mean_ref, _ = sde_step( - prev, v_ref, st, sn, prev_sample=sample, noise_scale=self.sde_noise_scale) + _, _, mean_ref, _ = sde_step(prev, v_ref, st, sn, prev_sample=sample, noise_scale=self.sde_noise_scale) ratio = float(np.exp(np.clip(logp_cur - float(rec["sde_logprob"]), -10.0, 10.0))) clipped = float(np.clip(ratio, 1.0 - self.ppo_clip, 1.0 + self.ppo_clip)) - ppo = max(-advantage * ratio, -advantage * clipped) # FlowGRPO PPO (per element) - kl = float(np.mean((mean_cur - mean_ref) ** 2) / (2.0 * max(eff_std ** 2, 1e-8))) + ppo = max(-advantage * ratio, -advantage * clipped) # FlowGRPO PPO (per element) + kl = float(np.mean((mean_cur - mean_ref)**2) / (2.0 * max(eff_std**2, 1e-8))) ppo_losses.append(ppo) kls.append(kl) ratios.append(ratio) v_target = _ml_velocity(prev, sample, st, sn, noise_scale=self.sde_noise_scale) # PG direction - _, gn = self.dit.mse_grad_step(prev, emb, st, v_target, - self.dit_lr * float(np.clip(advantage, -1.0, 1.0))) + _, gn = self.dit.mse_grad_step(prev, emb, st, v_target, self.dit_lr * float(np.clip(advantage, -1.0, 1.0))) gnorms.append(gn) - m = lambda xs: float(np.mean(xs)) if xs else 0.0 # noqa: E731 + m = lambda xs: float(np.mean(xs)) if xs else 0.0 # noqa: E731 return m(ppo_losses), m(kls), m(ratios), m(gnorms) def managed_train_step(self, batch: dict, iteration: int) -> tuple[dict, dict]: self.iteration = iteration prompts = batch["prompts"] seeds = batch.get("seeds", list(range(len(prompts)))) - neg_base = cached_text_encode(self.student, "") # shared negative embed + neg_base = cached_text_encode(self.student, "") # shared negative embed lm_losses, dit_losses, kls, ratios = [], [], [], [] lm_gnorms, dit_gnorms, all_reward, all_adv = [], [], [], [] for pi, prompt in enumerate(prompts): - base_emb = cached_text_encode(self.student, prompt) # encoded ONCE for the group + base_emb = cached_text_encode(self.student, prompt) # encoded ONCE for the group rng = np.random.default_rng((seeds[pi] + 1) * 7919 + iteration) - action, logp_action = self.llm.sample_refinement(rng) # the LM's RL action (sampled) + action, logp_action = self.llm.sample_refinement(rng) # the LM's RL action (sampled) refined_emb = self.llm.refined_embed(base_emb, action) - samples = [] # (latent, behavior, emb, reward, refined?) + samples = [] # (latent, behavior, emb, reward, refined?) for k in range(self.K): - refined = k >= self.num_skip # partial refinement + refined = k >= self.num_skip # partial refinement emb = refined_emb if refined else base_emb act = action if refined else None res = self._rollout_sample(prompt, emb, neg_base, seed=seeds[pi] * 1000 + k) @@ -166,8 +176,8 @@ class UnifiedRLMethod(TrainingMethod): samples.append((latent, res.behavior, emb, reward, refined)) rewards = np.array([s[3] for s in samples], dtype="float64") - adv = np.clip((rewards - rewards.mean()) / (rewards.std() + 1e-4), - -self.adv_clip_max, self.adv_clip_max) # group-relative advantage + adv = np.clip((rewards - rewards.mean()) / (rewards.std() + 1e-4), -self.adv_clip_max, + self.adv_clip_max) # group-relative advantage all_reward.extend(rewards.tolist()) all_adv.extend(adv.tolist()) @@ -175,7 +185,7 @@ class UnifiedRLMethod(TrainingMethod): refined_adv = [adv[k] for k in range(self.K) if samples[k][4]] lm_adv = float(np.mean(refined_adv)) if refined_adv else 0.0 lm_kl = self.llm.kl_to(self.reference.component("llm")) - gn_lm = self.llm.pg_step(action, lm_adv, self.llm_lr) # REAL REINFORCE step + gn_lm = self.llm.pg_step(action, lm_adv, self.llm_lr) # REAL REINFORCE step lm_losses.append(-lm_adv * logp_action + self.beta * lm_kl) lm_gnorms.append(gn_lm) @@ -190,11 +200,11 @@ class UnifiedRLMethod(TrainingMethod): ratios.append(ratio) dit_gnorms.append(gn) - # publish both experts' weights — independent versions + cache scopes (design_v3 §10) + # publish both experts' weights — independent versions + cache scopes lm_ver = self.llm_sync.apply(self.llm, self.llm, self.student) dit_ver = self.dit_sync.apply(self.dit, self.dit, self.student) if self.joint else "frozen" - p_target = float(self.llm._probs()[self.target_action]) # LM learning signal + p_target = float(self.llm._probs()[self.target_action]) # LM learning signal metrics = { "lm_pg_loss": float(np.mean(lm_losses)), "grad_norm/llm": float(np.mean(lm_gnorms)), diff --git a/v2/training/methods/workflow_rl.py b/v2/training/methods/workflow_rl.py index c81e53fd3..9da0a80a1 100644 --- a/v2/training/methods/workflow_rl.py +++ b/v2/training/methods/workflow_rl.py @@ -1,18 +1,16 @@ -"""WorkflowRLMethod — end-to-end RL over a cross-model workflow (design_v3 §10 + §13). +"""WorkflowRLMethod — end-to-end RL over a cross-model workflow. -The stress test of the training-plane / Workflow boundary: train BOTH stages of the T2I→I2V workflow -(``flux-t2i`` and ``wan-i2v`` — two *separate* models) from **one final-video reward**. The rollout -spans two instances: roll out the T2I diffusion loop with SDE capture → decode an image → condition -and roll out the I2V diffusion loop with SDE capture → decode a video → score it. One group-relative -advantage from the final video then drives a **FlowGRPO PPO update on each stage's transformer**, with -**two independent WeightSyncPlans on two different instances**. +Train BOTH stages of the T2I→I2V workflow (``flux-t2i`` and ``wan-i2v`` — two separate models) from one +final-video reward. The rollout spans two instances: roll out the T2I diffusion loop with SDE capture → +decode an image → condition and roll out the I2V diffusion loop with SDE capture → decode a video → score +it. One group-relative advantage from the final video drives a FlowGRPO PPO update on each stage's +transformer, with two independent WeightSyncPlans on two different instances. -The interesting part is the credit assignment: the reward is at the *end* (the video), yet it credits an -*earlier* model (the T2I image generator) — end-to-end RL across a model boundary. The T2I stage learns -to produce images that lead to better videos. This probes whether "rollout == serve + capture" holds for -a *workflow* (multi-instance), not just a single card's loops — and it does, with no new primitive: the -method drives the same shared loops the engine serves, on each instance, and threads the artifact between -them exactly as the serve-time Workflow does. +The interesting part is credit assignment: the reward is at the end (the video) yet credits an earlier +model (the T2I generator), so the T2I stage learns to produce images that lead to better videos. This +shows "rollout == serve + capture" holds for a multi-instance workflow with no new primitive — the method +drives the same shared loops the engine serves and threads the artifact between them as the serve-time +Workflow does. """ from __future__ import annotations @@ -28,8 +26,7 @@ from v2.training.weight_sync import WeightRole, WeightSyncPlan from v2.training.methods.base import TrainingMethod, new_instance -def _flowgrpo_ppo_step(dit, ref_dit, behavior, emb, advantage, *, lr, ppo_clip, beta, noise_scale, - sde_step=None): +def _flowgrpo_ppo_step(dit, ref_dit, behavior, emb, advantage, *, lr, ppo_clip, beta, noise_scale, sde_step=None): """One FlowGRPO PPO pass over an SDE trajectory: per-step ratio + KL, then move the policy toward the max-likelihood velocity of each realized sample (advantage-weighted). Returns (ppo_loss, kl, ratio, grad_norm). Shared by both stages. ``sde_step`` is the platform-resolved SDE @@ -42,20 +39,18 @@ def _flowgrpo_ppo_step(dit, ref_dit, behavior, emb, advantage, *, lr, ppo_clip, prev, sample = rec["prev"], rec["sample"] st, sn = float(rec["sigma_t"]), float(rec["sigma_next"]) v_cur = dit(prev, emb, st) - _, logp_cur, mean_cur, eff_std = sde_step( - prev, v_cur, st, sn, prev_sample=sample, noise_scale=noise_scale) + _, logp_cur, mean_cur, eff_std = sde_step(prev, v_cur, st, sn, prev_sample=sample, noise_scale=noise_scale) v_ref = ref_dit(prev, emb, st) - _, _, mean_ref, _ = sde_step( - prev, v_ref, st, sn, prev_sample=sample, noise_scale=noise_scale) + _, _, mean_ref, _ = sde_step(prev, v_ref, st, sn, prev_sample=sample, noise_scale=noise_scale) ratio = float(np.exp(np.clip(logp_cur - float(rec["sde_logprob"]), -10.0, 10.0))) clipped = float(np.clip(ratio, 1.0 - ppo_clip, 1.0 + ppo_clip)) ppo_losses.append(max(-advantage * ratio, -advantage * clipped)) - kls.append(float(np.mean((mean_cur - mean_ref) ** 2) / (2.0 * max(eff_std ** 2, 1e-8)))) + kls.append(float(np.mean((mean_cur - mean_ref)**2) / (2.0 * max(eff_std**2, 1e-8)))) ratios.append(ratio) v_target = flow_sde_ml_velocity(prev, sample, st, sn, noise_scale=noise_scale) _, gn = dit.mse_grad_step(prev, emb, st, v_target, lr * float(np.clip(advantage, -1.0, 1.0))) gnorms.append(gn) - m = lambda xs: float(np.mean(xs)) if xs else 0.0 # noqa: E731 + m = lambda xs: float(np.mean(xs)) if xs else 0.0 # noqa: E731 return m(ppo_losses), m(kls), m(ratios), m(gnorms) @@ -63,11 +58,22 @@ class WorkflowRLMethod(TrainingMethod): name = "workflow_rl" consistency = ConsistencyLevel.C2 - def __init__(self, t2i_student, i2v_student, t2i_reference, i2v_reference, *, - num_samples_per_prompt: int = 4, rollout_steps: int = 4, t2i_lr: float = 0.02, - i2v_lr: float = 0.02, beta: float = 0.01, adv_clip_max: float = 5.0, - ppo_clip: float = 1e-2, sde_noise_scale: float = 0.7, **kw): - super().__init__(t2i_student, lr=t2i_lr, **kw) # base "student" = the T2I stage + def __init__(self, + t2i_student, + i2v_student, + t2i_reference, + i2v_reference, + *, + num_samples_per_prompt: int = 4, + rollout_steps: int = 4, + t2i_lr: float = 0.02, + i2v_lr: float = 0.02, + beta: float = 0.01, + adv_clip_max: float = 5.0, + ppo_clip: float = 1e-2, + sde_noise_scale: float = 0.7, + **kw): + super().__init__(t2i_student, lr=t2i_lr, **kw) # base "student" = the T2I stage self.t2i, self.i2v = t2i_student, i2v_student self.t2i_ref, self.i2v_ref = t2i_reference, i2v_reference self.K = num_samples_per_prompt @@ -76,8 +82,8 @@ class WorkflowRLMethod(TrainingMethod): self.beta, self.adv_clip_max, self.ppo_clip = beta, adv_clip_max, ppo_clip self.sde_noise_scale = sde_noise_scale # one WeightSyncPlan per stage's generator — on two DIFFERENT instances (independent versions) - self.t2i_sync = WeightSyncPlan(role=WeightRole.STUDENT, components=("transformer",)) - self.i2v_sync = WeightSyncPlan(role=WeightRole.STUDENT, components=("transformer",)) + self.t2i_sync = WeightSyncPlan(role=WeightRole.STUDENT, components=("transformer", )) + self.i2v_sync = WeightSyncPlan(role=WeightRole.STUDENT, components=("transformer", )) self.t2i_ref.component("transformer").copy_from(self.t2i.component("transformer")) self.i2v_ref.component("transformer").copy_from(self.i2v.component("transformer")) @@ -85,20 +91,32 @@ class WorkflowRLMethod(TrainingMethod): return True def get_grad_clip_targets(self, iteration: int = 0) -> dict: - return {"t2i.transformer": self.t2i.component("transformer"), - "i2v.transformer": self.i2v.component("transformer")} + return { + "t2i.transformer": self.t2i.component("transformer"), + "i2v.transformer": self.i2v.component("transformer") + } def _sde_params(self, **extra): - return DiffusionParams(num_steps=self.rollout_steps, guidance_scale=1.0, sde_rollout=True, - sde_noise_scale=self.sde_noise_scale, **extra) + return DiffusionParams(num_steps=self.rollout_steps, + guidance_scale=1.0, + sde_rollout=True, + sde_noise_scale=self.sde_noise_scale, + **extra) def _rollout_t2i(self, prompt, seed): emb = cached_text_encode(self.t2i, prompt) neg = cached_text_encode(self.t2i, "") - req = make_request(TaskType.T2I, self.t2i.card.model_id, prompt, + req = make_request(TaskType.T2I, + self.t2i.card.model_id, + prompt, diffusion=self._sde_params(num_frames=1, seed=seed)) - res = rollout_loop(self.t2i, "t2i_denoise", req, - slots={"text_embeds": emb, "neg_text_embeds": neg}, + res = rollout_loop(self.t2i, + "t2i_denoise", + req, + slots={ + "text_embeds": emb, + "neg_text_embeds": neg + }, profile=ExecutionProfile.ROLLOUT) image = self.t2i.component("vae").decode(res.outputs["latents"]) return res, emb, image @@ -106,19 +124,26 @@ class WorkflowRLMethod(TrainingMethod): def _rollout_i2v(self, prompt, image, seed): emb = cached_text_encode(self.i2v, prompt) neg = cached_text_encode(self.i2v, "") - cond = self.i2v.component("vae").encode(np.asarray(image, dtype="float32")) # image conditioning + cond = self.i2v.component("vae").encode(np.asarray(image, dtype="float32")) # image conditioning emb = (np.asarray(emb, dtype="float32") + float(np.tanh(np.mean(cond)))).astype("float32") - req = make_request(TaskType.I2V, self.i2v.card.model_id, prompt, + req = make_request(TaskType.I2V, + self.i2v.card.model_id, + prompt, diffusion=self._sde_params(num_frames=81, seed=seed)) - res = rollout_loop(self.i2v, "i2v_denoise", req, - slots={"text_embeds": emb, "neg_text_embeds": neg}, + res = rollout_loop(self.i2v, + "i2v_denoise", + req, + slots={ + "text_embeds": emb, + "neg_text_embeds": neg + }, profile=ExecutionProfile.ROLLOUT) video = self.i2v.component("vae").decode(res.outputs["latents"]) return res, emb, video @staticmethod def _reward(video): - return float(np.tanh(-np.std(np.asarray(video, dtype="float64")))) # lower-spread video = better + return float(np.tanh(-np.std(np.asarray(video, dtype="float64")))) # lower-spread video = better def managed_train_step(self, batch: dict, iteration: int) -> tuple[dict, dict]: self.iteration = iteration @@ -130,37 +155,47 @@ class WorkflowRLMethod(TrainingMethod): samples = [] for k in range(self.K): seed = seeds[pi] * 1000 + k - res_t, emb_t, image = self._rollout_t2i(prompt, seed) # stage 1 rollout (SDE) - res_i, emb_i, video = self._rollout_i2v(prompt, image, seed) # stage 2 rollout (SDE) + res_t, emb_t, image = self._rollout_t2i(prompt, seed) # stage 1 rollout (SDE) + res_i, emb_i, video = self._rollout_i2v(prompt, image, seed) # stage 2 rollout (SDE) samples.append((res_t, emb_t, res_i, emb_i, self._reward(video))) rewards = np.array([s[4] for s in samples], dtype="float64") - adv = np.clip((rewards - rewards.mean()) / (rewards.std() + 1e-4), - -self.adv_clip_max, self.adv_clip_max) # one final-video advantage + adv = np.clip((rewards - rewards.mean()) / (rewards.std() + 1e-4), -self.adv_clip_max, + self.adv_clip_max) # one final-video advantage all_reward.extend(rewards.tolist()) for k, (res_t, emb_t, res_i, emb_i, _r) in enumerate(samples): a = float(adv[k]) # the SAME final advantage credits BOTH stages (end-to-end across the model boundary) - pt, _kt, rt, gt = _flowgrpo_ppo_step( - self.t2i.component("transformer"), self.t2i_ref.component("transformer"), - res_t.behavior, emb_t, a, lr=self.t2i_lr, ppo_clip=self.ppo_clip, - beta=self.beta, noise_scale=self.sde_noise_scale, - sde_step=self.t2i.platform.kernels.get(FLOW_SDE_STEP)) - pi_, _ki, ri, gi = _flowgrpo_ppo_step( - self.i2v.component("transformer"), self.i2v_ref.component("transformer"), - res_i.behavior, emb_i, a, lr=self.i2v_lr, ppo_clip=self.ppo_clip, - beta=self.beta, noise_scale=self.sde_noise_scale, - sde_step=self.i2v.platform.kernels.get(FLOW_SDE_STEP)) - t2i_loss.append(pt); i2v_loss.append(pi_) - t2i_gn.append(gt); i2v_gn.append(gi) + pt, _kt, rt, gt = _flowgrpo_ppo_step(self.t2i.component("transformer"), + self.t2i_ref.component("transformer"), + res_t.behavior, + emb_t, + a, + lr=self.t2i_lr, + ppo_clip=self.ppo_clip, + beta=self.beta, + noise_scale=self.sde_noise_scale, + sde_step=self.t2i.platform.kernels.get(FLOW_SDE_STEP)) + pi_, _ki, ri, gi = _flowgrpo_ppo_step(self.i2v.component("transformer"), + self.i2v_ref.component("transformer"), + res_i.behavior, + emb_i, + a, + lr=self.i2v_lr, + ppo_clip=self.ppo_clip, + beta=self.beta, + noise_scale=self.sde_noise_scale, + sde_step=self.i2v.platform.kernels.get(FLOW_SDE_STEP)) + t2i_loss.append(pt) + i2v_loss.append(pi_) + t2i_gn.append(gt) + i2v_gn.append(gi) ratios.extend([rt, ri]) - t2i_ver = self.t2i_sync.apply(self.t2i.component("transformer"), - self.t2i.component("transformer"), self.t2i) - i2v_ver = self.i2v_sync.apply(self.i2v.component("transformer"), - self.i2v.component("transformer"), self.i2v) - m = lambda xs: float(np.mean(xs)) if xs else 0.0 # noqa: E731 + t2i_ver = self.t2i_sync.apply(self.t2i.component("transformer"), self.t2i.component("transformer"), self.t2i) + i2v_ver = self.i2v_sync.apply(self.i2v.component("transformer"), self.i2v.component("transformer"), self.i2v) + m = lambda xs: float(np.mean(xs)) if xs else 0.0 # noqa: E731 metrics = { "reward_mean": m(all_reward), "grad_norm/t2i": m(t2i_gn), @@ -177,5 +212,5 @@ class WorkflowRLMethod(TrainingMethod): def build_workflow_rl(t2i_card, i2v_card, **kw) -> WorkflowRLMethod: """Four roles: trainable T2I + I2V students and their frozen references (two instances each side).""" - return WorkflowRLMethod(new_instance(t2i_card), new_instance(i2v_card), - new_instance(t2i_card), new_instance(i2v_card), **kw) + return WorkflowRLMethod(new_instance(t2i_card), new_instance(i2v_card), new_instance(t2i_card), + new_instance(i2v_card), **kw) diff --git a/v2/training/rewards.py b/v2/training/rewards.py index bab9c8224..1facd4b6e 100644 --- a/v2/training/rewards.py +++ b/v2/training/rewards.py @@ -1,14 +1,12 @@ -"""Reward layer (design_v3 §10; verl-omni's shape): weighted multi-reward, duck-typed scorers. - -> Landed ``train/methods/rl/rewards/`` is the right seed: thin PickScore/CLIP scorers, duck-typed -> ``RewardScorer = Callable[[media, prompts], Tensor]``, weighted aggregation in MultiRewardScorer. +"""Reward layer: weighted multi-reward over duck-typed scorers +(``RewardScorer = Callable[[media, prompts], Tensor]``). Toy scorers here are deterministic functions of (media, prompt) — enough to exercise group-relative -advantages and the reward→advantage→update loop without a real reward model. +advantages and the reward->advantage->update loop without a real reward model. """ from __future__ import annotations -from typing import Callable +from collections.abc import Callable import numpy as np @@ -21,7 +19,7 @@ RewardScorer = Callable[[list, list], np.ndarray] def toy_pickscore(media: list, prompts: list) -> np.ndarray: """Deterministic 'preference' score: blends a prompt-seeded target with media statistics.""" out = [] - for m, p in zip(media, prompts): + for m, p in zip(media, prompts, strict=False): target = np.random.default_rng(int(content_hash(p)[:8], 16)).standard_normal(1)[0] score = -abs(float(np.mean(np.asarray(m))) - 0.1 * target) out.append(score) @@ -37,7 +35,7 @@ _REGISTRY: dict[str, RewardScorer] = {"pickscore": toy_pickscore, "clipscore": t class MultiRewardScorer: - """Weighted aggregation of named scorers (design_v3 §10).""" + """Weighted aggregation of named scorers.""" def __init__(self, weights: dict[str, float]): self.weights = weights @@ -55,7 +53,7 @@ def build_multi_reward_scorer(config: dict[str, float]) -> MultiRewardScorer: class ServedRewardScorer: - """A reward scorer backed by a SERVED reward-model card (design_v3 §10). + """A reward scorer backed by a served reward-model card. Drop-in for ``MultiRewardScorer`` (same ``score(media, prompts) -> {"avg": ...}`` interface), but instead of a numpy heuristic it runs a reward MODEL through its ``score`` loop — so the K rollout @@ -70,8 +68,7 @@ class ServedRewardScorer: from v2._enums import ExecutionProfile from v2.request import TaskType, make_request from v2.training.rollout import rollout_loop - req = make_request(TaskType.REASON, self.inst.card.model_id, "") # task is cosmetic here - res = rollout_loop(self.inst, self.loop_id, req, slots={"media": list(media)}, - profile=ExecutionProfile.SERVE) + req = make_request(TaskType.REASON, self.inst.card.model_id, "") # task is cosmetic here + res = rollout_loop(self.inst, self.loop_id, req, slots={"media": list(media)}, profile=ExecutionProfile.SERVE) rewards = np.asarray(res.outputs["rewards"], dtype="float64") return {"served": rewards, "avg": rewards} diff --git a/v2/training/rollout.py b/v2/training/rollout.py index 3bd41917d..9404dae9e 100644 --- a/v2/training/rollout.py +++ b/v2/training/rollout.py @@ -1,16 +1,15 @@ -"""Rollout = serve loop + capture (design_v3 §10). The moat, made literal. +"""Rollout = serve loop + capture. ``rollout_loop`` drives the *same* ``Loop`` the engine serves, in the ROLLOUT execution profile, -through a minimal training context — so there is one numerics surface, not a second sampler. This -is the structural answer to the two-runtime tax: the landed RL stack vendors a bare-model loop -(``rl/common/sampling.py``); here the method drives the shared denoise/chunk loop directly. +through a minimal training context — so there is one numerics surface, not a second sampler. -``training`` imports ``loop``/``runtime`` here; the engine never imports ``training`` (the §10 rule). +``training`` imports ``loop``/``runtime`` here; the engine never imports ``training`` (one-way rule). """ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any, Callable +from typing import Any +from collections.abc import Callable from v2._enums import ExecutionProfile from v2.loop.contracts import StepResult, WorkPlan @@ -22,8 +21,7 @@ class _TrainingContext: """A minimal LoopContext for in-process (library-mode) rollout. Runs the kernel thunk directly; optional observers attach (e.g. ParityAligner for the consistency ladder).""" - def __init__(self, instance: Any, slots: dict[str, Any], profile: ExecutionProfile, - observers: list | None = None): + def __init__(self, instance: Any, slots: dict[str, Any], profile: ExecutionProfile, observers: list | None = None): self.instance = instance self.slots = slots self.profile = profile @@ -53,8 +51,13 @@ class _TrainingContext: o.observe(event, **kw) -def rollout_loop(instance: Any, loop_id: str, request: Any, *, slots: dict[str, Any] | None = None, - profile: ExecutionProfile = ExecutionProfile.ROLLOUT, observers: list | None = None): +def rollout_loop(instance: Any, + loop_id: str, + request: Any, + *, + slots: dict[str, Any] | None = None, + profile: ExecutionProfile = ExecutionProfile.ROLLOUT, + observers: list | None = None): """Drive the SHARED loop to completion in the given profile; return its LoopResult.""" ctx = _TrainingContext(instance, slots or {}, profile, observers) runner = LoopRunner(instance.loop(loop_id), ctx, request, instance) @@ -63,7 +66,7 @@ def rollout_loop(instance: Any, loop_id: str, request: Any, *, slots: dict[str, @dataclass class Trajectory: - """Our ``Sample`` (design_v3 §10): request + seeds + per-step records + reward + advantage.""" + """One rollout sample: request + seed + per-step records + reward + advantage.""" request_id: str prompt: str seed: int @@ -77,6 +80,7 @@ class Trajectory: class TrajectoryBuffer: + def __init__(self) -> None: self.items: list[Trajectory] = [] @@ -98,7 +102,7 @@ class TrajectoryBuffer: @dataclass class RolloutContext: - """What a ``RolloutFn`` receives (design_v3 §8.7): an engine/instance client + a seed source.""" + """What a ``RolloutFn`` receives: an engine/instance client + a seed source.""" instance: Any text_encode: Callable[[Any, Any, dict], None] base_seed: int = 0 diff --git a/v2/training/weight_sync.py b/v2/training/weight_sync.py index 01fbc56bc..60f6a604b 100644 --- a/v2/training/weight_sync.py +++ b/v2/training/weight_sync.py @@ -1,12 +1,12 @@ -"""WeightSyncPlan — versioned weight transfer with a declared role (design_v3 §10). +"""WeightSyncPlan — versioned weight transfer with a declared role. -> RL ships a *role*, not "the weights": student / EMA / decay-blended old policy is declared -> (the landed NFT behavior policy is the *old* copy, not the student — the plan must carry that). +A sync ships a *role*, not just "the weights": student / EMA / decay-blended old policy is declared +(e.g. an NFT behavior policy is the *old* copy, not the student — the plan must carry that). -The lifecycle (design_v3 §10): freeze admission → drain/boundary-stop → transfer → bump -``weights_version`` → invalidate incompatible caches/graphs → publish → resume. Computed from -three inputs (mesh specs + per-model layout adapters + transport), validated pre-flight, -CPU-testable on fake pools. Here the "transfer" copies/blends the toy DiT params. +Lifecycle: freeze admission -> drain/boundary-stop -> transfer -> bump ``weights_version`` -> +invalidate incompatible caches/graphs -> publish -> resume. Computed from three inputs (mesh specs ++ per-model layout adapters + transport), validated pre-flight, CPU-testable on fake pools. Here the +"transfer" copies/blends the toy DiT params. """ from __future__ import annotations @@ -17,12 +17,12 @@ from typing import Any class WeightRole(str, Enum): - STUDENT = "student" # the current learnable weights - EMA = "ema" # exponential moving average of student - OLD_POLICY = "old_policy" # decay-blended behavior policy (NFT samples from THIS, not student) - REFERENCE = "reference" # frozen KL anchor - TEACHER = "teacher" # frozen distillation teacher - CRITIC = "critic" # trainable fake-score / value head + STUDENT = "student" # the current learnable weights + EMA = "ema" # exponential moving average of student + OLD_POLICY = "old_policy" # decay-blended behavior policy (NFT samples from THIS, not student) + REFERENCE = "reference" # frozen KL anchor + TEACHER = "teacher" # frozen distillation teacher + CRITIC = "critic" # trainable fake-score / value head _ver = itertools.count(1) @@ -31,27 +31,27 @@ _ver = itertools.count(1) @dataclass class WeightSyncPlan: role: WeightRole - # three inputs (design_v3 §10): mesh specs (trainer↔engine), layout adapters, transport + # three inputs: mesh specs (trainer<->engine), layout adapters, transport mesh_specs: dict[str, Any] = field(default_factory=dict) layout_adapters: dict[str, Any] = field(default_factory=dict) - transport: str = "in_proc" # in_proc | shm | nccl | rdma - decay: float = 0.0 # for EMA / OLD_POLICY blend - # which component(s) this plan syncs — versioned + cache-invalidated in isolation (design_v3 §7.1). + transport: str = "in_proc" # in_proc | shm | nccl | rdma + decay: float = 0.0 # for EMA / OLD_POLICY blend + # which component(s) this plan syncs — versioned + cache-invalidated in isolation. # Joint RL (UniRL) ships ONE plan per trainable expert: ``("transformer",)`` and ``("llm",)`` — # so the LM weight sync does NOT flush the frozen text-encoder's feature cache, and vice-versa. - components: tuple[str, ...] = ("transformer",) + components: tuple[str, ...] = ("transformer", ) def apply(self, src_dit: Any, dst_dit: Any, dst_instance: Any = None) -> str: """Execute the lifecycle. Returns the new weights_version published on dst_instance.""" # 1) freeze admission + 2) drain/boundary-stop are scheduler ops (no-op in single-thread mini) # 3) transfer per role if self.role in (WeightRole.EMA, WeightRole.OLD_POLICY): - dst_dit.blend_from(src_dit, self.decay) # decay·dst + (1-decay)·src + dst_dit.blend_from(src_dit, self.decay) # decay·dst + (1-decay)·src else: - dst_dit.copy_from(src_dit) # hard copy (student push / reference init) + dst_dit.copy_from(src_dit) # hard copy (student push / reference init) # 4) bump version + 5) invalidate ONLY this plan's component caches/graphs + 6) publish. # Scoped to ``self.components`` so frozen text/vision encoders keep their feature caches, - # and the two experts of a joint-RL update version independently (design_v3 §7.1, §10). + # and the two experts of a joint-RL update version independently. version = f"w{next(_ver)}" if dst_instance is not None: dst_instance.set_weights_version(version, components=list(self.components)) @@ -59,11 +59,10 @@ class WeightSyncPlan: class WeightSyncController: - """The §10 hot weight-sync lifecycle over a *live, serving* resident instance, made explicit and - testable — the RL flywheel's hardest correctness surface (rollout while serving, swap weights, stay - correct): + """Hot weight-sync lifecycle over a *live, serving* resident instance — the RL flywheel's hardest + correctness surface (rollout while serving, swap weights, stay correct): - freeze() admission → (drain in-flight loops to a boundary) → sync() transfer+publish → resume. + freeze() admission -> (drain in-flight loops to a boundary) -> sync() transfer+publish -> resume. The load-bearing invariant: an in-flight request finishes on the weights it **started** with — the sync waits for it to drain, so its trajectory is never a half-and-half of two policies, and its @@ -82,14 +81,14 @@ class WeightSyncController: self.frozen = True def can_admit(self) -> bool: - return not self.frozen # admission gate: no new starts of the synced loop while frozen + return not self.frozen # admission gate: no new starts of the synced loop while frozen - def sync(self, src_dit: Any, *, component_id: str = "transformer", - role: WeightRole = WeightRole.STUDENT) -> str: + def sync(self, src_dit: Any, *, component_id: str = "transformer", role: WeightRole = WeightRole.STUDENT) -> str: """Transfer new weights into the resident component, bump its version, invalidate ONLY its caches, and resume admission. Call AFTER in-flight loops on this instance have drained.""" - version = WeightSyncPlan(role=role, components=(component_id,)).apply( - src_dit, self.instance.component(component_id), self.instance) + version = WeightSyncPlan(role=role, components=(component_id, )).apply(src_dit, + self.instance.component(component_id), + self.instance) self.frozen = False self.synced += 1 return version diff --git a/v2/transport/__init__.py b/v2/transport/__init__.py index f1fe07e28..7a1a11c11 100644 --- a/v2/transport/__init__.py +++ b/v2/transport/__init__.py @@ -1,4 +1,4 @@ -"""Transport plane — manifest-based, pluggable connectors with readiness + credit flow-control (§7.3).""" +"""Transport plane — manifest-based, pluggable connectors with readiness + credit flow-control.""" from __future__ import annotations from v2.transport.connector import ( @@ -12,5 +12,7 @@ from v2.transport.connector import ( ) from v2.transport.manifest import TransferManifest, payload_nbytes -__all__ = ["Connector", "CreditWindow", "InProcConnector", "ShmFakeConnector", "make_connector", - "KVConnector", "InProcKVConnector", "TransferManifest", "payload_nbytes"] +__all__ = [ + "Connector", "CreditWindow", "InProcConnector", "ShmFakeConnector", "make_connector", "KVConnector", + "InProcKVConnector", "TransferManifest", "payload_nbytes" +] diff --git a/v2/transport/connector.py b/v2/transport/connector.py index e249cb163..5b9fa7467 100644 --- a/v2/transport/connector.py +++ b/v2/transport/connector.py @@ -1,8 +1,8 @@ -"""Connectors + flow control (design_v3 §7.3; design.md §6.3.4). +"""Connectors + flow control. -Two mechanisms, both modeled (the design distinguishes them): - * ``chunk_ready`` readiness signal — vllm-omni's ``OmniConnectorBase`` (readiness-set parking); - * credit-based flow control — sglang-omni's Relay (a bounded in-flight window). +Two distinct mechanisms are modeled: + * ``chunk_ready`` readiness signal (readiness-set parking); + * credit-based flow control (a bounded in-flight window). Backends are pluggable: ``InProcConnector`` (zero-copy reference handoff — the single-node default) and ``ShmFakeConnector`` (copies on put to model a serialization boundary). A ``KVConnector`` exposes @@ -18,7 +18,7 @@ from v2.transport.manifest import TransferManifest, payload_nbytes class CreditWindow: - """Credit-based flow control (sglang-omni Relay): a bounded number of in-flight transfers.""" + """Credit-based flow control: a bounded number of in-flight transfers.""" def __init__(self, capacity: int = 8): self.capacity = capacity @@ -38,16 +38,27 @@ class CreditWindow: class Connector(Protocol): name: str - def put(self, key: str, value: Any, manifest: TransferManifest | None = None) -> None: ... - def chunk_ready(self, key: str) -> bool: ... - def get(self, key: str) -> Any: ... - def take(self, key: str) -> Any: ... - def acquire_credit(self, n: int = 1) -> bool: ... - def release_credit(self, n: int = 1) -> None: ... + def put(self, key: str, value: Any, manifest: TransferManifest | None = None) -> None: + ... + + def chunk_ready(self, key: str) -> bool: + ... + + def get(self, key: str) -> Any: + ... + + def take(self, key: str) -> Any: + ... + + def acquire_credit(self, n: int = 1) -> bool: + ... + + def release_credit(self, n: int = 1) -> None: + ... class InProcConnector: - """Zero-copy in-process handoff with readiness + credits (the single-node default, §6.3.4).""" + """Zero-copy in-process handoff with readiness + credits (the single-node default).""" def __init__(self, name: str = "inproc", credit_capacity: int = 8): self.name = name @@ -56,10 +67,10 @@ class InProcConnector: self.credits = CreditWindow(credit_capacity) self.transfers = 0 self.bytes_moved = 0 - self.parked = 0 # how many get()s waited on a not-yet-ready key (readiness parking) + self.parked = 0 # how many get()s waited on a not-yet-ready key (readiness parking) def _materialize(self, value: Any) -> Any: - return value # in-proc: pass by reference (zero overhead) + return value # in-proc: pass by reference (zero overhead) def put(self, key: str, value: Any, manifest: TransferManifest | None = None) -> None: self._store[key] = self._materialize(value) @@ -104,17 +115,27 @@ class ShmFakeConnector(InProcConnector): @runtime_checkable class KVConnector(Protocol): - """Cache-bearing edge protocol (design_v3 §7.3): scheduler-side query/alloc/finish + - worker-side save/load — the shape NIXL/LMCache/Mooncake/KVBM implement.""" + """Cache-bearing edge protocol: scheduler-side query/alloc/finish + worker-side save/load — + the shape NIXL/LMCache/Mooncake/KVBM implement.""" - def query(self, key: str) -> bool: ... - def alloc(self, key: str, nbytes: int) -> bool: ... - def finish(self, key: str) -> None: ... - def save(self, key: str, value: Any) -> None: ... - def load(self, key: str) -> Any: ... + def query(self, key: str) -> bool: + ... + + def alloc(self, key: str, nbytes: int) -> bool: + ... + + def finish(self, key: str) -> None: + ... + + def save(self, key: str, value: Any) -> None: + ... + + def load(self, key: str) -> Any: + ... class InProcKVConnector: + def __init__(self, capacity_bytes: int = 1 << 30): self.capacity_bytes = capacity_bytes self.used = 0 diff --git a/v2/transport/manifest.py b/v2/transport/manifest.py index da3f8462a..154f53086 100644 --- a/v2/transport/manifest.py +++ b/v2/transport/manifest.py @@ -1,11 +1,11 @@ -"""TransferManifest — typed description of a cross-pool tensor/artifact move (design_v3 §7.3; designv2). +"""TransferManifest — typed description of a cross-pool tensor/artifact move. Transfers are scheduled WorkUnits, not side effects: every cross-pool edge produces a manifest the TransferScheduler accounts for (a ``transfer`` WorkUnit kind), and the connector moves the payload. """ from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any @@ -24,9 +24,9 @@ class TransferManifest: keys: tuple[str, ...] producer_id: str consumer_id: str - src_location: str = "in_proc" # "in_proc" | "shm" | "nccl" | "nixl" | "disk" + src_location: str = "in_proc" # "in_proc" | "shm" | "nccl" | "nixl" | "disk" dst_location: str = "in_proc" - lifetime: str = "request" # "request" | "session" | "persistent" + lifetime: str = "request" # "request" | "session" | "persistent" cache_key: Any = None priority: int = 0 nbytes: int = 0 diff --git a/v2/video_generator.py b/v2/video_generator.py index 81d284f59..61cbfdd39 100644 --- a/v2/video_generator.py +++ b/v2/video_generator.py @@ -1,19 +1,18 @@ """v2 ``VideoGenerator`` — a typed entrypoint mirroring ``fastvideo.entrypoints.VideoGenerator``, -wrapping the v2 ``Engine`` so a ``basic_dmd_new_api.py``-style script runs against the v2 -(recipe, runtime) substrate + the real torch backend on a GPU box. +wrapping the v2 ``Engine`` so an upstream-style script runs against the v2 (recipe, runtime) substrate +and the real torch backend on a GPU box. -It consumes the OFFICIAL ``fastvideo.api`` typed configs (``GeneratorConfig`` / ``GenerationRequest`` / -``EngineConfig`` / ``OutputConfig`` / ...) — so the only delta from the upstream example is importing -``VideoGenerator`` from ``v2`` instead of ``fastvideo``. ``from_config(GeneratorConfig)`` → +It consumes the official ``fastvideo.api`` typed configs (``GeneratorConfig`` / ``GenerationRequest`` / +``EngineConfig`` / ``OutputConfig`` / ...), so the only delta from the upstream example is importing +``VideoGenerator`` from ``v2`` instead of ``fastvideo``. Flow: ``from_config(GeneratorConfig)`` → ``generate(GenerationRequest)`` → ``GenerationResult`` with ``.video_path``. All torch / fastvideo / huggingface imports are lazy (inside methods) so ``import v2`` stays CPU-clean. -Model dispatch lives in the SHARED ``v2/registry.py`` (so a CLI / server resolve identically): an -explicit HF-id registry PRIMARY + architecture inference FALLBACK, mirroring fastvideo's -``fastvideo/registry.py``. Wired: Wan2.1, Wan2.2-TI2V-5B, Wan2.2-A14B (MoE), SF-causal Wan, LTX-2 -(two-stage distilled + single-stage base/2.3). Multi-GPU / CPU-FSDP offload / VSA are accepted for API -parity but the bring-up runs single-GPU, resident, on TORCH_SDPA (no fastvideo-kernel) — see GPU_BRINGUP.md. +Model dispatch lives in the shared ``v2/registry.py`` (HF-id registry primary + architecture-inference +fallback) so a CLI or server resolve identically. Multi-GPU / CPU-FSDP offload / VSA are accepted for +API parity but the bring-up runs single-GPU, resident, on TORCH_SDPA (no fastvideo-kernel) — see +GPU_BRINGUP.md. """ from __future__ import annotations @@ -27,8 +26,8 @@ from typing import Any def _resolve_default(key: str, generic: Any, kwargs: dict, sp: Any, sd: Any, card_attr: str | None = None) -> Any: """Resolve one sampling param with precedence kwargs > SamplingParam override > card ``SamplingDefaults`` > generic fallback. ``card_attr`` aliases the card field name when it differs - from the request key (e.g. ``num_inference_steps`` -> ``SamplingDefaults.num_steps``). Pure + - torch-free so it's unit-testable without a GPU.""" + from the request key (e.g. ``num_inference_steps`` -> ``SamplingDefaults.num_steps``). Pure and + torch-free, so it's unit-testable without a GPU.""" if key in kwargs and kwargs[key] is not None: return kwargs[key] if sp is not None and getattr(sp, key, None) is not None: @@ -233,9 +232,9 @@ class VideoGenerator: image_path = os.path.join(out_dir, f"{stem}.png") import imageio.v2 as imageio imageio.imwrite(image_path, img) - # Audio (T2VS / Stable Audio): the AudioArtifact carries the raw stereo waveform; save a sibling .wav at the - # vocoder's real rate (read from the built audio_vae adapter — the artifact's default rate is a - # placeholder). A plain T2V request produces no audio artifact, so this is skipped. + # Audio (T2VS / Stable Audio): the AudioArtifact carries the raw stereo waveform; save a sibling + # .wav at the vocoder's real rate (read from the built audio_vae adapter — the artifact's default + # rate is a placeholder). A plain T2V request produces no audio artifact, so this is skipped. audio = audio_sr = audio_path = None art = out.artifacts.get("audio") if hasattr(out.artifacts, "get") else None wav = getattr(art, "samples", None) if art is not None else None