Spaces:
Running on Zero
Running on Zero
Streaming CWM demo
Browse files- .gitattributes +6 -0
- README.md +24 -2
- app.py +238 -112
- examples/ex03.safetensors +3 -0
- examples/ex03_anchor.png +3 -0
- examples/ex03_semantic.npz +3 -0
- examples/ex06.safetensors +3 -0
- examples/ex06_anchor.png +3 -0
- examples/ex06_semantic.npz +3 -0
- examples/ex08.safetensors +3 -0
- examples/ex08_anchor.png +3 -0
- examples/ex08_semantic.npz +3 -0
- examples/ex29.safetensors +3 -0
- examples/ex29_anchor.png +3 -0
- examples/ex29_semantic.npz +3 -0
- examples/ex31.safetensors +3 -0
- examples/ex31_anchor.png +3 -0
- examples/ex31_semantic.npz +3 -0
- examples/ex36.safetensors +3 -0
- examples/ex36_anchor.png +3 -0
- examples/ex36_semantic.npz +3 -0
- examples/scenes.json +32 -0
- examples/silence_audio.safetensors +3 -0
- materialize_hf_examples.py +0 -369
- requirements.txt +1 -5
- streaming_runtime.py +187 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
examples/ex03_anchor.png filter=lfs diff=lfs merge=lfs -text
|
| 37 |
+
examples/ex06_anchor.png filter=lfs diff=lfs merge=lfs -text
|
| 38 |
+
examples/ex08_anchor.png filter=lfs diff=lfs merge=lfs -text
|
| 39 |
+
examples/ex29_anchor.png filter=lfs diff=lfs merge=lfs -text
|
| 40 |
+
examples/ex31_anchor.png filter=lfs diff=lfs merge=lfs -text
|
| 41 |
+
examples/ex36_anchor.png filter=lfs diff=lfs merge=lfs -text
|
README.md
CHANGED
|
@@ -4,10 +4,32 @@ emoji: 🌍
|
|
| 4 |
colorFrom: blue
|
| 5 |
colorTo: red
|
| 6 |
sdk: gradio
|
| 7 |
-
sdk_version:
|
| 8 |
app_file: app.py
|
| 9 |
short_description: Causal chunk-by-chunk world-model video streaming
|
| 10 |
python_version: "3.12"
|
| 11 |
startup_duration_timeout: 1h
|
|
|
|
|
|
|
|
|
|
| 12 |
---
|
| 13 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 4 |
colorFrom: blue
|
| 5 |
colorTo: red
|
| 6 |
sdk: gradio
|
| 7 |
+
sdk_version: 6.30.0
|
| 8 |
app_file: app.py
|
| 9 |
short_description: Causal chunk-by-chunk world-model video streaming
|
| 10 |
python_version: "3.12"
|
| 11 |
startup_duration_timeout: 1h
|
| 12 |
+
models:
|
| 13 |
+
- Luffuly/streaming-code-world-model
|
| 14 |
+
- Comfy-Org/MiniMax-H3
|
| 15 |
---
|
| 16 |
+
|
| 17 |
+
# Streaming Code World Model
|
| 18 |
+
|
| 19 |
+
Demo of [Streaming Code World Model](https://github.com/TingtingLiao/streaming-code-world-model)
|
| 20 |
+
([checkpoints](https://huggingface.co/Luffuly/streaming-code-world-model)): causal, chunk-by-chunk video generation
|
| 21 |
+
from a live control proxy, built on [Code World Model](https://github.com/buaacyw/code-world-model)
|
| 22 |
+
(MiniMax-H3 Ref2VA + CWM LoRA).
|
| 23 |
+
|
| 24 |
+
Sampling follows `rtcwm.eval.open_world generate --steps 4 --rule x0 --shift 5 --kv-reuse --fp8` with
|
| 25 |
+
`streaming_dmd_4step.safetensors`. Adaptations for ZeroGPU (RTX PRO 6000 Blackwell): FlashAttention-3 / flex-attention
|
| 26 |
+
replaced by PyTorch SDPA over the same block-causal key prefixes, FP8 block linears (the authors' `--fp8`), and the
|
| 27 |
+
per-block AdaLN projections precomputed for the sampler's fixed set of timesteps.
|
| 28 |
+
|
| 29 |
+
Scenes are released Code World Model inference examples
|
| 30 |
+
([NTU-yiwen/code-world-model-inference-examples-40](https://huggingface.co/datasets/NTU-yiwen/code-world-model-inference-examples-40)):
|
| 31 |
+
their start frame, the "Appearance and scene" caption of the first window and the Depth + Semantic-ID proxy, pre-encoded
|
| 32 |
+
(Qwen3-VL-32B states, VAE latents) exactly as `rtcwm.eval.open_world prepare` does.
|
| 33 |
+
|
| 34 |
+
Vendored code: `musubi_tuner/` (kohya-ss/musubi-tuner @ c4dc6b8 + the CWM `minimax-h3-retake34` and headless patches,
|
| 35 |
+
Apache-2.0), `cwm_h3_inference/` (buaacyw/code-world-model, Apache-2.0), `rtcwm/` (TingtingLiao/streaming-code-world-model).
|
app.py
CHANGED
|
@@ -1,124 +1,250 @@
|
|
| 1 |
-
"""
|
| 2 |
-
|
| 3 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 4 |
import os
|
| 5 |
import sys
|
| 6 |
|
| 7 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 8 |
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
|
| 15 |
-
from pathlib import Path
|
| 16 |
|
| 17 |
-
|
| 18 |
-
|
| 19 |
-
|
| 20 |
-
|
| 21 |
-
|
| 22 |
-
|
| 23 |
|
| 24 |
-
H3_REPO, H3_REV = "Comfy-Org/MiniMax-H3", "0543966fbdce5ba05709a8f2031c94bdba629b4a"
|
| 25 |
-
QWEN_REPO, QWEN_REV = "Qwen/Qwen3-VL-32B-Instruct", "0cfaf48183f594c314753d30a4c4974bc75f3ccb"
|
| 26 |
-
EX_REPO, EX_REV = "NTU-yiwen/code-world-model-inference-examples-40", "f8307bf50cb19bed191682734fdf444738160f91"
|
| 27 |
-
EXAMPLES = [3, 6, 8, 29, 31, 36]
|
| 28 |
-
CHUNKS = 24
|
| 29 |
-
CLIP_FRAMES = 17
|
| 30 |
-
|
| 31 |
-
models = Path(snapshot_download(H3_REPO, revision=H3_REV, local_dir="/tmp/h3", allow_patterns=[
|
| 32 |
-
"text_encoders/qwen3vl_32b_minimax_h3_bf16.safetensors", "vae/minimax_h3_video_vae_fp16.safetensors",
|
| 33 |
-
"vae/minimax_h3_audio_vae_fp32.safetensors"]))
|
| 34 |
-
proc_dir = Path(snapshot_download(QWEN_REPO, revision=QWEN_REV, local_dir="/tmp/qwen_proc", allow_patterns=[
|
| 35 |
-
"chat_template.json", "config.json", "generation_config.json", "merges.txt", "preprocessor_config.json",
|
| 36 |
-
"tokenizer_config.json", "tokenizer.json", "video_preprocessor_config.json", "vocab.json"]))
|
| 37 |
-
src = Path(snapshot_download(EX_REPO, repo_type="dataset", revision=EX_REV, allow_patterns=[
|
| 38 |
-
"manifest.json", "semantic_id_palette.json"] + [f"{i}/*" for i in EXAMPLES] + [f"{i}/**" for i in EXAMPLES]))
|
| 39 |
-
|
| 40 |
-
import materialize_hf_examples as mat
|
| 41 |
-
|
| 42 |
-
REL = Path("/tmp/release")
|
| 43 |
-
for i in EXAMPLES:
|
| 44 |
-
mat.materialize_example(src, REL, i, require_config=False, frame_limit=(CHUNKS + 1) * CLIP_FRAMES)
|
| 45 |
-
print("materialized", flush=True)
|
| 46 |
-
|
| 47 |
-
from musubi_tuner.minimax_h3.audio_vae import load_audio_vae
|
| 48 |
-
from musubi_tuner.minimax_h3.text_encoder import load_h3_processor, load_h3_text_encoder
|
| 49 |
-
from musubi_tuner.minimax_h3.video_vae import encode_video_condition, load_video_vae
|
| 50 |
-
|
| 51 |
-
vae = load_video_vae(models / "vae/minimax_h3_video_vae_fp16.safetensors", device="cpu", dtype=torch.float16).eval().requires_grad_(False)
|
| 52 |
-
vae.to("cuda")
|
| 53 |
-
avae = load_audio_vae(models / "vae/minimax_h3_audio_vae_fp32.safetensors", device="cpu", dtype=torch.float32).eval()
|
| 54 |
-
avae.to("cuda")
|
| 55 |
-
proc = load_h3_processor(str(proc_dir))
|
| 56 |
-
qwen = load_h3_text_encoder(str(models / "text_encoders/qwen3vl_32b_minimax_h3_bf16.safetensors"),
|
| 57 |
-
processor_path=str(proc_dir), device="cpu", dtype=torch.bfloat16).eval()
|
| 58 |
-
qwen.to("cuda")
|
| 59 |
-
print("models loaded", flush=True)
|
| 60 |
-
|
| 61 |
-
|
| 62 |
-
def _prepare_pixels(frames): # musubi_tuner.minimax_h3_generate_video._prepare_pixels
|
| 63 |
-
if frames.dtype == torch.uint8:
|
| 64 |
-
frames = frames.float().div(127.5).sub(1.0)
|
| 65 |
-
else:
|
| 66 |
-
frames = frames.float().mul(2.0).sub(1.0)
|
| 67 |
-
return frames.permute(3, 0, 1, 2).unsqueeze(0).contiguous()
|
| 68 |
-
|
| 69 |
-
|
| 70 |
-
@spaces.GPU(duration=600, size="xlarge")
|
| 71 |
-
def prep() -> str:
|
| 72 |
-
"""Encode the example caches and return a zip."""
|
| 73 |
-
from cwm_h3_inference.cache import _load_anchor
|
| 74 |
-
from cwm_h3_inference.config import _canonical_caption
|
| 75 |
-
from cwm_h3_inference.constants import TARGET_HEIGHT, TARGET_WIDTH, WINDOW_AUDIO_SAMPLES
|
| 76 |
-
from cwm_h3_inference.duv import load_duv_frame
|
| 77 |
-
from rtcwm.eval.open_world import static_caption
|
| 78 |
-
from rtcwm.train.preencode import static_text_states
|
| 79 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 80 |
dev = torch.device("cuda")
|
| 81 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 82 |
with torch.inference_mode():
|
| 83 |
-
|
| 84 |
-
|
| 85 |
-
|
| 86 |
-
|
| 87 |
-
|
| 88 |
-
|
| 89 |
-
|
| 90 |
-
|
| 91 |
-
duv = np.stack([load_duv_frame(cond, f, for_qwen=False) for f in range(k * CLIP_FRAMES)])
|
| 92 |
-
duv = np.concatenate([duv, np.repeat(duv[-1:], CLIP_FRAMES, 0)], 0)
|
| 93 |
-
with torch.inference_mode():
|
| 94 |
-
ref = encode_video_condition(vae, _prepare_pixels(torch.from_numpy(duv)).to(dev, torch.float16))[0]
|
| 95 |
-
ref = ref.float().cpu()[:, :5 * k].contiguous()
|
| 96 |
-
picture1 = _load_anchor(ex / "first_frame.png", size=(TARGET_WIDTH, TARGET_HEIGHT))
|
| 97 |
-
png = out / f"ex{i:02d}_anchor.png"
|
| 98 |
-
Image.fromarray(picture1).save(png)
|
| 99 |
-
with torch.inference_mode():
|
| 100 |
-
lat = encode_video_condition(vae, _prepare_pixels(torch.from_numpy(picture1).unsqueeze(0)).to(dev, torch.float16))
|
| 101 |
-
caption = _canonical_caption(static_caption(prompt), f"ow_example_{i:02d}/static_c00")
|
| 102 |
-
hidden, tags = static_text_states(proc, qwen, picture1, png, cond, 0, caption)
|
| 103 |
-
sem = np.stack([np.asarray(Image.open(cond / f"{f:06d}.semantic_id.png")) for f in range(k * CLIP_FRAMES)])
|
| 104 |
-
np.savez_compressed(out / f"ex{i:02d}_semantic.npz", semantic_id=sem.astype(np.uint8))
|
| 105 |
-
save_file({"duv": ref, "anchor": lat[0].float().cpu().contiguous(),
|
| 106 |
-
"hidden": hidden.to("cpu", torch.bfloat16).contiguous(), "tags": tags.to("cpu", torch.int64).contiguous()},
|
| 107 |
-
str(out / f"ex{i:02d}.safetensors"), metadata={"chunks": str(k), "caption": caption})
|
| 108 |
-
log.append(dict(id=i, chunks=k, duv=list(ref.shape), anchor=list(lat.shape), hidden=list(hidden.shape)))
|
| 109 |
-
print(log[-1], flush=True)
|
| 110 |
-
(out / "log.json").write_text(json.dumps(log))
|
| 111 |
-
z = tempfile.NamedTemporaryFile(suffix=".zip", delete=False).name
|
| 112 |
-
with zipfile.ZipFile(z, "w", zipfile.ZIP_STORED) as zf:
|
| 113 |
-
for p in out.iterdir():
|
| 114 |
-
zf.write(p, p.name)
|
| 115 |
-
return z
|
| 116 |
|
| 117 |
|
| 118 |
-
|
| 119 |
-
|
| 120 |
-
|
| 121 |
-
|
| 122 |
-
|
| 123 |
|
| 124 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Streaming Code World Model -- causal, chunk-by-chunk world-model video on ZeroGPU.
|
| 2 |
+
|
| 3 |
+
Reference implementation: https://github.com/TingtingLiao/streaming-code-world-model (rtcwm), sampling exactly as
|
| 4 |
+
`python -m rtcwm.eval.open_world generate --steps 4 --rule x0 --shift 5 --kv-reuse --fp8` with the
|
| 5 |
+
`streaming_dmd_4step` LoRA on the released MiniMax-H3 Ref2VA model. Hardware adaptations (RTX PRO 6000, sm_120):
|
| 6 |
+
* FlashAttention-3 / flex_attention -> torch SDPA over the same key prefixes (rtcwm.model.attention.dense_df);
|
| 7 |
+
* the blocks' qkv / out / fc1 / fc2 in FP8 exactly as rtcwm.model.fp8 (the authors' --fp8 option);
|
| 8 |
+
* each block's AdaLN projection only ever sees the time embedding of a handful of fixed times (condition rows
|
| 9 |
+
1.0 / 0.999 and the 4 sampling steps), so its outputs are precomputed once at startup instead of keeping the
|
| 10 |
+
26 GB of AdaLN weights on the GPU.
|
| 11 |
+
"""
|
| 12 |
import os
|
| 13 |
import sys
|
| 14 |
|
| 15 |
+
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
|
| 16 |
+
APP_DIR = os.path.dirname(os.path.abspath(__file__))
|
| 17 |
+
sys.path.insert(0, APP_DIR)
|
| 18 |
+
|
| 19 |
+
import spaces # noqa: E402 (before torch)
|
| 20 |
+
|
| 21 |
+
import json # noqa: E402
|
| 22 |
+
import random # noqa: E402
|
| 23 |
+
import tempfile # noqa: E402
|
| 24 |
+
import time # noqa: E402
|
| 25 |
+
from pathlib import Path # noqa: E402
|
| 26 |
+
|
| 27 |
+
import av # noqa: E402
|
| 28 |
+
import gradio as gr # noqa: E402
|
| 29 |
+
import numpy as np # noqa: E402
|
| 30 |
+
import torch # noqa: E402
|
| 31 |
+
from huggingface_hub import hf_hub_download # noqa: E402
|
| 32 |
+
from PIL import Image # noqa: E402
|
| 33 |
+
from safetensors import safe_open # noqa: E402
|
| 34 |
+
from safetensors.torch import load_file # noqa: E402
|
| 35 |
+
|
| 36 |
+
from musubi_tuner.minimax_h3.model import load_h3_transformer # noqa: E402
|
| 37 |
+
from musubi_tuner.minimax_h3.video_vae import load_video_vae # noqa: E402
|
| 38 |
+
from rtcwm.model import sampler as rt_sampler # noqa: E402
|
| 39 |
+
from rtcwm.model.layout import CHUNK_LATENTS, CLIP_FRAMES, TARGET_LATENTS # noqa: E402
|
| 40 |
+
from rtcwm.model.packing import audio_frames # noqa: E402
|
| 41 |
|
| 42 |
+
H3_REPO, H3_REV = "Comfy-Org/MiniMax-H3", "0543966fbdce5ba05709a8f2031c94bdba629b4a"
|
| 43 |
+
LORA_REPO, LORA_FILE = "Luffuly/streaming-code-world-model", "streaming_dmd_4step.safetensors"
|
| 44 |
+
STEPS, SHIFT, RULE = 4, 5.0, "x0" # the released 4-step DMD sampling (model card)
|
| 45 |
+
MAX_CHUNKS = 24
|
| 46 |
+
FPS = 24
|
| 47 |
+
EX_DIR = Path(APP_DIR) / "examples"
|
|
|
|
| 48 |
|
| 49 |
+
# --------------------------------------------------------------------------------------------- examples
|
| 50 |
+
SCENES = json.loads((EX_DIR / "scenes.json").read_text()) # name -> {id, chunks, caption}
|
| 51 |
+
SILENCE = load_file(str(EX_DIR / "silence_audio.safetensors"))["audio"]
|
| 52 |
+
PALETTE = np.array([[0, 0, 0], [135, 206, 235], [30, 90, 200], [150, 120, 80], [90, 90, 90], [60, 160, 60],
|
| 53 |
+
[200, 120, 80], [220, 200, 60], [230, 40, 40], [200, 80, 200], [40, 220, 220], [250, 250, 250]],
|
| 54 |
+
dtype=np.uint8) # rtcwm.eval.open_world.PALETTE
|
| 55 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 56 |
|
| 57 |
+
def _scene_tensors(name: str) -> dict:
|
| 58 |
+
with safe_open(str(EX_DIR / f"ex{SCENES[name]['id']:02d}.safetensors"), "pt") as f:
|
| 59 |
+
return {k: f.get_tensor(k) for k in f.keys()}
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
SCENE_DATA = {name: _scene_tensors(name) for name in SCENES}
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def load_window(name: str, j: int) -> dict:
|
| 66 |
+
"""rtcwm.train.dataset.load_window(..., start=0, with_target=False) on the precomputed caches."""
|
| 67 |
+
d = SCENE_DATA[name]
|
| 68 |
+
duv = d["duv"][:, CHUNK_LATENTS * j:CHUNK_LATENTS * j + TARGET_LATENTS]
|
| 69 |
+
assert duv.shape[1] == TARGET_LATENTS, (name, j, duv.shape)
|
| 70 |
+
ref = torch.nn.functional.pad(duv.unsqueeze(0), (0, 1, 0, 0), value=0.0)
|
| 71 |
+
audio = SILENCE[..., :audio_frames(TARGET_LATENTS)].unsqueeze(0).clone()
|
| 72 |
+
return dict(audio=audio, anchor=d["anchor"].unsqueeze(0), ref=ref, hidden=d["hidden"].unsqueeze(0),
|
| 73 |
+
tags=d["tags"].unsqueeze(0), stream_start=j == 0)
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
# --------------------------------------------------------------------------------------------- load
|
| 77 |
+
from streaming_runtime import prepare_transformer # noqa: E402
|
| 78 |
+
|
| 79 |
+
dit = hf_hub_download(H3_REPO, "diffusion_models/minimax_h3_ref2va_bf16.safetensors", revision=H3_REV, local_dir="/tmp/h3")
|
| 80 |
+
lora_path = hf_hub_download(LORA_REPO, LORA_FILE)
|
| 81 |
+
transformer = load_h3_transformer(dit, device="cpu", dtype=torch.bfloat16, attn_mode="torch")
|
| 82 |
+
transformer = prepare_transformer(transformer, load_file(lora_path)).to("cuda")
|
| 83 |
+
vae_path = hf_hub_download(H3_REPO, "vae/minimax_h3_video_vae_fp16.safetensors", revision=H3_REV, local_dir="/tmp/h3")
|
| 84 |
+
vae = load_video_vae(vae_path, device="cpu", dtype=torch.float16).eval().requires_grad_(False).to("cuda")
|
| 85 |
+
print("models ready", flush=True)
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
# --------------------------------------------------------------------------------------------- helpers
|
| 89 |
+
def _write_mp4(frames: np.ndarray, crf: int = 18) -> str:
|
| 90 |
+
path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
|
| 91 |
+
with av.open(path, "w") as c:
|
| 92 |
+
s = c.add_stream("libx264", rate=FPS)
|
| 93 |
+
s.width, s.height, s.pix_fmt = frames.shape[2], frames.shape[1], "yuv420p"
|
| 94 |
+
s.options = {"crf": str(crf)}
|
| 95 |
+
for f in frames:
|
| 96 |
+
for p in s.encode(av.VideoFrame.from_ndarray(np.ascontiguousarray(f), format="rgb24")):
|
| 97 |
+
c.mux(p)
|
| 98 |
+
for p in s.encode():
|
| 99 |
+
c.mux(p)
|
| 100 |
+
return path
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def _to_uint8(px: torch.Tensor) -> np.ndarray:
|
| 104 |
+
"""[3,F,H,W] in [-1,1] -> [F,H,W,3] uint8 (rtcwm.eval.video.decode)."""
|
| 105 |
+
return ((px.float().clamp(-1, 1) + 1) * 127.5).round().byte().permute(1, 2, 3, 0).cpu().numpy()
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
_PROXY_CACHE = {}
|
| 109 |
+
|
| 110 |
+
|
| 111 |
+
def proxy_video(name: str, chunks: int) -> str:
|
| 112 |
+
"""The Semantic-ID half of the proxy the stream follows, coloured as in rtcwm.eval.open_world.proxy_frames."""
|
| 113 |
+
key = (name, int(chunks))
|
| 114 |
+
if key not in _PROXY_CACHE:
|
| 115 |
+
sem = np.load(EX_DIR / f"ex{SCENES[name]['id']:02d}_semantic.npz")["semantic_id"][: int(chunks) * CLIP_FRAMES]
|
| 116 |
+
frames = PALETTE[np.clip(sem, 0, 11)]
|
| 117 |
+
frames = np.stack([np.asarray(Image.fromarray(f).resize((672, 384), Image.NEAREST)) for f in frames])
|
| 118 |
+
_PROXY_CACHE[key] = _write_mp4(frames, crf=23)
|
| 119 |
+
return _PROXY_CACHE[key]
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
def _duration(scene, chunks=8, seed=0, *args, **kwargs):
|
| 123 |
+
return int(40 + 14 * int(chunks))
|
| 124 |
+
|
| 125 |
+
|
| 126 |
+
# --------------------------------------------------------------------------------------------- inference
|
| 127 |
+
@spaces.GPU(duration=_duration, size="xlarge")
|
| 128 |
+
def generate(scene: str, chunks: int = 8, seed: int = 0):
|
| 129 |
+
"""Stream a world-model video chunk by chunk from a released Code World Model scene.
|
| 130 |
+
|
| 131 |
+
Each 17-frame chunk (0.71 s at 24 fps) is generated causally from the start image, its caption, the
|
| 132 |
+
Depth + Semantic-ID proxy up to that chunk and the previous chunks (KV cache, 6-chunk attention span),
|
| 133 |
+
with the 4-step DMD streaming model.
|
| 134 |
+
|
| 135 |
+
Args:
|
| 136 |
+
scene: one of the released example scenes.
|
| 137 |
+
chunks: number of 0.71 s chunks to stream (1-24).
|
| 138 |
+
seed: random seed.
|
| 139 |
+
|
| 140 |
+
Returns:
|
| 141 |
+
(latest-chunk preview frame, final mp4 path) -- yielded progressively.
|
| 142 |
+
"""
|
| 143 |
+
chunks = max(1, min(MAX_CHUNKS, int(chunks)))
|
| 144 |
dev = torch.device("cuda")
|
| 145 |
+
gen = torch.Generator().manual_seed(int(seed))
|
| 146 |
+
t_start = time.time()
|
| 147 |
+
torch.cuda.reset_peak_memory_stats()
|
| 148 |
+
events = []
|
| 149 |
+
|
| 150 |
+
def on_chunk(k, secs):
|
| 151 |
+
events.append((k, secs))
|
| 152 |
+
print(f"{scene}: chunk {k + 1}/{chunks} {secs:.1f}s", flush=True)
|
| 153 |
+
|
| 154 |
+
# stream the chunks (rtcwm.model.sampler.generate_stream), previewing each finished chunk
|
| 155 |
+
import queue
|
| 156 |
+
import threading
|
| 157 |
+
q = queue.Queue()
|
| 158 |
+
result = {}
|
| 159 |
+
|
| 160 |
+
def preview_chunk(k, secs):
|
| 161 |
+
on_chunk(k, secs)
|
| 162 |
+
q.put(k)
|
| 163 |
+
|
| 164 |
+
def worker():
|
| 165 |
+
try:
|
| 166 |
+
with torch.inference_mode():
|
| 167 |
+
result["lat"] = rt_sampler.generate_stream(
|
| 168 |
+
transformer, lambda j: load_window(scene, j), chunks, STEPS, gen, dev, RULE, SHIFT,
|
| 169 |
+
history_noise=0.0, kv_reuse=True, on_chunk=preview_chunk)
|
| 170 |
+
except Exception as e: # noqa: BLE001
|
| 171 |
+
result["err"] = e
|
| 172 |
+
q.put(None)
|
| 173 |
+
|
| 174 |
+
th = threading.Thread(target=worker, daemon=True)
|
| 175 |
+
th.start()
|
| 176 |
+
while True:
|
| 177 |
+
k = q.get()
|
| 178 |
+
if k is None:
|
| 179 |
+
break
|
| 180 |
+
yield gr.update(label=f"Streaming… chunk {k + 1}/{chunks} ({(k + 1) * CLIP_FRAMES / FPS:.1f} s generated)"), None
|
| 181 |
+
th.join()
|
| 182 |
+
if "err" in result:
|
| 183 |
+
raise result["err"]
|
| 184 |
+
lat = result["lat"]
|
| 185 |
+
gen_s = time.time() - t_start
|
| 186 |
with torch.inference_mode():
|
| 187 |
+
px = vae.decode(lat.to(dev, torch.float16))[0]
|
| 188 |
+
frames = _to_uint8(px)
|
| 189 |
+
path = _write_mp4(frames)
|
| 190 |
+
peak = torch.cuda.max_memory_allocated() / 2**30
|
| 191 |
+
print(f"{scene}: {chunks} chunks, {len(frames)} frames, stream {gen_s:.1f}s, total {time.time() - t_start:.1f}s, "
|
| 192 |
+
f"peak {peak:.1f} GiB", flush=True)
|
| 193 |
+
yield gr.update(label=f"Done: {chunks} chunks ({len(frames) / FPS:.1f} s) in {time.time() - t_start:.0f} s",
|
| 194 |
+
value=Image.fromarray(frames[-1]).resize((672, 384))), path
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 195 |
|
| 196 |
|
| 197 |
+
# --------------------------------------------------------------------------------------------- UI
|
| 198 |
+
def show_scene(scene: str, chunks: int = 8):
|
| 199 |
+
"""Preview a scene's start image, caption and proxy."""
|
| 200 |
+
s = SCENES[scene]
|
| 201 |
+
return str(EX_DIR / f"ex{s['id']:02d}_anchor.png"), s["caption"], proxy_video(scene, chunks)
|
| 202 |
|
| 203 |
+
|
| 204 |
+
CSS = """
|
| 205 |
+
#col-container { max-width: 1280px; margin: 0 auto; }
|
| 206 |
+
.dark .gradio-container { color: var(--body-text-color); }
|
| 207 |
+
"""
|
| 208 |
+
NAMES = list(SCENES)
|
| 209 |
+
|
| 210 |
+
with gr.Blocks() as demo:
|
| 211 |
+
with gr.Column(elem_id="col-container"):
|
| 212 |
+
gr.Markdown(
|
| 213 |
+
"# Streaming Code World Model\n"
|
| 214 |
+
"Causal, chunk-by-chunk video generation from a live control proxy: each 0.71 s chunk is generated from "
|
| 215 |
+
"the start image, its caption, the Depth + Semantic-ID proxy **so far** and the previous chunks — never "
|
| 216 |
+
"the future. 4-step DMD streaming LoRA on MiniMax-H3 Ref2VA + "
|
| 217 |
+
"[Code World Model](https://github.com/buaacyw/code-world-model). \n"
|
| 218 |
+
"[GitHub](https://github.com/TingtingLiao/streaming-code-world-model) · "
|
| 219 |
+
"[Checkpoints](https://huggingface.co/Luffuly/streaming-code-world-model)"
|
| 220 |
+
)
|
| 221 |
+
with gr.Row():
|
| 222 |
+
with gr.Column(scale=1):
|
| 223 |
+
scene = gr.Dropdown(NAMES, value=NAMES[0], label="Scene (released CWM example)")
|
| 224 |
+
chunks = gr.Slider(1, MAX_CHUNKS, value=8, step=1, label="Chunks to stream (0.71 s each)")
|
| 225 |
+
run = gr.Button("Stream", variant="primary")
|
| 226 |
+
anchor = gr.Image(label="Start image (Picture 1)", interactive=False, height=240)
|
| 227 |
+
proxy = gr.Video(label="Control proxy (Semantic-ID view)", interactive=False, autoplay=True, loop=True,
|
| 228 |
+
height=240)
|
| 229 |
+
with gr.Accordion("Advanced", open=False):
|
| 230 |
+
seed = gr.Number(value=0, precision=0, label="Seed")
|
| 231 |
+
randomize = gr.Checkbox(value=False, label="Randomize seed")
|
| 232 |
+
caption = gr.Textbox(label="Static caption (pre-encoded by Qwen3-VL-32B)", lines=6,
|
| 233 |
+
interactive=False)
|
| 234 |
+
with gr.Column(scale=2):
|
| 235 |
+
status = gr.Image(label="Stream status", interactive=False, height=200)
|
| 236 |
+
video = gr.Video(label="Generated stream (1344×768, 24 fps)", autoplay=True)
|
| 237 |
+
|
| 238 |
+
def _pick_seed(s, r):
|
| 239 |
+
return random.randint(0, 2**31 - 1) if r else int(s)
|
| 240 |
+
|
| 241 |
+
scene.change(show_scene, [scene, chunks], [anchor, caption, proxy], api_name=False)
|
| 242 |
+
chunks.release(show_scene, [scene, chunks], [anchor, caption, proxy], api_name=False)
|
| 243 |
+
run.click(_pick_seed, [seed, randomize], seed, api_name=False).then(
|
| 244 |
+
generate, [scene, chunks, seed], [status, video], api_name="generate")
|
| 245 |
+
demo.load(show_scene, [scene, chunks], [anchor, caption, proxy], api_name=False)
|
| 246 |
+
|
| 247 |
+
gr.Examples(examples=[[n, 8] for n in NAMES], inputs=[scene, chunks], outputs=[status, video], fn=generate,
|
| 248 |
+
cache_examples=True, cache_mode="lazy", label="Scenes")
|
| 249 |
+
|
| 250 |
+
demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)
|
examples/ex03.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:4785caba442438656cd34056a7da119d2966878d67745868c61edcd86d21ea98
|
| 3 |
+
size 32910376
|
examples/ex03_anchor.png
ADDED
|
Git LFS Details
|
examples/ex03_semantic.npz
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:10d31cb12bf7c4b20df70213e8f6a17bab0127741b1e049326229ec1d18844e5
|
| 3 |
+
size 707731
|
examples/ex06.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:352418f73ffa28f6b08c418ba61a029c3c27878f2e32563d06539d12a9c0239b
|
| 3 |
+
size 32438672
|
examples/ex06_anchor.png
ADDED
|
Git LFS Details
|
examples/ex06_semantic.npz
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:a03e31f973b6582cbab5236eb83d5ee13232b954751fa81ab4fdeb27f699fd93
|
| 3 |
+
size 340520
|
examples/ex08.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:8cff117933fd681f8ae9ca002d15423c24c223a214ea9273d1bcce29b680436d
|
| 3 |
+
size 31997816
|
examples/ex08_anchor.png
ADDED
|
Git LFS Details
|
examples/ex08_semantic.npz
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:1e803072d468f6b652253bf0bbb65c24071883f11a84261cbd3e7f1dbc315029
|
| 3 |
+
size 947066
|
examples/ex29.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:e8841affedb78a159c92822cf487781b96ffd81455d97dc5d9812bdc2a9b2f90
|
| 3 |
+
size 32725840
|
examples/ex29_anchor.png
ADDED
|
Git LFS Details
|
examples/ex29_semantic.npz
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:ce08de120367af8cd161e880ee37789c1b8ab5c807e1dbf20f12cec3297ca7fc
|
| 3 |
+
size 675227
|
examples/ex31.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:c8fce85a06bfa466a187baa680ef268e5fec2b067c3b5033660b530697f675a7
|
| 3 |
+
size 32028552
|
examples/ex31_anchor.png
ADDED
|
Git LFS Details
|
examples/ex31_semantic.npz
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:5c1eb24ccd1df760a4d6b439ed35821b1f7d81889cb0988a95a16f555fd9e55a
|
| 3 |
+
size 591886
|
examples/ex36.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:0e76cbb7cb2b884ca2e19eb4d0912a3a937358d0b778f5e09267e1bbaefa4fb6
|
| 3 |
+
size 34253600
|
examples/ex36_anchor.png
ADDED
|
Git LFS Details
|
examples/ex36_semantic.npz
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:14aef84569137ffa9ee0365cc71a868519d95bde238eee94ab718e7cc3f560ef
|
| 3 |
+
size 902838
|
examples/scenes.json
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"Clockwork courier \u00b7 retro-futurist promenade": {
|
| 3 |
+
"id": 3,
|
| 4 |
+
"chunks": 24,
|
| 5 |
+
"caption": "Appearance and scene:\n\nThe protagonist is a friendly nonhuman clockwork exposition courier with a compact, round-bodied mechanical build. Its warm ivory enamel torso is egg-shaped and edged by fine soft-charcoal seams, tiny rivets, circular coral fittings, and small brass mechanisms. A clear bubble helm surrounds its powder-blue head, and pointed coral-and-ivory receiver fins project from either side. Two coiled brass winding keys rise from its upper back. Its short articulated arms and legs have cream enamel shells, exposed brass pivot rings, mitten-like mechanical hands, and rounded coral-soled feet; these are stable identity features. A broad coral utility sash wraps diagonally around its waist and supports a square coral satchel with a brass clasp. At the beginning, it leans forward in a running stride near the lower center-left, facing away and toward screen-right, with one arm reaching forward, the other lowered, one foot planted behind, and the opposite foot lifted. Its long deep-teal printed shadow falls ahead-right across the butter-yellow roadway.\n\nThe directly approached vehicle later in view is the coral open-cockpit roadster, a compact atomic-age touring car rendered in the same layered print language. It has a low pointed nose, paired cream headlamp circles, smooth coral fenders, powder-blue windshield glass, a deep-teal open cabin, silver-gray wheel rims with coral hubs, and a tapering rear deck with restrained fin forms. When it first becomes visible ahead-right, its nose faces screen-left along the roadside while its rear remains screen-right beside the solar-topiary boundary. It stays at rideable road scale, and its solid simplified surfaces retain coarse stipple, slight ink misregistration, and flat charcoal undershadows.\n\nA powder-blue finned sedan briefly approaches from farther up the roadway. It has a rounded cream roof, broad powder-blue side panels, coral hub details, and a compact soft-charcoal windshield, all simplified into screen-printed shapes. No other moving identity is visible at the opening.\n\nThe setting is an optimistic retro-futurist exposition district printed on warm fibrous paper. A sweeping cream ceramic pavilion with a thick coral rim and four vertical coral oval windows fills the left side. Beyond it stands a dense deep-teal modular solar-topiary wall patterned with powder-blue panel grids, circular recesses, and clipped geometric foliage. A powder-blue pavilion and a tall cypress-shaped antenna sculpture rise behind the wall, while slender coral signal poles punctuate the fairground. Terraced pale-pink mooncrete with round teal shrubs and angular agave forms borders the road opening on the right. The broad butter-yellow road curves uphill through the exhibits, edged by powder-blue curbs and flat coral accents under a clear stippled blue sky. Halftone dots, coarse three-ink grain, imperfect color overlap, and flat layered shadows remain visible across every surface. The trailing camera begins low and wide behind the courier, preserving its lower center-left placement and the open roadway ahead, then follows forward and turns gradually toward the approaching roadster."
|
| 6 |
+
},
|
| 7 |
+
"Olivewood automaton \u00b7 Mediterranean cliffside": {
|
| 8 |
+
"id": 6,
|
| 9 |
+
"chunks": 24,
|
| 10 |
+
"caption": "Appearance and scene:\n\nThe protagonist is an olivewood courier automaton with a compact, athletic humanoid build. Its small cream-colored ceramic sun-mask head bears a crisp cobalt sun emblem and rests above a short dark-bronze neck. A ribbed, amphora-shaped olivewood torso narrows into a banded bronze waist, with curved shoulder plates, jointed wooden upper limbs, dark-bronze elbows, articulated hands, wooden shins, and round bronze knees. The ceramic head, amphora-shaped torso, articulated dark-bronze joints, and carefully fitted wooden limbs are stable identity features. Loose indigo cloth trousers gather below the waist and bunch around the knees, while simple rope-soled sandals wrap its feet with narrow ochre ties. Fine black hatching follows the wood grain, dense crosshatching darkens the joints, and small cobalt accents visually connect the ceramic head with the trousers. Bright sunlight strikes the upper wood and ceramic while the forward lean leaves narrow warm-black shadows beneath the arms and around the waist.\n\nAt the opening instant, the olivewood courier automaton occupies the central-lower slope exactly as a running figure seen from behind. It faces uphill toward the upper center, with its torso inclined forward, one sandal planted on pale limestone, the rear foot lifted, both knees bent, and both arms pumping close to its sides. The immense village-cliff fills the right and upper area as a continuous overhanging mass of honey-colored stone and chalk-white lime-plaster dwellings. Tiny cobalt shutters, recessed windows, terracotta ledges, painted ceramic drainage jars, scattered olive leaves, and bright bougainvillea punctuate the facade without breaking its solid lower boundary. Its shadowed underside follows the dark rock shelf above the runner. A broad sun-bleached limestone lane remains open on the left, passing around the projecting village-cliff and climbing into the caldera. Sparse grasses, yellow flowers, purple herbs, and low olive scrub emerge from cracks across the foreground. Farther left, stepped olive terraces and small white dwellings cover dry ochre hills beneath a clear turquoise sky crossed by a thin halftone cloud.\n\nThe entire landscape is rendered as a bold comic-book print with forceful warm-black contours, dry-brush shadows, visible Ben-Day dots, slight color-registration offsets, and dense crosshatching across stone and wood. The limited chalk-white, ochre, terracotta, cobalt, olive-green, turquoise, indigo, and warm-black palette remains consistent from foreground plants to distant terraces. Bright late-afternoon light produces crisp graphic shadows. The trailing camera begins low behind the olivewood courier automaton, showing its full running body against the towering village-cliff, then follows its route around the projecting stone while gradually rising and widening toward the open lane."
|
| 11 |
+
},
|
| 12 |
+
"Lotus carriage \u00b7 lantern-lit night road": {
|
| 13 |
+
"id": 8,
|
| 14 |
+
"chunks": 24,
|
| 15 |
+
"caption": "Appearance and scene:\n\nThe protagonist is a low, driveable festival carriage shaped like a hand-sculpted glazed-clay lotus bud, with a layered coral-and-jade body, a warm lantern chamber beneath its petals, four exposed dark ceramic spool wheels, and delicate brass-colored clay fittings. The compact body stays low over its wheels, and overlapping coral petals wrap around jade-green inner panels beneath fine raised brass-colored ribs. Curled fittings ornament the rear panel above two small red taillanterns, while the tapered front remains enclosed beneath the converging petals. Subtle fingerprints, modeled seams, uneven glazed edges, and small kiln-darkened recesses are stable identity features. Its warm inner light softly brightens the lower petals and nearby dust. No rider or other occupant is visible. At 0.00s, its unmistakable lanterned rear is nearest the high trailing camera in the lower foreground, while its tapered front points away toward the upper-left reach of the broad road.\n\nAn extremely tiny closed festival wagon occupies the remote upper-right road beyond the river. It has a rounded dark-clay cabin, four minute spool wheels, tiny amber lantern windows, and a compact enclosed festival-wagon silhouette. It remains far smaller than the coral-and-jade lotus carriage and is the only other mobile identity visible.\n\nA wide road of moon-lavender clay dust curves ahead through pinched moss-green clay rosettes and low sculpted shrubs. The road surface is open and traversable, with shallow wheel grooves, scattered clay pebbles, compressed patches, and gently modeled shoulders that remain level with the surrounding ground. A glossy teal resin-clay river follows its exact channel along the road's right side, catching broken cool moon reflections and warm lantern streaks. A low, nearly straight bridge spans the water near the center distance, retaining its modest deck, short supports, broad openings, and small attached lanterns. Slender bamboo-look clay poles carry hanging amber lanterns along both banks and beside the road. Blank parchment-colored shadow-play screens stand among the shrubs without readable markings. Layered violet-brown clay hills, clustered miniature trees, and dark sculpted clouds recede beneath a midnight-indigo sky and a bright textured moon. Warm lantern light, cool moonlight, miniature-set depth, fingerprints, and stop-motion sculpting imperfections remain visible throughout. The high trailing camera begins directly behind and above the coral-and-jade lotus carriage, showing its lanterned rear, the open bend ahead, the river, and the bridge, then follows smoothly while yawing with the carriage's changing heading."
|
| 16 |
+
},
|
| 17 |
+
"Crystal-cavern wayfarer \u00b7 dark forest": {
|
| 18 |
+
"id": 29,
|
| 19 |
+
"chunks": 24,
|
| 20 |
+
"caption": "Appearance and scene:\n\nThe protagonist is a rugged adult crystal-cavern wayfarer with a sturdy, athletic human build. His weathered olive-tan skin, strong angular profile, silver-black hair braided into a thick rope down his back, loose gray-streaked temple hairs, and rough dark stubble are stable identity features. A charcoal wool traveling mantle wraps his shoulders beneath a shallow hood, with frayed edges, coarse hand stitching, and a long uneven cloak panel hanging behind him. Under it he wears a fitted dark leather tunic over a smoke-gray undershirt, a broad utility belt with practical iron buckles, layered leather bracers, close leather gloves, reinforced dark trousers, and heavy climbing boots. A compact translucent quartz climbing tool rests in an iron setting at his right hip, secured tightly to the belt. Dust dulls the lower cloak and boots, while cool mineral light glints along the tool, worn leather, buckles, and individual silver strands in his braid. He begins at lower center-left, seen from behind and in three-quarter profile with his torso leaning uphill toward screen-right, his right knee deeply bent and lifted into the next step, his left leg supporting him lower on the incline, and his relaxed right hand hanging beside the quartz tool.\n\nThe only other living identity is the distant prism-winged cave ray high in the upper middle-right, a tiny graceful flyer with a dark tapered body, broad translucent blue-violet mineral fins held in a level horizontal posture, and a fine trailing tail.\n\nAn open calcite path rises from the gritty foreground through center and climbs toward screen-right, its irregular surface formed from worn limestone, damp mineral crust, embedded pebbles, and scattered translucent shards. Dark granular shelves border both sides of the route. Tall smoky-quartz columns emerge sparsely from these shelves, including a broad weathered pillar directly beyond the wayfarer and slender blade-like crystals along the right edge. The ground falls into a recessed shadowed basin at screen-left, where small amber crystal growths catch warm reflected light beneath a heavy black limestone overhang. Farther away, massive blue-veined cavern walls and restrained clusters of sapphire crystal recede through cool haze. A small faceted crystal sanctuary stands embedded in distant rock left of center, partially screened by columns but marked by a pale luminous doorway. Sapphire light pours through fissures in the immense vault and crosses warmer amber reflections on the path, revealing damp stone, mineral staining, airborne grit, natural cloth fibers, skin texture, and imperfect metal. The atmosphere remains that of a physically real crystal-lit cavern captured with restrained cinematic exposure, selective focus, subtle grain, slight motion softness, and mild optical aberration. The medium-wide camera begins near waist height behind-left of the wayfarer, preserving the opening uphill composition, then tracks his advance with a gentle rightward pan and holds its ground as he turns, changing naturally from a trailing view to a frontal view."
|
| 21 |
+
},
|
| 22 |
+
"White stag \u00b7 moonlit dirt road": {
|
| 23 |
+
"id": 31,
|
| 24 |
+
"chunks": 24,
|
| 25 |
+
"caption": "Appearance and scene:\n\nThe protagonist is a large adult white stag with a graceful, athletic cervine build. Its long back, deep chest, narrow waist, powerful haunches, slim lower legs, and dark cloven hooves form a physically credible running body. Dense ivory winter fur covers its entire body, with longer uneven guard hairs along the spine, shoulders, rump, and neck. A short pale tail lies above the rounded haunches. Its elongated head has a charcoal nose, dark eyes, alert pointed ears, and slender branching antlers. Faint strands of old black mourning ribbon hang from several antler tines and trail lightly during motion. The dense ivory winter fur, charcoal nose, alert ears, slender antlers, and old black mourning ribbons are stable identity features. Cold moonlight gives the raised fur a restrained pearl sheen while leaving soft blue-gray shadows between the legs and beneath the torso.\n\nAt the opening instant, the large white stag occupies the lower center foreground, seen directly from behind and oriented away from the camera toward the upper-right bend. It is already mid-bound on the compacted silver-gray chalk road, with its torso lifted, one foreleg flexed beneath the chest, its hind legs divided between push-off and recovery, and pale dust scattering immediately below the rear hooves. Its antlers rise clearly above the head, and the narrow ribbons hang from the tines without obscuring their shape.\n\nThe broad chalk road fills the foreground with two worn tracks, scattered pale stones, shallow ruts, and darker strips of compressed dust. A lower branch crosses toward screen-left while the unobstructed main track climbs gently toward screen-right between rough grass verges. Ancient wind-shaped yews gather beside the upper-right bend, while isolated pale birches and low shrubs punctuate the open rolling hills. Abandoned stone keeper lodges and roofless conservatory remnants stand across the middle distance, partially surrounded by low ground mist. A delicate weathered wrought-iron aqueduct extends along the far screen-left hillside. The bright full moon hangs high at screen-left in a deep indigo sky, casting physically plausible cool illumination, long soft-edged shadows, and restrained highlights across fur, chalk, bark, grass, stone, and iron. The air remains clear above the hollows, with subtle high-sensitivity grain and natural motion softness around the fastest hooves and ribbon tips. The documentary wildlife camera begins in a wide rear-facing composition near the stag\u2019s shoulder height, keeping the animal low and central while preserving the moon, road junction, yew-lined bend, ruined lodges, aqueduct, and open hills."
|
| 26 |
+
},
|
| 27 |
+
"Route surveyor \u00b7 rainy desert highway": {
|
| 28 |
+
"id": 36,
|
| 29 |
+
"chunks": 24,
|
| 30 |
+
"caption": "Appearance and scene:\n\nThe protagonist is a solitary female route surveyor with a sturdy, practical build. She stands left of center with her back to the camera, facing away along the wet road. Her dark brown hair is gathered into a loose high bun, with several damp strands resting around the nape of her neck; this tied hair, her compact build, and her field-ready silhouette are stable identity features. She wears a charcoal weather shell with a raised folded hood, long sleeves, dark gloves, and matching weatherproof cargo trousers. The trousers and her lower legs continue below the bottom edge, leaving her practical boots, feet, and ground contact out of view at the opening. A pale ceramic field pack covers most of her back, its scuffed hard shell shaped into broad upper shoulders and a rounded lower compartment. Gray webbing passes over both shoulders, narrow seams divide the pack into fitted panels, and a horizontal zippered pocket crosses its lower face. Compact surveying equipment is secured along both sides: a short cylindrical canister hangs low on her right, while a slim upright instrument with a capped circular sensor rises beside her left shoulder. Moisture dulls the shell fabric and gives the ceramic pack a cold, softly reflective sheen.\n\nThe nearest machine is the low enclosed electric bog rover in the right midground, broadside across the far lanes and motionless near the junction. It has a low armored nose, a dark enclosed cabin with angular windows, a long charcoal utility body, and four large deeply treaded wheels under flared guards. Damp moss staining gathers around the lower doors, wheel housings, and sill panels, while rainwater glints across the roof and window edges. The low enclosed electric bog rover faces screen-left toward the curving road, while its rear remains screen-right near the guardrail. Its broad stance, enclosed cabin, four-wheel running structure, and clear front-to-rear organization remain constant. Only faint suspension settling and intermittent beads of water moving down its body disturb its stillness.\n\nThe other visible identities are ordered with the solitary female route surveyor nearest, the low enclosed electric bog rover in the middle distance, the squat six-wheeled maintenance crawler farther ahead near the road center, and the distant black triangular weather drone farthest overhead. The squat six-wheeled maintenance crawler has a compact moss-stained utility body, six low wheels, a dark rear panel, and two small upper work lamps; it is seen nearly end-on from behind and points away along the road. The distant black triangular weather drone is a tiny matte-black wedge with short swept edges, held high against the cloud ceiling. These identities retain their distinct scales and positions throughout the view.\n\nWet charcoal asphalt spreads from the foreground into a broad relay junction, with pale lane lines shining through shallow rainwater and curving left around the basalt plain. Irregular volcanic shoulders carry sparse emerald moss, silver bunchgrass, and black-barked shrubs twisted by wind. A dark rock rise fills the near screen-left verge. Basalt cable stanchions stand along the roads beneath taut lines, while exactly two inclined fog-harvesting vanes occupy the right distance. An elevated circular barometer array rises near the center horizon, and slender vertical-axis wind harvesters repeat along the far ridge. Steel-gray clouds hang over layered mountains, cold cyan haze fills the valleys, and a narrow trace of fading amber remains beneath the cloud bank. Fine mist softens the distance while the wet ground preserves crisp restrained reflections. The trailing camera begins near shoulder height behind the solitary female route surveyor, holding the opening layout with her cropped lower legs, the two rovers, the high weather drone, and the stable near-to-far ordering clearly visible."
|
| 31 |
+
}
|
| 32 |
+
}
|
examples/silence_audio.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:e54157313bc0cf6f7757785e94cf45a537916441c9fa0a083671eef8e38258f1
|
| 3 |
+
size 53072
|
materialize_hf_examples.py
DELETED
|
@@ -1,369 +0,0 @@
|
|
| 1 |
-
#!/usr/bin/env python3
|
| 2 |
-
"""Expand the published compact CWM examples into inference-ready inputs."""
|
| 3 |
-
|
| 4 |
-
from __future__ import annotations
|
| 5 |
-
|
| 6 |
-
import argparse
|
| 7 |
-
import hashlib
|
| 8 |
-
import json
|
| 9 |
-
import math
|
| 10 |
-
import os
|
| 11 |
-
from pathlib import Path
|
| 12 |
-
import shutil
|
| 13 |
-
import tempfile
|
| 14 |
-
from typing import Any, Iterable
|
| 15 |
-
|
| 16 |
-
import numpy as np
|
| 17 |
-
from PIL import Image
|
| 18 |
-
|
| 19 |
-
|
| 20 |
-
WIDTH = 336
|
| 21 |
-
HEIGHT = 192
|
| 22 |
-
WINDOW_FRAMES = 124
|
| 23 |
-
STRIDE_FRAMES = 90
|
| 24 |
-
NEAR_METRES = 0.3
|
| 25 |
-
FAR_METRES = 256.0
|
| 26 |
-
METRIC_DEPTH = "float32_le_metric_first_surface_metres"
|
| 27 |
-
FIXED_LOG_DEPTH = "uint16_le_fixed_log_gray16_codes"
|
| 28 |
-
SUPPORTED_DEPTH_FORMATS = {METRIC_DEPTH, FIXED_LOG_DEPTH}
|
| 29 |
-
RELEASE_MODELS = {
|
| 30 |
-
"audio_vae": "../../models/vae/minimax_h3_audio_vae_fp32.safetensors",
|
| 31 |
-
"dit": "../../models/diffusion_models/minimax_h3_ref2va_bf16.safetensors",
|
| 32 |
-
"lora": "../../models/lora/minimax-h3-cwm.safetensors",
|
| 33 |
-
"processor": "../../models/processors/Qwen3-VL-32B-Instruct",
|
| 34 |
-
"text_encoder": "../../models/text_encoders/qwen3vl_32b_minimax_h3_bf16.safetensors",
|
| 35 |
-
"video_vae": "../../models/vae/minimax_h3_video_vae_fp16.safetensors",
|
| 36 |
-
}
|
| 37 |
-
RELEASE_SEEDS = {
|
| 38 |
-
1: 9942001,
|
| 39 |
-
2: 9800002,
|
| 40 |
-
3: 9700005,
|
| 41 |
-
4: 9700006,
|
| 42 |
-
5: 9700007,
|
| 43 |
-
6: 9700013,
|
| 44 |
-
7: 9700020,
|
| 45 |
-
8: 9700021,
|
| 46 |
-
9: 9700030,
|
| 47 |
-
10: 9700036,
|
| 48 |
-
11: 9700037,
|
| 49 |
-
12: 9700040,
|
| 50 |
-
13: 9700044,
|
| 51 |
-
14: 9700045,
|
| 52 |
-
15: 9700046,
|
| 53 |
-
16: 9700047,
|
| 54 |
-
17: 9700051,
|
| 55 |
-
18: 9700054,
|
| 56 |
-
19: 9700055,
|
| 57 |
-
20: 9700057,
|
| 58 |
-
21: 9700063,
|
| 59 |
-
22: 9700065,
|
| 60 |
-
23: 9700069,
|
| 61 |
-
24: 9700071,
|
| 62 |
-
25: 9700078,
|
| 63 |
-
26: 9700088,
|
| 64 |
-
27: 9700094,
|
| 65 |
-
28: 9600014,
|
| 66 |
-
29: 9600020,
|
| 67 |
-
30: 9600023,
|
| 68 |
-
31: 9600028,
|
| 69 |
-
32: 9600040,
|
| 70 |
-
33: 9600042,
|
| 71 |
-
34: 9300038,
|
| 72 |
-
35: 9400059,
|
| 73 |
-
36: 9300066,
|
| 74 |
-
37: 9300069,
|
| 75 |
-
38: 9800006,
|
| 76 |
-
39: 9800010,
|
| 77 |
-
40: 9800013,
|
| 78 |
-
}
|
| 79 |
-
|
| 80 |
-
|
| 81 |
-
def require(condition: bool, message: str) -> None:
|
| 82 |
-
if not condition:
|
| 83 |
-
raise ValueError(message)
|
| 84 |
-
|
| 85 |
-
|
| 86 |
-
def _read_json(path: Path) -> Any:
|
| 87 |
-
return json.loads(path.read_text(encoding="utf-8"))
|
| 88 |
-
|
| 89 |
-
|
| 90 |
-
def _canonical_prompt(value: str) -> str:
|
| 91 |
-
return value.strip().replace("\r\n", "\n").replace("\r", "\n")
|
| 92 |
-
|
| 93 |
-
|
| 94 |
-
def build_release_config(
|
| 95 |
-
example_root: Path, metadata: dict[str, Any], example_id: int
|
| 96 |
-
) -> dict[str, Any]:
|
| 97 |
-
windows = []
|
| 98 |
-
for prompt_item in metadata["prompts"]:
|
| 99 |
-
prompt = _canonical_prompt(
|
| 100 |
-
(example_root / prompt_item["file"]).read_text(encoding="utf-8")
|
| 101 |
-
)
|
| 102 |
-
# Preserve the text representation used by the released configurations.
|
| 103 |
-
if example_id != 2:
|
| 104 |
-
prompt = prompt.replace("\n", "\r\n")
|
| 105 |
-
windows.append(
|
| 106 |
-
{
|
| 107 |
-
"prompt": prompt,
|
| 108 |
-
"start_frame": prompt_item["conditioning_frame_start"],
|
| 109 |
-
}
|
| 110 |
-
)
|
| 111 |
-
name = f"example_{example_id:02d}"
|
| 112 |
-
return {
|
| 113 |
-
"anchor_image": "first_frame.png",
|
| 114 |
-
"attention_mode": "flash3",
|
| 115 |
-
"blocks_to_swap": 48,
|
| 116 |
-
"cache_dir": f"../../outputs/{name}/cache",
|
| 117 |
-
"condition_root": "conditions",
|
| 118 |
-
"expected_lora_modules": 200,
|
| 119 |
-
"models": RELEASE_MODELS,
|
| 120 |
-
"output": f"../../outputs/{name}/output.mp4",
|
| 121 |
-
"seed": RELEASE_SEEDS[example_id],
|
| 122 |
-
"steps": 20,
|
| 123 |
-
"windows": windows,
|
| 124 |
-
}
|
| 125 |
-
|
| 126 |
-
|
| 127 |
-
def ensure_release_config(
|
| 128 |
-
config_path: Path, example_root: Path, metadata: dict[str, Any], example_id: int
|
| 129 |
-
) -> None:
|
| 130 |
-
if not config_path.exists():
|
| 131 |
-
payload = json.dumps(
|
| 132 |
-
build_release_config(example_root, metadata, example_id),
|
| 133 |
-
ensure_ascii=False,
|
| 134 |
-
indent=2,
|
| 135 |
-
sort_keys=True,
|
| 136 |
-
)
|
| 137 |
-
_atomic_bytes(config_path, (payload + "\n").encode("utf-8"))
|
| 138 |
-
validate_config(config_path, example_root, metadata, example_id)
|
| 139 |
-
|
| 140 |
-
|
| 141 |
-
def fixed_log_codes_to_metric(codes: np.ndarray) -> np.ndarray:
|
| 142 |
-
"""Choose float32 depths that round-trip to every published uint16 code."""
|
| 143 |
-
require(codes.dtype == np.dtype("uint16"), f"fixed-log depth must be uint16, got {codes.dtype}")
|
| 144 |
-
metric = np.zeros(codes.shape, dtype="<f4")
|
| 145 |
-
valid = codes > 0
|
| 146 |
-
if bool(np.any(valid)):
|
| 147 |
-
log_near = math.log(NEAR_METRES)
|
| 148 |
-
log_far = math.log(FAR_METRES)
|
| 149 |
-
normalized = codes[valid].astype(np.float64) / 65535.0
|
| 150 |
-
metric[valid] = np.exp(log_far - normalized * (log_far - log_near)).astype(np.float32)
|
| 151 |
-
return np.ascontiguousarray(metric)
|
| 152 |
-
|
| 153 |
-
|
| 154 |
-
def metric_to_fixed_log_codes(metric: np.ndarray) -> np.ndarray:
|
| 155 |
-
"""Reference implementation matching the inference loader."""
|
| 156 |
-
metric = np.asarray(metric, dtype=np.float32)
|
| 157 |
-
valid = metric > np.float32(0.001)
|
| 158 |
-
codes = np.zeros(metric.shape, dtype=np.uint16)
|
| 159 |
-
if bool(np.any(valid)):
|
| 160 |
-
clipped = np.clip(metric[valid].astype(np.float64), NEAR_METRES, FAR_METRES)
|
| 161 |
-
normalized = (math.log(FAR_METRES) - np.log(clipped)) / (
|
| 162 |
-
math.log(FAR_METRES) - math.log(NEAR_METRES)
|
| 163 |
-
)
|
| 164 |
-
codes[valid] = np.floor(np.clip(normalized, 0.0, 1.0) * 65535.0 + 0.5).astype(
|
| 165 |
-
np.uint16
|
| 166 |
-
)
|
| 167 |
-
return codes
|
| 168 |
-
|
| 169 |
-
|
| 170 |
-
def _atomic_bytes(path: Path, payload: bytes) -> None:
|
| 171 |
-
path.parent.mkdir(parents=True, exist_ok=True)
|
| 172 |
-
if path.exists():
|
| 173 |
-
require(path.read_bytes() == payload, f"existing file differs: {path}")
|
| 174 |
-
return
|
| 175 |
-
descriptor, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
|
| 176 |
-
temporary = Path(temporary_name)
|
| 177 |
-
try:
|
| 178 |
-
with os.fdopen(descriptor, "wb") as handle:
|
| 179 |
-
handle.write(payload)
|
| 180 |
-
handle.flush()
|
| 181 |
-
os.fsync(handle.fileno())
|
| 182 |
-
temporary.replace(path)
|
| 183 |
-
finally:
|
| 184 |
-
temporary.unlink(missing_ok=True)
|
| 185 |
-
|
| 186 |
-
|
| 187 |
-
def _atomic_semantic_png(path: Path, semantic_id: np.ndarray) -> None:
|
| 188 |
-
semantic_id = np.ascontiguousarray(semantic_id, dtype=np.uint8)
|
| 189 |
-
if path.exists():
|
| 190 |
-
with Image.open(path) as image:
|
| 191 |
-
existing = np.asarray(image.convert("L"), dtype=np.uint8)
|
| 192 |
-
require(np.array_equal(existing, semantic_id), f"existing Semantic-ID differs: {path}")
|
| 193 |
-
return
|
| 194 |
-
path.parent.mkdir(parents=True, exist_ok=True)
|
| 195 |
-
descriptor, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".png", dir=path.parent)
|
| 196 |
-
os.close(descriptor)
|
| 197 |
-
temporary = Path(temporary_name)
|
| 198 |
-
try:
|
| 199 |
-
Image.fromarray(semantic_id, mode="L").save(temporary, format="PNG", compress_level=6)
|
| 200 |
-
temporary.replace(path)
|
| 201 |
-
finally:
|
| 202 |
-
temporary.unlink(missing_ok=True)
|
| 203 |
-
|
| 204 |
-
|
| 205 |
-
def _copy_exact(source: Path, destination: Path) -> None:
|
| 206 |
-
require(source.is_file(), f"missing source file: {source}")
|
| 207 |
-
payload = source.read_bytes()
|
| 208 |
-
_atomic_bytes(destination, payload)
|
| 209 |
-
|
| 210 |
-
|
| 211 |
-
def validate_metadata(example_root: Path, metadata: dict[str, Any], example_id: int) -> None:
|
| 212 |
-
require(metadata.get("id") == example_id, f"example {example_id}: metadata id mismatch")
|
| 213 |
-
require(metadata.get("fps") == 24, f"example {example_id}: fps must be 24")
|
| 214 |
-
require(metadata.get("resolution") == [WIDTH, HEIGHT], f"example {example_id}: bad resolution")
|
| 215 |
-
require(metadata.get("depth_format") in SUPPORTED_DEPTH_FORMATS, f"example {example_id}: unsupported depth format")
|
| 216 |
-
frames = metadata.get("frames")
|
| 217 |
-
require(type(frames) is int and frames >= WINDOW_FRAMES, f"example {example_id}: invalid frame count")
|
| 218 |
-
prompts = metadata.get("prompts")
|
| 219 |
-
require(isinstance(prompts, list) and prompts, f"example {example_id}: prompts missing")
|
| 220 |
-
require(frames == WINDOW_FRAMES + STRIDE_FRAMES * (len(prompts) - 1), f"example {example_id}: frame/window mismatch")
|
| 221 |
-
for index, item in enumerate(prompts):
|
| 222 |
-
require(item.get("index") == index, f"example {example_id}: prompt index mismatch")
|
| 223 |
-
require(item.get("conditioning_frame_start") == index * STRIDE_FRAMES, f"example {example_id}: prompt start mismatch")
|
| 224 |
-
require(item.get("conditioning_frame_end_exclusive") == index * STRIDE_FRAMES + WINDOW_FRAMES, f"example {example_id}: prompt end mismatch")
|
| 225 |
-
prompt_path = example_root / item.get("file", "")
|
| 226 |
-
require(prompt_path.is_file(), f"example {example_id}: missing prompt {prompt_path}")
|
| 227 |
-
chunks = metadata.get("condition_chunks")
|
| 228 |
-
require(isinstance(chunks, list) and chunks, f"example {example_id}: condition chunks missing")
|
| 229 |
-
cursor = 0
|
| 230 |
-
for item in chunks:
|
| 231 |
-
start = item.get("start_frame")
|
| 232 |
-
end = item.get("end_frame_exclusive")
|
| 233 |
-
require(start == cursor and type(end) is int and end > start, f"example {example_id}: non-contiguous chunks")
|
| 234 |
-
require(item.get("frames") == end - start, f"example {example_id}: chunk frame count mismatch")
|
| 235 |
-
chunk_path = example_root / item.get("file", "")
|
| 236 |
-
require(chunk_path.is_file(), f"example {example_id}: missing chunk {chunk_path}")
|
| 237 |
-
cursor = end
|
| 238 |
-
require(cursor == frames, f"example {example_id}: chunks do not cover all frames")
|
| 239 |
-
|
| 240 |
-
|
| 241 |
-
def validate_config(config_path: Path, example_root: Path, metadata: dict[str, Any], example_id: int) -> None:
|
| 242 |
-
require(config_path.is_file(), f"missing released config: {config_path}")
|
| 243 |
-
config = _read_json(config_path)
|
| 244 |
-
require(config.get("anchor_image") == "first_frame.png", f"example {example_id}: config anchor path mismatch")
|
| 245 |
-
require(config.get("condition_root") == "conditions", f"example {example_id}: config condition path mismatch")
|
| 246 |
-
require(config.get("seed") == RELEASE_SEEDS[example_id], f"example {example_id}: released seed mismatch")
|
| 247 |
-
require(config.get("steps") == 20, f"example {example_id}: steps mismatch")
|
| 248 |
-
require(config.get("attention_mode") == "flash3", f"example {example_id}: attention mode mismatch")
|
| 249 |
-
require(config.get("blocks_to_swap") == 48, f"example {example_id}: block swap count mismatch")
|
| 250 |
-
require(config.get("expected_lora_modules") == 200, f"example {example_id}: LoRA module count mismatch")
|
| 251 |
-
require(config.get("models") == RELEASE_MODELS, f"example {example_id}: model paths mismatch")
|
| 252 |
-
name = f"example_{example_id:02d}"
|
| 253 |
-
require(config.get("cache_dir") == f"../../outputs/{name}/cache", f"example {example_id}: cache path mismatch")
|
| 254 |
-
require(config.get("output") == f"../../outputs/{name}/output.mp4", f"example {example_id}: output path mismatch")
|
| 255 |
-
windows = config.get("windows")
|
| 256 |
-
prompt_items = metadata["prompts"]
|
| 257 |
-
require(isinstance(windows, list) and len(windows) == len(prompt_items), f"example {example_id}: config window count mismatch")
|
| 258 |
-
for window, prompt_item in zip(windows, prompt_items, strict=True):
|
| 259 |
-
require(window.get("start_frame") == prompt_item["conditioning_frame_start"], f"example {example_id}: config window start mismatch")
|
| 260 |
-
published = (example_root / prompt_item["file"]).read_text(encoding="utf-8")
|
| 261 |
-
require(_canonical_prompt(window.get("prompt", "")) == _canonical_prompt(published), f"example {example_id}: config prompt differs from published prompt")
|
| 262 |
-
|
| 263 |
-
|
| 264 |
-
def _iter_chunks(example_root: Path, metadata: dict[str, Any]) -> Iterable[tuple[int, np.ndarray, np.ndarray]]:
|
| 265 |
-
depth_format = metadata["depth_format"]
|
| 266 |
-
for chunk in metadata["condition_chunks"]:
|
| 267 |
-
path = example_root / chunk["file"]
|
| 268 |
-
with np.load(path, allow_pickle=False) as payload:
|
| 269 |
-
require(set(payload.files) == {"depth", "semantic_id"}, f"unexpected arrays in {path}")
|
| 270 |
-
depth = payload["depth"]
|
| 271 |
-
semantic_id = payload["semantic_id"]
|
| 272 |
-
expected_shape = (chunk["frames"], HEIGHT, WIDTH)
|
| 273 |
-
require(depth.shape == expected_shape, f"bad depth shape in {path}: {depth.shape}")
|
| 274 |
-
require(semantic_id.shape == expected_shape, f"bad Semantic-ID shape in {path}: {semantic_id.shape}")
|
| 275 |
-
require(semantic_id.dtype == np.dtype("uint8"), f"bad Semantic-ID dtype in {path}: {semantic_id.dtype}")
|
| 276 |
-
require(int(semantic_id.max(initial=0)) <= 11, f"Semantic-ID outside [0,11] in {path}")
|
| 277 |
-
if depth_format == METRIC_DEPTH:
|
| 278 |
-
require(depth.dtype == np.dtype("float32"), f"bad metric depth dtype in {path}: {depth.dtype}")
|
| 279 |
-
require(bool(np.all(np.isfinite(depth))) and not bool(np.any(depth < 0)), f"invalid metric depth in {path}")
|
| 280 |
-
metric = np.ascontiguousarray(depth, dtype="<f4")
|
| 281 |
-
else:
|
| 282 |
-
require(depth.dtype == np.dtype("uint16"), f"bad fixed-log depth dtype in {path}: {depth.dtype}")
|
| 283 |
-
metric = fixed_log_codes_to_metric(depth)
|
| 284 |
-
require(np.array_equal(metric_to_fixed_log_codes(metric), depth), f"fixed-log round-trip failed in {path}")
|
| 285 |
-
start = chunk["start_frame"]
|
| 286 |
-
for offset in range(chunk["frames"]):
|
| 287 |
-
yield start + offset, metric[offset], semantic_id[offset]
|
| 288 |
-
|
| 289 |
-
|
| 290 |
-
def materialize_example(
|
| 291 |
-
source_root: Path,
|
| 292 |
-
output_root: Path,
|
| 293 |
-
example_id: int,
|
| 294 |
-
*,
|
| 295 |
-
require_config: bool,
|
| 296 |
-
frame_limit: int | None = None,
|
| 297 |
-
) -> dict[str, Any]:
|
| 298 |
-
example_root = source_root / str(example_id)
|
| 299 |
-
metadata_path = example_root / "metadata.json"
|
| 300 |
-
require(metadata_path.is_file(), f"missing metadata: {metadata_path}")
|
| 301 |
-
metadata = _read_json(metadata_path)
|
| 302 |
-
validate_metadata(example_root, metadata, example_id)
|
| 303 |
-
destination = output_root / "examples" / f"example_{example_id:02d}"
|
| 304 |
-
config_path = destination / "config.json"
|
| 305 |
-
if require_config:
|
| 306 |
-
ensure_release_config(config_path, example_root, metadata, example_id)
|
| 307 |
-
_copy_exact(example_root / metadata["first_frame"], destination / "first_frame.png")
|
| 308 |
-
written = 0
|
| 309 |
-
for ordinal, depth, semantic_id in _iter_chunks(example_root, metadata):
|
| 310 |
-
if frame_limit is not None and ordinal >= frame_limit:
|
| 311 |
-
break
|
| 312 |
-
condition_root = destination / "conditions"
|
| 313 |
-
_atomic_bytes(condition_root / f"{ordinal:06d}.depth.f32", np.ascontiguousarray(depth, dtype="<f4").tobytes(order="C"))
|
| 314 |
-
_atomic_semantic_png(condition_root / f"{ordinal:06d}.semantic_id.png", semantic_id)
|
| 315 |
-
written += 1
|
| 316 |
-
expected = min(metadata["frames"], frame_limit) if frame_limit is not None else metadata["frames"]
|
| 317 |
-
require(written == expected, f"example {example_id}: materialized {written}, expected {expected}")
|
| 318 |
-
return {
|
| 319 |
-
"id": example_id,
|
| 320 |
-
"depth_format": metadata["depth_format"],
|
| 321 |
-
"frames": metadata["frames"],
|
| 322 |
-
"materialized_frames": written,
|
| 323 |
-
"windows": len(metadata["prompts"]),
|
| 324 |
-
}
|
| 325 |
-
|
| 326 |
-
|
| 327 |
-
def parse_ids(values: list[int] | None) -> list[int]:
|
| 328 |
-
ids = list(range(1, 41)) if not values else sorted(set(values))
|
| 329 |
-
require(ids and all(1 <= value <= 40 for value in ids), "example ids must be in [1,40]")
|
| 330 |
-
return ids
|
| 331 |
-
|
| 332 |
-
|
| 333 |
-
def main(argv: list[str] | None = None) -> int:
|
| 334 |
-
parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False)
|
| 335 |
-
parser.add_argument("--source", required=True, type=Path, help="downloaded code-world-model-inference-examples-40 directory")
|
| 336 |
-
parser.add_argument("--output", required=True, type=Path, help="release root containing models/, examples/, and outputs/")
|
| 337 |
-
parser.add_argument("--example", action="append", type=int, dest="examples", help="materialize one public example id; repeat as needed")
|
| 338 |
-
parser.add_argument("--without-config-check", action="store_true", help=argparse.SUPPRESS)
|
| 339 |
-
parser.add_argument("--frame-limit", type=int, help=argparse.SUPPRESS)
|
| 340 |
-
args = parser.parse_args(argv)
|
| 341 |
-
source = args.source.resolve()
|
| 342 |
-
output = args.output.resolve()
|
| 343 |
-
manifest = _read_json(source / "manifest.json")
|
| 344 |
-
require(manifest.get("schema") == "code_world_model_inference_examples_v3", "unsupported dataset manifest")
|
| 345 |
-
require(manifest.get("count") == 40 and manifest.get("expected_count") == 40, "dataset manifest does not contain 40 examples")
|
| 346 |
-
results = []
|
| 347 |
-
for example_id in parse_ids(args.examples):
|
| 348 |
-
result = materialize_example(
|
| 349 |
-
source,
|
| 350 |
-
output,
|
| 351 |
-
example_id,
|
| 352 |
-
require_config=not args.without_config_check,
|
| 353 |
-
frame_limit=args.frame_limit,
|
| 354 |
-
)
|
| 355 |
-
results.append(result)
|
| 356 |
-
print(f"example_{example_id:02d}: {result['materialized_frames']} frames", flush=True)
|
| 357 |
-
receipt = {
|
| 358 |
-
"schema": "cwm_h3_materialized_examples_v1",
|
| 359 |
-
"source_manifest_sha256": hashlib.sha256((source / "manifest.json").read_bytes()).hexdigest(),
|
| 360 |
-
"examples": results,
|
| 361 |
-
}
|
| 362 |
-
if args.frame_limit is None and len(results) == 40:
|
| 363 |
-
path = output / "examples" / "materialization.json"
|
| 364 |
-
_atomic_bytes(path, (json.dumps(receipt, indent=2, sort_keys=True) + "\n").encode("utf-8"))
|
| 365 |
-
return 0
|
| 366 |
-
|
| 367 |
-
|
| 368 |
-
if __name__ == "__main__":
|
| 369 |
-
raise SystemExit(main())
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
requirements.txt
CHANGED
|
@@ -1,8 +1,4 @@
|
|
| 1 |
-
|
| 2 |
-
torchvision==0.24.1
|
| 3 |
-
transformers==4.57.6
|
| 4 |
-
accelerate==1.6.0
|
| 5 |
-
huggingface-hub==0.34.3
|
| 6 |
safetensors
|
| 7 |
av
|
| 8 |
numpy
|
|
|
|
| 1 |
+
accelerate
|
|
|
|
|
|
|
|
|
|
|
|
|
| 2 |
safetensors
|
| 3 |
av
|
| 4 |
numpy
|
streaming_runtime.py
ADDED
|
@@ -0,0 +1,187 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Runtime adaptations of the rtcwm streaming sampler for ZeroGPU (RTX PRO 6000, sm_120). See app.py."""
|
| 2 |
+
import time
|
| 3 |
+
|
| 4 |
+
import torch
|
| 5 |
+
import torch.nn.functional as F
|
| 6 |
+
from torch import nn
|
| 7 |
+
|
| 8 |
+
from musubi_tuner.minimax_h3 import model as h3model
|
| 9 |
+
from musubi_tuner.minimax_h3.packing import _build_modulation_segments
|
| 10 |
+
from rtcwm.model import attention as rt_attention
|
| 11 |
+
from rtcwm.model import packing as rt_packing
|
| 12 |
+
from rtcwm.model import sampler as rt_sampler
|
| 13 |
+
|
| 14 |
+
STEPS, SHIFT = 4, 5.0
|
| 15 |
+
|
| 16 |
+
class FP8Linear(nn.Module):
|
| 17 |
+
"""rtcwm.model.fp8.FP8Linear (per-output-channel e4m3 weights, per-row e4m3 activations, torch._scaled_mm),
|
| 18 |
+
with an emulated fallback should the rowwise-scaled fp8 GEMM be unavailable on this GPU."""
|
| 19 |
+
use_scaled_mm = True
|
| 20 |
+
|
| 21 |
+
def __init__(self, weight_bf16: torch.Tensor):
|
| 22 |
+
super().__init__()
|
| 23 |
+
w = weight_bf16.float()
|
| 24 |
+
scale = (w.abs().amax(dim=1, keepdim=True) / 448.0).clamp_min(1e-12)
|
| 25 |
+
self.register_buffer("weight_fp8", (w / scale).to(torch.float8_e4m3fn))
|
| 26 |
+
self.register_buffer("weight_scale", scale.t().contiguous())
|
| 27 |
+
self.out_features = w.shape[0]
|
| 28 |
+
|
| 29 |
+
def forward(self, x):
|
| 30 |
+
shape = x.shape
|
| 31 |
+
x2 = x.reshape(-1, shape[-1])
|
| 32 |
+
scale = (x2.abs().amax(dim=1, keepdim=True).float() / 448.0).clamp_min(1e-12)
|
| 33 |
+
xq = (x2.float() / scale).to(torch.float8_e4m3fn)
|
| 34 |
+
if FP8Linear.use_scaled_mm:
|
| 35 |
+
try:
|
| 36 |
+
y = torch._scaled_mm(xq, self.weight_fp8.t(), scale_a=scale, scale_b=self.weight_scale,
|
| 37 |
+
out_dtype=torch.bfloat16)
|
| 38 |
+
return y.reshape(*shape[:-1], self.out_features)
|
| 39 |
+
except Exception as e: # noqa: BLE001
|
| 40 |
+
print(f"_scaled_mm unavailable ({e!r}); using the emulated FP8 GEMM", flush=True)
|
| 41 |
+
FP8Linear.use_scaled_mm = False
|
| 42 |
+
y = (xq.to(torch.bfloat16) @ self.weight_fp8.to(torch.bfloat16).t()).float() * scale * self.weight_scale
|
| 43 |
+
return y.to(torch.bfloat16).reshape(*shape[:-1], self.out_features)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
_ADALN_T = [None] # the unique times of the forward in flight (set by run_blocks)
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def _tkey(t: float) -> int:
|
| 50 |
+
return int(round(float(t) * 1_000_000))
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
class CachedAdaln(nn.Module):
|
| 54 |
+
"""A DiT block's AdalnProj evaluated once per known time: same output as
|
| 55 |
+
linear(silu(time_embedder(t).bf16)), looked up by t."""
|
| 56 |
+
|
| 57 |
+
def __init__(self, table: torch.Tensor, keys: list, modalities: int, expand: int, hidden: int):
|
| 58 |
+
super().__init__()
|
| 59 |
+
self.register_buffer("table", table) # [n_t, expand*hidden*modalities] bf16
|
| 60 |
+
self.index = {k: i for i, k in enumerate(keys)}
|
| 61 |
+
self.modalities, self.expand, self.hidden = modalities, expand, hidden
|
| 62 |
+
|
| 63 |
+
def forward(self, _temb):
|
| 64 |
+
rows = [self.index[_tkey(t)] for t in _ADALN_T[0]]
|
| 65 |
+
p = self.table[rows]
|
| 66 |
+
p = p.reshape(p.shape[0] * self.modalities, self.expand * self.hidden)
|
| 67 |
+
return p.chunk(self.expand, dim=-1)
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
def sampling_times() -> list:
|
| 71 |
+
"""Every model time a 4-step x0 stream evaluates: condition rows (1.0 text / audio / clean history, 0.999
|
| 72 |
+
Picture 1 + proxy) and 1 - sigma of each step, computed like the sampler does (float32 rows)."""
|
| 73 |
+
sig = rt_sampler.schedule(STEPS, "cpu", SHIFT)
|
| 74 |
+
ts = [1.0, 0.999] + [1.0 - float(sig[i]) for i in range(STEPS)]
|
| 75 |
+
return sorted({float(torch.tensor(t, dtype=torch.float32)) for t in ts})
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def prepare_transformer(tr, lora: dict, fp8: bool = True, cache_adaln: bool = True):
|
| 79 |
+
"""Merge the streaming LoRA into the blocks' linears (musubi merge_to, multiplier 1, bf16), put them in FP8
|
| 80 |
+
(rtcwm.model.fp8), precompute the blocks' AdaLN outputs for the sampling times, and install the rtcwm
|
| 81 |
+
attention (with the SDPA stand-ins)."""
|
| 82 |
+
t0 = time.time()
|
| 83 |
+
tr.requires_grad_(False)
|
| 84 |
+
used = set()
|
| 85 |
+
times = sampling_times()
|
| 86 |
+
with torch.no_grad():
|
| 87 |
+
temb = tr.time_embedder(torch.tensor(times, dtype=torch.float32)).to(torch.bfloat16)
|
| 88 |
+
act = F.silu(temb).float()
|
| 89 |
+
keys = [_tkey(t) for t in times]
|
| 90 |
+
for i, blk in enumerate(tr.blocks):
|
| 91 |
+
# LoRA merge (musubi LoRAModule.merge_to, multiplier 1, bf16) + FP8 (rtcwm.model.fp8.quantize_blocks)
|
| 92 |
+
for parent, prefix, names in ((blk.attn, "attn", ("qkv_proj", "out_proj")), (blk.mlp, "mlp", ("fc1", "fc2"))):
|
| 93 |
+
for n in names:
|
| 94 |
+
lin = getattr(parent, n)
|
| 95 |
+
assert lin.bias is None
|
| 96 |
+
key = f"lora_unet_blocks_{i}_{prefix}_{n}"
|
| 97 |
+
down = lora[f"{key}.lora_down.weight"].float()
|
| 98 |
+
up = lora[f"{key}.lora_up.weight"].float()
|
| 99 |
+
scale = float(lora[f"{key}.alpha"].float()) / down.shape[0]
|
| 100 |
+
w = (lin.weight.float() + (up @ down) * scale).to(torch.bfloat16)
|
| 101 |
+
if fp8:
|
| 102 |
+
setattr(parent, n, FP8Linear(w))
|
| 103 |
+
else:
|
| 104 |
+
lin.weight.data = w
|
| 105 |
+
used.update({f"{key}.lora_down.weight", f"{key}.lora_up.weight", f"{key}.alpha"})
|
| 106 |
+
del w, lin
|
| 107 |
+
if not cache_adaln:
|
| 108 |
+
continue
|
| 109 |
+
ad = blk.adaln_proj
|
| 110 |
+
table = (act @ ad.linear.weight.float().t() + ad.linear.bias.float()).to(torch.bfloat16)
|
| 111 |
+
blk.adaln_proj = CachedAdaln(table, keys, ad.modalities, ad.expand, ad.hidden_size)
|
| 112 |
+
del ad
|
| 113 |
+
if i % 10 == 0:
|
| 114 |
+
print(f"block {i} prepared ({time.time() - t0:.0f}s)", flush=True)
|
| 115 |
+
assert used == set(lora), f"unused LoRA keys: {sorted(set(lora) - used)[:5]}"
|
| 116 |
+
print(f"merged {len(used) // 3} LoRA modules, FP8 blocks ready in {time.time() - t0:.0f}s", flush=True)
|
| 117 |
+
install(tr)
|
| 118 |
+
return tr.eval()
|
| 119 |
+
|
| 120 |
+
|
| 121 |
+
def install(tr):
|
| 122 |
+
rt_attention.install(tr)
|
| 123 |
+
h3model.Attention.forward = _attention_forward
|
| 124 |
+
|
| 125 |
+
|
| 126 |
+
# --------------------------------------------------------------------------------------------- attention
|
| 127 |
+
def _sdpa(q, k, v):
|
| 128 |
+
"""FlashAttention-3 stand-in: q, k, v [B,S,H,D] -> [B,S,H,D], unmasked."""
|
| 129 |
+
return F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2)
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
rt_attention._fa3 = _sdpa
|
| 133 |
+
|
| 134 |
+
|
| 135 |
+
class DenseMask:
|
| 136 |
+
"""rtcwm.model.masks.df_mask as prefix attention calls (rtcwm.model.attention.dense_df) -- for the condition
|
| 137 |
+
rows alone (n_chunks 0) or together with chunk 0 (n_chunks 1)."""
|
| 138 |
+
|
| 139 |
+
def __init__(self, head_rows, duv_rows, audio_rows, n_chunks, static_sees_first):
|
| 140 |
+
self.head_rows, self.duv_rows, self.audio_rows = head_rows, duv_rows, audio_rows
|
| 141 |
+
self.n_chunks, self.static_sees_first = n_chunks, static_sees_first
|
| 142 |
+
|
| 143 |
+
|
| 144 |
+
def _df_mask(head_rows, duv_rows, audio_rows, n_chunks, device, first_chunk=0, static_sees_first=True):
|
| 145 |
+
assert first_chunk == 0
|
| 146 |
+
return DenseMask(head_rows, duv_rows, audio_rows, n_chunks, static_sees_first)
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
rt_sampler.df_mask = _df_mask
|
| 150 |
+
_rtcwm_forward = rt_attention._forward
|
| 151 |
+
|
| 152 |
+
|
| 153 |
+
def _attention_forward(self, hidden_states, rotation_table=None):
|
| 154 |
+
mask, kv = rt_attention._state["mask"], rt_attention._state["kv"]
|
| 155 |
+
if not isinstance(mask, DenseMask):
|
| 156 |
+
return _rtcwm_forward(self, hidden_states, rotation_table)
|
| 157 |
+
b, s, _ = hidden_states.shape
|
| 158 |
+
q, k, v = self.qkv_proj(hidden_states).split(self.inner_dim, dim=-1)
|
| 159 |
+
q = self.q_norm(q.reshape(b, s, self.num_heads, self.head_dim))
|
| 160 |
+
k = self.k_norm(k.reshape(b, s, self.num_heads, self.head_dim))
|
| 161 |
+
v = v.reshape(b, s, self.num_heads, self.head_dim)
|
| 162 |
+
if rotation_table is not None:
|
| 163 |
+
q = rt_attention.rope(q, rotation_table)
|
| 164 |
+
k = rt_attention.rope(k, rotation_table)
|
| 165 |
+
q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)
|
| 166 |
+
if kv is not None and kv.get("store") is not None:
|
| 167 |
+
kv["store"].append((k.transpose(1, 2).contiguous(), v.transpose(1, 2).contiguous()))
|
| 168 |
+
out = rt_attention.dense_df(q, k, v, mask)
|
| 169 |
+
return self.out_proj(out.reshape(b, s, self.inner_dim))
|
| 170 |
+
|
| 171 |
+
|
| 172 |
+
def run_blocks(model, hidden, t_rows, tags, positions, *, mask=None, kv=None, train=False, offload=False):
|
| 173 |
+
"""rtcwm.model.packing.run_blocks (inference), publishing the forward's times to CachedAdaln."""
|
| 174 |
+
uniq, inv = torch.unique(t_rows, sorted=True, return_inverse=True)
|
| 175 |
+
segments = (_build_modulation_segments(3 * inv + tags),)
|
| 176 |
+
temb = model.time_embedder(uniq.to(hidden.device)).to(model.dtype)
|
| 177 |
+
rot = model._rotation_table(positions.unsqueeze(0).to(hidden.device), hidden.dtype)
|
| 178 |
+
_ADALN_T[0] = uniq.tolist()
|
| 179 |
+
for block in model.blocks:
|
| 180 |
+
with rt_attention.active(mask, kv):
|
| 181 |
+
hidden = block(hidden, temb, segments, rot)
|
| 182 |
+
return hidden, temb, inv
|
| 183 |
+
|
| 184 |
+
|
| 185 |
+
rt_packing.run_blocks = run_blocks
|
| 186 |
+
rt_sampler.run_blocks = run_blocks
|
| 187 |
+
|