SolitaryThinker
a9e5f6ee7a
Directory rename (git mv, history preserved) plus rewrite of all references: absolute imports in tests, the zero-dep runner, docstrings, comments, and the README. No behavior change. Run: python3 -m pytest v2/tests/ -q ; python3 v2/run_tests.py ; python3 -m v2.examples
159 lines
7.4 KiB
Python
159 lines
7.4 KiB
Python
"""Worked examples (design_v3 §15), runnable as a demo:
|
||
|
||
python3 -m v2.examples
|
||
|
||
Each prints what it demonstrates. This doubles as living documentation of the API.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import numpy as np
|
||
|
||
from .models import build_default_engine, build_omni_engine
|
||
from .models.wan21 import build_wan21_card
|
||
from .parity import assert_interleave_parity
|
||
from .request import DiffusionParams, OutputSpec, SamplingParams, TaskType, make_request
|
||
from .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)
|
||
|
||
|
||
def example_a_text_to_video(eng) -> None:
|
||
print("\n(a) Text → video, one instance (Wan2.1-1.3B)")
|
||
out = eng.run(_t2v("wan2.1-1.3b", "a cat surfing a wave", 7))
|
||
print(f" video {out.artifacts['video'].frames.shape} "
|
||
f"denoise_steps={out.metrics['denoise_steps']:.0f} gpu_s={out.metrics['gpu_seconds']:.2e}")
|
||
|
||
|
||
def example_b_ltx2_two_stage(eng) -> None:
|
||
print("\n(b) LTX2.3 two-stage distilled (base 8-step → upsample → refine 3-step), shared transformer")
|
||
out = eng.run(_t2v("ltx2.3-distilled", "a neon city at night", 2))
|
||
print(f" video {out.artifacts['video'].frames.shape} "
|
||
f"base={out.metrics['base_steps']:.0f} refine={out.metrics['refine_steps']:.0f}")
|
||
|
||
|
||
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})))
|
||
print(f" latents {out.artifacts['latents'].latent.shape} chunks={out.metrics['chunks']:.0f} "
|
||
f"streamed_chunks={out.metrics.get('stream_chunks', 0)}")
|
||
|
||
|
||
def example_c2_interleave_gate(eng) -> None:
|
||
print("\n(c2) Interleave parity gate — serial == interleaved, bit-identical (the §9.3 obligation)")
|
||
reqs = [_t2v("wan2.1-1.3b", "alpha", 11), _t2v("wan2.1-1.3b", "beta", 22)]
|
||
divs = assert_interleave_parity(eng, reqs)
|
||
print(f" divergences: {divs or 'NONE — gate PASSES ✓'}")
|
||
|
||
|
||
def example_d_rl_rollout() -> None:
|
||
print("\n(d) RL rollout (DiffusionNFT): the SAME denoise loop + behavior capture (train ≡ serve)")
|
||
nft = build_diffusion_nft(build_wan21_card(), num_video_per_prompt=4, num_inner_timesteps=2)
|
||
loss, m = nft.managed_train_step({"prompts": ["a red car", "a blue boat"], "seeds": [1, 2]}, 0)
|
||
fc = nft.old.caches.stats()["feature"]
|
||
print(f" policy_loss={loss['policy_loss']:.3f} kl={loss['kl_div_loss']:.5f} "
|
||
f"reward_mean={m['reward_mean']:.3f} consistency={nft.consistency_level().value} (likelihood-free)")
|
||
print(f" shared-prompt feature-cache reuse: {fc['hits']} hits / {fc['misses']} misses "
|
||
f"(K samples encode the prompt once — the 24× reduction)")
|
||
|
||
|
||
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)))
|
||
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)))
|
||
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): "
|
||
f"{dict(eng.admission.metrics.by_kind)}")
|
||
|
||
|
||
async def _serving_demo() -> None:
|
||
import asyncio
|
||
|
||
from .deploy import DynamoWorkerAdapter, FakeDynamoRuntime, LocalFleet, build_deployment_card
|
||
from .models.wan21 import build_wan21_card, build_wan_t2v_program
|
||
from .runtime import AsyncEngine, PoolSet, wan_t2v_disaggregated
|
||
from .serving import OmniOpenAIServer
|
||
|
||
eng = build_default_engine()
|
||
build_omni_engine(eng)
|
||
ae = AsyncEngine(eng)
|
||
|
||
# disaggregated pools: encoder → denoiser → decoder
|
||
card = build_wan21_card()
|
||
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)))
|
||
print(f" disaggregated T2V (enc→den→dec): video={out.artifacts['video'].frames.shape} "
|
||
f"cross-pool transfers={out.metrics['transfers']:.0f}")
|
||
|
||
# our own OpenAI server over a real socket
|
||
server = OmniOpenAIServer(ae, engine_id="worker-0")
|
||
host, port = await server.serve(port=0)
|
||
|
||
async def http(method, path, body=b""):
|
||
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()
|
||
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)")
|
||
await server.close()
|
||
|
||
# our own fleet router + Dynamo adapter (frontable, not relied upon) — same DeploymentCard
|
||
dcard = build_deployment_card("worker-0", [card])
|
||
fleet = LocalFleet("least_loaded")
|
||
fleet.register("worker-0", ae, dcard)
|
||
routed = fleet.route(make_request(TaskType.T2V, "wan2.1-1.3b", "x", diffusion=DiffusionParams(num_steps=2)))
|
||
dyn = FakeDynamoRuntime()
|
||
dyn.register_worker(DynamoWorkerAdapter(ae, dcard))
|
||
print(f" LocalFleet routes to '{routed.worker_id}'; Dynamo adapter registered "
|
||
f"({len(dyn.registry)} worker) — both consume the SAME DeploymentCard")
|
||
|
||
|
||
def example_h_serving_and_fleet() -> None:
|
||
import asyncio
|
||
print("\n(h) Serving + fleet (OUR OWN — Dynamo-optional): async engine, role pools, OpenAI server")
|
||
asyncio.run(_serving_demo())
|
||
|
||
|
||
def main() -> None:
|
||
print("=" * 78)
|
||
print("v2 — worked examples (design_v3 §6,§13-16). One runtime, many loops.")
|
||
print("=" * 78)
|
||
eng = build_default_engine()
|
||
print(f"registered (recipe, runtime) cards: {list(eng._registry)}")
|
||
example_a_text_to_video(eng)
|
||
example_b_ltx2_two_stage(eng)
|
||
example_c_causal_streaming(eng)
|
||
example_c2_interleave_gate(eng)
|
||
example_d_rl_rollout()
|
||
example_g_omni_mot()
|
||
example_h_serving_and_fleet()
|
||
print("\n" + "=" * 78)
|
||
print("All examples ran on CPU with numpy toy components. The architecture (cards, driven")
|
||
print("loops, scheduler, caches, parity, training-on-shared-loops) is real; the neural")
|
||
print("forwards are toys. On a GPU box, swap ComponentSpec.factory for the torch adapters.")
|
||
print("=" * 78)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|